From 35f4206bbad836f6a09e24711eccb39d2c955e7f Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Sat, 12 Sep 2026 11:43:02 +0700 Subject: [PATCH] feat: run length on binary input natively Register datafusion-spark's length function in the native session so length, char_length, character_length and len return the byte count for binary input and the character count for strings, always as Int32. Drop the BinaryType fallback from CometLength. The native Parquet scan casts dictionary-typed columns to the required Spark type before any expression sees them, so the replaced character_length's dictionary handling is not needed; a Rust roundtrip test and a Scala test over an Arrow-written dictionary-typed file pin that. --- .../expression-audits/string_funcs.md | 2 +- native/core/src/execution/jni_api.rs | 36 +++++++++ native/core/src/parquet/schema_adapter.rs | 46 +++++++++++- .../org/apache/comet/serde/strings.scala | 9 +-- .../expressions/string/bit_length.sql | 3 +- .../sql-tests/expressions/string/length.sql | 71 ++++++++++++++++++ .../expressions/string/octet_length.sql | 3 +- .../dictionary-string-binary.parquet | Bin 0 -> 751 bytes .../comet/CometStringExpressionSuite.scala | 32 ++++++++ 9 files changed, 187 insertions(+), 15 deletions(-) create mode 100644 spark/src/test/resources/test-data/dictionary-string-binary.parquet diff --git a/docs/source/contributor-guide/expression-audits/string_funcs.md b/docs/source/contributor-guide/expression-audits/string_funcs.md index d9a4c92bd69..12721470b4a 100644 --- a/docs/source/contributor-guide/expression-audits/string_funcs.md +++ b/docs/source/contributor-guide/expression-audits/string_funcs.md @@ -124,7 +124,7 @@ ## length - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `(StringType|BinaryType) -> IntegerType`; eval returns `numChars` for strings and `.length` for binary. `BinaryType` input falls back via `Unsupported` (DataFusion's `character_length` accepts string types only). +- Spark 3.5.8 (audited 2026-05-27): baseline. `(StringType|BinaryType) -> IntegerType`; eval returns `numChars` for strings and `.length` for binary. `BinaryType` input runs natively since Comet registers `datafusion-spark`'s `length`, which returns the byte count for binary and the character count for strings (audit refreshed 2026-09-12). - Spark 4.0.1 (audited 2026-05-27): `inputTypes` widened to `StringTypeWithCollation(supportsTrimCollation = true)`; semantics unchanged. Non-default collations not honoured by Comet ([#4496](https://github.com/apache/datafusion-comet/issues/4496)). ## lower diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 65a2d68ec18..b57fd8c5a5f 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -68,6 +68,7 @@ use datafusion_spark::function::math::trigonometry::SparkSec; use datafusion_spark::function::math::width_bucket::SparkWidthBucket; use datafusion_spark::function::string::char::CharFunc; use datafusion_spark::function::string::concat::SparkConcat; +use datafusion_spark::function::string::length::SparkLengthFunc; use datafusion_spark::function::string::luhn_check::SparkLuhnCheck; use datafusion_spark::function::string::space::SparkSpace; use datafusion_spark::function::string::substring::SparkSubstring; @@ -775,6 +776,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) { session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitShift::right_unsigned())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSoundex::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSubstring::default())); + session_ctx.register_udf(ScalarUDF::new_from_impl(SparkLengthFunc::default())); } /// Prepares arrow arrays for output. @@ -1589,7 +1591,11 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( #[cfg(test)] mod tests { use super::*; + use arrow::datatypes::{DataType, Field}; use datafusion::execution::memory_pool::{MemoryConsumer, UnboundedMemoryPool}; + use datafusion::execution::FunctionRegistry; + use datafusion::logical_expr::type_coercion::functions::fields_with_udf; + use datafusion::logical_expr::ReturnFieldArgs; fn entry_count(thread_id: u64) -> usize { get_thread_memory_pools() @@ -1650,4 +1656,34 @@ mod tests { drop(pool); assert!(weak.upgrade().is_none()); } + + #[test] + fn length_resolves_to_spark_length_for_string_and_binary() { + let ctx = SessionContext::new(); + register_datafusion_spark_function(&ctx); + for name in ["length", "char_length", "character_length"] { + let udf = ctx.udf(name).unwrap(); + for input in [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + ] { + let arg = Arc::new(Field::new("arg0", input.clone(), true)); + // The Spark signature accepts the type as-is, so no cast is injected. + let coerced = fields_with_udf(std::slice::from_ref(&arg), udf.as_ref()) + .unwrap_or_else(|e| panic!("{name}({input}) rejected: {e}")); + assert_eq!(coerced[0].data_type(), &input); + let ret = udf + .return_field_from_args(ReturnFieldArgs { + arg_fields: &[arg], + scalar_arguments: &[None], + }) + .unwrap(); + assert_eq!(ret.data_type(), &DataType::Int32, "{name}({input})"); + } + } + } } diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index 5a4b3e16305..694e3ea6e87 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -1546,14 +1546,15 @@ mod test { use arrow::array::cast::AsArray; use arrow::array::UInt32Array; use arrow::array::{ - Array, ArrayRef, BinaryArray, Date32Array, Decimal128Array, FixedSizeListArray, - Float32Array, Float64Array, Int32Array, Int64Array, LargeListArray, ListArray, MapArray, - StringArray, StructArray, TimestampMicrosecondArray, TimestampMillisecondArray, + Array, ArrayRef, BinaryArray, Date32Array, Decimal128Array, DictionaryArray, + FixedSizeListArray, Float32Array, Float64Array, Int32Array, Int64Array, LargeListArray, + ListArray, MapArray, StringArray, StructArray, TimestampMicrosecondArray, + TimestampMillisecondArray, }; use arrow::buffer::OffsetBuffer; use arrow::datatypes::SchemaRef; use arrow::datatypes::{ - DataType, Field, Fields, Int64Type, Schema, TimeUnit, TimestampMicrosecondType, + DataType, Field, Fields, Int32Type, Int64Type, Schema, TimeUnit, TimestampMicrosecondType, }; use arrow::record_batch::RecordBatch; use datafusion::common::DataFusionError; @@ -2180,6 +2181,43 @@ mod test { stream.next().await.unwrap() } + /// A file whose Arrow schema declares dictionary-typed string and binary columns reads + /// back as plain `Utf8` and `Binary` when the required schema asks for them, so scalar + /// functions downstream, `length` among them, never see a `Dictionary` array. + #[tokio::test] + async fn dictionary_columns_are_unwrapped_to_required_types() -> Result<(), DataFusionError> { + let strings: DictionaryArray = + vec![Some("hello"), Some("\u{e9}"), None, Some("")] + .into_iter() + .collect(); + let keys = Int32Array::from(vec![Some(0), Some(1), None, Some(0)]); + let values = BinaryArray::from(vec![&b"hello"[..], &[0xff, 0x00][..]]); + let binaries = DictionaryArray::::try_new(keys, Arc::new(values))?; + let dictionary = + |value: DataType| DataType::Dictionary(Box::new(DataType::Int32), Box::new(value)); + let schema = Arc::new(Schema::new(vec![ + Field::new("s", dictionary(DataType::Utf8), true), + Field::new("b", dictionary(DataType::Binary), true), + ])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(strings), Arc::new(binaries)])?; + let required_schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("b", DataType::Binary, true), + ])); + + let read = roundtrip(&batch, required_schema).await?; + + assert_eq!(read.column(0).data_type(), &DataType::Utf8); + assert_eq!(read.column(1).data_type(), &DataType::Binary); + let strings = read.column(0).as_string::(); + assert_eq!(strings.value(1), "\u{e9}"); + assert!(strings.is_null(2)); + let binaries = read.column(1).as_binary::(); + assert_eq!(binaries.value(1), &[0xff, 0x00]); + assert!(binaries.is_null(2)); + Ok(()) + } + /// Build a one-column batch `s: struct` holding `values`, for the nested /// conversion tests (#5671). fn struct_batch(field: Field, values: ArrayRef) -> Result { diff --git a/spark/src/main/scala/org/apache/comet/serde/strings.scala b/spark/src/main/scala/org/apache/comet/serde/strings.scala index 5adb6aab7db..2b0b1319aa1 100644 --- a/spark/src/main/scala/org/apache/comet/serde/strings.scala +++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala @@ -82,14 +82,7 @@ object CometUpper extends CometCaseConversionBase[Upper]("upper") object CometLower extends CometCaseConversionBase[Lower]("lower") -object CometLength extends CometScalarFunction[Length]("length") { - override def getUnsupportedReasons(): Seq[String] = Seq("`BinaryType` input is not supported") - - override def getSupportLevel(expr: Length): SupportLevel = expr.child.dataType match { - case _: BinaryType => Unsupported(Some("Length on BinaryType is not supported")) - case _ => Compatible() - } -} +object CometLength extends CometScalarFunction[Length]("length") object CometBitLength extends CometScalarFunction[BitLength]("bit_length") { override def getUnsupportedReasons(): Seq[String] = Seq("`BinaryType` input is not supported") diff --git a/spark/src/test/resources/sql-tests/expressions/string/bit_length.sql b/spark/src/test/resources/sql-tests/expressions/string/bit_length.sql index a2327120163..64eb762aa0a 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/bit_length.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/bit_length.sql @@ -29,7 +29,8 @@ query SELECT bit_length('hello'), bit_length(''), bit_length(NULL) -- BinaryType input falls back to Spark; the native DataFusion impl rejects Binary at runtime, --- so the serde gates Binary as Unsupported (matching the existing CometLength shape). +-- so the serde gates Binary as Unsupported; `length` no longer needs the gate since it runs on +-- the datafusion-spark kernel, which handles binary input. statement CREATE TABLE test_bit_length_binary(b binary) USING parquet diff --git a/spark/src/test/resources/sql-tests/expressions/string/length.sql b/spark/src/test/resources/sql-tests/expressions/string/length.sql index 1e1e8ccc7a5..dfba742badf 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/length.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/length.sql @@ -15,6 +15,8 @@ -- specific language governing permissions and limitations -- under the License. +-- Config: spark.comet.shuffle.mode=native + statement CREATE TABLE test_length(s string) USING parquet @@ -27,3 +29,72 @@ SELECT length(s), char_length(s) FROM test_length -- literal arguments query SELECT length('hello'), length(''), length(NULL) + +-- BinaryType input runs natively. length counts bytes for binary and characters for strings, +-- so the same text gives a different answer through each column. +statement +CREATE TABLE test_length_binary(s string, b binary, h string, r struct) USING parquet + +statement +INSERT INTO test_length_binary VALUES + ('hello', X'68656C6C6F', '68656C6C6F', named_struct('b', X'68656C6C6F')), + (CAST(X'C3A9' AS STRING), X'C3A9', 'C3A9', named_struct('b', X'C3A9')), + (CAST(X'F09F9880' AS STRING), X'F09F9880', 'F09F9880', named_struct('b', X'F09F9880')), + ('', X'', '', named_struct('b', X'')), + (NULL, NULL, NULL, named_struct('b', CAST(NULL AS BINARY))) + +query +SELECT s, length(s), length(b) FROM test_length_binary + +-- binary nested in a struct field +query +SELECT length(r.b) FROM test_length_binary + +-- binary produced by unhex rather than read from the table +query +SELECT length(unhex(h)) FROM test_length_binary + +-- substring on binary yields binary, so length counts the bytes of the slice +query +SELECT length(substring(b, 1, 2)), length(substring(b, 2)) FROM test_length_binary + +-- Spark parses char_length and character_length to the same Length expression, so they accept binary +query +SELECT char_length(b), character_length(b) FROM test_length_binary + +-- literal arguments +query +SELECT length(X'00FF'), length(X''), length(CAST(NULL AS BINARY)), length(unhex('C3A9')) + +-- bytes that are not valid UTF-8 and embedded NUL bytes count as bytes, never as text +query +SELECT length(X'FF'), length(X'0000'), length(X'C3'), length(X'FFFE0000') + +-- a string cast to binary counts its UTF-8 bytes, so the two lengths differ on multi-byte text +query +SELECT length(CAST(s AS BINARY)), length(s) FROM test_length_binary + +-- binary inside an array element and a map value +query +SELECT length(array(b, X'01')[0]), length(map(1, b)[1]) FROM test_length_binary + +-- the Int32 result takes part in arithmetic, a filter and a native aggregate +query +SELECT length(b) + 1, length(b) * 2 FROM test_length_binary WHERE length(b) >= 0 + +query +SELECT length(b) AS n, count(*) FROM test_length_binary GROUP BY length(b) + +-- the binary column crosses a native shuffle before length reads it +query +SELECT length(b), length(r.b) FROM test_length_binary DISTRIBUTE BY b + +-- a column that is NULL on every row +statement +CREATE TABLE test_length_all_null(b binary) USING parquet + +statement +INSERT INTO test_length_all_null VALUES (CAST(NULL AS BINARY)), (CAST(NULL AS BINARY)) + +query +SELECT length(b), char_length(b) FROM test_length_all_null diff --git a/spark/src/test/resources/sql-tests/expressions/string/octet_length.sql b/spark/src/test/resources/sql-tests/expressions/string/octet_length.sql index e650950347c..5fd82780c3d 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/octet_length.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/octet_length.sql @@ -29,7 +29,8 @@ query SELECT octet_length('hello'), octet_length(''), octet_length(NULL) -- BinaryType input falls back to Spark; the native DataFusion impl rejects Binary at runtime, --- so the serde gates Binary as Unsupported (matching the existing CometLength shape). +-- so the serde gates Binary as Unsupported; `length` no longer needs the gate since it runs on +-- the datafusion-spark kernel, which handles binary input. statement CREATE TABLE test_octet_length_binary(b binary) USING parquet diff --git a/spark/src/test/resources/test-data/dictionary-string-binary.parquet b/spark/src/test/resources/test-data/dictionary-string-binary.parquet new file mode 100644 index 0000000000000000000000000000000000000000..1bc3c58dfc76e81b059bd22d45ccd23f51f635be GIT binary patch literal 751 zcmZva&r2Io5XWbqaS0)ZD9*YM_mG1wg~p$dP}Ce|ts247Y)X+{icvx$sEHeEN>85r z8$1+x>7~%42akpR1N|fV2lyZ8%&d`0v#|4K=FNQO`(WzSSjG%iaCH~62$KLaopAu@ zw2qIvhMnI(*g%H=yKmQ*2snj)33FIL7u}#@USuTi!^8xa;b<0 z53}AM12*vU)M&ijTTNf#-4 zuug8?QjSzm6g5PTDy5pHRG)e>jcWWf?N=o0)>D(LhfmfsNcCvbJGET*_?x1+lO8Q6L)8a~bVQFF6mUuB_H`aL7>b&X?1p$}l literal 0 HcmV?d00001 diff --git a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala index c58403487ac..e2e5989f24d 100644 --- a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala @@ -442,6 +442,38 @@ class CometStringExpressionSuite extends CometTestBase with CometCodegenAssertio } } + test("length on binary input runs natively") { + // repeated values so the parquet writer can dictionary-encode the column + val data = (0 until 1000).map { i => + val b: Array[Byte] = i % 5 match { + case 0 => "hello".getBytes("UTF-8") + case 1 => Array(0xc3.toByte, 0xa9.toByte) + case 2 => Array.empty[Byte] + case 3 => null + case 4 => Array(0x00.toByte, 0xff.toByte) + } + Tuple1(b) + } + Seq(true, false).foreach { dictionary => + withParquetTable(data, "tbl", withDictionary = dictionary) { + checkSparkAnswerAndOperator( + "SELECT length(_1), char_length(_1), character_length(_1) FROM tbl") + checkSparkAnswerAndOperator("SELECT length(_1) FROM tbl WHERE length(_1) > 2") + } + } + } + + test("length on dictionary-typed string and binary columns from an Arrow-written file") { + // The file's Arrow schema declares both columns as dictionary, + // so the native reader hands length a dictionary array unless the scan unwraps it. + withTempView("dict") { + readResourceParquetFile("test-data/dictionary-string-binary.parquet").createTempView("dict") + checkSparkAnswerAndOperator( + "SELECT length(s), length(b), char_length(s), character_length(b), s, b FROM dict") + checkSparkAnswerAndOperator("SELECT length(b) + length(s) FROM dict WHERE length(b) > 0") + } + } + // Simplified version of "filter pushdown - StringPredicate" that does not generate dictionaries test("string predicate filter") { Seq(false, true).foreach { pushdown =>