Feat: add support for elt expression

native/core/src/execution/jni_api.rs

@@ -52,6 +52,7 @@ use datafusion_spark::function::math::hex::SparkHex;
 use datafusion_spark::function::math::width_bucket::SparkWidthBucket;
 use datafusion_spark::function::string::char::CharFunc;
 use datafusion_spark::function::string::concat::SparkConcat;
+use datafusion_spark::function::string::elt::SparkElt;
 use futures::poll;
 use futures::stream::StreamExt;
 use jni::objects::JByteBuffer;
@@ -353,6 +354,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) {
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitwiseNot::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkHex::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkWidthBucket::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkElt::default()));
 }
 
 /// Prepares arrow arrays for output.

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -172,7 +172,8 @@ object QueryPlanSerde extends Logging with CometExprShim {
     classOf[StringTrimRight] -> CometScalarFunction("rtrim"),
     classOf[Left] -> CometLeft,
     classOf[Substring] -> CometSubstring,
-    classOf[Upper] -> CometUpper)
+    classOf[Upper] -> CometUpper,
+    classOf[Elt] -> CometElt)
 
   private val bitwiseExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
     classOf[BitwiseAnd] -> CometBitwiseAnd,

spark/src/main/scala/org/apache/comet/serde/strings.scala

@@ -21,8 +21,8 @@ package org.apache.comet.serde
 
 import java.util.Locale
 
-import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Concat, Expression, InitCap, Left, Length, Like, Literal, Lower, RegExpReplace, RLike, StringLPad, StringRepeat, StringRPad, Substring, Upper}
-import org.apache.spark.sql.types.{BinaryType, DataTypes, LongType, StringType}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Concat, Elt, Expression, InitCap, Left, Length, Like, Literal, Lower, RegExpReplace, RLike, StringLPad, StringRepeat, StringRPad, Substring, Upper}
+import org.apache.spark.sql.types._
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.withInfo
@@ -289,6 +289,16 @@ object CometRegExpReplace extends CometExpressionSerde[RegExpReplace] {
   }
 }
 
+object CometElt extends CometScalarFunction[Elt]("elt") {
+
+  override def getSupportLevel(expr: Elt): SupportLevel = {
+    if (expr.failOnError) {
+      return Unsupported(Some("ANSI mode not supported"))
+    }
+    Compatible(None)
+  }
+}
+
 trait CommonStringExprs {
 
   def stringDecode(

spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala

@@ -23,8 +23,9 @@ import scala.util.Random
 
 import org.apache.parquet.hadoop.ParquetOutputFormat
 import org.apache.spark.sql.{CometTestBase, DataFrame}
+import org.apache.spark.sql.functions.lit
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
+import org.apache.spark.sql.types.{DataTypes, StringType, StructField, StructType}
 
 import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator}
 
@@ -391,4 +392,43 @@ class CometStringExpressionSuite extends CometTestBase {
     }
   }
 
+  test("elt") {
+    val wrongNumArgsWithoutSuggestionExceptionMsg =
+      "[WRONG_NUM_ARGS.WITHOUT_SUGGESTION] The `elt` requires > 1 parameters but the actual number is 1."
+    withSQLConf(
+      SQLConf.ANSI_ENABLED.key -> "false",
+      SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+        "org.apache.spark.sql.catalyst.optimizer.ConstantFolding") {
+      val r = new Random(42)
+      val fieldsCount = 10
+      val indexes = Seq.range(1, fieldsCount)
+      val edgeCasesIndexes = Seq(-1, 0, -100, fieldsCount + 100)
+      val schema = indexes
+        .foldLeft(new StructType())((schema, idx) =>
+          schema.add(s"c$idx", StringType, nullable = true))
+      val df = FuzzDataGenerator.generateDataFrame(
+        r,
+        spark,
+        schema,
+        100,
+        DataGenOptions(maxStringLength = 6))
+      df.withColumn(
+        "idx",
+        lit(Random.shuffle(indexes ++ edgeCasesIndexes).headOption.getOrElse(-1)))
+        .createOrReplaceTempView("t1")
+      checkSparkAnswerAndOperator(
+        sql(s"SELECT elt(idx, ${schema.fieldNames.mkString(",")}) FROM t1"))
+      checkSparkAnswerAndOperator(
+        sql(s"SELECT elt(cast(null as int), ${schema.fieldNames.mkString(",")}) FROM t1"))
+      checkSparkAnswerMaybeThrows(sql("SELECT elt(1) FROM t1")) match {
+        case (Some(spark), Some(comet)) =>
+          assert(spark.getMessage.contains(wrongNumArgsWithoutSuggestionExceptionMsg))
+          assert(comet.getMessage.contains(wrongNumArgsWithoutSuggestionExceptionMsg))
+        case (spark, comet) =>
+          fail(
+            s"Expected Spark and Comet to throw exception, but got\nSpark: $spark\nComet: $comet")
+      }
+      checkSparkAnswerAndOperator("SELECT elt(2, 'a', 'b', 'c')")
+    }
+  }
 }

spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala

@@ -76,7 +76,8 @@ object CometStringExpressionBenchmark extends CometBenchmarkBase {
     StringExprConfig("substring", "select substring(c1, 1, 100) from parquetV1Table"),
     StringExprConfig("translate", "select translate(c1, '123456', 'aBcDeF') from parquetV1Table"),
     StringExprConfig("trim", "select trim(c1) from parquetV1Table"),
-    StringExprConfig("upper", "select upper(c1) from parquetV1Table"))
+    StringExprConfig("upper", "select upper(c1) from parquetV1Table"),
+    StringExprConfig("elt", "select elt(2, c1, c1) from parquetV1Table"))
 
   override def runCometBenchmark(mainArgs: Array[String]): Unit = {
     runBenchmarkWithTable("String expressions", 1024) { v =>