tmater commented on code in PR #16714: URL: https://github.com/apache/iceberg/pull/16714#discussion_r3568567515
########## spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/SparkVariantExtractionReaders.java: ########## @@ -0,0 +1,483 @@ +/* + * 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.data; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.time.ZoneId; +import java.util.List; +import org.apache.iceberg.expressions.PathUtil; +import org.apache.iceberg.parquet.ParquetValueReader; +import org.apache.iceberg.parquet.ParquetVariantExtractionReaders; +import org.apache.iceberg.parquet.ParquetVariantExtractionReaders.VariantExtractionField; +import org.apache.iceberg.parquet.ParquetVariantExtractionReaders.VariantExtractionRow; +import org.apache.iceberg.parquet.ParquetVariantReaders.DelegatingValueReader; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.variants.PhysicalType; +import org.apache.iceberg.variants.VariantMetadata; +import org.apache.iceberg.variants.VariantValue; +import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.MessageType; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow; +import org.apache.spark.sql.catalyst.expressions.variant.VariantCastArgs; +import org.apache.spark.sql.internal.SQLConf; +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.types.TimestampNTZType; +import org.apache.spark.sql.types.TimestampType; +import org.apache.spark.unsafe.types.UTF8String; +import org.apache.spark.unsafe.types.VariantVal; + +/** Parquet readers that materialize Spark variant extraction struct rows from shredded variants. */ +public class SparkVariantExtractionReaders { + // Empty variant metadata dictionary, reused for delegating primitive-leaf casts (primitives do + // not reference the dictionary, so the per-row dictionary need not be serialized). + private static final byte[] EMPTY_VARIANT_METADATA = serializeMetadata(VariantMetadata.empty()); + + private SparkVariantExtractionReaders() {} + + public static ParquetValueReader<InternalRow> buildStructReader( + MessageType fileSchema, + GroupType variantGroup, + List<String> variantColumnPath, + StructType extractionStruct) { + List<VariantExtractionField> fields = Lists.newArrayList(); + int numFields = 0; + for (StructField field : extractionStruct.fields()) { + numFields = Math.max(numFields, Integer.parseInt(field.name()) + 1); + } + + DataType[] targetTypes = new DataType[numFields]; + boolean[] failOnError = new boolean[numFields]; + String[] timeZoneIds = new String[numFields]; + for (StructField field : extractionStruct.fields()) { + int ordinal = Integer.parseInt(field.name()); + targetTypes[ordinal] = field.dataType(); + failOnError[ordinal] = SparkVariantExtractionUtil.failOnError(field); + timeZoneIds[ordinal] = SparkVariantExtractionUtil.timeZoneId(field); + fields.add( + new VariantExtractionField( + ordinal, + SparkVariantExtractionUtil.isPlaceholderExtraction(field), + PathUtil.parse(SparkVariantExtractionUtil.extractionPath(field)))); + } + + ParquetValueReader<VariantExtractionRow> parquetReader = + ParquetVariantExtractionReaders.buildRowReader( + fileSchema, variantGroup, variantColumnPath, fields); + + return new SparkVariantExtractionStructReader( + parquetReader, targetTypes, failOnError, timeZoneIds); + } + + /** Counts leaf Parquet columns referenced by a reader tree. */ + public static int leafColumnCount(ParquetValueReader<?> reader) { + return ParquetVariantExtractionReaders.leafColumnCount(reader); + } + + /** Visible for unit tests in {@code org.apache.iceberg.spark.data}. */ + static Object toSparkValueForTests(VariantValue value, DataType targetType) { + return toSparkValueForTests(value, targetType, false); + } + + /** Visible for unit tests in {@code org.apache.iceberg.spark.data}. */ + static Object toSparkValueForTests(VariantValue value, DataType targetType, boolean failOnError) { + // Tests pass standalone primitive values, so an empty metadata dictionary and UTC suffice. + return toSparkValue(value, targetType, failOnError, VariantMetadata.empty(), "UTC"); + } + + private static class SparkVariantExtractionStructReader + extends DelegatingValueReader<VariantExtractionRow, InternalRow> { + private final DataType[] targetTypes; + private final boolean[] failOnError; + private final String[] timeZoneIds; + + private SparkVariantExtractionStructReader( + ParquetValueReader<VariantExtractionRow> reader, + DataType[] targetTypes, + boolean[] failOnError, + String[] timeZoneIds) { + super(reader); + this.targetTypes = targetTypes; + this.failOnError = failOnError; + this.timeZoneIds = timeZoneIds; + } + + @Override + public InternalRow read(InternalRow reuse) { + VariantExtractionRow row = readFromDelegate(null); + int numFields = row.numFields(); + GenericInternalRow result = + reuse instanceof GenericInternalRow + ? (GenericInternalRow) reuse + : new GenericInternalRow(numFields); + + if (row.metadata() == null) { + // SQL NULL variant: match Spark variant_get semantics by nulling every extraction slot, + // including full-variant placeholder slots. Distinct from a non-null variant where only + // missing or JSON-null paths produce per-field SQL NULLs. + for (int i = 0; i < numFields; i += 1) { + result.setNullAt(i); + } + return result; + } + + for (int i = 0; i < numFields; i += 1) { + if (row.placeholder(i)) { + result.setBoolean(i, true); + } else { + Object sparkValue = + toSparkValue( + row.value(i), targetTypes[i], failOnError[i], row.metadata(), timeZoneIds[i]); + if (sparkValue == null) { + result.setNullAt(i); + } else { + result.update(i, sparkValue); + } + } + } + + return result; + } + } + + /** + * Materializes a single extracted shredded value as the Spark {@code targetType}. + * + * <p>Uses a narrow inline conversion for the common, stable pairs ({@link #inlineOwned}) and + * delegates everything else to Spark's {@code VariantGet.cast}. The inline path exists purely for + * performance: it reads the shredded primitive directly, whereas delegating must serialize the + * value back into a Spark {@code VariantVal} and run a {@code Cast} — on the order of ~1 KB of + * transient allocation per value (vs. tens of bytes inline), which is significant on wide scans. + * It is intentionally kept to pairs that are either trivially Spark-identical (identity unwraps, + * integer range-checks) or that Spark's variant cannot represent at all and so <em>must</em> be + * handled here ({@code TIME}, nanos timestamps). Every cross-type, lossy, or zone-sensitive cast Review Comment: I missed the nanosecond timestamp part earlier. Looking at the current code, the pushed-down selective reader handles `TIMESTAMPTZ_NANOS` / `TIMESTAMPNTZ_NANOS` inline and converts nanos to Spark microseconds, while the comment also notes that Spark’s Variant cannot represent those values. That seems like it could introduce a semantic mismatch once this is used through the pushdown path: the same `variant_get(..., 'timestamp')` query could fail when Spark evaluates the Variant normally, but succeed when Iceberg accepts the extraction pushdown. -- 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]
