Examples with JoinDTO io.hops.hopsworks.common.featurestore.query.join.JoinDTO used on opensource projects

Example 6 with JoinDTO

use of io.hops.hopsworks.common.featurestore.query.join.JoinDTO in project hopsworks by logicalclocks.

the class TestConstructorController method testExtractFeaturesBothSides.

@Test
public void testExtractFeaturesBothSides() throws Exception {
    Mockito.when(featuregroupController.getFeatures(Mockito.any(), Mockito.any(), Mockito.any())).thenReturn(fg1FeaturesDTO, fg2FeaturesDTO);
    Mockito.when(featuregroupFacade.findById(Mockito.any())).thenReturn(Optional.of(fg1), Optional.of(fg2));
    Mockito.when(featurestoreFacade.getHiveDbName(Mockito.any())).thenReturn("fg1", "fg2");
    FeaturegroupDTO fg1 = new FeaturegroupDTO();
    fg1.setId(1);
    FeaturegroupDTO fg2 = new FeaturegroupDTO();
    fg2.setId(2);
    List<FeatureGroupFeatureDTO> requestedFeatures = new ArrayList<>();
    requestedFeatures.add(new FeatureGroupFeatureDTO("*"));
    QueryDTO rightQueryDTO = new QueryDTO(fg2, requestedFeatures);
    JoinDTO joinDTO = new JoinDTO(rightQueryDTO, null, null);
    QueryDTO queryDTO = new QueryDTO(fg1, requestedFeatures, Arrays.asList(joinDTO));
    Map<Integer, String> fgAliasLookup = new HashMap<>();
    Map<Integer, Featuregroup> fgLookup = new HashMap<>();
    Map<Integer, List<Feature>> availableFeatureLookup = new HashMap<>();
    queryController.populateFgLookupTables(queryDTO, 1, fgAliasLookup, fgLookup, availableFeatureLookup, project, user, null);
    Query query = queryController.convertQueryDTO(queryDTO, fgAliasLookup, fgLookup, availableFeatureLookup, false);
    List<Feature> extractedFeatures = target.collectFeatures(query);
    // Make sure both features have been returned.
    // It's going to be 3 as the feature "pr" will be identified as primary key and joining key
    // so it's not going to be duplicated
    Assert.assertEquals(3, extractedFeatures.size());
    // Make sure the method sets the feature group name
    Assert.assertTrue(extractedFeatures.get(0).getFgAlias(false).equals("fg1") || extractedFeatures.get(0).getFgAlias(false).equals("fg2"));
}