rdblue commented on code in PR #5305: URL: https://github.com/apache/iceberg/pull/5305#discussion_r930240935
########## spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/TruncateFunction.java: ########## @@ -0,0 +1,435 @@ +/* + * 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.BigDecimal; +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.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; + +/** + * Implementation of {@link UnboundFunction} that matches the <b>truncate</b> transformation. + * This unbound function is registered with the {@link org.apache.iceberg.spark.SparkCatalog} + * such that the function can be used as {@code truncate(width, col)} or {@code truncate(2, col)}. + * <p> + * Specific {@link BoundFunction} implementations are resolved via the analyzer + * depending on their input types. + */ +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"; + } + } + + // tinyint Review Comment: Nit: this comment duplicates what is already in the class 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
