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..0fe886916cf5 --- /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 testDeduplicationOfTablesAcrossRows() { + 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(); + } +}