diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriver.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriver.java new file mode 100644 index 000000000000..5da44a6c51a1 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriver.java @@ -0,0 +1,438 @@ +/* + * 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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; + +import com.google.auto.value.AutoValue; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.MapCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VoidCoder; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.state.MapState; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.Deduplicate; +import org.apache.beam.sdk.transforms.Distinct; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.DoFn.StateId; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Reshuffle; +import org.apache.beam.sdk.transforms.Sample; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.transforms.WithKeys; +import org.apache.beam.sdk.transforms.display.DisplayData; +import org.apache.beam.sdk.transforms.windowing.AfterPane; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindows; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +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.apache.beam.sdk.values.ValueInSingleWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.exceptions.NoSuchTableException; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A driver transform that extracts table identifiers from incoming {@link Row}s, deduplicates them + * per window, optionally bounds the cache size up to {@code maximumCacheSize}, loads their + * declarative metadata from the Iceberg catalog, and emits {@link KV} pairs of table identifier + * strings to {@link SerializableTableSpec}. This is intended to be used in Beam pipelines that may + * utilize a large number of workers to handle Iceberg writes, where having every worker thread + * query for table metadata results in an excessive amount of requests and a high level of + * redundancy. + * + *

Can also be materialized into a broadcasted {@link PCollectionView} via {@link + * #asView(IcebergCatalogConfig, DynamicDestinations)}. By default, the cache size is uncapped. If + * {@code maximumCacheSize} is configured and the number of distinct tables in a window exceeds it, + * up to {@code maximumCacheSize} tables are sampled into the broadcasted view, while remaining + * destinations fall back to worker-local catalog loading. + * + *

For unbounded streaming pipelines in {@link GlobalWindows}, {@link Deduplicate} is used to + * deduplicate table identifiers over the configured {@code refreshInterval} (defaulting to {@link + * #DEFAULT_REFRESH_INTERVAL}), allowing periodic refresh of table metadata when schemas evolve. + */ +@Internal +@AutoValue +public abstract class TableMetadataDriver + extends PTransform, PCollection>> { + + public static final Duration DEFAULT_REFRESH_INTERVAL = Duration.standardMinutes(5); + public static final int DEFAULT_POLLING_BUCKETS = 1; + + public abstract IcebergCatalogConfig getCatalogConfig(); + + public abstract DynamicDestinations getDynamicDestinations(); + + public abstract @Nullable Integer getMaximumCacheSize(); + + public abstract @Nullable Duration getRefreshInterval(); + + /** + * Returns the number of parallel buckets/workers used to query the Iceberg catalog, or {@code + * null} for default. + */ + public abstract @Nullable Integer getPollingBuckets(); + + public static Builder builder() { + return new AutoValue_TableMetadataDriver.Builder(); + } + + public abstract Builder toBuilder(); + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setCatalogConfig(IcebergCatalogConfig catalogConfig); + + public abstract Builder setDynamicDestinations(DynamicDestinations dynamicDestinations); + + public abstract Builder setMaximumCacheSize(@Nullable Integer maximumCacheSize); + + public abstract Builder setRefreshInterval(@Nullable Duration refreshInterval); + + /** + * Sets the number of parallel buckets (worker tasks) used to query the Iceberg catalog. + * + *

Defaults to {@link #DEFAULT_POLLING_BUCKETS} (1), which serializes all catalog lookups to + * avoid overwhelming catalog metastores (e.g. Hive Metastore, REST catalog). For pipelines + * writing to a large number of distinct dynamic tables (e.g. hundreds of tables per window), + * consider increasing this value (e.g. 5–10) to parallelize catalog lookups while still + * bounding load. + */ + public abstract Builder setPollingBuckets(@Nullable Integer pollingBuckets); + + abstract TableMetadataDriver autoBuild(); + + public TableMetadataDriver build() { + TableMetadataDriver driver = autoBuild(); + Integer maxCacheSize = driver.getMaximumCacheSize(); + if (maxCacheSize != null) { + Preconditions.checkArgument( + maxCacheSize > 0, "maximumCacheSize must be greater than 0, got %s", maxCacheSize); + } + Duration refreshInterval = driver.getRefreshInterval(); + if (refreshInterval != null) { + Preconditions.checkArgument( + refreshInterval.isLongerThan(Duration.ZERO), + "refreshInterval must be positive, got %s", + refreshInterval); + } + Integer pollingBuckets = driver.getPollingBuckets(); + if (pollingBuckets != null) { + Preconditions.checkArgument( + pollingBuckets > 0, "pollingBuckets must be greater than 0, got %s", pollingBuckets); + } + return driver; + } + } + + /** + * Helper that applies {@link TableMetadataDriver} and creates an uncapped {@link PCollectionView} + * of {@link Map} of table identifier strings to {@link SerializableTableSpec}. + */ + public static PTransform, PCollectionView>> + asView(IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations) { + return asView(catalogConfig, dynamicDestinations, null, null, null); + } + + /** + * Helper that applies {@link TableMetadataDriver} with an optional {@code maximumCacheSize} limit + * and creates a {@link PCollectionView} of {@link Map} of table identifier strings to {@link + * SerializableTableSpec}. + * + * @param catalogConfig the catalog configuration used to poll metadata. + * @param dynamicDestinations destination strategy extracting table IDs from rows. + * @param maximumCacheSize optional maximum distinct tables to poll and broadcast per window (null + * for uncapped). + */ + public static PTransform, PCollectionView>> + asView( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Integer maximumCacheSize) { + return asView(catalogConfig, dynamicDestinations, maximumCacheSize, null, null); + } + + /** + * Helper that applies {@link TableMetadataDriver} with an optional {@code maximumCacheSize} limit + * and custom {@code refreshInterval}, creating a {@link PCollectionView} of {@link Map} of table + * identifier strings to {@link SerializableTableSpec}. + * + * @param catalogConfig the catalog configuration used to poll metadata. + * @param dynamicDestinations destination strategy extracting table IDs from rows. + * @param maximumCacheSize optional maximum distinct tables to poll and broadcast per window (null + * for uncapped). + * @param refreshInterval optional refresh interval for streaming global window triggers. + */ + public static PTransform, PCollectionView>> + asView( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Integer maximumCacheSize, + @Nullable Duration refreshInterval) { + return asView(catalogConfig, dynamicDestinations, maximumCacheSize, refreshInterval, null); + } + + /** + * Helper that applies {@link TableMetadataDriver} with an optional {@code maximumCacheSize} + * limit, custom {@code refreshInterval}, and custom {@code pollingBuckets}, creating a {@link + * PCollectionView} of {@link Map} of table identifier strings to {@link SerializableTableSpec}. + * + * @param catalogConfig the catalog configuration used to poll metadata. + * @param dynamicDestinations destination strategy extracting table IDs from rows. + * @param maximumCacheSize optional maximum distinct tables to poll and broadcast per window (null + * for uncapped). + * @param refreshInterval optional refresh interval for streaming global window triggers. + * @param pollingBuckets optional number of parallel buckets/workers for catalog polling. + */ + public static PTransform, PCollectionView>> + asView( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Integer maximumCacheSize, + @Nullable Duration refreshInterval, + @Nullable Integer pollingBuckets) { + return new PTransform, PCollectionView>>() { + @Override + public PCollectionView> expand(PCollection input) { + boolean isStreaming = input.isBounded() == PCollection.IsBounded.UNBOUNDED; + + PCollection> specs = + input.apply( + "GenerateTableMetadata", + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .setMaximumCacheSize(maximumCacheSize) + .setRefreshInterval(refreshInterval) + .setPollingBuckets(pollingBuckets) + .build()); + + if (isStreaming) { + return specs + .apply("KeyForGlobalCache", WithKeys.of((Void) null)) + .setCoder(KvCoder.of(VoidCoder.of(), specs.getCoder())) + .apply("AccumulateCacheMap", ParDo.of(new AccumulateTableMetadataMapDoFn())) + .setCoder(MapCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())) + .apply( + "StreamingCacheWindow", + Window.>into(new GlobalWindows()) + .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1))) + .discardingFiredPanes()) + .apply( + "CreateMetadataSingletonView", + Combine.globally(new MapMergerFn()).asSingletonView()); + } + + return specs.apply("CreateTableMetadataView", View.asMap()); + } + }; + } + + @Override + public PCollection> expand(PCollection input) { + PCollection tableIds = + input + .apply("ExtractTableIds", ParDo.of(new ExtractTableIdsDoFn(getDynamicDestinations()))) + .setCoder(StringUtf8Coder.of()) + .apply("MetadataGlobalWindow", Window.into(new GlobalWindows())); + + boolean isStreaming = input.isBounded() == PCollection.IsBounded.UNBOUNDED; + + PCollection distinctTableIds; + if (isStreaming) { + Duration customInterval = getRefreshInterval(); + Duration interval = + checkNotNull(customInterval != null ? customInterval : DEFAULT_REFRESH_INTERVAL); + distinctTableIds = + tableIds.apply( + "DeduplicateTableIds", Deduplicate.values().withDuration(interval)); + } else { + distinctTableIds = tableIds.apply("DistinctTableIds", Distinct.create()); + } + + PCollection cachedTableIds; + Integer maxCacheSize = getMaximumCacheSize(); + if (maxCacheSize != null) { + if (isStreaming) { + throw new UnsupportedOperationException( + "maximumCacheSize is currently not supported for unbounded streaming pipelines."); + } + cachedTableIds = distinctTableIds.apply("CapCacheSize", Sample.any(maxCacheSize)); + } else { + cachedTableIds = distinctTableIds; + } + + @Nullable Integer configuredBuckets = getPollingBuckets(); + int pollingBuckets = configuredBuckets != null ? configuredBuckets : DEFAULT_POLLING_BUCKETS; + PCollection pollingTableIds = + cachedTableIds.apply( + "ReshufflePollingBuckets", + Reshuffle.viaRandomKey().withNumBuckets(pollingBuckets)); + + PCollection> specs = + pollingTableIds + .apply("PollTableMetadata", ParDo.of(new CatalogPollingDoFn(getCatalogConfig()))) + .setCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())); + + if (isStreaming) { + return specs.apply( + "ApplyStreamingTrigger", + Window.>into(new GlobalWindows()) + .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1))) + .discardingFiredPanes()); + } + return specs; + } + + @Override + public void populateDisplayData(DisplayData.Builder builder) { + super.populateDisplayData(builder); + builder.addIfNotNull( + DisplayData.item("maximumCacheSize", getMaximumCacheSize()) + .withLabel("Maximum Cache Size")); + builder.addIfNotNull( + DisplayData.item("refreshInterval", getRefreshInterval()) + .withLabel("Table Metadata Refresh Interval")); + builder.addIfNotNull( + DisplayData.item("pollingBuckets", getPollingBuckets()) + .withLabel("Catalog Polling Buckets")); + } + + static class ExtractTableIdsDoFn extends DoFn { + private final DynamicDestinations dynamicDestinations; + + ExtractTableIdsDoFn(DynamicDestinations dynamicDestinations) { + this.dynamicDestinations = dynamicDestinations; + } + + @ProcessElement + public void processElement( + @Element Row element, + BoundedWindow window, + PaneInfo paneInfo, + @Timestamp Instant timestamp, + OutputReceiver out) { + String tableIdentifier = + dynamicDestinations.getTableStringIdentifier( + ValueInSingleWindow.of(element, timestamp, window, paneInfo)); + if (tableIdentifier != null && !tableIdentifier.trim().isEmpty()) { + out.output(tableIdentifier.trim()); + } + } + } + + static class CatalogPollingDoFn extends DoFn> { + private static final Logger LOG = LoggerFactory.getLogger(CatalogPollingDoFn.class); + private static final Counter TABLES_POLLED_COUNTER = + Metrics.counter(TableMetadataDriver.class, "tablesPolled"); + private static final Counter TABLES_SKIPPED_MISSING_COUNTER = + Metrics.counter(TableMetadataDriver.class, "tablesSkippedMissing"); + + private final IcebergCatalogConfig catalogConfig; + + CatalogPollingDoFn(IcebergCatalogConfig catalogConfig) { + this.catalogConfig = catalogConfig; + } + + @ProcessElement + public void processElement( + @Element String tableIdString, OutputReceiver> out) { + try { + TableIdentifier tableId = IcebergUtils.parseTableIdentifier(tableIdString); + Table table = catalogConfig.catalog().loadTable(tableId); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableIdString, table); + TABLES_POLLED_COUNTER.inc(); + out.output(KV.of(tableIdString, spec)); + } catch (NoSuchTableException e) { + LOG.info( + "Table '{}' does not exist in catalog. Skipping metadata emission for side-input view.", + tableIdString); + TABLES_SKIPPED_MISSING_COUNTER.inc(); + } catch (IllegalArgumentException e) { + LOG.warn( + "Failed to parse table identifier '{}'. Skipping metadata emission for side-input view.", + tableIdString, + e); + TABLES_SKIPPED_MISSING_COUNTER.inc(); + } + } + } + + static class AccumulateTableMetadataMapDoFn + extends DoFn< + KV>, Map> { + @StateId("tableCache") + private final StateSpec> cacheStateSpec = + StateSpecs.map(StringUtf8Coder.of(), SerializableTableSpec.getCoder()); + + @ProcessElement + public void processElement( + @Element KV> element, + @StateId("tableCache") MapState cacheState, + OutputReceiver> out) { + KV kv = element.getValue(); + cacheState.put(kv.getKey(), kv.getValue()); + + Map mapSnapshot = new HashMap<>(); + for (Map.Entry entry : cacheState.entries().read()) { + mapSnapshot.put(entry.getKey(), entry.getValue()); + } + out.output(Collections.unmodifiableMap(mapSnapshot)); + } + } + + static class MapMergerFn extends Combine.BinaryCombineFn> { + @Override + public Map apply( + Map left, Map right) { + if (left == null || left.isEmpty()) { + return right != null ? right : Collections.emptyMap(); + } + if (right == null || right.isEmpty()) { + return left; + } + Map merged = new HashMap<>(left); + merged.putAll(right); + return Collections.unmodifiableMap(merged); + } + + @Override + public Map identity() { + return Collections.emptyMap(); + } + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriverTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriverTest.java new file mode 100644 index 000000000000..37c80631ff58 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriverTest.java @@ -0,0 +1,1028 @@ +/* + * 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 static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestStream; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.display.DisplayData; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +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.apache.beam.sdk.values.TimestampedValue; +import org.apache.beam.sdk.values.ValueInSingleWindow; +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.FileFormat; +import org.apache.iceberg.PartitionKey; +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.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Types; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +public class TableMetadataDriverTest implements Serializable { + + @Rule public transient TestPipeline pipeline = TestPipeline.create(); + @Rule public transient TemporaryFolder tempFolder = new TemporaryFolder(); + + private String warehouseLocation; + private IcebergCatalogConfig catalogConfig; + + private static final Schema BEAM_SCHEMA = + Schema.builder() + .addInt64Field("id") + .addStringField("data") + .addNullableStringField("dest") + .build(); + + private static final org.apache.iceberg.Schema ICEBERG_SCHEMA = + IcebergUtils.beamSchemaToIcebergSchema( + Schema.builder().addInt64Field("id").addStringField("data").build()); + + private static final TableIdentifier TABLE_ID = TableIdentifier.of("default", "table"); + + private static final DynamicDestinations SINGLE_TABLE_DYNAMIC_DESTINATIONS = + DynamicDestinations.singleTable(TABLE_ID, BEAM_SCHEMA); + + private static final DynamicDestinations DYNAMIC_DESTINATIONS = + new DynamicDestinations() { + @Override + public Schema getDataSchema() { + return BEAM_SCHEMA; + } + + @Override + public Row getData(Row element) { + return element; + } + + @Override + public IcebergDestination instantiateDestination(String destination) { + return IcebergDestination.builder() + .setTableIdentifier(IcebergUtils.parseTableIdentifier(destination)) + .build(); + } + + @Override + public String getTableStringIdentifier(ValueInSingleWindow element) { + return element.getValue().getString("dest"); + } + }; + + @Before + public void setUp() throws Exception { + warehouseLocation = "file:" + tempFolder.newFolder().getAbsolutePath(); + catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogName("hadoop") + .setCatalogProperties(ImmutableMap.of("type", "hadoop", "warehouse", warehouseLocation)) + .build(); + } + + private Catalog getCatalog() { + return CatalogUtil.loadCatalog( + CatalogUtil.ICEBERG_CATALOG_HADOOP, + "hadoop", + ImmutableMap.of(CatalogProperties.WAREHOUSE_LOCATION, warehouseLocation), + new Configuration()); + } + + @Test + public void testSingleTableExtractionAndSpecOutput() { + Table realTable = getCatalog().createTable(TABLE_ID, ICEBERG_SCHEMA); + + List rows = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + rows.add( + Row.withSchema(BEAM_SCHEMA) + .withFieldValue("id", (long) i) + .withFieldValue("data", "val_" + i) + .withFieldValue("dest", null) + .build()); + } + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .build()); + + String expectedTableIdString = IcebergUtils.tableIdentifierToString(TABLE_ID); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(1, list.size()); + KV kv = list.get(0); + assertEquals(expectedTableIdString, kv.getKey()); + SerializableTableSpec spec = kv.getValue(); + assertNotNull(spec); + assertEquals(realTable.name(), spec.getName()); + assertEquals(realTable.location(), spec.getLocation()); + assertEquals(realTable.schema().asStruct(), spec.getSchema().asStruct()); + assertEquals(realTable.spec(), spec.getPartitionSpec()); + assertNotNull(spec.getFileIO()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testMultipleDynamicDestinationsExtraction() { + Catalog catalog = getCatalog(); + TableIdentifier tableA = TableIdentifier.of("default", "table_a"); + TableIdentifier tableB = TableIdentifier.of("default", "table_b"); + TableIdentifier tableC = TableIdentifier.of("default", "table_c"); + + catalog.createTable(tableA, ICEBERG_SCHEMA); + catalog.createTable(tableB, ICEBERG_SCHEMA); + catalog.createTable(tableC, ICEBERG_SCHEMA); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.table_a").build(), + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.table_b").build(), + Row.withSchema(BEAM_SCHEMA).addValues(3L, "v3", "default.table_c").build(), + Row.withSchema(BEAM_SCHEMA).addValues(4L, "v4", "default.table_a").build(), + Row.withSchema(BEAM_SCHEMA).addValues(5L, "v5", "default.table_b").build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(3, list.size()); + Map map = + list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); + assertTrue(map.containsKey("default.table_a")); + assertTrue(map.containsKey("default.table_b")); + assertTrue(map.containsKey("default.table_c")); + return null; + }); + + pipeline.run(); + } + + @Test + public void testWindowedDeduplication() { + Catalog catalog = getCatalog(); + TableIdentifier table1 = TableIdentifier.of("default", "t1"); + TableIdentifier table2 = TableIdentifier.of("default", "t2"); + + catalog.createTable(table1, ICEBERG_SCHEMA); + catalog.createTable(table2, ICEBERG_SCHEMA); + + List rows = new ArrayList<>(); + for (int i = 0; i < 100; i++) { + String dest = (i % 2 == 0) ? "default.t1" : "default.t2"; + rows.add(Row.withSchema(BEAM_SCHEMA).addValues((long) i, "val_" + i, dest).build()); + } + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + Map map = + list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); + assertTrue(map.containsKey("default.t1")); + assertTrue(map.containsKey("default.t2")); + return null; + }); + + pipeline.run(); + } + + @Test + public void testUnboundedGlobalWindowStreamingDeduplication() { + Catalog catalog = getCatalog(); + TableIdentifier table1 = TableIdentifier.of("default", "stream_t1"); + TableIdentifier table2 = TableIdentifier.of("default", "stream_t2"); + + catalog.createTable(table1, ICEBERG_SCHEMA); + catalog.createTable(table2, ICEBERG_SCHEMA); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.stream_t1").build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.stream_t2").build(); + Row row3 = Row.withSchema(BEAM_SCHEMA).addValues(3L, "v3", "default.stream_t1").build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(row1) + .addElements(row2) + .addElements(row3) + .advanceProcessingTime(Duration.standardSeconds(5)) + .advanceWatermarkToInfinity(); + + PCollection input = pipeline.apply("StreamInput", stream); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(2)) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + Map map = + list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); + assertTrue(map.containsKey("default.stream_t1")); + assertTrue(map.containsKey("default.stream_t2")); + return null; + }); + + pipeline.run(); + } + + @Test + public void testMetadataRefreshedAcrossIntervals() { + Catalog catalog = getCatalog(); + TableIdentifier tableId = TableIdentifier.of("default", "evolving_table"); + catalog.createTable(tableId, ICEBERG_SCHEMA); + + Row row1 = + Row.withSchema(BEAM_SCHEMA).addValues(1L, "initial_data", "default.evolving_table").build(); + Row row2 = + Row.withSchema(BEAM_SCHEMA) + .addValues(2L, "trigger_update", "default.evolving_table") + .build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(row1) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(row2) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceWatermarkToInfinity(); + + PCollection input = + pipeline + .apply("StreamInput", stream) + .apply( + "EvolveSchemaOnTriggerRow", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(@Element Row row, OutputReceiver out) { + if ("trigger_update".equals(row.getString("data"))) { + Table table = + catalogConfig + .catalog() + .loadTable( + IcebergUtils.parseTableIdentifier("default.evolving_table")); + table + .updateSchema() + .addColumn("new_col", Types.StringType.get()) + .commit(); + } + out.output(row); + } + })) + .setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(2)) + .build()); + + // Downstream consumer transform verifying that updated metadata is received + PCollection consumerReceivedSchemas = + specs.apply( + "ConsumerTransform", + ParDo.of( + new DoFn, String>() { + @ProcessElement + public void processElement( + @Element KV element, + OutputReceiver out) { + boolean hasNewCol = element.getValue().getSchema().findField("new_col") != null; + out.output(hasNewCol ? "UPDATED_SCHEMA" : "INITIAL_SCHEMA"); + } + })); + + PAssert.that(consumerReceivedSchemas).containsInAnyOrder("INITIAL_SCHEMA", "UPDATED_SCHEMA"); + + pipeline.run(); + } + + @Test + public void testMetadataRefreshedAcrossIntervalsAsSideInput() { + Catalog catalog = getCatalog(); + TableIdentifier tableId = TableIdentifier.of("default", "evolving_side_input_table"); + catalog.createTable(tableId, ICEBERG_SCHEMA); + + String tableIdStr = "default.evolving_side_input_table"; + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1L, "initial_data", tableIdStr).build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2L, "trigger_update", tableIdStr).build(); + Row row3 = Row.withSchema(BEAM_SCHEMA).addValues(3L, "post_update_data", tableIdStr).build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(row1) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(row2) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(row3) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceWatermarkToInfinity(); + + PCollection input = + pipeline + .apply("StreamInput", stream) + .apply( + "EvolveSchemaOnTriggerRow", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(@Element Row row, OutputReceiver out) { + if ("trigger_update".equals(row.getString("data"))) { + Table table = + catalogConfig + .catalog() + .loadTable( + IcebergUtils.parseTableIdentifier( + "default.evolving_side_input_table")); + table + .updateSchema() + .addColumn("new_col", Types.StringType.get()) + .commit(); + } + out.output(row); + } + })) + .setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", + TableMetadataDriver.asView( + catalogConfig, DYNAMIC_DESTINATIONS, null, Duration.standardSeconds(2))); + + PCollection consumerObserved = + input.apply( + "ConsumeSideInput", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element Row row, OutputReceiver out, ProcessContext c) { + if ("trigger_update".equals(row.getString("data"))) { + return; + } + Map viewMap = c.sideInput(metadataView); + SerializableTableSpec spec = viewMap.get(row.getString("dest")); + assertNotNull("Expected spec in side input view", spec); + + SideInputTable sideInputTable = new SideInputTable(spec); + boolean hasNewCol = sideInputTable.schema().findField("new_col") != null; + out.output( + row.getString("data") + + ":" + + (hasNewCol ? "UPDATED_SCHEMA" : "INITIAL_SCHEMA")); + } + }) + .withSideInputs(metadataView)); + + PAssert.that(consumerObserved) + .containsInAnyOrder("initial_data:INITIAL_SCHEMA", "post_update_data:UPDATED_SCHEMA"); + + pipeline.run(); + } + + @Test + public void testMetadataRefreshedAcrossIntervalsAsSideInputWithMultipleTables() { + Catalog catalog = getCatalog(); + TableIdentifier tableA = TableIdentifier.of("default", "multi_table_a"); + TableIdentifier tableB = TableIdentifier.of("default", "multi_table_b"); + catalog.createTable(tableA, ICEBERG_SCHEMA); + catalog.createTable(tableB, ICEBERG_SCHEMA); + + String tableAStr = "default.multi_table_a"; + String tableBStr = "default.multi_table_b"; + + Row rowSeedA = Row.withSchema(BEAM_SCHEMA).addValues(0L, "seed_a", tableAStr).build(); + Row rowSeedB = Row.withSchema(BEAM_SCHEMA).addValues(0L, "seed_b", tableBStr).build(); + Row rowA1 = Row.withSchema(BEAM_SCHEMA).addValues(1L, "a1", tableAStr).build(); + Row rowB1 = Row.withSchema(BEAM_SCHEMA).addValues(2L, "b1", tableBStr).build(); + Row rowTriggerUpdateA = + Row.withSchema(BEAM_SCHEMA).addValues(3L, "trigger_update_a", tableAStr).build(); + Row rowA2 = Row.withSchema(BEAM_SCHEMA).addValues(4L, "a2", tableAStr).build(); + Row rowB2 = Row.withSchema(BEAM_SCHEMA).addValues(5L, "b2", tableBStr).build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(rowSeedA, rowSeedB) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(rowA1, rowB1) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(rowTriggerUpdateA) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(rowA2, rowB2) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceWatermarkToInfinity(); + + PCollection input = + pipeline + .apply("StreamInput", stream) + .apply( + "EvolveSchemaOnTriggerRow", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(@Element Row row, OutputReceiver out) { + if ("trigger_update_a".equals(row.getString("data"))) { + Table table = + catalogConfig + .catalog() + .loadTable( + IcebergUtils.parseTableIdentifier("default.multi_table_a")); + table + .updateSchema() + .addColumn("new_col_a", Types.StringType.get()) + .commit(); + } + out.output(row); + } + })) + .setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", + TableMetadataDriver.asView( + catalogConfig, DYNAMIC_DESTINATIONS, null, Duration.standardSeconds(2))); + + PCollection consumerObserved = + input.apply( + "ConsumeSideInput", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element Row row, OutputReceiver out, ProcessContext c) { + String data = row.getString("data"); + if ("seed_a".equals(data) + || "seed_b".equals(data) + || "trigger_update_a".equals(data)) { + return; + } + Map viewMap = c.sideInput(metadataView); + SerializableTableSpec spec = viewMap.get(row.getString("dest")); + assertNotNull( + "Expected table " + row.getString("dest") + " in side input view", + spec); + + SideInputTable sideInputTable = new SideInputTable(spec); + boolean hasNewColA = sideInputTable.schema().findField("new_col_a") != null; + out.output( + row.getString("data") + + ":" + + (hasNewColA ? "UPDATED_SCHEMA" : "INITIAL_SCHEMA")); + } + }) + .withSideInputs(metadataView)); + + PAssert.that(consumerObserved) + .containsInAnyOrder( + "a1:INITIAL_SCHEMA", "b1:INITIAL_SCHEMA", "a2:UPDATED_SCHEMA", "b2:INITIAL_SCHEMA"); + + pipeline.run(); + } + + @Test + public void testMaximumCacheSizeInStreamingThrowsUnsupportedOperationException() { + pipeline.enableAbandonedNodeEnforcement(false); + Row row = Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.test_table").build(); + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(row) + .advanceWatermarkToInfinity(); + + PCollection input = pipeline.apply("StreamInput", stream); + + assertThrows( + UnsupportedOperationException.class, + () -> + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setMaximumCacheSize(5) + .build())); + } + + @Test + public void testMalformedTableIdentifierSkippedWithoutFailingBundle() { + Row validRow = Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.valid_table").build(); + Row malformedRow = + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.invalid..name///").build(); + + getCatalog().createTable(TableIdentifier.of("default", "valid_table"), ICEBERG_SCHEMA); + + PCollection input = + pipeline.apply(Create.of(validRow, malformedRow)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(1, list.size()); + assertEquals("default.valid_table", list.get(0).getKey()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testMaximumCacheSizeCap() { + Catalog catalog = getCatalog(); + for (int i = 1; i <= 6; i++) { + catalog.createTable(TableIdentifier.of("default", "cap_table_" + i), ICEBERG_SCHEMA); + } + + List rows = new ArrayList<>(); + for (int i = 1; i <= 6; i++) { + rows.add( + Row.withSchema(BEAM_SCHEMA) + .addValues((long) i, "v_" + i, "default.cap_table_" + i) + .build()); + } + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + int maxCacheSize = 3; + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setMaximumCacheSize(maxCacheSize) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(maxCacheSize, list.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testUncappedByDefault() { + Catalog catalog = getCatalog(); + for (int i = 1; i <= 10; i++) { + catalog.createTable(TableIdentifier.of("default", "uncapped_table_" + i), ICEBERG_SCHEMA); + } + + List rows = new ArrayList<>(); + for (int i = 1; i <= 10; i++) { + rows.add( + Row.withSchema(BEAM_SCHEMA) + .addValues((long) i, "v_" + i, "default.uncapped_table_" + i) + .build()); + } + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + // Without setting maximumCacheSize, all 10 distinct tables are emitted + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(10, list.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testNonExistentTableIsSkippedWithoutFailingBundle() { + Catalog catalog = getCatalog(); + TableIdentifier validTable = TableIdentifier.of("default", "existing_table"); + catalog.createTable(validTable, ICEBERG_SCHEMA); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.existing_table").build(), + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.non_existent_table").build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + // Only the existing table is emitted; the non-existent table is skipped without failing bundle + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(1, list.size()); + assertEquals("default.existing_table", list.get(0).getKey()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testFiltersNullAndBlankTableIdentifiers() { + TableIdentifier validTableId = TableIdentifier.of("default", "valid_dest_table"); + getCatalog().createTable(validTableId, ICEBERG_SCHEMA); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", null).build(), + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "").build(), + Row.withSchema(BEAM_SCHEMA).addValues(3L, "v3", " ").build(), + Row.withSchema(BEAM_SCHEMA).addValues(4L, "v4", "default.valid_dest_table").build(), + Row.withSchema(BEAM_SCHEMA) + .addValues(5L, "v5", " default.valid_dest_table ") + .build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(1, list.size()); + assertEquals("default.valid_dest_table", list.get(0).getKey()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testInvalidMaximumCacheSizeThrowsException() { + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setMaximumCacheSize(0) + .build()); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setMaximumCacheSize(-5) + .build()); + } + + @Test + public void testInvalidRefreshIntervalThrowsException() { + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.ZERO) + .build()); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(-5)) + .build()); + } + + @Test + public void testInvalidPollingBucketsThrowsException() { + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setPollingBuckets(0) + .build()); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setPollingBuckets(-2) + .build()); + } + + @Test + public void testConfigurablePollingBuckets() { + Catalog catalog = getCatalog(); + TableIdentifier table1 = TableIdentifier.of("default", "bucket_t1"); + TableIdentifier table2 = TableIdentifier.of("default", "bucket_t2"); + catalog.createTable(table1, ICEBERG_SCHEMA); + catalog.createTable(table2, ICEBERG_SCHEMA); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.bucket_t1").build(), + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.bucket_t2").build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setPollingBuckets(2) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + Map map = + list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); + assertTrue(map.containsKey("default.bucket_t1")); + assertTrue(map.containsKey("default.bucket_t2")); + return null; + }); + + pipeline.run(); + } + + @Test + public void testWindowPreservation() { + Catalog catalog = getCatalog(); + TableIdentifier tableW1 = TableIdentifier.of("default", "table_w1"); + TableIdentifier tableW2 = TableIdentifier.of("default", "table_w2"); + + catalog.createTable(tableW1, ICEBERG_SCHEMA); + catalog.createTable(tableW2, ICEBERG_SCHEMA); + + Instant t1 = new Instant(1000); + Instant t2 = new Instant(70000); + + PCollection input = + pipeline + .apply( + Create.timestamped( + TimestampedValue.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.table_w1").build(), + t1), + TimestampedValue.of( + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.table_w2").build(), + t2))) + .setCoder(RowCoder.of(BEAM_SCHEMA)) + .apply(Window.into(FixedWindows.of(Duration.standardMinutes(1)))); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testEmptyInputProducesEmptyOutput() { + getCatalog().createTable(TABLE_ID, ICEBERG_SCHEMA); + + PCollection input = pipeline.apply(Create.empty(RowCoder.of(BEAM_SCHEMA))); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs).empty(); + + pipeline.run(); + } + + @Test + public void testDisplayData() { + TableMetadataDriver driver = + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setMaximumCacheSize(42) + .setRefreshInterval(Duration.standardMinutes(10)) + .setPollingBuckets(3) + .build(); + + DisplayData displayData = DisplayData.from(driver); + Map items = displayData.asMap(); + + assertNotNull(displayData); + boolean hasCacheSize = false; + boolean hasRefreshInterval = false; + boolean hasPollingBuckets = false; + for (DisplayData.Item item : items.values()) { + if ("maximumCacheSize".equals(item.getKey())) { + assertEquals(42L, item.getValue()); + hasCacheSize = true; + } + if ("refreshInterval".equals(item.getKey())) { + hasRefreshInterval = true; + } + if ("pollingBuckets".equals(item.getKey())) { + assertEquals(3L, item.getValue()); + hasPollingBuckets = true; + } + } + assertTrue(hasCacheSize); + assertTrue(hasRefreshInterval); + assertTrue(hasPollingBuckets); + } + + @Test + public void testViewAsMapIntegration() { + PartitionSpec partitionSpec = PartitionSpec.builderFor(ICEBERG_SCHEMA).identity("data").build(); + getCatalog().createTable(TABLE_ID, ICEBERG_SCHEMA, partitionSpec); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(10L, "partition_val_a", null).build(), + Row.withSchema(BEAM_SCHEMA).addValues(20L, "partition_val_b", null).build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", + TableMetadataDriver.asView(catalogConfig, SINGLE_TABLE_DYNAMIC_DESTINATIONS)); + + String expectedTableIdString = IcebergUtils.tableIdentifierToString(TABLE_ID); + + PCollection writtenFiles = + input.apply( + "WriteWithSideInputTable", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element Row row, OutputReceiver out, ProcessContext c) + throws Exception { + Map viewMap = c.sideInput(metadataView); + SerializableTableSpec spec = viewMap.get(expectedTableIdString); + assertNotNull(spec); + + SideInputTable sideInputTable = new SideInputTable(spec); + PartitionKey partitionKey = + new PartitionKey(sideInputTable.spec(), sideInputTable.schema()); + Record record = GenericRecord.create(sideInputTable.schema()); + record.setField("id", row.getInt64("id")); + record.setField("data", row.getString("data")); + partitionKey.partition(record); + + RecordWriter writer = + new RecordWriter( + sideInputTable, + FileFormat.PARQUET, + "side_input_test_file_" + row.getInt64("id"), + partitionKey, + ImmutableMap.of()); + writer.write(record); + writer.close(); + + out.output(writer.getDataFile().path().toString()); + } + }) + .withSideInputs(metadataView)); + + PAssert.that(writtenFiles) + .satisfies( + files -> { + List paths = ImmutableList.copyOf(files); + assertEquals(2, paths.size()); + return null; + }); + + pipeline.run(); + } +}