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 new file mode 100644 index 000000000000..cf5d6310a1d7 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriver.java @@ -0,0 +1,539 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; + +import com.google.auto.value.AutoValue; +import java.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.NullableCoder; +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; +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; +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.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; +import org.apache.beam.sdk.transforms.windowing.AfterPane; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindows; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.transforms.windowing.Repeatedly; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.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; +import org.apache.iceberg.exceptions.NoSuchTableException; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A driver transform that extracts table identifiers from incoming {@link Row}s, deduplicates them + * per window, optionally bounds the cache size up to {@code maximumCacheSize} (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. 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>> { + + 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(); + + public abstract @Nullable Integer getMaximumCacheSize(); + + public abstract @Nullable Duration getRefreshInterval(); + + /** + * Returns the number of parallel buckets/workers used to query the Iceberg catalog, or {@code + * null} for default. + */ + public abstract @Nullable Integer getPollingBuckets(); + + public static Builder builder() { + return new AutoValue_TableMetadataDriver.Builder(); + } + + public abstract Builder toBuilder(); + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setCatalogConfig(IcebergCatalogConfig catalogConfig); + + public abstract Builder setDynamicDestinations(DynamicDestinations dynamicDestinations); + + public abstract Builder setMaximumCacheSize(@Nullable Integer maximumCacheSize); + + public abstract Builder setRefreshInterval(@Nullable Duration refreshInterval); + + /** + * Sets the number of parallel buckets (worker tasks) used to query the Iceberg catalog. + * + *

Defaults to {@link #DEFAULT_POLLING_BUCKETS} (1), which serializes all catalog lookups to + * avoid overwhelming catalog metastores (e.g. Hive Metastore, REST catalog). For pipelines + * writing to a large number of distinct dynamic tables (e.g. hundreds of tables per window), + * consider increasing this value (e.g. 5–10) to parallelize catalog lookups while still + * bounding load. + */ + public abstract Builder setPollingBuckets(@Nullable Integer pollingBuckets); + + abstract TableMetadataDriver autoBuild(); + + public TableMetadataDriver build() { + TableMetadataDriver driver = autoBuild(); + Integer maxCacheSize = driver.getMaximumCacheSize(); + if (maxCacheSize != null) { + Preconditions.checkArgument( + maxCacheSize > 0, "maximumCacheSize must be greater than 0, got %s", maxCacheSize); + } + Duration refreshInterval = driver.getRefreshInterval(); + if (refreshInterval != null) { + Preconditions.checkArgument( + refreshInterval.isLongerThan(Duration.ZERO), + "refreshInterval must be positive, got %s", + refreshInterval); + } + Integer pollingBuckets = driver.getPollingBuckets(); + if (pollingBuckets != null) { + Preconditions.checkArgument( + pollingBuckets > 0, "pollingBuckets must be greater than 0, got %s", pollingBuckets); + } + return driver; + } + } + + /** + * Helper that applies this {@link TableMetadataDriver} and creates a {@link PCollectionView} of + * {@link Map} of table identifier strings to {@link SerializableTableSpec}. + */ + public PTransform, PCollectionView>> + asView() { + return asView(this, null, null); + } + + /** + * Helper that applies {@link TableMetadataDriver} with default configuration 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 builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .build() + .asView(); + } + + @VisibleForTesting + static PTransform, PCollectionView>> asView( + TableMetadataDriver driver, @Nullable Clock clock) { + return asView(driver, null, clock); + } + + @VisibleForTesting + static PTransform, PCollectionView>> asView( + TableMetadataDriver driver, @Nullable Duration cacheTtl, @Nullable Clock clock) { + Preconditions.checkNotNull(driver, "driver must not be null"); + return new PTransform, PCollectionView>>() { + @Override + public PCollectionView> expand(PCollection input) { + boolean isStreaming = input.isBounded() == PCollection.IsBounded.UNBOUNDED; + + Duration customInterval = driver.getRefreshInterval(); + Duration interval = customInterval != null ? customInterval : DEFAULT_REFRESH_INTERVAL; + + PCollection> specs = + input.apply("GenerateTableMetadata", driver); + + if (isStreaming) { + AccumulateTableMetadataMapDoFn accumulateDoFn = + new AccumulateTableMetadataMapDoFn(interval, cacheTtl, clock); + return specs + .apply("KeyForGlobalCache", WithKeys.of((Void) null)) + .setCoder(KvCoder.of(VoidCoder.of(), specs.getCoder())) + .apply("AccumulateCacheMap", ParDo.of(accumulateDoFn)) + .setCoder(MapCoder.of(StringUtf8Coder.of(), SerializableTableSpec.getCoder())) + .apply( + "StreamingCacheWindow", + Window.>into(new GlobalWindows()) + .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1))) + .discardingFiredPanes()) + .apply( + "CreateMetadataSingletonView", + Combine.globally(new MapMergerFn()).asSingletonView()); + } + + return specs + .apply( + "FilterValidSpecsForView", + Filter.by( + (SerializableFunction, Boolean>) + kv -> kv.getValue() != null)) + .apply("CreateTableMetadataView", View.asMap()); + } + }; + } + + @Override + public PCollection> expand(PCollection input) { + PCollection tableIds = + input + .apply("ExtractTableIds", ParDo.of(new ExtractTableIdsDoFn(getDynamicDestinations()))) + .setCoder(StringUtf8Coder.of()) + .apply("MetadataGlobalWindow", Window.into(new GlobalWindows())); + + boolean isStreaming = input.isBounded() == PCollection.IsBounded.UNBOUNDED; + + PCollection distinctTableIds; + if (isStreaming) { + Duration customInterval = getRefreshInterval(); + Duration interval = + checkNotNull(customInterval != null ? customInterval : DEFAULT_REFRESH_INTERVAL); + distinctTableIds = + tableIds.apply( + "DeduplicateTableIds", Deduplicate.values().withDuration(interval)); + } else { + distinctTableIds = tableIds.apply("DistinctTableIds", Distinct.create()); + } + + PCollection cachedTableIds; + Integer maxCacheSize = getMaximumCacheSize(); + if (maxCacheSize != null) { + if (isStreaming) { + throw new UnsupportedOperationException( + "maximumCacheSize is currently not supported for unbounded streaming pipelines."); + } + cachedTableIds = distinctTableIds.apply("CapCacheSize", Sample.any(maxCacheSize)); + } else { + cachedTableIds = distinctTableIds; + } + + @Nullable Integer configuredBuckets = getPollingBuckets(); + int pollingBuckets = configuredBuckets != null ? configuredBuckets : DEFAULT_POLLING_BUCKETS; + PCollection pollingTableIds = + cachedTableIds.apply( + "ReshufflePollingBuckets", + Reshuffle.viaRandomKey().withNumBuckets(pollingBuckets)); + + PCollection> specs = + pollingTableIds + .apply("PollTableMetadata", ParDo.of(new CatalogPollingDoFn(getCatalogConfig()))) + .setCoder( + KvCoder.of( + StringUtf8Coder.of(), NullableCoder.of(SerializableTableSpec.getCoder()))); + + if (isStreaming) { + return specs.apply( + "ApplyStreamingTrigger", + Window.>into(new GlobalWindows()) + .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1))) + .discardingFiredPanes()); + } + return specs; + } + + @Override + public void populateDisplayData(DisplayData.Builder builder) { + super.populateDisplayData(builder); + builder.addIfNotNull( + DisplayData.item("maximumCacheSize", getMaximumCacheSize()) + .withLabel("Maximum Cache Size")); + builder.addIfNotNull( + DisplayData.item("refreshInterval", getRefreshInterval()) + .withLabel("Table Metadata Refresh Interval")); + builder.addIfNotNull( + DisplayData.item("pollingBuckets", getPollingBuckets()) + .withLabel("Catalog Polling Buckets")); + } + + static class ExtractTableIdsDoFn extends DoFn { + private final DynamicDestinations dynamicDestinations; + + ExtractTableIdsDoFn(DynamicDestinations dynamicDestinations) { + this.dynamicDestinations = dynamicDestinations; + } + + @ProcessElement + public void processElement( + @Element Row element, + BoundedWindow window, + PaneInfo paneInfo, + @Timestamp Instant timestamp, + OutputReceiver out) { + String tableIdentifier = + dynamicDestinations.getTableStringIdentifier( + ValueInSingleWindow.of(element, timestamp, window, paneInfo)); + if (tableIdentifier != null && !tableIdentifier.trim().isEmpty()) { + out.output(tableIdentifier.trim()); + } + } + } + + static class CatalogPollingDoFn + extends DoFn> { + private static final Logger LOG = LoggerFactory.getLogger(CatalogPollingDoFn.class); + private static final Counter TABLES_POLLED_COUNTER = + Metrics.counter(TableMetadataDriver.class, "tablesPolled"); + private static final Counter TABLES_SKIPPED_MISSING_COUNTER = + Metrics.counter(TableMetadataDriver.class, "tablesSkippedMissing"); + private 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; + + CatalogPollingDoFn(IcebergCatalogConfig catalogConfig) { + this.catalogConfig = catalogConfig; + } + + @ProcessElement + public void processElement( + @Element String tableIdString, + OutputReceiver> out) { + TableIdentifier tableId; + try { + tableId = IcebergUtils.parseTableIdentifier(tableIdString); + } catch (IllegalArgumentException e) { + LOG.warn( + "Failed to parse table identifier '{}'. Emitting empty metadata signal for side-input view.", + tableIdString, + e); + TABLES_PARSE_FAILED_COUNTER.inc(); + out.output(KV.of(tableIdString, null)); + return; + } + + Table table; + try { + table = catalogConfig.catalog().loadTable(tableId); + } catch (NoSuchTableException e) { + LOG.info( + "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; + try { + spec = SerializableTableSpec.fromTable(tableIdString, table); + } catch (IllegalArgumentException e) { + LOG.warn( + "Failed to create SerializableTableSpec for table '{}'. Emitting empty metadata signal for side-input view.", + tableIdString, + e); + TABLES_SPEC_CREATION_FAILED_COUNTER.inc(); + out.output(KV.of(tableIdString, null)); + return; + } + TABLES_POLLED_COUNTER.inc(); + out.output(KV.of(tableIdString, spec)); + } + } + + static class AccumulateTableMetadataMapDoFn + extends DoFn< + 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 = + StateSpecs.map(StringUtf8Coder.of(), SerializableTableSpec.getCoder()); + + @StateId("lastSeen") + private final StateSpec> lastSeenStateSpec = + StateSpecs.map(StringUtf8Coder.of(), VarLongCoder.of()); + + private final Duration refreshInterval; + private final Duration cacheTtl; + private final Clock clock; + + AccumulateTableMetadataMapDoFn( + Duration refreshInterval, @Nullable Duration cacheTtl, @Nullable 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, + @StateId("tableCache") MapState cacheState, + @StateId("lastSeen") MapState lastSeenState, + OutputReceiver> out) { + long now = clock.currentTimeMillis(); + KV kv = element.getValue(); + String tableId = kv.getKey(); + @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); + } + + // Populate a local map with the state for constant-time lookups during + // cache entry expiration. + Map lastSeenMap = new HashMap<>(); + for (Map.Entry entry : lastSeenState.entries().read()) { + lastSeenMap.put(entry.getKey(), entry.getValue()); + } + + long expirationCutoff = now - cacheTtl.getMillis(); + List expiredTables = new ArrayList<>(); + Map mapSnapshot = new HashMap<>(); + + for (Map.Entry entry : cacheState.entries().read()) { + 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)); + } + } + + 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( + 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); + for (Map.Entry entry : right.entrySet()) { + String tableId = entry.getKey(); + SerializableTableSpec rightSpec = entry.getValue(); + SerializableTableSpec leftSpec = merged.get(tableId); + if (leftSpec == null || isNewer(rightSpec, leftSpec)) { + merged.put(tableId, rightSpec); + } + } + 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/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 new file mode 100644 index 000000000000..c46b155921b1 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriverTest.java @@ -0,0 +1,1501 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; +import 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; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestStream; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.display.DisplayData; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TimestampedValue; +import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.CatalogProperties; +import org.apache.iceberg.CatalogUtil; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionKey; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Types; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +public class TableMetadataDriverTest implements Serializable { + + @Rule public transient TestPipeline pipeline = TestPipeline.create(); + @Rule public transient TemporaryFolder tempFolder = new TemporaryFolder(); + + private String warehouseLocation; + private IcebergCatalogConfig catalogConfig; + + private static final Schema BEAM_SCHEMA = + Schema.builder() + .addInt64Field("id") + .addStringField("data") + .addNullableStringField("dest") + .build(); + + private static final org.apache.iceberg.Schema ICEBERG_SCHEMA = + IcebergUtils.beamSchemaToIcebergSchema( + Schema.builder().addInt64Field("id").addStringField("data").build()); + + private static final TableIdentifier TABLE_ID = TableIdentifier.of("default", "table"); + + private static final DynamicDestinations SINGLE_TABLE_DYNAMIC_DESTINATIONS = + DynamicDestinations.singleTable(TABLE_ID, BEAM_SCHEMA); + + private static final DynamicDestinations DYNAMIC_DESTINATIONS = + new DynamicDestinations() { + @Override + public Schema getDataSchema() { + return BEAM_SCHEMA; + } + + @Override + public Row getData(Row element) { + return element; + } + + @Override + public IcebergDestination instantiateDestination(String destination) { + return IcebergDestination.builder() + .setTableIdentifier(IcebergUtils.parseTableIdentifier(destination)) + .build(); + } + + @Override + public String getTableStringIdentifier(ValueInSingleWindow element) { + return element.getValue().getString("dest"); + } + }; + + @Before + public void setUp() throws Exception { + warehouseLocation = "file:" + tempFolder.newFolder().getAbsolutePath(); + catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogName("hadoop") + .setCatalogProperties(ImmutableMap.of("type", "hadoop", "warehouse", warehouseLocation)) + .build(); + } + + private Catalog getCatalog() { + return CatalogUtil.loadCatalog( + CatalogUtil.ICEBERG_CATALOG_HADOOP, + "hadoop", + ImmutableMap.of(CatalogProperties.WAREHOUSE_LOCATION, warehouseLocation), + new Configuration()); + } + + @Test + public void testSingleTableExtractionAndSpecOutput() { + Table realTable = getCatalog().createTable(TABLE_ID, ICEBERG_SCHEMA); + + List rows = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + rows.add( + Row.withSchema(BEAM_SCHEMA) + .withFieldValue("id", (long) i) + .withFieldValue("data", "val_" + i) + .withFieldValue("dest", null) + .build()); + } + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .build()); + + String expectedTableIdString = IcebergUtils.tableIdentifierToString(TABLE_ID); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = + ImmutableList.copyOf(elements); + assertEquals(1, list.size()); + KV kv = list.get(0); + assertEquals(expectedTableIdString, kv.getKey()); + SerializableTableSpec spec = kv.getValue(); + assertNotNull(spec); + assertEquals(realTable.name(), spec.getName()); + assertEquals(realTable.location(), spec.getLocation()); + assertEquals(realTable.schema().asStruct(), spec.getSchema().asStruct()); + assertEquals(realTable.spec(), spec.getPartitionSpec()); + assertNotNull(spec.getFileIO()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testMultipleDynamicDestinationsExtraction() { + Catalog catalog = getCatalog(); + TableIdentifier tableA = TableIdentifier.of("default", "table_a"); + TableIdentifier tableB = TableIdentifier.of("default", "table_b"); + TableIdentifier tableC = TableIdentifier.of("default", "table_c"); + + catalog.createTable(tableA, ICEBERG_SCHEMA); + catalog.createTable(tableB, ICEBERG_SCHEMA); + catalog.createTable(tableC, ICEBERG_SCHEMA); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.table_a").build(), + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.table_b").build(), + Row.withSchema(BEAM_SCHEMA).addValues(3L, "v3", "default.table_c").build(), + Row.withSchema(BEAM_SCHEMA).addValues(4L, "v4", "default.table_a").build(), + Row.withSchema(BEAM_SCHEMA).addValues(5L, "v5", "default.table_b").build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = + ImmutableList.copyOf(elements); + assertEquals(3, list.size()); + Map map = + list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); + assertTrue(map.containsKey("default.table_a")); + assertTrue(map.containsKey("default.table_b")); + assertTrue(map.containsKey("default.table_c")); + return null; + }); + + pipeline.run(); + } + + @Test + public void testWindowedDeduplication() { + Catalog catalog = getCatalog(); + TableIdentifier table1 = TableIdentifier.of("default", "t1"); + TableIdentifier table2 = TableIdentifier.of("default", "t2"); + + catalog.createTable(table1, ICEBERG_SCHEMA); + catalog.createTable(table2, ICEBERG_SCHEMA); + + List rows = new ArrayList<>(); + for (int i = 0; i < 100; i++) { + String dest = (i % 2 == 0) ? "default.t1" : "default.t2"; + rows.add(Row.withSchema(BEAM_SCHEMA).addValues((long) i, "val_" + i, dest).build()); + } + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = + ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + Map map = + list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); + assertTrue(map.containsKey("default.t1")); + assertTrue(map.containsKey("default.t2")); + return null; + }); + + pipeline.run(); + } + + @Test + public void testUnboundedGlobalWindowStreamingDeduplication() { + Catalog catalog = getCatalog(); + TableIdentifier table1 = TableIdentifier.of("default", "stream_t1"); + TableIdentifier table2 = TableIdentifier.of("default", "stream_t2"); + + catalog.createTable(table1, ICEBERG_SCHEMA); + catalog.createTable(table2, ICEBERG_SCHEMA); + + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.stream_t1").build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.stream_t2").build(); + Row row3 = Row.withSchema(BEAM_SCHEMA).addValues(3L, "v3", "default.stream_t1").build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(row1) + .addElements(row2) + .addElements(row3) + .advanceProcessingTime(Duration.standardSeconds(5)) + .advanceWatermarkToInfinity(); + + PCollection input = pipeline.apply("StreamInput", stream); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(2)) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = + ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + Map map = + list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); + assertTrue(map.containsKey("default.stream_t1")); + assertTrue(map.containsKey("default.stream_t2")); + return null; + }); + + pipeline.run(); + } + + @Test + public void testMetadataRefreshedAcrossIntervals() { + Catalog catalog = getCatalog(); + TableIdentifier tableId = TableIdentifier.of("default", "evolving_table"); + catalog.createTable(tableId, ICEBERG_SCHEMA); + + Row row1 = + Row.withSchema(BEAM_SCHEMA).addValues(1L, "initial_data", "default.evolving_table").build(); + Row row2 = + Row.withSchema(BEAM_SCHEMA) + .addValues(2L, "trigger_update", "default.evolving_table") + .build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(row1) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(row2) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceWatermarkToInfinity(); + + PCollection input = + pipeline + .apply("StreamInput", stream) + .apply( + "EvolveSchemaOnTriggerRow", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(@Element Row row, OutputReceiver out) { + if ("trigger_update".equals(row.getString("data"))) { + Table table = + catalogConfig + .catalog() + .loadTable( + IcebergUtils.parseTableIdentifier("default.evolving_table")); + table + .updateSchema() + .addColumn("new_col", Types.StringType.get()) + .commit(); + } + out.output(row); + } + })) + .setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(2)) + .build()); + + // Downstream consumer transform verifying that updated metadata is received + PCollection consumerReceivedSchemas = + specs.apply( + "ConsumerTransform", + ParDo.of( + new DoFn, String>() { + @ProcessElement + public void processElement( + @Element KV element, + OutputReceiver out) { + SerializableTableSpec spec = checkNotNull(element.getValue()); + boolean hasNewCol = spec.getSchema().findField("new_col") != null; + out.output(hasNewCol ? "UPDATED_SCHEMA" : "INITIAL_SCHEMA"); + } + })); + + PAssert.that(consumerReceivedSchemas).containsInAnyOrder("INITIAL_SCHEMA", "UPDATED_SCHEMA"); + + pipeline.run(); + } + + @Test + public void testMetadataRefreshedAcrossIntervalsAsSideInput() { + Catalog catalog = getCatalog(); + TableIdentifier tableId = TableIdentifier.of("default", "evolving_side_input_table"); + catalog.createTable(tableId, ICEBERG_SCHEMA); + + String tableIdStr = "default.evolving_side_input_table"; + Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1L, "initial_data", tableIdStr).build(); + Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2L, "trigger_update", tableIdStr).build(); + Row row3 = Row.withSchema(BEAM_SCHEMA).addValues(3L, "post_update_data", tableIdStr).build(); + + TestStream stream = + TestStream.create(RowCoder.of(BEAM_SCHEMA)) + .advanceWatermarkTo(new Instant(0)) + .addElements(row1) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(row2) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(row3) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceWatermarkToInfinity(); + + PCollection input = + pipeline + .apply("StreamInput", stream) + .apply( + "EvolveSchemaOnTriggerRow", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(@Element Row row, OutputReceiver out) { + if ("trigger_update".equals(row.getString("data"))) { + Table table = + catalogConfig + .catalog() + .loadTable( + IcebergUtils.parseTableIdentifier( + "default.evolving_side_input_table")); + table + .updateSchema() + .addColumn("new_col", Types.StringType.get()) + .commit(); + } + out.output(row); + } + })) + .setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(2)) + .build() + .asView()); + + 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.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(2)) + .build() + .asView()); + + 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 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.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(2)) + .build() + .asView()); + + 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.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(2)) + .build() + .asView()); + + 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); + 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 -> { + 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; + }); + + pipeline.run(); + } + + @Test + public void testMaximumCacheSizeCap() { + Catalog catalog = getCatalog(); + for (int i = 1; i <= 6; i++) { + catalog.createTable(TableIdentifier.of("default", "cap_table_" + i), ICEBERG_SCHEMA); + } + + List rows = new ArrayList<>(); + for (int i = 1; i <= 6; i++) { + rows.add( + Row.withSchema(BEAM_SCHEMA) + .addValues((long) i, "v_" + i, "default.cap_table_" + i) + .build()); + } + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + int maxCacheSize = 3; + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setMaximumCacheSize(maxCacheSize) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = + ImmutableList.copyOf(elements); + assertEquals(maxCacheSize, list.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testUncappedByDefault() { + Catalog catalog = getCatalog(); + for (int i = 1; i <= 10; i++) { + catalog.createTable(TableIdentifier.of("default", "uncapped_table_" + i), ICEBERG_SCHEMA); + } + + List rows = new ArrayList<>(); + for (int i = 1; i <= 10; i++) { + rows.add( + Row.withSchema(BEAM_SCHEMA) + .addValues((long) i, "v_" + i, "default.uncapped_table_" + i) + .build()); + } + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + // Without setting maximumCacheSize, all 10 distinct tables are emitted + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = + ImmutableList.copyOf(elements); + assertEquals(10, list.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testNonExistentTableIsSkippedWithoutFailingBundle() { + Catalog catalog = getCatalog(); + TableIdentifier validTable = TableIdentifier.of("default", "existing_table"); + catalog.createTable(validTable, ICEBERG_SCHEMA); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.existing_table").build(), + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.non_existent_table").build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + // Both existing and missing table entries are emitted; missing table has null spec + PAssert.that(specs) + .satisfies( + elements -> { + 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; + }); + + pipeline.run(); + } + + @Test + public void testFiltersNullAndBlankTableIdentifiers() { + TableIdentifier validTableId = TableIdentifier.of("default", "valid_dest_table"); + getCatalog().createTable(validTableId, ICEBERG_SCHEMA); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", null).build(), + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "").build(), + Row.withSchema(BEAM_SCHEMA).addValues(3L, "v3", " ").build(), + Row.withSchema(BEAM_SCHEMA).addValues(4L, "v4", "default.valid_dest_table").build(), + Row.withSchema(BEAM_SCHEMA) + .addValues(5L, "v5", " default.valid_dest_table ") + .build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = + ImmutableList.copyOf(elements); + assertEquals(1, list.size()); + assertEquals("default.valid_dest_table", list.get(0).getKey()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testInvalidMaximumCacheSizeThrowsException() { + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setMaximumCacheSize(0) + .build()); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setMaximumCacheSize(-5) + .build()); + } + + @Test + public void testInvalidRefreshIntervalThrowsException() { + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.ZERO) + .build()); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setRefreshInterval(Duration.standardSeconds(-5)) + .build()); + } + + @Test + public void testInvalidPollingBucketsThrowsException() { + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setPollingBuckets(0) + .build()); + + assertThrows( + IllegalArgumentException.class, + () -> + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setPollingBuckets(-2) + .build()); + } + + @Test + public void testConfigurablePollingBuckets() { + Catalog catalog = getCatalog(); + TableIdentifier table1 = TableIdentifier.of("default", "bucket_t1"); + TableIdentifier table2 = TableIdentifier.of("default", "bucket_t2"); + catalog.createTable(table1, ICEBERG_SCHEMA); + catalog.createTable(table2, ICEBERG_SCHEMA); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.bucket_t1").build(), + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.bucket_t2").build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setPollingBuckets(2) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = + ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + Map map = + list.stream().collect(ImmutableMap.toImmutableMap(KV::getKey, KV::getValue)); + assertTrue(map.containsKey("default.bucket_t1")); + assertTrue(map.containsKey("default.bucket_t2")); + return null; + }); + + pipeline.run(); + } + + @Test + public void testWindowPreservation() { + Catalog catalog = getCatalog(); + TableIdentifier tableW1 = TableIdentifier.of("default", "table_w1"); + TableIdentifier tableW2 = TableIdentifier.of("default", "table_w2"); + + catalog.createTable(tableW1, ICEBERG_SCHEMA); + catalog.createTable(tableW2, ICEBERG_SCHEMA); + + Instant t1 = new Instant(1000); + Instant t2 = new Instant(70000); + + PCollection input = + pipeline + .apply( + Create.timestamped( + TimestampedValue.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", "default.table_w1").build(), + t1), + TimestampedValue.of( + Row.withSchema(BEAM_SCHEMA).addValues(2L, "v2", "default.table_w2").build(), + t2))) + .setCoder(RowCoder.of(BEAM_SCHEMA)) + .apply(Window.into(FixedWindows.of(Duration.standardMinutes(1)))); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs) + .satisfies( + elements -> { + List> list = + ImmutableList.copyOf(elements); + assertEquals(2, list.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testEmptyInputProducesEmptyOutput() { + getCatalog().createTable(TABLE_ID, ICEBERG_SCHEMA); + + PCollection input = pipeline.apply(Create.empty(RowCoder.of(BEAM_SCHEMA))); + + PCollection> specs = + input.apply( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .build()); + + PAssert.that(specs).empty(); + + pipeline.run(); + } + + @Test + public void testDisplayData() { + TableMetadataDriver driver = + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(SINGLE_TABLE_DYNAMIC_DESTINATIONS) + .setMaximumCacheSize(42) + .setRefreshInterval(Duration.standardMinutes(10)) + .setPollingBuckets(3) + .build(); + + DisplayData displayData = DisplayData.from(driver); + Map items = displayData.asMap(); + + assertNotNull(displayData); + boolean hasCacheSize = false; + boolean hasRefreshInterval = false; + boolean hasPollingBuckets = false; + for (DisplayData.Item item : items.values()) { + if ("maximumCacheSize".equals(item.getKey())) { + assertEquals(42L, item.getValue()); + hasCacheSize = true; + } + if ("refreshInterval".equals(item.getKey())) { + hasRefreshInterval = true; + } + if ("pollingBuckets".equals(item.getKey())) { + assertEquals(3L, item.getValue()); + hasPollingBuckets = true; + } + } + assertTrue(hasCacheSize); + assertTrue(hasRefreshInterval); + assertTrue(hasPollingBuckets); + } + + @Test + public void testViewAsMapIntegration() { + PartitionSpec partitionSpec = PartitionSpec.builderFor(ICEBERG_SCHEMA).identity("data").build(); + getCatalog().createTable(TABLE_ID, ICEBERG_SCHEMA, partitionSpec); + + List rows = + ImmutableList.of( + Row.withSchema(BEAM_SCHEMA).addValues(10L, "partition_val_a", null).build(), + Row.withSchema(BEAM_SCHEMA).addValues(20L, "partition_val_b", null).build()); + + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", + TableMetadataDriver.asView(catalogConfig, SINGLE_TABLE_DYNAMIC_DESTINATIONS)); + + String expectedTableIdString = IcebergUtils.tableIdentifierToString(TABLE_ID); + + PCollection writtenFiles = + input.apply( + "WriteWithSideInputTable", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element Row row, OutputReceiver out, ProcessContext c) + throws Exception { + Map viewMap = c.sideInput(metadataView); + SerializableTableSpec spec = viewMap.get(expectedTableIdString); + assertNotNull(spec); + + SideInputTable sideInputTable = new SideInputTable(spec); + PartitionKey partitionKey = + new PartitionKey(sideInputTable.spec(), sideInputTable.schema()); + Record record = GenericRecord.create(sideInputTable.schema()); + record.setField("id", row.getInt64("id")); + record.setField("data", row.getString("data")); + partitionKey.partition(record); + + RecordWriter writer = + new RecordWriter( + sideInputTable, + FileFormat.PARQUET, + "side_input_test_file_" + row.getInt64("id"), + partitionKey, + ImmutableMap.of()); + writer.write(record); + writer.close(); + + out.output(writer.getDataFile().path().toString()); + } + }) + .withSideInputs(metadataView)); + + PAssert.that(writtenFiles) + .satisfies( + files -> { + List paths = ImmutableList.copyOf(files); + assertEquals(2, paths.size()); + return null; + }); + + pipeline.run(); + } + + @Test + public void testMapMergerFnCommutativeAndTimestampAware() { + TableIdentifier tableIdA = TableIdentifier.of("default", "merge_table_a"); + Table realTableA = getCatalog().createTable(tableIdA, ICEBERG_SCHEMA); + SerializableTableSpec baseA = SerializableTableSpec.fromTable(tableIdA, realTableA); + SerializableTableSpec specAOld = baseA.toBuilder().setLastUpdatedMillis(1000L).build(); + SerializableTableSpec specANew = baseA.toBuilder().setLastUpdatedMillis(2000L).build(); + + TableIdentifier tableIdB = TableIdentifier.of("default", "merge_table_b"); + Table realTableB = getCatalog().createTable(tableIdB, ICEBERG_SCHEMA); + SerializableTableSpec baseB = SerializableTableSpec.fromTable(tableIdB, realTableB); + SerializableTableSpec specB = baseB.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 baseSpec = SerializableTableSpec.fromTable(tableId, realTable); + SerializableTableSpec specSchema0 = + baseSpec.toBuilder().setLastUpdatedMillis(1000L).setSchemaId(0).build(); + SerializableTableSpec specSchema1 = + baseSpec.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); + } + + 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(20000L); + } + out.output(row); + } + })) + .setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", + TableMetadataDriver.asView( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(refreshInterval) + .build(), + 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(); + } + + @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 testAsViewWithDriverInstance() { + TableIdentifier tableId = TableIdentifier.of("default", "as_view_driver_test"); + getCatalog().createTable(tableId, ICEBERG_SCHEMA); + String tableStr = IcebergUtils.tableIdentifierToString(tableId); + + List rows = + ImmutableList.of(Row.withSchema(BEAM_SCHEMA).addValues(1L, "v1", tableStr).build()); + PCollection input = pipeline.apply(Create.of(rows)).setCoder(RowCoder.of(BEAM_SCHEMA)); + + TableMetadataDriver driver = + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .build(); + + PCollectionView> metadataView = + input.apply("CreateMetadataView", driver.asView()); + + PCollection hasTable = + input.apply( + "CheckSideInput", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(OutputReceiver out, ProcessContext c) { + out.output(c.sideInput(metadataView).containsKey(tableStr)); + } + }) + .withSideInputs(metadataView)); + + PAssert.that(hasTable).containsInAnyOrder(true); + pipeline.run(); + } + + @Test + public void testStreamingMissingTableSignalDoesNotEvictBeforeExpirationCutoff() { + TableIdentifier tableId = TableIdentifier.of("default", "dropped_table"); + getCatalog().createTable(tableId, ICEBERG_SCHEMA); + String tableStr = IcebergUtils.tableIdentifierToString(tableId); + + Duration refreshInterval = Duration.standardSeconds(2); + ControllableTestClock.setTime(1000L); + ControllableTestClock testClock = new ControllableTestClock(); + + 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(); + Row rowTriggerEvict = + Row.withSchema(BEAM_SCHEMA).addValues(4L, "trigger_evict", tableStr).build(); + Row rowAfterCutoff = + Row.withSchema(BEAM_SCHEMA).addValues(5L, "after_cutoff", 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)) + .addElements(rowPostDrop) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(rowTriggerEvict) + .advanceProcessingTime(Duration.standardSeconds(3)) + .addElements(rowAfterCutoff) + .advanceProcessingTime(Duration.standardSeconds(3)) + .advanceWatermarkToInfinity(); + + PCollection input = + pipeline + .apply("StreamInput", stream) + .apply( + "ControlClockAndCatalogOnTriggerRows", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(@Element Row row, OutputReceiver out) { + String data = row.getString("data"); + if ("trigger_drop".equals(data)) { + ControllableTestClock.setTime(2000L); + catalogConfig + .catalog() + .dropTable( + IcebergUtils.parseTableIdentifier("default.dropped_table")); + } else if ("trigger_evict".equals(data)) { + ControllableTestClock.setTime(20000L); + } + out.output(row); + } + })) + .setCoder(RowCoder.of(BEAM_SCHEMA)); + + PCollectionView> metadataView = + input.apply( + "CreateMetadataView", + TableMetadataDriver.asView( + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(DYNAMIC_DESTINATIONS) + .setRefreshInterval(refreshInterval) + .build(), + 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 ("trigger_drop".equals(data) || "trigger_evict".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=true", "after_cutoff:hasTable=false"); + + pipeline.run(); + } +}