Simplified stages #245

Author: andrus

RELEASE-NOTES.md

@@ -2,6 +2,7 @@
 
 * #239 Remove APIs deprecated in 3.0.0 
 * #244 Simplify StringConverter
+* #245 Simplified stages
 * #246 Missing stage: "SourceKeysStage.CONVERT_SOURCE_ROWS"
 * #247 Missing stage: "DeleteStage.DELETE_TARGET"
 

link-move/src/main/java/com/nhl/link/move/CreateBuilder.java

@@ -3,8 +3,10 @@
 import com.nhl.link.move.extractor.model.ExtractorName;
 import com.nhl.link.move.runtime.task.create.CreateSegment;
 import com.nhl.link.move.runtime.task.create.CreateStage;
+import org.dflib.DataFrame;
 
 import java.util.function.BiConsumer;
+import java.util.function.Function;
 
 /**
  * A builder of an {@link LmTask} that performs fast "create" synchronization without any source/target key matching.
@@ -50,9 +52,20 @@ default CreateBuilder sourceExtractor(String location) {
     CreateBuilder batchSize(int batchSize);
 
     /**
-     * Adds a callback invoked for each processed segment after the specified stage in the "create" processing pipeline.
+     * Adds a callback invoked for each data segment after the specified stage in the "create" processing pipeline.
      *
      * @since 3.0.0
      */
-    CreateBuilder stage(CreateStage stageType, BiConsumer<Execution, CreateSegment> callback);
+    CreateBuilder stage(CreateStage stage, BiConsumer<Execution, CreateSegment> callback);
+
+    /**
+     * Adds a callback invoked for each date segment after the specified stage in the "create" pipeline was processed.
+     * The result of that stage is passed to the transformer argument. The value returned from the transformer overrides
+     * the previous result for the stage.
+     *
+     * @since 4.0.0
+     */
+    default CreateBuilder stage(CreateStage stage, Function<DataFrame, DataFrame> transformer) {
+        return stage(stage, (e, s) -> s.postProcess(stage, transformer));
+    }
 }

link-move/src/main/java/com/nhl/link/move/CreateOrUpdateBuilder.java

@@ -5,8 +5,10 @@
 import com.nhl.link.move.runtime.task.createorupdate.CreateOrUpdateSegment;
 import com.nhl.link.move.runtime.task.createorupdate.CreateOrUpdateStage;
 import org.apache.cayenne.exp.property.Property;
+import org.dflib.DataFrame;
 
 import java.util.function.BiConsumer;
+import java.util.function.Function;
 
 /**
  * A builder of an {@link LmTask} that performs create-or-update synchronization.
@@ -87,5 +89,16 @@ default CreateOrUpdateBuilder sourceExtractor(String location) {
      *
      * @since 3.0.0
      */
-    CreateOrUpdateBuilder stage(CreateOrUpdateStage stageType, BiConsumer<Execution, CreateOrUpdateSegment> callback);
+    CreateOrUpdateBuilder stage(CreateOrUpdateStage stage, BiConsumer<Execution, CreateOrUpdateSegment> callback);
+
+    /**
+     * Adds a callback invoked for each date segment after the specified stage in the "create-or-update" pipeline was processed.
+     * The result of that stage is passed to the transformer argument. The value returned from the transformer overrides
+     * the previous result for the stage.
+     *
+     * @since 4.0.0
+     */
+    default CreateOrUpdateBuilder stage(CreateOrUpdateStage stage, Function<DataFrame, DataFrame> transformer) {
+        return stage(stage, (e, s) -> s.postProcess(stage, transformer));
+    }
 }

link-move/src/main/java/com/nhl/link/move/DeleteBuilder.java

@@ -6,8 +6,10 @@
 import com.nhl.link.move.runtime.task.delete.DeleteStage;
 import org.apache.cayenne.exp.Expression;
 import org.apache.cayenne.exp.property.Property;
