diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 70efdf8559..f128825f36 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -546,74 +546,74 @@ expression-level). The `outer` variants are wired but marked `Incompatible`; the ## string_funcs -| Function | Status | Implementation | Notes | -| --- | --- | --- | --- | -| `ascii` | ✅ | Native | | -| `base64` | ✅ | Native | | -| `bit_length` | ✅ | Native | | -| `btrim` | ✅ | — | | -| `char` | ✅ | Native | | -| `char_length` | ✅ | Native | | -| `character_length` | ✅ | Native | | -| `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` | ✅ | Native | | -| `contains` | ✅ | — | | -| `decode` | ✅ | — | | -| `elt` | ✅ | Codegen dispatch | | -| `encode` | 🔜 | — | Lowers to `StaticInvoke(encode)` (not allowlisted); falls back | -| `endswith` | ✅ | — | | -| `find_in_set` | ✅ | Codegen dispatch | | -| `format_number` | ✅ | Codegen dispatch | | -| `format_string` | ✅ | Codegen dispatch | | -| `initcap` | ✅ | Hybrid | | -| `instr` | ✅ | Native | | -| `lcase` | ✅ | Hybrid | | -| `left` | ✅ | Native | | -| `len` | ✅ | Native | | -| `length` | ✅ | Native | | -| `levenshtein` | ✅ | Native | | -| `locate` | ✅ | Codegen dispatch | | -| `lower` | ✅ | Hybrid | | -| `lpad` | ✅ | — | | -| `ltrim` | ✅ | Native | | -| `luhn_check` | ✅ | — | Native via `StaticInvoke` (tests: luhn_check.sql) | -| `mask` | ✅ | — | Routed through the JVM codegen dispatcher | -| `octet_length` | ✅ | Native | | -| `overlay` | ✅ | Codegen dispatch | | -| `position` | ✅ | Codegen dispatch | | -| `printf` | ✅ | Codegen dispatch | | -| `regexp_count` | ✅ | — | Runs natively (rewrites to `size(regexp_extract_all(...))`) | -| `regexp_extract` | ✅ | Native | | -| `regexp_extract_all` | ✅ | Native | | -| `regexp_instr` | ✅ | Codegen dispatch | Routed through the JVM codegen dispatcher | -| `regexp_replace` | ✅ | Hybrid | | -| `regexp_substr` | ✅ | — | Runs natively (rewrites to `nullif(regexp_extract(...), '')`) | -| `repeat` | ✅ | Native | | -| `replace` | ✅ | Hybrid | | -| `right` | ✅ | Native | | -| `rpad` | ✅ | — | | -| `rtrim` | ✅ | Native | | -| `soundex` | ✅ | Native | | -| `space` | ✅ | Native | | -| `split` | ✅ | Hybrid | | -| `split_part` | ✅ | — | Spark 4.0+ | -| `startswith` | ✅ | — | | -| `substr` | ✅ | Native | | -| `substring` | ✅ | Native | | -| `substring_index` | ✅ | Native | | -| `to_binary` | ✅ | — | Hex form accelerated; other formats fall back | -| `to_char` | ✅ | Codegen dispatch | | -| `to_number` | ✅ | Codegen dispatch | | -| `to_varchar` | ✅ | Codegen dispatch | | -| `translate` | ✅ | Native | DataFusion's `translate` iterates over Unicode graphemes (Spark uses code points) and substitutes U+0000 instead of treating it as a deletion sentinel, so the native path is opt-in via allowIncompatible | -| `trim` | ✅ | Native | | -| `try_to_binary` | ✅ | — | Runs natively (rewrites to `try_eval(to_binary(...))`) | -| `try_to_number` | ✅ | Codegen dispatch | Routed through the JVM codegen dispatcher | -| `ucase` | ✅ | Hybrid | | -| `unbase64` | ✅ | Codegen dispatch | | -| `upper` | ✅ | Hybrid | | +| Function | Status | Implementation | Notes | +| --- | --- |------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `ascii` | ✅ | Native | | +| `base64` | ✅ | Native | | +| `bit_length` | ✅ | Native | | +| `btrim` | ✅ | — | | +| `char` | ✅ | Native | | +| `char_length` | ✅ | Native | | +| `character_length` | ✅ | Native | | +| `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` | ✅ | Native | | +| `contains` | ✅ | — | | +| `decode` | ✅ | — | | +| `elt` | ✅ | Hybrid | Codegen dispatch when spark.sql.ansi.enabled=true | +| `encode` | 🔜 | — | Lowers to `StaticInvoke(encode)` (not allowlisted); falls back | +| `endswith` | ✅ | — | | +| `find_in_set` | ✅ | Codegen dispatch | | +| `format_number` | ✅ | Codegen dispatch | | +| `format_string` | ✅ | Codegen dispatch | | +| `initcap` | ✅ | Hybrid | | +| `instr` | ✅ | Native | | +| `lcase` | ✅ | Hybrid | | +| `left` | ✅ | Native | | +| `len` | ✅ | Native | | +| `length` | ✅ | Native | | +| `levenshtein` | ✅ | Native | | +| `locate` | ✅ | Codegen dispatch | | +| `lower` | ✅ | Hybrid | | +| `lpad` | ✅ | — | | +| `ltrim` | ✅ | Native | | +| `luhn_check` | ✅ | — | Native via `StaticInvoke` (tests: luhn_check.sql) | +| `mask` | ✅ | — | Routed through the JVM codegen dispatcher | +| `octet_length` | ✅ | Native | | +| `overlay` | ✅ | Codegen dispatch | | +| `position` | ✅ | Codegen dispatch | | +| `printf` | ✅ | Codegen dispatch | | +| `regexp_count` | ✅ | — | Runs natively (rewrites to `size(regexp_extract_all(...))`) | +| `regexp_extract` | ✅ | Native | | +| `regexp_extract_all` | ✅ | Native | | +| `regexp_instr` | ✅ | Codegen dispatch | Routed through the JVM codegen dispatcher | +| `regexp_replace` | ✅ | Hybrid | | +| `regexp_substr` | ✅ | — | Runs natively (rewrites to `nullif(regexp_extract(...), '')`) | +| `repeat` | ✅ | Native | | +| `replace` | ✅ | Hybrid | | +| `right` | ✅ | Native | | +| `rpad` | ✅ | — | | +| `rtrim` | ✅ | Native | | +| `soundex` | ✅ | Native | | +| `space` | ✅ | Native | | +| `split` | ✅ | Hybrid | | +| `split_part` | ✅ | — | Spark 4.0+ | +| `startswith` | ✅ | — | | +| `substr` | ✅ | Native | | +| `substring` | ✅ | Native | | +| `substring_index` | ✅ | Native | | +| `to_binary` | ✅ | — | Hex form accelerated; other formats fall back | +| `to_char` | ✅ | Codegen dispatch | | +| `to_number` | ✅ | Codegen dispatch | | +| `to_varchar` | ✅ | Codegen dispatch | | +| `translate` | ✅ | Native | DataFusion's `translate` iterates over Unicode graphemes (Spark uses code points) and substitutes U+0000 instead of treating it as a deletion sentinel, so the native path is opt-in via allowIncompatible | +| `trim` | ✅ | Native | | +| `try_to_binary` | ✅ | — | Runs natively (rewrites to `try_eval(to_binary(...))`) | +| `try_to_number` | ✅ | Codegen dispatch | Routed through the JVM codegen dispatcher | +| `ucase` | ✅ | Hybrid | | +| `unbase64` | ✅ | Codegen dispatch | | +| `upper` | ✅ | Hybrid | | --- diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 09e5cf6cd5..276b5a54ba 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -650,6 +650,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) { session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitShift::right_unsigned())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSoundex::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSubstring::default())); + session_ctx.register_udf(ScalarUDF::new_from_impl(SparkElt::default())); } /// Prepares arrow arrays for output. @@ -1223,6 +1224,7 @@ pub extern "system" fn Java_org_apache_comet_Native_getRustThreadId( use crate::execution::columnar_to_row::ColumnarToRowContext; use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; use datafusion_spark::function::math::bin::SparkBin; +use datafusion_spark::function::string::elt::SparkElt; use datafusion_spark::function::string::soundex::SparkSoundex; /// Initialize a native columnar to row converter. 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 7422dadd97..c78936d9ac 100644 --- a/spark/src/main/scala/org/apache/comet/serde/strings.scala +++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala @@ -20,7 +20,7 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, Base64, BitLength, Cast, Concat, ConcatWs, Contains, Elt, Empty2Null, EndsWith, Expression, FindInSet, FormatNumber, FormatString, GetJsonObject, InitCap, Left, Length, Levenshtein, Like, Literal, Lower, Mask, OctetLength, Overlay, RegExpExtract, RegExpExtractAll, RegExpInStr, RegExpReplace, Right, RLike, SoundEx, StartsWith, StringLocate, StringLPad, StringRepeat, StringReplace, StringRPad, StringSplit, StringTranslate, Substring, SubstringIndex, ToCharacter, ToNumber, TryToNumber, UnBase64, Upper} -import org.apache.spark.sql.types.{BinaryType, DataTypes, IntegerType, LongType, StringType} +import org.apache.spark.sql.types._ import org.apache.comet.CometConf import org.apache.comet.serde.ExprOuterClass.Expr @@ -679,9 +679,20 @@ object CometGetJsonObject extends CometCodegenDispatch[GetJsonObject] with Nativ } } +object CometElt extends CometScalarFunction[Elt]("elt") with CodegenDispatchFallback { + override def getIncompatibleReasons(): Seq[String] = Seq( + "The elt function does not respect spark.sql.ansi.enabled=true") + + override def getSupportLevel(expr: Elt): SupportLevel = { + if (expr.failOnError) { + return Incompatible(Some("The elt function does not respect spark.sql.ansi.enabled=true")) + } + Compatible() + } +} + // Expressions routed through the JVM codegen dispatcher: no native implementation, so Spark's own // doGenCode runs inside the Comet pipeline, matching Spark exactly. -object CometElt extends CometCodegenDispatch[Elt] object CometFindInSet extends CometCodegenDispatch[FindInSet] diff --git a/spark/src/test/resources/sql-tests/expressions/string/elt.sql b/spark/src/test/resources/sql-tests/expressions/string/elt.sql index fd526e2978..31f5720cff 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/elt.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/elt.sql @@ -15,8 +15,6 @@ -- specific language governing permissions and limitations -- under the License. --- Routes elt through the codegen dispatcher so behavior matches Spark exactly. - statement CREATE TABLE test_elt(n int) USING parquet @@ -24,8 +22,26 @@ statement INSERT INTO test_elt VALUES (1), (2), (3), (NULL) query -SELECT n, elt(n, 'a', 'b') FROM test_elt +SELECT elt(0, 'a', 'b', 'c'), elt(-1, 'a', 'b', 'c'), elt(4, 'a', 'b', 'c') + +query +SELECT elt(n + 1, 'x', 'y', 'z', 'w') FROM test_elt + +query +SELECT elt(1, 'a', NULL, 'c'), elt(2, 'a', NULL, 'c'), elt(3, 'a', NULL, 'c') + +statement +CREATE TABLE test_elt_edge (idx int, v1 string, v2 string, v3 string) USING parquet + +statement +INSERT INTO test_elt_edge VALUES + (1, 'foo', 'bar', 'baz'), + (2, 'foo', NULL, 'baz'), + (3, NULL, 'bar', 'baz'), + (4, 'foo', 'bar', 'baz'), + (NULL, 'foo', 'bar', 'baz'), + (0, 'foo', 'bar', 'baz'), + (-1, 'foo', 'bar', 'baz') --- literal arguments query -SELECT elt(1, 'scala', 'java'), elt(2, 'scala', 'java'), elt(2, 'a', 'b', 'c') +SELECT elt(idx, v1, v2, v3) FROM test_elt_edge \ No newline at end of file diff --git a/spark/src/test/resources/sql-tests/expressions/string/elt_ansi.sql b/spark/src/test/resources/sql-tests/expressions/string/elt_ansi.sql new file mode 100644 index 0000000000..09d3baa483 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/string/elt_ansi.sql @@ -0,0 +1,21 @@ +-- 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. + +-- Config: spark.sql.ansi.enabled=true + +query expect_error(The index 100 is out of bounds) +SELECT elt(100, 'a', 'b', 'c')