feast-dev · zhilingc · May 15, 2020 · May 15, 2020 · May 18, 2020 · May 18, 2020
diff --git a/core/src/main/java/feast/core/job/JobUpdateTask.java b/core/src/main/java/feast/core/job/JobUpdateTask.java
@@ -102,7 +102,18 @@ public Job call() {
 
   boolean requiresUpdate(Job job) {
     // If set of feature sets has changed
-    return !Sets.newHashSet(featureSets).equals(Sets.newHashSet(job.getFeatureSets()));
+    if (!Sets.newHashSet(featureSets).equals(Sets.newHashSet(job.getFeatureSets()))) {
+      return false;
+    }
+
+    // If job is not currently populating the required version
+    for (FeatureSet featureSet : featureSets) {
+      if (job.getFeatureSetVersion(featureSet.getId()) != featureSet.getVersion()) {
+        return true;
+      }
+    }
+
+    return false;
   }
 
   private Job createJob() {
@@ -112,10 +123,17 @@ private Job createJob() {
 
   /** Start or update the job to ingest data to the sink. */
   private Job startJob(String jobId) {
-
     Job job =
-        new Job(
-            jobId, "", jobManager.getRunnerType(), source, store, featureSets, JobStatus.PENDING);
+        Job.builder()
+            .id(jobId)
+            .extId("")
+            .runner(jobManager.getRunnerType())
+            .source(source)
+            .store(store)
+            .status(JobStatus.PENDING)
+            .build();
+    job.addAllFeatureSets(featureSets);
+
     try {
       logAudit(Action.SUBMIT, job, "Building graph and submitting to %s", runnerName);
 
@@ -142,7 +160,7 @@ private Job startJob(String jobId) {
 
   /** Update the given job */
   private Job updateJob(Job job) {
-    job.setFeatureSets(featureSets);
+    job.addAllFeatureSets(featureSets);
     job.setStore(store);
     logAudit(Action.UPDATE, job, "Updating job %s for runner %s", job.getId(), runnerName);
     return jobManager.updateJob(job);

diff --git a/core/src/main/java/feast/core/model/FeatureSet.java b/core/src/main/java/feast/core/model/FeatureSet.java
@@ -42,6 +42,9 @@ public class FeatureSet extends AbstractTimestampEntity {
   // Id of the featureSet, defined as project/feature_set_name:feature_set_version
   @Id @GeneratedValue private long id;
 
+  // Version of the feature set. Only incremented when a new schema is applied.
+  private long version = 0;
+
   // Name of the featureSet
   @Column(name = "name", nullable = false)
   private String name;
@@ -207,6 +210,9 @@ public void updateFromProto(FeatureSetProto.FeatureSet featureSetProto)
       Feature newFeature = Feature.fromProto(featureSpec);
       addFeature(newFeature);
     }
+
+    // 5. Increment the version
+    this.version += 1;
   }
 
   public void addEntities(List<Entity> entities) {

diff --git a/core/src/main/java/feast/core/model/Job.java b/core/src/main/java/feast/core/model/Job.java
@@ -20,16 +20,17 @@
 import feast.core.job.Runner;
 import feast.proto.core.FeatureSetProto;
 import feast.proto.core.IngestionJobProto;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
 import javax.persistence.*;
 import javax.persistence.Entity;
 import lombok.AllArgsConstructor;
+import lombok.Builder;
 import lombok.Getter;
 import lombok.Setter;
 
 /** Contains information about a run job. */
 @AllArgsConstructor
+@Builder
 @Getter
 @Setter
 @Entity
@@ -69,14 +70,29 @@ public class Job extends AbstractTimestampEntity {
         @Index(name = "idx_jobs_feature_sets_job_id", columnList = "job_id"),
         @Index(name = "idx_jobs_feature_sets_feature_sets_id", columnList = "feature_sets_id")
       })
-  private List<FeatureSet> featureSets;
+  @Builder.Default
+  private Set<FeatureSet> featureSets = new HashSet<>();
+
+  @ElementCollection
+  @CollectionTable(
+      name = "feature_sets_version_mapping",
+      joinColumns = {@JoinColumn(name = "job_id", referencedColumnName = "id")})
+  @MapKeyColumn(name = "feature_set_id")
+  @Column(name = "feature_sets_version")
+  @Builder.Default
+  private Map<Long, Long> featureSetVersionMap = new HashMap<>();
 
   @Enumerated(EnumType.STRING)
   @Column(name = "status", length = 16)
   private JobStatus status;
 
   public Job() {
     super();
+
+    // Need to add defaults here since Builder.Default masks defaults
+    // from this constructor
+    this.featureSetVersionMap = new HashMap<>();
+    this.featureSets = new HashSet<>();
   }
 
   public boolean hasTerminated() {
@@ -91,6 +107,21 @@ public String getSinkName() {
     return store.getName();
   }
 
+  public Long getFeatureSetVersion(Long id) {
+    return featureSetVersionMap.getOrDefault(id, -1L);
+  }
+
+  public void addFeatureSet(FeatureSet featureSet) {
+    featureSets.add(featureSet);
+    featureSetVersionMap.put(featureSet.getId(), featureSet.getVersion());
+  }
+
+  public void addAllFeatureSets(List<FeatureSet> allFeatureSets) {
+    for (FeatureSet featureSet : allFeatureSets) {
+      this.addFeatureSet(featureSet);
+    }
+  }
+
   /**
    * Convert a job model to ingestion job proto
    *

diff --git a/core/src/main/java/feast/core/service/JobCoordinatorService.java b/core/src/main/java/feast/core/service/JobCoordinatorService.java
@@ -29,7 +29,6 @@
 import feast.core.model.Store;
 import feast.proto.core.CoreServiceProto.ListStoresRequest.Filter;
 import feast.proto.core.CoreServiceProto.ListStoresResponse;
-import feast.proto.core.FeatureSetProto.FeatureSetStatus;
 import feast.proto.core.StoreProto;
 import feast.proto.core.StoreProto.Store.Subscription;
 import java.util.ArrayList;
@@ -121,9 +120,6 @@ public void Poll() throws InvalidProtocolBufferException {
 
     log.info("Creating/Updating {} jobs...", jobUpdateTasks.size());
     startOrUpdateJobs(jobUpdateTasks);
-
-    log.info("Updating feature set status");
-    updateFeatureSetStatuses(jobUpdateTasks);
   }
 
   void startOrUpdateJobs(List<JobUpdateTask> tasks) {
@@ -132,49 +128,22 @@ void startOrUpdateJobs(List<JobUpdateTask> tasks) {
     tasks.forEach(ecs::submit);
 
     int completedTasks = 0;
+    List<Job> completedJobs = new ArrayList<>();
     while (completedTasks < tasks.size()) {
       try {
         Job job = ecs.take().get();
         if (job != null) {
-          jobRepository.saveAndFlush(job);
+          completedJobs.add(job);
         }
       } catch (ExecutionException | InterruptedException e) {
         log.warn("Unable to start or update job: {}", e.getMessage());
       }
       completedTasks++;
     }
+    jobRepository.saveAll(completedJobs);
     executorService.shutdown();
   }
 
-  // TODO: make this more efficient
-  private void updateFeatureSetStatuses(List<JobUpdateTask> jobUpdateTasks) {
-    Set<FeatureSet> ready = new HashSet<>();
-    Set<FeatureSet> pending = new HashSet<>();
-    for (JobUpdateTask task : jobUpdateTasks) {
-      getJob(task.getSource(), task.getStore())
-          .ifPresent(
-              job -> {
-                if (job.isRunning()) {
-                  ready.addAll(job.getFeatureSets());
-                } else {
-                  pending.addAll(job.getFeatureSets());
-                }
-              });
-    }
-    ready.removeAll(pending);
-    ready.forEach(
-        fs -> {
-          fs.setStatus(FeatureSetStatus.STATUS_READY);
-          featureSetRepository.save(fs);
-        });
-    pending.forEach(
-        fs -> {
-          fs.setStatus(FeatureSetStatus.STATUS_PENDING);
-          featureSetRepository.save(fs);
-        });
-    featureSetRepository.flush();
-  }
-
   @Transactional
   public Optional<Job> getJob(Source source, Store store) {
     List<Job> jobs =

diff --git a/core/src/main/java/feast/core/service/JobService.java b/core/src/main/java/feast/core/service/JobService.java
@@ -34,6 +34,7 @@
 import feast.proto.core.CoreServiceProto.RestartIngestionJobResponse;
 import feast.proto.core.CoreServiceProto.StopIngestionJobRequest;
 import feast.proto.core.CoreServiceProto.StopIngestionJobResponse;
+import feast.proto.core.FeatureSetProto.FeatureSetStatus;
 import feast.proto.core.FeatureSetReferenceProto.FeatureSetReference;
 import feast.proto.core.IngestionJobProto;
 import java.util.ArrayList;
@@ -181,6 +182,10 @@ public RestartIngestionJobResponse restartJob(RestartIngestionJobRequest request
             job.getId(), job.getExtId(), job.getRunner()));
     // sync job status & update job model in job repository
     job = this.syncJobStatus(jobManager, job);
+
+    for (FeatureSet featureSet : job.getFeatureSets()) {
+      featureSet.setStatus(FeatureSetStatus.STATUS_PENDING);
+    }
     this.jobRepository.saveAndFlush(job);
 
     return RestartIngestionJobResponse.newBuilder().build();
@@ -226,6 +231,11 @@ public StopIngestionJobResponse stopJob(StopIngestionJobRequest request)
 
     // sync job status & update job model in job repository
     job = this.syncJobStatus(jobManager, job);
+
+    for (FeatureSet featureSet : job.getFeatureSets()) {
+      featureSet.setStatus(FeatureSetStatus.STATUS_PENDING);
+    }
+
     this.jobRepository.saveAndFlush(job);
 
     return StopIngestionJobResponse.newBuilder().build();

diff --git a/core/src/main/java/feast/core/service/SpecService.java b/core/src/main/java/feast/core/service/SpecService.java
@@ -21,13 +21,11 @@
 
 import com.google.protobuf.InvalidProtocolBufferException;
 import feast.core.dao.FeatureSetRepository;
+import feast.core.dao.JobRepository;
 import feast.core.dao.ProjectRepository;
 import feast.core.dao.StoreRepository;
 import feast.core.exception.RetrievalException;
-import feast.core.model.FeatureSet;
-import feast.core.model.Project;
-import feast.core.model.Source;
-import feast.core.model.Store;
+import feast.core.model.*;
 import feast.core.validators.FeatureSetValidator;
 import feast.proto.core.CoreServiceProto.ApplyFeatureSetResponse;
 import feast.proto.core.CoreServiceProto.ApplyFeatureSetResponse.Status;
@@ -46,7 +44,9 @@
 import feast.proto.core.StoreProto;
 import feast.proto.core.StoreProto.Store.Subscription;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
@@ -63,17 +63,20 @@ public class SpecService {
   private final FeatureSetRepository featureSetRepository;
   private final ProjectRepository projectRepository;
   private final StoreRepository storeRepository;
+  private final JobRepository jobRepository;
   private final Source defaultSource;
 
   @Autowired
   public SpecService(
       FeatureSetRepository featureSetRepository,
       StoreRepository storeRepository,
       ProjectRepository projectRepository,
+      JobRepository jobRepository,
       Source defaultSource) {
     this.featureSetRepository = featureSetRepository;
     this.storeRepository = storeRepository;
     this.projectRepository = projectRepository;
+    this.jobRepository = jobRepository;
     this.defaultSource = defaultSource;
   }
 
@@ -83,8 +86,9 @@ public SpecService(
    * is omitted, the latest feature set will be provided.
    *
    * @param request: GetFeatureSetRequest Request containing filter parameters.
-   * @return Returns a GetFeatureSetResponse containing a feature set..
+   * @return Returns a GetFeatureSetResponse containing a feature set.
    */
+  @Transactional
   public GetFeatureSetResponse getFeatureSet(GetFeatureSetRequest request)
       throws InvalidProtocolBufferException {
 
@@ -108,6 +112,9 @@ public GetFeatureSetResponse getFeatureSet(GetFeatureSetRequest request)
       throw new RetrievalException(
           String.format("Feature set with name \"%s\" could not be found.", request.getName()));
     }
+
+    checkAndUpdateStatus(featureSet);
+
     return GetFeatureSetResponse.newBuilder().setFeatureSet(featureSet.toProto()).build();
   }
 
@@ -126,6 +133,7 @@ public GetFeatureSetResponse getFeatureSet(GetFeatureSetRequest request)
    * @param filter filter containing the desired featureSet name
    * @return ListFeatureSetsResponse with list of featureSets found matching the filter
    */
+  @Transactional
   public ListFeatureSetsResponse listFeatureSets(ListFeatureSetsRequest.Filter filter)
       throws InvalidProtocolBufferException {
     String name = filter.getFeatureSetName();
@@ -181,6 +189,7 @@ public ListFeatureSetsResponse listFeatureSets(ListFeatureSetsRequest.Filter fil
     ListFeatureSetsResponse.Builder response = ListFeatureSetsResponse.newBuilder();
     if (featureSets.size() > 0) {
       for (FeatureSet featureSet : featureSets) {
+        checkAndUpdateStatus(featureSet);
         response.addFeatureSets(featureSet.toProto());
       }
     }
@@ -330,4 +339,50 @@ public UpdateStoreResponse updateStore(UpdateStoreRequest updateStoreRequest)
         .setStore(updateStoreRequest.getStore())
         .build();
   }
+
+  /**
+   * Checks the status of the given feature set. If there are jobs populating values from this
+   * feature set, and if for each source and sink pair, there is at least 1 job running, it sets the
+   * feature set's status to STATUS_READY and updates the db with the new status.
+   *
+   * @param featureSet {@link FeatureSet}
+   */
+  private void checkAndUpdateStatus(FeatureSet featureSet) throws InvalidProtocolBufferException {
+    // check if the job is ready
+    List<Job> jobsForFeatureSet =
+        jobRepository.findByFeatureSetsIn(Collections.singletonList(featureSet));
+    if (jobsForFeatureSet.size() != 0) {
+
+      List<List<Job>> jobsGroupedBySourceAndSink =
+          jobsForFeatureSet.stream()
+              .collect(
+                  Collectors.groupingBy(Job::getSource, Collectors.groupingBy(Job::getSinkName)))
+              .values()
+              .stream()
+              .flatMap(map -> map.values().stream())
+              .collect(Collectors.toList());
+
+      Long featureSetVersion = featureSet.getVersion();
+      for (List<Job> jobs : jobsGroupedBySourceAndSink) {
+        long jobsRunning =
+            jobs.stream()
+                .filter(job -> job.getStatus() == JobStatus.RUNNING)
+                .filter(
+                    job -> job.getFeatureSetVersion(featureSet.getId()).equals(featureSetVersion))
+                .count();
+        if (jobsRunning == 0) {
+          if (featureSet.getStatus() == FeatureSetStatus.STATUS_READY) {
+            featureSet.setStatus(FeatureSetStatus.STATUS_PENDING);
+            featureSetRepository.save(featureSet);
+          }
+          return;
+        }
+      }
+
+      if (featureSet.getStatus() != FeatureSetStatus.STATUS_READY) {
+        featureSet.setStatus(FeatureSetStatus.STATUS_READY);
+        featureSetRepository.save(featureSet);
+      }
+    }
+  }
 }