OnlineServingService implements ServingService { @Override public GetOnlineFeaturesResponse getOnlineFeatures(GetOnlineFeaturesRequest request) { try (Scope scope = tracer.buildSpan("getOnlineFeatures").startActive(true)) { List<EntityRow> entityRows = request.getEntityRowsList(); Map<EntityRow, Map<String, Value>> entityValuesMap = entityRows.stream().collect(Collectors.toMap(row -> row, row -> new HashMap<>())); Map<EntityRow, Map<String, FieldStatus>> entityStatusesMap = entityRows.stream().collect(Collectors.toMap(row -> row, row -> new HashMap<>())); List<List<Optional<FeatureRow>>> logFeatureRows = new LinkedList<>(); if (!request.getOmitEntitiesInResponse()) { entityRows.forEach( entityRow -> { Map<String, Value> valueMap = entityRow.getFieldsMap(); entityValuesMap.get(entityRow).putAll(valueMap); entityStatusesMap.get(entityRow).putAll(getMetadataMap(valueMap, false, false)); }); } List<FeatureSetRequest> featureSetRequests = specService.getFeatureSets(request.getFeaturesList(), request.getProject()); for (FeatureSetRequest featureSetRequest : featureSetRequests) { List<Optional<FeatureRow>> featureRows = retriever.getOnlineFeatures(entityRows, featureSetRequest); if (featureRows.size() != entityRows.size()) { throw Status.INTERNAL .withDescription( "The no. of FeatureRow obtained from OnlineRetriever" + "does not match no. of entityRow passed.") .asRuntimeException(); } Streams.zip(entityRows.stream(), featureRows.stream(), Pair::of) .forEach( entityFeaturePair -> { EntityRow entityRow = entityFeaturePair.getLeft(); Optional<FeatureRow> featureRow = entityFeaturePair.getRight(); boolean isOutsideMaxAge = checkOutsideMaxAge(featureSetRequest, entityRow, featureRow); Map<String, Value> valueMap = unpackValueMap(featureRow, featureSetRequest, isOutsideMaxAge); entityValuesMap.get(entityRow).putAll(valueMap); boolean isNotFound = featureRow.isEmpty(); Map<String, FieldStatus> statusMap = getMetadataMap(valueMap, isNotFound, isOutsideMaxAge); entityStatusesMap.get(entityRow).putAll(statusMap); populateCountMetrics(statusMap, featureSetRequest); }); populateRequestCountMetrics(featureSetRequest); logFeatureRows.add(featureRows); } if (scope != null) { logFeatureRowsTrace(scope, logFeatureRows, featureSetRequests); } List<FieldValues> fieldValuesList = entityRows.stream() .map( entityRow -> { return FieldValues.newBuilder() .putAllFields(entityValuesMap.get(entityRow)) .putAllStatuses(entityStatusesMap.get(entityRow)) .build(); }) .collect(Collectors.toList()); return GetOnlineFeaturesResponse.newBuilder().addAllFieldValues(fieldValuesList).build(); } } OnlineServingService(
      OnlineRetriever retriever, CachedSpecService specService, Tracer tracer); @Override GetFeastServingInfoResponse getFeastServingInfo(
      GetFeastServingInfoRequest getFeastServingInfoRequest); @Override GetOnlineFeaturesResponse getOnlineFeatures(GetOnlineFeaturesRequest request); @Override GetBatchFeaturesResponse getBatchFeatures(GetBatchFeaturesRequest getFeaturesRequest); @Override GetJobResponse getJob(GetJobRequest getJobRequest);  }