+import org.dflib.DataFrame;
 
 import java.util.function.BiConsumer;
+import java.util.function.Function;
 
 /**
  * since 1.3
@@ -50,5 +52,16 @@ default DeleteBuilder sourceMatchExtractor(String location) {
      *
      * @since 3.0.0
      */
-    DeleteBuilder stage(DeleteStage stageType, BiConsumer<Execution, DeleteSegment> callback);
+    DeleteBuilder stage(DeleteStage stage, BiConsumer<Execution, DeleteSegment> callback);
+
+    /**
+     * Adds a callback invoked for each date segment after the specified stage in the "delete" pipeline was processed.
+     * The result of that stage is passed to the transformer argument. The value returned from the transformer overrides
+     * the previous result for the stage.
+     *
+     * @since 4.0.0
+     */
+    default DeleteBuilder stage(DeleteStage stage, Function<DataFrame, DataFrame> transformer) {
+        return stage(stage, (e, s) -> s.postProcess(stage, transformer));
+    }
 }

link-move/src/main/java/com/nhl/link/move/SourceKeysBuilder.java

@@ -4,8 +4,10 @@
 import com.nhl.link.move.mapper.Mapper;
 import com.nhl.link.move.runtime.task.sourcekeys.SourceKeysSegment;
 import com.nhl.link.move.runtime.task.sourcekeys.SourceKeysStage;
+import org.dflib.DataFrame;
 
 import java.util.function.BiConsumer;
+import java.util.function.Function;
 
 /**
  * A builder of an {@link LmTask} that extracts all the keys from the source
@@ -62,5 +64,16 @@ default SourceKeysBuilder sourceExtractor(String location) {
      *
      * @since 3.0.0
      */
-    SourceKeysBuilder stage(SourceKeysStage stageType, BiConsumer<Execution, SourceKeysSegment> callback);
+    SourceKeysBuilder stage(SourceKeysStage stage, BiConsumer<Execution, SourceKeysSegment> callback);
+
+    /**
+     * Adds a callback invoked for each date segment after the specified stage in the "source keys" pipeline was processed.
+     * The result of that stage is passed to the transformer argument. The value returned from the transformer overrides
+     * the previous result for the stage.
+     *
+     * @since 4.0.0
+     */
+    default SourceKeysBuilder stage(SourceKeysStage stage, Function<DataFrame, DataFrame> transformer) {
+        return stage(stage, (e, s) -> s.postProcess(stage, transformer));
+    }
 }

link-move/src/main/java/com/nhl/link/move/runtime/task/BaseTaskBuilder.java

