Pass runtime configs & variables to BeamSqlSeekableTable (#28253)

* closing #28145 * remove unsupported parameter & add parameter type to context * remove unnecessary semicolon
apache · Sep 5, 2023 · 3364271 · 3364271
1 parent f6c6400
commit 3364271
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/...extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/BeamSqlSeekableTable.java b/...extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/BeamSqlSeekableTable.java
@@ -19,6 +19,8 @@
 
 import java.io.Serializable;
 import java.util.List;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.values.Row;
 
 /**
@@ -27,11 +29,17 @@
  */
 public interface BeamSqlSeekableTable extends Serializable {
   /** prepare the instance. */
-  default void setUp() {};
+  default void setUp() {}
+
+  default void startBundle(
+      DoFn<Row, Row>.StartBundleContext context, PipelineOptions pipelineOptions) {}
+
+  default void finishBundle(
+      DoFn<Row, Row>.FinishBundleContext context, PipelineOptions pipelineOptions) {}
 
   /** return a list of {@code Row} with given key set. */
   List<Row> seekRow(Row lookupSubRow);
 
   /** cleanup resources of the instance. */
-  default void tearDown() {};
+  default void tearDown() {}
 }
diff --git a/...l/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamJoinTransforms.java b/...l/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamJoinTransforms.java
@@ -26,6 +26,7 @@
 import org.apache.beam.sdk.extensions.sql.impl.utils.SerializableRexFieldAccess;
 import org.apache.beam.sdk.extensions.sql.impl.utils.SerializableRexInputRef;
 import org.apache.beam.sdk.extensions.sql.impl.utils.SerializableRexNode;
+import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.transforms.DoFn;
@@ -155,6 +156,20 @@ public void setup() {
                       seekableTable.setUp();
                     }
 
+                    @StartBundle
+                    public void startBundle(
+                        DoFn<Row, Row>.StartBundleContext context,
+                        PipelineOptions pipelineOptions) {
+                      seekableTable.startBundle(context, pipelineOptions);
+                    }
+
+                    @FinishBundle
+                    public void finishBundle(
+                        DoFn<Row, Row>.FinishBundleContext context,
+                        PipelineOptions pipelineOptions) {
+                      seekableTable.finishBundle(context, pipelineOptions);
+                    }
+
                     @ProcessElement
                     public void processElement(ProcessContext context) {
                       Row factRow = context.element();