[IcebergIO] Support sort order on dynamic table creation (#38269)

* [IcebergIO] Support sort order on dynamic table creation

* update changes and trigger IT

* fix changes.md

* address review feedback

Author: dejii

.github/trigger_files/IO_Iceberg_Integration_Tests.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run.",
-    "modification": 4
+    "modification": 5
 }

CHANGES.md

@@ -65,6 +65,7 @@
 ## I/Os
 
 * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* IcebergIO: support declaring a table's sort order on dynamic table creation via the new `sort_fields` config ([#38269](https://github.com/apache/beam/issues/38269)).
 
 ## New Features / Improvements
 

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

@@ -93,9 +93,11 @@
 import org.apache.iceberg.PartitionKey;
 import org.apache.iceberg.PartitionSpec;
 import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.SortOrder;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.TableProperties;
 import org.apache.iceberg.avro.Avro;
+import org.apache.iceberg.catalog.Catalog;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.iceberg.exceptions.AlreadyExistsException;
 import org.apache.iceberg.exceptions.NoSuchTableException;
@@ -145,19 +147,22 @@ public class AddFiles extends PTransform<PCollection<String>, PCollectionRowTupl
   private final int manifestFileSize;
   private final @Nullable String locationPrefix;
   private final @Nullable List<String> partitionFields;
+  private final @Nullable List<String> sortFields;
   private final @Nullable Map<String, String> tableProps;
 
   public AddFiles(
       IcebergCatalogConfig catalogConfig,
       String tableIdentifier,
       @Nullable String locationPrefix,
       @Nullable List<String> partitionFields,
+      @Nullable List<String> sortFields,
       @Nullable Map<String, String> tableProps,
       @Nullable Integer manifestFileSize,
       @Nullable Duration intervalTrigger) {
     this.catalogConfig = catalogConfig;
     this.tableIdentifier = tableIdentifier;
     this.partitionFields = partitionFields;
+    this.sortFields = sortFields;
     this.tableProps = tableProps;
     this.intervalTrigger = intervalTrigger;
     this.manifestFileSize =
@@ -193,6 +198,7 @@ public PCollectionRowTuple expand(PCollection<String> input) {
                         tableIdentifier,
                         locationPrefix,
                         partitionFields,
+                        sortFields,
                         tableProps))
                 .withOutputTags(DATA_FILES, TupleTagList.of(ERRORS)));
     SchemaCoder<SerializableDataFile> sdfCoder;
@@ -267,6 +273,7 @@ static class ConvertToDataFile extends DoFn<String, SerializableDataFile> {
     public static final TupleTag<SerializableDataFile> DATA_FILES = new TupleTag<>();
     private final @Nullable String prefix;
     private final @Nullable List<String> partitionFields;
+    private final @Nullable List<String> sortFields;
     private final @Nullable Map<String, String> tableProps;
     private transient @MonotonicNonNull ExecutorService executor;
     private transient @MonotonicNonNull LinkedList<Future<ProcessResult>> activeTasks;
@@ -281,11 +288,13 @@ public ConvertToDataFile(
         String identifier,
         @Nullable String prefix,
         @Nullable List<String> partitionFields,
+        @Nullable List<String> sortFields,
         @Nullable Map<String, String> tableProps) {
       this.catalogConfig = catalogConfig;
       this.identifier = identifier;
       this.prefix = prefix;
       this.partitionFields = partitionFields;
+      this.sortFields = sortFields;
       this.tableProps = tableProps;
     }
 
@@ -522,12 +531,18 @@ private Table getOrCreateTable(String filePath, FileFormat format) throws IOExce
         try {
           org.apache.iceberg.Schema schema = getSchema(filePath, format);
           PartitionSpec spec = PartitionUtils.toPartitionSpec(partitionFields, schema);
+          SortOrder sortOrder = SortOrderUtils.toSortOrder(sortFields, schema);
 
-          return tableProps == null
-              ? catalogConfig.catalog().createTable(TableIdentifier.parse(identifier), schema, spec)
-              : catalogConfig
+          Catalog.TableBuilder builder =
+              catalogConfig
                   .catalog()
-                  .createTable(TableIdentifier.parse(identifier), schema, spec, tableProps);
+                  .buildTable(tableId, schema)
+                  .withPartitionSpec(spec)
+                  .withSortOrder(sortOrder);
+          if (tableProps != null) {
+            builder.withProperties(tableProps);
+          }
+          return builder.create();
         } catch (AlreadyExistsException e2) { // if table already exists, just load it
           return catalogConfig.catalog().loadTable(TableIdentifier.parse(identifier));
         }

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

@@ -106,6 +106,15 @@ public static Builder builder() {
             + " please visit https://iceberg.apache.org/docs/latest/configuration/#table-properties.")
     public abstract @Nullable Map<String, String> getTableProperties();
 
+    @SchemaFieldDescription(
+        "Fields used to set the table's sort order, applied when the table is created. "
+            + "Each entry has the form `<term> [asc|desc] [nulls first|nulls last]`, where `<term>` "
+            + "is a field name or one of the partition transforms (e.g. `bucket(col, 4)`, `day(ts)`). "
+            + "Direction defaults to ascending; null order defaults to nulls-first for ascending and "
+            + "nulls-last for descending.\n"
+            + "For more information on sort orders, please visit https://iceberg.apache.org/spec/#sort-orders.")
+    public abstract @Nullable List<String> getSortFields();
+
     @SchemaFieldDescription("This option specifies whether and where to output unwritable rows.")
     public abstract @Nullable ErrorHandling getErrorHandling();
 
@@ -127,6 +136,8 @@ public abstract static class Builder {
 
       public abstract Builder setTableProperties(Map<String, String> props);
 
+      public abstract Builder setSortFields(List<String> sortFields);
+
       public abstract Builder setErrorHandling(ErrorHandling errorHandling);
 
       public abstract Configuration build();
@@ -172,6 +183,7 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
                       configuration.getTable(),
                       configuration.getLocationPrefix(),
                       configuration.getPartitionFields(),
+                      configuration.getSortFields(),
                       configuration.getTableProperties(),
                       configuration.getManifestFileSize(),
                       frequency != null ? Duration.standardSeconds(frequency) : null));

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

