add setJoinSubsetType to inject joinSubsetType to BeamSqlSeekableTable

apache · Sep 15, 2023 · e7ceff4 · e7ceff4
1 parent 5df59a9
commit e7ceff4
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 0 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
@@ -20,6 +20,7 @@
 import java.io.Serializable;
 import java.util.List;
 import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.values.Row;
 
@@ -28,6 +29,7 @@
  * FROM FACT_TABLE JOIN LOOKUP_TABLE ON ...}.
  */
 public interface BeamSqlSeekableTable extends Serializable {
+  default void setJoinSubsetType(Schema joinSubsetType) {}
   /** prepare the instance. */
   default void setUp() {}
 

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
@@ -116,6 +116,7 @@ public JoinAsLookup(
       this.outputSchema = outputSchema;
       this.factColOffset = factColOffset;
       joinFieldsMapping(joinCondition, factColOffset, lkpColOffset);
+      this.seekableTable.setJoinSubsetType(joinSubsetType);
     }
 
     private void joinFieldsMapping(RexNode joinCondition, int factColOffset, int lkpColOffset) {

diff --git a/...test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSideInputLookupJoinRelTest.java b/...test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSideInputLookupJoinRelTest.java
@@ -34,6 +34,7 @@
 import org.apache.beam.sdk.values.POutput;
 import org.apache.beam.sdk.values.Row;
 import org.hamcrest.core.StringContains;
+import org.junit.Assert;
 import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
@@ -48,10 +49,17 @@ public class BeamSideInputLookupJoinRelTest extends BaseRelTest {
   /** Test table for JOIN-AS-LOOKUP. */
   public static class SiteLookupTable extends SchemaBaseBeamTable implements BeamSqlSeekableTable {
 
+    private Schema joinSubsetType;
+
     public SiteLookupTable(Schema schema) {
       super(schema);
     }
 
+    @Override
+    public void setJoinSubsetType(Schema joinSubsetType) {
+      this.joinSubsetType = joinSubsetType;
+    }
+
     @Override
     public PCollection.IsBounded isBounded() {
       return PCollection.IsBounded.BOUNDED;
@@ -69,6 +77,7 @@ public POutput buildIOWriter(PCollection<Row> input) {
 
     @Override
     public List<Row> seekRow(Row lookupSubRow) {
+      Assert.assertNotNull(joinSubsetType);
       if (lookupSubRow.getInt32("site_id") == 2) {
         return Arrays.asList(Row.withSchema(getSchema()).addValues(2, "SITE1").build());
       }