Feat: impl elt function

Author: Kazantsev Maksim

native/core/src/execution/jni_api.rs

@@ -51,6 +51,7 @@ use datafusion_spark::function::math::expm1::SparkExpm1;
 use datafusion_spark::function::math::hex::SparkHex;
 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;
@@ -351,6 +352,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) {
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkConcat::default()));
     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(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,7 +21,7 @@ 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.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.{BinaryType, DataTypes, LongType, StringType}
 
 import org.apache.comet.CometConf
@@ -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("failOnError=true is not supported"))
+    }
+    Compatible(None)
+  }
+}
+
 trait CommonStringExprs {
 
   def stringDecode(

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

@@ -22,9 +22,10 @@ package org.apache.comet
 import scala.util.Random
 
 import org.apache.parquet.hadoop.ParquetOutputFormat
-import org.apache.spark.sql.{CometTestBase, DataFrame}
+import org.apache.spark.sql.{functions, 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,38 @@ class CometStringExpressionSuite extends CometTestBase {
     }
   }
 
+  test("elt") {
+    withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") {
+      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(s"SELECT elt(1) FROM t1")) match {
+        case (Some(spark), Some(comet)) =>
+          assert(spark.getMessage.contains("WRONG_NUM_ARGS.WITHOUT_SUGGESTION"))
+          assert(comet.getMessage.contains("WRONG_NUM_ARGS.WITHOUT_SUGGESTION"))
+        case (spark, comet) =>
+          fail(
+            s"Expected Spark and Comet to throw exception, but got\nSpark: $spark\nComet: $comet")
+      }
+    }
+  }
+
 }