Expose Public API for Iceberg Side Input Cache

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java

@@ -29,8 +29,10 @@
 import org.apache.beam.sdk.options.StreamingOptions;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.display.DisplayData;
 import org.apache.beam.sdk.values.PBegin;
 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.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicates;
@@ -390,6 +392,7 @@ public static WriteRows writeRows(IcebergCatalogConfig catalog) {
         .setCatalogConfig(catalog)
         .setDistributionMode(DistributionMode.NONE)
         .setAutoSharding(false)
+        .setUsingSideInputTableCache(false)
         .build();
   }
 
@@ -416,6 +419,14 @@ public abstract static class WriteRows extends PTransform<PCollection<Row>, Iceb
 
     abstract @Nullable List<String> getSortFields();
 
+    abstract boolean getUsingSideInputTableCache();
+
+    abstract @Nullable Integer getMaximumCacheSize();
+
+    abstract @Nullable Duration getTableRefreshInterval();
+
+    abstract @Nullable Integer getPollingBuckets();
+
     abstract Builder toBuilder();
 
     @AutoValue.Builder
@@ -440,6 +451,14 @@ abstract static class Builder {
 
       abstract Builder setSortFields(List<String> sortFields);
 
+      abstract Builder setUsingSideInputTableCache(boolean usingSideInputTableCache);
+
+      abstract Builder setMaximumCacheSize(@Nullable Integer maximumCacheSize);
+
+      abstract Builder setTableRefreshInterval(@Nullable Duration refreshInterval);
+
+      abstract Builder setPollingBuckets(@Nullable Integer pollingBuckets);
+
       abstract WriteRows build();
     }
 
