You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/09/08 16:22:20 UTC
(datafusion-comet) branch main updated: feat: support native concat_ws with string arrays (#5725)
This is an automated email from the ASF dual-hosted git repository.
sunchao 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 8e6846850c feat: support native concat_ws with string arrays (#5725)
8e6846850c is described below
commit 8e6846850c525506dd2b9194f2014e8acd2ab60a
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Wed Sep 9 00:22:15 2026 +0800
feat: support native concat_ws with string arrays (#5725)
* feat: execute concat_ws with array arguments natively
* fix: handle concat_ws runtime scalars and retain string kernel
* bench: compare concat_ws native and old dispatcher paths
---
.../expression-audits/string_funcs.md | 7 +-
docs/source/user-guide/latest/expressions.md | 2 +-
native/spark-expr/src/comet_scalar_funcs.rs | 3 +
native/spark-expr/src/string_funcs/concat_ws.rs | 64 ++++++++++++
native/spark-expr/src/string_funcs/mod.rs | 2 +
native/spark-expr/tests/spark_expr_reg.rs | 69 +++++++++++++
.../scala/org/apache/comet/serde/strings.scala | 38 +++----
.../sql-tests/expressions/string/concat_ws.sql | 21 ++--
.../apache/comet/CometStringExpressionSuite.scala | 46 ++++++---
.../sql/benchmark/CometConcatWsBenchmark.scala | 109 +++++++++++++++++++++
10 files changed, 308 insertions(+), 53 deletions(-)
diff --git a/docs/source/contributor-guide/expression-audits/string_funcs.md b/docs/source/contributor-guide/expression-audits/string_funcs.md
index c1db191fdc..d9a4c92bd6 100644
--- a/docs/source/contributor-guide/expression-audits/string_funcs.md
+++ b/docs/source/contributor-guide/expression-audits/string_funcs.md
@@ -66,9 +66,10 @@
## concat_ws
-- Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8.
-- Spark 3.5.8 (audited 2026-05-27): baseline. `Seq[Expression] -> StringType`; NULL separator yields NULL, NULL element values are skipped, children can be `StringType` or `ArrayType(StringType)`. Comet serde rewrites a NULL-literal separator to a NULL of the result type and bails out on all-foldable inputs so Spark's `ConstantFolding` handles them; otherwise delegates to DataFusion `concat_ws`.
-- Spark 4.0.1 (audited 2026-05-27): `inputTypes` widened to `StringTypeWithCollation` / `AbstractArrayType`; `dataType` becomes `children.head.dataType` (collation-derived). Semantics unchanged for `UTF8_BINARY`. Non-default collations not honoured by Comet ([#4496](https://github.com/apache/datafusion-comet/issues/4496)).
+- Spark 3.4.3 (audited 2026-09-06): identical to 3.5.8.
+- Spark 3.5.8 (audited 2026-09-06): accepts ordered mixtures of strings and string arrays, skips null arguments and array elements, and returns NULL for a NULL separator. Comet retains DataFusion's string kernel and uses DataFusion Spark's `SparkConcatWs` for arrays and separator-only inputs, with native column coverage for multiple arrays, empty arrays, nulls, empty strings, Unicode, and varying separators. The existing all-foldable codegen fallback remains. A native adapter evaluates r [...]
+- Spark 4.0.1 (audited 2026-09-06): collation-aware input and result types; concatenation still uses `UTF8String.concatWs` without collation-dependent comparisons.
+- Spark 4.1.1 (audited 2026-09-06): same concatenation and null semantics as 4.0.1.
## contains
diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
index 2cf0894ceb..1fb979fd3c 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -561,7 +561,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
| `chr` | β
| Native | |
| `collate` | π | β | Spark collation (umbrella [#2190](https://github.com/apache/datafusion-comet/issues/2190)) |
| `collation` | β
| β | Constant-folded to a literal (Spark 4.0+) |
-| `concat_ws` | β
| Hybrid | Array arguments route through the JVM codegen dispatcher; string arguments run natively ([details](compatibility/expressions/string.md)) |
+| `concat_ws` | β
| Hybrid | Mixed string and array arguments run natively; all-foldable arguments use codegen dispatch |
| `contains` | β
| β | |
| `decode` | β
| β | |
| `elt` | β
| Codegen dispatch | |
diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs
index 2bfbba8f8a..b5820144ea 100644
--- a/native/spark-expr/src/comet_scalar_funcs.rs
+++ b/native/spark-expr/src/comet_scalar_funcs.rs
@@ -118,6 +118,9 @@ pub fn create_comet_physical_fun_with_eval_mode(
) -> Result<Arc<ScalarUDF>, DataFusionError> {
let fail_on_error = fail_on_error.unwrap_or(false);
match fun_name {
+ "spark_concat_ws" => Ok(Arc::new(ScalarUDF::new_from_impl(
+ crate::string_funcs::CometConcatWs::default(),
+ ))),
"ceil" => {
make_comet_scalar_udf!("ceil", spark_ceil, data_type)
}
diff --git a/native/spark-expr/src/string_funcs/concat_ws.rs b/native/spark-expr/src/string_funcs/concat_ws.rs
new file mode 100644
index 0000000000..799fb92be0
--- /dev/null
+++ b/native/spark-expr/src/string_funcs/concat_ws.rs
@@ -0,0 +1,64 @@
+// 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::datatypes::DataType;
+use datafusion::common::{exec_err, Result, ScalarValue};
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
+use datafusion_spark::function::string::concat_ws::SparkConcatWs;
+
+/// Adapt SparkConcatWs to runtime scalars, including non-foldable scalar subqueries.
+#[derive(Debug, Default, PartialEq, Eq, Hash)]
+pub struct CometConcatWs(SparkConcatWs);
+
+impl ScalarUDFImpl for CometConcatWs {
+ fn name(&self) -> &str {
+ self.0.name()
+ }
+
+ fn signature(&self) -> &Signature {
+ self.0.signature()
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ self.0.return_type(arg_types)
+ }
+
+ fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+ self.0.coerce_types(arg_types)
+ }
+
+ fn invoke_with_args(&self, mut args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ if args.args.is_empty() {
+ return exec_err!("concat_ws requires a separator");
+ }
+ let all_scalar = args
+ .args
+ .iter()
+ .all(|arg| matches!(arg, ColumnarValue::Scalar(_)));
+ // The upstream kernel expands scalars to one row, independently of number_rows.
+ // Evaluate once and let DataFusion broadcast the scalar to the enclosing batch.
+ if all_scalar {
+ args.number_rows = 1;
+ }
+ match (all_scalar, self.0.invoke_with_args(args)?) {
+ (true, ColumnarValue::Array(array)) => Ok(ColumnarValue::Scalar(
+ ScalarValue::try_from_array(&array, 0)?,
+ )),
+ (_, result) => Ok(result),
+ }
+ }
+}
diff --git a/native/spark-expr/src/string_funcs/mod.rs b/native/spark-expr/src/string_funcs/mod.rs
index dc51cfea1b..94afc12d05 100644
--- a/native/spark-expr/src/string_funcs/mod.rs
+++ b/native/spark-expr/src/string_funcs/mod.rs
@@ -16,6 +16,7 @@
// under the License.
mod base64;
+mod concat_ws;
mod contains;
mod get_json_object;
mod levenshtein;
@@ -26,6 +27,7 @@ mod split;
mod unbase64;
pub use base64::spark_base64;
+pub use concat_ws::CometConcatWs;
pub use contains::SparkContains;
pub use get_json_object::spark_get_json_object;
pub use levenshtein::spark_levenshtein;
diff --git a/native/spark-expr/tests/spark_expr_reg.rs b/native/spark-expr/tests/spark_expr_reg.rs
index 1049f409b0..4931255029 100644
--- a/native/spark-expr/tests/spark_expr_reg.rs
+++ b/native/spark-expr/tests/spark_expr_reg.rs
@@ -25,6 +25,75 @@ mod tests {
use datafusion_comet_spark_expr::create_comet_physical_fun;
use datafusion_comet_spark_expr::register_all_comet_functions;
+ #[test]
+ fn test_concat_ws_runtime_scalars() -> Result<()> {
+ use arrow::datatypes::Field;
+ use datafusion::common::{config::ConfigOptions, ScalarValue};
+ use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs};
+ use std::sync::Arc;
+
+ let state = SessionStateBuilder::new().build();
+ let udf = create_comet_physical_fun("spark_concat_ws", DataType::Utf8, &state, None)?;
+ let value = ScalarValue::Utf8(Some("a".into()));
+ let list = ScalarValue::List(ScalarValue::new_list(
+ &[
+ value.clone(),
+ ScalarValue::Utf8(None),
+ ScalarValue::Utf8(Some("".into())),
+ ],
+ &DataType::Utf8,
+ true,
+ ));
+ for (input, expected) in [(value, "a"), (list, "a,")] {
+ for number_rows in [8, 1, 0] {
+ for separator in [Some(",".to_string()), None] {
+ let expected =
+ ScalarValue::Utf8(separator.as_ref().map(|_| expected.to_string()));
+ let result = udf.invoke_with_args(ScalarFunctionArgs {
+ args: vec![
+ ColumnarValue::Scalar(ScalarValue::Utf8(separator)),
+ ColumnarValue::Scalar(input.clone()),
+ ],
+ arg_fields: vec![],
+ number_rows,
+ return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
+ config_options: Arc::new(ConfigOptions::default()),
+ })?;
+ let ColumnarValue::Scalar(result) = result else {
+ panic!("expected scalar output");
+ };
+ assert_eq!(result, expected);
+ }
+ }
+ }
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn test_concat_ws_array_arguments() -> Result<()> {
+ use arrow::array::{AsArray, StringArray};
+
+ let ctx = SessionContext::new();
+ let state = ctx.state();
+ let udf = create_comet_physical_fun("spark_concat_ws", DataType::Utf8, &state, None)?;
+ ctx.register_udf(udf.as_ref().clone());
+ let results = ctx
+ .sql(
+ "SELECT concat_ws(sep, make_array(a, NULL, b), 'tail', make_array(b, a))
+ FROM (VALUES ('|', '', 'Γ©'), ('', NULL, 'x'), (NULL, 'a', 'b'))
+ AS t(sep, a, b)",
+ )
+ .await?
+ .collect()
+ .await?;
+ assert_eq!(results.len(), 1);
+ assert_eq!(
+ results[0].column(0).as_string::<i32>(),
+ &StringArray::from(vec![Some("|Γ©|tail|Γ©|"), Some("xtailx"), None])
+ );
+ Ok(())
+ }
+
#[tokio::test]
async fn test_udf_registration() -> Result<()> {
// 1. Setup session with UDF registration of existing Spark-compatible expression
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 32cb6b6641..4665c95461 100644
--- a/spark/src/main/scala/org/apache/comet/serde/strings.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala
@@ -297,43 +297,31 @@ object CometConcat
}
object CometConcatWs extends CometExpressionSerde[ConcatWs] with CodegenDispatchFallback {
-
- // Spark's ConcatWs accepts `array<string>` arguments after the separator and flattens their
- // elements into the list of strings to join (null elements are skipped, like null strings).
- // DataFusion's `concat_ws` accepts only string arguments and fails at execution time when it is
- // handed a list, so these calls have no native path. Because this serde mixes in
- // `CodegenDispatchFallback`, they run through the JVM codegen dispatcher (Spark's own
- // `ConcatWs.doGenCode` inside the Comet pipeline) instead, and fall back to Spark only when the
- // dispatcher is disabled. See https://github.com/apache/datafusion-comet/issues/5675.
- private val arrayArgumentReason =
- "`concat_ws` with `array<string>` arguments: Spark flattens the array elements into the " +
- "strings to join, which DataFusion's `concat_ws` does not support " +
- "(https://github.com/apache/datafusion-comet/issues/5675)"
-
- override def getUnsupportedReasons(): Seq[String] = Seq(arrayArgumentReason)
+ override def getUnsupportedReasons(): Seq[String] = Seq("all arguments are foldable")
override def getSupportLevel(expr: ConcatWs): SupportLevel = expr.children.headOption match {
- // A NULL separator converts directly to a NULL result, so it stays supported.
case Some(Literal(null, _)) => Compatible()
- case _ if expr.children.exists(_.dataType.isInstanceOf[ArrayType]) =>
- Unsupported(Some(arrayArgumentReason))
- // Decline all-literal args so that Spark's ConstantFolding handles them (it normally folds
- // them before they reach Comet). With the dispatcher enabled they run through Spark's own
- // generated code in-pipeline; otherwise the projection falls back to Spark.
case _ if expr.children.forall(_.foldable) =>
+ // Preserve the existing ConstantFolding/codegen behavior for foldable expressions.
+ // Non-foldable runtime scalars are handled by the native adapter.
Unsupported(Some("all arguments are foldable"))
case _ => Compatible()
}
override def convert(expr: ConcatWs, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = {
expr.children.headOption match {
- // Match Spark behavior: when the separator is NULL, the result of concat_ws is NULL.
case Some(Literal(null, _)) =>
- val nullLiteral = Literal.create(null, expr.dataType)
- exprToProtoInternal(nullLiteral, inputs, binding)
-
+ exprToProtoInternal(Literal.create(null, expr.dataType), inputs, binding)
+ case _
+ if expr.children.length == 1 ||
+ expr.children.exists(_.dataType.isInstanceOf[ArrayType]) =>
+ val childExprs = expr.children.map(exprToProtoInternal(_, inputs, binding))
+ scalarFunctionExprToProtoWithReturnType(
+ "spark_concat_ws",
+ expr.dataType,
+ false,
+ childExprs: _*)
case _ =>
- // For all other cases, use the generic scalar function implementation.
CometScalarFunction[ConcatWs]("concat_ws").convert(expr, inputs, binding)
}
}
diff --git a/spark/src/test/resources/sql-tests/expressions/string/concat_ws.sql b/spark/src/test/resources/sql-tests/expressions/string/concat_ws.sql
index fe6df743f7..801c7f8271 100644
--- a/spark/src/test/resources/sql-tests/expressions/string/concat_ws.sql
+++ b/spark/src/test/resources/sql-tests/expressions/string/concat_ws.sql
@@ -40,17 +40,11 @@ INSERT INTO names VALUES(1, 'James', 'B', 'Taylor'), (2, 'Smith', 'C', 'Davis'),
query
SELECT concat_ws(' ', first_name, middle_initial, last_name) FROM names
--- literal + literal + literal (declined natively when all args are foldable; routed through the
--- JVM codegen dispatcher, or falls back to Spark when the dispatcher is disabled)
+-- All-scalar inputs retain codegen dispatch.
query spark_answer_only
SELECT concat_ws(',', 'hello', 'world'), concat_ws(',', '', ''), concat_ws(',', NULL, 'b', 'c'), concat_ws(NULL, 'a', 'b')
--- https://github.com/apache/datafusion-comet/issues/5675
--- Spark accepts array<string> arguments after the separator and flattens their elements into the
--- strings to join (skipping null elements). DataFusion's concat_ws rejects list arguments, so
--- ConcatWs mixes in CodegenDispatchFallback: these calls route through the JVM codegen dispatcher
--- (Spark's own ConcatWs.doGenCode inside the Comet pipeline) and stay native while matching Spark
--- exactly.
+-- Mixed string and array arguments execute natively.
statement
CREATE TABLE test_concat_ws_array(arr array<string>, s string) USING parquet
@@ -85,3 +79,14 @@ SELECT concat_ws(NULL, arr, s) FROM test_concat_ws_array
-- literal array arguments
query
SELECT concat_ws(',', array('a', 'b'), 'c'), concat_ws(',', array('x', NULL, 'y')), concat_ws(',', CAST(array() AS array<string>), 'w')
+
+-- Non-foldable separators, including empty strings and NULL.
+query
+SELECT concat_ws(s, arr, 'tail', arr), concat_ws(s) FROM test_concat_ws_array
+
+-- Empty strings count as elements; null strings, arrays, and elements do not.
+query
+SELECT concat_ws(',', array('', NULL, 'Γ©'), s, CAST(NULL AS array<string>), array(NULL, '', 'δΈη')) FROM test_concat_ws_array
+
+query
+SELECT concat_ws(','), concat_ws(NULL), concat_ws(',', array(NULL, NULL)), concat_ws(',', CAST(NULL AS array<string>))
diff --git a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
index 42d9b11f63..e67125e389 100644
--- a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
@@ -726,18 +726,34 @@ class CometStringExpressionSuite extends CometTestBase with CometCodegenAssertio
// scalastyle:on
}
+ test("concat_ws with scalar subqueries over a multi-row batch") {
+ withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") {
+ withParquetTable((1 to 32).map(i => (i, "row")), "fact") {
+ withParquetTable(Seq(Tuple1("a"), Tuple1("b")), "lookup") {
+ for (subquery <- Seq(
+ "(SELECT max(_1) FROM lookup)",
+ "(SELECT max(_1) FROM lookup WHERE _1 = 'missing')")) {
+ checkSparkAnswerAndOperator(s"SELECT _1, concat_ws($subquery) FROM fact")
+ checkSparkAnswerAndOperator(s"SELECT _1, concat_ws(',', $subquery) FROM fact")
+ checkSparkAnswerAndOperator(
+ s"SELECT _1, concat_ws(',', array('x', NULL, ''), $subquery) FROM fact")
+ }
+ }
+ }
+ }
+ }
+
test("concat_ws with array<string> arguments") {
- // https://github.com/apache/datafusion-comet/issues/5675
- // Spark flattens array<string> arguments into the strings to join (skipping null elements).
- // DataFusion's concat_ws rejects list arguments, so these calls run through the JVM codegen
- // dispatcher (Spark's own doGenCode inside the Comet pipeline) instead of the native path.
val data: Seq[(Seq[String], String)] = Seq(
(Seq("a", "b"), "c d"),
(Seq("x", null, "y"), "z"),
(Seq("only"), ""),
(Seq.empty[String], "w"),
(null, "v"),
- (Seq("p", "q"), null))
+ (Seq("p", "q"), null),
+ (Seq(null, "", "\u00e9"), "|"),
+ (Seq(null, null), ""),
+ (null, null))
withParquetTable(data, "tbl") {
val arrayArgQueries = Seq(
"SELECT concat_ws(',', _1, _2) FROM tbl",
@@ -745,18 +761,16 @@ class CometStringExpressionSuite extends CometTestBase with CometCodegenAssertio
"SELECT concat_ws(',', _1) FROM tbl",
"SELECT concat_ws('-', _1, _2, _1) FROM tbl",
"SELECT concat_ws(',', split(_2, ' ')) FROM tbl",
- "SELECT concat_ws(',', split(_2, ' '), _1) FROM tbl")
- for (query <- arrayArgQueries) {
- // Spark's answer, the whole plan stays in Comet, and the codegen dispatcher actually ran.
- assertCodegenRan {
- checkSparkAnswerAndOperator(query)
- }
- }
- // With the dispatcher disabled there is no in-pipeline path, so the projection falls back
- // to Spark with the serde's reason instead of failing at native execution.
- withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") {
+ "SELECT concat_ws(',', split(_2, ' '), _1) FROM tbl",
+ "SELECT concat_ws(_2, _1, 'tail', _1) FROM tbl",
+ "SELECT concat_ws('', _1, _2, array('x', NULL, 'y')) FROM tbl",
+ "SELECT concat_ws(NULL, _1, _2) FROM tbl",
+ "SELECT concat_ws(_2) FROM tbl")
+ withSQLConf(
+ CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false",
+ "spark.comet.expression.StringSplit.allowIncompatible" -> "true") {
for (query <- arrayArgQueries) {
- checkSparkAnswerAndFallbackReason(query, "`concat_ws` with `array<string>` arguments")
+ checkSparkAnswerAndOperator(query)
}
}
// A NULL separator produces NULL regardless of the argument types and stays native.
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometConcatWsBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometConcatWsBenchmark.scala
new file mode 100644
index 0000000000..6db28b3072
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometConcatWsBenchmark.scala
@@ -0,0 +1,109 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.benchmark
+
+import org.apache.spark.sql.execution.WholeStageCodegenExec
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.CometConf
+import org.apache.comet.udf.codegen.CometScalaUDFCodegen
+
+/**
+ * Matched Parquet inputs for native concat_ws versus Spark whole-stage codegen. Pass `strings` to
+ * run only the previously supported string cases against an older native library. `long-arrays`
+ * selects width 128 and lengths up to 32. `dispatch` requires the pre-native concat_ws serde on
+ * the classpath and verifies that the old default dispatcher actually ran.
+ */
+object CometConcatWsBenchmark extends CometBenchmarkBase {
+ override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+ val rows = 65536
+ val stringsOnly = mainArgs.contains("strings")
+ val longArraysOnly = mainArgs.contains("long-arrays")
+ val dispatch = mainArgs.contains("dispatch")
+ val cometConfigs = Map(
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString)
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> "8192",
+ CometConf.COMET_BATCH_SIZE.key -> "8192") {
+ for (width <- Seq(8, 128); arrayLength <- Seq(2, 8, 32)
+ if (!stringsOnly || arrayLength == 8) &&
+ (!longArraysOnly || (width == 128 && arrayLength == 32))) {
+ withTempPath { dir =>
+ withTempTable("concat_inputs") {
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .sql(s"""WITH strings AS (
+ | SELECT id, CASE WHEN id % 13 = 0 THEN NULL
+ | ELSE lpad(CAST(id AS STRING), $width, 'x') END AS c1
+ | FROM range($rows)), arrays AS (
+ | SELECT *, CASE WHEN id % 11 = 0 THEN NULL ELSE
+ | transform(sequence(1, CAST(id % $arrayLength + 1 AS INT)),
+ | j -> CASE WHEN j % 5 = 0 THEN NULL
+ | ELSE concat(c1, CAST(j AS STRING)) END) END AS a1
+ | FROM strings)
+ | SELECT c1, a1, reverse(a1) AS a2,
+ | CASE WHEN id % 17 = 0 THEN NULL WHEN id % 2 = 0 THEN '|'
+ | ELSE '--' END AS sep FROM arrays""".stripMargin)
+ .coalesce(1)
+ .write
+ .parquet(dir.getCanonicalPath)
+ }
+ spark.read.parquet(dir.getCanonicalPath).createOrReplaceTempView("concat_inputs")
+ val expressions =
+ (if (arrayLength == 8) Seq("strings" -> "c1, c1") else Seq.empty) ++
+ (if (stringsOnly) Seq.empty else Seq("mixed" -> "a1, c1, a2, 'tail'"))
+ for ((shape, args) <- expressions; separator <- Seq("' '", "sep")) {
+ val query = s"SELECT concat_ws($separator, $args) FROM concat_inputs"
+ val name = s"concat_ws $shape width=$width arrayLength=$arrayLength sep=$separator"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ val plan = spark.sql(query).queryExecution.executedPlan
+ require(plan.find(_.isInstanceOf[WholeStageCodegenExec]).nonEmpty)
+ require(plan.toString.contains("concat_ws"))
+ // scalastyle:off println
+ println(s"$name Spark plan:\n$plan")
+ // scalastyle:on println
+ }
+ withSQLConf(cometConfigs.toSeq: _*) {
+ CometScalaUDFCodegen.resetStats()
+ val df = spark.sql(query)
+ df.noop()
+ val stats = CometScalaUDFCodegen.stats()
+ require(
+ (stats.compileCount + stats.cacheHitCount > 0) == dispatch,
+ stats.toString)
+ val plan = df.queryExecution.executedPlan
+ require(findFirstNonCometOperator(plan).isEmpty, plan.toString)
+ // scalastyle:off println
+ println(s"$name Comet plan (dispatch=$dispatch, stats=$stats):\n$plan")
+ // scalastyle:on println
+ }
+ runBenchmark(name) {
+ runExpressionBenchmark(name, rows, query, cometConfigs)
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org