@@ -31,8 +31,8 @@ public BaseTaskBuilder(LmLogger logger) {
     /**
      * @since 3.0.0
      */
-    public T stage(U stageType, BiConsumer<Execution, S> callback) {
-        stageListenersBuilder.addStageCallback(stageType, callback);
+    public T stage(U stage, BiConsumer<Execution, S> callback) {
+        stageListenersBuilder.addStageCallback(stage, callback);
         return (T) this;
     }
 

link-move/src/main/java/com/nhl/link/move/runtime/task/common/DataSegment.java

@@ -1,4 +1,51 @@
 package com.nhl.link.move.runtime.task.common;
 
-public interface DataSegment {
+import org.dflib.DataFrame;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+
+public abstract class DataSegment<U extends TaskStageType> {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(DataSegment.class);
+
+    protected final Map<U, DataFrame> stageResults;
+
+    public DataSegment() {
+        // while results are mutable, stages are sequential (and DataSegment is assumed to be single-threaded), so no
+        // need for concurrency
+        this.stageResults = new HashMap<>();
+    }
+
+    /**
+     * @since 4.0.0
+     */
+    public DataFrame get(U stage) {
+        return stageResults.get(stage);
+    }
+
+    /**
+     * @since 4.0.0
+     */
+    public void set(U stage, DataFrame df) {
+        stageResults.put(stage, df);
+    }
+
+    /**
+     * Retrieves data from the specified stage and, if it is present, applies the processor and saves the result back
+     * replacing the previous result of this stage.
+     *
+     * @since 4.0.0
+     */
+    public void postProcess(U stage, Function<DataFrame, DataFrame> processor) {
+        DataFrame df = get(stage);
+        if (df != null) {
+            set(stage, processor.apply(df));
+        } else {
+            LOGGER.warn("Skipping postprocessing. Stage {} has no results", stage);
+        }
+    }
 }

link-move/src/main/java/com/nhl/link/move/runtime/task/create/CreateSegment.java

@@ -1,27 +1,21 @@
 package com.nhl.link.move.runtime.task.create;
 
 import com.nhl.link.move.RowAttribute;
-import org.dflib.DataFrame;
 import com.nhl.link.move.runtime.task.common.DataSegment;
 import org.apache.cayenne.ObjectContext;
+import org.dflib.DataFrame;
 
 /**
  * @since 2.6
  */
-public class CreateSegment implements DataSegment {
+public class CreateSegment extends DataSegment<CreateStage> {
 
     public static final String TARGET_COLUMN = "$lm_target";
     public static final String TARGET_CREATED_COLUMN = "$lm_target_created";
 
     private final ObjectContext context;
     private final RowAttribute[] sourceRowsHeader;
 
-    private DataFrame sourceRows;
-    private DataFrame sources;
-    private DataFrame mapped;
-    private DataFrame fksResolved;
-    private DataFrame merged;
-
     public CreateSegment(ObjectContext context, RowAttribute[] sourceRowsHeader) {
         this.sourceRowsHeader = sourceRowsHeader;
         this.context = context;
@@ -36,56 +30,56 @@ public RowAttribute[] getSourceRowsHeader() {
     }
 
     public DataFrame getSourceRows() {
-        return sourceRows;
+        return get(CreateStage.EXTRACT_SOURCE_ROWS);
     }
 
     /**
      * @since 2.17
      */
-    public CreateSegment setSourceRows(DataFrame sourceRows) {
-        this.sourceRows = sourceRows;
+    public CreateSegment setSourceRows(DataFrame df) {
+        set(CreateStage.EXTRACT_SOURCE_ROWS, df);
         return this;
     }
 
     public DataFrame getSources() {
-        return sources;
+        return get(CreateStage.CONVERT_SOURCE_ROWS);
     }
 
-    public CreateSegment setSources(DataFrame translatedSegment) {
-        this.sources = translatedSegment;
-        return this;
-    }
-
-    public DataFrame getMerged() {
-        return merged;
-    }
-
-    public CreateSegment setMerged(DataFrame merged) {
-        this.merged = merged;
+    public CreateSegment setSources(DataFrame df) {
+        set(CreateStage.CONVERT_SOURCE_ROWS, df);
         return this;
     }
 
     public DataFrame getMapped() {
-        return mapped;
+        return get(CreateStage.MAP_TARGET);
     }
 
-    public CreateSegment setMapped(DataFrame mapped) {
-        this.mapped = mapped;
+    public CreateSegment setMapped(DataFrame df) {
+        set(CreateStage.MAP_TARGET, df);
         return this;
     }
 
     /**
      * @since 2.12
      */
     public DataFrame getFksResolved() {
-        return fksResolved;
+        return get(CreateStage.RESOLVE_FK_VALUES);
     }
 
     /**
      * @since 2.12
      */
-    public CreateSegment setFksResolved(DataFrame fksResolved) {
-        this.fksResolved = fksResolved;
+    public CreateSegment setFksResolved(DataFrame df) {
+        set(CreateStage.RESOLVE_FK_VALUES, df);
+        return this;
+    }
+
+    public DataFrame getMerged() {
+        return get(CreateStage.MERGE_TARGET);
+    }
+
+    public CreateSegment setMerged(DataFrame df) {
+        set(CreateStage.MERGE_TARGET, df);
         return this;
     }
 }