@@ -476,11 +495,11 @@ public WriteRows withDirectWriteByteLimit(Integer directWriteByteLimit) {
      * Defines distribution of write data. Supported distributions:
      *
      * <ol>
-     *   <li>{@link DistributionMode.NONE}: don't shuffle rows (default)
-     *   <li>{@link DistributionMode.HASH}: shuffle rows by partition key before writing data
+     *   <li>{@link DistributionMode#NONE}: don't shuffle rows (default)
+     *   <li>{@link DistributionMode#HASH}: shuffle rows by partition key before writing data
      * </ol>
      *
-     * {@link DistributionMode.RANGE} is not supported yet
+     * {@link DistributionMode#RANGE} is not supported yet
      */
     public WriteRows withDistributionMode(DistributionMode mode) {
       return toBuilder().setDistributionMode(mode).build();
@@ -524,6 +543,68 @@ public WriteRows withSortOrder(List<String> sortFields) {
       return toBuilder().setSortFields(sortFields).build();
     }
 
+    /**
+     * Enables expirable side-input caching of Iceberg table metadata across workers.
+     *
+     * <p>When enabled, a driver transform periodically polls the Iceberg catalog and broadcasts
+     * lightweight table specifications as a side input. Workers construct in-memory {@link Table}
+     * representations without issuing remote catalog RPCs, drastically reducing catalog load.
+     */
+    public WriteRows withSideInputTableCache() {
+      return toBuilder().setUsingSideInputTableCache(true).build();
+    }
+
+    /**
+     * Sets the maximum number of distinct table metadata specifications to broadcast in the
+     * side-input cache. Any tables exceeding this limit fall back to worker-local catalog loading.
+     *
+     * <p><b>Note:</b> This option is only supported for bounded (batch) pipelines. Calling this on
+     * an unbounded streaming pipeline will throw an exception at pipeline construction.
+     */
+    public WriteRows withMaximumCacheSize(int maximumCacheSize) {
+      Preconditions.checkArgument(maximumCacheSize > 0, "maximumCacheSize must be greater than 0");
+      return toBuilder().setMaximumCacheSize(maximumCacheSize).build();
+    }
+
+    /**
+     * Sets the interval at which table metadata is refreshed from the Iceberg catalog.
+     *
+     * <p>Applicable for unbounded streaming pipelines. Defaults to 5 minutes.
+     */
+    public WriteRows withTableRefreshInterval(Duration refreshInterval) {
+      Preconditions.checkNotNull(refreshInterval, "refreshInterval must not be null");
+      Preconditions.checkArgument(
+          refreshInterval.isLongerThan(Duration.ZERO), "refreshInterval must be greater than 0");
+      return toBuilder().setTableRefreshInterval(refreshInterval).build();
+    }
+
+    /**
+     * Sets the number of parallel buckets/workers used to query the Iceberg catalog during
+     * refreshes. Defaults to 1 to serialize catalog queries and protect catalogs from connection
+     * spikes.
+     */
+    public WriteRows withPollingBuckets(int pollingBuckets) {
+      Preconditions.checkArgument(pollingBuckets > 0, "pollingBuckets must be greater than 0");
+      return toBuilder().setPollingBuckets(pollingBuckets).build();
+    }
+
+    @Override
+    public void populateDisplayData(DisplayData.Builder builder) {
+      super.populateDisplayData(builder);
+      builder.add(
+          DisplayData.item("usingSideInputTableCache", getUsingSideInputTableCache())
+              .withLabel("Using Side-Input Table Cache"));
+      builder.addIfNotNull(
+          DisplayData.item("maximumCacheSize", getMaximumCacheSize())
+              .withLabel("Maximum Cache Size"));
+      builder.addIfNotNull(
+          DisplayData.item("tableRefreshInterval", getTableRefreshInterval())
+              .withLabel("Table Refresh Interval"));
+      builder.addIfNotNull(
+          DisplayData.item("pollingBuckets", getPollingBuckets())
+              .withLabel("Catalog Polling Buckets"));
+    }
+
     @Override
     public IcebergWriteResult expand(PCollection<Row> input) {
       List<?> allToArgs = Arrays.asList(getTableIdentifier(), getDynamicDestinations());
@@ -549,6 +630,35 @@ public IcebergWriteResult expand(PCollection<Row> input) {
             "Must only provide direct write limit for unbounded pipelines.");
       }
 
+      boolean hasSideInputOptions =
+          getMaximumCacheSize() != null
+              || getTableRefreshInterval() != null
+              || getPollingBuckets() != null;
+      Preconditions.checkArgument(
+          getUsingSideInputTableCache() || !hasSideInputOptions,
+          "Cannot specify side-input cache sub-options (maximumCacheSize, "
+              + "tableRefreshInterval, pollingBuckets) without enabling side-input table cache via withSideInputTableCache().");
+
+      PCollectionView<Map<String, SerializableTableSpec>> metadataView = null;
+      if (getUsingSideInputTableCache()) {
+        TableMetadataDriver.Builder driverBuilder =
+            TableMetadataDriver.builder()
+                .setCatalogConfig(getCatalogConfig())
+                .setDynamicDestinations(destinations);
+
+        if (getMaximumCacheSize() != null) {
+          driverBuilder.setMaximumCacheSize(getMaximumCacheSize());
+        }
+        if (getTableRefreshInterval() != null) {
+          driverBuilder.setRefreshInterval(getTableRefreshInterval());
+        }
+        if (getPollingBuckets() != null) {
+          driverBuilder.setPollingBuckets(getPollingBuckets());
+        }
+
+        metadataView = input.apply("GenerateTableMetadataView", driverBuilder.build().asView());
+      }
+
       switch (getDistributionMode()) {
         case NONE:
           Preconditions.checkArgument(
@@ -563,20 +673,23 @@ public IcebergWriteResult expand(PCollection<Row> input) {
                       destinations,
                       getTriggeringFrequency(),
                       getDirectWriteByteLimit(),
-                      getWriteProperties()));
+                      getWriteProperties(),
+                      metadataView));
         case HASH:
           return input
               .apply(
                   "AssignDestinationAndPartition",
-                  new AssignDestinationsAndPartitions(destinations, getCatalogConfig()))
+                  new AssignDestinationsAndPartitions(
+                      destinations, getCatalogConfig(), metadataView))
               .apply(
                   "Write Rows to Partitions",
                   new WriteToPartitions(
                       getCatalogConfig(),
                       destinations,
                       getTriggeringFrequency(),
                       getAutoSharding(),
-                      getWriteProperties()));
+                      getWriteProperties(),
+                      metadataView));
         default:
           throw new UnsupportedOperationException(
               "Unsupported distribution mode: " + getDistributionMode());

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProvider.java

@@ -163,6 +163,23 @@ public static Builder builder() {
             + "'write.parquet.bloom-filter-enabled.column.<col>').")
     public abstract @Nullable Map<String, String> getWriteProperties();
 
+    @SchemaFieldDescription(
+        "Enables expirable side-input caching of Iceberg table metadata across workers to reduce catalog load.")
+    public abstract @Nullable Boolean getUsingSideInputTableCache();
+
+    @SchemaFieldDescription(
+        "For a streaming pipeline, sets the interval in seconds at which table metadata is refreshed from the catalog.")
+    public abstract @Nullable Integer getTableRefreshIntervalSeconds();
+
+    @SchemaFieldDescription(
+        "For a batch pipeline, sets the maximum number of table metadata specs to cache in memory. "
+            + "Tables exceeding this limit fall back to worker-local catalog loading.")
+    public abstract @Nullable Integer getMaximumCacheSize();
+
+    @SchemaFieldDescription(
+        "Sets the number of parallel buckets/workers used to query the Iceberg catalog during refreshes. Defaults to 1.")
+    public abstract @Nullable Integer getPollingBuckets();
+
     @AutoValue.Builder
     public abstract static class Builder {
       public abstract Builder setTable(String table);
@@ -195,6 +212,14 @@ public abstract static class Builder {
 
       public abstract Builder setWriteProperties(Map<String, String> writeProperties);
 
+      public abstract Builder setUsingSideInputTableCache(Boolean usingSideInputTableCache);
+
+      public abstract Builder setTableRefreshIntervalSeconds(Integer tableRefreshIntervalSeconds);
+
+      public abstract Builder setMaximumCacheSize(Integer maximumCacheSize);
+
+      public abstract Builder setPollingBuckets(Integer pollingBuckets);
+
       public abstract Configuration build();
     }
 
@@ -291,6 +316,38 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
         writeTransform = writeTransform.withWriteProperties(writeProperties);
       }
 
+      boolean hasSideInputOptions =
+          configuration.getTableRefreshIntervalSeconds() != null
+              || configuration.getMaximumCacheSize() != null
+              || configuration.getPollingBuckets() != null;
+
+      if (!Boolean.TRUE.equals(configuration.getUsingSideInputTableCache())
+          && hasSideInputOptions) {
+        throw new IllegalArgumentException(
+            "Cannot specify side-input cache sub-options (table_refresh_interval_seconds, "
+                + "maximum_cache_size, polling_buckets) without explicitly setting using_side_input_table_cache to true.");
+      }
+
+      boolean enableSideInputCache =
+          Boolean.TRUE.equals(configuration.getUsingSideInputTableCache());
+
+      if (enableSideInputCache) {
+        writeTransform = writeTransform.withSideInputTableCache();
+        @Nullable Integer refreshSec = configuration.getTableRefreshIntervalSeconds();
+        if (refreshSec != null) {
+          writeTransform =
+              writeTransform.withTableRefreshInterval(Duration.standardSeconds(refreshSec));
+        }
+        @Nullable Integer maxCacheSize = configuration.getMaximumCacheSize();
+        if (maxCacheSize != null) {
+          writeTransform = writeTransform.withMaximumCacheSize(maxCacheSize);
+        }
+        @Nullable Integer pollingBuckets = configuration.getPollingBuckets();
+        if (pollingBuckets != null) {
+          writeTransform = writeTransform.withPollingBuckets(pollingBuckets);
+        }
+      }
+
       // TODO: support dynamic destinations
       IcebergWriteResult result = rows.apply(writeTransform);
 

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/TableMetadataDriver.java

@@ -276,7 +276,7 @@ public PCollectionView<Map<String, SerializableTableSpec>> expand(PCollection<Ro
     Integer maxCacheSize = getMaximumCacheSize();
     if (maxCacheSize != null) {
       if (isStreaming) {
-        throw new UnsupportedOperationException(
+        throw new IllegalArgumentException(
             "maximumCacheSize is currently not supported for unbounded streaming pipelines.");
       }
       cachedTableIds = distinctTableIds.apply("CapCacheSize", Sample.any(maxCacheSize));