From 768b3e90f261c7aea58bdb98dc698b90deeeae34 Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sun, 14 Dec 2025 16:24:01 +0400 Subject: [PATCH 1/8] impl map_from_entries --- native/core/src/execution/jni_api.rs | 2 + .../apache/comet/serde/QueryPlanSerde.scala | 3 +- .../scala/org/apache/comet/serde/maps.scala | 29 +++++++++++- .../comet/CometMapExpressionSuite.scala | 45 +++++++++++++++++++ 4 files changed, 77 insertions(+), 2 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index a24d9930597..4f53cea3e68 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -46,6 +46,7 @@ use datafusion_spark::function::datetime::date_add::SparkDateAdd; use datafusion_spark::function::datetime::date_sub::SparkDateSub; use datafusion_spark::function::hash::sha1::SparkSha1; use datafusion_spark::function::hash::sha2::SparkSha2; +use datafusion_spark::function::map::map_from_entries::MapFromEntries; use datafusion_spark::function::math::expm1::SparkExpm1; use datafusion_spark::function::string::char::CharFunc; use datafusion_spark::function::string::concat::SparkConcat; @@ -337,6 +338,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) { session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha1::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkConcat::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitwiseNot::default())); + session_ctx.register_udf(ScalarUDF::new_from_impl(MapFromEntries::default())); } /// Prepares arrow arrays for output. diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 54df2f1688d..a99cf3824bf 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -125,7 +125,8 @@ object QueryPlanSerde extends Logging with CometExprShim { classOf[MapKeys] -> CometMapKeys, classOf[MapEntries] -> CometMapEntries, classOf[MapValues] -> CometMapValues, - classOf[MapFromArrays] -> CometMapFromArrays) + classOf[MapFromArrays] -> CometMapFromArrays, + classOf[MapFromEntries] -> CometMapFromEntries) private val structExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map( classOf[CreateNamedStruct] -> CometCreateNamedStruct, diff --git a/spark/src/main/scala/org/apache/comet/serde/maps.scala b/spark/src/main/scala/org/apache/comet/serde/maps.scala index 2e217f6af0b..498aa3594cf 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -19,9 +19,12 @@ package org.apache.comet.serde +import scala.annotation.tailrec + import org.apache.spark.sql.catalyst.expressions._ -import org.apache.spark.sql.types.{ArrayType, MapType} +import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, MapType, StructType} +import org.apache.comet.serde.CometArrayReverse.containsBinary import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType} object CometMapKeys extends CometExpressionSerde[MapKeys] { @@ -89,3 +92,27 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { optExprWithInfo(mapFromArraysExpr, expr, expr.children: _*) } } + +object CometMapFromEntries extends CometScalarFunction[MapFromEntries]("map_from_entries") { + val keyUnsupportedReason = "Using BinaryType as Map keys is not allowed in map_from_entries" + val valueUnsupportedReason = "Using BinaryType as Map values is not allowed in map_from_entries" + + private def containsBinary(dataType: DataType): Boolean = { + dataType match { + case BinaryType => true + case StructType(fields) => fields.exists(field => containsBinary(field.dataType)) + case ArrayType(elementType, _) => containsBinary(elementType) + case _ => false + } + } + + override def getSupportLevel(expr: MapFromEntries): SupportLevel = { + if (containsBinary(expr.dataType.keyType)) { + return Incompatible(Some(keyUnsupportedReason)) + } + if (containsBinary(expr.dataType.valueType)) { + return Incompatible(Some(valueUnsupportedReason)) + } + Compatible(None) + } +} diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 88c13391a67..01b9744ed6f 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -25,7 +25,9 @@ import org.apache.hadoop.fs.Path import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.BinaryType +import org.apache.comet.serde.CometMapFromEntries import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometMapExpressionSuite extends CometTestBase { @@ -125,4 +127,47 @@ class CometMapExpressionSuite extends CometTestBase { } } + test("map_from_entries") { + withTempDir { dir => + val path = new Path(dir.toURI.toString, "test.parquet") + val filename = path.toString + val random = new Random(42) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val schemaGenOptions = + SchemaGenOptions( + generateArray = true, + generateStruct = true, + primitiveTypes = SchemaGenOptions.defaultPrimitiveTypes.filterNot(_ == BinaryType)) + val dataGenOptions = DataGenOptions(allowNull = false, generateNegativeZero = false) + ParquetGenerator.makeParquetFile( + random, + spark, + filename, + 100, + schemaGenOptions, + dataGenOptions) + } + val df = spark.read.parquet(filename) + df.createOrReplaceTempView("t1") + for (field <- df.schema.fieldNames) { + checkSparkAnswerAndOperator( + spark.sql(s"SELECT map_from_entries(array(struct($field as a, $field as b))) FROM t1")) + } + } + } + + test("map_from_entries - fallback for binary type") { + val table = "t2" + withTable(table) { + sql( + s"create table $table using parquet as select cast(array() as array) as c1 from range(10)") + checkSparkAnswerAndFallbackReason( + sql(s"select map_from_entries(array(struct(c1, 0))) from $table"), + CometMapFromEntries.keyUnsupportedReason) + checkSparkAnswerAndFallbackReason( + sql(s"select map_from_entries(array(struct(0, c1))) from $table"), + CometMapFromEntries.valueUnsupportedReason) + } + } + } From c68c3428676b5d991e7ba9e13464bf2ce1ec84e8 Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Tue, 16 Dec 2025 16:10:43 +0400 Subject: [PATCH 2/8] Revert "impl map_from_entries" This reverts commit 768b3e90f261c7aea58bdb98dc698b90deeeae34. --- native/core/src/execution/jni_api.rs | 2 - .../apache/comet/serde/QueryPlanSerde.scala | 3 +- .../scala/org/apache/comet/serde/maps.scala | 29 +----------- .../comet/CometMapExpressionSuite.scala | 45 ------------------- 4 files changed, 2 insertions(+), 77 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 4f53cea3e68..a24d9930597 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -46,7 +46,6 @@ use datafusion_spark::function::datetime::date_add::SparkDateAdd; use datafusion_spark::function::datetime::date_sub::SparkDateSub; use datafusion_spark::function::hash::sha1::SparkSha1; use datafusion_spark::function::hash::sha2::SparkSha2; -use datafusion_spark::function::map::map_from_entries::MapFromEntries; use datafusion_spark::function::math::expm1::SparkExpm1; use datafusion_spark::function::string::char::CharFunc; use datafusion_spark::function::string::concat::SparkConcat; @@ -338,7 +337,6 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) { session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha1::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkConcat::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitwiseNot::default())); - session_ctx.register_udf(ScalarUDF::new_from_impl(MapFromEntries::default())); } /// Prepares arrow arrays for output. diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index a99cf3824bf..54df2f1688d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -125,8 +125,7 @@ object QueryPlanSerde extends Logging with CometExprShim { classOf[MapKeys] -> CometMapKeys, classOf[MapEntries] -> CometMapEntries, classOf[MapValues] -> CometMapValues, - classOf[MapFromArrays] -> CometMapFromArrays, - classOf[MapFromEntries] -> CometMapFromEntries) + classOf[MapFromArrays] -> CometMapFromArrays) private val structExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map( classOf[CreateNamedStruct] -> CometCreateNamedStruct, diff --git a/spark/src/main/scala/org/apache/comet/serde/maps.scala b/spark/src/main/scala/org/apache/comet/serde/maps.scala index 498aa3594cf..2e217f6af0b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -19,12 +19,9 @@ package org.apache.comet.serde -import scala.annotation.tailrec - import org.apache.spark.sql.catalyst.expressions._ -import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, MapType, StructType} +import org.apache.spark.sql.types.{ArrayType, MapType} -import org.apache.comet.serde.CometArrayReverse.containsBinary import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType} object CometMapKeys extends CometExpressionSerde[MapKeys] { @@ -92,27 +89,3 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { optExprWithInfo(mapFromArraysExpr, expr, expr.children: _*) } } - -object CometMapFromEntries extends CometScalarFunction[MapFromEntries]("map_from_entries") { - val keyUnsupportedReason = "Using BinaryType as Map keys is not allowed in map_from_entries" - val valueUnsupportedReason = "Using BinaryType as Map values is not allowed in map_from_entries" - - private def containsBinary(dataType: DataType): Boolean = { - dataType match { - case BinaryType => true - case StructType(fields) => fields.exists(field => containsBinary(field.dataType)) - case ArrayType(elementType, _) => containsBinary(elementType) - case _ => false - } - } - - override def getSupportLevel(expr: MapFromEntries): SupportLevel = { - if (containsBinary(expr.dataType.keyType)) { - return Incompatible(Some(keyUnsupportedReason)) - } - if (containsBinary(expr.dataType.valueType)) { - return Incompatible(Some(valueUnsupportedReason)) - } - Compatible(None) - } -} diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 01b9744ed6f..88c13391a67 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -25,9 +25,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.BinaryType -import org.apache.comet.serde.CometMapFromEntries import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometMapExpressionSuite extends CometTestBase { @@ -127,47 +125,4 @@ class CometMapExpressionSuite extends CometTestBase { } } - test("map_from_entries") { - withTempDir { dir => - val path = new Path(dir.toURI.toString, "test.parquet") - val filename = path.toString - val random = new Random(42) - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val schemaGenOptions = - SchemaGenOptions( - generateArray = true, - generateStruct = true, - primitiveTypes = SchemaGenOptions.defaultPrimitiveTypes.filterNot(_ == BinaryType)) - val dataGenOptions = DataGenOptions(allowNull = false, generateNegativeZero = false) - ParquetGenerator.makeParquetFile( - random, - spark, - filename, - 100, - schemaGenOptions, - dataGenOptions) - } - val df = spark.read.parquet(filename) - df.createOrReplaceTempView("t1") - for (field <- df.schema.fieldNames) { - checkSparkAnswerAndOperator( - spark.sql(s"SELECT map_from_entries(array(struct($field as a, $field as b))) FROM t1")) - } - } - } - - test("map_from_entries - fallback for binary type") { - val table = "t2" - withTable(table) { - sql( - s"create table $table using parquet as select cast(array() as array) as c1 from range(10)") - checkSparkAnswerAndFallbackReason( - sql(s"select map_from_entries(array(struct(c1, 0))) from $table"), - CometMapFromEntries.keyUnsupportedReason) - checkSparkAnswerAndFallbackReason( - sql(s"select map_from_entries(array(struct(0, c1))) from $table"), - CometMapFromEntries.valueUnsupportedReason) - } - } - } From 45be7e60a9dbee759118b6be18871d9109575739 Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Thu, 20 Aug 2026 10:55:15 +0400 Subject: [PATCH 3/8] work --- native/spark-expr/Cargo.toml | 6 +- native/spark-expr/benches/split.rs | 96 ++++++ native/spark-expr/src/string_funcs/split.rs | 314 +++++++++++--------- 3 files changed, 273 insertions(+), 143 deletions(-) create mode 100644 native/spark-expr/benches/split.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6faa9fec4ec..b81d3bb91ed 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -222,4 +222,8 @@ harness = false [[bench]] name = "cast_int_to_decimal" -harness = false \ No newline at end of file +harness = false + +[[bench]] +name = "split" +harness = false diff --git a/native/spark-expr/benches/split.rs b/native/spark-expr/benches/split.rs new file mode 100644 index 00000000000..948c1526009 --- /dev/null +++ b/native/spark-expr/benches/split.rs @@ -0,0 +1,96 @@ +use std::sync::Arc; +use arrow::array::StringArray; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use datafusion::logical_expr::ColumnarValue; +use datafusion::common::ScalarValue; +use datafusion_comet_spark_expr::spark_split; +use std::hint::black_box; + +fn generate_string_array(num_rows: usize, pattern_type: &str) -> Arc { + let mut builder = arrow::array::StringBuilder::with_capacity(num_rows, num_rows * 64); + + for i in 0..num_rows { + if i % 20 == 0 { + builder.append_null(); + continue; + } + + match pattern_type { + "csv" => builder.append_value(format!("field1_{i},field2_{i},field3_{i},field4_{i},field5_{i}")), + "whitespace_regex" => builder.append_value(format!("word1_{i} word2_{i} \t word3_{i} word4_{i}")), + "trailing_delimiters" => builder.append_value(format!("data_{i},,,,")), + _ => unreachable!(), + } + } + + Arc::new(builder.finish()) +} + +fn bench_spark_split(c: &mut Criterion) { + let mut group = c.benchmark_group("spark_split"); + let batch_sizes = [1024, 8192]; + + for &size in &batch_sizes { + group.throughput(Throughput::Elements(size as u64)); + + { + let array = generate_string_array(size, "csv"); + let args = vec![ + ColumnarValue::Array(array), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string()))), + ]; + + group.bench_with_input( + BenchmarkId::new("literal_char_default_limit", size), + &args, + |b, args| { + b.iter(|| { + black_box(spark_split(black_box(args)).unwrap()); + }); + }, + ); + } + + { + let array = generate_string_array(size, "whitespace_regex"); + let args = vec![ + ColumnarValue::Array(array), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(r"\s+".to_string()))), + ]; + + group.bench_with_input( + BenchmarkId::new("regex_pattern", size), + &args, + |b, args| { + b.iter(|| { + black_box(spark_split(black_box(args)).unwrap()); + }); + }, + ); + } + + { + let array = generate_string_array(size, "trailing_delimiters"); + let args = vec![ + ColumnarValue::Array(array), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string()))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(0))), + ]; + + group.bench_with_input( + BenchmarkId::new("literal_char_limit_0", size), + &args, + |b, args| { + b.iter(|| { + black_box(spark_split(black_box(args)).unwrap()); + }); + }, + ); + } + } + + group.finish(); +} + +criterion_group!(benches, bench_spark_split); +criterion_main!(benches); diff --git a/native/spark-expr/src/string_funcs/split.rs b/native/spark-expr/src/string_funcs/split.rs index 967a898440d..015e74a6411 100644 --- a/native/spark-expr/src/string_funcs/split.rs +++ b/native/spark-expr/src/string_funcs/split.rs @@ -15,10 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ - Array, ArrayBuilder, ArrayRef, GenericListArray, GenericStringArray, GenericStringBuilder, - ListArray, NullBufferBuilder, OffsetSizeTrait, -}; +use arrow::array::{Array, ArrayBuilder, ArrayRef, BufferBuilder, GenericListArray, GenericStringArray, GenericStringBuilder, ListArray, NullBufferBuilder, OffsetSizeTrait}; use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field}; use datafusion::common::{ @@ -97,9 +94,30 @@ pub fn spark_split(args: &[ColumnarValue]) -> DataFusionResult { } }; - let result = split_string(string.as_ref().unwrap(), pattern_str, limit)?; - let string_array = GenericStringArray::::from(result); - let list_array = create_list_array(Arc::new(string_array)); + let s = string.as_ref().unwrap(); + let regex = Regex::new(pattern_str).map_err(|e| { + DataFusionError::Execution(format!("Invalid regex pattern '{}': {}", pattern_str, e)) + })?; + + let mut str_offsets = BufferBuilder::::new(8); + let mut str_values = BufferBuilder::::new(s.len()); + str_offsets.append(0); + + let mut scratch = Vec::new(); + push_split_parts(s, ®ex, limit, &mut str_offsets, &mut str_values, &mut scratch); + + let item_offsets_buffer = OffsetBuffer::new(str_offsets.finish().into()); + let item_values_buffer = str_values.finish(); + + let string_array_values = unsafe { + GenericStringArray::::new_unchecked( + item_offsets_buffer, + item_values_buffer, + None, + ) + }; + + let list_array = create_list_array(Arc::new(string_array_values)); Ok(ColumnarValue::Scalar(ScalarValue::List(Arc::new( list_array, @@ -160,9 +178,33 @@ pub fn spark_split_sql(args: &[ColumnarValue]) -> DataFusionResult return exec_err!("split_sql delimiter must be a string"), }; - let result = split_sql_string(string.as_ref().unwrap(), delimiter); - let string_array = GenericStringArray::::from(result); - let list_array = create_list_array(Arc::new(string_array)); + let string = string.clone().unwrap(); + let mut str_offsets = BufferBuilder::::new(8); + let mut str_values = BufferBuilder::::new(string.len()); + str_offsets.append(0); + + if delimiter.is_empty() { + for ch in string.chars() { + append_str(ch.encode_utf8(&mut [0; 4]), &mut str_offsets, &mut str_values); + } + } else { + for p in string.split(delimiter) { + append_str(p, &mut str_offsets, &mut str_values); + } + } + + let item_offsets_buffer = OffsetBuffer::new(str_offsets.finish().into()); + let item_values_buffer = str_values.finish(); + + let string_array_values = unsafe { + GenericStringArray::::new_unchecked( + item_offsets_buffer, + item_values_buffer, + None, + ) + }; + + let list_array = create_list_array(Arc::new(string_array_values)); Ok(ColumnarValue::Scalar(ScalarValue::List(Arc::new( list_array, @@ -172,11 +214,62 @@ pub fn spark_split_sql(args: &[ColumnarValue]) -> DataFusionResult bool { + !pattern.chars().any(|c| matches!(c, + '.' | '^' | '$' | '*' | '+' | '?' | '(' | ')' | '[' | ']' | '{' | '}' | '|' | '\\' + )) +} + +#[inline] +fn push_split_literal<'a, O: OffsetSizeTrait>( + string: &'a str, + delimiter: &str, + limit: i32, + offsets: &mut BufferBuilder, + values: &mut BufferBuilder, + scratch: &mut Vec<&'a str>, +) { + if limit == 0 { + scratch.clear(); + scratch.extend(string.split(delimiter)); + while scratch.last().is_some_and(|s| s.is_empty()) { + scratch.pop(); + } + if scratch.is_empty() { + append_str("", offsets, values); + } else { + for &p in scratch.iter() { + append_str(p, offsets, values); + } + } + } else if limit > 0 { + let cap = (limit - 1) as usize; + let mut count = 0; + let mut last_end = 0; + for (start, part) in string.match_indices(delimiter) { + if count >= cap { + break; + } + append_str(&string[last_end..start], offsets, values); + last_end = start + delimiter.len(); + count += 1; + } + append_str(&string[last_end..], offsets, values); + } else { + for p in string.split(delimiter) { + append_str(p, offsets, values); + } + } +} + fn split_array( string_array: &dyn arrow::array::Array, pattern: &str, limit: i32, ) -> DataFusionResult { + if is_regex_literal(pattern) { + spli + } // Compile regex once for the entire array let regex = Regex::new(pattern).map_err(|e| { DataFusionError::Execution(format!("Invalid regex pattern '{}': {}", pattern, e)) @@ -285,22 +378,38 @@ fn split_generic( limit: i32, ) -> DataFusionResult { let len = string_array.len(); - let mut offsets: Vec = Vec::with_capacity(len + 1); - let mut values_builder = GenericStringBuilder::::new(); - offsets.push(O::usize_as(0)); + let mut list_offsets: Vec = Vec::with_capacity(len + 1); + + let estimated_items = (len * 4).max(16); + let bytes_capacity = string_array.value_data().len(); + + let mut str_offsets = BufferBuilder::::new(estimated_items + 1); + let mut str_values = BufferBuilder::::new(bytes_capacity); + str_offsets.append(O::usize_as(0)); + + let mut scratch = Vec::new(); + list_offsets.push(O::usize_as(0)); - // Bulk-NULL: output null mask equals input's, so reuse it instead of - // tracking per-row in a NullBufferBuilder. Null rows contribute no parts - // (offset does not advance) and the cloned NullBuffer marks them. for i in 0..len { if !string_array.is_null(i) { let s = string_array.value(i); - push_split_parts(s, regex, limit, &mut values_builder); + push_split_parts(s, regex, limit, &mut str_offsets, &mut str_values, &mut scratch); } - offsets.push(O::usize_as(values_builder.len())); + list_offsets.push(O::usize_as(str_offsets.len() - 1)); } - let values_array = Arc::new(values_builder.finish()) as ArrayRef; + let item_offsets_buffer = OffsetBuffer::new(str_offsets.finish().into()); + let item_values_buffer = str_values.finish(); + + let string_array_values = unsafe { + GenericStringArray::::new_unchecked( + item_offsets_buffer, + item_values_buffer, + None, + ) + }; + let values_array = Arc::new(string_array_values) as ArrayRef; + let item_type = if O::IS_LARGE { DataType::LargeUtf8 } else { @@ -309,7 +418,7 @@ fn split_generic( let field = Arc::new(Field::new("item", item_type, false)); let list_array = GenericListArray::::new( field, - OffsetBuffer::new(offsets.into()), + OffsetBuffer::new(list_offsets.into()), values_array, string_array.nulls().cloned(), ); @@ -323,7 +432,11 @@ fn split_sql_generic_scalar( ) -> DataFusionResult { let len = string_array.len(); let mut offsets: Vec = Vec::with_capacity(len + 1); - let mut values_builder = GenericStringBuilder::::new(); + + let estimated_items = (len * 4).max(16); + let bytes_capacity = string_array.value_data().len(); + let mut values_builder = GenericStringBuilder::::with_capacity(estimated_items, bytes_capacity); + offsets.push(O::usize_as(0)); for i in 0..len { @@ -356,7 +469,11 @@ fn split_sql_generic_scalar_array( ) -> DataFusionResult { let len = delimiter_array.len(); let mut offsets: Vec = Vec::with_capacity(len + 1); - let mut values_builder = GenericStringBuilder::::new(); + + let estimated_items = (len * 4).max(16); + let bytes_capacity = string.len() * len; + let mut values_builder = GenericStringBuilder::::with_capacity(estimated_items, bytes_capacity); + let mut nulls = NullBufferBuilder::new(len); offsets.push(O::usize_as(0)); @@ -393,7 +510,12 @@ fn split_sql_generic_array( ) -> DataFusionResult { let len = string_array.len(); let mut offsets: Vec = Vec::with_capacity(len + 1); - let mut values_builder = GenericStringBuilder::::new(); + + let estimated_items = (len * 4).max(16); + let bytes_capacity = string_array.value_data().len(); + let mut values_builder = GenericStringBuilder::::with_capacity(estimated_items, bytes_capacity); + + let mut nulls = NullBufferBuilder::new(len); offsets.push(O::usize_as(0)); @@ -428,45 +550,52 @@ fn split_sql_generic_array( Ok(ColumnarValue::Array(Arc::new(list_array))) } -/// Push the splits of `string` into `builder`. Avoids materializing an -/// intermediate `Vec` — appends each `&str` slice from the regex -/// iterator directly (the builder copies into its own buffer). -fn push_split_parts( - string: &str, +#[inline] +fn append_str( + s: &str, + offsets: &mut BufferBuilder, + values: &mut BufferBuilder, +) { + values.append_slice(s.as_bytes()); + offsets.append(O::usize_as(values.len())); +} + +#[inline] +fn push_split_parts<'a, O: OffsetSizeTrait>( + string: &'a str, regex: &Regex, limit: i32, - builder: &mut GenericStringBuilder, + offsets: &mut BufferBuilder, + values: &mut BufferBuilder, + scratch: &mut Vec<&'a str>, ) { if limit == 0 { - // limit = 0: split all, drop trailing empties. Need to know the end - // before pushing, so collect borrowed slices first (no string copies). - let mut parts: Vec<&str> = regex.split(string).collect(); - while parts.last().is_some_and(|s| s.is_empty()) { - parts.pop(); + scratch.clear(); + scratch.extend(regex.split(string)); + while scratch.last().is_some_and(|s| s.is_empty()) { + scratch.pop(); } - if parts.is_empty() { - builder.append_value(""); + if scratch.is_empty() { + append_str("", offsets, values); } else { - for p in parts { - builder.append_value(p); + for &p in scratch.iter() { + append_str(p, offsets, values); } } } else if limit > 0 { - // limit > 0: at most limit-1 splits. let mut last_end = 0; let cap = (limit - 1) as usize; for (count, mat) in regex.find_iter(string).enumerate() { if count >= cap { break; } - builder.append_value(&string[last_end..mat.start()]); + append_str(&string[last_end..mat.start()], offsets, values); last_end = mat.end(); } - builder.append_value(&string[last_end..]); + append_str(&string[last_end..], offsets, values); } else { - // limit < 0: split all, keep trailing empties. for p in regex.split(string) { - builder.append_value(p); + append_str(p, offsets, values); } } } @@ -485,56 +614,6 @@ fn push_split_sql_parts( } } -fn split_string(string: &str, pattern: &str, limit: i32) -> DataFusionResult> { - let regex = Regex::new(pattern).map_err(|e| { - DataFusionError::Execution(format!("Invalid regex pattern '{}': {}", pattern, e)) - })?; - - Ok(split_with_regex(string, ®ex, limit)) -} - -fn split_with_regex(string: &str, regex: &Regex, limit: i32) -> Vec { - if limit == 0 { - // limit = 0: split as many times as possible, discard trailing empty strings - let mut parts: Vec = regex.split(string).map(|s| s.to_string()).collect(); - // Remove trailing empty strings - while parts.last().is_some_and(|s| s.is_empty()) { - parts.pop(); - } - if parts.is_empty() { - vec!["".to_string()] - } else { - parts - } - } else if limit > 0 { - // limit > 0: at most limit-1 splits (array length <= limit) - let mut parts: Vec = Vec::new(); - let mut last_end = 0; - - for (count, mat) in regex.find_iter(string).enumerate() { - if count >= (limit - 1) as usize { - break; - } - parts.push(string[last_end..mat.start()].to_string()); - last_end = mat.end(); - } - // Add the remaining string - parts.push(string[last_end..].to_string()); - parts - } else { - // limit < 0: split as many times as possible, keep trailing empty strings - regex.split(string).map(|s| s.to_string()).collect() - } -} - -fn split_sql_string(string: &str, delimiter: &str) -> Vec { - if delimiter.is_empty() { - vec![string.to_string()] - } else { - string.split(delimiter).map(|s| s.to_string()).collect() - } -} - fn create_list_array(values: ArrayRef) -> ListArray { let field = Arc::new(Field::new("item", DataType::Utf8, false)); let offsets = vec![0i32, values.len() as i32]; @@ -615,30 +694,6 @@ mod tests { assert!(matches!(result, ColumnarValue::Array(_))); } - #[test] - fn test_split_regex() { - let parts = split_string("foo123bar456baz", r"\d+", -1).unwrap(); - assert_eq!(parts, vec!["foo", "bar", "baz"]); - } - - #[test] - fn test_split_limit_positive() { - let parts = split_string("a,b,c,d,e", ",", 3).unwrap(); - assert_eq!(parts, vec!["a", "b", "c,d,e"]); - } - - #[test] - fn test_split_limit_zero() { - let parts = split_string("a,b,c,,", ",", 0).unwrap(); - assert_eq!(parts, vec!["a", "b", "c"]); - } - - #[test] - fn test_split_limit_negative() { - let parts = split_string("a,b,c,,", ",", -1).unwrap(); - assert_eq!(parts, vec!["a", "b", "c", "", ""]); - } - #[test] fn test_split_with_nulls() { // Test that NULL inputs produce NULL outputs (not empty arrays) @@ -669,31 +724,6 @@ mod tests { } } - #[test] - fn test_split_empty_string() { - // Test that empty string input produces array with single empty string - let parts = split_string("", ",", -1).unwrap(); - assert_eq!(parts, vec![""]); - } - - #[test] - fn test_split_sql_literal_delimiter() { - let parts = split_sql_string("a.b.", "."); - assert_eq!(parts, vec!["a", "b", ""]); - } - - #[test] - fn test_split_sql_empty_delimiter() { - let parts = split_sql_string("abc", ""); - assert_eq!(parts, vec!["abc"]); - } - - #[test] - fn test_split_sql_keeps_regex_chars_literal() { - let parts = split_sql_string("a.b.c", "."); - assert_eq!(parts, vec!["a", "b", "c"]); - } - #[test] fn test_split_sql_scalar_nulls_return_typed_null_list() { let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string()))); From 89161832a1f36245375fd6c079f712b2cf61ce25 Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sat, 22 Aug 2026 21:21:57 +0400 Subject: [PATCH 4/8] fmt --- native/spark-expr/benches/split.rs | 46 ++++-- native/spark-expr/src/string_funcs/split.rs | 153 ++++++++++++++++---- 2 files changed, 157 insertions(+), 42 deletions(-) diff --git a/native/spark-expr/benches/split.rs b/native/spark-expr/benches/split.rs index 948c1526009..edc1f64e785 100644 --- a/native/spark-expr/benches/split.rs +++ b/native/spark-expr/benches/split.rs @@ -1,11 +1,29 @@ -use std::sync::Arc; +// 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::StringArray; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; -use datafusion::logical_expr::ColumnarValue; use datafusion::common::ScalarValue; +use datafusion::logical_expr::ColumnarValue; use datafusion_comet_spark_expr::spark_split; use std::hint::black_box; +use std::sync::Arc; + fn generate_string_array(num_rows: usize, pattern_type: &str) -> Arc { let mut builder = arrow::array::StringBuilder::with_capacity(num_rows, num_rows * 64); @@ -16,8 +34,12 @@ fn generate_string_array(num_rows: usize, pattern_type: &str) -> Arc builder.append_value(format!("field1_{i},field2_{i},field3_{i},field4_{i},field5_{i}")), - "whitespace_regex" => builder.append_value(format!("word1_{i} word2_{i} \t word3_{i} word4_{i}")), + "csv" => builder.append_value(format!( + "field1_{i},field2_{i},field3_{i},field4_{i},field5_{i}" + )), + "whitespace_regex" => { + builder.append_value(format!("word1_{i} word2_{i} \t word3_{i} word4_{i}")) + } "trailing_delimiters" => builder.append_value(format!("data_{i},,,,")), _ => unreachable!(), } @@ -55,18 +77,14 @@ fn bench_spark_split(c: &mut Criterion) { let array = generate_string_array(size, "whitespace_regex"); let args = vec![ ColumnarValue::Array(array), - ColumnarValue::Scalar(ScalarValue::Utf8(Some(r"\s+".to_string()))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(r"\s".to_string()))), ]; - group.bench_with_input( - BenchmarkId::new("regex_pattern", size), - &args, - |b, args| { - b.iter(|| { - black_box(spark_split(black_box(args)).unwrap()); - }); - }, - ); + group.bench_with_input(BenchmarkId::new("regex_pattern", size), &args, |b, args| { + b.iter(|| { + black_box(spark_split(black_box(args)).unwrap()); + }); + }); } { diff --git a/native/spark-expr/src/string_funcs/split.rs b/native/spark-expr/src/string_funcs/split.rs index 015e74a6411..fa0f2be22cf 100644 --- a/native/spark-expr/src/string_funcs/split.rs +++ b/native/spark-expr/src/string_funcs/split.rs @@ -15,7 +15,10 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, ArrayBuilder, ArrayRef, BufferBuilder, GenericListArray, GenericStringArray, GenericStringBuilder, ListArray, NullBufferBuilder, OffsetSizeTrait}; +use arrow::array::{ + Array, ArrayBuilder, ArrayRef, BufferBuilder, GenericListArray, GenericStringArray, + GenericStringBuilder, ListArray, NullBufferBuilder, OffsetSizeTrait, +}; use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field}; use datafusion::common::{ @@ -96,7 +99,10 @@ pub fn spark_split(args: &[ColumnarValue]) -> DataFusionResult { let s = string.as_ref().unwrap(); let regex = Regex::new(pattern_str).map_err(|e| { - DataFusionError::Execution(format!("Invalid regex pattern '{}': {}", pattern_str, e)) + DataFusionError::Execution(format!( + "Invalid regex pattern '{}': {}", + pattern_str, e + )) })?; let mut str_offsets = BufferBuilder::::new(8); @@ -104,7 +110,14 @@ pub fn spark_split(args: &[ColumnarValue]) -> DataFusionResult { str_offsets.append(0); let mut scratch = Vec::new(); - push_split_parts(s, ®ex, limit, &mut str_offsets, &mut str_values, &mut scratch); + push_split_parts( + s, + ®ex, + limit, + &mut str_offsets, + &mut str_values, + &mut scratch, + ); let item_offsets_buffer = OffsetBuffer::new(str_offsets.finish().into()); let item_values_buffer = str_values.finish(); @@ -185,7 +198,11 @@ pub fn spark_split_sql(args: &[ColumnarValue]) -> DataFusionResult DataFusionResult bool { - !pattern.chars().any(|c| matches!(c, - '.' | '^' | '$' | '*' | '+' | '?' | '(' | ')' | '[' | ']' | '{' | '}' | '|' | '\\' - )) + !pattern.chars().any(|c| { + matches!( + c, + '.' | '^' | '$' | '*' | '+' | '?' | '(' | ')' | '[' | ']' | '{' | '}' | '|' | '\\' + ) + }) } #[inline] @@ -246,7 +266,7 @@ fn push_split_literal<'a, O: OffsetSizeTrait>( let cap = (limit - 1) as usize; let mut count = 0; let mut last_end = 0; - for (start, part) in string.match_indices(delimiter) { + for (start, _) in string.match_indices(delimiter) { if count >= cap { break; } @@ -262,25 +282,97 @@ fn push_split_literal<'a, O: OffsetSizeTrait>( } } -fn split_array( - string_array: &dyn arrow::array::Array, +fn split_generic_literal( + string_array: &GenericStringArray, pattern: &str, limit: i32, ) -> DataFusionResult { - if is_regex_literal(pattern) { - spli + let len = string_array.len(); + let mut list_offsets: Vec = Vec::with_capacity(len + 1); + + let estimated_items = (len * 4).max(16); + let bytes_capacity = string_array.value_data().len(); + + let mut str_offsets = BufferBuilder::::new(estimated_items + 1); + let mut str_values = BufferBuilder::::new(bytes_capacity); + str_offsets.append(O::usize_as(0)); + + let mut scratch = Vec::new(); + list_offsets.push(O::usize_as(0)); + + for i in 0..len { + if !string_array.is_null(i) { + let s = string_array.value(i); + push_split_literal( + s, + pattern, + limit, + &mut str_offsets, + &mut str_values, + &mut scratch, + ); + } + list_offsets.push(O::usize_as(str_offsets.len() - 1)); } - // Compile regex once for the entire array - let regex = Regex::new(pattern).map_err(|e| { - DataFusionError::Execution(format!("Invalid regex pattern '{}': {}", pattern, e)) - })?; + let item_offsets_buffer = OffsetBuffer::new(str_offsets.finish().into()); + let item_values_buffer = str_values.finish(); + + let string_array_values = unsafe { + GenericStringArray::::new_unchecked(item_offsets_buffer, item_values_buffer, None) + }; + let values_array = Arc::new(string_array_values) as ArrayRef; + + let item_type = if O::IS_LARGE { + DataType::LargeUtf8 + } else { + DataType::Utf8 + }; + let field = Arc::new(Field::new("item", item_type, false)); + let list_array = GenericListArray::::new( + field, + OffsetBuffer::new(list_offsets.into()), + values_array, + string_array.nulls().cloned(), + ); + + Ok(ColumnarValue::Array(Arc::new(list_array))) +} + +fn split_array( + string_array: &dyn Array, + pattern: &str, + limit: i32, +) -> DataFusionResult { + let is_literal = is_regex_literal(pattern); match string_array.data_type() { DataType::Utf8 => { - split_generic::(as_generic_string_array::(string_array)?, ®ex, limit) + let string_array = as_generic_string_array::(string_array)?; + if is_literal { + split_generic_literal::(string_array, pattern, limit) + } else { + let regex = Regex::new(pattern).map_err(|e| { + DataFusionError::Execution(format!( + "Invalid regex pattern '{}': {}", + pattern, e + )) + })?; + split_generic::(string_array, ®ex, limit) + } } DataType::LargeUtf8 => { - split_generic::(as_generic_string_array::(string_array)?, ®ex, limit) + let string_array = as_generic_string_array::(string_array)?; + if is_literal { + split_generic_literal::(string_array, pattern, limit) + } else { + let regex = Regex::new(pattern).map_err(|e| { + DataFusionError::Execution(format!( + "Invalid regex pattern '{}': {}", + pattern, e + )) + })?; + split_generic::(string_array, ®ex, limit) + } } _ => exec_err!( "split expects Utf8 or LargeUtf8 string array, got {:?}", @@ -393,7 +485,14 @@ fn split_generic( for i in 0..len { if !string_array.is_null(i) { let s = string_array.value(i); - push_split_parts(s, regex, limit, &mut str_offsets, &mut str_values, &mut scratch); + push_split_parts( + s, + regex, + limit, + &mut str_offsets, + &mut str_values, + &mut scratch, + ); } list_offsets.push(O::usize_as(str_offsets.len() - 1)); } @@ -402,11 +501,7 @@ fn split_generic( let item_values_buffer = str_values.finish(); let string_array_values = unsafe { - GenericStringArray::::new_unchecked( - item_offsets_buffer, - item_values_buffer, - None, - ) + GenericStringArray::::new_unchecked(item_offsets_buffer, item_values_buffer, None) }; let values_array = Arc::new(string_array_values) as ArrayRef; @@ -435,7 +530,8 @@ fn split_sql_generic_scalar( let estimated_items = (len * 4).max(16); let bytes_capacity = string_array.value_data().len(); - let mut values_builder = GenericStringBuilder::::with_capacity(estimated_items, bytes_capacity); + let mut values_builder = + GenericStringBuilder::::with_capacity(estimated_items, bytes_capacity); offsets.push(O::usize_as(0)); @@ -472,7 +568,8 @@ fn split_sql_generic_scalar_array( let estimated_items = (len * 4).max(16); let bytes_capacity = string.len() * len; - let mut values_builder = GenericStringBuilder::::with_capacity(estimated_items, bytes_capacity); + let mut values_builder = + GenericStringBuilder::::with_capacity(estimated_items, bytes_capacity); let mut nulls = NullBufferBuilder::new(len); offsets.push(O::usize_as(0)); @@ -513,8 +610,8 @@ fn split_sql_generic_array( let estimated_items = (len * 4).max(16); let bytes_capacity = string_array.value_data().len(); - let mut values_builder = GenericStringBuilder::::with_capacity(estimated_items, bytes_capacity); - + let mut values_builder = + GenericStringBuilder::::with_capacity(estimated_items, bytes_capacity); let mut nulls = NullBufferBuilder::new(len); offsets.push(O::usize_as(0)); From eb04df86605f78b32cff293d5dd1557b97353620 Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sat, 22 Aug 2026 21:59:03 +0400 Subject: [PATCH 5/8] complete --- native/spark-expr/src/lib.rs | 1 + native/spark-expr/src/string_funcs/split.rs | 143 ++++++++++++++++---- 2 files changed, 117 insertions(+), 27 deletions(-) diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 2b5c29befc6..60239e10776 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -18,6 +18,7 @@ // The clippy throws an error if the reference clone not wrapped into `Arc::clone` // The lint makes easier for code reader/reviewer separate references clones from more heavyweight ones #![deny(clippy::clone_on_ref_ptr)] +extern crate core; mod error; mod query_context; diff --git a/native/spark-expr/src/string_funcs/split.rs b/native/spark-expr/src/string_funcs/split.rs index fa0f2be22cf..a235ca333f0 100644 --- a/native/spark-expr/src/string_funcs/split.rs +++ b/native/spark-expr/src/string_funcs/split.rs @@ -14,7 +14,6 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. - use arrow::array::{ Array, ArrayBuilder, ArrayRef, BufferBuilder, GenericListArray, GenericStringArray, GenericStringBuilder, ListArray, NullBufferBuilder, OffsetSizeTrait, @@ -98,26 +97,49 @@ pub fn spark_split(args: &[ColumnarValue]) -> DataFusionResult { }; let s = string.as_ref().unwrap(); - let regex = Regex::new(pattern_str).map_err(|e| { - DataFusionError::Execution(format!( - "Invalid regex pattern '{}': {}", - pattern_str, e - )) - })?; let mut str_offsets = BufferBuilder::::new(8); let mut str_values = BufferBuilder::::new(s.len()); str_offsets.append(0); let mut scratch = Vec::new(); - push_split_parts( - s, - ®ex, - limit, - &mut str_offsets, - &mut str_values, - &mut scratch, - ); + if is_regex_literal(pattern_str) { + let mut chars = pattern_str.chars(); + if let (Some(ch), None) = (chars.next(), chars.next()) { + push_split_char( + s, + ch, + limit, + &mut str_offsets, + &mut str_values, + &mut scratch, + ); + } else { + push_split_literal( + s, + pattern_str, + limit, + &mut str_offsets, + &mut str_values, + &mut scratch, + ); + } + } else { + let regex = Regex::new(pattern_str).map_err(|e| { + DataFusionError::Execution(format!( + "Invalid regex pattern '{}': {}", + pattern_str, e + )) + })?; + push_split_parts( + s, + ®ex, + limit, + &mut str_offsets, + &mut str_values, + &mut scratch, + ); + } let item_offsets_buffer = OffsetBuffer::new(str_offsets.finish().into()); let item_values_buffer = str_values.finish(); @@ -282,6 +304,50 @@ fn push_split_literal<'a, O: OffsetSizeTrait>( } } +#[inline] +fn push_split_char<'a, O: OffsetSizeTrait>( + string: &'a str, + delimiter: char, + limit: i32, + offsets: &mut BufferBuilder, + values: &mut BufferBuilder, + scratch: &mut Vec<&'a str>, +) { + if limit == 0 { + scratch.clear(); + scratch.extend(string.split(delimiter)); // std::str::split(char) + while scratch.last().is_some_and(|s| s.is_empty()) { + scratch.pop(); + } + if scratch.is_empty() { + append_str("", offsets, values); + } else { + for &p in scratch.iter() { + append_str(p, offsets, values); + } + } + } else if limit > 0 { + let cap = (limit - 1) as usize; + let mut count = 0; + let mut last_end = 0; + // match_indices(char) в std работает через быстрый поиск символа + for (start, _) in string.match_indices(delimiter) { + if count >= cap { + break; + } + append_str(&string[last_end..start], offsets, values); + last_end = start + delimiter.len_utf8(); + count += 1; + } + append_str(&string[last_end..], offsets, values); + } else { + // limit < 0 + for p in string.split(delimiter) { + append_str(p, offsets, values); + } + } +} + fn split_generic_literal( string_array: &GenericStringArray, pattern: &str, @@ -300,19 +366,42 @@ fn split_generic_literal( let mut scratch = Vec::new(); list_offsets.push(O::usize_as(0)); - for i in 0..len { - if !string_array.is_null(i) { - let s = string_array.value(i); - push_split_literal( - s, - pattern, - limit, - &mut str_offsets, - &mut str_values, - &mut scratch, - ); + let mut chars = pattern.chars(); + let single_char = match (chars.next(), chars.next()) { + (Some(ch), None) => Some(ch), + _ => None, + }; + + if let Some(ch) = single_char { + for i in 0..len { + if !string_array.is_null(i) { + let s = string_array.value(i); + push_split_char( + s, + ch, + limit, + &mut str_offsets, + &mut str_values, + &mut scratch, + ); + } + list_offsets.push(O::usize_as(str_offsets.len() - 1)); + } + } else { + for i in 0..len { + if !string_array.is_null(i) { + let s = string_array.value(i); + push_split_literal( + s, + pattern, + limit, + &mut str_offsets, + &mut str_values, + &mut scratch, + ); + } + list_offsets.push(O::usize_as(str_offsets.len() - 1)); } - list_offsets.push(O::usize_as(str_offsets.len() - 1)); } let item_offsets_buffer = OffsetBuffer::new(str_offsets.finish().into()); From 5648d488db4522fd965ec9b7a8554ffcb39a9d2d Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sat, 22 Aug 2026 21:59:55 +0400 Subject: [PATCH 6/8] complete --- native/spark-expr/src/string_funcs/split.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/native/spark-expr/src/string_funcs/split.rs b/native/spark-expr/src/string_funcs/split.rs index a235ca333f0..24ec1e0885a 100644 --- a/native/spark-expr/src/string_funcs/split.rs +++ b/native/spark-expr/src/string_funcs/split.rs @@ -315,7 +315,7 @@ fn push_split_char<'a, O: OffsetSizeTrait>( ) { if limit == 0 { scratch.clear(); - scratch.extend(string.split(delimiter)); // std::str::split(char) + scratch.extend(string.split(delimiter)); while scratch.last().is_some_and(|s| s.is_empty()) { scratch.pop(); } @@ -330,7 +330,6 @@ fn push_split_char<'a, O: OffsetSizeTrait>( let cap = (limit - 1) as usize; let mut count = 0; let mut last_end = 0; - // match_indices(char) в std работает через быстрый поиск символа for (start, _) in string.match_indices(delimiter) { if count >= cap { break; From 4f382219f7c1dbddeff784501a5b1b518c859f6b Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sat, 22 Aug 2026 22:06:51 +0400 Subject: [PATCH 7/8] complete --- native/spark-expr/src/lib.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 60239e10776..2b5c29befc6 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -18,7 +18,6 @@ // The clippy throws an error if the reference clone not wrapped into `Arc::clone` // The lint makes easier for code reader/reviewer separate references clones from more heavyweight ones #![deny(clippy::clone_on_ref_ptr)] -extern crate core; mod error; mod query_context; From 4a2e50f4c3b01f360c9472e26b9dc700924e77c1 Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Mon, 24 Aug 2026 21:36:33 +0400 Subject: [PATCH 8/8] address comments --- native/spark-expr/src/string_funcs/split.rs | 111 ++++++++++++-------- 1 file changed, 69 insertions(+), 42 deletions(-) diff --git a/native/spark-expr/src/string_funcs/split.rs b/native/spark-expr/src/string_funcs/split.rs index 24ec1e0885a..d83dccdc077 100644 --- a/native/spark-expr/src/string_funcs/split.rs +++ b/native/spark-expr/src/string_funcs/split.rs @@ -16,9 +16,9 @@ // under the License. use arrow::array::{ Array, ArrayBuilder, ArrayRef, BufferBuilder, GenericListArray, GenericStringArray, - GenericStringBuilder, ListArray, NullBufferBuilder, OffsetSizeTrait, + GenericStringBuilder, ListArray, NullBufferBuilder, OffsetSizeTrait, StringArray, }; -use arrow::buffer::OffsetBuffer; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{DataType, Field}; use datafusion::common::{ cast::as_generic_string_array, exec_err, DataFusionError, Result as DataFusionResult, @@ -212,42 +212,40 @@ pub fn spark_split_sql(args: &[ColumnarValue]) -> DataFusionResult return exec_err!("split_sql delimiter must be a string"), }; - let string = string.clone().unwrap(); - let mut str_offsets = BufferBuilder::::new(8); - let mut str_values = BufferBuilder::::new(string.len()); - str_offsets.append(0); + + let mut offsets_builder = BufferBuilder::::new(2); + let mut values_builder = BufferBuilder::::new(string.len()); + + offsets_builder.append(0); if delimiter.is_empty() { - for ch in string.chars() { - append_str( - ch.encode_utf8(&mut [0; 4]), - &mut str_offsets, - &mut str_values, - ); - } + values_builder.append_slice(string.as_bytes()); + offsets_builder.append(string.len() as i32); } else { - for p in string.split(delimiter) { - append_str(p, &mut str_offsets, &mut str_values); + let mut offset = 0i32; + for part in string.split(delimiter.as_str()) { + values_builder.append_slice(part.as_bytes()); + offset += part.len() as i32; + offsets_builder.append(offset); } } - let item_offsets_buffer = OffsetBuffer::new(str_offsets.finish().into()); - let item_values_buffer = str_values.finish(); + let offsets_buffer = offsets_builder.finish(); + let values_buffer = values_builder.finish(); - let string_array_values = unsafe { - GenericStringArray::::new_unchecked( - item_offsets_buffer, - item_values_buffer, - None, - ) - }; + let list_field = Arc::new(Field::new("item", DataType::Utf8, true)); + let values_array = Arc::new(StringArray::try_new( + OffsetBuffer::new(offsets_buffer.into()), + values_buffer, + None, + )?); - let list_array = create_list_array(Arc::new(string_array_values)); + let list_offsets = + OffsetBuffer::new(ScalarBuffer::from(vec![0i32, values_array.len() as i32])); + let list_array = ListArray::try_new(list_field, list_offsets, values_array, None)?; - Ok(ColumnarValue::Scalar(ScalarValue::List(Arc::new( - list_array, - )))) + Ok(ColumnarValue::Array(Arc::new(list_array))) } _ => exec_err!("split_sql expects string arguments"), } @@ -286,15 +284,13 @@ fn push_split_literal<'a, O: OffsetSizeTrait>( } } else if limit > 0 { let cap = (limit - 1) as usize; - let mut count = 0; let mut last_end = 0; - for (start, _) in string.match_indices(delimiter) { + for (count, (start, _)) in string.match_indices(delimiter).enumerate() { if count >= cap { break; } append_str(&string[last_end..start], offsets, values); last_end = start + delimiter.len(); - count += 1; } append_str(&string[last_end..], offsets, values); } else { @@ -328,15 +324,13 @@ fn push_split_char<'a, O: OffsetSizeTrait>( } } else if limit > 0 { let cap = (limit - 1) as usize; - let mut count = 0; let mut last_end = 0; - for (start, _) in string.match_indices(delimiter) { + for (count, (start, _)) in string.match_indices(delimiter).enumerate() { if count >= cap { break; } append_str(&string[last_end..start], offsets, values); last_end = start + delimiter.len_utf8(); - count += 1; } append_str(&string[last_end..], offsets, values); } else { @@ -802,12 +796,7 @@ fn push_split_sql_parts( fn create_list_array(values: ArrayRef) -> ListArray { let field = Arc::new(Field::new("item", DataType::Utf8, false)); let offsets = vec![0i32, values.len() as i32]; - ListArray::new( - field, - arrow::buffer::OffsetBuffer::new(offsets.into()), - values, - None, - ) + ListArray::new(field, OffsetBuffer::new(offsets.into()), values, None) } fn new_null_list_array(len: usize) -> ArrayRef { @@ -831,7 +820,7 @@ fn new_null_list_array_with_offset(len: usize) -> ArrayRef { Arc::new(GenericListArray::::new( field, - arrow::buffer::OffsetBuffer::new(offsets.into()), + OffsetBuffer::new(offsets.into()), values, Some(nulls), )) @@ -845,7 +834,7 @@ fn new_null_list_array_value(len: usize) -> ListArray { ListArray::new( field, - arrow::buffer::OffsetBuffer::new(offsets.into()), + OffsetBuffer::new(offsets.into()), values, Some(nulls), ) @@ -969,6 +958,44 @@ mod tests { } } + #[test] + fn test_split_sql_empty_delimiter_scalar() { + let input = ColumnarValue::Scalar(ScalarValue::Utf8(Some("hello world".to_string()))); + let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8(Some("".to_string()))); + + let result = spark_split_sql(&[input, delimiter]) + .unwrap() + .into_array(1) + .unwrap(); + let list_array = result.as_any().downcast_ref::().unwrap(); + + assert_eq!(list_array.len(), 1); + let values = list_array.value(0); + let str_array = values.as_any().downcast_ref::().unwrap(); + + assert_eq!(str_array.len(), 1); + assert_eq!(str_array.value(0), "hello world"); + } + + #[test] + fn test_split_sql_empty_string_and_empty_delimiter_scalar() { + let input = ColumnarValue::Scalar(ScalarValue::Utf8(Some("".to_string()))); + let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8(Some("".to_string()))); + + let result = spark_split_sql(&[input, delimiter]) + .unwrap() + .into_array(1) + .unwrap(); + let list_array = result.as_any().downcast_ref::().unwrap(); + + assert_eq!(list_array.len(), 1); + let values = list_array.value(0); + let str_array = values.as_any().downcast_ref::().unwrap(); + + assert_eq!(str_array.len(), 1); + assert_eq!(str_array.value(0), ""); + } + fn assert_list_value(list_array: &ListArray, row: usize, expected: &[&str]) { let value = list_array.value(row); let strings = value.as_any().downcast_ref::().unwrap();