sunchao commented on code in PR #5407: URL: https://github.com/apache/datafusion-comet/pull/5407#discussion_r3866319241
########## native/core/src/parquet/cast_column/variant.rs: ########## @@ -0,0 +1,1386 @@ +// 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. +use arrow::{ + array::{ + make_array, Array, ArrayRef, AsArray, BinaryArray, BinaryBuilder, ListLikeArray, + StructArray, + }, + buffer::NullBuffer, + compute::{cast, cast_with_options}, + datatypes::{DataType, FieldRef, TimeUnit}, + error::ArrowError, +}; +use datafusion::common::{ + format::DEFAULT_CAST_OPTIONS, DataFusionError, Result as DataFusionResult, +}; +use parquet::variant::{ + unshred_variant, BorrowedShreddingState, ListBuilder, MetadataBuilder, ObjectBuilder, + ParentState, ReadOnlyMetadataBuilder, ValueBuilder, Variant, VariantArray, VariantBuilder, + VariantDecimal4, VariantDecimal8, VariantMetadata, WritableMetadataBuilder, +}; +use std::{ + collections::HashSet, + panic::{catch_unwind, AssertUnwindSafe}, + sync::Arc, +}; + +pub(super) fn normalize_variant_array( + array: &ArrayRef, + target_field: &FieldRef, +) -> DataFusionResult<ArrayRef> { + let DataType::Struct(fields) = target_field.data_type() else { + return Err(DataFusionError::Execution( + "Variant extension field must use Struct storage".to_string(), + )); + }; + if fields.len() != 2 + || fields[0].name() != "value" + || fields[1].name() != "metadata" + || fields + .iter() + .any(|field| field.data_type() != &DataType::Binary) + { + return Err(DataFusionError::Execution( + "Variant output must contain Binary children [value, metadata]".to_string(), + )); + } + + let array = normalize_variant_storage(array)?; + let variant = VariantArray::try_new(array.as_ref())?; + let was_shredded = variant.typed_value_field().is_some(); + let unshredded = unshred_variant_for_spark(&variant)?; + let value = unshredded.value_field().ok_or_else(|| { + DataFusionError::Execution("Unshredded Variant is missing its value field".to_string()) + })?; + let value = cast(value.as_ref(), &DataType::Binary)?; + let metadata = cast(unshredded.metadata_field().as_ref(), &DataType::Binary)?; + let (value, metadata) = if was_shredded { + rebuild_shredded_variant_for_spark(&variant, &value, &metadata, unshredded.inner().nulls())? + } else { + let value = reorder_variant_values( + &value, + &metadata, + unshredded.inner().nulls(), + VariantObjectKeyOrder::SparkUtf16, + false, + )?; + (value, metadata) + }; + let output = StructArray::try_new( + fields.clone(), + vec![value, metadata], + unshredded.inner().nulls().cloned(), + )?; + Ok(Arc::new(output)) +} + +fn unshred_variant_for_spark(variant: &VariantArray) -> DataFusionResult<VariantArray> { + let first_error = match unshred_variant(variant) { + Ok(array) => return Ok(array), + Err(error) => DataFusionError::from(error), + }; + + if let Ok(prepared) = prepare_variant_for_unshredding(variant, None) { + if let Ok(array) = unshred_variant(&prepared) { + return Ok(array); + } + } + let Some(metadata) = canonicalize_spark_empty_key_metadata(variant)? else { + return Err(first_error); + }; + let Ok(prepared) = prepare_variant_for_unshredding(variant, Some(metadata.as_binary::<i32>())) + else { + return Err(first_error); + }; + match unshred_variant(&prepared) { + Ok(array) => Ok(array), + Err(_) => Err(first_error), + } +} + +fn normalize_variant_type(data_type: &DataType) -> Option<DataType> { + fn normalize_field(field: &FieldRef) -> Option<FieldRef> { + normalize_variant_type(field.data_type()) + .map(|data_type| Arc::new(field.as_ref().clone().with_data_type(data_type))) + } + + match data_type { + DataType::Dictionary(_, value_type) => { + Some(normalize_variant_type(value_type).unwrap_or_else(|| value_type.as_ref().clone())) + } + DataType::UInt8 => Some(DataType::Int16), + DataType::UInt16 => Some(DataType::Int32), + DataType::UInt32 => Some(DataType::Int64), + // Spark reads Parquet UINT_64 as Decimal(20, 0). This is lossless for the full range and + // lets the existing Spark-compatible rebuild choose the Variant decimal width per value. + DataType::UInt64 => Some(DataType::Decimal128(20, 0)), + DataType::Timestamp(TimeUnit::Millisecond, timezone) => { + Some(DataType::Timestamp(TimeUnit::Microsecond, timezone.clone())) + } + DataType::FixedSizeList(field, _) => Some(DataType::List( + normalize_field(field).unwrap_or_else(|| Arc::clone(field)), + )), + DataType::List(field) => normalize_field(field).map(DataType::List), + DataType::LargeList(field) => normalize_field(field).map(DataType::LargeList), + DataType::ListView(field) => normalize_field(field).map(DataType::ListView), + DataType::LargeListView(field) => normalize_field(field).map(DataType::LargeListView), + DataType::Struct(fields) => { + let mut changed = false; + let fields = fields + .iter() + .map(|field| match normalize_field(field) { + Some(field) => { + changed = true; + field + } + None => Arc::clone(field), + }) + .collect::<Vec<_>>(); + changed.then(|| DataType::Struct(fields.into())) + } + _ => None, Review Comment: [P2] Preserve binary semantics for unannotated fixed-length children This leaves `FixedSizeBinary` unchanged, although Spark reads an unannotated Parquet `FIXED_LEN_BYTE_ARRAY` typed child as Binary. On this exact head with Spark 4.0.4, an ordinary Parquet `v` group containing binary metadata and `typed_value: FIXED_LEN_BYTE_ARRAY(16)` (bytes `00..0f`, no UUID annotation or `ARROW:schema`) returns `BINARY` and `"AAECAwQFBgcICQoLDA0ODw=="` under both Spark readers, but an asserted `CometNativeScanExec` returns `UUID` and `"00010203-0405-0607-0809-0a0b0c0d0e0f"`. Lengths 8 and 20 instead fail with `Illegal shredded value type: FixedSizeBinary(...)` while Spark succeeds; ordinary BINARY is a passing control. These scans are admitted with `allowReadingShredded=true`, pushdown disabled and nanosAsLong false. Please preserve the physical binary interpretation before unshredding, or fall back for these inputs. ########## spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala: ########## @@ -85,6 +98,435 @@ abstract class ParquetReadSuite extends CometTestBase { } } + test("native scan projects Variant through a Spark-compatible vector") { + assume(CometSparkSessionExtensions.isSpark40Plus, "VariantType requires Spark 4.0+") + + Seq(false, true).foreach { shredded => + withTable("variant_projection") { + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + "spark.sql.variant.writeShredding.enabled" -> shredded.toString, + "spark.sql.variant.forceShreddingSchemaForTest" -> "a BIGINT") { + sql("CREATE TABLE variant_projection(id INT, v VARIANT, tail STRING) USING parquet") + sql("""INSERT INTO variant_projection VALUES + |(1, parse_json('{"b": "hello", "a": 10}'), 'object'), + |(2, parse_json('[1, true, "x"]'), 'array'), + |(3, parse_json('42'), 'scalar'), + |(4, parse_json('null'), 'json-null'), + |(5, CAST(NULL AS VARIANT), 'sql-null')""".stripMargin) + } + + val queries = Seq( + "SELECT v FROM variant_projection" -> 0, + "SELECT id, v, tail FROM variant_projection" -> 1) + var expected = Seq.empty[Seq[Seq[Any]]] + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + "spark.sql.variant.allowReadingShredded" -> "true") { + expected = queries.map { case (query, variantOrdinal) => + normalizedVariantRows(sql(query), variantOrdinal) + } + } + + // Phase A handles only whole values; Spark's pushed VariantStruct remains a fallback. + withSQLConf( + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true", + "spark.sql.variant.allowReadingShredded" -> "true", + "spark.sql.variant.pushVariantIntoScan" -> "false") { + val plans = queries.zip(expected).map { case ((query, variantOrdinal), expectedRows) => + val df = sql(query) + assert(normalizedVariantRows(df, variantOrdinal) == expectedRows) + df.queryExecution.executedPlan + } + + if (!shredded) { + queries.foreach { case (query, _) => checkSparkAnswerAndOperator(sql(query)) } + } + + plans.foreach { cometPlan => + assert(collect(cometPlan) { case _: CometNativeScanExec => true }.size == 1) + assert(collect(cometPlan) { case _: CometNativeColumnarToRowExec => true }.isEmpty) + assert(collect(cometPlan) { case _: CometColumnarToRowExec => true }.nonEmpty) + } + + val scan = collect(plans.head) { case scan: CometNativeScanExec => scan }.head + + val summaries = scan + .executeColumnar() + .mapPartitions { batches => + batches.map { batch => + try { + val vector = batch.column(0) + val struct = vector.asInstanceOf[CometStructVector] + val field = struct.getValueVector.getField + val getVariant = struct.getClass.getMethod("getVariant", Integer.TYPE) + val values = (0 until batch.numRows()).map { rowId => + if (struct.isNullAt(rowId)) { + None + } else { + Some(getVariant.invoke(struct, Int.box(rowId)).toString) + } + } + ( + Utils.isVariantType(struct.dataType()), + field.getName, + field.isNullable, + field.getMetadata.get("ARROW:extension:name"), + field.getChildren.asScala.map(_.getName).toSeq, + Seq(struct.getChild(0).dataType(), struct.getChild(1).dataType()), + values) + } finally { + batch.close() + } + } + } + .collect() + + assert(summaries.nonEmpty) + summaries.foreach { + case (isVariant, name, nullable, extension, children, childTypes, _) => + assert(isVariant) + assert(name == "v") + assert(nullable) + assert(extension == "arrow.parquet.variant") + assert(children == Seq("value", "metadata")) + assert(childTypes == Seq(BinaryType, BinaryType)) + } + val values = summaries.flatMap(_._7) + assert(values.count(_.isEmpty) == 1) + assert( + values.flatten.toSet == + Set("{\"a\":10,\"b\":\"hello\"}", "[1,true,\"x\"]", "42", "null")) + } + } + } + } + + test("native scan honors Spark's shredded Variant reader configuration") { + assume(CometSparkSessionExtensions.isSpark40Plus, "VariantType requires Spark 4.0+") + + withTable("variant_reader_mode") { + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + "spark.sql.variant.allowReadingShredded" -> "true", + "spark.sql.variant.writeShredding.enabled" -> "true", + "spark.sql.variant.forceShreddingSchemaForTest" -> "a BIGINT") { + sql("CREATE TABLE variant_reader_mode(v VARIANT) USING parquet") + sql("INSERT INTO variant_reader_mode VALUES (parse_json('{\"a\":1}'))") + } + + val key = "spark.sql.variant.allowReadingShredded" + val conf = SQLConf.get + val original = conf.getAllConfs.get(key) + try { + Seq(None, Some(false), Some(true)).foreach { configured => + configured match { + case Some(value) => conf.setConfString(key, value.toString) + case None => conf.unsetConf(key) + } + val allowShredded = conf.getConfString(key).toBoolean + + withSQLConf( + "spark.sql.variant.pushVariantIntoScan" -> "false", + "spark.sql.variant.inferShreddingSchema" -> "false") { + val result = sql("SELECT v FROM variant_reader_mode") + val nativeScans = collect(result.queryExecution.executedPlan) { + case _: CometNativeScanExec => true + } + + if (allowShredded) { + assert(normalizedVariantRows(result, 0) == Seq(Seq("{\"a\":1}"))) + assert(nativeScans.size == 1) + } else { + assert(nativeScans.isEmpty) + val error = intercept[SparkException](result.collect()).getCause + assert(error.isInstanceOf[AnalysisException]) + assert( + error.asInstanceOf[AnalysisException].getErrorClass == + "INVALID_VARIANT_FROM_PARQUET.WRONG_NUM_FIELDS") + } + } + } + } finally { + original match { + case Some(value) => conf.setConfString(key, value) + case None => conf.unsetConf(key) + } + } + } + } + + test("nanosAsLong Variant children fall back to Spark") { + assume(CometSparkSessionExtensions.isSpark40Plus, "VariantType requires Spark 4.0+") + + withTempDir { dir => + val rawNanos = Seq(1704067200123000000L, 1704067200123456789L) + val variant = sql("SELECT parse_json('0')").head().get(0) + val metadata = + variant.getClass.getMethod("getMetadata").invoke(variant).asInstanceOf[Array[Byte]] + + Seq(true, false).foreach { adjustedToUtc => + val path = new Path(dir.toURI.toString, s"nanos-$adjustedToUtc.parquet") + val parquetSchema = MessageTypeParser.parseMessageType(s"""message root { + | optional group v { + | optional binary value; + | required binary metadata; + | optional int64 typed_value (TIMESTAMP(NANOS,$adjustedToUtc)); + | } + |} + |""".stripMargin) + val writer = ExampleParquetWriter + .builder(path) + .withType(parquetSchema) + .withConf(spark.sessionState.newHadoopConf()) + .build() + + try { + rawNanos.foreach { value => + val row = new SimpleGroup(parquetSchema) + val group = row.addGroup(0) + group.add(1, Binary.fromConstantByteArray(metadata)) + group.add(2, value) + writer.write(row) + } + } finally { + writer.close() + } + } + + withTable("variant_nanos_as_long") { + sql(s"""CREATE TABLE variant_nanos_as_long(v VARIANT) + |USING parquet LOCATION '${dir.getCanonicalPath}'""".stripMargin) + withSQLConf( + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true", + "spark.sql.legacy.parquet.nanosAsLong" -> "true", + "spark.sql.variant.allowReadingShredded" -> "true", + "spark.sql.variant.pushVariantIntoScan" -> "false", + "spark.sql.variant.inferShreddingSchema" -> "false") { + val result = sql("SELECT v FROM variant_nanos_as_long") + assert( + normalizedVariantRows(result, 0) == + rawNanos.flatMap(value => Seq.fill(2)(Seq(value.toString)))) + + val plan = result.queryExecution.executedPlan + assert(collect(plan) { case _: CometNativeScanExec => true }.isEmpty) + val fallbackReasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert( + fallbackReasons.exists(_.contains("spark.sql.legacy.parquet.nanosAsLong")), + s"Expected nanosAsLong fallback, found: ${fallbackReasons.mkString(", ")}") + } + } + } + } + + test("native scan preserves Variant existence default pairing") { + assume(CometSparkSessionExtensions.isSpark40Plus, "VariantType requires Spark 4.0+") + + withTable("variant_defaults") { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql("CREATE TABLE variant_defaults(v VARIANT DEFAULT parse_json('1')) USING parquet") + sql("INSERT INTO variant_defaults VALUES (parse_json('42'))") + sql("ALTER TABLE variant_defaults ADD COLUMNS(n INT DEFAULT 7)") + } + + withSQLConf("spark.sql.variant.pushVariantIntoScan" -> "false") { Review Comment: [P2] Opt both native-default tests into the permissive reader This scope and the native-read scope in `native scan fills a Variant existence default for an old Parquet file` only disable pushdown. Spark 4.0 registers `allowReadingShredded=false`, so the corrected production guard now falls back: the pairing test fails its native-scan count, and the missing-column test reaches Spark's unsupported vectorized Variant-default assignment. Both fail in the exact-head Spark 4.0.4 focused run; the current [macOS scans job](https://github.com/apache/datafusion-comet/actions/runs/33004921883/job/98299172370) also records both failures. An external subclass reusing these unchanged tests and changing only `sparkConf` to set `allowReadingShredded=true` passes 2/2. Please explicitly set that option in both native-read scopes and retain the separate unset/false/true strict-reader regression. ########## spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala: ########## @@ -102,6 +118,31 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with CometTypeS withFallbackReason(scanExec, "Full native scan disabled because ignoreMissingFiles enabled") } + if (serializeExistenceDefaultValues(scanExec.requiredSchema, scanExec.output).isEmpty) { + withFallbackReason(scanExec, unsupportedDefaultReason) + } + + val hasVariant = + scanExec.requiredSchema.fields.exists(field => isVariantType(field.dataType)) + + // Spark's strict mode validates the legacy two-field Variant layout and reports SPARK-47546 + // errors itself: https://issues.apache.org/jira/browse/SPARK-47546 + if (hasVariant && + !SQLConf.get.getConfString("spark.sql.variant.allowReadingShredded").toBoolean) { + withFallbackReason( + scanExec, + "Full native scan disabled because Spark's strict unshredded Variant reader is enabled") + } + + // Spark interprets TIMESTAMP(NANOS) leaves as raw longs in this legacy mode. The logical + // Variant schema does not expose whether a file contains such a shredded child, so preserve + // Spark semantics by falling back before reading any requested Variant value. + if (hasVariant && SQLConf.get.legacyParquetNanosAsLong) { Review Comment: [P2] Honor disabled timestamp-NTZ inference inside Variant `nanosAsLong` is not the only reader setting that changes a shredded leaf's logical type. With `spark.sql.parquet.inferTimestampNTZ.enabled=false`, Spark 4.0.4's vectorized reader treats `TIMESTAMP(MICROS,false)` as `TIMESTAMP`, while this admitted native Variant path preserves Arrow's absent timezone and returns `TIMESTAMP_NTZ`. I reproduced this with raw value `1704067200123456`, `allowReadingShredded=true`, pushdown disabled and nanosAsLong false: `schema_of_variant(v)` differs, and in `America/Los_Angeles` the STRING results are `2023-12-31 16:00:00.123456` versus native `2024-01-01 00:00:00.123456`. Every Comet query asserted one native scan. MILLIS reproduces too; default-inference and adjusted-to-UTC controls agree. The baseline here is explicitly vectorized Spark; its row reader ignores this setting for Variant. Please propagate this interpretation or gate Variant scans when it cannot be honored. -- 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]