@@ -21,7 +21,9 @@
 import java.util.List;
 import java.util.Map;
 import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.annotations.SchemaIgnore;
 import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.SortOrder;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.checkerframework.dataflow.qual.Pure;
 
@@ -41,6 +43,16 @@ public PartitionSpec getPartitionSpec() {
   @Pure
   public abstract @Nullable List<String> getPartitionFields();
 
+  /** Sort order to apply when the table is dynamically created. */
+  @Pure
+  @SchemaIgnore
+  public SortOrder getSortOrder() {
+    return SortOrderUtils.toSortOrder(getSortFields(), getSchema());
+  }
+
+  @Pure
+  public abstract @Nullable List<String> getSortFields();
+
   @Pure
   public abstract @Nullable Map<String, String> getTableProperties();
 
@@ -55,6 +67,8 @@ public abstract static class Builder {
 
     public abstract Builder setPartitionFields(@Nullable List<String> partitionFields);
 
+    public abstract Builder setSortFields(@Nullable List<String> sortFields);
+
     public abstract Builder setTableProperties(@Nullable Map<String, String> tableProperties);
 
     @Pure

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

@@ -134,6 +134,16 @@ public static Builder builder() {
             + " please visit https://iceberg.apache.org/docs/latest/configuration/#table-properties.")
     public abstract @Nullable Map<String, String> getTableProperties();
 
+    @SchemaFieldDescription(
+        "Fields used to set the table's sort order, applied when the table is created. "
+            + "Each entry has the form `<term> [asc|desc] [nulls first|nulls last]`, where `<term>` "
+            + "is a field name or one of the partition transforms (e.g. `bucket(col, 4)`, `day(ts)`). "
+            + "Direction defaults to ascending; null order defaults to nulls-first for ascending and "
+            + "nulls-last for descending. Note: this sets the table's declared sort order as metadata; "
+            + "it does not cause Beam to physically sort records before writing.\n"
+            + "For more information on sort orders, please visit https://iceberg.apache.org/spec/#sort-orders.")
+    public abstract @Nullable List<String> getSortFields();
+
     @AutoValue.Builder
     public abstract static class Builder {
       public abstract Builder setTable(String table);
@@ -158,6 +168,8 @@ public abstract static class Builder {
 
       public abstract Builder setTableProperties(Map<String, String> tableProperties);
 
+      public abstract Builder setSortFields(List<String> sortFields);
+
       public abstract Configuration build();
     }
 
@@ -223,6 +235,7 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
                       FileFormat.PARQUET.toString(),
                       rows.getSchema(),
                       configuration.getPartitionFields(),
+                      configuration.getSortFields(),
                       configuration.getTableProperties(),
                       configuration.getDrop(),
                       configuration.getKeep(),

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

@@ -34,19 +34,22 @@ class PortableIcebergDestinations implements DynamicDestinations {
   private final String fileFormat;
 
   private final @Nullable List<String> partitionFields;
+  private final @Nullable List<String> sortFields;
   private final @Nullable Map<String, String> tableProperties;
 
   public PortableIcebergDestinations(
       String destinationTemplate,
       String fileFormat,
       Schema inputSchema,
       @Nullable List<String> partitionFields,
+      @Nullable List<String> sortFields,
       @Nullable Map<String, String> tableProperties,
       @Nullable List<String> fieldsToDrop,
       @Nullable List<String> fieldsToKeep,
       @Nullable String onlyField) {
     this.interpolator = new RowStringInterpolator(destinationTemplate, inputSchema);
     this.partitionFields = partitionFields;
+    this.sortFields = sortFields;
     this.tableProperties = tableProperties;
     RowFilter rf = new RowFilter(inputSchema);
 
@@ -86,6 +89,7 @@ public IcebergDestination instantiateDestination(String dest) {
             IcebergTableCreateConfig.builder()
                 .setSchema(getDataSchema())
                 .setPartitionFields(partitionFields)
+                .setSortFields(sortFields)
                 .setTableProperties(tableProperties)
                 .build())
         .setFileFormat(FileFormat.fromString(fileFormat))

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

@@ -49,6 +49,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.Table;
 import org.apache.iceberg.catalog.Catalog;
 import org.apache.iceberg.catalog.Namespace;
@@ -329,6 +330,7 @@ Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) {
     @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()
@@ -357,13 +359,20 @@ Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) {
       } catch (NoSuchTableException e) { // Otherwise, create the table
         org.apache.iceberg.Schema tableSchema = IcebergUtils.beamSchemaToIcebergSchema(dataSchema);
         try {
-          table = catalog.createTable(identifier, tableSchema, partitionSpec, tableProperties);
+          table =
+              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);
         } catch (AlreadyExistsException ignored) {
           // race condition: another worker already created this table