This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 333a7787e94 Implement Iceberg Side-Input Table Cache Integration with
Fallback (#40080)
333a7787e94 is described below
commit 333a7787e94117440c82226a3acc3436ef338ac2
Author: Jack McCluskey <[email protected]>
AuthorDate: Wed Sep 16 09:57:24 2026 -0400
Implement Iceberg Side-Input Table Cache Integration with Fallback (#40080)
* Implement Iceberg Side-Input Table Cache Integration with Fallback
* review comments
* clean up extra identifier round trip
---
.../iceberg/AssignDestinationsAndPartitions.java | 75 +++-
.../beam/sdk/io/iceberg/RecordWriterManager.java | 56 ++-
.../sdk/io/iceberg/WriteDirectRowsToFiles.java | 44 ++-
.../sdk/io/iceberg/WriteGroupedRowsToFiles.java | 47 ++-
.../io/iceberg/WritePartitionedRowsToFiles.java | 66 +++-
.../beam/sdk/io/iceberg/WriteToDestinations.java | 31 +-
.../beam/sdk/io/iceberg/WriteToPartitions.java | 21 +-
.../sdk/io/iceberg/WriteUngroupedRowsToFiles.java | 72 +++-
.../AssignDestinationsAndPartitionsTest.java | 200 +++++++++++
.../sdk/io/iceberg/RecordWriterManagerTest.java | 104 ++++++
.../sdk/io/iceberg/WriteWithMetadataViewTest.java | 397 +++++++++++++++++++++
11 files changed, 1062 insertions(+), 51 deletions(-)
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 a744ff93097..4167ad00624 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.BoundedWindow;
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<Map<String, SerializableTableSpec>>
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<Map<String, SerializableTableSpec>>
metadataView) {
this.dynamicDestinations = dynamicDestinations;
this.catalogConfig = catalogConfig;
+ this.metadataView = metadataView;
}
@Override
public PCollection<KV<Row, Row>> expand(PCollection<Row> input) {
+ ParDo.SingleOutput<Row, KV<Row, Row>> 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 @@ class AssignDestinationsAndPartitions
private transient @MonotonicNonNull Map<String, PartitionKey>
partitionKeys;
private transient @MonotonicNonNull Map<String, BeamRowWrapper> wrappers;
private transient @MonotonicNonNull Map<String, Instant> lastRefreshTimes;
+ private transient @MonotonicNonNull Map<String, Integer> cachedSpecIds;
private final DynamicDestinations dynamicDestinations;
private final IcebergCatalogConfig catalogConfig;
+ private final @Nullable PCollectionView<Map<String,
SerializableTableSpec>> metadataView;
AssignDoFn(DynamicDestinations dynamicDestinations, IcebergCatalogConfig
catalogConfig) {
+ this(dynamicDestinations, catalogConfig, null);
+ }
+
+ AssignDoFn(
+ DynamicDestinations dynamicDestinations,
+ IcebergCatalogConfig catalogConfig,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView) {
this.dynamicDestinations = dynamicDestinations;
this.catalogConfig = catalogConfig;
+ this.metadataView = metadataView;
}
@Setup
@@ -97,10 +123,12 @@ class AssignDestinationsAndPartitions
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,71 @@ class AssignDestinationsAndPartitions
dynamicDestinations.getTableStringIdentifier(
ValueInSingleWindow.of(element, timestamp, window, paneInfo));
+ SerializableTableSpec tableSpec = null;
+ if (metadataView != null) {
+ Map<String, SerializableTableSpec> viewMap = c.sideInput(metadataView);
+ if (viewMap != 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)
{
-
+ if (tableSpec != null) {
+ spec = tableSpec.getPartitionSpec();
+ if (data.getSchema().getFieldCount() ==
tableSpec.getSchema().columns().size()) {
+ schema = tableSpec.getSchema();
+ }
+ checkStateNotNull(cachedSpecIds).put(tableIdentifier,
tableSpec.getSpecId());
+ } else if (createConfig != null && createConfig.getPartitionFields()
!= null) {
spec =
PartitionUtils.toPartitionSpec(createConfig.getPartitionFields(),
data.getSchema());
-
} 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 0995e7a6102..25e5a13da43 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.ZoneOffset;
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 @@ class RecordWriterManager implements AutoCloseable {
private final long maxFileSize;
private final int maxNumWriters;
private final @Nullable Map<String, String> writeProperties;
+ private volatile @Nullable Map<String, SerializableTableSpec>
sideInputTableSpecs;
@VisibleForTesting int openWriters = 0;
@VisibleForTesting
@@ -263,7 +265,7 @@ class RecordWriterManager implements AutoCloseable {
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 @@ class RecordWriterManager implements AutoCloseable {
long maxFileSize,
int maxNumWriters,
@Nullable Map<String, String> writeProperties) {
+ this(catalogConfig, filePrefix, maxFileSize, maxNumWriters,
writeProperties, null);
+ }
+
+ RecordWriterManager(
+ IcebergCatalogConfig catalogConfig,
+ String filePrefix,
+ long maxFileSize,
+ int maxNumWriters,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable Map<String, SerializableTableSpec> sideInputTableSpecs) {
this.catalogConfig = catalogConfig;
this.filePrefix = filePrefix;
this.maxFileSize = maxFileSize;
this.maxNumWriters = maxNumWriters;
this.writeProperties = writeProperties;
+ this.sideInputTableSpecs = sideInputTableSpecs;
+ }
+
+ @VisibleForTesting
+ void setSideInputTableSpecs(@Nullable Map<String, SerializableTableSpec>
sideInputTableSpecs) {
+ this.sideInputTableSpecs = sideInputTableSpecs;
}
/**
@@ -291,12 +309,29 @@ class RecordWriterManager implements AutoCloseable {
* using the Iceberg API.
*/
@VisibleForTesting
- Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) {
+ Table getOrCreateTable(
+ IcebergDestination destination,
+ Schema dataSchema,
+ @Nullable Map<String, SerializableTableSpec> 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<String, String> 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 @@ class RecordWriterManager implements AutoCloseable {
}
}
+ /**
+ * Fetches the appropriate {@link RecordWriter} for this destination and
partition and writes the
+ * record, optionally updating the side-input table specs map.
+ *
+ * <p>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> icebergDestination,
+ Row row,
+ @Nullable Map<String, SerializableTableSpec> 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 e03085e6be7..8bafe9eeaef 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.BoundedWindow;
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<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView;
WriteDirectRowsToFiles(
IcebergCatalogConfig catalogConfig,
@@ -48,19 +50,39 @@ class WriteDirectRowsToFiles
String filePrefix,
long maxBytesPerFile,
@Nullable Map<String, String> writeProperties) {
+ this(catalogConfig, dynamicDestinations, filePrefix, maxBytesPerFile,
writeProperties, null);
+ }
+
+ WriteDirectRowsToFiles(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ String filePrefix,
+ long maxBytesPerFile,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.maxBytesPerFile = maxBytesPerFile;
this.writeProperties = writeProperties;
+ this.metadataView = metadataView;
}
@Override
public PCollection<FileWriteResult> expand(PCollection<KV<String, Row>>
input) {
- return input.apply(
+ ParDo.SingleOutput<KV<String, Row>, 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<KV<String,
Row>, FileWriteResult> {
@@ -70,6 +92,7 @@ class WriteDirectRowsToFiles
private final String filePrefix;
private final long maxFileSize;
private final @Nullable Map<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String,
SerializableTableSpec>> metadataView;
private transient @Nullable RecordWriterManager recordWriterManager;
WriteDirectRowsToFilesDoFn(
@@ -78,11 +101,22 @@ class WriteDirectRowsToFiles
long maxFileSize,
String filePrefix,
@Nullable Map<String, String> writeProperties) {
+ this(catalogConfig, dynamicDestinations, maxFileSize, filePrefix,
writeProperties, null);
+ }
+
+ WriteDirectRowsToFilesDoFn(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ long maxFileSize,
+ String filePrefix,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
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 @@ class WriteDirectRowsToFiles
@ProcessElement
public void processElement(
- @SuppressWarnings("unused") ProcessContext context,
+ ProcessContext context,
@Element KV<String, Row> element,
BoundedWindow window,
PaneInfo paneInfo)
@@ -104,8 +138,10 @@ class WriteDirectRowsToFiles
IcebergDestination destination =
dynamicDestinations.instantiateDestination(tableIdentifier);
WindowedValue<IcebergDestination> windowedDestination =
WindowedValues.of(destination, window.maxTimestamp(), window,
paneInfo);
+ Map<String, SerializableTableSpec> 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 e74715a7eeb..a2ed40c87f9 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.transforms.windowing.PaneInfo;
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<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView;
WriteGroupedRowsToFiles(
IcebergCatalogConfig catalogConfig,
@@ -49,20 +51,40 @@ class WriteGroupedRowsToFiles
String filePrefix,
long maxBytesPerFile,
@Nullable Map<String, String> writeProperties) {
+ this(catalogConfig, dynamicDestinations, filePrefix, maxBytesPerFile,
writeProperties, null);
+ }
+
+ WriteGroupedRowsToFiles(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ String filePrefix,
+ long maxBytesPerFile,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.maxBytesPerFile = maxBytesPerFile;
this.writeProperties = writeProperties;
+ this.metadataView = metadataView;
}
@Override
public PCollection<FileWriteResult> expand(
PCollection<KV<ShardedKey<String>, Iterable<Row>>> input) {
- return input.apply(
+ ParDo.SingleOutput<KV<ShardedKey<String>, Iterable<Row>>, 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 @@ class WriteGroupedRowsToFiles
private final String filePrefix;
private final long maxFileSize;
private final @Nullable Map<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String,
SerializableTableSpec>> metadataView;
WriteGroupedRowsToFilesDoFn(
IcebergCatalogConfig catalogConfig,
@@ -80,11 +103,22 @@ class WriteGroupedRowsToFiles
long maxFileSize,
String filePrefix,
@Nullable Map<String, String> writeProperties) {
+ this(catalogConfig, dynamicDestinations, maxFileSize, filePrefix,
writeProperties, null);
+ }
+
+ WriteGroupedRowsToFilesDoFn(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ long maxFileSize,
+ String filePrefix,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.maxFileSize = maxFileSize;
this.writeProperties = writeProperties;
+ this.metadataView = metadataView;
}
@ProcessElement
@@ -99,10 +133,17 @@ class WriteGroupedRowsToFiles
IcebergDestination destination =
dynamicDestinations.instantiateDestination(tableIdentifier);
WindowedValue<IcebergDestination> windowedDestination =
WindowedValues.of(destination, window.maxTimestamp(), window,
paneInfo);
+ Map<String, SerializableTableSpec> 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 338a2162080..881d2577fad 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.AssignDestinationsAndPartitions.PAR
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.PTransform;
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<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView;
WritePartitionedRowsToFiles(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
String filePrefix,
@Nullable Map<String, String> writeProperties) {
+ this(catalogConfig, dynamicDestinations, filePrefix, writeProperties,
null);
+ }
+
+ WritePartitionedRowsToFiles(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ String filePrefix,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.writeProperties = writeProperties;
+ this.metadataView = metadataView;
}
@Override
@@ -81,10 +94,19 @@ class WritePartitionedRowsToFiles
((KvCoder<Row, Iterable<Row>>)
input.getCoder()).getValueCoder())
.getElemCoder())
.getSchema();
- return input.apply(
+ ParDo.SingleOutput<KV<Row, Iterable<Row>>, 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<KV<Row, Iterable<Row>>,
FileWriteResult> {
@@ -94,6 +116,7 @@ class WritePartitionedRowsToFiles
private final String filePrefix;
private final Schema dataSchema;
private final @Nullable Map<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String,
SerializableTableSpec>> metadataView;
private transient @MonotonicNonNull Map<TableIdentifier, Integer> specIds;
private transient @MonotonicNonNull Map<TableIdentifier, Map<String,
PartitionField>>
partitionFieldMaps;
@@ -104,11 +127,22 @@ class WritePartitionedRowsToFiles
String filePrefix,
Schema dataSchema,
@Nullable Map<String, String> writeProperties) {
+ this(catalogConfig, dynamicDestinations, filePrefix, dataSchema,
writeProperties, null);
+ }
+
+ WriteDoFn(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ String filePrefix,
+ Schema dataSchema,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.dataSchema = dataSchema;
this.writeProperties = writeProperties;
+ this.metadataView = metadataView;
}
@Setup
@@ -119,13 +153,17 @@ class WritePartitionedRowsToFiles
@ProcessElement
public void processElement(
- @Element KV<Row, Iterable<Row>> element,
OutputReceiver<FileWriteResult> out)
+ ProcessContext c,
+ @Element KV<Row, Iterable<Row>> element,
+ OutputReceiver<FileWriteResult> 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<String, SerializableTableSpec> sideInputs =
+ metadataView != null ? c.sideInput(metadataView) : null;
+ Table table = getOrCreateTable(destination, dataSchema, sideInputs);
partitionPath =
getPartitionDataPath(
partitionPath,
getPartitionFieldMap(destination.getTableIdentifier(), table));
@@ -176,14 +214,30 @@ class WritePartitionedRowsToFiles
return partitionFieldMap;
}
- Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) {
+ Table getOrCreateTable(
+ IcebergDestination destination,
+ Schema dataSchema,
+ @Nullable Map<String, SerializableTableSpec> 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<String, String> 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 @@ class WritePartitionedRowsToFiles
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 684ef350a20..6d8f08a2ae5 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.KV;
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<PCollection<KV<String, Row>>, Icebe
private final String filePrefix;
private final @Nullable Integer directWriteByteLimit;
private final @Nullable Map<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView;
WriteToDestinations(
IcebergCatalogConfig catalogConfig,
@@ -67,11 +69,28 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
@Nullable Duration triggeringFrequency,
@Nullable Integer directWriteByteLimit,
@Nullable Map<String, String> writeProperties) {
+ this(
+ catalogConfig,
+ dynamicDestinations,
+ triggeringFrequency,
+ directWriteByteLimit,
+ writeProperties,
+ null);
+ }
+
+ WriteToDestinations(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ @Nullable Duration triggeringFrequency,
+ @Nullable Integer directWriteByteLimit,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
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 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
dynamicDestinations,
filePrefix,
DEFAULT_MAX_BYTES_PER_FILE,
- writeProperties));
+ writeProperties,
+ metadataView));
}
private PCollection<FileWriteResult>
applyUserTriggering(PCollection<FileWriteResult> input) {
@@ -168,7 +188,8 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
dynamicDestinations,
filePrefix,
DEFAULT_MAX_BYTES_PER_FILE,
- writeProperties));
+ writeProperties,
+ metadataView));
PCollection<FileWriteResult> groupedFileWrites =
groupAndWriteRecords(smallBatches);
@@ -204,7 +225,8 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
dynamicDestinations,
filePrefix,
DEFAULT_MAX_BYTES_PER_FILE,
- writeProperties));
+ writeProperties,
+ metadataView));
// Then write the rest by shuffling on the destination
PCollection<FileWriteResult> writeGroupedResult =
@@ -218,7 +240,8 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
dynamicDestinations,
filePrefix,
DEFAULT_MAX_BYTES_PER_FILE,
- writeProperties));
+ writeProperties,
+ metadataView));
return PCollectionList.of(writeUngroupedResult.getWrittenFiles())
.and(writeGroupedResult)
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
index 84b9c6a9c0f..7e9d56df98a 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
@@ -36,6 +36,7 @@ import org.apache.beam.sdk.transforms.windowing.Repeatedly;
import org.apache.beam.sdk.transforms.windowing.Window;
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.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Duration;
@@ -48,6 +49,7 @@ class WriteToPartitions extends
PTransform<PCollection<KV<Row, Row>>, IcebergWri
private final String filePrefix;
private final boolean autoSharding;
private final @Nullable Map<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView;
WriteToPartitions(
IcebergCatalogConfig catalogConfig,
@@ -55,6 +57,22 @@ class WriteToPartitions extends
PTransform<PCollection<KV<Row, Row>>, IcebergWri
@Nullable Duration triggeringFrequency,
boolean autoSharding,
@Nullable Map<String, String> writeProperties) {
+ this(
+ catalogConfig,
+ dynamicDestinations,
+ triggeringFrequency,
+ autoSharding,
+ writeProperties,
+ null);
+ }
+
+ WriteToPartitions(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ @Nullable Duration triggeringFrequency,
+ boolean autoSharding,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView) {
this.dynamicDestinations = dynamicDestinations;
this.catalogConfig = catalogConfig;
this.triggeringFrequency = triggeringFrequency;
@@ -62,6 +80,7 @@ class WriteToPartitions extends
PTransform<PCollection<KV<Row, Row>>, IcebergWri
this.filePrefix = UUID.randomUUID().toString();
this.autoSharding = autoSharding;
this.writeProperties = writeProperties;
+ this.metadataView = metadataView;
}
private PCollection<KV<Row, Iterable<Row>>>
groupByPartition(PCollection<KV<Row, Row>> input) {
@@ -100,7 +119,7 @@ class WriteToPartitions extends
PTransform<PCollection<KV<Row, Row>>, IcebergWri
PCollection<FileWriteResult> 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 7c780e6395d..8eb462f159b 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.util.ShardedKey;
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<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView;
WriteUngroupedRowsToFiles(
IcebergCatalogConfig catalogConfig,
@@ -80,29 +82,46 @@ class WriteUngroupedRowsToFiles
String filePrefix,
long maxBytesPerFile,
@Nullable Map<String, String> writeProperties) {
+ this(catalogConfig, dynamicDestinations, filePrefix, maxBytesPerFile,
writeProperties, null);
+ }
+
+ WriteUngroupedRowsToFiles(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ String filePrefix,
+ long maxBytesPerFile,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
metadataView) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.maxBytesPerFile = maxBytesPerFile;
this.writeProperties = writeProperties;
+ this.metadataView = metadataView;
}
@Override
public Result expand(PCollection<KV<String, Row>> 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<KV<String, Row>, 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 @@ class WriteUngroupedRowsToFiles
private final DynamicDestinations dynamicDestinations;
private final IcebergCatalogConfig catalogConfig;
private final @Nullable Map<String, String> writeProperties;
+ private final @Nullable PCollectionView<Map<String,
SerializableTableSpec>> metadataView;
private transient @Nullable RecordWriterManager recordWriterManager;
private int spilledShardNumber;
@@ -206,12 +226,31 @@ class WriteUngroupedRowsToFiles
int maximumWritersPerBundle,
long maxFileSize,
@Nullable Map<String, String> writeProperties) {
+ this(
+ catalogConfig,
+ dynamicDestinations,
+ filename,
+ maximumWritersPerBundle,
+ maxFileSize,
+ writeProperties,
+ null);
+ }
+
+ public WriteUngroupedRowsToFilesDoFn(
+ IcebergCatalogConfig catalogConfig,
+ DynamicDestinations dynamicDestinations,
+ String filename,
+ int maximumWritersPerBundle,
+ long maxFileSize,
+ @Nullable Map<String, String> writeProperties,
+ @Nullable PCollectionView<Map<String, SerializableTableSpec>>
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 @@ class WriteUngroupedRowsToFiles
@ProcessElement
public void processElement(
+ ProcessContext c,
@Element KV<String, Row> element,
BoundedWindow window,
PaneInfo paneInfo,
@@ -235,12 +275,16 @@ class WriteUngroupedRowsToFiles
WindowedValue<IcebergDestination> windowedDestination =
WindowedValues.of(destination, window.maxTimestamp(), window,
paneInfo);
+ Map<String, SerializableTableSpec> 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 00000000000..a753c7b947a
--- /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<Row> input =
+ pipeline.apply("CreateInput", Create.of(row1,
row2).withRowSchema(BEAM_SCHEMA));
+
+ PCollection<KV<Row, Row>> assigned =
+ input.apply(new AssignDestinationsAndPartitions(dynamicDestinations,
catalogConfig));
+
+ PCollection<String> 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<Map<String, SerializableTableSpec>> 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<Row> input =
+ pipeline.apply("CreateInput",
Create.of(row1).withRowSchema(BEAM_SCHEMA));
+
+ PCollection<KV<Row, Row>> assigned =
+ input.apply(
+ new AssignDestinationsAndPartitions(dynamicDestinations,
catalogConfig, metadataView));
+
+ PCollection<String> 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<Map<String, SerializableTableSpec>> 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<Row> input =
+ pipeline.apply("CreateInput",
Create.of(row1).withRowSchema(BEAM_SCHEMA));
+
+ PCollection<KV<Row, Row>> assigned =
+ input.apply(
+ new AssignDestinationsAndPartitions(dynamicDestinations,
catalogConfig, metadataView));
+
+ PCollection<String> 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 03b3560f746..98597ba4c4f 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 class RecordWriterManagerTest {
}
}
}
+
+ @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<String, SerializableTableSpec> 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<String, SerializableTableSpec> 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<IcebergDestination> dest =
WindowedValues.valueInGlobalWindow(destination);
+
+ Map<String, SerializableTableSpec> 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<SerializableDataFile> 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 00000000000..75677ff97bc
--- /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<Map<String, SerializableTableSpec>> 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<KV<String, Row>> 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<String> 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<Map<String, SerializableTableSpec>> 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<String> shardedKey = ShardedKey.of(tableIdString, new byte[]
{0});
+ PCollection<KV<ShardedKey<String>, Iterable<Row>>> input =
+ pipeline.apply(
+ "CreateGroupedInput",
+ Create.of(KV.of(shardedKey, (Iterable<Row>) ImmutableList.of(row1,
row2)))
+ .withCoder(
+ KvCoder.of(
+ ShardedKey.Coder.of(StringUtf8Coder.of()),
+ IterableCoder.of(RowCoder.of(BEAM_SCHEMA)))));
+
+ PCollection<FileWriteResult> writtenFiles =
+ input.apply(
+ new WriteGroupedRowsToFiles(
+ catalogConfig, dynamicDestinations, "prefix", 1024L * 1024L,
null, metadataView));
+
+ PCollection<String> 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<Map<String, SerializableTableSpec>> 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<KV<String, Row>> input =
+ pipeline.apply(
+ "CreateDirectInput",
+ Create.of(KV.of(tableIdString, row1))
+ .withCoder(KvCoder.of(StringUtf8Coder.of(),
RowCoder.of(BEAM_SCHEMA))));
+
+ PCollection<FileWriteResult> writtenFiles =
+ input.apply(
+ new WriteDirectRowsToFiles(
+ catalogConfig, dynamicDestinations, "prefix", 1024L * 1024L,
null, metadataView));
+
+ PCollection<String> 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<Map<String, SerializableTableSpec>> 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<KV<Row, Iterable<Row>>> input =
+ pipeline.apply(
+ "CreatePartitionedInput",
+ Create.of(KV.of(partitionRow, (Iterable<Row>)
ImmutableList.of(dataRow)))
+ .withCoder(
+ KvCoder.of(
+
RowCoder.of(AssignDestinationsAndPartitions.OUTPUT_SCHEMA),
+ IterableCoder.of(RowCoder.of(BEAM_SCHEMA)))));
+
+ PCollection<FileWriteResult> writtenFiles =
+ input.apply(
+ new WritePartitionedRowsToFiles(
+ catalogConfig, dynamicDestinations, "prefix", null,
metadataView));
+
+ PCollection<String> 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<Map<String, SerializableTableSpec>> 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<KV<String, Row>> 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<Record> 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<Map<String, SerializableTableSpec>> 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<Row> input =
+ pipeline.apply("CreateRows", Create.of(row1,
row2).withRowSchema(BEAM_SCHEMA));
+
+ PCollection<KV<Row, Row>> 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<Record> 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<Map<String, SerializableTableSpec>> 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<KV<String, Row>> 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<String> 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();
+ }
+}
