[GitHub] [iceberg] rdblue commented on a diff in pull request #5431: Spark: Support truncate in FunctionCatalog

[GitBox <gi...@apache.org>]

rdblue commented on code in PR #5431:
URL: https://github.com/apache/iceberg/pull/5431#discussion_r938252617


##########
spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/TruncateFunction.java:
##########
@@ -0,0 +1,395 @@
+/*
+ * 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.spark.functions;
+
+import java.math.BigInteger;
+import java.nio.ByteBuffer;
+import java.nio.CharBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.util.ByteBuffers;
+import org.apache.iceberg.util.TruncateUtil;
+import org.apache.iceberg.util.UnicodeUtil;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.catalog.functions.BoundFunction;
+import org.apache.spark.sql.connector.catalog.functions.ScalarFunction;
+import org.apache.spark.sql.connector.catalog.functions.UnboundFunction;
+import org.apache.spark.sql.types.BinaryType;
+import org.apache.spark.sql.types.ByteType;
+import org.apache.spark.sql.types.CharType;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.types.DecimalType;
+import org.apache.spark.sql.types.IntegerType;
+import org.apache.spark.sql.types.LongType;
+import org.apache.spark.sql.types.ShortType;
+import org.apache.spark.sql.types.StringType;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.types.VarcharType;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/**
+ * A function for use in SQL that truncates input to a given width according to the rules of the
+ * existing Truncate partition transform, e.g. {@code SELECT system.truncate(1, 'abc')} will return
+ * the String 'a'.
+ *
+ * <p>Note that for performance reasons, the given input width is not validated in the
+ * implementations used in code-gen. The width must remain non-negative to give meaningful results.
+ */
+public class TruncateFunction implements UnboundFunction {
+  private static final List<DataType> truncateableAtomicTypes =
+      ImmutableList.of(
+          DataTypes.ByteType,
+          DataTypes.ShortType,
+          DataTypes.IntegerType,
+          DataTypes.LongType,
+          DataTypes.StringType,
+          DataTypes.BinaryType);
+
+  private static void validateTruncationFieldType(DataType dt) {
+    if (truncateableAtomicTypes.stream().noneMatch(type -> type.sameType(dt))
+        && !(dt instanceof DecimalType)) {
+      String expectedTypes =
+          "[ByteType, ShortType, IntegerType, LongType, StringType, BinaryType, DecimalType]";
+      throw new UnsupportedOperationException(
+          String.format(
+              "Invalid input type to truncate. Expected one of %s, but found %s",
+              expectedTypes, dt));
+    }
+  }
+
+  private static void validateTruncationWidthType(DataType widthType) {
+    if (!DataTypes.IntegerType.sameType(widthType)
+        && !DataTypes.ShortType.sameType(widthType)
+        && !DataTypes.ByteType.sameType(widthType)) {
+      throw new UnsupportedOperationException(
+          "Expected truncation width to be one of [ByteType, ShortType, IntegerType], but found "
+              + widthType);
+    }
+  }
+
+  @Override
+  public BoundFunction bind(StructType inputType) {
+    if (inputType.fields().length != 2) {
+      throw new UnsupportedOperationException(
+          String.format(
+              "Invalid input type. Expected 2 fields but found %s", inputType.fields().length));
+    }
+
+    StructField widthField = inputType.apply(0);
+    StructField toTruncateField = inputType.apply(1);
+
+    validateTruncationFieldType(toTruncateField.dataType());
+    validateTruncationWidthType(widthField.dataType());
+
+    DataType toTruncateDataType = toTruncateField.dataType();
+    if (toTruncateDataType instanceof ByteType) {
+      return new TruncateTinyInt();
+    } else if (toTruncateDataType instanceof ShortType) {
+      return new TruncateSmallInt();
+    } else if (toTruncateDataType instanceof IntegerType) {
+      return new TruncateInt();
+    } else if (toTruncateDataType instanceof LongType) {
+      return new TruncateBigInt();
+    } else if (toTruncateDataType instanceof DecimalType) {
+      return new TruncateDecimal(
+          ((DecimalType) toTruncateDataType).precision(),
+          ((DecimalType) toTruncateDataType).scale());
+    } else if (toTruncateDataType instanceof StringType
+        || toTruncateDataType instanceof VarcharType
+        || toTruncateDataType instanceof CharType) {
+      return new TruncateString();
+    } else if (toTruncateDataType instanceof BinaryType) {
+      return new TruncateBinary();
+    } else {
+      throw new UnsupportedOperationException("Cannot truncate type: " + toTruncateDataType);
+    }
+  }
+
+  @Override
+  public String description() {
+    return "Truncate - The Iceberg truncate function used for truncate partition transformations.\n"
+        + "\tCalled with the truncation width as the first argument: e.g. system.truncate(width, col)";
+  }
+
+  @Override
+  public String name() {
+    return "truncate";
+  }
+
+  public abstract static class TruncateBase<T> implements ScalarFunction<T> {
+    @Override
+    public String name() {
+      return "truncate";
+    }
+  }
+
+  public static class TruncateTinyInt extends TruncateBase<Byte> {
+    public static byte invoke(int width, byte value) {
+      return TruncateUtil.truncateByte(width, value);
+    }
+
+    @Override
+    public DataType[] inputTypes() {
+      return new DataType[] {DataTypes.IntegerType, DataTypes.ByteType};
+    }
+
+    @Override
+    public DataType resultType() {
+      return DataTypes.ByteType;
+    }
+
+    @Override
+    public String canonicalName() {
+      return "iceberg.truncate[width](tinyint)";

Review Comment:
   Looks like you forgot to actually embed width in the name.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org