This is an automated email from the ASF dual-hosted git repository.
ahmedabu98 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 70a522364e9 (IcebergIO) Support PartitionSpec/SortOrder on dynamic
table creation via IcebergIO (#39408)
70a522364e9 is described below
commit 70a522364e9a5766ddb48e2cadbf1b9471b6d205
Author: Claire McGinty <[email protected]>
AuthorDate: Mon Jul 27 16:34:21 2026 -0400
(IcebergIO) Support PartitionSpec/SortOrder on dynamic table creation via
IcebergIO (#39408)
* (IcebergIO) Support SortOrder on dynamic table creation via IcebergIO
* withPartitionSpec -> withPartitionFields
---
.../beam/sdk/io/iceberg/DynamicDestinations.java | 12 ++++-
.../org/apache/beam/sdk/io/iceberg/IcebergIO.java | 33 ++++++++++++-
.../io/iceberg/OneTableDynamicDestinations.java | 33 +++++++++++--
.../io/iceberg/WritePartitionedRowsToFiles.java | 13 +++++-
.../beam/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 0185758c8ae..ba9ec771c11 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<Row> 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<String> partitionFields,
+ @Nullable List<String> 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 04feea5037d..abdc2a179b5 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 class IcebergIO {
abstract @Nullable Map<String, String> getWriteProperties();
+ abstract @Nullable List<String> getPartitionFields();
+
+ abstract @Nullable List<String> getSortFields();
+
abstract Builder toBuilder();
@AutoValue.Builder
@@ -429,6 +433,10 @@ public class IcebergIO {
abstract Builder setWriteProperties(Map<String, String> writeProperties);
+ abstract Builder setPartitionFields(List<String> partitionFields);
+
+ abstract Builder setSortFields(List<String> sortFields);
+
abstract WriteRows build();
}
@@ -483,6 +491,26 @@ public class IcebergIO {
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)`
+ *
+ * <p>See: https://iceberg.apache.org/spec/#partitioning
+ */
+ public WriteRows withPartitionFields(List<String> 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`
+ *
+ * <p>See: https://iceberg.apache.org/spec/#sorting
+ */
+ public WriteRows withSortOrder(List<String> sortFields) {
+ return toBuilder().setSortFields(sortFields).build();
+ }
+
@Override
public IcebergWriteResult expand(PCollection<Row> input) {
List<?> allToArgs = Arrays.asList(getTableIdentifier(),
getDynamicDestinations());
@@ -494,7 +522,10 @@ public class IcebergIO {
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 afca43fca15..a9399558948 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.Externalizable;
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.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.Vi
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<String> partitionFields;
+ private transient @Nullable List<String> sortFields;
@VisibleForTesting
TableIdentifier getTableIdentifier() {
@@ -46,9 +50,15 @@ class OneTableDynamicDestinations implements
DynamicDestinations, Externalizable
return tableId;
}
- OneTableDynamicDestinations(TableIdentifier tableId, Schema rowSchema) {
+ OneTableDynamicDestinations(
+ TableIdentifier tableId,
+ Schema rowSchema,
+ @Nullable List<String> partitionFields,
+ @Nullable List<String> sortFields) {
this.tableIdString = IcebergUtils.tableIdentifierToString(tableId);
this.rowSchema = rowSchema;
+ this.partitionFields = partitionFields;
+ this.sortFields = sortFields;
}
@Override
@@ -68,9 +78,18 @@ class OneTableDynamicDestinations implements
DynamicDestinations, Externalizable
@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 @@ class OneTableDynamicDestinations implements
DynamicDestinations, Externalizable
// 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<String>) in.readObject();
+ sortFields = (List<String>) 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 84e7d05b124..d1a08980fa9 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.DataFiles;
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 @@ class WritePartitionedRowsToFiles
@Nullable IcebergTableCreateConfig createConfig =
destination.getTableCreateConfig();
PartitionSpec partitionSpec =
createConfig != null ? createConfig.getPartitionSpec() :
PartitionSpec.unpartitioned();
+ SortOrder sortOrder =
+ createConfig != null ? createConfig.getSortOrder() :
SortOrder.unsorted();
Map<String, String> tableProperties =
createConfig != null && createConfig.getTableProperties() != null
? createConfig.getTableProperties()
@@ -217,13 +220,19 @@ class WritePartitionedRowsToFiles
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 6580f4eeaf6..384e11b761b 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.CatalogUtil;
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 class IcebergIOWriteTest implements Serializable {
.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<String, String> catalogProps =
+ ImmutableMap.<String, String>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)
+ .withPartitionFields(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<Record> writtenRecords =
ImmutableList.copyOf(IcebergGenerics.read(table).build());
+ assertThat(writtenRecords,
Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray()));
+ }
}
