cloud-fan commented on a change in pull request #32777: URL: https://github.com/apache/spark/pull/32777#discussion_r648071740
########## File path: sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java ########## @@ -0,0 +1,980 @@ +/* + * 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.spark.sql.execution.datasources.parquet; + +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.Dictionary; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.LogicalTypeAnnotation.IntLogicalTypeAnnotation; +import org.apache.parquet.schema.LogicalTypeAnnotation.DecimalLogicalTypeAnnotation; +import org.apache.parquet.schema.LogicalTypeAnnotation.TimestampLogicalTypeAnnotation; +import org.apache.parquet.schema.PrimitiveType; + +import org.apache.spark.sql.catalyst.util.DateTimeUtils; +import org.apache.spark.sql.catalyst.util.RebaseDateTime; +import org.apache.spark.sql.execution.datasources.DataSourceUtils; +import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException; +import org.apache.spark.sql.execution.vectorized.WritableColumnVector; +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.DecimalType; + +import java.math.BigInteger; +import java.time.ZoneId; +import java.time.ZoneOffset; +import java.util.Arrays; + +public class ParquetVectorUpdaterFactory { + private static final ZoneId UTC = ZoneOffset.UTC; + + private final LogicalTypeAnnotation logicalTypeAnnotation; + // The timezone conversion to apply to int96 timestamps. Null if no conversion. + private final ZoneId convertTz; + private final String datetimeRebaseMode; + private final String int96RebaseMode; + + ParquetVectorUpdaterFactory( + LogicalTypeAnnotation logicalTypeAnnotation, + ZoneId convertTz, + String datetimeRebaseMode, + String int96RebaseMode) { + this.logicalTypeAnnotation = logicalTypeAnnotation; + this.convertTz = convertTz; + this.datetimeRebaseMode = datetimeRebaseMode; + this.int96RebaseMode = int96RebaseMode; + } + + public ParquetVectorUpdater getUpdater(ColumnDescriptor descriptor, DataType sparkType) { + PrimitiveType.PrimitiveTypeName typeName = descriptor.getPrimitiveType().getPrimitiveTypeName(); + + switch (typeName) { + case BOOLEAN: + if (sparkType == DataTypes.BooleanType) { + return new BooleanUpdater(); + } + throw constructConvertNotSupportedException(descriptor, sparkType); + + case INT32: + if (sparkType == DataTypes.IntegerType || canReadAsIntDecimal(descriptor, sparkType)) { + return new IntegerUpdater(); + } else if (sparkType == DataTypes.LongType) { + // In `ParquetToSparkSchemaConverter`, we map parquet UINT32 to our LongType. + // For unsigned int32, it stores as plain signed int32 in Parquet when dictionary + // fallbacks. We read them as long values. + return new UnsignedIntegerUpdater(); + } else if (sparkType == DataTypes.ByteType) { + return new ByteUpdater(); + } else if (sparkType == DataTypes.ShortType) { + return new ShortUpdater(); + } else if (sparkType == DataTypes.DateType) { + if ("CORRECTED".equals(datetimeRebaseMode)) { + return new IntegerUpdater(); + } else { + boolean failIfRebase = "EXCEPTION".equals(datetimeRebaseMode); + return new IntegerWithRebaseUpdater(failIfRebase); + } + } + throw constructConvertNotSupportedException(descriptor, sparkType); + + case INT64: + // This is where we implement support for the valid type conversions. + if (sparkType == DataTypes.LongType || canReadAsLongDecimal(descriptor, sparkType)) { + return new LongUpdater(DecimalType.is32BitDecimalType(sparkType)); + } else if (isUnsignedIntTypeMatched(64)) { + // In `ParquetToSparkSchemaConverter`, we map parquet UINT64 to our Decimal(20, 0). + // For unsigned int64, it stores as plain signed int64 in Parquet when dictionary + // fallbacks. We read them as decimal values. + return new UnsignedLongUpdater(); + } else if (isTimestampTypeMatched(LogicalTypeAnnotation.TimeUnit.MICROS)) { + if ("CORRECTED".equals(datetimeRebaseMode)) { + return new LongUpdater(false); + } else { + boolean failIfRebase = "EXCEPTION".equals(datetimeRebaseMode); + return new LongWithRebaseUpdater(failIfRebase); + } + } else if (isTimestampTypeMatched(LogicalTypeAnnotation.TimeUnit.MILLIS)) { + if ("CORRECTED".equals(datetimeRebaseMode)) { + return new LongAsMicrosUpdater(); + } else { + final boolean failIfRebase = "EXCEPTION".equals(datetimeRebaseMode); + return new LongAsMicrosRebaseUpdater(failIfRebase); + } + } + throw constructConvertNotSupportedException(descriptor, sparkType); + + case FLOAT: + if (sparkType == DataTypes.FloatType) { + return new FloatUpdater(); + } + throw constructConvertNotSupportedException(descriptor, sparkType); + + case DOUBLE: + if (sparkType == DataTypes.DoubleType) { + return new DoubleUpdater(); + } + throw constructConvertNotSupportedException(descriptor, sparkType); + + case INT96: + case BINARY: + if (sparkType == DataTypes.StringType || sparkType == DataTypes.BinaryType || + canReadAsBinaryDecimal(descriptor, sparkType)) { + return new BinaryUpdater(); + } else if (sparkType == DataTypes.TimestampType) { Review comment: Yea let's fix them. -- 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. For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
