diff --git a/sdks/java/io/iceberg/build.gradle b/sdks/java/io/iceberg/build.gradle index b3228004ed3b..7278ca56159e 100644 --- a/sdks/java/io/iceberg/build.gradle +++ b/sdks/java/io/iceberg/build.gradle @@ -50,6 +50,7 @@ dependencies { implementation library.java.avro implementation library.java.slf4j_api implementation library.java.joda_time + implementation library.java.jackson_databind implementation "org.apache.parquet:parquet-column:$parquet_version" implementation "org.apache.parquet:parquet-hadoop:$parquet_version" implementation "org.apache.parquet:parquet-common:$parquet_version" diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfig.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfig.java index 8fe05cb5b21f..fd836dc30ddd 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfig.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfig.java @@ -80,6 +80,21 @@ public Catalog catalog() { return CATALOG_CACHE.computeIfAbsent(this, IcebergCatalogConfig::buildCatalog); } + /** + * Constructs and returns a new {@link Configuration} populated with properties from {@link + * #getConfigProperties()}. + */ + public Configuration getHadoopConfiguration() { + Configuration config = new Configuration(); + Map confProps = getConfigProperties(); + if (confProps != null) { + for (Map.Entry prop : confProps.entrySet()) { + config.set(prop.getKey(), prop.getValue()); + } + } + return config; + } + private static Catalog buildCatalog(IcebergCatalogConfig catalogConfig) { String catalogName = catalogConfig.getCatalogName(); if (catalogName == null) { @@ -89,14 +104,7 @@ private static Catalog buildCatalog(IcebergCatalogConfig catalogConfig) { if (catalogProps == null) { catalogProps = Maps.newHashMap(); } - Map confProps = catalogConfig.getConfigProperties(); - if (confProps == null) { - confProps = Maps.newHashMap(); - } - Configuration config = new Configuration(); - for (Map.Entry prop : confProps.entrySet()) { - config.set(prop.getKey(), prop.getValue()); - } + Configuration config = catalogConfig.getHadoopConfiguration(); return CatalogUtil.buildIcebergCatalog(catalogName, catalogProps, config); } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java index 25e5a13da43e..e7e79c495535 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java @@ -27,7 +27,6 @@ import java.time.format.DateTimeFormatter; import java.time.temporal.ChronoUnit; import java.util.ArrayList; -import java.util.Collections; import java.util.List; import java.util.Map; import java.util.UUID; @@ -318,9 +317,7 @@ Table getOrCreateTable( if (sideInputTableSpecs != null && sideInputTableSpecs.containsKey(tableIdString)) { SerializableTableSpec spec = sideInputTableSpecs.get(tableIdString); if (spec != null) { - Map catalogProperties = catalogConfig.getCatalogProperties(); - return new SideInputTable( - spec, catalogProperties != null ? catalogProperties : Collections.emptyMap()); + return new SideInputTable(spec, catalogConfig); } } return TableCache.getAndRefreshIfStale( 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 e89a341f55c3..ea61b496cc1c 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 @@ -19,7 +19,9 @@ import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import com.fasterxml.jackson.databind.JsonNode; import com.google.auto.value.AutoValue; +import java.io.IOException; import java.io.Serializable; import java.util.Collections; import java.util.List; @@ -33,6 +35,8 @@ import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; import org.apache.beam.sdk.schemas.annotations.SchemaIgnore; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.hadoop.conf.Configurable; +import org.apache.hadoop.conf.Configuration; import org.apache.iceberg.EncryptedKeyParser; import org.apache.iceberg.HasTableOperations; import org.apache.iceberg.PartitionSpec; @@ -47,6 +51,7 @@ import org.apache.iceberg.encryption.EncryptedKey; import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.FileIOParser; +import org.apache.iceberg.util.JsonUtil; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; @@ -185,7 +190,14 @@ public Map getSortOrders() { if (local == null) { ImmutableMap.Builder builder = ImmutableMap.builder(); for (Map.Entry entry : getSortOrdersJson().entrySet()) { - builder.put(entry.getKey(), SortOrderParser.fromJson(getSchema(), entry.getValue())); + try { + JsonNode node = JsonUtil.mapper().readTree(entry.getValue()); + builder.put( + entry.getKey(), SortOrderParser.fromJson(getSchema(), node, getOrderId())); + } catch (IOException e) { + throw new IllegalArgumentException( + "Failed to parse sort order JSON for orderId " + entry.getKey(), e); + } } cachedSortOrders = local = builder.build(); } @@ -224,17 +236,34 @@ public TableIdentifier getTableIdentifier() { return local; } + /** Returns a cached {@link FileIO} instance for this table using default configuration. */ @SchemaIgnore public FileIO getFileIO() { + return getFileIO(null); + } + + /** + * Returns a cached {@link FileIO} instance for this table, configured with the provided Hadoop + * {@link Configuration} if supported. + */ + @SchemaIgnore + public FileIO getFileIO(@Nullable Configuration conf) { FileIO local = cachedFileIO; if (local == null) { synchronized (this) { local = cachedFileIO; if (local == null) { - cachedFileIO = local = FileIOParser.fromJson(getFileIoJson()); + cachedFileIO = + local = + conf != null + ? FileIOParser.fromJson(getFileIoJson(), conf) + : FileIOParser.fromJson(getFileIoJson()); } } } + if (conf != null && local instanceof Configurable) { + ((Configurable) local).setConf(conf); + } return local; } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SideInputTable.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SideInputTable.java index aa51571c4cbd..c5a4ddc3d82d 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SideInputTable.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SideInputTable.java @@ -25,6 +25,7 @@ import java.util.Objects; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.apache.hadoop.conf.Configuration; import org.apache.iceberg.AppendFiles; import org.apache.iceberg.DeleteFiles; import org.apache.iceberg.ExpireSnapshots; @@ -61,6 +62,7 @@ import org.apache.iceberg.encryption.PlaintextEncryptionManager; import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.LocationProvider; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A lightweight adapter that implements {@link Table} backed by a {@link SerializableTableSpec}. @@ -80,14 +82,23 @@ public class SideInputTable implements Table { private final SerializableTableSpec spec; private final EncryptionManager encryptionManager; private final LocationProvider locationProvider; + private final @Nullable Configuration hadoopConf; public SideInputTable(SerializableTableSpec spec) { - this(spec, Collections.emptyMap()); + this(spec, Collections.emptyMap(), null); } public SideInputTable(SerializableTableSpec spec, Map catalogProperties) { + this(spec, catalogProperties, null); + } + + public SideInputTable( + SerializableTableSpec spec, + Map catalogProperties, + @Nullable Configuration hadoopConf) { this.spec = checkNotNull(spec, "spec must not be null"); checkNotNull(catalogProperties, "catalogProperties must not be null"); + this.hadoopConf = hadoopConf; this.locationProvider = LocationProviders.locationsFor(spec.getLocation(), spec.getProperties()); @@ -102,10 +113,27 @@ public SideInputTable(SerializableTableSpec spec, Map catalogPro } public SideInputTable(SerializableTableSpec spec, EncryptionManager encryptionManager) { + this(spec, encryptionManager, null); + } + + public SideInputTable( + SerializableTableSpec spec, + EncryptionManager encryptionManager, + @Nullable Configuration hadoopConf) { this.spec = checkNotNull(spec, "spec must not be null"); this.encryptionManager = checkNotNull(encryptionManager, "encryptionManager must not be null"); this.locationProvider = LocationProviders.locationsFor(spec.getLocation(), spec.getProperties()); + this.hadoopConf = hadoopConf; + } + + public SideInputTable(SerializableTableSpec spec, IcebergCatalogConfig catalogConfig) { + this( + spec, + checkNotNull(catalogConfig, "catalogConfig must not be null").getCatalogProperties() != null + ? catalogConfig.getCatalogProperties() + : Collections.emptyMap(), + catalogConfig.getHadoopConfiguration()); } public SerializableTableSpec getTableSpec() { @@ -164,7 +192,7 @@ public LocationProvider locationProvider() { @Override public FileIO io() { - return spec.getFileIO(); + return spec.getFileIO(hadoopConf); } @Override @@ -363,6 +391,7 @@ public String toString() { return MoreObjects.toStringHelper(this) .add("spec", spec) .add("encryptionManager", encryptionManager) + .add("hadoopConf", hadoopConf) .toString(); } } 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 679463ab5db5..d7442e3d23b5 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 @@ -24,6 +24,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import org.apache.beam.sdk.annotations.Internal; @@ -252,19 +253,24 @@ public PCollectionView> expand(PCollection> expand(PCollection input) { + boolean isStreaming = input.isBounded() == PCollection.IsBounded.UNBOUNDED; + + Duration customInterval = getRefreshInterval(); + Duration interval = + checkNotNull(customInterval != null ? customInterval : DEFAULT_REFRESH_INTERVAL); + PCollection tableIds = input - .apply("ExtractTableIds", ParDo.of(new ExtractTableIdsDoFn(getDynamicDestinations()))) + .apply( + "ExtractTableIds", + ParDo.of( + new ExtractTableIdsDoFn( + getDynamicDestinations(), isStreaming ? interval : null))) .setCoder(StringUtf8Coder.of()) .apply("MetadataGlobalWindow", Window.into(new GlobalWindows())); - boolean isStreaming = input.isBounded() == PCollection.IsBounded.UNBOUNDED; - PCollection distinctTableIds; if (isStreaming) { - Duration customInterval = getRefreshInterval(); - Duration interval = - checkNotNull(customInterval != null ? customInterval : DEFAULT_REFRESH_INTERVAL); distinctTableIds = tableIds.apply( "DeduplicateTableIds", Deduplicate.values().withDuration(interval)); @@ -323,10 +329,61 @@ public void populateDisplayData(DisplayData.Builder builder) { } static class ExtractTableIdsDoFn extends DoFn { + private static final int DEFAULT_LOCAL_CACHE_MAX_SIZE = 10_000; + + private static volatile @Nullable Clock globalTestClock; + private final DynamicDestinations dynamicDestinations; + private final @Nullable Duration refreshInterval; + private transient @Nullable Clock clock; + private transient @Nullable LinkedHashMap lastEmittedCache; ExtractTableIdsDoFn(DynamicDestinations dynamicDestinations) { + this(dynamicDestinations, null); + } + + ExtractTableIdsDoFn( + DynamicDestinations dynamicDestinations, @Nullable Duration refreshInterval) { this.dynamicDestinations = dynamicDestinations; + this.refreshInterval = refreshInterval; + } + + @VisibleForTesting + void setClock(@Nullable Clock clock) { + this.clock = clock; + } + + @VisibleForTesting + static void setGlobalTestClock(@Nullable Clock clock) { + globalTestClock = clock; + } + + @Setup + public void setup() { + initCache(); + } + + private void initCache() { + if (lastEmittedCache == null && refreshInterval != null) { + this.lastEmittedCache = + new LinkedHashMap(16, 0.75f, true) { + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + return size() > DEFAULT_LOCAL_CACHE_MAX_SIZE; + } + }; + } + } + + private long getNow() { + if (clock != null) { + return clock.currentTimeMillis(); + } + Clock global = globalTestClock; + if (global != null) { + return global.currentTimeMillis(); + } + return System.currentTimeMillis(); } @ProcessElement @@ -340,7 +397,19 @@ public void processElement( dynamicDestinations.getTableStringIdentifier( ValueInSingleWindow.of(element, timestamp, window, paneInfo)); if (tableIdentifier != null && !tableIdentifier.trim().isEmpty()) { - out.output(tableIdentifier.trim()); + Map cache = lastEmittedCache; + Duration interval = refreshInterval; + if (cache != null && interval != null) { + long now = getNow(); + Long lastEmitted = cache.get(tableIdentifier); + long minInterval = Math.max(1L, interval.getMillis() / 2); + if (lastEmitted == null || (now - lastEmitted) >= minInterval) { + cache.put(tableIdentifier, now); + out.output(tableIdentifier); + } + } else { + out.output(tableIdentifier); + } } } } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java index 881d2577fad6..fbb6cb08aca4 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java @@ -22,7 +22,6 @@ import static org.apache.beam.sdk.io.iceberg.RecordWriterManager.getPartitionDataPath; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; -import java.util.Collections; import java.util.Map; import java.util.UUID; import org.apache.beam.sdk.coders.IterableCoder; @@ -223,9 +222,7 @@ Table getOrCreateTable( if (sideInputTableSpecs != null && sideInputTableSpecs.containsKey(tableIdString)) { SerializableTableSpec spec = sideInputTableSpecs.get(tableIdString); if (spec != null) { - Map catalogProperties = catalogConfig.getCatalogProperties(); - return new SideInputTable( - spec, catalogProperties != null ? catalogProperties : Collections.emptyMap()); + return new SideInputTable(spec, catalogConfig); } } return TableCache.getAndRefreshIfStale( diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfigTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfigTest.java new file mode 100644 index 000000000000..4e0e9c95b50e --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfigTest.java @@ -0,0 +1,56 @@ +/* + * 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.assertNull; + +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.hadoop.conf.Configuration; +import org.junit.Test; + +public class IcebergCatalogConfigTest { + + @Test + public void testGetHadoopConfigurationWhenPropertiesNull() { + IcebergCatalogConfig config = + IcebergCatalogConfig.builder().setCatalogName("test_catalog").build(); + + Configuration hadoopConf = config.getHadoopConfiguration(); + assertNotNull(hadoopConf); + assertNull(hadoopConf.get("non.existent.key")); + } + + @Test + public void testGetHadoopConfigurationPopulatesProperties() { + IcebergCatalogConfig config = + IcebergCatalogConfig.builder() + .setCatalogName("test_catalog") + .setConfigProperties( + ImmutableMap.of( + "fs.defaultFS", "file:///test/path", + "custom.hadoop.key", "custom-hadoop-val")) + .build(); + + Configuration hadoopConf = config.getHadoopConfiguration(); + assertNotNull(hadoopConf); + assertEquals("file:///test/path", hadoopConf.get("fs.defaultFS")); + assertEquals("custom-hadoop-val", hadoopConf.get("custom.hadoop.key")); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java index 246f8410ec3c..289fabc69918 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java @@ -607,6 +607,12 @@ public void processElement(@Element Row row, OutputReceiver out) { catalogConfig.catalog().loadTable(IcebergUtils.parseTableIdentifier(tableIdString)); if (table.spec().isUnpartitioned()) { table.updateSpec().addField("city").commit(); + // Ensure worker-local table ID cache TTL (interval / 2 = 500ms) has elapsed + try { + Thread.sleep(700); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } } } out.output(row); 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 a2f2abadfe49..9045b910c650 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 @@ -21,6 +21,7 @@ import static org.apache.iceberg.types.Types.NestedField.required; 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 static org.mockito.Mockito.mock; @@ -42,6 +43,7 @@ import org.apache.beam.sdk.schemas.SchemaCoder; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.hadoop.conf.Configurable; import org.apache.hadoop.conf.Configuration; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.CatalogUtil; @@ -350,4 +352,79 @@ public void testConcurrentGetterInitializationThreadSafety() throws Exception { executor.shutdown(); } } + + @Test + public void testHistoricalSortOrderWithDroppedColumn() { + TableIdentifier tableId = TableIdentifier.of("default", "historical_sort_table"); + Schema v1Schema = + new Schema( + required(1, "id", Types.LongType.get()), + optional(2, "name", Types.StringType.get()), + optional(3, "dropped_col", Types.StringType.get())); + + SortOrder v1SortOrder = + SortOrder.builderFor(v1Schema) + .sortBy("dropped_col", SortDirection.ASC, NullOrder.NULLS_FIRST) + .build(); + + Table table = catalog.buildTable(tableId, v1Schema).withSortOrder(v1SortOrder).create(); + + int v1OrderId = table.sortOrder().orderId(); + + // First replace sort order with one referencing the remaining fields, + // making v1SortOrder a historical sort order + table.replaceSortOrder().asc("id").commit(); + int v2OrderId = table.sortOrder().orderId(); + + // Then evolve schema by deleting the column that was part of the original sort order + table.updateSchema().deleteColumn("dropped_col").commit(); + + SerializableTableSpec spec = SerializableTableSpec.fromTable(table); + + assertEquals(v2OrderId, spec.getOrderId()); + assertEquals(table.sortOrder(), spec.getSortOrder()); + assertEquals(1, spec.getSortOrder().fields().get(0).sourceId()); + + // Calling getSortOrders() should successfully bind historical sort orders + // without failing with ValidationException: Cannot find source column + Map sortOrders = spec.getSortOrders(); + assertNotNull(sortOrders); + assertTrue(sortOrders.containsKey(v1OrderId)); + assertTrue(sortOrders.containsKey(v2OrderId)); + + SortOrder historicalOrder = spec.getSortOrder(v1OrderId); + assertNotNull(historicalOrder); + assertEquals(v1OrderId, historicalOrder.orderId()); + } + + @Test + public void testFileIOWithHadoopConfiguration() { + TableIdentifier tableId = TableIdentifier.of("default", "hadoop_conf_table"); + Table table = catalog.createTable(tableId, TestFixtures.SCHEMA); + + SerializableTableSpec spec = SerializableTableSpec.fromTable(table); + + Configuration conf = new Configuration(); + conf.set("custom.test.prop", "test-value-123"); + + FileIO fileIO = spec.getFileIO(conf); + assertNotNull(fileIO); + assertTrue("FileIO must be Configurable", fileIO instanceof Configurable); + assertEquals("test-value-123", ((Configurable) fileIO).getConf().get("custom.test.prop")); + + // Also verify that zero-arg getFileIO() returns cached instance + assertEquals(fileIO, spec.getFileIO()); + + // Verify that calling zero-arg getFileIO() first does not pollute cache: + // subsequent getFileIO(conf) must update configuration on Configurable FileIO + SerializableTableSpec spec2 = SerializableTableSpec.fromTable(table); + FileIO fileIO2 = spec2.getFileIO(); + assertNotNull(fileIO2); + assertNull(((Configurable) fileIO2).getConf().get("custom.test.prop")); + + FileIO configuredFileIO2 = spec2.getFileIO(conf); + assertEquals(fileIO2, configuredFileIO2); + assertEquals( + "test-value-123", ((Configurable) configuredFileIO2).getConf().get("custom.test.prop")); + } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SideInputTableTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SideInputTableTest.java index 663c818b5879..6b1aa1843f0b 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SideInputTableTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SideInputTableTest.java @@ -25,6 +25,7 @@ import java.util.Map; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.hadoop.conf.Configurable; import org.apache.hadoop.conf.Configuration; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.CatalogUtil; @@ -40,6 +41,7 @@ import org.apache.iceberg.data.GenericRecord; import org.apache.iceberg.data.Record; import org.apache.iceberg.encryption.PlaintextEncryptionManager; +import org.apache.iceberg.io.FileIO; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -236,4 +238,50 @@ public void testEqualsHashCodeAndToString() { assertNotNull(table1a.toString()); assertTrue(table1a.toString().contains("SideInputTable")); } + + @Test + public void testSideInputTableWithHadoopConfiguration() { + TableIdentifier tableId = TableIdentifier.of("default", "side_input_hadoop_conf_table"); + Table realTable = catalog.createTable(tableId, TestFixtures.SCHEMA); + SerializableTableSpec spec = SerializableTableSpec.fromTable(tableId, realTable); + + Configuration conf = new Configuration(); + conf.set("custom.sideinput.prop", "custom-value-456"); + + // Test constructor accepting hadoopConf directly + SideInputTable tableWithConf = new SideInputTable(spec, ImmutableMap.of(), conf); + FileIO io = tableWithConf.io(); + assertNotNull(io); + assertTrue("FileIO must implement Configurable", io instanceof Configurable); + assertEquals("custom-value-456", ((Configurable) io).getConf().get("custom.sideinput.prop")); + + // Test constructor accepting IcebergCatalogConfig + IcebergCatalogConfig catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogName("test_catalog") + .setConfigProperties(ImmutableMap.of("catalog.conf.prop", "val-789")) + .build(); + + SerializableTableSpec specFromConfig = SerializableTableSpec.fromTable(tableId, realTable); + SideInputTable tableFromCatalogConfig = new SideInputTable(specFromConfig, catalogConfig); + FileIO ioFromConfig = tableFromCatalogConfig.io(); + assertNotNull(ioFromConfig); + assertTrue("FileIO must implement Configurable", ioFromConfig instanceof Configurable); + assertEquals("val-789", ((Configurable) ioFromConfig).getConf().get("catalog.conf.prop")); + + // Test constructor accepting EncryptionManager and hadoopConf + SerializableTableSpec specWithEncryption = SerializableTableSpec.fromTable(tableId, realTable); + SideInputTable tableWithEncryptionAndConf = + new SideInputTable(specWithEncryption, PlaintextEncryptionManager.instance(), conf); + FileIO ioWithEncryption = tableWithEncryptionAndConf.io(); + assertNotNull(ioWithEncryption); + assertTrue("FileIO must implement Configurable", ioWithEncryption instanceof Configurable); + assertEquals( + "custom-value-456", + ((Configurable) ioWithEncryption).getConf().get("custom.sideinput.prop")); + + // Test constructor rejecting null catalogConfig + assertThrows( + NullPointerException.class, () -> new SideInputTable(spec, (IcebergCatalogConfig) null)); + } } 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 cb627f150630..48f86ff3239b 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 @@ -18,6 +18,7 @@ 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.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; @@ -36,10 +37,13 @@ import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.testing.TestStream; import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.Deduplicate; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.display.DisplayData; import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; @@ -62,8 +66,10 @@ import org.apache.iceberg.data.Record; import org.apache.iceberg.types.Types; import org.checkerframework.checker.nullness.qual.Nullable; +import org.hamcrest.Matchers; import org.joda.time.Duration; import org.joda.time.Instant; +import org.junit.After; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -118,6 +124,19 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { } }; + 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(); + } + } + @Before public void setUp() throws Exception { warehouseLocation = "file:" + tempFolder.newFolder().getAbsolutePath(); @@ -126,6 +145,13 @@ public void setUp() throws Exception { .setCatalogName("hadoop") .setCatalogProperties(ImmutableMap.of("type", "hadoop", "warehouse", warehouseLocation)) .build(); + ControllableTestClock.setTime(1000L); + TableMetadataDriver.ExtractTableIdsDoFn.setGlobalTestClock(new ControllableTestClock()); + } + + @After + public void tearDown() { + TableMetadataDriver.ExtractTableIdsDoFn.setGlobalTestClock(null); } private Catalog getCatalog() { @@ -321,6 +347,9 @@ public void testMetadataRefreshedAcrossIntervals() { TableIdentifier tableId = TableIdentifier.of("default", "evolving_table"); catalog.createTable(tableId, ICEBERG_SCHEMA); + Duration refreshInterval = Duration.standardSeconds(2); + ControllableTestClock.setTime(1000L); + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1L, "initial_data", "default.evolving_table").build(); Row row2 = @@ -347,6 +376,7 @@ public void testMetadataRefreshedAcrossIntervals() { @ProcessElement public void processElement(@Element Row row, OutputReceiver out) { if ("trigger_update".equals(row.getString("data"))) { + ControllableTestClock.setTime(5000L); Table table = catalogConfig .catalog() @@ -367,7 +397,7 @@ public void processElement(@Element Row row, OutputReceiver out) { TableMetadataDriver.builder() .setCatalogConfig(catalogConfig) .setDynamicDestinations(DYNAMIC_DESTINATIONS) - .setRefreshInterval(Duration.standardSeconds(2)) + .setRefreshInterval(refreshInterval) .build()); // Downstream consumer transform verifying that updated metadata is received @@ -398,6 +428,10 @@ public void testMetadataRefreshedAcrossIntervalsAsSideInput() { catalog.createTable(tableId, ICEBERG_SCHEMA); String tableIdStr = "default.evolving_side_input_table"; + Duration refreshInterval = Duration.standardSeconds(2); + ControllableTestClock.setTime(1000L); + ControllableTestClock testClock = new ControllableTestClock(); + 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(); @@ -422,7 +456,9 @@ public void testMetadataRefreshedAcrossIntervalsAsSideInput() { new DoFn() { @ProcessElement public void processElement(@Element Row row, OutputReceiver out) { - if ("trigger_update".equals(row.getString("data"))) { + String data = row.getString("data"); + if ("trigger_update".equals(data)) { + ControllableTestClock.setTime(5000L); Table table = catalogConfig .catalog() @@ -433,6 +469,8 @@ public void processElement(@Element Row row, OutputReceiver out) { .updateSchema() .addColumn("new_col", Types.StringType.get()) .commit(); + } else if ("post_update_data".equals(data)) { + ControllableTestClock.setTime(10000L); } out.output(row); } @@ -442,12 +480,13 @@ public void processElement(@Element Row row, OutputReceiver out) { PCollectionView> metadataView = input.apply( "CreateMetadataView", - TableMetadataDriver.builder() - .setCatalogConfig(catalogConfig) - .setDynamicDestinations(DYNAMIC_DESTINATIONS) - .setRefreshInterval(Duration.standardSeconds(2)) - .build() - .asView()); + TableMetadataDriver.asView( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(refreshInterval) + .build(), + testClock)); PCollection consumerObserved = input.apply( @@ -490,6 +529,9 @@ public void testMetadataRefreshedAcrossIntervalsAsSideInputWithMultipleTables() String tableAStr = "default.multi_table_a"; String tableBStr = "default.multi_table_b"; + Duration refreshInterval = Duration.standardSeconds(2); + 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(); @@ -523,6 +565,7 @@ public void testMetadataRefreshedAcrossIntervalsAsSideInputWithMultipleTables() @ProcessElement public void processElement(@Element Row row, OutputReceiver out) { if ("trigger_update_a".equals(row.getString("data"))) { + ControllableTestClock.setTime(5000L); Table table = catalogConfig .catalog() @@ -541,12 +584,13 @@ public void processElement(@Element Row row, OutputReceiver out) { PCollectionView> metadataView = input.apply( "CreateMetadataView", - TableMetadataDriver.builder() - .setCatalogConfig(catalogConfig) - .setDynamicDestinations(DYNAMIC_DESTINATIONS) - .setRefreshInterval(Duration.standardSeconds(2)) - .build() - .asView()); + TableMetadataDriver.asView( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(refreshInterval) + .build(), + testClock)); PCollection consumerObserved = input.apply( @@ -884,10 +928,7 @@ public void testFiltersNullAndBlankTableIdentifiers() { 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()); + Row.withSchema(BEAM_SCHEMA).addValues(4L, "v4", "default.valid_dest_table").build()); PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); @@ -1234,19 +1275,6 @@ public void testMapMergerFnTieBreaksBySchemaIdCommutatively() { 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"); @@ -1498,4 +1526,137 @@ public void processElement( pipeline.run(); } + + @Test + public void testExtractTableIdsWorkerLocalPreFiltering() { + TestStream.Builder streamBuilder = TestStream.create(BEAM_SCHEMA); + for (int i = 0; i < 1000; i++) { + streamBuilder = + streamBuilder.addElements( + Row.withSchema(BEAM_SCHEMA).addValues((long) i, "data", "default.table").build()); + } + TestStream testStream = streamBuilder.advanceWatermarkToInfinity(); + + PCollection tableIds = + pipeline + .apply(testStream) + .apply( + ParDo.of( + new TableMetadataDriver.ExtractTableIdsDoFn( + SINGLE_TABLE_DYNAMIC_DESTINATIONS, Duration.standardMinutes(5)))); + + // With worker-local pre-filtering, 1,000 rows emit at most 1 string per worker thread + // rather than 1,000 strings + PAssert.that(tableIds) + .satisfies( + actual -> { + List list = ImmutableList.copyOf(actual); + for (String id : list) { + assertEquals("default.table", id); + } + assertThat(list.size(), Matchers.lessThanOrEqualTo(100)); + return null; + }); + + PCollection distinctIds = + tableIds.apply(Deduplicate.values().withDuration(Duration.standardMinutes(5))); + PAssert.that(distinctIds).containsInAnyOrder("default.table"); + pipeline.run(); + } + + @Test + public void testExtractTableIdsMultipleTables() { + TestStream.Builder streamBuilder = TestStream.create(BEAM_SCHEMA); + for (int i = 0; i < 100; i++) { + streamBuilder = + streamBuilder.addElements( + Row.withSchema(BEAM_SCHEMA).addValues((long) i, "data", "default.table_a").build(), + Row.withSchema(BEAM_SCHEMA) + .addValues((long) (i + 100), "data", "default.table_b") + .build()); + } + TestStream testStream = streamBuilder.advanceWatermarkToInfinity(); + + PCollection tableIds = + pipeline + .apply(testStream) + .apply( + ParDo.of( + new TableMetadataDriver.ExtractTableIdsDoFn( + DYNAMIC_DESTINATIONS, Duration.standardMinutes(5)))); + + PCollection distinctIds = + tableIds.apply(Deduplicate.values().withDuration(Duration.standardMinutes(5))); + PAssert.that(distinctIds).containsInAnyOrder("default.table_a", "default.table_b"); + pipeline.run(); + } + + @Test + public void testExtractTableIdsCacheExpiration() { + TableMetadataDriver.ExtractTableIdsDoFn doFn = + new TableMetadataDriver.ExtractTableIdsDoFn( + SINGLE_TABLE_DYNAMIC_DESTINATIONS, Duration.standardMinutes(10)); + ControllableTestClock testClock = new ControllableTestClock(); + ControllableTestClock.setTime(1000L); + doFn.setClock(testClock); + + List outputs = new ArrayList<>(); + DoFn.OutputReceiver receiver = + new DoFn.OutputReceiver() { + @Override + public void output(String output) { + outputs.add(output); + } + + @Override + public void outputWithTimestamp(String output, Instant timestamp) { + outputs.add(output); + } + + @Override + public org.apache.beam.sdk.values.OutputBuilder builder(String output) { + throw new UnsupportedOperationException(); + } + }; + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1L, "data", "default.table").build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2L, "data", "default.table").build(); + Row row3 = Row.withSchema(BEAM_SCHEMA).addValues(3L, "data", "default.table").build(); + + doFn.setup(); + doFn.processElement(row1, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING, Instant.now(), receiver); + assertEquals(1, outputs.size()); + assertEquals("default.table", outputs.get(0)); + + // Second element should be suppressed by worker-local cache + doFn.processElement(row2, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING, Instant.now(), receiver); + assertEquals(1, outputs.size()); + + // Advance clock beyond interval / 2 (5 minutes) + ControllableTestClock.setTime(1000L + Duration.standardMinutes(6).getMillis()); + + // Third element arrives after expiration -> should be emitted + doFn.processElement(row3, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING, Instant.now(), receiver); + assertEquals(2, outputs.size()); + assertEquals("default.table", outputs.get(1)); + } + + @Test + public void testExtractTableIdsIgnoresNullAndWhitespace() { + List rows = new ArrayList<>(); + rows.add(Row.withSchema(BEAM_SCHEMA).addValues(1L, "data", (String) null).build()); + rows.add(Row.withSchema(BEAM_SCHEMA).addValues(2L, "data", "").build()); + rows.add(Row.withSchema(BEAM_SCHEMA).addValues(3L, "data", " ").build()); + + PCollection tableIds = + pipeline + .apply(Create.of(rows)) + .apply( + ParDo.of( + new TableMetadataDriver.ExtractTableIdsDoFn( + DYNAMIC_DESTINATIONS, Duration.standardMinutes(5)))); + + PAssert.that(tableIds).empty(); + pipeline.run(); + } }