From acde4e4d8c005ebc60450cd504292778c92f8b71 Mon Sep 17 00:00:00 2001 From: Claire McGinty Date: Tue, 21 Jul 2026 09:14:25 -0400 Subject: [PATCH 1/2] (IcebergIO) Support SortOrder on dynamic table creation via IcebergIO --- .../sdk/io/iceberg/DynamicDestinations.java | 12 ++++- .../apache/beam/sdk/io/iceberg/IcebergIO.java | 33 +++++++++++- .../iceberg/OneTableDynamicDestinations.java | 33 ++++++++++-- .../iceberg/WritePartitionedRowsToFiles.java | 13 ++++- .../sdk/io/iceberg/IcebergIOWriteTest.java | 54 +++++++++++++++++++ 5 files changed, 138 insertions(+), 7 deletions(-) diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/DynamicDestinations.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/DynamicDestinations.java index 0185758c8aeb..ba9ec771c11d 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/DynamicDestinations.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/DynamicDestinations.java @@ -18,10 +18,12 @@ package org.apache.beam.sdk.io.iceberg; import java.io.Serializable; +import java.util.List; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.ValueInSingleWindow; import org.apache.iceberg.catalog.TableIdentifier; +import org.checkerframework.checker.nullness.qual.Nullable; public interface DynamicDestinations extends Serializable { @@ -34,6 +36,14 @@ public interface DynamicDestinations extends Serializable { String getTableStringIdentifier(ValueInSingleWindow element); static DynamicDestinations singleTable(TableIdentifier tableId, Schema inputSchema) { - return new OneTableDynamicDestinations(tableId, inputSchema); + return new OneTableDynamicDestinations(tableId, inputSchema, null, null); + } + + static DynamicDestinations singleTable( + TableIdentifier tableId, + Schema inputSchema, + @Nullable List partitionFields, + @Nullable List sortFields) { + return new OneTableDynamicDestinations(tableId, inputSchema, partitionFields, sortFields); } } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java index 04feea5037d1..9b2946d239db 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java @@ -409,6 +409,10 @@ public abstract static class WriteRows extends PTransform, Iceb abstract @Nullable Map getWriteProperties(); + abstract @Nullable List getPartitionFields(); + + abstract @Nullable List getSortFields(); + abstract Builder toBuilder(); @AutoValue.Builder @@ -429,6 +433,10 @@ abstract static class Builder { abstract Builder setWriteProperties(Map writeProperties); + abstract Builder setPartitionFields(List partitionFields); + + abstract Builder setSortFields(List sortFields); + abstract WriteRows build(); } @@ -483,6 +491,26 @@ public WriteRows withWriteProperties(Map writeProperties) { return toBuilder().setWriteProperties(writeProperties).build(); } + /** + * Defines the desired Partition Spec to be applied when the Iceberg table must be dynamically + * created, e.g. `bucket(id_field, 32)` or `day(timestamp_field)` + * + *

See: https://iceberg.apache.org/spec/#partitioning + */ + public WriteRows withPartitionSpec(List partitionFields) { + return toBuilder().setPartitionFields(partitionFields).build(); + } + + /** + * Defines the desired Sort Order to be applied when the Iceberg table must be dynamically + * created, e.g. `int_field desc` or `bucket(modulo_5, 4) asc nulls last` + * + *

See: https://iceberg.apache.org/spec/#sorting + */ + public WriteRows withSortOrder(List sortFields) { + return toBuilder().setSortFields(sortFields).build(); + } + @Override public IcebergWriteResult expand(PCollection input) { List allToArgs = Arrays.asList(getTableIdentifier(), getDynamicDestinations()); @@ -494,7 +522,10 @@ public IcebergWriteResult expand(PCollection input) { if (destinations == null) { destinations = DynamicDestinations.singleTable( - Preconditions.checkNotNull(getTableIdentifier()), input.getSchema()); + Preconditions.checkNotNull(getTableIdentifier()), + input.getSchema(), + getPartitionFields(), + getSortFields()); } // Assign destinations before re-windowing to global in WriteToDestinations because diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/OneTableDynamicDestinations.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/OneTableDynamicDestinations.java index afca43fca15c..a93995589483 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/OneTableDynamicDestinations.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/OneTableDynamicDestinations.java @@ -23,6 +23,7 @@ import java.io.IOException; import java.io.ObjectInput; import java.io.ObjectOutput; +import java.util.List; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.ValueInSingleWindow; @@ -30,6 +31,7 @@ import org.apache.iceberg.FileFormat; import org.apache.iceberg.catalog.TableIdentifier; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; +import org.checkerframework.checker.nullness.qual.Nullable; class OneTableDynamicDestinations implements DynamicDestinations, Externalizable { // TableId represented as String for serializability @@ -37,6 +39,8 @@ class OneTableDynamicDestinations implements DynamicDestinations, Externalizable private transient @MonotonicNonNull TableIdentifier tableId; private transient @MonotonicNonNull Schema rowSchema; + private transient @Nullable List partitionFields; + private transient @Nullable List sortFields; @VisibleForTesting TableIdentifier getTableIdentifier() { @@ -46,9 +50,15 @@ TableIdentifier getTableIdentifier() { return tableId; } - OneTableDynamicDestinations(TableIdentifier tableId, Schema rowSchema) { + OneTableDynamicDestinations( + TableIdentifier tableId, + Schema rowSchema, + @Nullable List partitionFields, + @Nullable List sortFields) { this.tableIdString = IcebergUtils.tableIdentifierToString(tableId); this.rowSchema = rowSchema; + this.partitionFields = partitionFields; + this.sortFields = sortFields; } @Override @@ -68,9 +78,18 @@ public String getTableStringIdentifier(ValueInSingleWindow element) { @Override public IcebergDestination instantiateDestination(String unused) { + @Nullable IcebergTableCreateConfig createConfig = null; + if (partitionFields != null || sortFields != null) { + createConfig = + IcebergTableCreateConfig.builder() + .setSchema(checkStateNotNull(rowSchema)) + .setPartitionFields(partitionFields) + .setSortFields(sortFields) + .build(); + } return IcebergDestination.builder() .setTableIdentifier(getTableIdentifier()) - .setTableCreateConfig(null) + .setTableCreateConfig(createConfig) .setFileFormat(FileFormat.PARQUET) .build(); } @@ -78,14 +97,22 @@ public IcebergDestination instantiateDestination(String unused) { // Need a public default constructor for custom serialization public OneTableDynamicDestinations() {} + @SuppressWarnings("nullness") @Override public void writeExternal(ObjectOutput out) throws IOException { out.writeUTF(checkStateNotNull(tableIdString)); + out.writeObject(rowSchema); + out.writeObject(partitionFields); + out.writeObject(sortFields); } + @SuppressWarnings("nullness") @Override - public void readExternal(ObjectInput in) throws IOException { + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { tableIdString = in.readUTF(); tableId = IcebergUtils.parseTableIdentifier(tableIdString); + rowSchema = (Schema) in.readObject(); + partitionFields = (List) in.readObject(); + sortFields = (List) in.readObject(); } } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java index 84e7d05b124a..d1a08980fa9d 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java @@ -39,6 +39,7 @@ import org.apache.iceberg.PartitionField; import org.apache.iceberg.PartitionKey; import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.SortOrder; import org.apache.iceberg.StructLike; import org.apache.iceberg.Table; import org.apache.iceberg.catalog.Catalog; @@ -189,6 +190,8 @@ private Table loadOrCreateTable( @Nullable IcebergTableCreateConfig createConfig = destination.getTableCreateConfig(); PartitionSpec partitionSpec = createConfig != null ? createConfig.getPartitionSpec() : PartitionSpec.unpartitioned(); + SortOrder sortOrder = + createConfig != null ? createConfig.getSortOrder() : SortOrder.unsorted(); Map tableProperties = createConfig != null && createConfig.getTableProperties() != null ? createConfig.getTableProperties() @@ -217,13 +220,19 @@ private Table loadOrCreateTable( org.apache.iceberg.Schema tableSchema = IcebergUtils.beamSchemaToIcebergSchema(dataSchema); try { Table table = - catalog.createTable(identifier, tableSchema, partitionSpec, tableProperties); + catalog + .buildTable(identifier, tableSchema) + .withPartitionSpec(partitionSpec) + .withSortOrder(sortOrder) + .withProperties(tableProperties) + .create(); LOG.info( "Created Iceberg table '{}' with schema: {}\n" - + ", partition spec: {}, table properties: {}", + + ", partition spec: {}, sort order: {}, table properties: {}", identifier, tableSchema, partitionSpec, + sortOrder, tableProperties); return table; } catch (AlreadyExistsException ignored) { diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java index 6580f4eeaf62..5500edbadbc0 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java @@ -76,7 +76,10 @@ import org.apache.iceberg.DataFile; import org.apache.iceberg.DistributionMode; import org.apache.iceberg.FileFormat; +import org.apache.iceberg.NullOrder; import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.SortDirection; +import org.apache.iceberg.SortOrder; import org.apache.iceberg.Table; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.SupportsNamespaces; @@ -775,4 +778,55 @@ public void process(@Element KV>> sums) { .getCommitted(); assertEquals(5L, numWaves); } + + @Test + public void testCreateTableWithPartitionSpecAndSortOrder() { + TableIdentifier tableId = + TableIdentifier.of( + "default", "sorted_partitioned_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + + Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(TestFixtures.SCHEMA); + + Map catalogProps = + ImmutableMap.builder() + .put("type", CatalogUtil.ICEBERG_CATALOG_TYPE_HADOOP) + .put("warehouse", warehouse.location) + .build(); + + IcebergCatalogConfig catalog = + IcebergCatalogConfig.builder() + .setCatalogName("name") + .setCatalogProperties(catalogProps) + .build(); + + testPipeline + .apply("Records To Add", Create.of(TestFixtures.asRows(TestFixtures.FILE1SNAPSHOT1))) + .setRowSchema(beamSchema) + .apply( + "Append To Table", + writeTransform(catalog, tableId) + .withPartitionSpec(ImmutableList.of("bucket(id, 2)")) + .withSortOrder(ImmutableList.of("data asc nulls first"))); + + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(tableId); + + // verify partition spec + PartitionSpec spec = table.spec(); + assertTrue(spec.isPartitioned()); + assertEquals(1, spec.fields().size()); + assertEquals("id_bucket", spec.fields().get(0).name()); + + // verify sort order + SortOrder sortOrder = table.sortOrder(); + assertTrue(sortOrder.isSorted()); + assertEquals(1, sortOrder.fields().size()); + assertEquals(SortDirection.ASC, sortOrder.fields().get(0).direction()); + assertEquals(NullOrder.NULLS_FIRST, sortOrder.fields().get(0).nullOrder()); + + // verify data was written correctly + List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); + assertThat(writtenRecords, Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray())); + } } From 13750153e00ee80b1d4d6a64fe290b36a37331a1 Mon Sep 17 00:00:00 2001 From: Claire McGinty Date: Mon, 27 Jul 2026 14:37:24 -0400 Subject: [PATCH 2/2] withPartitionSpec -> withPartitionFields --- .../src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java | 2 +- .../java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java index 9b2946d239db..abdc2a179b58 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java @@ -497,7 +497,7 @@ public WriteRows withWriteProperties(Map writeProperties) { * *

See: https://iceberg.apache.org/spec/#partitioning */ - public WriteRows withPartitionSpec(List partitionFields) { + public WriteRows withPartitionFields(List partitionFields) { return toBuilder().setPartitionFields(partitionFields).build(); } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java index 5500edbadbc0..384e11b761bb 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java @@ -805,7 +805,7 @@ public void testCreateTableWithPartitionSpecAndSortOrder() { .apply( "Append To Table", writeTransform(catalog, tableId) - .withPartitionSpec(ImmutableList.of("bucket(id, 2)")) + .withPartitionFields(ImmutableList.of("bucket(id, 2)")) .withSortOrder(ImmutableList.of("data asc nulls first"))); testPipeline.run().waitUntilFinish();