[Feature][Transform-V2] Spark support transform with multi-table (#8340)

apache · Dec 23, 2024 · e128ccc · e128ccc
1 parent a53d809
commit e128ccc
Show file tree

Hide file tree

Showing 17 changed files with 84 additions and 132 deletions.
diff --git a/...ain/java/org/apache/seatunnel/core/starter/spark/execution/TransformExecuteProcessor.java b/...ain/java/org/apache/seatunnel/core/starter/spark/execution/TransformExecuteProcessor.java
@@ -22,9 +22,9 @@
 import org.apache.seatunnel.api.common.JobContext;
 import org.apache.seatunnel.api.configuration.ReadonlyConfig;
 import org.apache.seatunnel.api.configuration.util.ConfigValidator;
+import org.apache.seatunnel.api.table.catalog.CatalogTable;
 import org.apache.seatunnel.api.table.factory.TableTransformFactory;
 import org.apache.seatunnel.api.table.factory.TableTransformFactoryContext;
-import org.apache.seatunnel.api.table.type.SeaTunnelDataType;
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
 import org.apache.seatunnel.api.transform.SeaTunnelFlatMapTransform;
 import org.apache.seatunnel.api.transform.SeaTunnelMapTransform;
@@ -34,23 +34,20 @@
 import org.apache.seatunnel.plugin.discovery.seatunnel.SeaTunnelFactoryDiscovery;
 import org.apache.seatunnel.plugin.discovery.seatunnel.SeaTunnelTransformPluginDiscovery;
 import org.apache.seatunnel.translation.spark.execution.DatasetTableInfo;
-import org.apache.seatunnel.translation.spark.serialization.SeaTunnelRowConverter;
-import org.apache.seatunnel.translation.spark.utils.TypeConverterUtils;
+import org.apache.seatunnel.translation.spark.execution.MultiTableManager;
 
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.spark.api.java.function.FlatMapFunction;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder;
 import org.apache.spark.sql.catalyst.encoders.RowEncoder;
-import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema;
-import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.catalyst.expressions.GenericRow;
 
 import lombok.extern.slf4j.Slf4j;
 
 import java.net.URL;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -141,7 +138,7 @@ public List<DatasetTableInfo> execute(List<DatasetTableInfo> upstreamDataStreams
                         pluginOutputIdentifier,
                         new DatasetTableInfo(
                                 inputDataset,
-                                Collections.singletonList(transform.getProducedCatalogTable()),
+                                transform.getProducedCatalogTables(),
                                 pluginOutputIdentifier));
             } catch (Exception e) {
                 throw new TaskExecuteException(
@@ -155,55 +152,52 @@ public List<DatasetTableInfo> execute(List<DatasetTableInfo> upstreamDataStreams
     }
 
     private Dataset<Row> sparkTransform(SeaTunnelTransform transform, DatasetTableInfo tableInfo) {
+        MultiTableManager inputManager =
+                new MultiTableManager(tableInfo.getCatalogTables().toArray(new CatalogTable[0]));
+        MultiTableManager outputManager =
+                new MultiTableManager(
+                        (CatalogTable[])
+                                transform.getProducedCatalogTables().toArray(new CatalogTable[0]));
         Dataset<Row> stream = tableInfo.getDataset();
-        SeaTunnelDataType<?> inputDataType =
-                tableInfo.getCatalogTables().get(0).getSeaTunnelRowType();
-        SeaTunnelDataType<?> outputDataTYpe =
-                transform.getProducedCatalogTable().getSeaTunnelRowType();
-        StructType outputSchema = (StructType) TypeConverterUtils.parcel(outputDataTYpe);
-        SeaTunnelRowConverter inputRowConverter = new SeaTunnelRowConverter(inputDataType);
-        SeaTunnelRowConverter outputRowConverter = new SeaTunnelRowConverter(outputDataTYpe);
-        ExpressionEncoder<Row> encoder = RowEncoder.apply(outputSchema);
-
+        ExpressionEncoder<Row> encoder = RowEncoder.apply(outputManager.getTableSchema());
         return stream.flatMap(
-                        new TransformMapPartitionsFunction(
-                                transform, inputRowConverter, outputRowConverter),
+                        new TransformMapPartitionsFunction(transform, inputManager, outputManager),
                         encoder)
                 .filter(Objects::nonNull);
     }
 
     private static class TransformMapPartitionsFunction implements FlatMapFunction<Row, Row> {
         private SeaTunnelTransform<SeaTunnelRow> transform;
-        private SeaTunnelRowConverter inputRowConverter;
-        private SeaTunnelRowConverter outputRowConverter;
+        private MultiTableManager inputManager;
+        private MultiTableManager outputManager;
 
         public TransformMapPartitionsFunction(
                 SeaTunnelTransform<SeaTunnelRow> transform,
-                SeaTunnelRowConverter inputRowConverter,
-                SeaTunnelRowConverter outputRowConverter) {
+                MultiTableManager inputManager,
+                MultiTableManager outputManager) {
             this.transform = transform;
-            this.inputRowConverter = inputRowConverter;
-            this.outputRowConverter = outputRowConverter;
+            this.inputManager = inputManager;
+            this.outputManager = outputManager;
         }
 
         @Override
         public Iterator<Row> call(Row row) throws Exception {
             List<Row> rows = new ArrayList<>();
 
-            SeaTunnelRow seaTunnelRow = inputRowConverter.unpack((GenericRowWithSchema) row);
+            SeaTunnelRow seaTunnelRow = inputManager.reconvert((GenericRow) row);
             if (transform instanceof SeaTunnelFlatMapTransform) {
                 List<SeaTunnelRow> seaTunnelRows =
                         ((SeaTunnelFlatMapTransform<SeaTunnelRow>) transform).flatMap(seaTunnelRow);
                 if (CollectionUtils.isNotEmpty(seaTunnelRows)) {
                     for (SeaTunnelRow seaTunnelRowTransform : seaTunnelRows) {
-                        rows.add(outputRowConverter.parcel(seaTunnelRowTransform));
+                        rows.add(outputManager.convert(seaTunnelRowTransform));
                     }
                 }
             } else if (transform instanceof SeaTunnelMapTransform) {
                 SeaTunnelRow seaTunnelRowTransform =
                         ((SeaTunnelMapTransform<SeaTunnelRow>) transform).map(seaTunnelRow);
                 if (seaTunnelRowTransform != null) {
-                    rows.add(outputRowConverter.parcel(seaTunnelRowTransform));
+                    rows.add(outputManager.convert(seaTunnelRowTransform));
                 }
             }
             return rows.iterator();

diff --git a/...transforms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestCopyIT.java b/...transforms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestCopyIT.java
@@ -17,9 +17,7 @@
 
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -35,10 +33,6 @@ public void testCopy(TestContainer container) throws IOException, InterruptedExc
         Assertions.assertEquals(0, execResult.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testCopyMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...forms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestEmbeddingIT.java b/...forms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestEmbeddingIT.java
@@ -93,10 +93,6 @@ public void testEmbedding(TestContainer container) throws IOException, Interrupt
         Assertions.assertEquals(0, execResult.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testEmbeddingMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...ansforms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestFilterIT.java b/...ansforms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestFilterIT.java
@@ -17,9 +17,7 @@
 
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -35,10 +33,6 @@ public void testFilter(TestContainer container) throws IOException, InterruptedE
         Assertions.assertEquals(0, execResult.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testFilterMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...s-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestFilterRowKindIT.java b/...s-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestFilterRowKindIT.java
@@ -17,9 +17,7 @@
 
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -43,10 +41,6 @@ public void testFilterRowKind(TestContainer container)
         Assertions.assertEquals(0, execResult3.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testFilterRowKindMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...-transforms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestLLMIT.java b/...-transforms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestLLMIT.java
@@ -18,9 +18,7 @@
 package org.apache.seatunnel.e2e.transform;
 
 import org.apache.seatunnel.e2e.common.TestResource;
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
@@ -90,10 +88,6 @@ public void testLLMWithOpenAI(TestContainer container)
         Assertions.assertEquals(0, execResult.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testLLMWithOpenAIMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...t-1/src/test/java/org/apache/seatunnel/e2e/transform/TestRowKindExtractorTransformIT.java b/...t-1/src/test/java/org/apache/seatunnel/e2e/transform/TestRowKindExtractorTransformIT.java
@@ -17,9 +17,7 @@
 
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -40,10 +38,6 @@ public void testRowKindExtractorTransform(TestContainer container)
         Assertions.assertEquals(0, execResult2.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testRowKindExtractorMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...ransforms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestSplitIT.java b/...ransforms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestSplitIT.java
@@ -17,9 +17,7 @@
 
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -35,10 +33,6 @@ public void testSplit(TestContainer container) throws IOException, InterruptedEx
         Assertions.assertEquals(0, execResult.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testSplitMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestDynamicCompileIT.java b/...-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestDynamicCompileIT.java
@@ -19,9 +19,7 @@
 
 import org.apache.seatunnel.e2e.common.TestResource;
 import org.apache.seatunnel.e2e.common.container.ContainerExtendedFactory;
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 import org.apache.seatunnel.e2e.common.junit.TestContainerExtension;
 
 import org.junit.jupiter.api.AfterAll;
@@ -118,10 +116,6 @@ public void testDynamicSingleCompileJava(TestContainer container)
         Assertions.assertEquals(0, execResult.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testDynamicSingleCompileJavaMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...rms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestFieldMapperIT.java b/...rms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestFieldMapperIT.java
@@ -17,9 +17,7 @@
 
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -39,10 +37,6 @@ public void testFieldMapper(TestContainer container) throws IOException, Interru
         Assertions.assertEquals(0, execResult1.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testFieldMapperMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestJsonPathTransformIT.java b/...-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestJsonPathTransformIT.java
@@ -16,9 +16,7 @@
  */
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -33,10 +31,6 @@ public void testBasicType(TestContainer container) throws Exception {
         Assertions.assertEquals(0, execResult.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testBasicTypeMultiTable(TestContainer container) throws Exception {
         Container.ExecResult execResult =

diff --git a/...sforms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestMetadataIT.java b/...sforms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestMetadataIT.java
@@ -17,9 +17,7 @@
 
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -29,10 +27,6 @@
 
 public class TestMetadataIT extends TestSuiteBase {
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testMetadataMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...ansforms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestRenameIT.java b/...ansforms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestRenameIT.java
@@ -17,9 +17,7 @@
 
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -29,10 +27,6 @@
 
 public class TestRenameIT extends TestSuiteBase {
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testRenameMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...nsforms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestReplaceIT.java b/...nsforms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestReplaceIT.java
@@ -17,9 +17,7 @@
 
 package org.apache.seatunnel.e2e.transform;
 
-import org.apache.seatunnel.e2e.common.container.EngineType;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
-import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestTemplate;
@@ -35,10 +33,6 @@ public void testReplace(TestContainer container) throws IOException, Interrupted
         Assertions.assertEquals(0, execResult.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testReplaceMultiTable(TestContainer container)
             throws IOException, InterruptedException {

diff --git a/...-transforms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestSQLIT.java b/...-transforms-v2-e2e-part-2/src/test/java/org/apache/seatunnel/e2e/transform/TestSQLIT.java
@@ -79,10 +79,6 @@ public void testSQLTransform(TestContainer container) throws IOException, Interr
         Assertions.assertEquals(0, splitSql.getExitCode());
     }
 
-    @DisabledOnContainer(
-            value = {},
-            type = {EngineType.SPARK},
-            disabledReason = "Currently SPARK do not multi table transform")
     @TestTemplate
     public void testSQLTransformMultiTable(TestContainer container)
             throws IOException, InterruptedException {