HIVE-27918: Iceberg: Push transforms for clustering during table writes

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

@@ -38,7 +38,6 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.apache.commons.collections.MapUtils;
@@ -76,7 +75,6 @@
 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.FetchOperator;
-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;
@@ -120,7 +118,6 @@
 import org.apache.hadoop.hive.serde2.objectinspector.StructField;
 import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
 import org.apache.hadoop.hive.serde2.typeinfo.PrimitiveTypeInfo;
-import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoFactory;
 import org.apache.hadoop.io.Writable;
 import org.apache.hadoop.io.WritableComparable;
 import org.apache.hadoop.mapred.InputFormat;
@@ -217,26 +214,6 @@ public class HiveIcebergStorageHandler implements HiveStoragePredicateHandler, H
 
   public static final String TABLE_DEFAULT_LOCATION = "TABLE_DEFAULT_LOCATION";
 
-  /**
-   * Function template for producing a custom sort expression function:
-   * Takes the source column index and the bucket count to creat a function where Iceberg bucket UDF is used to build
-   * the sort expression, e.g. iceberg_bucket(_col2, 5)
-   */
-  private static final transient BiFunction<Integer, Integer, Function<List<ExprNodeDesc>, ExprNodeDesc>>
-      BUCKET_SORT_EXPR =
-          (idx, bucket) -> cols -> {
-            try {
-              ExprNodeDesc icebergBucketSourceCol = cols.get(idx);
-              return ExprNodeGenericFuncDesc.newInstance(new GenericUDFIcebergBucket(), "iceberg_bucket",
-                  Lists.newArrayList(
-                      icebergBucketSourceCol,
-                      new ExprNodeConstantDesc(TypeInfoFactory.intTypeInfo, bucket)
-                  ));
-            } catch (UDFArgumentException e) {
-              throw new RuntimeException(e);
-            }
-          };
-
   private static final List<VirtualColumn> ACID_VIRTUAL_COLS = ImmutableList.of(VirtualColumn.PARTITION_SPEC_ID,
       VirtualColumn.PARTITION_HASH, VirtualColumn.FILE_PATH, VirtualColumn.ROW_POSITION);
   private static final List<FieldSchema> ACID_VIRTUAL_COLS_AS_FIELD_SCHEMA = ACID_VIRTUAL_COLS.stream()
@@ -731,14 +708,9 @@ private void addCustomSortExpr(Table table,  org.apache.hadoop.hive.ql.metadata.
     int offset = (shouldOverwrite(hmsTable, writeOperation) ?
         ACID_VIRTUAL_COLS_AS_FIELD_SCHEMA : acidSelectColumns(hmsTable, writeOperation)).size();
 
-    for (TransformSpec spec : transformSpecs) {
-      int order = fieldOrderMap.get(spec.getColumnName());
-      if (TransformSpec.TransformType.BUCKET.equals(spec.getTransformType())) {
-        customSortExprs.add(BUCKET_SORT_EXPR.apply(order + offset, spec.getTransformParam().get()));
-      } else {
-        customSortExprs.add(cols -> cols.get(order + offset).clone());
-      }
-    }
+    customSortExprs.addAll(transformSpecs.stream().map(spec ->
+        IcebergTransformSortFunctionUtil.getCustomSortExprs(spec, fieldOrderMap.get(spec.getColumnName()) + offset)
+    ).collect(Collectors.toList()));
   }
 
   @Override

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

