You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/08/18 22:48:44 UTC
(datafusion-comet) branch main updated: fix: preserve Spark semantics for dictionary-encoded Parquet inputs and reject dictionary targets (#5234)
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new a28ac348f5 fix: preserve Spark semantics for dictionary-encoded Parquet inputs and reject dictionary targets (#5234)
a28ac348f5 is described below
commit a28ac348f5a68500d96fe872b49e453d96e4bab6
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Wed Aug 19 06:48:40 2026 +0800
fix: preserve Spark semantics for dictionary-encoded Parquet inputs and reject dictionary targets (#5234)
* refactor: use Arrow casts for dictionary conversions
* fix: handle dictionary sources and reject dictionary targets
---
native/core/src/parquet/parquet_support.rs | 31 +-----
native/spark-expr/src/conversion_funcs/cast.rs | 119 +++++++--------------
.../apache/comet/parquet/ParquetReadSuite.scala | 59 ++++++++++
3 files changed, 99 insertions(+), 110 deletions(-)
diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs
index 56176b1637..5b22afa260 100644
--- a/native/core/src/parquet/parquet_support.rs
+++ b/native/core/src/parquet/parquet_support.rs
@@ -22,10 +22,10 @@ use arrow::compute::can_cast_types;
use arrow::datatypes::{FieldRef, Fields};
use arrow::{
array::{
- cast::AsArray, new_null_array, types::Int32Type, types::TimestampMicrosecondType, Array,
- ArrayRef, DictionaryArray, StructArray,
+ cast::AsArray, new_null_array, types::TimestampMicrosecondType, Array, ArrayRef,
+ StructArray,
},
- compute::{cast_with_options, take, CastOptions},
+ compute::{cast_with_options, CastOptions},
datatypes::{DataType, TimeUnit},
util::display::FormatOptions,
};
@@ -169,31 +169,6 @@ fn parquet_convert_array(
parquet_options: &SparkParquetOptions,
) -> DataFusionResult<ArrayRef> {
use DataType::*;
- let from_type = array.data_type().clone();
-
- let array = match &from_type {
- Dictionary(key_type, value_type)
- if key_type.as_ref() == &Int32
- && (value_type.as_ref() == &Utf8 || value_type.as_ref() == &LargeUtf8) =>
- {
- let dict_array = array
- .as_any()
- .downcast_ref::<DictionaryArray<Int32Type>>()
- .expect("Expected a dictionary array");
-
- let casted_dictionary = DictionaryArray::<Int32Type>::new(
- dict_array.keys().clone(),
- parquet_convert_array(Arc::clone(dict_array.values()), to_type, parquet_options)?,
- );
-
- let casted_result = match to_type {
- Dictionary(_, _) => Arc::new(casted_dictionary.clone()),
- _ => take(casted_dictionary.values().as_ref(), dict_array.keys(), None)?,
- };
- return Ok(casted_result);
- }
- _ => array,
- };
let from_type = array.data_type();
// Try Comet specific handlers first, then arrow-rs cast if supported,
diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs
index 504ea6695e..b1b71e8267 100644
--- a/native/spark-expr/src/conversion_funcs/cast.rs
+++ b/native/spark-expr/src/conversion_funcs/cast.rs
@@ -42,18 +42,17 @@ use crate::{cast_whole_num_to_binary, BinaryOutputStyle};
use crate::{EvalMode, SparkError};
use arrow::array::builder::{GenericStringBuilder, StringBuilder};
use arrow::array::{
- new_null_array, BinaryBuilder, DictionaryArray, GenericByteArray, ListArray, MapArray,
- StringArray, StructArray,
+ new_null_array, BinaryBuilder, GenericByteArray, ListArray, MapArray, StringArray, StructArray,
};
-use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Schema};
+use arrow::datatypes::{DataType, Schema};
use arrow::datatypes::{Field, Fields, GenericBinaryType};
use arrow::error::ArrowError;
use arrow::{
array::{
- cast::AsArray, types::Int32Type, Array, ArrayRef, Int16Array, Int32Array, Int64Array,
- Int8Array, OffsetSizeTrait, PrimitiveArray,
+ cast::AsArray, Array, ArrayRef, Int16Array, Int32Array, Int64Array, Int8Array,
+ OffsetSizeTrait,
},
- compute::{cast_with_options, take, CastOptions},
+ compute::{cast_with_options, CastOptions},
record_batch::RecordBatch,
util::display::FormatOptions,
};
@@ -213,40 +212,6 @@ pub fn spark_cast(
Ok(result)
}
-// copied from datafusion common scalar/mod.rs
-fn dict_from_values<K: ArrowDictionaryKeyType>(
- values_array: ArrayRef,
-) -> datafusion::common::Result<ArrayRef> {
- // Create a key array with `size` elements of 0..array_len for all
- // non-null value elements
- let key_array: PrimitiveArray<K> = (0..values_array.len())
- .map(|index| {
- if values_array.is_valid(index) {
- let native_index = K::Native::from_usize(index).ok_or_else(|| {
- DataFusionError::Internal(format!(
- "Can not create index of type {} from value {}",
- K::DATA_TYPE,
- index
- ))
- })?;
- Ok(Some(native_index))
- } else {
- Ok(None)
- }
- })
- .collect::<datafusion::common::Result<Vec<_>>>()?
- .into_iter()
- .collect();
-
- // create a new DictionaryArray
- //
- // Note: this path could be made faster by using the ArrayData
- // APIs and skipping validation, if it every comes up in
- // performance traces.
- let dict_array = DictionaryArray::<K>::try_new(key_array, values_array)?;
- Ok(Arc::new(dict_array))
-}
-
pub(crate) fn cast_array(
array: ArrayRef,
to_type: &DataType,
@@ -255,6 +220,12 @@ pub(crate) fn cast_array(
use DataType::*;
let from_type = array.data_type().clone();
+ // Spark's SQL data-type grammar cannot express Dictionary as a cast target:
+ // https://github.com/apache/spark/blob/v4.2.0/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4#L1477-L1525
+ if matches!(to_type, Dictionary(_, _)) {
+ return internal_err!("Spark cannot specify dictionary types as cast targets");
+ }
+
if &from_type == to_type {
return Ok(Arc::new(array));
}
@@ -269,47 +240,18 @@ pub(crate) fn cast_array(
.with_timestamp_format(TIMESTAMP_FORMAT),
};
- let array = match &from_type {
- Dictionary(key_type, value_type)
- if key_type.as_ref() == &Int32
- && (value_type.as_ref() == &Utf8
- || value_type.as_ref() == &LargeUtf8
- || value_type.as_ref() == &Binary
- || value_type.as_ref() == &LargeBinary) =>
- {
- let dict_array = array
- .as_any()
- .downcast_ref::<DictionaryArray<Int32Type>>()
- .expect("Expected a dictionary array");
-
- let casted_result = match to_type {
- Dictionary(_, to_value_type) => {
- let casted_dictionary = DictionaryArray::<Int32Type>::new(
- dict_array.keys().clone(),
- cast_array(Arc::clone(dict_array.values()), to_value_type, cast_options)?,
- );
- Arc::new(casted_dictionary.clone())
- }
- _ => {
- let casted_dictionary = DictionaryArray::<Int32Type>::new(
- dict_array.keys().clone(),
- cast_array(Arc::clone(dict_array.values()), to_type, cast_options)?,
- );
- take(casted_dictionary.values().as_ref(), dict_array.keys(), None)?
- }
- };
+ // Spark infers Parquet schemas from its own metadata or the Parquet MessageType, not
+ // ARROW:schema, so Arrow can expose a dictionary source while Spark requests its value type:
+ // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala#L585-L599
+ if let Dictionary(_, value_type) = &from_type {
+ if matches!(value_type.as_ref(), Utf8 | LargeUtf8 | Binary | LargeBinary) {
+ let dictionary = array.as_any_dictionary();
+ let values = cast_array(Arc::clone(dictionary.values()), to_type, cast_options)?;
+ let dictionary = dictionary.with_values(values);
+ let casted_result = cast_with_options(&dictionary, to_type, &native_cast_options)?;
return Ok(spark_cast_postprocess(casted_result, &from_type, to_type));
}
- _ => {
- if let Dictionary(_, _) = to_type {
- let dict_array = dict_from_values::<Int32Type>(array)?;
- let casted_result = cast_array(dict_array, to_type, cast_options)?;
- return Ok(spark_cast_postprocess(casted_result, &from_type, to_type));
- } else {
- array
- }
- }
- };
+ }
let cast_result = match (&from_type, to_type) {
// Null arrays carry no concrete values, so Arrow's native cast can change only the
@@ -905,10 +847,23 @@ fn cast_binary_to_string<O: OffsetSizeTrait>(
#[cfg(test)]
mod tests {
use super::*;
- use arrow::array::{BinaryArray, ListArray, NullArray, StringArray};
+ use arrow::array::{BinaryArray, ListArray, NullArray, PrimitiveArray, StringArray};
use arrow::buffer::OffsetBuffer;
- use arrow::datatypes::TimestampMicrosecondType;
- use arrow::datatypes::{Field, Fields};
+ use arrow::datatypes::{Field, Fields, Int32Type, TimestampMicrosecondType};
+
+ #[test]
+ fn test_cast_to_dictionary_is_rejected() {
+ let error = cast_array(
+ Arc::new(StringArray::from(vec!["a"])),
+ &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
+ &SparkCastOptions::new(EvalMode::Legacy, "UTC", false),
+ )
+ .unwrap_err();
+
+ assert!(error
+ .to_string()
+ .contains("Spark cannot specify dictionary types as cast targets"));
+ }
#[test]
fn test_cast_binary_to_string_replaces_invalid_utf8_jvm_compatibly() {
diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
index 684c4d6581..b256c917a1 100644
--- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
@@ -22,6 +22,7 @@ package org.apache.comet.parquet
import java.io.File
import java.math.{BigDecimal, BigInteger}
import java.time.{ZoneId, ZoneOffset}
+import java.util.{Base64, Collections}
import scala.reflect.ClassTag
import scala.reflect.runtime.universe.TypeTag
@@ -29,8 +30,11 @@ import scala.reflect.runtime.universe.TypeTag
import org.scalactic.source.Position
import org.scalatest.Tag
+import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field => ArrowField, FieldType, Schema => ArrowSchema}
import org.apache.hadoop.fs.Path
import org.apache.parquet.example.data.simple.SimpleGroup
+import org.apache.parquet.hadoop.example.ExampleParquetWriter
+import org.apache.parquet.io.api.Binary
import org.apache.parquet.schema.MessageTypeParser
import org.apache.spark.SparkException
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
@@ -81,6 +85,61 @@ abstract class ParquetReadSuite extends CometTestBase {
}
}
+ // Spark ignores ARROW:schema during Parquet schema inference:
+ // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala#L585-L599
+ // With binaryAsString, Spark maps unannotated BINARY to StringType:
+ // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala#L344-L350
+ test("native scan casts Arrow dictionary binary values with Spark semantics") {
+ withTempDir { dir =>
+ val path = new Path(dir.toURI.toString, "dictionary-binary.parquet")
+ val parquetSchema = MessageTypeParser.parseMessageType("""message root {
+ | optional binary value;
+ |}
+ |""".stripMargin)
+ val arrowField = new ArrowField(
+ "value",
+ new FieldType(
+ true,
+ ArrowType.Binary.INSTANCE,
+ new DictionaryEncoding(0L, false, new ArrowType.Int(32, true))),
+ Collections.emptyList[ArrowField]())
+ val arrowSchema = new ArrowSchema(Collections.singletonList(arrowField))
+ val metadata = Collections.singletonMap(
+ "ARROW:schema",
+ Base64.getEncoder.encodeToString(arrowSchema.serializeAsMessage()))
+ val writer = ExampleParquetWriter
+ .builder(path)
+ .withType(parquetSchema)
+ .withDictionaryEncoding(true)
+ .withExtraMetaData(metadata)
+ .withConf(spark.sessionState.newHadoopConf())
+ .build()
+
+ try {
+ Seq(
+ Array[Byte](0x66, 0x80.toByte, 0x6f),
+ Array[Byte](0x66, 0x80.toByte, 0x6f),
+ Array[Byte](0x76, 0x61, 0x6c, 0x69, 0x64)).foreach { bytes =>
+ val row = new SimpleGroup(parquetSchema)
+ row.add(0, Binary.fromConstantByteArray(bytes))
+ writer.write(row)
+ }
+ } finally {
+ writer.close()
+ }
+
+ withSQLConf(SQLConf.PARQUET_BINARY_AS_STRING.key -> "true") {
+ withParquetTable(path.toString, "dictionary_binary") {
+ val (_, cometPlan) =
+ checkSparkAnswerAndOperator(sql("SELECT value FROM dictionary_binary"))
+ assert(
+ collect(cometPlan) { case scan: CometNativeScanExec => scan }.nonEmpty,
+ "Expected a CometNativeScanExec")
+ }
+ }
+ }
+ }
+
test("basic data types") {
Seq(7, 1024).foreach { batchSize =>
withSQLConf(CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org