Refactor stores

ingestion/src/main/java/feast/ingestion/boot/ImportJobModule.java

@@ -28,12 +28,12 @@
 import feast.ingestion.service.SpecService.Builder;
 import feast.ingestion.service.SpecService.UnsupportedBuilder;
 import feast.specs.ImportSpecProto.ImportSpec;
-import feast.storage.ErrorsStore;
-import feast.storage.ServingStore;
-import feast.storage.WarehouseStore;
-import feast.storage.service.ErrorsStoreService;
-import feast.storage.service.ServingStoreService;
-import feast.storage.service.WarehouseStoreService;
+import feast.store.errors.FeatureErrorsFactory;
+import feast.store.errors.FeatureErrorsFactoryService;
+import feast.store.serving.FeatureServingFactory;
+import feast.store.serving.FeatureServingFactoryService;
+import feast.store.warehouse.FeatureWarehouseFactory;
+import feast.store.warehouse.FeatureWarehouseFactoryService;
 import java.util.List;
 import org.apache.beam.sdk.options.PipelineOptions;
 
@@ -78,19 +78,19 @@ Specs provideSpecs(SpecService.Builder specService) {
 
   @Provides
   @Singleton
-  List<WarehouseStore> provideWarehouseStores() {
-    return WarehouseStoreService.getAll();
+  List<FeatureWarehouseFactory> provideWarehouseStores() {
+    return FeatureWarehouseFactoryService.getAll();
   }
 
   @Provides
   @Singleton
-  List<ServingStore> provideServingStores() {
-    return ServingStoreService.getAll();
+  List<FeatureServingFactory> provideServingStores() {
+    return FeatureServingFactoryService.getAll();
   }
 
   @Provides
   @Singleton
-  List<ErrorsStore> provideErrorsStores() {
-    return ErrorsStoreService.getAll();
+  List<FeatureErrorsFactory> provideErrorsStores() {
+    return FeatureErrorsFactoryService.getAll();
   }
 }

ingestion/src/main/java/feast/ingestion/transform/ErrorsStoreTransform.java

@@ -23,10 +23,10 @@
 import com.google.inject.Inject;
 import feast.ingestion.model.Specs;
 import feast.ingestion.options.ImportJobPipelineOptions;
-import feast.ingestion.transform.FeatureIO.Write;
 import feast.specs.StorageSpecProto.StorageSpec;
-import feast.storage.ErrorsStore;
-import feast.storage.noop.NoOpIO;
+import feast.store.FeatureStoreWrite;
+import feast.store.NoOpIO;
+import feast.store.errors.FeatureErrorsFactory;
 import feast.types.FeatureRowExtendedProto.FeatureRowExtended;
 import java.util.List;
 import lombok.extern.slf4j.Slf4j;
@@ -35,18 +35,19 @@
 import org.apache.hadoop.hbase.util.Strings;
 
 @Slf4j
-public class ErrorsStoreTransform extends FeatureIO.Write {
+public class ErrorsStoreTransform extends FeatureStoreWrite {
 
   private String errorsStoreType;
   private StorageSpec errorsStoreSpec;
   private Specs specs;
-  private List<ErrorsStore> errorsStores;
+  private List<FeatureErrorsFactory> errorsStoreFactories;
 
   @Inject
   public ErrorsStoreTransform(
-      ImportJobPipelineOptions options, Specs specs, List<ErrorsStore> errorsStores) {
+      ImportJobPipelineOptions options, Specs specs,
+      List<FeatureErrorsFactory> errorsStoreFactories) {
     this.specs = specs;
-    this.errorsStores = errorsStores;
+    this.errorsStoreFactories = errorsStoreFactories;
     this.errorsStoreType = options.getErrorsStoreType();
 
     if (!Strings.isEmpty(errorsStoreType)) {
@@ -60,7 +61,7 @@ public ErrorsStoreTransform(
 
   @Override
   public PDone expand(PCollection<FeatureRowExtended> input) {
-    Write write;
+    FeatureStoreWrite write;
     if (Strings.isEmpty(errorsStoreType)) {
       write = new NoOpIO.Write();
     } else {
@@ -70,11 +71,11 @@ public PDone expand(PCollection<FeatureRowExtended> input) {
     return PDone.in(input.getPipeline());
   }
 
-  ErrorsStore getErrorStore() {
+  FeatureErrorsFactory getErrorStore() {
     checkArgument(!errorsStoreType.isEmpty(), "Errors store type not provided");
-    for (ErrorsStore errorsStore : errorsStores) {
-      if (errorsStore.getType().equals(errorsStoreType)) {
-        return errorsStore;
+    for (FeatureErrorsFactory errorsStoreFactory : errorsStoreFactories) {
+      if (errorsStoreFactory.getType().equals(errorsStoreType)) {
+        return errorsStoreFactory;
       }
     }
     throw new IllegalArgumentException("Errors store type not found");

ingestion/src/main/java/feast/ingestion/transform/ServingStoreTransform.java

@@ -21,21 +21,20 @@
 import feast.ingestion.metrics.FeastMetrics;
 import feast.ingestion.model.Specs;
 import feast.ingestion.values.PFeatureRows;
-import feast.storage.ServingStore;
+import feast.store.serving.FeatureServingFactory;
+import java.util.List;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 
-import java.util.List;
-
 @Slf4j
 public class ServingStoreTransform extends PTransform<PFeatureRows, PFeatureRows> {
 
-  private List<ServingStore> stores;
+  private List<FeatureServingFactory> stores;
   private Specs specs;
 
   @Inject
-  public ServingStoreTransform(List<ServingStore> stores, Specs specs) {
+  public ServingStoreTransform(List<FeatureServingFactory> stores, Specs specs) {
     this.stores = stores;
     this.specs = specs;
   }
@@ -50,7 +49,8 @@ public PFeatureRows expand(PFeatureRows input) {
 
     output.getMain().apply("metrics.store.lag", ParDo.of(FeastMetrics.lagUpdateDoFn()));
     output.getMain().apply("metrics.store.main", ParDo.of(FeastMetrics.incrDoFn("serving_stored")));
-    output.getErrors().apply("metrics.store.errors", ParDo.of(FeastMetrics.incrDoFn("serving_errors")));
+    output.getErrors()
+        .apply("metrics.store.errors", ParDo.of(FeastMetrics.incrDoFn("serving_errors")));
     return output;
   }
 }

ingestion/src/main/java/feast/ingestion/transform/SplitOutputByStore.java

@@ -20,12 +20,12 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Lists;
 import feast.ingestion.model.Specs;
-import feast.ingestion.transform.FeatureIO.Write;
 import feast.ingestion.transform.SplitFeatures.MultiOutputSplit;
 import feast.ingestion.values.PFeatureRows;
 import feast.specs.FeatureSpecProto.FeatureSpec;
 import feast.specs.StorageSpecProto.StorageSpec;
-import feast.storage.FeatureStore;
+import feast.store.FeatureStoreWrite;
+import feast.store.FeatureStoreFactory;
 import feast.types.FeatureRowExtendedProto.FeatureRowExtended;
 import java.util.Collection;
 import java.util.HashMap;
@@ -46,20 +46,20 @@
 @Slf4j
 public class SplitOutputByStore extends PTransform<PFeatureRows, PFeatureRows> {
 
-  private Collection<? extends FeatureStore> stores;
+  private Collection<? extends FeatureStoreFactory> stores;
   private SerializableFunction<FeatureSpec, String> selector;
   private Specs specs;
 
   @Override
   public PFeatureRows expand(PFeatureRows input) {
-    Map<String, Write> transforms = getFeatureStoreTransforms();
+    Map<String, FeatureStoreWrite> transforms = getFeatureStoreTransforms();
     Set<String> keys = transforms.keySet();
 
     log.info(String.format("Splitting on keys = [%s]", String.join(",", keys)));
     MultiOutputSplit<String> splitter = new MultiOutputSplit<>(selector, keys, specs);
     PCollectionTuple splits = input.getMain().apply(splitter);
 
-    Map<TupleTag<FeatureRowExtended>, Write> taggedTransforms = new HashMap<>();
+    Map<TupleTag<FeatureRowExtended>, FeatureStoreWrite> taggedTransforms = new HashMap<>();
     for (String key : transforms.keySet()) {
       TupleTag<FeatureRowExtended> tag = splitter.getStrategy().getTag(key);
       taggedTransforms.put(tag, transforms.get(key));
@@ -72,17 +72,17 @@ public PFeatureRows expand(PFeatureRows input) {
         input.getErrors());
   }
 
-  private Map<String, FeatureStore> getStoresMap() {
-    Map<String, FeatureStore> storesMap = new HashMap<>();
-    for (FeatureStore servingStore : stores) {
+  private Map<String, FeatureStoreFactory> getStoresMap() {
+    Map<String, FeatureStoreFactory> storesMap = new HashMap<>();
+    for (FeatureStoreFactory servingStore : stores) {
       storesMap.put(servingStore.getType(), servingStore);
     }
     return storesMap;
   }
 
-  private Map<String, Write> getFeatureStoreTransforms() {
-    Map<String, FeatureStore> storesMap = getStoresMap();
-    Map<String, Write> transforms = new HashMap<>();
+  private Map<String, FeatureStoreWrite> getFeatureStoreTransforms() {
+    Map<String, FeatureStoreFactory> storesMap = getStoresMap();
+    Map<String, FeatureStoreWrite> transforms = new HashMap<>();
     Map<String, StorageSpec> storageSpecs = specs.getStorageSpecs();
     for (String storeId : storageSpecs.keySet()) {
       StorageSpec storageSpec = storageSpecs.get(storeId);
@@ -109,14 +109,14 @@ private Map<String, Write> getFeatureStoreTransforms() {
   public static class WriteTags extends
       PTransform<PCollectionTuple, PCollection<FeatureRowExtended>> {
 
-    private Map<TupleTag<FeatureRowExtended>, Write> transforms;
+    private Map<TupleTag<FeatureRowExtended>, FeatureStoreWrite> transforms;
     private TupleTag<FeatureRowExtended> mainTag;
 
     @Override
     public PCollection<FeatureRowExtended> expand(PCollectionTuple tuple) {
       List<PCollection<FeatureRowExtended>> outputList = Lists.newArrayList();
       for (TupleTag<FeatureRowExtended> tag : transforms.keySet()) {
-        Write write = transforms.get(tag);
+        FeatureStoreWrite write = transforms.get(tag);
         Preconditions.checkNotNull(write, String.format("Null transform for tag=%s", tag.getId()));
         PCollection<FeatureRowExtended> input = tuple.get(tag);
         input.apply(String.format("Write to %s", tag.getId()), write);

ingestion/src/main/java/feast/ingestion/transform/WarehouseStoreTransform.java

@@ -18,23 +18,23 @@
 package feast.ingestion.transform;
 
 import com.google.inject.Inject;
+import feast.ingestion.metrics.FeastMetrics;
+import feast.ingestion.model.Specs;
+import feast.ingestion.values.PFeatureRows;
+import feast.store.warehouse.FeatureWarehouseFactory;
 import java.util.List;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
-import feast.ingestion.metrics.FeastMetrics;
-import feast.ingestion.model.Specs;
-import feast.ingestion.values.PFeatureRows;
-import feast.storage.WarehouseStore;
 
 @Slf4j
 public class WarehouseStoreTransform extends PTransform<PFeatureRows, PFeatureRows> {
 
-  private List<WarehouseStore> stores;
+  private List<FeatureWarehouseFactory> stores;
   private Specs specs;
 
   @Inject
-  public WarehouseStoreTransform(List<WarehouseStore> stores, Specs specs) {
+  public WarehouseStoreTransform(List<FeatureWarehouseFactory> stores, Specs specs) {
     this.stores = stores;
     this.specs = specs;
   }
@@ -48,8 +48,10 @@ public PFeatureRows expand(PFeatureRows input) {
                 stores,
                 (featureSpec) -> featureSpec.getDataStores().getWarehouse().getId(),
                 specs));
-    output.getMain().apply("metrics.store.main", ParDo.of(FeastMetrics.incrDoFn("warehouse_stored")));
-    output.getErrors().apply("metrics.store.errors", ParDo.of(FeastMetrics.incrDoFn("warehouse_errors")));
+    output.getMain()
+        .apply("metrics.store.main", ParDo.of(FeastMetrics.incrDoFn("warehouse_stored")));
+    output.getErrors()
+        .apply("metrics.store.errors", ParDo.of(FeastMetrics.incrDoFn("warehouse_errors")));
     return output;
   }
 }

ingestion/src/main/java/feast/ingestion/transform/fn/ValidateFeatureRowsDoFn.java

@@ -31,10 +31,10 @@
 import feast.specs.ImportSpecProto.Field;
 import feast.specs.ImportSpecProto.ImportSpec;
 import feast.specs.StorageSpecProto.StorageSpec;
-import feast.storage.ServingStore;
-import feast.storage.WarehouseStore;
-import feast.storage.service.ServingStoreService;
-import feast.storage.service.WarehouseStoreService;
+import feast.store.serving.FeatureServingFactory;
+import feast.store.serving.FeatureServingFactoryService;
+import feast.store.warehouse.FeatureWarehouseFactory;
+import feast.store.warehouse.FeatureWarehouseFactoryService;
 import feast.types.FeatureProto.Feature;
 import feast.types.FeatureRowProto.FeatureRow;
 import feast.types.GranularityProto.Granularity.Enum;
@@ -66,10 +66,10 @@ public void setup() {
         featureIds.add(field.getFeatureId());
       }
     }
-    for (ServingStore store : ServingStoreService.getAll()) {
+    for (FeatureServingFactory store : FeatureServingFactoryService.getAll()) {
       supportedServingTypes.add(store.getType());
     }
-    for (WarehouseStore store : WarehouseStoreService.getAll()) {
+    for (FeatureWarehouseFactory store : FeatureWarehouseFactoryService.getAll()) {
       supportedWarehouseTypes.add(store.getType());
     }
   }

ingestion/src/main/java/feast/source/bigquery/BigQueryToFeatureRowFn.java

@@ -23,7 +23,7 @@
 import com.google.cloud.bigquery.StandardSQLTypeName;
 import com.google.common.collect.Maps;
 import com.google.protobuf.Timestamp;
-import feast.storage.bigquery.ValueBigQueryBuilder;
+import feast.store.warehouse.bigquery.ValueBigQueryBuilder;
 import java.util.Map;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.io.gcp.bigquery.SchemaAndRecord;
@@ -34,8 +34,6 @@
 import feast.types.FeatureProto.Feature;
 import feast.types.FeatureRowProto.FeatureRow;
 import feast.types.ValueProto.Value;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /**
  * This is a serializable function used with the BigQueryIO for fetching feature rows directly from