@@ -0,0 +1,158 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iceberg.mr.hive;
+
+import java.util.List;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
+import org.apache.hadoop.hive.ql.parse.TransformSpec;
+import org.apache.hadoop.hive.ql.plan.ExprNodeConstantDesc;
+import org.apache.hadoop.hive.ql.plan.ExprNodeDesc;
+import org.apache.hadoop.hive.ql.plan.ExprNodeGenericFuncDesc;
+import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoFactory;
+import org.apache.iceberg.mr.hive.udf.GenericUDFIcebergBucket;
+import org.apache.iceberg.mr.hive.udf.GenericUDFIcebergDay;
+import org.apache.iceberg.mr.hive.udf.GenericUDFIcebergHour;
+import org.apache.iceberg.mr.hive.udf.GenericUDFIcebergMonth;
+import org.apache.iceberg.mr.hive.udf.GenericUDFIcebergTruncate;
+import org.apache.iceberg.mr.hive.udf.GenericUDFIcebergYear;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+
+/**
+ * A utility class which provides Iceberg transform sort functions.
+ */
+public final class IcebergTransformSortFunctionUtil {
+
+  private IcebergTransformSortFunctionUtil() {
+    // not called
+  }
+
+  /**
+   * Function template for producing a custom sort expression function:
+   * Takes the source column index and the bucket count to create a function where Iceberg transform UDF is used to
+   * build the sort expression, e.g. iceberg_bucket(_col2, 5)
+   */
+  private static final transient BiFunction<Integer, Integer, Function<List<ExprNodeDesc>, ExprNodeDesc>>
+      BUCKET_SORT_EXPR =
+          (idx, bucket) -> cols -> {
+            try {
+              ExprNodeDesc icebergBucketSourceCol = cols.get(idx);
+              return ExprNodeGenericFuncDesc.newInstance(new GenericUDFIcebergBucket(), "iceberg_bucket",
+                  Lists.newArrayList(
+                      icebergBucketSourceCol,
+                      new ExprNodeConstantDesc(TypeInfoFactory.intTypeInfo, bucket)
+                  ));
+            } catch (UDFArgumentException e) {
+              throw new RuntimeException(e);
+            }
+          };
+
+  private static final transient BiFunction<Integer, Integer, Function<List<ExprNodeDesc>, ExprNodeDesc>>
+      TRUNCATE_SORT_EXPR =
+          (idx, truncateLength) -> cols -> {
+            try {
+              ExprNodeDesc icebergTruncateSourceCol = cols.get(idx);
+              return ExprNodeGenericFuncDesc.newInstance(new GenericUDFIcebergTruncate(), "iceberg_truncate",
+                  Lists.newArrayList(
+                      icebergTruncateSourceCol,
+                      new ExprNodeConstantDesc(TypeInfoFactory.intTypeInfo, truncateLength)
+                  ));
+            } catch (UDFArgumentException e) {
+              throw new RuntimeException(e);
+            }
+          };
+
+  private static final transient Function<Integer, Function<List<ExprNodeDesc>, ExprNodeDesc>>
+      YEAR_SORT_EXPR =
+          idx -> cols -> {
+            try {
+              ExprNodeDesc icebergYearSourceCol = cols.get(idx);
+              return ExprNodeGenericFuncDesc.newInstance(new GenericUDFIcebergYear(), "iceberg_year",
+                  Lists.newArrayList(
+                      icebergYearSourceCol
+                  ));
+            } catch (UDFArgumentException e) {
+              throw new RuntimeException(e);
+            }
+          };
+
+  private static final transient Function<Integer, Function<List<ExprNodeDesc>, ExprNodeDesc>>
+      MONTH_SORT_EXPR =
+          idx -> cols -> {
+            try {
+              ExprNodeDesc icebergMonthSourceCol = cols.get(idx);
+              return ExprNodeGenericFuncDesc.newInstance(new GenericUDFIcebergMonth(), "iceberg_month",
+                  Lists.newArrayList(
+                      icebergMonthSourceCol
+                  ));
+            } catch (UDFArgumentException e) {
+              throw new RuntimeException(e);
+            }
+          };
+
+  private static final transient Function<Integer, Function<List<ExprNodeDesc>, ExprNodeDesc>>
+      DAY_SORT_EXPR =
+          idx -> cols -> {
+            try {
+              ExprNodeDesc icebergDaySourceCol = cols.get(idx);
+              return ExprNodeGenericFuncDesc.newInstance(new GenericUDFIcebergDay(), "iceberg_day",
+                  Lists.newArrayList(
+                      icebergDaySourceCol
+                  ));
+            } catch (UDFArgumentException e) {
+              throw new RuntimeException(e);
+            }
+          };
+
+  private static final transient Function<Integer, Function<List<ExprNodeDesc>, ExprNodeDesc>>
+      HOUR_SORT_EXPR =
+          idx -> cols -> {
+            try {
+              ExprNodeDesc icebergHourSourceCol = cols.get(idx);
+              return ExprNodeGenericFuncDesc.newInstance(new GenericUDFIcebergHour(), "iceberg_hour",
+                  Lists.newArrayList(
+                      icebergHourSourceCol
+                  ));
+            } catch (UDFArgumentException e) {
+              throw new RuntimeException(e);
+            }
+          };
+
+  public static Function<List<ExprNodeDesc>, ExprNodeDesc> getCustomSortExprs(TransformSpec spec, int index) {
+    switch (spec.getTransformType()) {
+      case BUCKET:
+        return BUCKET_SORT_EXPR.apply(index, spec.getTransformParam().get());
+      case TRUNCATE:
+        return TRUNCATE_SORT_EXPR.apply(index, spec.getTransformParam().get());
+      case YEAR:
+        return YEAR_SORT_EXPR.apply(index);
+      case MONTH:
+        return MONTH_SORT_EXPR.apply(index);
+      case DAY:
+        return DAY_SORT_EXPR.apply(index);
+      case HOUR:
+        return HOUR_SORT_EXPR.apply(index);
+      default:
+        return cols -> cols.get(index).clone();
+    }
+  }
+
+}

iceberg/iceberg-handler/src/main/java/org/apache/iceberg/mr/hive/GenericUDFIcebergBucket.java → iceberg/iceberg-handler/src/main/java/org/apache/iceberg/mr/hive/udf/GenericUDFIcebergBucket.java

@@ -16,7 +16,7 @@
  * limitations under the License.
  */
 
