You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by co...@apache.org on 2025/09/25 01:27:15 UTC
(datafusion-comet) branch main updated: feat:add_additional_char_support_rpad (#2436)
This is an automated email from the ASF dual-hosted git repository.
comphead 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 abf5b6912 feat:add_additional_char_support_rpad (#2436)
abf5b6912 is described below
commit abf5b6912605825a1868b872bc4d6189721e3754
Author: B Vadlamani <11...@users.noreply.github.com>
AuthorDate: Wed Sep 24 18:27:11 2025 -0700
feat:add_additional_char_support_rpad (#2436)
* add_additional_char_support_rpad
* rebase
* rebase
---
.../char_varchar_utils/read_side_padding.rs | 78 +++++++++++++++++++++-
.../scala/org/apache/comet/serde/strings.scala | 16 ++---
.../org/apache/comet/CometExpressionSuite.scala | 10 +++
3 files changed, 92 insertions(+), 12 deletions(-)
diff --git a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs
index 40735de7a..166bb6ddf 100644
--- a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs
+++ b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs
@@ -25,6 +25,7 @@ use datafusion::common::{cast::as_generic_string_array, DataFusionError, ScalarV
use datafusion::physical_plan::ColumnarValue;
use std::sync::Arc;
+const SPACE: &str = " ";
/// Similar to DataFusion `rpad`, but not to truncate when the string is already longer than length
pub fn spark_read_side_padding(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
spark_read_side_padding2(args, false)
@@ -46,11 +47,13 @@ fn spark_read_side_padding2(
array,
truncate,
ColumnarValue::Scalar(ScalarValue::Int32(Some(*length))),
+ SPACE,
),
DataType::LargeUtf8 => spark_read_side_padding_internal::<i64>(
array,
truncate,
ColumnarValue::Scalar(ScalarValue::Int32(Some(*length))),
+ SPACE,
),
// Dictionary support required for SPARK-48498
DataType::Dictionary(_, value_type) => {
@@ -60,12 +63,57 @@ fn spark_read_side_padding2(
dict.values(),
truncate,
ColumnarValue::Scalar(ScalarValue::Int32(Some(*length))),
+ SPACE,
)?
} else {
spark_read_side_padding_internal::<i64>(
dict.values(),
truncate,
ColumnarValue::Scalar(ScalarValue::Int32(Some(*length))),
+ SPACE,
+ )?
+ };
+ // col consists of an array, so arg of to_array() is not used. Can be anything
+ let values = col.to_array(0)?;
+ let result = DictionaryArray::try_new(dict.keys().clone(), values)?;
+ Ok(ColumnarValue::Array(make_array(result.into())))
+ }
+ other => Err(DataFusionError::Internal(format!(
+ "Unsupported data type {other:?} for function rpad/read_side_padding",
+ ))),
+ }
+ }
+ [ColumnarValue::Array(array), ColumnarValue::Scalar(ScalarValue::Int32(Some(length))), ColumnarValue::Scalar(ScalarValue::Utf8(Some(string)))] =>
+ {
+ match array.data_type() {
+ DataType::Utf8 => spark_read_side_padding_internal::<i32>(
+ array,
+ truncate,
+ ColumnarValue::Scalar(ScalarValue::Int32(Some(*length))),
+ string,
+ ),
+ DataType::LargeUtf8 => spark_read_side_padding_internal::<i64>(
+ array,
+ truncate,
+ ColumnarValue::Scalar(ScalarValue::Int32(Some(*length))),
+ string,
+ ),
+ // Dictionary support required for SPARK-48498
+ DataType::Dictionary(_, value_type) => {
+ let dict = as_dictionary_array::<Int32Type>(array);
+ let col = if value_type.as_ref() == &DataType::Utf8 {
+ spark_read_side_padding_internal::<i32>(
+ dict.values(),
+ truncate,
+ ColumnarValue::Scalar(ScalarValue::Int32(Some(*length))),
+ SPACE,
+ )?
+ } else {
+ spark_read_side_padding_internal::<i64>(
+ dict.values(),
+ truncate,
+ ColumnarValue::Scalar(ScalarValue::Int32(Some(*length))),
+ SPACE,
)?
};
// col consists of an array, so arg of to_array() is not used. Can be anything
@@ -83,16 +131,37 @@ fn spark_read_side_padding2(
array,
truncate,
ColumnarValue::Array(Arc::<dyn Array>::clone(array_int)),
+ SPACE,
),
DataType::LargeUtf8 => spark_read_side_padding_internal::<i64>(
array,
truncate,
ColumnarValue::Array(Arc::<dyn Array>::clone(array_int)),
+ SPACE,
),
other => Err(DataFusionError::Internal(format!(
"Unsupported data type {other:?} for function rpad/read_side_padding",
))),
},
+ [ColumnarValue::Array(array), ColumnarValue::Array(array_int), ColumnarValue::Scalar(ScalarValue::Utf8(Some(string)))] => {
+ match array.data_type() {
+ DataType::Utf8 => spark_read_side_padding_internal::<i32>(
+ array,
+ truncate,
+ ColumnarValue::Array(Arc::<dyn Array>::clone(array_int)),
+ string,
+ ),
+ DataType::LargeUtf8 => spark_read_side_padding_internal::<i64>(
+ array,
+ truncate,
+ ColumnarValue::Array(Arc::<dyn Array>::clone(array_int)),
+ string,
+ ),
+ other => Err(DataFusionError::Internal(format!(
+ "Unsupported data type {other:?} for function rpad/read_side_padding",
+ ))),
+ }
+ }
other => Err(DataFusionError::Internal(format!(
"Unsupported arguments {other:?} for function rpad/read_side_padding",
))),
@@ -103,6 +172,7 @@ fn spark_read_side_padding_internal<T: OffsetSizeTrait>(
array: &ArrayRef,
truncate: bool,
pad_type: ColumnarValue,
+ pad_string: &str,
) -> Result<ColumnarValue, DataFusionError> {
let string_array = as_generic_string_array::<T>(array)?;
match pad_type {
@@ -120,6 +190,7 @@ fn spark_read_side_padding_internal<T: OffsetSizeTrait>(
string.parse().unwrap(),
length.unwrap() as usize,
truncate,
+ pad_string,
)?),
_ => builder.append_null(),
}
@@ -140,6 +211,7 @@ fn spark_read_side_padding_internal<T: OffsetSizeTrait>(
string.parse().unwrap(),
length,
truncate,
+ pad_string,
)?),
_ => builder.append_null(),
}
@@ -153,10 +225,10 @@ fn add_padding_string(
string: String,
length: usize,
truncate: bool,
+ pad_string: &str,
) -> Result<String, DataFusionError> {
// It looks Spark's UTF8String is closer to chars rather than graphemes
// https://stackoverflow.com/a/46290728
- let space_string = " ".repeat(length);
let char_len = string.chars().count();
if length <= char_len {
if truncate {
@@ -176,6 +248,8 @@ fn add_padding_string(
Ok(string)
}
} else {
- Ok(string + &space_string[char_len..])
+ let pad_needed = length - char_len;
+ let pad: String = pad_string.chars().cycle().take(pad_needed).collect();
+ Ok(string + &pad)
}
}
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 4d183f6fa..36df9ed1c 100644
--- a/spark/src/main/scala/org/apache/comet/serde/strings.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala
@@ -159,16 +159,12 @@ object CometStringRPad extends CometExpressionSerde[StringRPad] {
expr: StringRPad,
inputs: Seq[Attribute],
binding: Boolean): Option[Expr] = {
- expr.pad match {
- case Literal(str, DataTypes.StringType) if str.toString == " " =>
- scalarFunctionExprToProto(
- "rpad",
- exprToProtoInternal(expr.str, inputs, binding),
- exprToProtoInternal(expr.len, inputs, binding))
- case _ =>
- withInfo(expr, "StringRPad with non-space characters is not supported")
- None
- }
+
+ scalarFunctionExprToProto(
+ "rpad",
+ exprToProtoInternal(expr.str, inputs, binding),
+ exprToProtoInternal(expr.len, inputs, binding),
+ exprToProtoInternal(expr.pad, 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 55b5121dd..600f4e45b 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -415,6 +415,16 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("RPAD with character support other than default space") {
+ val data = Seq(("IfIWasARoadIWouldBeBent", 10), ("hi", 2))
+ withParquetTable(data, "t1") {
+ val res = sql(
+ """ select rpad(_1,_2,'?'), rpad(_1,_2,'??') , rpad(_1,2, '??'), hex(rpad(unhex('aabb'), 5)),
+ rpad(_1, 5, '??') from t1 order by _1 """.stripMargin)
+ checkSparkAnswerAndOperator(res)
+ }
+ }
+
test("dictionary arithmetic") {
// TODO: test ANSI mode
withSQLConf(SQLConf.ANSI_ENABLED.key -> "false", "parquet.enable.dictionary" -> "true") {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org