You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by mb...@apache.org on 2025/08/04 14:46:00 UTC
(datafusion-comet) branch main updated: Chore: implement string_space as ScalarUDFImpl (#2041)
This is an automated email from the ASF dual-hosted git repository.
mbutrovich 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 7247d9c33 Chore: implement string_space as ScalarUDFImpl (#2041)
7247d9c33 is described below
commit 7247d9c33db0e9179c2b0be2fe3b1d853bcc0809
Author: Kazantsev Maksim <ka...@yandex.ru>
AuthorDate: Mon Aug 4 18:45:54 2025 +0400
Chore: implement string_space as ScalarUDFImpl (#2041)
---
native/core/src/execution/planner.rs | 9 +-
native/proto/src/proto/expr.proto | 1 -
native/spark-expr/src/comet_scalar_funcs.rs | 2 +
native/spark-expr/src/kernels/strings.rs | 62 --------
native/spark-expr/src/string_funcs/mod.rs | 2 +-
native/spark-expr/src/string_funcs/string_space.rs | 167 +++++++++++++--------
.../org/apache/comet/serde/QueryPlanSerde.scala | 11 +-
.../org/apache/comet/CometExpressionSuite.scala | 16 --
.../apache/comet/CometStringExpressionSuite.scala | 16 ++
9 files changed, 128 insertions(+), 158 deletions(-)
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 057090563..5556bca53 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -104,8 +104,8 @@ use datafusion_comet_spark_expr::monotonically_increasing_id::MonotonicallyIncre
use datafusion_comet_spark_expr::{
ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct,
GetArrayStructFields, GetStructField, IfExpr, ListExtract, NormalizeNaNAndZero, RLike,
- RandExpr, RandnExpr, SparkCastOptions, Stddev, StringSpaceExpr, SubstringExpr, SumDecimal,
- TimestampTruncExpr, ToJson, UnboundColumn, Variance,
+ RandExpr, RandnExpr, SparkCastOptions, Stddev, SubstringExpr, SumDecimal, TimestampTruncExpr,
+ ToJson, UnboundColumn, Variance,
};
use itertools::Itertools;
use jni::objects::GlobalRef;
@@ -546,11 +546,6 @@ impl PhysicalPlanner {
len as u64,
)))
}
- ExprStruct::StringSpace(expr) => {
- let child = self.create_expr(expr.child.as_ref().unwrap(), input_schema)?;
-
- Ok(Arc::new(StringSpaceExpr::new(child)))
- }
ExprStruct::Like(expr) => {
let left =
self.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?;
diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto
index 50864e582..8b193ba84 100644
--- a/native/proto/src/proto/expr.proto
+++ b/native/proto/src/proto/expr.proto
@@ -45,7 +45,6 @@ message Expr {
BinaryExpr or = 18;
SortOrder sort_order = 19;
Substring substring = 20;
- UnaryExpr string_space = 21;
Hour hour = 22;
Minute minute = 23;
Second second = 24;
diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs
index d3f50b2fa..8a08b9336 100644
--- a/native/spark-expr/src/comet_scalar_funcs.rs
+++ b/native/spark-expr/src/comet_scalar_funcs.rs
@@ -22,6 +22,7 @@ use crate::{
spark_decimal_integral_div, spark_floor, spark_hex, spark_isnan, spark_make_decimal,
spark_read_side_padding, spark_round, spark_rpad, spark_unhex, spark_unscaled_value,
SparkBitwiseCount, SparkBitwiseGet, SparkBitwiseNot, SparkChrFunc, SparkDateTrunc,
+ SparkStringSpace,
};
use arrow::datatypes::DataType;
use datafusion::common::{DataFusionError, Result as DataFusionResult};
@@ -175,6 +176,7 @@ fn all_scalar_functions() -> Vec<Arc<ScalarUDF>> {
Arc::new(ScalarUDF::new_from_impl(SparkBitwiseCount::default())),
Arc::new(ScalarUDF::new_from_impl(SparkBitwiseGet::default())),
Arc::new(ScalarUDF::new_from_impl(SparkDateTrunc::default())),
+ Arc::new(ScalarUDF::new_from_impl(SparkStringSpace::default())),
]
}
diff --git a/native/spark-expr/src/kernels/strings.rs b/native/spark-expr/src/kernels/strings.rs
index f9e0b1311..25fbf6446 100644
--- a/native/spark-expr/src/kernels/strings.rs
+++ b/native/spark-expr/src/kernels/strings.rs
@@ -21,33 +21,11 @@ use std::sync::Arc;
use arrow::{
array::*,
- buffer::MutableBuffer,
compute::kernels::substring::{substring as arrow_substring, substring_by_char},
datatypes::{DataType, Int32Type},
};
use datafusion::common::DataFusionError;
-/// Returns an ArrayRef with a string consisting of `length` spaces.
-///
-/// # Preconditions
-///
-/// - elements in `length` must not be negative
-pub fn string_space(length: &dyn Array) -> Result<ArrayRef, DataFusionError> {
- match length.data_type() {
- DataType::Int32 => {
- let array = length.as_any().downcast_ref::<Int32Array>().unwrap();
- Ok(generic_string_space::<i32>(array))
- }
- DataType::Dictionary(_, _) => {
- let dict = as_dictionary_array::<Int32Type>(length);
- let values = string_space(dict.values())?;
- let result = DictionaryArray::try_new(dict.keys().clone(), values)?;
- Ok(Arc::new(result))
- }
- dt => panic!("Unsupported input type for function 'string_space': {dt:?}"),
- }
-}
-
pub fn substring(array: &dyn Array, start: i64, length: u64) -> Result<ArrayRef, DataFusionError> {
match array.data_type() {
DataType::LargeUtf8 => substring_by_char(
@@ -82,43 +60,3 @@ pub fn substring(array: &dyn Array, start: i64, length: u64) -> Result<ArrayRef,
dt => panic!("Unsupported input type for function 'substring': {dt:?}"),
}
}
-
-fn generic_string_space<OffsetSize: OffsetSizeTrait>(length: &Int32Array) -> ArrayRef {
- let array_len = length.len();
- let mut offsets = MutableBuffer::new((array_len + 1) * std::mem::size_of::<OffsetSize>());
- let mut length_so_far = OffsetSize::zero();
-
- // compute null bitmap (copy)
- let null_bit_buffer = length.to_data().nulls().map(|b| b.buffer().clone());
-
- // Gets slice of length array to access it directly for performance.
- let length_data = length.to_data();
- let lengths = length_data.buffers()[0].typed_data::<i32>();
- let total = lengths.iter().map(|l| *l as usize).sum::<usize>();
- let mut values = MutableBuffer::new(total);
-
- offsets.push(length_so_far);
-
- let blank = " ".as_bytes()[0];
- values.resize(total, blank);
-
- (0..array_len).for_each(|i| {
- let current_len = lengths[i] as usize;
-
- length_so_far += OffsetSize::from_usize(current_len).unwrap();
- offsets.push(length_so_far);
- });
-
- let data = unsafe {
- ArrayData::new_unchecked(
- GenericStringArray::<OffsetSize>::DATA_TYPE,
- array_len,
- None,
- null_bit_buffer,
- 0,
- vec![offsets.into(), values.into()],
- vec![],
- )
- };
- make_array(data)
-}
diff --git a/native/spark-expr/src/string_funcs/mod.rs b/native/spark-expr/src/string_funcs/mod.rs
index f4250c2d8..18d30b111 100644
--- a/native/spark-expr/src/string_funcs/mod.rs
+++ b/native/spark-expr/src/string_funcs/mod.rs
@@ -20,5 +20,5 @@ mod string_space;
mod substring;
pub use chr::SparkChrFunc;
-pub use string_space::StringSpaceExpr;
+pub use string_space::SparkStringSpace;
pub use substring::SubstringExpr;
diff --git a/native/spark-expr/src/string_funcs/string_space.rs b/native/spark-expr/src/string_funcs/string_space.rs
index 14da7a520..101328f22 100644
--- a/native/spark-expr/src/string_funcs/string_space.rs
+++ b/native/spark-expr/src/string_funcs/string_space.rs
@@ -15,94 +15,137 @@
// specific language governing permissions and limitations
// under the License.
-#![allow(deprecated)]
-
-use crate::kernels::strings::string_space;
-use arrow::datatypes::{DataType, Schema};
-use arrow::record_batch::RecordBatch;
-use datafusion::common::DataFusionError;
-use datafusion::logical_expr::ColumnarValue;
-use datafusion::physical_expr::PhysicalExpr;
-use std::{
- any::Any,
- fmt::{Display, Formatter},
- hash::Hash,
- sync::Arc,
+use arrow::array::{
+ as_dictionary_array, make_array, Array, ArrayData, ArrayRef, DictionaryArray,
+ GenericStringArray, Int32Array, OffsetSizeTrait,
};
+use arrow::buffer::MutableBuffer;
+use arrow::datatypes::{DataType, Int32Type};
+use datafusion::common::{exec_err, internal_datafusion_err, DataFusionError, Result};
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+use std::{any::Any, sync::Arc};
-#[derive(Debug, Eq)]
-pub struct StringSpaceExpr {
- pub child: Arc<dyn PhysicalExpr>,
-}
-
-impl Hash for StringSpaceExpr {
- fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
- self.child.hash(state);
- }
-}
-
-impl PartialEq for StringSpaceExpr {
- fn eq(&self, other: &Self) -> bool {
- self.child.eq(&other.child)
- }
+#[derive(Debug)]
+pub struct SparkStringSpace {
+ signature: Signature,
+ aliases: Vec<String>,
}
-impl StringSpaceExpr {
- pub fn new(child: Arc<dyn PhysicalExpr>) -> Self {
- Self { child }
+impl Default for SparkStringSpace {
+ fn default() -> Self {
+ Self::new()
}
}
-impl Display for StringSpaceExpr {
- fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
- write!(f, "StringSpace [child: {}] ", self.child)
+impl SparkStringSpace {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::user_defined(Volatility::Immutable),
+ aliases: vec![],
+ }
}
}
-impl PhysicalExpr for StringSpaceExpr {
+impl ScalarUDFImpl for SparkStringSpace {
fn as_any(&self) -> &dyn Any {
self
}
- fn fmt_sql(&self, _: &mut Formatter<'_>) -> std::fmt::Result {
- unimplemented!()
+ fn name(&self) -> &str {
+ "string_space"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
}
- fn data_type(&self, input_schema: &Schema) -> datafusion::common::Result<DataType> {
- match self.child.data_type(input_schema)? {
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ Ok(match &arg_types[0] {
DataType::Dictionary(key_type, _) => {
- Ok(DataType::Dictionary(key_type, Box::new(DataType::Utf8)))
+ DataType::Dictionary(key_type.clone(), Box::new(DataType::Utf8))
}
- _ => Ok(DataType::Utf8),
- }
+ _ => DataType::Utf8,
+ })
}
- fn nullable(&self, _: &Schema) -> datafusion::common::Result<bool> {
- Ok(true)
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ let args: [ColumnarValue; 1] = args
+ .args
+ .try_into()
+ .map_err(|_| internal_datafusion_err!("string_space expects exactly one argument"))?;
+ spark_string_space(&args)
}
- fn evaluate(&self, batch: &RecordBatch) -> datafusion::common::Result<ColumnarValue> {
- let arg = self.child.evaluate(batch)?;
- match arg {
- ColumnarValue::Array(array) => {
- let result = string_space(&array)?;
+ fn aliases(&self) -> &[String] {
+ &self.aliases
+ }
+}
- Ok(ColumnarValue::Array(result))
- }
- _ => Err(DataFusionError::Execution(
- "StringSpace(scalar) should be fold in Spark JVM side.".to_string(),
- )),
+pub fn spark_string_space(args: &[ColumnarValue; 1]) -> Result<ColumnarValue> {
+ match args {
+ [ColumnarValue::Array(array)] => {
+ let result = string_space(&array)?;
+
+ Ok(ColumnarValue::Array(result))
}
+ _ => exec_err!("StringSpace(scalar) should be fold in Spark JVM side."),
}
+}
- fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
- vec![&self.child]
+fn string_space(length: &dyn Array) -> std::result::Result<ArrayRef, DataFusionError> {
+ match length.data_type() {
+ DataType::Int32 => {
+ let array = length.as_any().downcast_ref::<Int32Array>().unwrap();
+ Ok(generic_string_space::<i32>(array))
+ }
+ DataType::Dictionary(_, _) => {
+ let dict = as_dictionary_array::<Int32Type>(length);
+ let values = string_space(dict.values())?;
+ let result = DictionaryArray::try_new(dict.keys().clone(), values)?;
+ Ok(Arc::new(result))
+ }
+ other => exec_err!("Unsupported input type for function 'string_space': {other:?}"),
}
+}
- fn with_new_children(
- self: Arc<Self>,
- children: Vec<Arc<dyn PhysicalExpr>>,
- ) -> datafusion::common::Result<Arc<dyn PhysicalExpr>> {
- Ok(Arc::new(StringSpaceExpr::new(Arc::clone(&children[0]))))
- }
+fn generic_string_space<OffsetSize: OffsetSizeTrait>(length: &Int32Array) -> ArrayRef {
+ let array_len = length.len();
+ let mut offsets = MutableBuffer::new((array_len + 1) * std::mem::size_of::<OffsetSize>());
+ let mut length_so_far = OffsetSize::zero();
+
+ // compute null bitmap (copy)
+ let null_bit_buffer = length.to_data().nulls().map(|b| b.buffer().clone());
+
+ // Gets slice of length array to access it directly for performance.
+ let length_data = length.to_data();
+ let lengths = length_data.buffers()[0].typed_data::<i32>();
+ let total = lengths.iter().map(|l| *l as usize).sum::<usize>();
+ let mut values = MutableBuffer::new(total);
+
+ offsets.push(length_so_far);
+
+ let blank = " ".as_bytes()[0];
+ values.resize(total, blank);
+
+ (0..array_len).for_each(|i| {
+ let current_len = lengths[i] as usize;
+
+ length_so_far += OffsetSize::from_usize(current_len).unwrap();
+ offsets.push(length_so_far);
+ });
+
+ let data = unsafe {
+ ArrayData::new_unchecked(
+ GenericStringArray::<OffsetSize>::DATA_TYPE,
+ array_len,
+ None,
+ null_bit_buffer,
+ 0,
+ vec![offsets.into(), values.into()],
+ vec![],
+ )
+ };
+ make_array(data)
}
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 022480cc9..b08cb89df 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -130,7 +130,8 @@ object QueryPlanSerde extends Logging with CometExprShim {
classOf[Rand] -> CometRand,
classOf[Randn] -> CometRandn,
classOf[SparkPartitionID] -> CometSparkPartitionId,
- classOf[MonotonicallyIncreasingID] -> CometMonotonicallyIncreasingId)
+ classOf[MonotonicallyIncreasingID] -> CometMonotonicallyIncreasingId,
+ classOf[StringSpace] -> UnaryScalarFuncSerde("string_space"))
/**
* Mapping of Spark aggregate expression class to Comet expression handler.
@@ -943,14 +944,6 @@ object QueryPlanSerde extends Logging with CometExprShim {
val valueExpr = exprToProtoInternal(value, inputs, binding)
scalarFunctionExprToProto("contains", attributeExpr, valueExpr)
- case StringSpace(child) =>
- createUnaryExpr(
- expr,
- child,
- inputs,
- binding,
- (builder, unaryExpr) => builder.setStringSpace(unaryExpr))
-
case Hour(child, timeZoneId) =>
val childExpr = exprToProtoInternal(child, inputs, binding)
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index 4f4554d9e..2015d3663 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -381,22 +381,6 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
- test("string_space") {
- withParquetTable((0 until 5).map(i => (i, i + 1)), "tbl") {
- checkSparkAnswerAndOperator("SELECT space(_1), space(_2) FROM tbl")
- }
- }
-
- test("string_space with dictionary") {
- val data = (0 until 1000).map(i => Tuple1(i % 5))
-
- withSQLConf("parquet.enable.dictionary" -> "true") {
- withParquetTable(data, "tbl") {
- checkSparkAnswerAndOperator("SELECT space(_1) FROM tbl")
- }
- }
- }
-
test("hour, minute, second") {
Seq(true, false).foreach { dictionaryEnabled =>
withTempDir { dir =>
diff --git a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
index 429dc7e1e..fac80853f 100644
--- a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
@@ -249,4 +249,20 @@ class CometStringExpressionSuite extends CometTestBase {
}
}
+ test("string_space") {
+ withParquetTable((0 until 5).map(i => (i, i + 1)), "tbl") {
+ checkSparkAnswerAndOperator("SELECT space(_1), space(_2) FROM tbl")
+ }
+ }
+
+ test("string_space with dictionary") {
+ val data = (0 until 1000).map(i => Tuple1(i % 5))
+
+ withSQLConf("parquet.enable.dictionary" -> "true") {
+ withParquetTable(data, "tbl") {
+ checkSparkAnswerAndOperator("SELECT space(_1) FROM tbl")
+ }
+ }
+ }
+
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org