HIVE-29133: Support Z-ordering for Iceberg tables via CREATE TABLE (#6138)

Author: kokila-19

iceberg/iceberg-handler/src/main/java/org/apache/iceberg/mr/InputFormatConfig.java

@@ -78,6 +78,10 @@ private InputFormatConfig() {
 
   public static final String CATALOG_CONFIG_PREFIX = "iceberg.catalog.";
 
+  public static final String SORT_ORDER = "sort.order";
+  public static final String SORT_COLUMNS = "sort.columns";
+  public static final String ZORDER = "ZORDER";
+
   public enum InMemoryDataModel {
     HIVE,
     GENERIC // Default data model is of Iceberg Generics

iceberg/iceberg-handler/src/main/java/org/apache/iceberg/mr/hive/BaseHiveIcebergMetaHook.java

@@ -19,6 +19,7 @@
 
 package org.apache.iceberg.mr.hive;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.util.Arrays;
 import java.util.Collection;
@@ -41,6 +42,8 @@
 import org.apache.hadoop.hive.metastore.api.hive_metastoreConstants;
 import org.apache.hadoop.hive.ql.ddl.misc.sortoder.SortFieldDesc;
 import org.apache.hadoop.hive.ql.ddl.misc.sortoder.SortFields;
+import org.apache.hadoop.hive.ql.ddl.misc.sortoder.ZOrderFieldDesc;
+import org.apache.hadoop.hive.ql.ddl.misc.sortoder.ZOrderFields;
 import org.apache.hadoop.hive.ql.util.NullOrdering;
 import org.apache.iceberg.BaseMetastoreTableOperations;
 import org.apache.iceberg.BaseTable;
@@ -74,6 +77,9 @@
 import org.slf4j.LoggerFactory;
 
 import static org.apache.iceberg.RowLevelOperationMode.MERGE_ON_READ;
+import static org.apache.iceberg.mr.InputFormatConfig.SORT_COLUMNS;
+import static org.apache.iceberg.mr.InputFormatConfig.SORT_ORDER;
+import static org.apache.iceberg.mr.InputFormatConfig.ZORDER;
 
 public class BaseHiveIcebergMetaHook implements HiveMetaHook {
   private static final Logger LOG = LoggerFactory.getLogger(BaseHiveIcebergMetaHook.class);
@@ -84,6 +90,7 @@ public class BaseHiveIcebergMetaHook implements HiveMetaHook {
   private static final Set<String> PARAMETERS_TO_REMOVE = ImmutableSet
       .of(InputFormatConfig.TABLE_SCHEMA, Catalogs.LOCATION, Catalogs.NAME, InputFormatConfig.PARTITION_SPEC);
   static final String ORC_FILES_ONLY = "iceberg.orc.files.only";
+  private static final String ZORDER_FIELDS_JSON_KEY = "zorderFields";
 
   protected final Configuration conf;
   protected Table icebergTable = null;
@@ -217,28 +224,82 @@ private void validateCatalogConfigsDefined() {
     }
   }
 
+  /**
+   *  Persists the table's write sort order based on the HMS property 'default-sort-order'
+   *  that is populated by the DDL layer.
+   * <p>
+   * Behaviour:
+   * - If the JSON represents Z-order, we remove DEFAULT_SORT_ORDER
+   *   as Iceberg does not have Z-order support in its spec.
+   *   So, we persist Z-order metadata in {@link org.apache.iceberg.mr.InputFormatConfig#SORT_ORDER}
+   *   and {@link org.apache.iceberg.mr.InputFormatConfig#SORT_COLUMNS} to be used by Hive Writer.
+   * <p>
+   * - Otherwise, the JSON is a list of SortFields; we convert it to Iceberg
+   *   SortOrder JSON and keep it in DEFAULT_SORT_ORDER for Iceberg to use it.
+   */
   private void setSortOrder(org.apache.hadoop.hive.metastore.api.Table hmsTable, Schema schema,
       Properties properties) {
-    String sortOderJSONString = hmsTable.getParameters().get(TableProperties.DEFAULT_SORT_ORDER);
-    SortFields sortFields = null;
-    if (!Strings.isNullOrEmpty(sortOderJSONString)) {
-      try {
-        sortFields = JSON_OBJECT_MAPPER.reader().readValue(sortOderJSONString, SortFields.class);
-      } catch (Exception e) {
-        LOG.warn("Can not read write order json: {}", sortOderJSONString, e);
-        return;
-      }
+    String sortOrderJSONString = hmsTable.getParameters().get(TableProperties.DEFAULT_SORT_ORDER);
+    if (Strings.isNullOrEmpty(sortOrderJSONString)) {
+      return;
+    }
+
+    if (isZOrderJSON(sortOrderJSONString)) {
+      properties.remove(TableProperties.DEFAULT_SORT_ORDER);
+      setZOrderSortOrder(sortOrderJSONString, properties, hmsTable.getTableName());
+      return;
+    }
+
+    try {
+      SortFields sortFields = JSON_OBJECT_MAPPER.reader().readValue(sortOrderJSONString, SortFields.class);
       if (sortFields != null && !sortFields.getSortFields().isEmpty()) {
-        SortOrder.Builder sortOderBuilder = SortOrder.builderFor(schema);
+        SortOrder.Builder sortOrderBuilder = SortOrder.builderFor(schema);
         sortFields.getSortFields().forEach(fieldDesc -> {
           NullOrder nullOrder = fieldDesc.getNullOrdering() == NullOrdering.NULLS_FIRST ?
               NullOrder.NULLS_FIRST : NullOrder.NULLS_LAST;
           SortDirection sortDirection = fieldDesc.getDirection() == SortFieldDesc.SortDirection.ASC ?
               SortDirection.ASC : SortDirection.DESC;
-          sortOderBuilder.sortBy(fieldDesc.getColumnName(), sortDirection, nullOrder);
+          sortOrderBuilder.sortBy(fieldDesc.getColumnName(), sortDirection, nullOrder);
         });
-        properties.put(TableProperties.DEFAULT_SORT_ORDER, SortOrderParser.toJson(sortOderBuilder.build()));
+        properties.put(TableProperties.DEFAULT_SORT_ORDER, SortOrderParser.toJson(sortOrderBuilder.build()));
+      }
+    } catch (Exception e) {
+      LOG.warn("Can not read write order json: {}", sortOrderJSONString);
+    }
+  }
+
+  /**
+   * Configures the Z-order sort order metadata in the given properties
+   * based on the specified Z-order fields.
+   *
+   * @param jsonString the JSON string representing sort orders
+   * @param properties the Properties object to store sort order metadata
+   * @param tableName name of the table
+   */
+  private void setZOrderSortOrder(String jsonString, Properties properties, String tableName) {
+    try {
+      ZOrderFields zorderFields = JSON_OBJECT_MAPPER.reader().readValue(jsonString, ZOrderFields.class);
+      if (zorderFields != null && !zorderFields.getZOrderFields().isEmpty()) {
+        List<String> columnNames = zorderFields.getZOrderFields().stream()
+            .map(ZOrderFieldDesc::getColumnName)
+            .collect(Collectors.toList());
+
+        properties.put(SORT_ORDER, ZORDER);
+        properties.put(SORT_COLUMNS, String.join(",", columnNames));
+
+        LOG.debug("Applying Z-ordering for Iceberg Table {} with Columns: {}", tableName, columnNames);
       }
+    } catch (Exception e) {
+      LOG.warn("Failed to parse Z-order sort order", e);
+    }
+  }
+
+  private boolean isZOrderJSON(String jsonString) {
+    try {
+      JsonNode node = JSON_OBJECT_MAPPER.readTree(jsonString);
+      return node.has(ZORDER_FIELDS_JSON_KEY);
+    } catch (Exception e) {
+      return false;
     }
   }
 

iceberg/iceberg-handler/src/main/java/org/apache/iceberg/mr/hive/HiveIcebergStorageHandler.java

@@ -82,6 +82,7 @@
 import org.apache.hadoop.hive.ql.ddl.table.create.like.CreateTableLikeDesc;
 import org.apache.hadoop.hive.ql.ddl.table.misc.properties.AlterTableSetPropertiesDesc;
 import org.apache.hadoop.hive.ql.exec.ColumnInfo;
+import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
 import org.apache.hadoop.hive.ql.exec.Utilities;
 import org.apache.hadoop.hive.ql.hooks.WriteEntity;
 import org.apache.hadoop.hive.ql.io.IOConstants;
@@ -184,6 +185,7 @@
 import org.apache.iceberg.mr.InputFormatConfig;
 import org.apache.iceberg.mr.hive.actions.HiveIcebergDeleteOrphanFiles;
 import org.apache.iceberg.mr.hive.plan.IcebergBucketFunction;
+import org.apache.iceberg.mr.hive.udf.GenericUDFIcebergZorder;
 import org.apache.iceberg.puffin.Blob;
 import org.apache.iceberg.puffin.BlobMetadata;
 import org.apache.iceberg.puffin.Puffin;
@@ -218,6 +220,9 @@
 import static org.apache.iceberg.SnapshotSummary.TOTAL_FILE_SIZE_PROP;
 import static org.apache.iceberg.SnapshotSummary.TOTAL_POS_DELETES_PROP;
 import static org.apache.iceberg.SnapshotSummary.TOTAL_RECORDS_PROP;
+import static org.apache.iceberg.mr.InputFormatConfig.SORT_COLUMNS;
+import static org.apache.iceberg.mr.InputFormatConfig.SORT_ORDER;
+import static org.apache.iceberg.mr.InputFormatConfig.ZORDER;
 
 public class HiveIcebergStorageHandler extends DefaultStorageHandler implements HiveStoragePredicateHandler {
   private static final Logger LOG = LoggerFactory.getLogger(HiveIcebergStorageHandler.class);
@@ -929,14 +934,64 @@ public DynamicPartitionCtx createDPContext(
       addCustomSortExpr(table, hmsTable, writeOperation, dpCtx, getSortTransformSpec(table));
     }
 
+    // Even if table has no explicit sort order, honor z-order if configured
+    Map<String, String> props = table.properties();
+    if (ZORDER.equalsIgnoreCase(props.getOrDefault(SORT_ORDER, ""))) {
+      addZOrderCustomExpr(props, dpCtx, table, hmsTable, writeOperation);
+    }
+
     return dpCtx;
   }
 
+  /**
+   * Adds a custom sort expression to the DynamicPartitionCtx that performs local Z-ordering on write.
+   *
+   * Behavior:
+   * - Reads Z-order properties from 'sort.order' and 'sort.columns' (comma-separated).
+   * - Resolves the referenced columns to their positions in the physical row (taking into account
+   *   ACID virtual columns offset for overwrite/update operations).
+   */
+  private void addZOrderCustomExpr(Map<String, String> props, DynamicPartitionCtx dpCtx, Table table,
+      org.apache.hadoop.hive.ql.metadata.Table hmsTable, Operation writeOperation) {
+    String colsProp = props.get(SORT_COLUMNS);
+    if (StringUtils.isNotBlank(colsProp)) {
+      List<String> zCols = Arrays.stream(colsProp.split(","))
+          .map(String::trim)
+          .filter(s -> !s.isEmpty())
+          .toList();
+
+      List<Types.NestedField> fields = table.schema().columns();
+      Map<String, Integer> fieldOrderMap = Maps.newHashMapWithExpectedSize(fields.size());
+      for (int i = 0; i < fields.size(); ++i) {
+        fieldOrderMap.put(fields.get(i).name(), i);
+      }
+      int offset = (shouldOverwrite(hmsTable, writeOperation) ?
+          ACID_VIRTUAL_COLS_AS_FIELD_SCHEMA : acidSelectColumns(hmsTable, writeOperation)).size();
+
+      List<Integer> zIndices = zCols.stream().map(col -> {
+        Integer base = fieldOrderMap.get(col);
+        Preconditions.checkArgument(base != null, "Z-order column not found in schema: %s", col);
+        return base + offset;
+      }).toList();
+
+      dpCtx.addCustomSortExpressions(Collections.singletonList(allCols -> {
+        List<ExprNodeDesc> args = zIndices.stream()
+            .map(allCols::get)
+            .toList();
+        try {
+          return ExprNodeGenericFuncDesc.newInstance(new GenericUDFIcebergZorder(), "iceberg_zorder", args);
+        } catch (UDFArgumentException e) {
+          throw new RuntimeException(e);
+        }
+      }));
+    }
+  }
+
   private void addCustomSortExpr(Table table,  org.apache.hadoop.hive.ql.metadata.Table hmsTable,
       Operation writeOperation, DynamicPartitionCtx dpCtx,
       List<TransformSpec> transformSpecs) {
-    Map<String, Integer> fieldOrderMap = Maps.newHashMap();
     List<Types.NestedField> fields = table.schema().columns();
+    Map<String, Integer> fieldOrderMap = Maps.newHashMapWithExpectedSize(fields.size());
     for (int i = 0; i < fields.size(); ++i) {
       fieldOrderMap.put(fields.get(i).name(), i);
     }