ahmedabu98 commented on code in PR #39883: URL: https://github.com/apache/beam/pull/39883#discussion_r3918082047
########## sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriver.java: ########## @@ -0,0 +1,296 @@ +/* + * 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.util.Map; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.transforms.Distinct; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Sample; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.transforms.display.DisplayData; +import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindows; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.transforms.windowing.Repeatedly; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.TableIdentifier; +import org.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}, loads their + * declarative metadata from the Iceberg catalog, and emits {@link KV} pairs of table identifier + * strings to {@link SerializableTableSpec}. This is intended to be used in Beam pipelines that may + * utilize a large number of workers to handle Iceberg writes, where having every worker thread + * query for table metadata results in an excessive amount of requests and a high level of + * redundancy. + * + * <p>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. + * + * <p>For unbounded streaming pipelines in {@link GlobalWindows}, an {@link AfterProcessingTime} + * trigger is automatically applied to fire deduplication and refresh table metadata at the + * configured {@code refreshInterval} (defaulting to {@link #DEFAULT_REFRESH_INTERVAL}). + */ +@Internal +@AutoValue +public abstract class TableMetadataDriver + extends PTransform<PCollection<Row>, PCollection<KV<String, SerializableTableSpec>>> { + + public static final Duration DEFAULT_REFRESH_INTERVAL = Duration.standardMinutes(5); + + public abstract IcebergCatalogConfig getCatalogConfig(); + + public abstract DynamicDestinations getDynamicDestinations(); + + public abstract @Nullable Integer getMaximumCacheSize(); + + public abstract @Nullable Duration getRefreshInterval(); + + public static Builder builder() { + return new AutoValue_TableMetadataDriver.Builder(); + } + + 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); + + 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); + } + return driver; + } + } + + /** + * Helper that applies {@link TableMetadataDriver} and creates an uncapped {@link PCollectionView} + * of {@link Map} of table identifier strings to {@link SerializableTableSpec}. + */ + public static PTransform<PCollection<Row>, PCollectionView<Map<String, SerializableTableSpec>>> + asView(IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations) { + return asView(catalogConfig, dynamicDestinations, null, null); + } + + /** + * Helper that applies {@link TableMetadataDriver} with an optional {@code maximumCacheSize} limit + * and creates a {@link PCollectionView} of {@link Map} of table identifier strings to {@link + * SerializableTableSpec}. + * + * @param catalogConfig the catalog configuration used to poll metadata. + * @param dynamicDestinations destination strategy extracting table IDs from rows. + * @param maximumCacheSize optional maximum distinct tables to poll and broadcast per window (null + * for uncapped). + */ + public static PTransform<PCollection<Row>, PCollectionView<Map<String, SerializableTableSpec>>> + asView( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Integer maximumCacheSize) { + return asView(catalogConfig, dynamicDestinations, maximumCacheSize, null); + } + + /** + * Helper that applies {@link TableMetadataDriver} with an optional {@code maximumCacheSize} limit + * and custom {@code refreshInterval}, creating a {@link PCollectionView} of {@link Map} of table + * identifier strings to {@link SerializableTableSpec}. + * + * @param catalogConfig the catalog configuration used to poll metadata. + * @param dynamicDestinations destination strategy extracting table IDs from rows. + * @param maximumCacheSize optional maximum distinct tables to poll and broadcast per window (null + * for uncapped). + * @param refreshInterval optional refresh interval for streaming global window triggers. + */ + public static PTransform<PCollection<Row>, PCollectionView<Map<String, SerializableTableSpec>>> + asView( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Integer maximumCacheSize, + @Nullable Duration refreshInterval) { + return new PTransform<PCollection<Row>, PCollectionView<Map<String, SerializableTableSpec>>>() { + @Override + public PCollectionView<Map<String, SerializableTableSpec>> expand(PCollection<Row> input) { + return input + .apply( + "GenerateTableMetadata", + TableMetadataDriver.builder() + .setCatalogConfig(catalogConfig) + .setDynamicDestinations(dynamicDestinations) + .setMaximumCacheSize(maximumCacheSize) + .setRefreshInterval(refreshInterval) + .build()) + .apply("CreateTableMetadataView", View.asMap()); + } + }; + } + + @Override + public PCollection<KV<String, SerializableTableSpec>> expand(PCollection<Row> input) { + PCollection<String> tableIds = + input + .apply("ExtractTableIds", ParDo.of(new ExtractTableIdsDoFn(getDynamicDestinations()))) + .setCoder(StringUtf8Coder.of()); + + boolean isUnboundedGlobal = + input.isBounded() == PCollection.IsBounded.UNBOUNDED + && input.getWindowingStrategy().getWindowFn() instanceof GlobalWindows; + + PCollection<String> triggeredTableIds; + if (isUnboundedGlobal) { + Duration customInterval = getRefreshInterval(); + Duration interval = + checkNotNull(customInterval != null ? customInterval : DEFAULT_REFRESH_INTERVAL); + triggeredTableIds = + tableIds.apply( + "ApplyStreamingTrigger", + Window.<String>into(new GlobalWindows()) + .triggering( + Repeatedly.forever( + AfterProcessingTime.pastFirstElementInPane().plusDelayOf(interval))) + .accumulatingFiredPanes()); + } else { + triggeredTableIds = tableIds; + } + + PCollection<String> distinctTableIds = + triggeredTableIds.apply("DistinctTableIds", Distinct.create()); Review Comment: We can leave it as a future improvement if we notice it's affecting throughput. I don't think it'll break update compatibility -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