-package org.apache.iceberg.mr.hive;
+package org.apache.iceberg.mr.hive.udf;
 
 import java.nio.ByteBuffer;
 import java.util.function.Function;

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

@@ -0,0 +1,129 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.iceberg.mr.hive.udf;
+
+import java.util.function.Function;
+import org.apache.hadoop.hive.ql.exec.Description;
+import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
+import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
+import org.apache.hadoop.hive.ql.metadata.HiveException;
+import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
+import org.apache.hadoop.hive.serde2.io.DateWritableV2;
+import org.apache.hadoop.hive.serde2.io.TimestampLocalTZWritable;
+import org.apache.hadoop.hive.serde2.io.TimestampWritableV2;
+import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
+import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorConverters;
+import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
+import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorConverter;
+import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
+import org.apache.hadoop.io.IntWritable;
+import org.apache.iceberg.transforms.Transforms;
+import org.apache.iceberg.types.Types;
+
+/**
+ * GenericUDFIcebergDay - UDF that wraps around Iceberg's day transform function
+ */
+@Description(name = "iceberg_day",
+    value = "_FUNC_(value) - " +
+     "Returns the bucket value calculated by Iceberg bucket transform function ",
+    extended = "Example:\n  > SELECT _FUNC_('2023-01-02', 5);\n  2")
+public class GenericUDFIcebergDay extends GenericUDF {
+  private final IntWritable result = new IntWritable();
+  private transient PrimitiveObjectInspector argumentOI;
+  private transient ObjectInspectorConverters.Converter converter;
+
+  @FunctionalInterface
+  private interface UDFEvalFunction<T> {
+    void apply(T argument) throws HiveException;
+  }
+
+  private transient UDFEvalFunction<DeferredObject> evaluator;
+
+  @Override
+  public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
+    if (arguments.length != 1) {
+      throw new UDFArgumentLengthException(
+        "ICEBERG_DAY requires 1 arguments (value), but got " + arguments.length);
+    }
+
+    if (arguments[0].getCategory() != ObjectInspector.Category.PRIMITIVE) {
+      throw new UDFArgumentException(
+        "ICEBERG_DAY first argument takes primitive types, got " + argumentOI.getTypeName());
+    }
+    argumentOI = (PrimitiveObjectInspector) arguments[0];
+
+    PrimitiveObjectInspector.PrimitiveCategory inputType = argumentOI.getPrimitiveCategory();
+    ObjectInspector outputOI;
+    switch (inputType) {
+      case DATE:
+        converter = new PrimitiveObjectInspectorConverter.DateConverter(argumentOI,
+          PrimitiveObjectInspectorFactory.writableDateObjectInspector);
+        Function<Object, Integer> dateTransform = Transforms.day().bind(Types.DateType.get());
+        evaluator = arg -> {
+          DateWritableV2 val = (DateWritableV2) converter.convert(arg.get());
+          result.set(dateTransform.apply(val.getDays()));
+        };
+        break;
+
+      case TIMESTAMP:
+        converter = new PrimitiveObjectInspectorConverter.TimestampConverter(argumentOI,
+          PrimitiveObjectInspectorFactory.writableTimestampObjectInspector);
+        Function<Object, Integer> timestampTransform = Transforms.day().bind(Types.TimestampType.withoutZone());
+        evaluator = arg -> {
+          TimestampWritableV2 val = (TimestampWritableV2) converter.convert(arg.get());
+          result.set(timestampTransform.apply(val.getNanos() / 1000L));
+        };
+        break;
+
+      case TIMESTAMPLOCALTZ:
+        converter = new PrimitiveObjectInspectorConverter.TimestampLocalTZConverter(argumentOI,
+          PrimitiveObjectInspectorFactory.writableTimestampTZObjectInspector);
+        Function<Object, Integer> timestampLocalTzTransform = Transforms.day().bind(Types.TimestampType.withZone());
+        evaluator = arg -> {
+          TimestampLocalTZWritable val = (TimestampLocalTZWritable) converter.convert(arg.get());
+          result.set(timestampLocalTzTransform.apply(val.getNanos() / 1000L));
+        };
+        break;
+
+      default:
+        throw new UDFArgumentException(
+          " ICEBERG_DAY() only takes DATE/TIMESTAMP/TIMESTAMPLOCALTZ" +
+          " types as first argument, got " + inputType);
+    }
+    outputOI = PrimitiveObjectInspectorFactory.writableIntObjectInspector;
+    return outputOI;
+  }
+
+  @Override
+  public Object evaluate(DeferredObject[] arguments) throws HiveException {
+    DeferredObject argument = arguments[0];
+    if (argument == null || argument.get() == null) {
+      return null;
+    } else {
+      evaluator.apply(argument);
+    }
+    return result;
+  }
+
+  @Override
+  public String getDisplayString(String[] children) {
+    return getStandardDisplayString("iceberg_day", children);
+  }
+
+}