From 83ddc7d178eb491f597d5a6ff248156aec5c6409 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 25 Aug 2026 15:02:05 +0000 Subject: [PATCH 01/12] 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 02/12] 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 03/12] 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 04/12] 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 05/12] 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 06/12] 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 07/12] 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(); From d9b260a044e04248429949a93f3f1cb3ae20f745 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 3 Sep 2026 14:05:07 +0000 Subject: [PATCH 08/12] Add reshuffle and polling buckets --- .../sdk/io/iceberg/TableMetadataDriver.java | 50 +++++++++++++- .../io/iceberg/TableMetadataDriverTest.java | 66 +++++++++++++++++++ 2 files changed, 113 insertions(+), 3 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 6148d80c9a22..c8b59b7a6170 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 @@ -31,6 +31,7 @@ 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.Reshuffle; import org.apache.beam.sdk.transforms.Sample; import org.apache.beam.sdk.transforms.View; import org.apache.beam.sdk.transforms.display.DisplayData; @@ -80,6 +81,7 @@ 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(); @@ -89,6 +91,8 @@ public abstract class TableMetadataDriver public abstract @Nullable Duration getRefreshInterval(); + public abstract @Nullable Integer getPollingBuckets(); + public static Builder builder() { return new AutoValue_TableMetadataDriver.Builder(); } @@ -105,6 +109,8 @@ public abstract static class Builder { public abstract Builder setRefreshInterval(@Nullable Duration refreshInterval); + public abstract Builder setPollingBuckets(@Nullable Integer pollingBuckets); + abstract TableMetadataDriver autoBuild(); public TableMetadataDriver build() { @@ -121,6 +127,11 @@ public TableMetadataDriver build() { "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; } } @@ -131,7 +142,7 @@ public TableMetadataDriver build() { */ public static PTransform, PCollectionView>> asView(IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations) { - return asView(catalogConfig, dynamicDestinations, null, null); + return asView(catalogConfig, dynamicDestinations, null, null, null); } /** @@ -149,7 +160,7 @@ public TableMetadataDriver build() { IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, @Nullable Integer maximumCacheSize) { - return asView(catalogConfig, dynamicDestinations, maximumCacheSize, null); + return asView(catalogConfig, dynamicDestinations, maximumCacheSize, null, null); } /** @@ -169,6 +180,28 @@ public TableMetadataDriver build() { 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) { @@ -180,6 +213,7 @@ public PCollectionView> expand(PCollection> expand(PCollection in 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 = - cachedTableIds + pollingTableIds .apply("PollTableMetadata", ParDo.of(new CatalogPollingDoFn(getCatalogConfig()))) .setCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())); @@ -241,6 +282,9 @@ public void populateDisplayData(DisplayData.Builder builder) { 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 { 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 accae66523fb..6333ecb0469c 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 @@ -565,6 +565,65 @@ public void testInvalidRefreshIntervalThrowsException() { .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(); @@ -634,6 +693,7 @@ public void testDisplayData() { .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) .setMaximumCacheSize(42) .setRefreshInterval(Duration.standardMinutes(10)) + .setPollingBuckets(3) .build(); DisplayData displayData = DisplayData.from(driver); @@ -642,6 +702,7 @@ public void testDisplayData() { 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()); @@ -650,9 +711,14 @@ public void testDisplayData() { 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 From f0515f1d9ed86bc90413a4f28d9f873e754b9498 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 3 Sep 2026 15:07:06 +0000 Subject: [PATCH 09/12] Side input tests + multiple tables, fix breakages --- .../sdk/io/iceberg/TableMetadataDriver.java | 121 ++++++++- .../io/iceberg/TableMetadataDriverTest.java | 239 ++++++++++++++++++ 2 files changed, 347 insertions(+), 13 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 c8b59b7a6170..5da44a6c51a1 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 @@ -20,20 +20,30 @@ 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; @@ -91,6 +101,10 @@ public abstract class TableMetadataDriver 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() { @@ -109,6 +123,15 @@ public abstract static class Builder { 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(); @@ -205,8 +228,10 @@ public TableMetadataDriver build() { return new PTransform, PCollectionView>>() { @Override public PCollectionView> expand(PCollection input) { - return input - .apply( + boolean isStreaming = input.isBounded() == PCollection.IsBounded.UNBOUNDED; + + PCollection> specs = + input.apply( "GenerateTableMetadata", TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -214,8 +239,25 @@ public PCollectionView> expand(PCollection>into(new GlobalWindows()) + .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1))) + .discardingFiredPanes()) + .apply( + "CreateMetadataSingletonView", + Combine.globally(new MapMergerFn()).asSingletonView()); + } + + return specs.apply("CreateTableMetadataView", View.asMap()); } }; } @@ -225,14 +267,13 @@ public PCollection> expand(PCollection in PCollection tableIds = input .apply("ExtractTableIds", ParDo.of(new ExtractTableIdsDoFn(getDynamicDestinations()))) - .setCoder(StringUtf8Coder.of()); + .setCoder(StringUtf8Coder.of()) + .apply("MetadataGlobalWindow", Window.into(new GlobalWindows())); - boolean isUnboundedGlobal = - input.isBounded() == PCollection.IsBounded.UNBOUNDED - && input.getWindowingStrategy().getWindowFn() instanceof GlobalWindows; + boolean isStreaming = input.isBounded() == PCollection.IsBounded.UNBOUNDED; PCollection distinctTableIds; - if (isUnboundedGlobal) { + if (isStreaming) { Duration customInterval = getRefreshInterval(); Duration interval = checkNotNull(customInterval != null ? customInterval : DEFAULT_REFRESH_INTERVAL); @@ -246,6 +287,10 @@ public PCollection> expand(PCollection in 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; @@ -263,12 +308,12 @@ public PCollection> expand(PCollection in .apply("PollTableMetadata", ParDo.of(new CatalogPollingDoFn(getCatalogConfig()))) .setCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())); - if (isUnboundedGlobal) { + if (isStreaming) { return specs.apply( - "ApplyStreamingViewTrigger", + "ApplyStreamingTrigger", Window.>into(new GlobalWindows()) .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1))) - .accumulatingFiredPanes()); + .discardingFiredPanes()); } return specs; } @@ -326,8 +371,8 @@ static class CatalogPollingDoFn extends DoFn> out) { - TableIdentifier tableId = IcebergUtils.parseTableIdentifier(tableIdString); try { + TableIdentifier tableId = IcebergUtils.parseTableIdentifier(tableIdString); Table table = catalogConfig.catalog().loadTable(tableId); SerializableTableSpec spec = SerializableTableSpec.fromTable(tableIdString, table); TABLES_POLLED_COUNTER.inc(); @@ -337,7 +382,57 @@ public void processElement( "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 index 6333ecb0469c..37c80631ff58 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 @@ -381,6 +381,245 @@ public void processElement( 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(); From adea35b8fc88132ea3f14fa89a800a82b218c6ec Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Fri, 4 Sep 2026 21:27:47 +0000 Subject: [PATCH 10/12] Explicit exception handling, time-based metadata merging --- .../sdk/io/iceberg/SerializableTableSpec.java | 7 ++ .../sdk/io/iceberg/TableMetadataDriver.java | 58 +++++++++++++-- .../io/iceberg/SerializableTableSpecTest.java | 6 ++ .../io/iceberg/TableMetadataDriverTest.java | 74 +++++++++++++++++++ 4 files changed, 137 insertions(+), 8 deletions(-) diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableTableSpec.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableTableSpec.java index c6ee4a976993..e89a341f55c3 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableTableSpec.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableTableSpec.java @@ -97,6 +97,9 @@ public abstract class SerializableTableSpec implements Serializable { @SchemaFieldNumber("11") public abstract List getEncryptedKeyJsons(); + @SchemaFieldNumber("12") + public abstract long getLastUpdatedMillis(); + private transient volatile @MonotonicNonNull Map cachedSchemas; private transient volatile @MonotonicNonNull Map cachedPartitionSpecs; private transient volatile @MonotonicNonNull Map cachedSortOrders; @@ -285,6 +288,8 @@ public abstract static class Builder { public abstract Builder setEncryptedKeyJsons(List encryptedKeyJsons); + public abstract Builder setLastUpdatedMillis(long lastUpdatedMillis); + @SchemaIgnore public Builder setFileIO(FileIO fileIO) { return setFileIoJson(FileIOParser.toJson(fileIO)); @@ -324,6 +329,7 @@ public static SerializableTableSpec fromTable(String tableIdentifierString, Tabl } TableMetadata metadata = ((HasTableOperations) table).operations().current(); + long lastUpdatedMillis = metadata != null ? metadata.lastUpdatedMillis() : 0L; List encryptedKeyJsons = Collections.emptyList(); if (metadata != null && metadata.encryptionKeys() != null) { encryptedKeyJsons = @@ -360,6 +366,7 @@ public static SerializableTableSpec fromTable(String tableIdentifierString, Tabl .setProperties(table.properties()) .setFileIoJson(FileIOParser.toJson(table.io())) .setEncryptedKeyJsons(encryptedKeyJsons) + .setLastUpdatedMillis(lastUpdatedMillis) .build(); } 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 5da44a6c51a1..062c6b0627b9 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 @@ -31,6 +31,7 @@ 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.ReadableState; import org.apache.beam.sdk.state.StateSpec; import org.apache.beam.sdk.state.StateSpecs; import org.apache.beam.sdk.transforms.Combine; @@ -371,24 +372,41 @@ static class CatalogPollingDoFn extends DoFn> out) { + TableIdentifier tableId; 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)); + tableId = IcebergUtils.parseTableIdentifier(tableIdString); + } catch (IllegalArgumentException e) { + LOG.warn( + "Failed to parse table identifier '{}'. Skipping metadata emission for side-input view.", + tableIdString, + e); + TABLES_SKIPPED_MISSING_COUNTER.inc(); + return; + } + + Table table; + try { + table = catalogConfig.catalog().loadTable(tableId); } catch (NoSuchTableException e) { LOG.info( "Table '{}' does not exist in catalog. Skipping metadata emission for side-input view.", tableIdString); TABLES_SKIPPED_MISSING_COUNTER.inc(); + return; + } + SerializableTableSpec spec; + try { + spec = SerializableTableSpec.fromTable(tableIdString, table); } catch (IllegalArgumentException e) { LOG.warn( - "Failed to parse table identifier '{}'. Skipping metadata emission for side-input view.", + "Failed to create SerializableTableSpec for table '{}'. Skipping metadata emission for side-input view.", tableIdString, e); TABLES_SKIPPED_MISSING_COUNTER.inc(); + return; } + TABLES_POLLED_COUNTER.inc(); + out.output(KV.of(tableIdString, spec)); } } @@ -405,7 +423,17 @@ public void processElement( @StateId("tableCache") MapState cacheState, OutputReceiver> out) { KV kv = element.getValue(); - cacheState.put(kv.getKey(), kv.getValue()); + String tableId = kv.getKey(); + SerializableTableSpec newSpec = kv.getValue(); + + ReadableState existingState = cacheState.get(tableId); + SerializableTableSpec existingSpec = existingState != null ? existingState.read() : null; + if (existingSpec == null + || newSpec.getLastUpdatedMillis() > existingSpec.getLastUpdatedMillis() + || (newSpec.getLastUpdatedMillis() == existingSpec.getLastUpdatedMillis() + && newSpec.getSchemaId() >= existingSpec.getSchemaId())) { + cacheState.put(tableId, newSpec); + } Map mapSnapshot = new HashMap<>(); for (Map.Entry entry : cacheState.entries().read()) { @@ -426,7 +454,21 @@ public Map apply( return left; } Map merged = new HashMap<>(left); - merged.putAll(right); + for (Map.Entry entry : right.entrySet()) { + String tableId = entry.getKey(); + SerializableTableSpec rightSpec = entry.getValue(); + SerializableTableSpec leftSpec = merged.get(tableId); + if (leftSpec == null) { + merged.put(tableId, rightSpec); + } else if (rightSpec.getLastUpdatedMillis() > leftSpec.getLastUpdatedMillis()) { + merged.put(tableId, rightSpec); + } else if (rightSpec.getLastUpdatedMillis() == leftSpec.getLastUpdatedMillis()) { + // Deterministic tie-breaker for strict commutativity + if (rightSpec.getSchemaId() > leftSpec.getSchemaId()) { + merged.put(tableId, rightSpec); + } + } + } return Collections.unmodifiableMap(merged); } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableTableSpecTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableTableSpecTest.java index 87a843db7a2f..3f7bd575a9eb 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableTableSpecTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableTableSpecTest.java @@ -45,6 +45,7 @@ import org.apache.hadoop.conf.Configuration; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.CatalogUtil; +import org.apache.iceberg.HasTableOperations; import org.apache.iceberg.NullOrder; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; @@ -132,6 +133,10 @@ public void testFromTableAndGettersUnpartitioned() { assertNotNull(spec.getEncryptedKeyJsons()); assertNotNull(spec.getEncryptedKeys()); assertTrue(spec.getEncryptedKeys().isEmpty()); + assertEquals( + ((HasTableOperations) table).operations().current().lastUpdatedMillis(), + spec.getLastUpdatedMillis()); + assertTrue(spec.getLastUpdatedMillis() > 0); } @Test @@ -280,6 +285,7 @@ public void testBeamSchemaCoderRoundtrip() throws Exception { assertEquals(original.getProperties(), decoded.getProperties()); assertEquals(original.getFileIoJson(), decoded.getFileIoJson()); assertEquals(original.getEncryptedKeyJsons(), decoded.getEncryptedKeyJsons()); + assertEquals(original.getLastUpdatedMillis(), decoded.getLastUpdatedMillis()); assertNotNull(decoded.getFileIO()); } 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 37c80631ff58..811a9cafc4c9 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 @@ -1025,4 +1025,78 @@ public void processElement( pipeline.run(); } + + @Test + public void testMapMergerFnCommutativeAndTimestampAware() { + TableIdentifier tableIdA = TableIdentifier.of("default", "merge_table_a"); + Table realTableA = getCatalog().createTable(tableIdA, ICEBERG_SCHEMA); + SerializableTableSpec specAOld = + SerializableTableSpec.fromTable(tableIdA, realTableA) + .toBuilder() + .setLastUpdatedMillis(1000L) + .build(); + SerializableTableSpec specANew = + SerializableTableSpec.fromTable(tableIdA, realTableA) + .toBuilder() + .setLastUpdatedMillis(2000L) + .build(); + + TableIdentifier tableIdB = TableIdentifier.of("default", "merge_table_b"); + Table realTableB = getCatalog().createTable(tableIdB, ICEBERG_SCHEMA); + SerializableTableSpec specB = + SerializableTableSpec.fromTable(tableIdB, realTableB) + .toBuilder() + .setLastUpdatedMillis(1500L) + .build(); + + TableMetadataDriver.MapMergerFn fn = new TableMetadataDriver.MapMergerFn(); + + Map map1 = ImmutableMap.of("tableA", specAOld, "tableB", specB); + Map map2 = ImmutableMap.of("tableA", specANew); + + // Left has old, right has new: right wins for tableA + Map merged1 = fn.apply(map1, map2); + assertEquals(2, merged1.size()); + assertEquals(2000L, merged1.get("tableA").getLastUpdatedMillis()); + assertEquals(1500L, merged1.get("tableB").getLastUpdatedMillis()); + + // Commutativity: left has new, right has old: left wins for tableA + Map merged2 = fn.apply(map2, map1); + assertEquals(2, merged2.size()); + assertEquals(2000L, merged2.get("tableA").getLastUpdatedMillis()); + assertEquals(1500L, merged2.get("tableB").getLastUpdatedMillis()); + + // Identical results in both merge directions + assertEquals(merged1, merged2); + } + + @Test + public void testMapMergerFnTieBreaksBySchemaIdCommutatively() { + TableIdentifier tableId = TableIdentifier.of("default", "tie_break_table"); + Table realTable = getCatalog().createTable(tableId, ICEBERG_SCHEMA); + SerializableTableSpec specSchema0 = + SerializableTableSpec.fromTable(tableId, realTable) + .toBuilder() + .setLastUpdatedMillis(1000L) + .setSchemaId(0) + .build(); + SerializableTableSpec specSchema1 = + SerializableTableSpec.fromTable(tableId, realTable) + .toBuilder() + .setLastUpdatedMillis(1000L) + .setSchemaId(1) + .build(); + + TableMetadataDriver.MapMergerFn fn = new TableMetadataDriver.MapMergerFn(); + + Map mapA = ImmutableMap.of("table", specSchema0); + Map mapB = ImmutableMap.of("table", specSchema1); + + Map mergedAB = fn.apply(mapA, mapB); + Map mergedBA = fn.apply(mapB, mapA); + + assertEquals(1, mergedAB.get("table").getSchemaId()); + assertEquals(1, mergedBA.get("table").getSchemaId()); + assertEquals(mergedAB, mergedBA); + } } From 7b304ef22e1c5bdbdc039627aeb5d27d22d9a7c6 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Fri, 4 Sep 2026 23:01:42 +0000 Subject: [PATCH 11/12] lastSeen impl --- .../sdk/io/iceberg/TableMetadataDriver.java | 93 +++++++++++++++- .../io/iceberg/TableMetadataDriverTest.java | 101 ++++++++++++++++++ 2 files changed, 191 insertions(+), 3 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 062c6b0627b9..5d990006a5c0 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 @@ -20,13 +20,17 @@ 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.io.Serializable; +import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.List; 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.VarLongCoder; import org.apache.beam.sdk.coders.VoidCoder; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; @@ -57,6 +61,7 @@ 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.annotations.VisibleForTesting; 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; @@ -94,6 +99,11 @@ public abstract class TableMetadataDriver public static final Duration DEFAULT_REFRESH_INTERVAL = Duration.standardMinutes(5); public static final int DEFAULT_POLLING_BUCKETS = 1; + @FunctionalInterface + public interface Clock extends Serializable { + long currentTimeMillis(); + } + public abstract IcebergCatalogConfig getCatalogConfig(); public abstract DynamicDestinations getDynamicDestinations(); @@ -226,11 +236,30 @@ public TableMetadataDriver build() { @Nullable Integer maximumCacheSize, @Nullable Duration refreshInterval, @Nullable Integer pollingBuckets) { + return asView( + catalogConfig, + dynamicDestinations, + maximumCacheSize, + refreshInterval, + pollingBuckets, + null); + } + + @VisibleForTesting + static PTransform, PCollectionView>> asView( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Integer maximumCacheSize, + @Nullable Duration refreshInterval, + @Nullable Integer pollingBuckets, + @Nullable Clock clock) { return new PTransform, PCollectionView>>() { @Override public PCollectionView> expand(PCollection input) { boolean isStreaming = input.isBounded() == PCollection.IsBounded.UNBOUNDED; + Duration interval = refreshInterval != null ? refreshInterval : DEFAULT_REFRESH_INTERVAL; + PCollection> specs = input.apply( "GenerateTableMetadata", @@ -238,15 +267,19 @@ public PCollectionView> expand(PCollection>, Map> { + private static final Logger LOG = LoggerFactory.getLogger(AccumulateTableMetadataMapDoFn.class); + private static final Counter TABLES_EVICTED_COUNTER = + Metrics.counter(TableMetadataDriver.class, "tablesEvictedUnused"); + @StateId("tableCache") private final StateSpec> cacheStateSpec = StateSpecs.map(StringUtf8Coder.of(), SerializableTableSpec.getCoder()); + @StateId("lastSeen") + private final StateSpec> lastSeenStateSpec = + StateSpecs.map(StringUtf8Coder.of(), VarLongCoder.of()); + + private final Duration refreshInterval; + private final Clock clock; + + AccumulateTableMetadataMapDoFn() { + this(DEFAULT_REFRESH_INTERVAL, System::currentTimeMillis); + } + + AccumulateTableMetadataMapDoFn(Duration refreshInterval) { + this(refreshInterval, System::currentTimeMillis); + } + + AccumulateTableMetadataMapDoFn(Duration refreshInterval, Clock clock) { + this.refreshInterval = refreshInterval != null ? refreshInterval : DEFAULT_REFRESH_INTERVAL; + this.clock = clock != null ? clock : System::currentTimeMillis; + } + @ProcessElement public void processElement( @Element KV> element, @StateId("tableCache") MapState cacheState, + @StateId("lastSeen") MapState lastSeenState, OutputReceiver> out) { + long now = clock.currentTimeMillis(); KV kv = element.getValue(); String tableId = kv.getKey(); SerializableTableSpec newSpec = kv.getValue(); @@ -434,11 +493,39 @@ public void processElement( && newSpec.getSchemaId() >= existingSpec.getSchemaId())) { cacheState.put(tableId, newSpec); } + lastSeenState.put(tableId, now); + Map lastSeenMap = new HashMap<>(); + for (Map.Entry entry : lastSeenState.entries().read()) { + lastSeenMap.put(entry.getKey(), entry.getValue()); + } + lastSeenMap.put(tableId, now); + + long expirationCutoff = now - refreshInterval.getMillis(); + List expiredTables = new ArrayList<>(); Map mapSnapshot = new HashMap<>(); + for (Map.Entry entry : cacheState.entries().read()) { - mapSnapshot.put(entry.getKey(), entry.getValue()); + String id = entry.getKey(); + Long lastSeen = lastSeenMap.get(id); + if (lastSeen == null) { + lastSeen = now; + lastSeenState.put(id, now); + } + if (lastSeen < expirationCutoff) { + expiredTables.add(id); + } else { + mapSnapshot.put(id, entry.getValue()); + } } + + for (String expired : expiredTables) { + cacheState.remove(expired); + lastSeenState.remove(expired); + TABLES_EVICTED_COUNTER.inc(); + LOG.info("Evicted unused table '{}' from side-input metadata cache.", expired); + } + out.output(Collections.unmodifiableMap(mapSnapshot)); } } 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 811a9cafc4c9..32c0a170c287 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 @@ -26,6 +26,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; import org.apache.beam.sdk.coders.RowCoder; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.testing.PAssert; @@ -1099,4 +1100,104 @@ public void testMapMergerFnTieBreaksBySchemaIdCommutatively() { assertEquals(1, mergedBA.get("table").getSchemaId()); assertEquals(mergedAB, mergedBA); } + + static class ControllableTestClock implements TableMetadataDriver.Clock { + private static final AtomicLong CURRENT_TIME = new AtomicLong(0L); + + public static void setTime(long millis) { + CURRENT_TIME.set(millis); + } + + @Override + public long currentTimeMillis() { + return CURRENT_TIME.get(); + } + } + + @Test + public void testUnusedTablesEvictedFromStreamingCache() { + TableIdentifier tableIdA = TableIdentifier.of("default", "evict_table_a"); + TableIdentifier tableIdB = TableIdentifier.of("default", "evict_table_b"); + getCatalog().createTable(tableIdA, ICEBERG_SCHEMA); + getCatalog().createTable(tableIdB, ICEBERG_SCHEMA); + + String tableAStr = IcebergUtils.tableIdentifierToString(tableIdA); + String tableBStr = IcebergUtils.tableIdentifierToString(tableIdB); + + Duration refreshInterval = Duration.standardSeconds(5); + ControllableTestClock.setTime(1000L); + ControllableTestClock testClock = new ControllableTestClock(); + + 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 rowTriggerEvictA = + Row.withSchema(BEAM_SCHEMA).addValues(3L, "trigger_evict_a", tableAStr).build(); + Row rowA2 = Row.withSchema(BEAM_SCHEMA).addValues(4L, "a2", tableAStr).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(6)) + .addElements(rowTriggerEvictA) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(rowA2) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceWatermarkToInfinity(); + + PCollection input = + pipeline + .apply("StreamInput", stream) + .apply( + "AdvanceClockOnTriggerRow", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(@Element Row row, OutputReceiver out) { + if ("trigger_evict_a".equals(row.getString("data"))) { + ControllableTestClock.setTime(7000L); + } + out.output(row); + } + })) + .setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", + TableMetadataDriver.asView( + catalogConfig, DYNAMIC_DESTINATIONS, null, refreshInterval, null, testClock)); + + 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_evict_a".equals(data)) { + return; + } + Map viewMap = c.sideInput(metadataView); + boolean hasA = viewMap.containsKey(tableAStr); + boolean hasB = viewMap.containsKey(tableBStr); + out.output(data + ":hasA=" + hasA + ",hasB=" + hasB); + } + }) + .withSideInputs(metadataView)); + + PAssert.that(consumerObserved) + .containsInAnyOrder( + "a1:hasA=true,hasB=true", "b1:hasA=true,hasB=true", "a2:hasA=true,hasB=false"); + + pipeline.run(); + } } From 7a5d1741f7d3ff5d37d0949a309709eb21cd100d Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Sat, 5 Sep 2026 00:01:22 +0000 Subject: [PATCH 12/12] Extra cleanup --- .../sdk/io/iceberg/TableMetadataDriver.java | 167 +++++++--- .../io/iceberg/TableMetadataDriverTest.java | 306 ++++++++++++++++-- 2 files changed, 391 insertions(+), 82 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 5d990006a5c0..c72c5e70339d 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 @@ -29,6 +29,7 @@ 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.NullableCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.coders.VarLongCoder; import org.apache.beam.sdk.coders.VoidCoder; @@ -43,10 +44,12 @@ 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.Filter; 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.SerializableFunction; import org.apache.beam.sdk.transforms.View; import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.transforms.display.DisplayData; @@ -74,27 +77,31 @@ /** * 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. + * per window, optionally bounds the cache size up to {@code maximumCacheSize} (batch pipelines + * only), loads their declarative metadata from the Iceberg catalog, and emits {@link KV} pairs of + * table identifier strings to {@link SerializableTableSpec} (or {@code null} if the table does not + * exist or fails to load). 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. + * destinations fall back to worker-local catalog loading. Note that {@code maximumCacheSize} is + * currently supported for bounded batch pipelines only. * *

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. + * Missing table signals ({@code null} specs) trigger side-input view materialization without + * caching the missing tables, ensuring downstream consumers are never blocked waiting for the side + * input. */ @Internal @AutoValue public abstract class TableMetadataDriver - extends PTransform, PCollection>> { + extends PTransform, PCollection>> { public static final Duration DEFAULT_REFRESH_INTERVAL = Duration.standardMinutes(5); public static final int DEFAULT_POLLING_BUCKETS = 1; @@ -253,6 +260,25 @@ static PTransform, PCollectionView, PCollectionView>> asView( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Integer maximumCacheSize, + @Nullable Duration refreshInterval, + @Nullable Integer pollingBuckets, + @Nullable Duration cacheTtl, + @Nullable Clock clock) { return new PTransform, PCollectionView>>() { @Override public PCollectionView> expand(PCollection input) { @@ -260,7 +286,7 @@ public PCollectionView> expand(PCollection> specs = + PCollection> specs = input.apply( "GenerateTableMetadata", TableMetadataDriver.builder() @@ -274,8 +300,8 @@ public PCollectionView> expand(PCollection> expand(PCollection, Boolean>) + kv -> kv.getValue() != null)) + .apply("CreateTableMetadataView", View.asMap()); } }; } @Override - public PCollection> expand(PCollection input) { + public PCollection> expand(PCollection input) { PCollection tableIds = input .apply("ExtractTableIds", ParDo.of(new ExtractTableIdsDoFn(getDynamicDestinations()))) @@ -337,15 +369,17 @@ public PCollection> expand(PCollection in "ReshufflePollingBuckets", Reshuffle.viaRandomKey().withNumBuckets(pollingBuckets)); - PCollection> specs = + PCollection> specs = pollingTableIds .apply("PollTableMetadata", ParDo.of(new CatalogPollingDoFn(getCatalogConfig()))) - .setCoder(KvCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())); + .setCoder( + KvCoder.of( + StringUtf8Coder.of(), NullableCoder.of(SerializableTableSpec.getCoder()))); if (isStreaming) { return specs.apply( "ApplyStreamingTrigger", - Window.>into(new GlobalWindows()) + Window.>into(new GlobalWindows()) .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1))) .discardingFiredPanes()); } @@ -389,12 +423,17 @@ public void processElement( } } - static class CatalogPollingDoFn extends DoFn> { + 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 static final Counter TABLES_PARSE_FAILED_COUNTER = + Metrics.counter(TableMetadataDriver.class, "tablesParseFailed"); + private static final Counter TABLES_SPEC_CREATION_FAILED_COUNTER = + Metrics.counter(TableMetadataDriver.class, "tablesSpecCreationFailed"); private final IcebergCatalogConfig catalogConfig; @@ -404,16 +443,18 @@ static class CatalogPollingDoFn extends DoFn> out) { + @Element String tableIdString, + OutputReceiver> out) { TableIdentifier tableId; try { tableId = IcebergUtils.parseTableIdentifier(tableIdString); } catch (IllegalArgumentException e) { LOG.warn( - "Failed to parse table identifier '{}'. Skipping metadata emission for side-input view.", + "Failed to parse table identifier '{}'. Emitting empty metadata signal for side-input view.", tableIdString, e); - TABLES_SKIPPED_MISSING_COUNTER.inc(); + TABLES_PARSE_FAILED_COUNTER.inc(); + out.output(KV.of(tableIdString, null)); return; } @@ -422,9 +463,10 @@ public void processElement( table = catalogConfig.catalog().loadTable(tableId); } catch (NoSuchTableException e) { LOG.info( - "Table '{}' does not exist in catalog. Skipping metadata emission for side-input view.", + "Table '{}' does not exist in catalog. Emitting empty metadata signal for side-input view.", tableIdString); TABLES_SKIPPED_MISSING_COUNTER.inc(); + out.output(KV.of(tableIdString, null)); return; } SerializableTableSpec spec; @@ -432,10 +474,11 @@ public void processElement( spec = SerializableTableSpec.fromTable(tableIdString, table); } catch (IllegalArgumentException e) { LOG.warn( - "Failed to create SerializableTableSpec for table '{}'. Skipping metadata emission for side-input view.", + "Failed to create SerializableTableSpec for table '{}'. Emitting empty metadata signal for side-input view.", tableIdString, e); - TABLES_SKIPPED_MISSING_COUNTER.inc(); + TABLES_SPEC_CREATION_FAILED_COUNTER.inc(); + out.output(KV.of(tableIdString, null)); return; } TABLES_POLLED_COUNTER.inc(); @@ -445,10 +488,12 @@ public void processElement( static class AccumulateTableMetadataMapDoFn extends DoFn< - KV>, Map> { + KV>, + Map> { private static final Logger LOG = LoggerFactory.getLogger(AccumulateTableMetadataMapDoFn.class); private static final Counter TABLES_EVICTED_COUNTER = Metrics.counter(TableMetadataDriver.class, "tablesEvictedUnused"); + static final int DEFAULT_TTL_MULTIPLIER = 3; @StateId("tableCache") private final StateSpec> cacheStateSpec = @@ -459,49 +504,66 @@ static class AccumulateTableMetadataMapDoFn StateSpecs.map(StringUtf8Coder.of(), VarLongCoder.of()); private final Duration refreshInterval; + private final Duration cacheTtl; private final Clock clock; AccumulateTableMetadataMapDoFn() { - this(DEFAULT_REFRESH_INTERVAL, System::currentTimeMillis); + this(DEFAULT_REFRESH_INTERVAL, null, System::currentTimeMillis); } AccumulateTableMetadataMapDoFn(Duration refreshInterval) { - this(refreshInterval, System::currentTimeMillis); + this(refreshInterval, null, System::currentTimeMillis); } AccumulateTableMetadataMapDoFn(Duration refreshInterval, Clock clock) { + this(refreshInterval, null, clock); + } + + AccumulateTableMetadataMapDoFn(Duration refreshInterval, @Nullable Duration cacheTtl) { + this(refreshInterval, cacheTtl, System::currentTimeMillis); + } + + AccumulateTableMetadataMapDoFn( + Duration refreshInterval, @Nullable Duration cacheTtl, Clock clock) { this.refreshInterval = refreshInterval != null ? refreshInterval : DEFAULT_REFRESH_INTERVAL; + this.cacheTtl = + cacheTtl != null ? cacheTtl : this.refreshInterval.multipliedBy(DEFAULT_TTL_MULTIPLIER); this.clock = clock != null ? clock : System::currentTimeMillis; } @ProcessElement public void processElement( - @Element KV> element, + @Element KV> element, @StateId("tableCache") MapState cacheState, @StateId("lastSeen") MapState lastSeenState, OutputReceiver> out) { long now = clock.currentTimeMillis(); - KV kv = element.getValue(); + KV kv = element.getValue(); String tableId = kv.getKey(); - SerializableTableSpec newSpec = kv.getValue(); - - ReadableState existingState = cacheState.get(tableId); - SerializableTableSpec existingSpec = existingState != null ? existingState.read() : null; - if (existingSpec == null - || newSpec.getLastUpdatedMillis() > existingSpec.getLastUpdatedMillis() - || (newSpec.getLastUpdatedMillis() == existingSpec.getLastUpdatedMillis() - && newSpec.getSchemaId() >= existingSpec.getSchemaId())) { - cacheState.put(tableId, newSpec); + @Nullable SerializableTableSpec newSpec = kv.getValue(); + + if (newSpec != null) { + ReadableState existingState = cacheState.get(tableId); + SerializableTableSpec existingSpec = existingState != null ? existingState.read() : null; + if (existingSpec == null || isNewer(newSpec, existingSpec)) { + cacheState.put(tableId, newSpec); + } + lastSeenState.put(tableId, now); + } else { + // Explicit missing table signal: immediately invalidate any cached entry + cacheState.remove(tableId); + lastSeenState.remove(tableId); } - lastSeenState.put(tableId, now); Map lastSeenMap = new HashMap<>(); for (Map.Entry entry : lastSeenState.entries().read()) { lastSeenMap.put(entry.getKey(), entry.getValue()); } - lastSeenMap.put(tableId, now); + if (newSpec != null) { + lastSeenMap.put(tableId, now); + } - long expirationCutoff = now - refreshInterval.getMillis(); + long expirationCutoff = now - cacheTtl.getMillis(); List expiredTables = new ArrayList<>(); Map mapSnapshot = new HashMap<>(); @@ -530,6 +592,22 @@ public void processElement( } } + static boolean isNewer(SerializableTableSpec candidate, SerializableTableSpec current) { + if (candidate.getLastUpdatedMillis() != current.getLastUpdatedMillis()) { + return candidate.getLastUpdatedMillis() > current.getLastUpdatedMillis(); + } + if (candidate.getSchemaId() != current.getSchemaId()) { + return candidate.getSchemaId() > current.getSchemaId(); + } + if (candidate.getSpecId() != current.getSpecId()) { + return candidate.getSpecId() > current.getSpecId(); + } + if (candidate.getOrderId() != current.getOrderId()) { + return candidate.getOrderId() > current.getOrderId(); + } + return candidate.getLocation().compareTo(current.getLocation()) > 0; + } + static class MapMergerFn extends Combine.BinaryCombineFn> { @Override public Map apply( @@ -545,15 +623,8 @@ public Map apply( String tableId = entry.getKey(); SerializableTableSpec rightSpec = entry.getValue(); SerializableTableSpec leftSpec = merged.get(tableId); - if (leftSpec == null) { - merged.put(tableId, rightSpec); - } else if (rightSpec.getLastUpdatedMillis() > leftSpec.getLastUpdatedMillis()) { + if (leftSpec == null || isNewer(rightSpec, leftSpec)) { merged.put(tableId, rightSpec); - } else if (rightSpec.getLastUpdatedMillis() == leftSpec.getLastUpdatedMillis()) { - // Deterministic tie-breaker for strict commutativity - if (rightSpec.getSchemaId() > leftSpec.getSchemaId()) { - merged.put(tableId, rightSpec); - } } } return Collections.unmodifiableMap(merged); 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 32c0a170c287..7e02e2b473e6 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 @@ -17,13 +17,16 @@ */ 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 static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import java.io.Serializable; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicLong; @@ -58,6 +61,7 @@ import org.apache.iceberg.data.GenericRecord; import org.apache.iceberg.data.Record; import org.apache.iceberg.types.Types; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.joda.time.Instant; import org.junit.Before; @@ -148,7 +152,7 @@ public void testSingleTableExtractionAndSpecOutput() { PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -160,9 +164,10 @@ public void testSingleTableExtractionAndSpecOutput() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); + List> list = + ImmutableList.copyOf(elements); assertEquals(1, list.size()); - KV kv = list.get(0); + KV kv = list.get(0); assertEquals(expectedTableIdString, kv.getKey()); SerializableTableSpec spec = kv.getValue(); assertNotNull(spec); @@ -198,7 +203,7 @@ public void testMultipleDynamicDestinationsExtraction() { PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -208,7 +213,8 @@ public void testMultipleDynamicDestinationsExtraction() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); + List> list = + ImmutableList.copyOf(elements); assertEquals(3, list.size()); Map map = list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); @@ -238,7 +244,7 @@ public void testWindowedDeduplication() { PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -248,7 +254,8 @@ public void testWindowedDeduplication() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); + List> list = + ImmutableList.copyOf(elements); assertEquals(2, list.size()); Map map = list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); @@ -284,7 +291,7 @@ public void testUnboundedGlobalWindowStreamingDeduplication() { PCollection input = pipeline.apply("StreamInput", stream); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -295,7 +302,8 @@ public void testUnboundedGlobalWindowStreamingDeduplication() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); + List> list = + ImmutableList.copyOf(elements); assertEquals(2, list.size()); Map map = list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); @@ -354,7 +362,7 @@ public void processElement(@Element Row row, OutputReceiver out) { })) .setCoder(RowCoder.of(BEAM_SCHEMA)); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -367,12 +375,13 @@ public void processElement(@Element Row row, OutputReceiver out) { specs.apply( "ConsumerTransform", ParDo.of( - new DoFn, String>() { + new DoFn, String>() { @ProcessElement public void processElement( - @Element KV element, + @Element KV element, OutputReceiver out) { - boolean hasNewCol = element.getValue().getSchema().findField("new_col") != null; + SerializableTableSpec spec = checkNotNull(element.getValue()); + boolean hasNewCol = spec.getSchema().findField("new_col") != null; out.output(hasNewCol ? "UPDATED_SCHEMA" : "INITIAL_SCHEMA"); } })); @@ -568,6 +577,118 @@ public void processElement( pipeline.run(); } + @Test + public void testStreamingNonExistentTableEmitsEmptyMapWithoutBlockingConsumer() { + Row row = + Row.withSchema(BEAM_SCHEMA) + .addValues(1L, "v1", "default.non_existent_streaming_table") + .build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(row) + .advanceWatermarkToInfinity(); + + PCollection input = pipeline.apply("StreamInput", stream); + + 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) { + Map viewMap = c.sideInput(metadataView); + assertNotNull("View map should not be null", viewMap); + assertTrue( + "View map should be empty when all polled tables do not exist", + viewMap.isEmpty()); + out.output("CONSUMER_UNBLOCKED_EMPTY_MAP"); + } + }) + .withSideInputs(metadataView)); + + PAssert.that(consumerObserved).containsInAnyOrder("CONSUMER_UNBLOCKED_EMPTY_MAP"); + + pipeline.run(); + } + + @Test + public void testStreamingMixedExistingAndNonExistentTables() { + Catalog catalog = getCatalog(); + TableIdentifier validTable = TableIdentifier.of("default", "mixed_valid_table"); + catalog.createTable(validTable, ICEBERG_SCHEMA); + + Row seedValidRow = + Row.withSchema(BEAM_SCHEMA) + .addValues(0L, "seed_valid", "default.mixed_valid_table") + .build(); + Row seedMissingRow = + Row.withSchema(BEAM_SCHEMA) + .addValues(0L, "seed_missing", "default.mixed_missing_table") + .build(); + Row validRow = + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.mixed_valid_table").build(); + Row missingRow = + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.mixed_missing_table").build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(seedValidRow, seedMissingRow) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(validRow, missingRow) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceWatermarkToInfinity(); + + PCollection input = pipeline.apply("StreamInput", stream); + + 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_valid".equals(data) || "seed_missing".equals(data)) { + return; + } + Map viewMap = c.sideInput(metadataView); + assertNotNull(viewMap); + String dest = row.getString("dest"); + if ("default.mixed_valid_table".equals(dest)) { + assertNotNull(viewMap.get(dest)); + out.output("VALID_TABLE_FOUND"); + } else { + assertTrue(!viewMap.containsKey(dest)); + out.output("MISSING_TABLE_NOT_FOUND"); + } + } + }) + .withSideInputs(metadataView)); + + PAssert.that(consumerObserved) + .containsInAnyOrder("VALID_TABLE_FOUND", "MISSING_TABLE_NOT_FOUND"); + + pipeline.run(); + } + @Test public void testMaximumCacheSizeInStreamingThrowsUnsupportedOperationException() { pipeline.enableAbandonedNodeEnforcement(false); @@ -602,7 +723,7 @@ public void testMalformedTableIdentifierSkippedWithoutFailingBundle() { PCollection input = pipeline.apply(Create.of(validRow, malformedRow)).setCoder(RowCoder.of(BEAM_SCHEMA)); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -612,9 +733,13 @@ public void testMalformedTableIdentifierSkippedWithoutFailingBundle() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); - assertEquals(1, list.size()); - assertEquals("default.valid_table", list.get(0).getKey()); + Map map = new HashMap<>(); + for (KV elem : elements) { + map.put(elem.getKey(), elem.getValue()); + } + assertEquals(2, map.size()); + assertNotNull(map.get("default.valid_table")); + assertNull(map.get("default.invalid..name///")); return null; }); @@ -639,7 +764,7 @@ public void testMaximumCacheSizeCap() { PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); int maxCacheSize = 3; - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -650,7 +775,8 @@ public void testMaximumCacheSizeCap() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); + List> list = + ImmutableList.copyOf(elements); assertEquals(maxCacheSize, list.size()); return null; }); @@ -676,7 +802,7 @@ public void testUncappedByDefault() { PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); // Without setting maximumCacheSize, all 10 distinct tables are emitted - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -686,7 +812,8 @@ public void testUncappedByDefault() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); + List> list = + ImmutableList.copyOf(elements); assertEquals(10, list.size()); return null; }); @@ -707,20 +834,24 @@ public void testNonExistentTableIsSkippedWithoutFailingBundle() { PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); - PCollection> specs = + 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 + // Both existing and missing table entries are emitted; missing table has null spec PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); - assertEquals(1, list.size()); - assertEquals("default.existing_table", list.get(0).getKey()); + Map map = new HashMap<>(); + for (KV elem : elements) { + map.put(elem.getKey(), elem.getValue()); + } + assertEquals(2, map.size()); + assertNotNull(map.get("default.existing_table")); + assertNull(map.get("default.non_existent_table")); return null; }); @@ -744,7 +875,7 @@ public void testFiltersNullAndBlankTableIdentifiers() { PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -754,7 +885,8 @@ public void testFiltersNullAndBlankTableIdentifiers() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); + List> list = + ImmutableList.copyOf(elements); assertEquals(1, list.size()); assertEquals("default.valid_dest_table", list.get(0).getKey()); return null; @@ -841,7 +973,7 @@ public void testConfigurablePollingBuckets() { PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -852,7 +984,8 @@ public void testConfigurablePollingBuckets() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); + List> list = + ImmutableList.copyOf(elements); assertEquals(2, list.size()); Map map = list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); @@ -889,7 +1022,7 @@ public void testWindowPreservation() { .setCoder(RowCoder.of(BEAM_SCHEMA)) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(1)))); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -899,7 +1032,8 @@ public void testWindowPreservation() { PAssert.that(specs) .satisfies( elements -> { - List> list = ImmutableList.copyOf(elements); + List> list = + ImmutableList.copyOf(elements); assertEquals(2, list.size()); return null; }); @@ -913,7 +1047,7 @@ public void testEmptyInputProducesEmptyOutput() { PCollection input = pipeline.apply(Create.empty(RowCoder.of(BEAM_SCHEMA))); - PCollection> specs = + PCollection> specs = input.apply( TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) @@ -1159,7 +1293,7 @@ public void testUnusedTablesEvictedFromStreamingCache() { @ProcessElement public void processElement(@Element Row row, OutputReceiver out) { if ("trigger_evict_a".equals(row.getString("data"))) { - ControllableTestClock.setTime(7000L); + ControllableTestClock.setTime(20000L); } out.output(row); } @@ -1200,4 +1334,108 @@ public void processElement( pipeline.run(); } + + @Test + public void testBatchAllNonExistentTablesEmitsEmptyMapWithoutBlockingConsumer() { + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.missing_1").build(), + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.missing_2").build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", TableMetadataDriver.asView(catalogConfig, DYNAMIC_DESTINATIONS)); + + PCollection consumerObserved = + input.apply( + "ConsumeSideInput", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element Row row, OutputReceiver out, ProcessContext c) { + Map viewMap = c.sideInput(metadataView); + out.output("size=" + viewMap.size()); + } + }) + .withSideInputs(metadataView)); + + PAssert.that(consumerObserved).containsInAnyOrder("size=0", "size=0"); + + pipeline.run(); + } + + @Test + public void testStreamingDroppedTableImmediatelyInvalidatedInCache() { + TableIdentifier tableId = TableIdentifier.of("default", "dropped_table"); + getCatalog().createTable(tableId, ICEBERG_SCHEMA); + String tableStr = IcebergUtils.tableIdentifierToString(tableId); + + Duration refreshInterval = Duration.standardSeconds(2); + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1L, "initial", tableStr).build(); + Row rowDrop = Row.withSchema(BEAM_SCHEMA).addValues(2L, "trigger_drop", tableStr).build(); + Row rowPostDrop = Row.withSchema(BEAM_SCHEMA).addValues(3L, "post_drop", tableStr).build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(row1) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(rowDrop) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(rowPostDrop) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceWatermarkToInfinity(); + + PCollection input = + pipeline + .apply("StreamInput", stream) + .apply( + "DropTableOnTriggerRow", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(@Element Row row, OutputReceiver out) { + if ("trigger_drop".equals(row.getString("data"))) { + catalogConfig + .catalog() + .dropTable( + IcebergUtils.parseTableIdentifier("default.dropped_table")); + } + out.output(row); + } + })) + .setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", + TableMetadataDriver.asView(catalogConfig, DYNAMIC_DESTINATIONS, null, refreshInterval)); + + 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 ("trigger_drop".equals(data)) { + return; + } + Map viewMap = c.sideInput(metadataView); + out.output(data + ":hasTable=" + viewMap.containsKey(tableStr)); + } + }) + .withSideInputs(metadataView)); + + PAssert.that(consumerObserved) + .containsInAnyOrder("initial:hasTable=true", "post_drop:hasTable=false"); + + pipeline.run(); + } }