fix: preserve Spark semantics for dictionary-encoded Parquet inputs and reject dictionary targets

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,

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
@@ -908,10 +850,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() {

spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala

@@ -22,15 +22,19 @@ 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
 
 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) {