@Test public void shouldReturnResponseWithValuesAndMetadataIfKeysPresent() { GetOnlineFeaturesRequest request = getOnlineFeaturesRequest( List.of( FeatureReference.newBuilder().setName("feature1").build(), FeatureReference.newBuilder().setName("feature2").setProject("project").build())); FeatureSetRequest featureSetRequest = FeatureSetRequest.newBuilder() .addAllFeatureReferences(request.getFeaturesList()) .setSpec(getFeatureSetSpec()) .build(); List<Optional<FeatureRow>> featureRows = List.of(Optional.of(testFeatureRows.get(0)), Optional.of(testFeatureRows.get(1))); when(specService.getFeatureSets(request.getFeaturesList(), "")) .thenReturn(Collections.singletonList(featureSetRequest)); when(retriever.getOnlineFeatures(request.getEntityRowsList(), featureSetRequest)) .thenReturn(featureRows); when(tracer.buildSpan(ArgumentMatchers.any())).thenReturn(Mockito.mock(SpanBuilder.class)); GetOnlineFeaturesResponse expected = GetOnlineFeaturesResponse.newBuilder() .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(1)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("a")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", intValue(1)) .putStatuses("feature1", FieldStatus.PRESENT) .putFields("project/feature2", intValue(1)) .putStatuses("project/feature2", FieldStatus.PRESENT) .build()) .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(2)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("b")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", intValue(2)) .putStatuses("feature1", FieldStatus.PRESENT) .putFields("project/feature2", intValue(2)) .putStatuses("project/feature2", FieldStatus.PRESENT) .build()) .build(); GetOnlineFeaturesResponse actual = onlineServingService.getOnlineFeatures(request); assertThat(actual, equalTo(expected)); }
@Test public void shouldReturnResponseWithUnsetValuesAndMetadataIfKeysNotPresent() { GetOnlineFeaturesRequest request = getOnlineFeaturesRequest( List.of( FeatureReference.newBuilder().setName("feature1").build(), FeatureReference.newBuilder().setName("feature2").setProject("project").build())); FeatureSetRequest featureSetRequest = FeatureSetRequest.newBuilder() .addAllFeatureReferences(request.getFeaturesList()) .setSpec(getFeatureSetSpec()) .build(); List<Optional<FeatureRow>> featureRows = List.of(Optional.of(testFeatureRows.get(0)), Optional.empty()); when(specService.getFeatureSets(request.getFeaturesList(), "")) .thenReturn(Collections.singletonList(featureSetRequest)); when(retriever.getOnlineFeatures(request.getEntityRowsList(), featureSetRequest)) .thenReturn(featureRows); when(tracer.buildSpan(ArgumentMatchers.any())).thenReturn(Mockito.mock(SpanBuilder.class)); GetOnlineFeaturesResponse expected = GetOnlineFeaturesResponse.newBuilder() .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(1)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("a")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", intValue(1)) .putStatuses("feature1", FieldStatus.PRESENT) .putFields("project/feature2", intValue(1)) .putStatuses("project/feature2", FieldStatus.PRESENT) .build()) .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(2)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("b")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", Value.newBuilder().build()) .putStatuses("feature1", FieldStatus.NOT_FOUND) .putFields("project/feature2", Value.newBuilder().build()) .putStatuses("project/feature2", FieldStatus.NOT_FOUND) .build()) .build(); GetOnlineFeaturesResponse actual = onlineServingService.getOnlineFeatures(request); assertThat(actual, equalTo(expected)); }
@Test public void shouldReturnResponseWithUnsetValuesAndMetadataIfMaxAgeIsExceeded() { GetOnlineFeaturesRequest request = getOnlineFeaturesRequest( List.of( FeatureReference.newBuilder().setName("feature1").build(), FeatureReference.newBuilder().setName("feature2").setProject("project").build())); List<Optional<FeatureRow>> featureRows = List.of(Optional.of(testFeatureRows.get(0)), Optional.of(testFeatureRows.get(2))); FeatureSetSpec spec = getFeatureSetSpec().toBuilder().setMaxAge(Duration.newBuilder().setSeconds(1)).build(); FeatureSetRequest featureSetRequest = FeatureSetRequest.newBuilder() .addAllFeatureReferences(request.getFeaturesList()) .setSpec(spec) .build(); when(specService.getFeatureSets(request.getFeaturesList(), "")) .thenReturn(Collections.singletonList(featureSetRequest)); when(retriever.getOnlineFeatures(request.getEntityRowsList(), featureSetRequest)) .thenReturn(featureRows); when(tracer.buildSpan(ArgumentMatchers.any())).thenReturn(Mockito.mock(SpanBuilder.class)); GetOnlineFeaturesResponse expected = GetOnlineFeaturesResponse.newBuilder() .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(1)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("a")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", intValue(1)) .putStatuses("feature1", FieldStatus.PRESENT) .putFields("project/feature2", intValue(1)) .putStatuses("project/feature2", FieldStatus.PRESENT) .build()) .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(2)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("b")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", Value.newBuilder().build()) .putStatuses("feature1", FieldStatus.OUTSIDE_MAX_AGE) .putFields("project/feature2", Value.newBuilder().build()) .putStatuses("project/feature2", FieldStatus.OUTSIDE_MAX_AGE) .build()) .build(); GetOnlineFeaturesResponse actual = onlineServingService.getOnlineFeatures(request); assertThat(actual, equalTo(expected)); }
@Test public void shouldFilterOutUndesiredRows() { GetOnlineFeaturesRequest request = getOnlineFeaturesRequest( List.of(FeatureReference.newBuilder().setName("feature1").build())); List<Optional<FeatureRow>> featureRows = List.of(Optional.of(testFeatureRows.get(0)), Optional.of(testFeatureRows.get(1))); FeatureSetRequest featureSetRequest = FeatureSetRequest.newBuilder() .addAllFeatureReferences(request.getFeaturesList()) .setSpec(getFeatureSetSpec()) .build(); when(specService.getFeatureSets(request.getFeaturesList(), "")) .thenReturn(Collections.singletonList(featureSetRequest)); when(retriever.getOnlineFeatures(request.getEntityRowsList(), featureSetRequest)) .thenReturn(featureRows); when(tracer.buildSpan(ArgumentMatchers.any())).thenReturn(Mockito.mock(SpanBuilder.class)); GetOnlineFeaturesResponse expected = GetOnlineFeaturesResponse.newBuilder() .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(1)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("a")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", intValue(1)) .putStatuses("feature1", FieldStatus.PRESENT) .build()) .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(2)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("b")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", intValue(2)) .putStatuses("feature1", FieldStatus.PRESENT) .build()) .build(); GetOnlineFeaturesResponse actual = onlineServingService.getOnlineFeatures(request); assertThat(actual, equalTo(expected)); }
@Test public void shouldApplyProjectOverrideInRequest() { GetOnlineFeaturesRequest request = getOnlineFeaturesRequest( List.of( FeatureReference.newBuilder().setName("feature1").build(), FeatureReference.newBuilder() .setName("feature2") .setProject("project") .build())) .toBuilder() .setProject("project") .build(); List<Optional<FeatureRow>> featureRows = List.of(Optional.of(testFeatureRows.get(0)), Optional.of(testFeatureRows.get(1))); FeatureSetRequest featureSetRequest = FeatureSetRequest.newBuilder() .addAllFeatureReferences(request.getFeaturesList()) .setSpec(getFeatureSetSpec()) .build(); when(specService.getFeatureSets(request.getFeaturesList(), "project")) .thenReturn(Collections.singletonList(featureSetRequest)); when(retriever.getOnlineFeatures(request.getEntityRowsList(), featureSetRequest)) .thenReturn(featureRows); when(tracer.buildSpan(ArgumentMatchers.any())).thenReturn(Mockito.mock(SpanBuilder.class)); GetOnlineFeaturesResponse expected = GetOnlineFeaturesResponse.newBuilder() .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(1)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("a")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", intValue(1)) .putStatuses("feature1", FieldStatus.PRESENT) .putFields("project/feature2", intValue(1)) .putStatuses("project/feature2", FieldStatus.PRESENT) .build()) .addFieldValues( FieldValues.newBuilder() .putFields("entity1", intValue(2)) .putStatuses("entity1", FieldStatus.PRESENT) .putFields("entity2", strValue("b")) .putStatuses("entity2", FieldStatus.PRESENT) .putFields("feature1", intValue(2)) .putStatuses("feature1", FieldStatus.PRESENT) .putFields("project/feature2", intValue(2)) .putStatuses("project/feature2", FieldStatus.PRESENT) .build()) .build(); GetOnlineFeaturesResponse actual = onlineServingService.getOnlineFeatures(request); assertThat(actual, equalTo(expected)); }