From 83ddc7d178eb491f597d5a6ff248156aec5c6409 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 25 Aug 2026 15:02:05 +0000 Subject: [PATCH 1/7] Implement Iceberg Table Metadata Driver transform --- .../sdk/io/iceberg/TableMetadataDriver.java | 195 ++++++ .../io/iceberg/TableMetadataDriverTest.java | 581 ++++++++++++++++++ 2 files changed, 776 insertions(+) create mode 100644 sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriver.java create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriverTest.java 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..6d9e00f1f12f --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriver.java @@ -0,0 +1,195 @@ +/* + * 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 com.google.auto.value.AutoValue; +import java.util.Map; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.transforms.Distinct; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Sample; +import org.apache.beam.sdk.transforms.View; +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.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.joda.time.Instant; + +/** + * A driver transform that extracts table identifiers from incoming {@link Row}s, deduplicates them + * per window, samples up to a maximum number of tables, loads their declarative metadata from the + * Iceberg catalog, and emits {@link KV} pairs of table identifier strings to {@link + * SerializableTableSpec}. + * + *

Can also be materialized into a broadcasted {@link PCollectionView} via {@link + * #asView(IcebergCatalogConfig, DynamicDestinations)}. If the number of distinct tables in a window + * exceeds {@code maxTables}, up to {@code maxTables} tables are sampled into the broadcasted view, + * while remaining destinations can fall back to worker-local catalog loading. + */ +@Internal +@AutoValue +public abstract class TableMetadataDriver + extends PTransform, PCollection>> { + + public static final int DEFAULT_MAX_TABLES = 100; + + public abstract IcebergCatalogConfig getCatalogConfig(); + + public abstract DynamicDestinations getDynamicDestinations(); + + public abstract int getMaxTables(); + + public static Builder builder() { + return new AutoValue_TableMetadataDriver.Builder().setMaxTables(DEFAULT_MAX_TABLES); + } + + 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 setMaxTables(int maxTables); + + abstract TableMetadataDriver autoBuild(); + + public TableMetadataDriver build() { + TableMetadataDriver driver = autoBuild(); + Preconditions.checkArgument( + driver.getMaxTables() > 0, + "maxTables must be greater than 0, got %s", + driver.getMaxTables()); + return driver; + } + } + + /** + * Helper that applies {@link TableMetadataDriver} and creates a {@link PCollectionView} of {@link + * Map} of table identifier strings to {@link SerializableTableSpec} using {@link + * #DEFAULT_MAX_TABLES}. + */ + public static PTransform, PCollectionView>> + asView(IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations) { + return asView(catalogConfig, dynamicDestinations, DEFAULT_MAX_TABLES); + } + + /** + * Helper that applies {@link TableMetadataDriver} with a custom {@code maxTables} 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 maxTables maximum distinct tables to poll and broadcast per window. + */ + public static PTransform, PCollectionView>> + asView( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + int maxTables) { + return new PTransform, PCollectionView>>() { + @Override + public PCollectionView> expand(PCollection input) { + return input + .apply( + "GenerateTableMetadata", + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .setMaxTables(maxTables) + .build()) + .apply("CreateTableMetadataView", View.asMap()); + } + }; + } + + @Override + public PCollection> expand(PCollection input) { + PCollection tableIds = + input + .apply("ExtractTableIds", ParDo.of(new ExtractTableIdsDoFn(getDynamicDestinations()))) + .setCoder(StringUtf8Coder.of()); + + PCollection distinctTableIds = tableIds.apply("DistinctTableIds", Distinct.create()); + + PCollection sampledTableIds = + distinctTableIds.apply("SampleTableIds", Sample.any(getMaxTables())); + + return sampledTableIds + .apply("PollTableMetadata", ParDo.of(new CatalogPollingDoFn(getCatalogConfig()))) + .setCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())); + } + + 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 Counter TABLES_POLLED_COUNTER = + Metrics.counter(TableMetadataDriver.class, "tablesPolled"); + + private final IcebergCatalogConfig catalogConfig; + + CatalogPollingDoFn(IcebergCatalogConfig catalogConfig) { + this.catalogConfig = catalogConfig; + } + + @ProcessElement + public void processElement( + @Element String tableIdString, OutputReceiver> out) { + 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)); + } + } +} 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..73c6ac26d1b8 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriverTest.java @@ -0,0 +1,581 @@ +/* + * 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.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +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.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()); + + @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() { + TableIdentifier tableId = TableIdentifier.of("default", "single_table"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_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(dynamicDestinations) + .build()); + + String expectedTableIdString = IcebergUtils.tableIdentifierToString(tableId); + + 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); + + DynamicDestinations dynamicDestinations = + 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"); + } + }; + + 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(dynamicDestinations) + .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); + + DynamicDestinations dynamicDestinations = + 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"); + } + }; + + 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(dynamicDestinations) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testMaxTablesCapSampling() { + Catalog catalog = getCatalog(); + for (int i = 1; i <= 6; i++) { + catalog.createTable(TableIdentifier.of("default", "cap_table_" + i), ICEBERG_SCHEMA); + } + + DynamicDestinations dynamicDestinations = + 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"); + } + }; + + 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 maxTables = 3; + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .setMaxTables(maxTables) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(maxTables, list.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testFiltersNullAndBlankTableIdentifiers() { + TableIdentifier validTableId = TableIdentifier.of("default", "valid_dest_table"); + getCatalog().createTable(validTableId, ICEBERG_SCHEMA); + + DynamicDestinations dynamicDestinations = + 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"); + } + }; + + 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(dynamicDestinations) + .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 testInvalidMaxTablesThrowsException() { + TableIdentifier tableId = TableIdentifier.of("default", "dummy_table"); + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .setMaxTables(0) + .build()); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .setMaxTables(-5) + .build()); + } + + @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); + + DynamicDestinations dynamicDestinations = + 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"); + } + }; + + 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(dynamicDestinations) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testEmptyInputProducesEmptyOutput() { + TableIdentifier tableId = TableIdentifier.of("default", "empty_input_table"); + getCatalog().createTable(tableId, ICEBERG_SCHEMA); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + PCollection input = pipeline.apply(Create.empty(RowCoder.of(BEAM_SCHEMA))); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .build()); + + PAssert.that(specs).empty(); + + pipeline.run(); + } + + @Test + public void testViewAsMapIntegration() { + TableIdentifier tableId = TableIdentifier.of("default", "view_integration_table"); + PartitionSpec partitionSpec = PartitionSpec.builderFor(ICEBERG_SCHEMA).identity("data").build(); + getCatalog().createTable(tableId, ICEBERG_SCHEMA, partitionSpec); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + 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, dynamicDestinations)); + + String expectedTableIdString = IcebergUtils.tableIdentifierToString(tableId); + + 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(); + } +} From 11fffa96309ab8d73aba4d4cf7bdec2cf7925f1d Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 27 Aug 2026 14:46:10 +0000 Subject: [PATCH 2/7] Rename unit test for clarity --- .../org/apache/beam/sdk/io/iceberg/TableMetadataDriverTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 73c6ac26d1b8..0fe886916cf5 100644 --- 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 @@ -216,7 +216,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { } @Test - public void testWindowedDeduplication() { + public void testDeduplicationOfTablesAcrossRows() { Catalog catalog = getCatalog(); TableIdentifier table1 = TableIdentifier.of("default", "t1"); TableIdentifier table2 = TableIdentifier.of("default", "t2"); From 57d0db146d16defbeff77b90059eb4dc0eba60c4 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Mon, 31 Aug 2026 15:57:40 +0000 Subject: [PATCH 3/7] Uncap cache size by default --- .../sdk/io/iceberg/TableMetadataDriver.java | 60 +++++++------- .../io/iceberg/TableMetadataDriverTest.java | 81 +++++++++++++++++-- 2 files changed, 106 insertions(+), 35 deletions(-) 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 index 6d9e00f1f12f..aa34d8ee562a 100644 --- 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 @@ -40,34 +40,34 @@ 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.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** * A driver transform that extracts table identifiers from incoming {@link Row}s, deduplicates them - * per window, samples up to a maximum number of tables, loads their declarative metadata from the - * Iceberg catalog, and emits {@link KV} pairs of table identifier strings to {@link - * SerializableTableSpec}. + * 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}. * *

Can also be materialized into a broadcasted {@link PCollectionView} via {@link - * #asView(IcebergCatalogConfig, DynamicDestinations)}. If the number of distinct tables in a window - * exceeds {@code maxTables}, up to {@code maxTables} tables are sampled into the broadcasted view, - * while remaining destinations can fall back to worker-local catalog loading. + * #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. */ @Internal @AutoValue public abstract class TableMetadataDriver extends PTransform, PCollection>> { - public static final int DEFAULT_MAX_TABLES = 100; - public abstract IcebergCatalogConfig getCatalogConfig(); public abstract DynamicDestinations getDynamicDestinations(); - public abstract int getMaxTables(); + public abstract @Nullable Integer getMaximumCacheSize(); public static Builder builder() { - return new AutoValue_TableMetadataDriver.Builder().setMaxTables(DEFAULT_MAX_TABLES); + return new AutoValue_TableMetadataDriver.Builder(); } public abstract Builder toBuilder(); @@ -78,44 +78,45 @@ public abstract static class Builder { public abstract Builder setDynamicDestinations(DynamicDestinations dynamicDestinations); - public abstract Builder setMaxTables(int maxTables); + public abstract Builder setMaximumCacheSize(@Nullable Integer maximumCacheSize); abstract TableMetadataDriver autoBuild(); public TableMetadataDriver build() { TableMetadataDriver driver = autoBuild(); - Preconditions.checkArgument( - driver.getMaxTables() > 0, - "maxTables must be greater than 0, got %s", - driver.getMaxTables()); + Integer maxCacheSize = driver.getMaximumCacheSize(); + if (maxCacheSize != null) { + Preconditions.checkArgument( + maxCacheSize > 0, "maximumCacheSize must be greater than 0, got %s", maxCacheSize); + } return driver; } } /** - * Helper that applies {@link TableMetadataDriver} and creates a {@link PCollectionView} of {@link - * Map} of table identifier strings to {@link SerializableTableSpec} using {@link - * #DEFAULT_MAX_TABLES}. + * 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, DEFAULT_MAX_TABLES); + return asView(catalogConfig, dynamicDestinations, null); } /** - * Helper that applies {@link TableMetadataDriver} with a custom {@code maxTables} limit and - * creates a {@link PCollectionView} of {@link Map} of table identifier strings to {@link + * 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 maxTables maximum distinct tables to poll and broadcast per window. + * @param maximumCacheSize optional maximum distinct tables to poll and broadcast per window (null + * for uncapped). */ public static PTransform, PCollectionView>> asView( IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, - int maxTables) { + @Nullable Integer maximumCacheSize) { return new PTransform, PCollectionView>>() { @Override public PCollectionView> expand(PCollection input) { @@ -125,7 +126,7 @@ public PCollectionView> expand(PCollection> expand(PCollection in PCollection distinctTableIds = tableIds.apply("DistinctTableIds", Distinct.create()); - PCollection sampledTableIds = - distinctTableIds.apply("SampleTableIds", Sample.any(getMaxTables())); + PCollection cachedTableIds; + Integer maxCacheSize = getMaximumCacheSize(); + if (maxCacheSize != null) { + cachedTableIds = distinctTableIds.apply("CapCacheSize", Sample.any(maxCacheSize)); + } else { + cachedTableIds = distinctTableIds; + } - return sampledTableIds + return cachedTableIds .apply("PollTableMetadata", ParDo.of(new CatalogPollingDoFn(getCatalogConfig()))) .setCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())); } 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 index 0fe886916cf5..d061824b3afd 100644 --- 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 @@ -216,7 +216,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { } @Test - public void testDeduplicationOfTablesAcrossRows() { + public void testWindowedDeduplication() { Catalog catalog = getCatalog(); TableIdentifier table1 = TableIdentifier.of("default", "t1"); TableIdentifier table2 = TableIdentifier.of("default", "t2"); @@ -269,6 +269,10 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { 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; }); @@ -276,7 +280,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { } @Test - public void testMaxTablesCapSampling() { + public void testMaximumCacheSizeCap() { Catalog catalog = getCatalog(); for (int i = 1; i <= 6; i++) { catalog.createTable(TableIdentifier.of("default", "cap_table_" + i), ICEBERG_SCHEMA); @@ -317,20 +321,81 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); - int maxTables = 3; + int maxCacheSize = 3; + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .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); + } + + DynamicDestinations dynamicDestinations = + 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"); + } + }; + + 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(dynamicDestinations) - .setMaxTables(maxTables) .build()); PAssert.that(specs) .satisfies( elements -> { List> list = ImmutableList.copyOf(elements); - assertEquals(maxTables, list.size()); + assertEquals(10, list.size()); return null; }); @@ -399,7 +464,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { } @Test - public void testInvalidMaxTablesThrowsException() { + public void testInvalidMaximumCacheSizeThrowsException() { TableIdentifier tableId = TableIdentifier.of("default", "dummy_table"); DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); @@ -409,7 +474,7 @@ public void testInvalidMaxTablesThrowsException() { TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) .setDynamicDestinations(dynamicDestinations) - .setMaxTables(0) + .setMaximumCacheSize(0) .build()); assertThrows( @@ -418,7 +483,7 @@ public void testInvalidMaxTablesThrowsException() { TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) .setDynamicDestinations(dynamicDestinations) - .setMaxTables(-5) + .setMaximumCacheSize(-5) .build()); } From c7b25dccc95baf4f373778f6140ee19581e72aab Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 1 Sep 2026 13:51:39 +0000 Subject: [PATCH 4/7] Handle NoSuchTableExceptions in CatalogPollingDoFn --- .../sdk/io/iceberg/TableMetadataDriver.java | 18 ++++-- .../io/iceberg/TableMetadataDriverTest.java | 58 +++++++++++++++++++ 2 files changed, 72 insertions(+), 4 deletions(-) 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 index aa34d8ee562a..0cbffe90ee3a 100644 --- 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 @@ -40,8 +40,11 @@ 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.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * A driver transform that extracts table identifiers from incoming {@link Row}s, deduplicates them @@ -179,6 +182,7 @@ public void processElement( } 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"); @@ -192,10 +196,16 @@ static class CatalogPollingDoFn extends DoFn> out) { 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)); + try { + 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.debug( + "Table '{}' does not exist in catalog. Skipping metadata emission for side-input view.", + tableIdString); + } } } } 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 index d061824b3afd..4cc5e726d020 100644 --- 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 @@ -402,6 +402,64 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { pipeline.run(); } + @Test + public void testNonExistentTableIsSkippedWithoutFailingBundle() { + Catalog catalog = getCatalog(); + TableIdentifier validTable = TableIdentifier.of("default", "existing_table"); + catalog.createTable(validTable, ICEBERG_SCHEMA); + + DynamicDestinations dynamicDestinations = + 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"); + } + }; + + 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(dynamicDestinations) + .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"); From 4bef0d86132b9ae4d63c443cdb7d34ac7f726e60 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 1 Sep 2026 19:32:43 +0000 Subject: [PATCH 5/7] unbounded global window support --- .../sdk/io/iceberg/TableMetadataDriver.java | 93 ++++++++++++- .../io/iceberg/TableMetadataDriverTest.java | 130 ++++++++++++++++++ 2 files changed, 219 insertions(+), 4 deletions(-) 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 index 0cbffe90ee3a..35788b5e43f5 100644 --- 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 @@ -17,6 +17,8 @@ */ 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.Map; import org.apache.beam.sdk.annotations.Internal; @@ -30,8 +32,13 @@ import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Sample; import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.transforms.display.DisplayData; +import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime; 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; @@ -42,6 +49,7 @@ 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; @@ -50,25 +58,36 @@ * 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}. + * 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}, an {@link AfterProcessingTime} + * trigger is automatically applied to fire deduplication and refresh table metadata at the + * configured {@code refreshInterval} (defaulting to {@link #DEFAULT_REFRESH_INTERVAL}). */ @Internal @AutoValue public abstract class TableMetadataDriver extends PTransform, PCollection>> { + public static final Duration DEFAULT_REFRESH_INTERVAL = Duration.standardMinutes(5); + public abstract IcebergCatalogConfig getCatalogConfig(); public abstract DynamicDestinations getDynamicDestinations(); public abstract @Nullable Integer getMaximumCacheSize(); + public abstract @Nullable Duration getRefreshInterval(); + public static Builder builder() { return new AutoValue_TableMetadataDriver.Builder(); } @@ -83,6 +102,8 @@ public abstract static class Builder { public abstract Builder setMaximumCacheSize(@Nullable Integer maximumCacheSize); + public abstract Builder setRefreshInterval(@Nullable Duration refreshInterval); + abstract TableMetadataDriver autoBuild(); public TableMetadataDriver build() { @@ -92,6 +113,13 @@ public TableMetadataDriver build() { 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); + } return driver; } } @@ -102,7 +130,7 @@ public TableMetadataDriver build() { */ public static PTransform, PCollectionView>> asView(IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations) { - return asView(catalogConfig, dynamicDestinations, null); + return asView(catalogConfig, dynamicDestinations, null, null); } /** @@ -120,6 +148,26 @@ public TableMetadataDriver build() { IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, @Nullable Integer maximumCacheSize) { + return asView(catalogConfig, dynamicDestinations, maximumCacheSize, 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 new PTransform, PCollectionView>>() { @Override public PCollectionView> expand(PCollection input) { @@ -130,6 +178,7 @@ public PCollectionView> expand(PCollection> expand(PCollection in .apply("ExtractTableIds", ParDo.of(new ExtractTableIdsDoFn(getDynamicDestinations()))) .setCoder(StringUtf8Coder.of()); - PCollection distinctTableIds = tableIds.apply("DistinctTableIds", Distinct.create()); + boolean isUnboundedGlobal = + input.isBounded() == PCollection.IsBounded.UNBOUNDED + && input.getWindowingStrategy().getWindowFn() instanceof GlobalWindows; + + PCollection triggeredTableIds; + if (isUnboundedGlobal) { + Duration customInterval = getRefreshInterval(); + Duration interval = + checkNotNull(customInterval != null ? customInterval : DEFAULT_REFRESH_INTERVAL); + triggeredTableIds = + tableIds.apply( + "ApplyStreamingTrigger", + Window.into(new GlobalWindows()) + .triggering( + Repeatedly.forever( + AfterProcessingTime.pastFirstElementInPane().plusDelayOf(interval))) + .accumulatingFiredPanes()); + } else { + triggeredTableIds = tableIds; + } + + PCollection distinctTableIds = + triggeredTableIds.apply("DistinctTableIds", Distinct.create()); PCollection cachedTableIds; Integer maxCacheSize = getMaximumCacheSize(); @@ -158,6 +229,17 @@ public PCollection> expand(PCollection in .setCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())); } + @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")); + } + static class ExtractTableIdsDoFn extends DoFn { private final DynamicDestinations dynamicDestinations; @@ -185,6 +267,8 @@ static class CatalogPollingDoFn extends DoFn element) { 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); + + DynamicDestinations dynamicDestinations = + 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"); + } + }; + + 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(dynamicDestinations) + .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 testMaximumCacheSizeCap() { Catalog catalog = getCatalog(); @@ -545,6 +619,30 @@ public void testInvalidMaximumCacheSizeThrowsException() { .build()); } + @Test + public void testInvalidRefreshIntervalThrowsException() { + TableIdentifier tableId = TableIdentifier.of("default", "dummy_table"); + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .setRefreshInterval(Duration.ZERO) + .build()); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .setRefreshInterval(Duration.standardSeconds(-5)) + .build()); + } + @Test public void testWindowPreservation() { Catalog catalog = getCatalog(); @@ -634,6 +732,38 @@ public void testEmptyInputProducesEmptyOutput() { pipeline.run(); } + @Test + public void testDisplayData() { + TableIdentifier tableId = TableIdentifier.of("default", "display_table"); + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + + TableMetadataDriver driver = + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .setMaximumCacheSize(42) + .setRefreshInterval(Duration.standardMinutes(10)) + .build(); + + DisplayData displayData = DisplayData.from(driver); + Map items = displayData.asMap(); + + assertNotNull(displayData); + boolean hasCacheSize = false; + boolean hasRefreshInterval = 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; + } + } + assertTrue(hasCacheSize); + assertTrue(hasRefreshInterval); + } + @Test public void testViewAsMapIntegration() { TableIdentifier tableId = TableIdentifier.of("default", "view_integration_table"); From cfeb3d5f1996232d73b5230cbffcfe0ab23d9d48 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 2 Sep 2026 14:40:00 +0000 Subject: [PATCH 6/7] Streamline test definitions --- .../io/iceberg/TableMetadataDriverTest.java | 291 ++++-------------- 1 file changed, 52 insertions(+), 239 deletions(-) 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 index 3c67adfe8978..8337b89165ec 100644 --- 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 @@ -82,6 +82,36 @@ public class TableMetadataDriverTest implements Serializable { 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(); @@ -102,10 +132,7 @@ private Catalog getCatalog() { @Test public void testSingleTableExtractionAndSpecOutput() { - TableIdentifier tableId = TableIdentifier.of("default", "single_table"); - Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA); - - DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + Table realTable = getCatalog().createTable(TABLE_ID, ICEBERG_SCHEMA); List rows = new ArrayList<>(); for (int i = 0; i < 5; i++) { @@ -123,10 +150,10 @@ public void testSingleTableExtractionAndSpecOutput() { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) .build()); - String expectedTableIdString = IcebergUtils.tableIdentifierToString(tableId); + String expectedTableIdString = IcebergUtils.tableIdentifierToString(TABLE_ID); PAssert.that(specs) .satisfies( @@ -159,31 +186,6 @@ public void testMultipleDynamicDestinationsExtraction() { catalog.createTable(tableB, ICEBERG_SCHEMA); catalog.createTable(tableC, ICEBERG_SCHEMA); - DynamicDestinations dynamicDestinations = - 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"); - } - }; - List rows = ImmutableList.of( Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.table_a").build(), @@ -198,7 +200,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) .build()); PAssert.that(specs) @@ -226,31 +228,6 @@ public void testWindowedDeduplication() { catalog.createTable(table1, ICEBERG_SCHEMA); catalog.createTable(table2, ICEBERG_SCHEMA); - DynamicDestinations dynamicDestinations = - 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"); - } - }; - List rows = new ArrayList<>(); for (int i = 0; i < 100; i++) { String dest = (i % 2 == 0) ? "default.t1" : "default.t2"; @@ -263,7 +240,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) .build()); PAssert.that(specs) @@ -290,31 +267,6 @@ public void testUnboundedGlobalWindowStreamingDeduplication() { catalog.createTable(table1, ICEBERG_SCHEMA); catalog.createTable(table2, ICEBERG_SCHEMA); - DynamicDestinations dynamicDestinations = - 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"); - } - }; - 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(); @@ -334,7 +286,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) .setRefreshInterval(Duration.standardSeconds(2)) .build()); @@ -360,31 +312,6 @@ public void testMaximumCacheSizeCap() { catalog.createTable(TableIdentifier.of("default", "cap_table_" + i), ICEBERG_SCHEMA); } - DynamicDestinations dynamicDestinations = - 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"); - } - }; - List rows = new ArrayList<>(); for (int i = 1; i <= 6; i++) { rows.add( @@ -400,7 +327,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) .setMaximumCacheSize(maxCacheSize) .build()); @@ -422,31 +349,6 @@ public void testUncappedByDefault() { catalog.createTable(TableIdentifier.of("default", "uncapped_table_" + i), ICEBERG_SCHEMA); } - DynamicDestinations dynamicDestinations = - 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"); - } - }; - List rows = new ArrayList<>(); for (int i = 1; i <= 10; i++) { rows.add( @@ -462,7 +364,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) .build()); PAssert.that(specs) @@ -482,31 +384,6 @@ public void testNonExistentTableIsSkippedWithoutFailingBundle() { TableIdentifier validTable = TableIdentifier.of("default", "existing_table"); catalog.createTable(validTable, ICEBERG_SCHEMA); - DynamicDestinations dynamicDestinations = - 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"); - } - }; - List rows = ImmutableList.of( Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.existing_table").build(), @@ -518,7 +395,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) .build()); // Only the existing table is emitted; the non-existent table is skipped without failing bundle @@ -539,31 +416,6 @@ public void testFiltersNullAndBlankTableIdentifiers() { TableIdentifier validTableId = TableIdentifier.of("default", "valid_dest_table"); getCatalog().createTable(validTableId, ICEBERG_SCHEMA); - DynamicDestinations dynamicDestinations = - 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"); - } - }; - List rows = ImmutableList.of( Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", null).build(), @@ -580,7 +432,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) .build()); PAssert.that(specs) @@ -597,15 +449,12 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { @Test public void testInvalidMaximumCacheSizeThrowsException() { - TableIdentifier tableId = TableIdentifier.of("default", "dummy_table"); - DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); - assertThrows( IllegalArgumentException.class, () -> TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) .setMaximumCacheSize(0) .build()); @@ -614,22 +463,19 @@ public void testInvalidMaximumCacheSizeThrowsException() { () -> TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) .setMaximumCacheSize(-5) .build()); } @Test public void testInvalidRefreshIntervalThrowsException() { - TableIdentifier tableId = TableIdentifier.of("default", "dummy_table"); - DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); - assertThrows( IllegalArgumentException.class, () -> TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) .setRefreshInterval(Duration.ZERO) .build()); @@ -638,7 +484,7 @@ public void testInvalidRefreshIntervalThrowsException() { () -> TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) .setRefreshInterval(Duration.standardSeconds(-5)) .build()); } @@ -652,31 +498,6 @@ public void testWindowPreservation() { catalog.createTable(tableW1, ICEBERG_SCHEMA); catalog.createTable(tableW2, ICEBERG_SCHEMA); - DynamicDestinations dynamicDestinations = - 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"); - } - }; - Instant t1 = new Instant(1000); Instant t2 = new Instant(70000); @@ -697,7 +518,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) .build()); PAssert.that(specs) @@ -713,10 +534,7 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { @Test public void testEmptyInputProducesEmptyOutput() { - TableIdentifier tableId = TableIdentifier.of("default", "empty_input_table"); - getCatalog().createTable(tableId, ICEBERG_SCHEMA); - - DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + getCatalog().createTable(TABLE_ID, ICEBERG_SCHEMA); PCollection input = pipeline.apply(Create.empty(RowCoder.of(BEAM_SCHEMA))); @@ -724,7 +542,7 @@ public void testEmptyInputProducesEmptyOutput() { input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) .build()); PAssert.that(specs).empty(); @@ -734,13 +552,10 @@ public void testEmptyInputProducesEmptyOutput() { @Test public void testDisplayData() { - TableIdentifier tableId = TableIdentifier.of("default", "display_table"); - DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); - TableMetadataDriver driver = TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) - .setDynamicDestinations(dynamicDestinations) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) .setMaximumCacheSize(42) .setRefreshInterval(Duration.standardMinutes(10)) .build(); @@ -766,11 +581,8 @@ public void testDisplayData() { @Test public void testViewAsMapIntegration() { - TableIdentifier tableId = TableIdentifier.of("default", "view_integration_table"); PartitionSpec partitionSpec = PartitionSpec.builderFor(ICEBERG_SCHEMA).identity("data").build(); - getCatalog().createTable(tableId, ICEBERG_SCHEMA, partitionSpec); - - DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, BEAM_SCHEMA); + getCatalog().createTable(TABLE_ID, ICEBERG_SCHEMA, partitionSpec); List rows = ImmutableList.of( @@ -781,9 +593,10 @@ public void testViewAsMapIntegration() { PCollectionView> metadataView = input.apply( - "CreateMetadataView", TableMetadataDriver.asView(catalogConfig, dynamicDestinations)); + "CreateMetadataView", + TableMetadataDriver.asView(catalogConfig, SINGLE_TABLE_DYNAMIC_DESTINATIONS)); - String expectedTableIdString = IcebergUtils.tableIdentifierToString(tableId); + String expectedTableIdString = IcebergUtils.tableIdentifierToString(TABLE_ID); PCollection writtenFiles = input.apply( From 4555fa888909355293785c3977d2b58f3966cf09 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 2 Sep 2026 15:47:51 +0000 Subject: [PATCH 7/7] add schema evolution test case, route through Deduplicate to re-emit panes --- .../sdk/io/iceberg/TableMetadataDriver.java | 41 +++++----- .../io/iceberg/TableMetadataDriverTest.java | 76 +++++++++++++++++++ 2 files changed, 98 insertions(+), 19 deletions(-) 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 index 35788b5e43f5..6148d80c9a22 100644 --- 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 @@ -26,6 +26,7 @@ import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; +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.PTransform; @@ -33,7 +34,7 @@ import org.apache.beam.sdk.transforms.Sample; import org.apache.beam.sdk.transforms.View; import org.apache.beam.sdk.transforms.display.DisplayData; -import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime; +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; @@ -69,9 +70,9 @@ * 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}, an {@link AfterProcessingTime} - * trigger is automatically applied to fire deduplication and refresh table metadata at the - * configured {@code refreshInterval} (defaulting to {@link #DEFAULT_REFRESH_INTERVAL}). + *

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 @@ -196,26 +197,18 @@ public PCollection> expand(PCollection in input.isBounded() == PCollection.IsBounded.UNBOUNDED && input.getWindowingStrategy().getWindowFn() instanceof GlobalWindows; - PCollection triggeredTableIds; + PCollection distinctTableIds; if (isUnboundedGlobal) { Duration customInterval = getRefreshInterval(); Duration interval = checkNotNull(customInterval != null ? customInterval : DEFAULT_REFRESH_INTERVAL); - triggeredTableIds = + distinctTableIds = tableIds.apply( - "ApplyStreamingTrigger", - Window.into(new GlobalWindows()) - .triggering( - Repeatedly.forever( - AfterProcessingTime.pastFirstElementInPane().plusDelayOf(interval))) - .accumulatingFiredPanes()); + "DeduplicateTableIds", Deduplicate.values().withDuration(interval)); } else { - triggeredTableIds = tableIds; + distinctTableIds = tableIds.apply("DistinctTableIds", Distinct.create()); } - PCollection distinctTableIds = - triggeredTableIds.apply("DistinctTableIds", Distinct.create()); - PCollection cachedTableIds; Integer maxCacheSize = getMaximumCacheSize(); if (maxCacheSize != null) { @@ -224,9 +217,19 @@ public PCollection> expand(PCollection in cachedTableIds = distinctTableIds; } - return cachedTableIds - .apply("PollTableMetadata", ParDo.of(new CatalogPollingDoFn(getCatalogConfig()))) - .setCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())); + PCollection> specs = + cachedTableIds + .apply("PollTableMetadata", ParDo.of(new CatalogPollingDoFn(getCatalogConfig()))) + .setCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())); + + if (isUnboundedGlobal) { + return specs.apply( + "ApplyStreamingViewTrigger", + Window.>into(new GlobalWindows()) + .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1))) + .accumulatingFiredPanes()); + } + return specs; } @Override 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 index 8337b89165ec..accae66523fb 100644 --- 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 @@ -56,6 +56,7 @@ 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; @@ -305,6 +306,81 @@ public void testUnboundedGlobalWindowStreamingDeduplication() { 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 testMaximumCacheSizeCap() { Catalog catalog = getCatalog();