feat(graph-size-per-shard): Integrating MultiTableReadWithUniformPartitions

Author: VardhanThigle

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/IoWrapper.java

@@ -18,6 +18,7 @@
 import com.google.cloud.teleport.v2.source.reader.io.row.SourceRow;
 import com.google.cloud.teleport.v2.source.reader.io.schema.SourceSchema;
 import com.google.cloud.teleport.v2.source.reader.io.schema.SourceTableReference;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.values.PBegin;
@@ -27,7 +28,8 @@
 public interface IoWrapper {
 
   /** Get a list of reader transforms. */
-  ImmutableMap<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>> getTableReaders();
+  ImmutableMap<ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
+      getTableReaders();
 
   /** Discover source schema. */
   SourceSchema discoverTableSchema();

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/cassandra/iowrapper/CassandraIOWrapperHelper.java

@@ -142,14 +142,16 @@ static SourceSchema getSourceSchema(
     return sourceSchemaBuilder.build();
   }
 
-  static ImmutableMap<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>>
+  static ImmutableMap<
+          ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
       getTableReaders(DataSource dataSource, SourceSchema sourceSchema) {
     /*
      * TODO(vardhanvthigle): Plugin alternate implementation if needed.
      */
     CassandraTableReaderFactory cassandraTableReaderFactory =
         new CassandraTableReaderFactoryCassandraIoImpl();
-    ImmutableMap.Builder<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>>
+    ImmutableMap.Builder<
+            ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
         tableReadersBuilder = ImmutableMap.builder();
     SourceSchemaReference sourceSchemaReference = sourceSchema.schemaReference();
     sourceSchema
@@ -165,7 +167,7 @@ static SourceSchema getSourceSchema(
               var tableReader =
                   cassandraTableReaderFactory.getTableReader(
                       dataSource.cassandra(), sourceSchemaReference, tableSchema);
-              tableReadersBuilder.put(sourceTableReference, tableReader);
+              tableReadersBuilder.put(ImmutableList.of(sourceTableReference), tableReader);
             });
     return tableReadersBuilder.build();
   }

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/cassandra/iowrapper/CassandraIoWrapper.java

