From 462dadab7c5fe36e314bc81243f57f84cab2f248 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 9 Sep 2026 15:04:34 +0000 Subject: [PATCH] Implement Iceberg Side-Input Table Cache Integration with Fallback --- .../AssignDestinationsAndPartitions.java | 85 +++- .../sdk/io/iceberg/RecordWriterManager.java | 56 ++- .../io/iceberg/WriteDirectRowsToFiles.java | 44 +- .../io/iceberg/WriteGroupedRowsToFiles.java | 47 ++- .../iceberg/WritePartitionedRowsToFiles.java | 66 ++- .../sdk/io/iceberg/WriteToDestinations.java | 31 +- .../sdk/io/iceberg/WriteToPartitions.java | 21 +- .../io/iceberg/WriteUngroupedRowsToFiles.java | 72 +++- .../AssignDestinationsAndPartitionsTest.java | 200 +++++++++ .../io/iceberg/RecordWriterManagerTest.java | 104 +++++ .../io/iceberg/WriteWithMetadataViewTest.java | 397 ++++++++++++++++++ 11 files changed, 1073 insertions(+), 50 deletions(-) create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitionsTest.java create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/WriteWithMetadataViewTest.java diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java index a744ff930975..58341ba4f95f 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java @@ -30,11 +30,13 @@ import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.ValueInSingleWindow; import org.apache.iceberg.PartitionKey; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; import org.apache.iceberg.exceptions.NoSuchTableException; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; @@ -51,6 +53,7 @@ class AssignDestinationsAndPartitions private final DynamicDestinations dynamicDestinations; private final IcebergCatalogConfig catalogConfig; + private final @Nullable PCollectionView> metadataView; static final String DESTINATION = "destination"; static final String PARTITION = "partition"; @@ -63,14 +66,27 @@ class AssignDestinationsAndPartitions public AssignDestinationsAndPartitions( DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig) { + this(dynamicDestinations, catalogConfig, null); + } + + public AssignDestinationsAndPartitions( + DynamicDestinations dynamicDestinations, + IcebergCatalogConfig catalogConfig, + @Nullable PCollectionView> metadataView) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; + this.metadataView = metadataView; } @Override public PCollection> expand(PCollection input) { + ParDo.SingleOutput> parDo = + ParDo.of(new AssignDoFn(dynamicDestinations, catalogConfig, metadataView)); + if (metadataView != null) { + parDo = parDo.withSideInputs(metadataView); + } return input - .apply(ParDo.of(new AssignDoFn(dynamicDestinations, catalogConfig))) + .apply(parDo) .setCoder( KvCoder.of( RowCoder.of(OUTPUT_SCHEMA), RowCoder.of(dynamicDestinations.getDataSchema()))); @@ -83,13 +99,23 @@ static class AssignDoFn extends DoFn> { private transient @MonotonicNonNull Map partitionKeys; private transient @MonotonicNonNull Map wrappers; private transient @MonotonicNonNull Map lastRefreshTimes; + private transient @MonotonicNonNull Map cachedSpecIds; private final DynamicDestinations dynamicDestinations; private final IcebergCatalogConfig catalogConfig; + private final @Nullable PCollectionView> metadataView; AssignDoFn(DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig) { + this(dynamicDestinations, catalogConfig, null); + } + + AssignDoFn( + DynamicDestinations dynamicDestinations, + IcebergCatalogConfig catalogConfig, + @Nullable PCollectionView> metadataView) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; + this.metadataView = metadataView; } @Setup @@ -97,10 +123,12 @@ public void setup() { this.wrappers = new HashMap<>(); this.partitionKeys = new HashMap<>(); this.lastRefreshTimes = new HashMap<>(); + this.cachedSpecIds = new HashMap<>(); } @ProcessElement public void processElement( + ProcessContext c, @Element Row element, BoundedWindow window, PaneInfo paneInfo, @@ -111,58 +139,83 @@ public void processElement( dynamicDestinations.getTableStringIdentifier( ValueInSingleWindow.of(element, timestamp, window, paneInfo)); + String canonicalTableId; + try { + canonicalTableId = + IcebergUtils.tableIdentifierToString( + IcebergUtils.parseTableIdentifier(tableIdentifier)); + } catch (Exception e) { + canonicalTableId = tableIdentifier.trim(); + } + + SerializableTableSpec tableSpec = null; + if (metadataView != null) { + Map viewMap = c.sideInput(metadataView); + if (viewMap != null) { + tableSpec = viewMap.get(canonicalTableId); + if (tableSpec == null) { + tableSpec = viewMap.get(tableIdentifier); + } + } + } + Row data = dynamicDestinations.getData(element); @Nullable PartitionKey partitionKey = checkStateNotNull(partitionKeys).get(tableIdentifier); - @Nullable BeamRowWrapper wrapper = checkStateNotNull(wrappers).get(tableIdentifier); - @Nullable Instant lastRefresh = checkStateNotNull(lastRefreshTimes).get(tableIdentifier); + @Nullable Integer cachedSpecId = checkStateNotNull(cachedSpecIds).get(tableIdentifier); Instant now = Instant.now(); + boolean specChanged = + tableSpec != null + && (cachedSpecId == null || !cachedSpecId.equals(tableSpec.getSpecId())); + boolean shouldRefresh = partitionKey == null || wrapper == null - || lastRefresh == null - || now.isAfter(lastRefresh.plus(REFRESH_INTERVAL)); + || specChanged + || (tableSpec == null + && (lastRefresh == null || now.isAfter(lastRefresh.plus(REFRESH_INTERVAL)))); if (shouldRefresh) { PartitionSpec spec = PartitionSpec.unpartitioned(); - Schema schema = IcebergUtils.beamSchemaToIcebergSchema(data.getSchema()); @Nullable IcebergTableCreateConfig createConfig = dynamicDestinations.instantiateDestination(tableIdentifier).getTableCreateConfig(); if (createConfig != null && createConfig.getPartitionFields() != null) { - spec = PartitionUtils.toPartitionSpec(createConfig.getPartitionFields(), data.getSchema()); - + } else if (tableSpec != null) { + spec = tableSpec.getPartitionSpec(); + if (data.getSchema().getFieldCount() == tableSpec.getSchema().columns().size()) { + schema = tableSpec.getSchema(); + } + checkStateNotNull(cachedSpecIds).put(tableIdentifier, tableSpec.getSpecId()); } else { - try { // see if table already exists with a spec - spec = + Table table = TableCache.getAndRefreshIfStale( - catalogConfig, IcebergUtils.parseTableIdentifier(tableIdentifier)) - .spec(); - + catalogConfig, IcebergUtils.parseTableIdentifier(tableIdentifier)); + spec = table.spec(); + if (data.getSchema().getFieldCount() == table.schema().columns().size()) { + schema = table.schema(); + } } catch (NoSuchTableException ignored) { // no partition to apply } } partitionKey = new PartitionKey(spec, schema); - wrapper = new BeamRowWrapper(data.getSchema(), schema.asStruct()); checkStateNotNull(partitionKeys).put(tableIdentifier, partitionKey); - checkStateNotNull(wrappers).put(tableIdentifier, wrapper); - checkStateNotNull(lastRefreshTimes).put(tableIdentifier, now); } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java index 0995e7a6102d..25e5a13da43e 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java @@ -27,6 +27,7 @@ import java.time.format.DateTimeFormatter; import java.time.temporal.ChronoUnit; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.UUID; @@ -251,6 +252,7 @@ static String getPartitionDataPath( private final long maxFileSize; private final int maxNumWriters; private final @Nullable Map writeProperties; + private volatile @Nullable Map sideInputTableSpecs; @VisibleForTesting int openWriters = 0; @VisibleForTesting @@ -263,7 +265,7 @@ static String getPartitionDataPath( RecordWriterManager( IcebergCatalogConfig catalogConfig, String filePrefix, long maxFileSize, int maxNumWriters) { - this(catalogConfig, filePrefix, maxFileSize, maxNumWriters, null); + this(catalogConfig, filePrefix, maxFileSize, maxNumWriters, null, null); } RecordWriterManager( @@ -272,11 +274,27 @@ static String getPartitionDataPath( long maxFileSize, int maxNumWriters, @Nullable Map writeProperties) { + this(catalogConfig, filePrefix, maxFileSize, maxNumWriters, writeProperties, null); + } + + RecordWriterManager( + IcebergCatalogConfig catalogConfig, + String filePrefix, + long maxFileSize, + int maxNumWriters, + @Nullable Map writeProperties, + @Nullable Map sideInputTableSpecs) { this.catalogConfig = catalogConfig; this.filePrefix = filePrefix; this.maxFileSize = maxFileSize; this.maxNumWriters = maxNumWriters; this.writeProperties = writeProperties; + this.sideInputTableSpecs = sideInputTableSpecs; + } + + @VisibleForTesting + void setSideInputTableSpecs(@Nullable Map sideInputTableSpecs) { + this.sideInputTableSpecs = sideInputTableSpecs; } /** @@ -291,12 +309,29 @@ static String getPartitionDataPath( * using the Iceberg API. */ @VisibleForTesting - Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) { + Table getOrCreateTable( + IcebergDestination destination, + Schema dataSchema, + @Nullable Map sideInputTableSpecs) { TableIdentifier identifier = destination.getTableIdentifier(); + String tableIdString = IcebergUtils.tableIdentifierToString(identifier); + if (sideInputTableSpecs != null && sideInputTableSpecs.containsKey(tableIdString)) { + SerializableTableSpec spec = sideInputTableSpecs.get(tableIdString); + if (spec != null) { + Map catalogProperties = catalogConfig.getCatalogProperties(); + return new SideInputTable( + spec, catalogProperties != null ? catalogProperties : Collections.emptyMap()); + } + } return TableCache.getAndRefreshIfStale( catalogConfig, identifier, () -> loadOrCreateTable(destination, dataSchema)); } + @VisibleForTesting + Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) { + return getOrCreateTable(destination, dataSchema, this.sideInputTableSpecs); + } + private Table loadOrCreateTable(IcebergDestination destination, Schema dataSchema) { Catalog catalog = catalogConfig.catalog(); TableIdentifier identifier = destination.getTableIdentifier(); @@ -354,6 +389,23 @@ private Table loadOrCreateTable(IcebergDestination destination, Schema dataSchem } } + /** + * Fetches the appropriate {@link RecordWriter} for this destination and partition and writes the + * record, optionally updating the side-input table specs map. + * + *

If the {@link RecordWriterManager} is saturated (i.e. has hit the maximum limit of open + * writers), the record is rejected and {@code false} is returned. + */ + public boolean write( + WindowedValue icebergDestination, + Row row, + @Nullable Map sideInputTableSpecs) { + if (sideInputTableSpecs != null) { + this.sideInputTableSpecs = sideInputTableSpecs; + } + return write(icebergDestination, row); + } + /** * Fetches the appropriate {@link RecordWriter} for this destination and partition and writes the * record. diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteDirectRowsToFiles.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteDirectRowsToFiles.java index e03085e6be78..8bafe9eeaef9 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteDirectRowsToFiles.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteDirectRowsToFiles.java @@ -26,6 +26,7 @@ import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; @@ -41,6 +42,7 @@ class WriteDirectRowsToFiles private final String filePrefix; private final long maxBytesPerFile; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; WriteDirectRowsToFiles( IcebergCatalogConfig catalogConfig, @@ -48,19 +50,39 @@ class WriteDirectRowsToFiles String filePrefix, long maxBytesPerFile, @Nullable Map writeProperties) { + this(catalogConfig, dynamicDestinations, filePrefix, maxBytesPerFile, writeProperties, null); + } + + WriteDirectRowsToFiles( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + long maxBytesPerFile, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxBytesPerFile = maxBytesPerFile; this.writeProperties = writeProperties; + this.metadataView = metadataView; } @Override public PCollection expand(PCollection> input) { - return input.apply( + ParDo.SingleOutput, FileWriteResult> parDo = ParDo.of( new WriteDirectRowsToFilesDoFn( - catalogConfig, dynamicDestinations, maxBytesPerFile, filePrefix, writeProperties))); + catalogConfig, + dynamicDestinations, + maxBytesPerFile, + filePrefix, + writeProperties, + metadataView)); + if (metadataView != null) { + parDo = parDo.withSideInputs(metadataView); + } + return input.apply(parDo); } private static class WriteDirectRowsToFilesDoFn extends DoFn, FileWriteResult> { @@ -70,6 +92,7 @@ private static class WriteDirectRowsToFilesDoFn extends DoFn, Fi private final String filePrefix; private final long maxFileSize; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; private transient @Nullable RecordWriterManager recordWriterManager; WriteDirectRowsToFilesDoFn( @@ -78,11 +101,22 @@ private static class WriteDirectRowsToFilesDoFn extends DoFn, Fi long maxFileSize, String filePrefix, @Nullable Map writeProperties) { + this(catalogConfig, dynamicDestinations, maxFileSize, filePrefix, writeProperties, null); + } + + WriteDirectRowsToFilesDoFn( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + long maxFileSize, + String filePrefix, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxFileSize = maxFileSize; this.writeProperties = writeProperties; + this.metadataView = metadataView; this.recordWriterManager = null; } @@ -95,7 +129,7 @@ public void startBundle() { @ProcessElement public void processElement( - @SuppressWarnings("unused") ProcessContext context, + ProcessContext context, @Element KV element, BoundedWindow window, PaneInfo paneInfo) @@ -104,8 +138,10 @@ public void processElement( IcebergDestination destination = dynamicDestinations.instantiateDestination(tableIdentifier); WindowedValue windowedDestination = WindowedValues.of(destination, window.maxTimestamp(), window, paneInfo); + Map sideInputs = + metadataView != null ? context.sideInput(metadataView) : null; Preconditions.checkNotNull(recordWriterManager) - .write(windowedDestination, element.getValue()); + .write(windowedDestination, element.getValue(), sideInputs); } @FinishBundle diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteGroupedRowsToFiles.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteGroupedRowsToFiles.java index e74715a7eebc..a2ed40c87f99 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteGroupedRowsToFiles.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteGroupedRowsToFiles.java @@ -27,6 +27,7 @@ import org.apache.beam.sdk.util.ShardedKey; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; @@ -42,6 +43,7 @@ class WriteGroupedRowsToFiles private final IcebergCatalogConfig catalogConfig; private final String filePrefix; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; WriteGroupedRowsToFiles( IcebergCatalogConfig catalogConfig, @@ -49,20 +51,40 @@ class WriteGroupedRowsToFiles String filePrefix, long maxBytesPerFile, @Nullable Map writeProperties) { + this(catalogConfig, dynamicDestinations, filePrefix, maxBytesPerFile, writeProperties, null); + } + + WriteGroupedRowsToFiles( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + long maxBytesPerFile, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxBytesPerFile = maxBytesPerFile; this.writeProperties = writeProperties; + this.metadataView = metadataView; } @Override public PCollection expand( PCollection, Iterable>> input) { - return input.apply( + ParDo.SingleOutput, Iterable>, FileWriteResult> parDo = ParDo.of( new WriteGroupedRowsToFilesDoFn( - catalogConfig, dynamicDestinations, maxBytesPerFile, filePrefix, writeProperties))); + catalogConfig, + dynamicDestinations, + maxBytesPerFile, + filePrefix, + writeProperties, + metadataView)); + if (metadataView != null) { + parDo = parDo.withSideInputs(metadataView); + } + return input.apply(parDo); } private static class WriteGroupedRowsToFilesDoFn @@ -73,6 +95,7 @@ private static class WriteGroupedRowsToFilesDoFn private final String filePrefix; private final long maxFileSize; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; WriteGroupedRowsToFilesDoFn( IcebergCatalogConfig catalogConfig, @@ -80,11 +103,22 @@ private static class WriteGroupedRowsToFilesDoFn long maxFileSize, String filePrefix, @Nullable Map writeProperties) { + this(catalogConfig, dynamicDestinations, maxFileSize, filePrefix, writeProperties, null); + } + + WriteGroupedRowsToFilesDoFn( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + long maxFileSize, + String filePrefix, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxFileSize = maxFileSize; this.writeProperties = writeProperties; + this.metadataView = metadataView; } @ProcessElement @@ -99,10 +133,17 @@ public void processElement( IcebergDestination destination = dynamicDestinations.instantiateDestination(tableIdentifier); WindowedValue windowedDestination = WindowedValues.of(destination, window.maxTimestamp(), window, paneInfo); + Map sideInputs = + metadataView != null ? c.sideInput(metadataView) : null; RecordWriterManager writer; try (RecordWriterManager openWriter = new RecordWriterManager( - catalogConfig, filePrefix, maxFileSize, Integer.MAX_VALUE, writeProperties)) { + catalogConfig, + filePrefix, + maxFileSize, + Integer.MAX_VALUE, + writeProperties, + sideInputs)) { writer = openWriter; for (Row e : element.getValue()) { writer.write(windowedDestination, e); diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java index 338a2162080b..881d2577fad6 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java @@ -22,6 +22,7 @@ import static org.apache.beam.sdk.io.iceberg.RecordWriterManager.getPartitionDataPath; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import java.util.Collections; import java.util.Map; import java.util.UUID; import org.apache.beam.sdk.coders.IterableCoder; @@ -33,6 +34,7 @@ import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.Row; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.iceberg.DataFiles; @@ -61,16 +63,27 @@ class WritePartitionedRowsToFiles private final IcebergCatalogConfig catalogConfig; private final String filePrefix; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; WritePartitionedRowsToFiles( IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, String filePrefix, @Nullable Map writeProperties) { + this(catalogConfig, dynamicDestinations, filePrefix, writeProperties, null); + } + + WritePartitionedRowsToFiles( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.writeProperties = writeProperties; + this.metadataView = metadataView; } @Override @@ -81,10 +94,19 @@ public PCollection expand(PCollection>> i ((KvCoder>) input.getCoder()).getValueCoder()) .getElemCoder()) .getSchema(); - return input.apply( + ParDo.SingleOutput>, FileWriteResult> parDo = ParDo.of( new WriteDoFn( - catalogConfig, dynamicDestinations, filePrefix, dataSchema, writeProperties))); + catalogConfig, + dynamicDestinations, + filePrefix, + dataSchema, + writeProperties, + metadataView)); + if (metadataView != null) { + parDo = parDo.withSideInputs(metadataView); + } + return input.apply(parDo); } private static class WriteDoFn extends DoFn>, FileWriteResult> { @@ -94,6 +116,7 @@ private static class WriteDoFn extends DoFn>, FileWriteRes private final String filePrefix; private final Schema dataSchema; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; private transient @MonotonicNonNull Map specIds; private transient @MonotonicNonNull Map> partitionFieldMaps; @@ -104,11 +127,22 @@ private static class WriteDoFn extends DoFn>, FileWriteRes String filePrefix, Schema dataSchema, @Nullable Map writeProperties) { + this(catalogConfig, dynamicDestinations, filePrefix, dataSchema, writeProperties, null); + } + + WriteDoFn( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + Schema dataSchema, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.dataSchema = dataSchema; this.writeProperties = writeProperties; + this.metadataView = metadataView; } @Setup @@ -119,13 +153,17 @@ public void setup() { @ProcessElement public void processElement( - @Element KV> element, OutputReceiver out) + ProcessContext c, + @Element KV> element, + OutputReceiver out) throws Exception { String tableIdentifier = checkStateNotNull(element.getKey().getString(DESTINATION)); String partitionPath = checkStateNotNull(element.getKey().getString(PARTITION)); IcebergDestination destination = dynamicDestinations.instantiateDestination(tableIdentifier); - Table table = getOrCreateTable(destination, dataSchema); + Map sideInputs = + metadataView != null ? c.sideInput(metadataView) : null; + Table table = getOrCreateTable(destination, dataSchema, sideInputs); partitionPath = getPartitionDataPath( partitionPath, getPartitionFieldMap(destination.getTableIdentifier(), table)); @@ -176,14 +214,30 @@ private Map getPartitionFieldMap( return partitionFieldMap; } - Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) { + Table getOrCreateTable( + IcebergDestination destination, + Schema dataSchema, + @Nullable Map sideInputTableSpecs) { TableIdentifier identifier = destination.getTableIdentifier(); + String tableIdString = IcebergUtils.tableIdentifierToString(identifier); + if (sideInputTableSpecs != null && sideInputTableSpecs.containsKey(tableIdString)) { + SerializableTableSpec spec = sideInputTableSpecs.get(tableIdString); + if (spec != null) { + Map catalogProperties = catalogConfig.getCatalogProperties(); + return new SideInputTable( + spec, catalogProperties != null ? catalogProperties : Collections.emptyMap()); + } + } return TableCache.getAndRefreshIfStale( catalogConfig, identifier, () -> loadOrCreateTable(catalogConfig.catalog(), destination, dataSchema)); } + Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) { + return getOrCreateTable(destination, dataSchema, null); + } + private Table loadOrCreateTable( Catalog catalog, IcebergDestination destination, Schema dataSchema) { TableIdentifier identifier = destination.getTableIdentifier(); @@ -207,7 +261,7 @@ private Table loadOrCreateTable( LOG.info("Created new namespace '{}'.", namespace); } catch (AlreadyExistsException ignored) { // race condition: another worker already created this namespace - LOG.info("Namespace `{}` already exists.", namespace); + LOG.info("Namespace '{}' already exists.", namespace); } } } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToDestinations.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToDestinations.java index 684ef350a20f..6d8f08a2ae56 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToDestinations.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToDestinations.java @@ -41,6 +41,7 @@ import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionList; import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; @@ -60,6 +61,7 @@ class WriteToDestinations extends PTransform>, Icebe private final String filePrefix; private final @Nullable Integer directWriteByteLimit; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; WriteToDestinations( IcebergCatalogConfig catalogConfig, @@ -67,11 +69,28 @@ class WriteToDestinations extends PTransform>, Icebe @Nullable Duration triggeringFrequency, @Nullable Integer directWriteByteLimit, @Nullable Map writeProperties) { + this( + catalogConfig, + dynamicDestinations, + triggeringFrequency, + directWriteByteLimit, + writeProperties, + null); + } + + WriteToDestinations( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Duration triggeringFrequency, + @Nullable Integer directWriteByteLimit, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; this.triggeringFrequency = triggeringFrequency; this.directWriteByteLimit = directWriteByteLimit; this.writeProperties = writeProperties; + this.metadataView = metadataView; // single unique prefix per write transform this.filePrefix = UUID.randomUUID().toString(); } @@ -119,7 +138,8 @@ private PCollection groupAndWriteRecords(PCollection applyUserTriggering(PCollection input) { @@ -168,7 +188,8 @@ private PCollection writeTriggeredWithBundleLifting( dynamicDestinations, filePrefix, DEFAULT_MAX_BYTES_PER_FILE, - writeProperties)); + writeProperties, + metadataView)); PCollection groupedFileWrites = groupAndWriteRecords(smallBatches); @@ -204,7 +225,8 @@ private PCollection writeUntriggered(PCollection writeGroupedResult = @@ -218,7 +240,8 @@ private PCollection writeUntriggered(PCollection>, IcebergWri private final String filePrefix; private final boolean autoSharding; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; WriteToPartitions( IcebergCatalogConfig catalogConfig, @@ -55,6 +57,22 @@ class WriteToPartitions extends PTransform>, IcebergWri @Nullable Duration triggeringFrequency, boolean autoSharding, @Nullable Map writeProperties) { + this( + catalogConfig, + dynamicDestinations, + triggeringFrequency, + autoSharding, + writeProperties, + null); + } + + WriteToPartitions( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Duration triggeringFrequency, + boolean autoSharding, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; this.triggeringFrequency = triggeringFrequency; @@ -62,6 +80,7 @@ class WriteToPartitions extends PTransform>, IcebergWri this.filePrefix = UUID.randomUUID().toString(); this.autoSharding = autoSharding; this.writeProperties = writeProperties; + this.metadataView = metadataView; } private PCollection>> groupByPartition(PCollection> input) { @@ -100,7 +119,7 @@ public IcebergWriteResult expand(PCollection> input) { PCollection writtenFiles = groupedRows.apply( new WritePartitionedRowsToFiles( - catalogConfig, dynamicDestinations, filePrefix, writeProperties)); + catalogConfig, dynamicDestinations, filePrefix, writeProperties, metadataView)); if (IcebergUtils.isUnbounded(input) && triggeringFrequency != null) { writtenFiles = diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteUngroupedRowsToFiles.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteUngroupedRowsToFiles.java index 7c780e6395df..8eb462f159b2 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteUngroupedRowsToFiles.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteUngroupedRowsToFiles.java @@ -34,6 +34,7 @@ import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.PInput; import org.apache.beam.sdk.values.POutput; import org.apache.beam.sdk.values.PValue; @@ -73,6 +74,7 @@ class WriteUngroupedRowsToFiles private final IcebergCatalogConfig catalogConfig; private final long maxBytesPerFile; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; WriteUngroupedRowsToFiles( IcebergCatalogConfig catalogConfig, @@ -80,29 +82,46 @@ class WriteUngroupedRowsToFiles String filePrefix, long maxBytesPerFile, @Nullable Map writeProperties) { + this(catalogConfig, dynamicDestinations, filePrefix, maxBytesPerFile, writeProperties, null); + } + + WriteUngroupedRowsToFiles( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + long maxBytesPerFile, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxBytesPerFile = maxBytesPerFile; this.writeProperties = writeProperties; + this.metadataView = metadataView; } @Override public Result expand(PCollection> input) { - PCollectionTuple resultTuple = - input.apply( - ParDo.of( - new WriteUngroupedRowsToFilesDoFn( - catalogConfig, - dynamicDestinations, - filePrefix, - DEFAULT_MAX_WRITERS_PER_BUNDLE, - maxBytesPerFile, - writeProperties)) - .withOutputTags( - WRITTEN_FILES_TAG, - TupleTagList.of(ImmutableList.of(WRITTEN_ROWS_TAG, SPILLED_ROWS_TAG)))); + ParDo.MultiOutput, FileWriteResult> parDo = + ParDo.of( + new WriteUngroupedRowsToFilesDoFn( + catalogConfig, + dynamicDestinations, + filePrefix, + DEFAULT_MAX_WRITERS_PER_BUNDLE, + maxBytesPerFile, + writeProperties, + metadataView)) + .withOutputTags( + WRITTEN_FILES_TAG, + TupleTagList.of(ImmutableList.of(WRITTEN_ROWS_TAG, SPILLED_ROWS_TAG))); + + if (metadataView != null) { + parDo = parDo.withSideInputs(metadataView); + } + + PCollectionTuple resultTuple = input.apply(parDo); return new Result( input.getPipeline(), @@ -196,6 +215,7 @@ private static class WriteUngroupedRowsToFilesDoFn private final DynamicDestinations dynamicDestinations; private final IcebergCatalogConfig catalogConfig; private final @Nullable Map writeProperties; + private final @Nullable PCollectionView> metadataView; private transient @Nullable RecordWriterManager recordWriterManager; private int spilledShardNumber; @@ -206,12 +226,31 @@ public WriteUngroupedRowsToFilesDoFn( int maximumWritersPerBundle, long maxFileSize, @Nullable Map writeProperties) { + this( + catalogConfig, + dynamicDestinations, + filename, + maximumWritersPerBundle, + maxFileSize, + writeProperties, + null); + } + + public WriteUngroupedRowsToFilesDoFn( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filename, + int maximumWritersPerBundle, + long maxFileSize, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filename = filename; this.maxWritersPerBundle = maximumWritersPerBundle; this.maxFileSize = maxFileSize; this.writeProperties = writeProperties; + this.metadataView = metadataView; } @StartBundle @@ -224,6 +263,7 @@ public void startBundle() { @ProcessElement public void processElement( + ProcessContext c, @Element KV element, BoundedWindow window, PaneInfo paneInfo, @@ -235,12 +275,16 @@ public void processElement( WindowedValue windowedDestination = WindowedValues.of(destination, window.maxTimestamp(), window, paneInfo); + Map sideInputs = + metadataView != null ? c.sideInput(metadataView) : null; + // Attempt to write record. If the writer is saturated and cannot accept // the record, spill it over to WriteGroupedRowsToFiles boolean writeSuccess; try { writeSuccess = - Preconditions.checkNotNull(recordWriterManager).write(windowedDestination, data); + Preconditions.checkNotNull(recordWriterManager) + .write(windowedDestination, data, sideInputs); } catch (Exception e) { try { Preconditions.checkNotNull(recordWriterManager).close(); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitionsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitionsTest.java new file mode 100644 index 000000000000..a753c7b947a9 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitionsTest.java @@ -0,0 +1,200 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg; + +import java.io.Serializable; +import java.util.Map; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TypeDescriptors; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.CatalogProperties; +import org.apache.iceberg.CatalogUtil; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.TableIdentifier; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link AssignDestinationsAndPartitions}. */ +@RunWith(JUnit4.class) +public class AssignDestinationsAndPartitionsTest implements Serializable { + + @Rule public transient TestPipeline pipeline = TestPipeline.create(); + @Rule public transient TemporaryFolder tempFolder = new TemporaryFolder(); + + private static final org.apache.beam.sdk.schemas.Schema BEAM_SCHEMA = + org.apache.beam.sdk.schemas.Schema.builder() + .addInt32Field("id") + .addStringField("name") + .addBooleanField("bool") + .build(); + + private static final org.apache.iceberg.Schema ICEBERG_SCHEMA = + IcebergUtils.beamSchemaToIcebergSchema(BEAM_SCHEMA); + + private static final PartitionSpec PARTITION_SPEC = + PartitionSpec.builderFor(ICEBERG_SCHEMA).truncate("name", 3).identity("bool").build(); + + private String warehouseLocation; + private IcebergCatalogConfig catalogConfig; + + @Before + public void setUp() throws Exception { + warehouseLocation = "file:" + tempFolder.newFolder().getAbsolutePath(); + catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogName("hadoop") + .setCatalogProperties(ImmutableMap.of("type", "hadoop", "warehouse", warehouseLocation)) + .build(); + TableCache.invalidateAll(); + } + + private Catalog getCatalog() { + return CatalogUtil.loadCatalog( + CatalogUtil.ICEBERG_CATALOG_HADOOP, + "hadoop", + ImmutableMap.of(CatalogProperties.WAREHOUSE_LOCATION, warehouseLocation), + new Configuration()); + } + + @Test + public void testAssignDestinationsWithoutMetadataViewFallsBackToTableCache() { + TableIdentifier tableId = TableIdentifier.of("default", "test_table_no_view"); + getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "bob", false).build(); + + PCollection input = + pipeline.apply("CreateInput", Create.of(row1, row2).withRowSchema(BEAM_SCHEMA)); + + PCollection> assigned = + input.apply(new AssignDestinationsAndPartitions(dynamicDestinations, catalogConfig)); + + PCollection partitionPaths = + assigned.apply( + "ExtractPartitionPaths", + MapElements.into(TypeDescriptors.strings()) + .via(kv -> kv.getKey().getString(AssignDestinationsAndPartitions.PARTITION))); + + PAssert.that(partitionPaths) + .containsInAnyOrder("name_trunc=ali/bool=true", "name_trunc=bob/bool=false"); + + pipeline.run(); + } + + @Test + public void testAssignDestinationsWithMetadataViewHit() { + TableIdentifier tableId = TableIdentifier.of("default", "test_table_view_hit"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + // Drop table from catalog and clear cache so that any catalog fallback would fail to find the + // spec + getCatalog().dropTable(tableId); + TableCache.invalidateAll(); + + PCollectionView> metadataView = + pipeline + .apply( + "CreateMetadata", + Create.of(KV.of(tableIdString, spec)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder()))) + .apply("AsView", View.asMap()); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + + PCollection input = + pipeline.apply("CreateInput", Create.of(row1).withRowSchema(BEAM_SCHEMA)); + + PCollection> assigned = + input.apply( + new AssignDestinationsAndPartitions(dynamicDestinations, catalogConfig, metadataView)); + + PCollection partitionPaths = + assigned.apply( + "ExtractPartitionPaths", + MapElements.into(TypeDescriptors.strings()) + .via(kv -> kv.getKey().getString(AssignDestinationsAndPartitions.PARTITION))); + + PAssert.that(partitionPaths).containsInAnyOrder("name_trunc=ali/bool=true"); + + pipeline.run(); + } + + @Test + public void testAssignDestinationsWithMetadataViewMissFallsBack() { + TableIdentifier tableId = TableIdentifier.of("default", "test_table_view_miss"); + getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + + TableIdentifier otherId = TableIdentifier.of("default", "test_other_table"); + SerializableTableSpec otherSpec = + SerializableTableSpec.fromTable(tableId, getCatalog().loadTable(tableId)); + String otherIdString = IcebergUtils.tableIdentifierToString(otherId); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + PCollectionView> metadataView = + pipeline + .apply( + "CreateMetadata", + Create.of(KV.of(otherIdString, otherSpec)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder()))) + .apply("AsView", View.asMap()); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + + PCollection input = + pipeline.apply("CreateInput", Create.of(row1).withRowSchema(BEAM_SCHEMA)); + + PCollection> assigned = + input.apply( + new AssignDestinationsAndPartitions(dynamicDestinations, catalogConfig, metadataView)); + + PCollection partitionPaths = + assigned.apply( + "ExtractPartitionPaths", + MapElements.into(TypeDescriptors.strings()) + .via(kv -> kv.getKey().getString(AssignDestinationsAndPartitions.PARTITION))); + + PAssert.that(partitionPaths).containsInAnyOrder("name_trunc=ali/bool=true"); + + pipeline.run(); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/RecordWriterManagerTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/RecordWriterManagerTest.java index 03b3560f746a..98597ba4c4f3 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/RecordWriterManagerTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/RecordWriterManagerTest.java @@ -1357,4 +1357,108 @@ public void testWritePropertiesAppliedToParquetFiles() throws IOException { } } } + + @Test + public void testGetOrCreateTableWithSideInputHit() { + TableIdentifier tableId = TableIdentifier.of("default", "test_side_input_hit"); + Table realTable = warehouse.createTable(tableId, ICEBERG_SCHEMA); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + Catalog mockCatalog = mock(Catalog.class); + IcebergCatalogConfig mockCatalogConfig = mockCatalogConfigFor(mockCatalog); + + IcebergDestination destination = + IcebergDestination.builder() + .setFileFormat(FileFormat.PARQUET) + .setTableIdentifier(tableId) + .build(); + + Map sideInputs = ImmutableMap.of(tableIdString, spec); + RecordWriterManager writerManager = + new RecordWriterManager(mockCatalogConfig, "test_prefix", 1024L, 1, null, sideInputs); + + Table resolvedTable = writerManager.getOrCreateTable(destination, BEAM_SCHEMA); + assertTrue(resolvedTable instanceof SideInputTable); + assertEquals(spec, ((SideInputTable) resolvedTable).getTableSpec()); + + // Verify catalog.loadTable was NEVER called + verify(mockCatalog, never()).loadTable(Mockito.any()); + } + + @Test + public void testGetOrCreateTableWithSideInputMissFallsBackToTableCache() { + TableIdentifier tableId = TableIdentifier.of("default", "test_side_input_miss"); + Table realTable = warehouse.createTable(tableId, ICEBERG_SCHEMA); + TableIdentifier otherId = TableIdentifier.of("default", "test_other_table"); + SerializableTableSpec otherSpec = SerializableTableSpec.fromTable(otherId, realTable); + + IcebergDestination destination = + IcebergDestination.builder() + .setFileFormat(FileFormat.PARQUET) + .setTableIdentifier(tableId) + .build(); + + Map sideInputs = + ImmutableMap.of(IcebergUtils.tableIdentifierToString(otherId), otherSpec); + RecordWriterManager writerManager = + new RecordWriterManager(catalogConfig, "test_prefix", 1024L, 1, null, sideInputs); + + Table resolvedTable = writerManager.getOrCreateTable(destination, BEAM_SCHEMA); + assertNotNull(resolvedTable); + assertFalse(resolvedTable instanceof SideInputTable); + assertEquals(realTable.location(), resolvedTable.location()); + } + + @Test + public void testGetOrCreateTableWithNullSideInputMapFallsBack() { + TableIdentifier tableId = TableIdentifier.of("default", "test_null_side_input"); + Table realTable = warehouse.createTable(tableId, ICEBERG_SCHEMA); + + IcebergDestination destination = + IcebergDestination.builder() + .setFileFormat(FileFormat.PARQUET) + .setTableIdentifier(tableId) + .build(); + + RecordWriterManager writerManager = + new RecordWriterManager(catalogConfig, "test_prefix", 1024L, 1); + + Table resolvedTable = writerManager.getOrCreateTable(destination, BEAM_SCHEMA, null); + assertNotNull(resolvedTable); + assertFalse(resolvedTable instanceof SideInputTable); + assertEquals(realTable.location(), resolvedTable.location()); + } + + @Test + public void testWriteWithSideInputTableProducesValidDataFiles() throws Exception { + TableIdentifier tableId = TableIdentifier.of("default", "test_side_input_write"); + Table realTable = warehouse.createTable(tableId, ICEBERG_SCHEMA); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + IcebergDestination destination = + IcebergDestination.builder() + .setFileFormat(FileFormat.PARQUET) + .setTableIdentifier(tableId) + .build(); + WindowedValue dest = WindowedValues.valueInGlobalWindow(destination); + + Map sideInputs = ImmutableMap.of(tableIdString, spec); + RecordWriterManager writerManager = + new RecordWriterManager( + catalogConfig, "test_side_input", Long.MAX_VALUE, 5, null, sideInputs); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "bob", false).build(); + + assertTrue(writerManager.write(dest, row1)); + assertTrue(writerManager.write(dest, row2)); + writerManager.close(); + + List dataFiles = writerManager.getSerializableDataFiles().get(dest); + assertNotNull(dataFiles); + assertEquals(1, dataFiles.size()); + assertEquals(2L, dataFiles.get(0).getRecordCount()); + } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/WriteWithMetadataViewTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/WriteWithMetadataViewTest.java new file mode 100644 index 000000000000..75677ff97bcd --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/WriteWithMetadataViewTest.java @@ -0,0 +1,397 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg; + +import static org.junit.Assert.assertEquals; + +import java.io.Serializable; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.coders.IterableCoder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.util.ShardedKey; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TypeDescriptors; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.CatalogProperties; +import org.apache.iceberg.CatalogUtil; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.IcebergGenerics; +import org.apache.iceberg.data.Record; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests verifying that file writers and orchestrators correctly resolve table metadata from {@link + * PCollectionView} of {@link SerializableTableSpec}. + */ +@RunWith(JUnit4.class) +public class WriteWithMetadataViewTest implements Serializable { + + @Rule public transient TestPipeline pipeline = TestPipeline.create(); + @Rule public transient TemporaryFolder tempFolder = new TemporaryFolder(); + + private static final org.apache.beam.sdk.schemas.Schema BEAM_SCHEMA = + org.apache.beam.sdk.schemas.Schema.builder() + .addInt32Field("id") + .addStringField("name") + .addBooleanField("bool") + .build(); + + private static final org.apache.iceberg.Schema ICEBERG_SCHEMA = + IcebergUtils.beamSchemaToIcebergSchema(BEAM_SCHEMA); + + private static final PartitionSpec PARTITION_SPEC = + PartitionSpec.builderFor(ICEBERG_SCHEMA).identity("bool").build(); + + private String warehouseLocation; + private IcebergCatalogConfig catalogConfig; + + @Before + public void setUp() throws Exception { + warehouseLocation = "file:" + tempFolder.newFolder().getAbsolutePath(); + catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogName("hadoop") + .setCatalogProperties(ImmutableMap.of("type", "hadoop", "warehouse", warehouseLocation)) + .build(); + TableCache.invalidateAll(); + } + + private Catalog getCatalog() { + return CatalogUtil.loadCatalog( + CatalogUtil.ICEBERG_CATALOG_HADOOP, + "hadoop", + ImmutableMap.of(CatalogProperties.WAREHOUSE_LOCATION, warehouseLocation), + new Configuration()); + } + + @Test + public void testWriteUngroupedRowsToFilesWithMetadataView() { + TableIdentifier tableId = TableIdentifier.of("default", "test_ungrouped"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + PCollectionView> metadataView = + pipeline + .apply( + "CreateMetadata", + Create.of(KV.of(tableIdString, spec)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder()))) + .apply("AsView", View.asMap()); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "bob", false).build(); + + PCollection> input = + pipeline.apply( + "CreateInput", + Create.of(KV.of(tableIdString, row1), KV.of(tableIdString, row2)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), RowCoder.of(BEAM_SCHEMA)))); + + WriteUngroupedRowsToFiles.Result result = + input.apply( + new WriteUngroupedRowsToFiles( + catalogConfig, dynamicDestinations, "prefix", 1024L * 1024L, null, metadataView)); + + PCollection tables = + result + .getWrittenFiles() + .apply( + MapElements.into(TypeDescriptors.strings()) + .via(f -> IcebergUtils.tableIdentifierToString(f.getTableIdentifier()))); + + PAssert.that(tables).containsInAnyOrder(tableIdString, tableIdString); + + PAssert.that(result.getWrittenRows()).containsInAnyOrder(row1, row2); + pipeline.run(); + } + + @Test + public void testWriteGroupedRowsToFilesWithMetadataView() { + TableIdentifier tableId = TableIdentifier.of("default", "test_grouped"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + PCollectionView> metadataView = + pipeline + .apply( + "CreateMetadata", + Create.of(KV.of(tableIdString, spec)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder()))) + .apply("AsView", View.asMap()); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "bob", false).build(); + + ShardedKey shardedKey = ShardedKey.of(tableIdString, new byte[] {0}); + PCollection, Iterable>> input = + pipeline.apply( + "CreateGroupedInput", + Create.of(KV.of(shardedKey, (Iterable) ImmutableList.of(row1, row2))) + .withCoder( + KvCoder.of( + ShardedKey.Coder.of(StringUtf8Coder.of()), + IterableCoder.of(RowCoder.of(BEAM_SCHEMA))))); + + PCollection writtenFiles = + input.apply( + new WriteGroupedRowsToFiles( + catalogConfig, dynamicDestinations, "prefix", 1024L * 1024L, null, metadataView)); + + PCollection tables = + writtenFiles.apply( + MapElements.into(TypeDescriptors.strings()) + .via(f -> IcebergUtils.tableIdentifierToString(f.getTableIdentifier()))); + + PAssert.that(tables).containsInAnyOrder(tableIdString, tableIdString); + pipeline.run(); + } + + @Test + public void testWriteDirectRowsToFilesWithMetadataView() { + TableIdentifier tableId = TableIdentifier.of("default", "test_direct"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + PCollectionView> metadataView = + pipeline + .apply( + "CreateMetadata", + Create.of(KV.of(tableIdString, spec)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder()))) + .apply("AsView", View.asMap()); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + + PCollection> input = + pipeline.apply( + "CreateDirectInput", + Create.of(KV.of(tableIdString, row1)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), RowCoder.of(BEAM_SCHEMA)))); + + PCollection writtenFiles = + input.apply( + new WriteDirectRowsToFiles( + catalogConfig, dynamicDestinations, "prefix", 1024L * 1024L, null, metadataView)); + + PCollection tables = + writtenFiles.apply( + MapElements.into(TypeDescriptors.strings()) + .via(f -> IcebergUtils.tableIdentifierToString(f.getTableIdentifier()))); + + PAssert.that(tables).containsInAnyOrder(tableIdString); + pipeline.run(); + } + + @Test + public void testWritePartitionedRowsToFilesWithMetadataView() { + TableIdentifier tableId = TableIdentifier.of("default", "test_partitioned"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + PCollectionView> metadataView = + pipeline + .apply( + "CreateMetadata", + Create.of(KV.of(tableIdString, spec)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder()))) + .apply("AsView", View.asMap()); + + Row partitionRow = + Row.withSchema(AssignDestinationsAndPartitions.OUTPUT_SCHEMA) + .addValues(tableIdString, "bool=true") + .build(); + Row dataRow = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + + PCollection>> input = + pipeline.apply( + "CreatePartitionedInput", + Create.of(KV.of(partitionRow, (Iterable) ImmutableList.of(dataRow))) + .withCoder( + KvCoder.of( + RowCoder.of(AssignDestinationsAndPartitions.OUTPUT_SCHEMA), + IterableCoder.of(RowCoder.of(BEAM_SCHEMA))))); + + PCollection writtenFiles = + input.apply( + new WritePartitionedRowsToFiles( + catalogConfig, dynamicDestinations, "prefix", null, metadataView)); + + PCollection tables = + writtenFiles.apply( + MapElements.into(TypeDescriptors.strings()) + .via(f -> IcebergUtils.tableIdentifierToString(f.getTableIdentifier()))); + + PAssert.that(tables).containsInAnyOrder(tableIdString); + pipeline.run(); + } + + @Test + public void testWriteToDestinationsUntriggeredWithMetadataView() { + TableIdentifier tableId = TableIdentifier.of("default", "test_destinations_end_to_end"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + PCollectionView> metadataView = + pipeline + .apply( + "CreateMetadata", + Create.of(KV.of(tableIdString, spec)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder()))) + .apply("AsView", View.asMap()); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "bob", false).build(); + + PCollection> input = + pipeline.apply( + "CreateInput", + Create.of(KV.of(tableIdString, row1), KV.of(tableIdString, row2)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), RowCoder.of(BEAM_SCHEMA)))); + + input.apply( + new WriteToDestinations( + catalogConfig, dynamicDestinations, null, null, null, metadataView)); + + pipeline.run(); + + // Verify records committed to table + realTable.refresh(); + List committed = ImmutableList.copyOf(IcebergGenerics.read(realTable).build()); + assertEquals(2, committed.size()); + } + + @Test + public void testWriteToPartitionsWithMetadataView() { + TableIdentifier tableId = TableIdentifier.of("default", "test_partitions_end_to_end"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + PCollectionView> metadataView = + pipeline + .apply( + "CreateMetadata", + Create.of(KV.of(tableIdString, spec)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder()))) + .apply("AsView", View.asMap()); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "bob", false).build(); + + PCollection input = + pipeline.apply("CreateRows", Create.of(row1, row2).withRowSchema(BEAM_SCHEMA)); + + PCollection> assigned = + input.apply( + new AssignDestinationsAndPartitions(dynamicDestinations, catalogConfig, metadataView)); + + assigned.apply( + new WriteToPartitions(catalogConfig, dynamicDestinations, null, false, null, metadataView)); + + pipeline.run(); + + // Verify records committed to table + realTable.refresh(); + List committed = ImmutableList.copyOf(IcebergGenerics.read(realTable).build()); + assertEquals(2, committed.size()); + } + + @Test + public void testWriteUngroupedRowsBypassesCatalogWhenUsingMetadataView() { + TableIdentifier tableId = TableIdentifier.of("default", "test_bypasses_catalog"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA, PARTITION_SPEC); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + + // Drop table from catalog and invalidate cache so catalog.loadTable() would fail if invoked + getCatalog().dropTable(tableId, false); + TableCache.invalidateAll(); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + PCollectionView> metadataView = + pipeline + .apply( + "CreateMetadata", + Create.of(KV.of(tableIdString, spec)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder()))) + .apply("AsView", View.asMap()); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "alice", true).build(); + + PCollection> input = + pipeline.apply( + "CreateInput", + Create.of(KV.of(tableIdString, row1)) + .withCoder(KvCoder.of(StringUtf8Coder.of(), RowCoder.of(BEAM_SCHEMA)))); + + WriteUngroupedRowsToFiles.Result result = + input.apply( + new WriteUngroupedRowsToFiles( + catalogConfig, dynamicDestinations, "prefix", 1024L * 1024L, null, metadataView)); + + PCollection tables = + result + .getWrittenFiles() + .apply( + MapElements.into(TypeDescriptors.strings()) + .via(f -> IcebergUtils.tableIdentifierToString(f.getTableIdentifier()))); + + PAssert.that(tables).containsInAnyOrder(tableIdString); + PAssert.that(result.getWrittenRows()).containsInAnyOrder(row1); + pipeline.run(); + } +}