@@ -36,7 +36,8 @@
 /** IOWrapper for Cassandra Source. */
 public final class CassandraIoWrapper implements IoWrapper {
   private SourceSchema sourceSchema;
-  private ImmutableMap<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>>
+  private ImmutableMap<
+          ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
       tableReaders;
 
   public CassandraIoWrapper(
@@ -75,7 +76,8 @@ public CassandraIoWrapper(
 
   /** Get a list of reader transforms for Cassandra source. */
   @Override
-  public ImmutableMap<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>>
+  public ImmutableMap<
+          ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
       getTableReaders() {
     return tableReaders;
   }

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/iowrapper/JdbcIoWrapper.java

@@ -24,6 +24,9 @@
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.iowrapper.config.TableConfig;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.rowmapper.JdbcSourceRowMapper;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.PartitionColumn;
+import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.TableIdentifier;
+import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.TableReadSpecification;
+import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.TableSplitSpecification;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.transforms.ReadWithUniformPartitions;
 import com.google.cloud.teleport.v2.source.reader.io.row.SourceRow;
 import com.google.cloud.teleport.v2.source.reader.io.schema.SchemaDiscovery;
@@ -62,7 +65,8 @@
 public final class JdbcIoWrapper implements IoWrapper {
   private static final Logger LOG = LoggerFactory.getLogger(JdbcIoWrapper.class);
 
-  private final ImmutableMap<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>>
+  private final ImmutableMap<
+          ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
       tableReaders;
   private final SourceSchema sourceSchema;
 
@@ -90,8 +94,9 @@ public static JdbcIoWrapper of(JdbcIOWrapperConfig config) throws SuitableIndexN
         autoInferTableConfigs(config, schemaDiscovery, DataSource.ofJdbc(dataSource));
     SourceSchema sourceSchema =
         getSourceSchema(config, schemaDiscovery, DataSource.ofJdbc(dataSource), tableConfigs);
-    ImmutableMap<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>> tableReaders =
-        buildTableReaders(config, tableConfigs, dataSourceConfiguration, sourceSchema);
+    ImmutableMap<ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
+        tableReaders =
+            buildTableReaders(config, tableConfigs, dataSourceConfiguration, sourceSchema);
     return new JdbcIoWrapper(tableReaders, sourceSchema);
   }
 
@@ -155,7 +160,8 @@ private static void setConnectionProperty(
    * @return Read transforms.
    */
   @Override
-  public ImmutableMap<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>>
+  public ImmutableMap<
+          ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
       getTableReaders() {
     return this.tableReaders;
   }
@@ -170,50 +176,55 @@ public SourceSchema discoverTableSchema() {
     return this.sourceSchema;
   }
 
-  static ImmutableMap<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>>
+  static ImmutableMap<
+          ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
       buildTableReaders(
           JdbcIOWrapperConfig config,
           ImmutableList<TableConfig> tableConfigs,
           DataSourceConfiguration dataSourceConfiguration,
           SourceSchema sourceSchema) {
+    if (config.readWithUniformPartitionsFeatureEnabled() && !tableConfigs.isEmpty()) {
+      return getMultiTableReadWithUniformPartitionIO(
+          config,
+          dataSourceConfiguration,
+          sourceSchema.schemaReference(),
+          tableConfigs,
+          sourceSchema);
+    }
     return tableConfigs.stream()
         .map(
             tableConfig -> {
               SourceTableSchema sourceTableSchema =
                   findSourceTableSchema(sourceSchema, tableConfig);
-              long estimatedRowSize = sourceTableSchema.estimatedRowSize();
-              Integer calculatedFetchSize =
-                  FetchSizeCalculator.getFetchSize(
-                      tableConfig,
-                      estimatedRowSize,
-                      config.workerMemoryBytes(),
-                      config.workerCores());
-              int fetchSize = calculatedFetchSize;
-              return Map.entry(
+              int fetchSize = getFetchSize(config, tableConfig, sourceTableSchema);
+              SourceTableReference sourceTableReference =
                   SourceTableReference.builder()
                       .setSourceSchemaReference(sourceSchema.schemaReference())
                       .setSourceTableName(delimitIdentifier(sourceTableSchema.tableName()))
                       .setSourceTableSchemaUUID(sourceTableSchema.tableSchemaUUID())
-                      .build(),
-                  (config.readWithUniformPartitionsFeatureEnabled())
-                      ? getReadWithUniformPartitionIO(
-                          config,
-                          dataSourceConfiguration,
-                          sourceSchema.schemaReference(),
-                          tableConfig,
-                          sourceTableSchema,
-                          fetchSize)
-                      : getJdbcIO(
-                          config,
-                          dataSourceConfiguration,
-                          sourceSchema.schemaReference(),
-                          tableConfig,
-                          sourceTableSchema,
-                          fetchSize));
+                      .build();
+              return Map.entry(
+                  ImmutableList.of(sourceTableReference),
+                  getJdbcIO(
+                      config,
+                      dataSourceConfiguration,
+                      sourceSchema.schemaReference(),
+                      tableConfig,
+                      sourceTableSchema,
+                      fetchSize));
             })
         .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
   }
 
+  private static int getFetchSize(
+      JdbcIOWrapperConfig config, TableConfig tableConfig, SourceTableSchema sourceTableSchema) {
+    return FetchSizeCalculator.getFetchSize(
+        tableConfig,
+        sourceTableSchema.estimatedRowSize(),
+        config.workerMemoryBytes(),
+        config.workerCores());
+  }
+
   static SourceTableSchema findSourceTableSchema(
       SourceSchema sourceSchema, TableConfig tableConfig) {
     return sourceSchema.tableSchemas().stream()
@@ -468,54 +479,98 @@ private static PTransform<PBegin, PCollection<SourceRow>> getJdbcIO(
   }
 
   /**
-   * Private helper to construct {@link ReadWithUniformPartitions} as per the reader configuration.
+   * Private helper to construct {@link ReadWithUniformPartitions} for multiple tables as per the
+   * reader configuration.
    *
    * @param config Configuration.
    * @param dataSourceConfiguration dataSourceConfiguration (which is derived earlier from the
    *     reader configuration)
-   * @param tableConfig discovered table configurations.
-   * @param sourceTableSchema schema of the source table.
-   * @return
+   * @param sourceSchemaReference reference for the source schema.
+   * @param tableConfigs list of discovered table configurations.
+   * @param sourceSchema schema of the source.
+   * @return a map with a single entry where the key is a list of all table references and the value
+   *     is the multi-table reader transform.
    */
-  private static PTransform<PBegin, PCollection<SourceRow>> getReadWithUniformPartitionIO(
-      JdbcIOWrapperConfig config,
-      DataSourceConfiguration dataSourceConfiguration,
-      SourceSchemaReference sourceSchemaReference,
-      TableConfig tableConfig,
-      SourceTableSchema sourceTableSchema,
-      int fetchSize) {
+  private static ImmutableMap<
+          ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
+      getMultiTableReadWithUniformPartitionIO(
+          JdbcIOWrapperConfig config,
+          DataSourceConfiguration dataSourceConfiguration,
+          SourceSchemaReference sourceSchemaReference,
+          ImmutableList<TableConfig> tableConfigs,
+          SourceSchema sourceSchema) {
+
+    ImmutableList.Builder<TableSplitSpecification> splitSpecsBuilder = ImmutableList.builder();
+    ImmutableMap.Builder<TableIdentifier, TableReadSpecification<SourceRow>> readSpecsBuilder =
+        ImmutableMap.builder();
+    ImmutableList.Builder<SourceTableReference> tableReferencesBuilder = ImmutableList.builder();
+
+    for (TableConfig tableConfig : tableConfigs) {
+      SourceTableSchema sourceTableSchema = findSourceTableSchema(sourceSchema, tableConfig);
+      int fetchSize = getFetchSize(config, tableConfig, sourceTableSchema);
+      TableIdentifier tableIdentifier =
+          TableIdentifier.builder()
+              .setTableName(delimitIdentifier(tableConfig.tableName()))
+              .build();
+
+      TableSplitSpecification.Builder tableSplitSpecificationBuilder =
+          TableSplitSpecification.builder()
+              .setTableIdentifier(tableIdentifier)
+              .setPartitionColumns(tableConfig.partitionColumns())
+              .setApproxRowCount(tableConfig.approxRowCount());
+      if (tableConfig.maxPartitions() != null) {
+        tableSplitSpecificationBuilder =
+            tableSplitSpecificationBuilder.setMaxPartitionsHint((long) tableConfig.maxPartitions());
+      }
+      if (config.splitStageCountHint() >= 0) {
+        tableSplitSpecificationBuilder =
+            tableSplitSpecificationBuilder.setSplitStagesCount((long) config.splitStageCountHint());
+      }
+      splitSpecsBuilder.add(tableSplitSpecificationBuilder.build());
+
+      TableReadSpecification.Builder<SourceRow> tableReadSpecificationBuilder =
+          TableReadSpecification.<SourceRow>builder()
+              .setFetchSize(fetchSize)
+              .setTableIdentifier(tableIdentifier)
+              .setRowMapper(
+                  new JdbcSourceRowMapper(
+                      config.valueMappingsProvider(),
+                      sourceSchemaReference,
+                      sourceTableSchema,
+                      config.shardID()));
+      if (config.maxFetchSize() != null) {
+        tableReadSpecificationBuilder =
+            tableReadSpecificationBuilder.setFetchSize(config.maxFetchSize());
+      }
+      readSpecsBuilder.put(tableIdentifier, tableReadSpecificationBuilder.build());
+
+      tableReferencesBuilder.add(
+          SourceTableReference.builder()
+              .setSourceSchemaReference(sourceSchemaReference)
+              .setSourceTableName(delimitIdentifier(sourceTableSchema.tableName()))
+              .setSourceTableSchemaUUID(sourceTableSchema.tableSchemaUUID())
+              .build());
+    }
 
-    ReadWithUniformPartitions.Builder<SourceRow> readWithUniformPartitionsBuilder =
+    ReadWithUniformPartitions<SourceRow> readWithUniformPartitions =
         ReadWithUniformPartitions.<SourceRow>builder()
-            .setTableName(delimitIdentifier(tableConfig.tableName()))
-            .setPartitionColumns(tableConfig.partitionColumns())
+            .setTableSplitSpecifications(splitSpecsBuilder.build())
+            .setTableReadSpecifications(readSpecsBuilder.build())
             .setDataSourceProviderFn(JdbcIO.PoolableDataSourceProvider.of(dataSourceConfiguration))
             .setDbAdapter(config.dialectAdapter())
-            .setApproxTotalRowCount(tableConfig.approxRowCount())
-            .setFetchSize(fetchSize)
-            .setRowMapper(
-                new JdbcSourceRowMapper(
-                    config.valueMappingsProvider(),
-                    sourceSchemaReference,
-                    sourceTableSchema,
-                    config.shardID()))
             .setWaitOn(config.waitOn())
             .setDbParallelizationForSplitProcess(config.dbParallelizationForSplitProcess())
             .setDbParallelizationForReads(config.dbParallelizationForReads())
-            .setAdditionalOperationsOnRanges(config.additionalOperationsOnRanges());
+            .setAdditionalOperationsOnRanges(config.additionalOperationsOnRanges())
+            .build();
 
-    if (config.splitStageCountHint() >= 0) {
-      readWithUniformPartitionsBuilder =
-          readWithUniformPartitionsBuilder.setSplitStageCountHint(config.splitStageCountHint());
-    }
+    LOG.info(
+        "Configured Multi-Table ReadWithUniformPartitions {} for tables {} with config {}",
+        readWithUniformPartitions,
+        tableConfigs.stream().map(TableConfig::tableName).collect(Collectors.toList()),
+        config);
 
-    if (tableConfig.maxPartitions() != null) {
-      readWithUniformPartitionsBuilder =
-          readWithUniformPartitionsBuilder.setMaxPartitionsHint((long) tableConfig.maxPartitions());
-    }
-    ReadWithUniformPartitions readWithUniformPartitions = readWithUniformPartitionsBuilder.build();
-    LOG.info("Configured ReadWithUniformPartitions {} for {}", readWithUniformPartitions, config);
-    return readWithUniformPartitions;
+    return ImmutableMap.of(tableReferencesBuilder.build(), readWithUniformPartitions);
   }
 
   /**
@@ -542,7 +597,8 @@ private static DataSourceConfiguration getDataSourceConfiguration(JdbcIOWrapperC
    *     Beam classes like {@link JdbcIO}
    */
   private JdbcIoWrapper(
-      ImmutableMap<SourceTableReference, PTransform<PBegin, PCollection<SourceRow>>> tableReaders,
+      ImmutableMap<ImmutableList<SourceTableReference>, PTransform<PBegin, PCollection<SourceRow>>>
+          tableReaders,
       SourceSchema sourceSchema) {
     this.tableReaders = tableReaders;
     this.sourceSchema = sourceSchema;

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/iowrapper/config/JdbcIOWrapperConfig.java

@@ -35,6 +35,7 @@
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.Wait.OnSignal;
 import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.commons.dbcp2.BasicDataSource;
 
@@ -166,7 +167,8 @@ public JdbcSchemaReference jdbcSourceSchemaReference() {
    * JdbcIOWrapperConfig#readWithUniformPartitionsFeatureEnabled()} is false. Defaults to null.
    */
   @Nullable
-  public abstract PTransform<PCollection<ImmutableList<Range>>, ?> additionalOperationsOnRanges();
+  public abstract PTransform<PCollection<KV<Integer, ImmutableList<Range>>>, ?>
+      additionalOperationsOnRanges();
 
   /**
    * Sets the {@code testOnBorrow} property. This property determines whether or not the pool will
@@ -377,7 +379,7 @@ public abstract Builder setTableVsPartitionColumns(
     public abstract Builder setDbParallelizationForReads(@Nullable Integer value);
 
     public abstract Builder setAdditionalOperationsOnRanges(
-        @Nullable PTransform<PCollection<ImmutableList<Range>>, ?> value);
+        PTransform<PCollection<KV<Integer, ImmutableList<Range>>>, ?> value);
 
     public abstract Builder setTestOnBorrow(Boolean value);
 

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/uniformsplitter/columnboundary/ColumnForBoundaryQuery.java

@@ -19,6 +19,7 @@
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.UniformSplitterDBAdapter;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.PartitionColumn;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.Range;
+import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.TableIdentifier;
 import java.io.Serializable;
 import javax.annotation.Nullable;
 
@@ -29,6 +30,9 @@
 @AutoValue
 public abstract class ColumnForBoundaryQuery implements Serializable {
 
+  /** Identifier of the table for a given boundary. */
+  public abstract TableIdentifier tableIdentifier();
+
   /**
    * @return column details.
    */
@@ -72,6 +76,8 @@ public Class columnClass() {
   @AutoValue.Builder
   public abstract static class Builder {
 
+    public abstract Builder setTableIdentifier(TableIdentifier value);
+
     public abstract Builder setPartitionColumn(PartitionColumn value);
 
     public abstract PartitionColumn.Builder partitionColumnBuilder();