feat: fix array_compact for Spark 4.0 and correct return type metadata (#3796)

Author: andygrove

native/spark-expr/src/array_funcs/array_compact.rs

@@ -132,6 +132,11 @@ fn compact_list<OffsetSize: OffsetSizeTrait>(
     );
     let mut valid = NullBufferBuilder::new(list_array.len());
 
+    // Use logical_nulls() instead of is_null() to correctly handle NullArray.
+    // NullArray::nulls() returns None (which makes is_null() return false),
+    // but logical_nulls() correctly reports all elements as null.
+    let value_nulls = values.logical_nulls();
+
     for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() {
         if list_array.is_null(row_index) {
             offsets.push(offsets[row_index]);
@@ -144,7 +149,8 @@ fn compact_list<OffsetSize: OffsetSizeTrait>(
         let mut copied = 0usize;
 
         for i in start..end {
-            if !values.is_null(i) {
+            let is_null = value_nulls.as_ref().map(|n| n.is_null(i)).unwrap_or(false);
+            if !is_null {
                 mutable.extend(0, i, i + 1);
                 copied += 1;
             }

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

@@ -295,6 +295,7 @@ object CometArrayRepeat extends CometExpressionSerde[ArrayRepeat] {
 }
 
 object CometArrayCompact extends CometExpressionSerde[Expression] {
+
   override def convert(
       expr: Expression,
       inputs: Seq[Attribute],

spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala

@@ -24,13 +24,13 @@ import org.apache.spark.sql.catalyst.expressions.json.StructsToJsonEvaluator
 import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.internal.types.StringTypeWithCollation
-import org.apache.spark.sql.types.{BinaryType, BooleanType, DataTypes, StringType}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, DataTypes, StringType}
 
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.expressions.{CometCast, CometEvalMode}
 import org.apache.comet.serde.{CommonStringExprs, Compatible, ExprOuterClass, Incompatible}
 import org.apache.comet.serde.ExprOuterClass.{BinaryOutputStyle, Expr}
-import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto}
+import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType}
 
 /**
  * `CometExprShim` acts as a shim for parsing expressions from different Spark versions.
@@ -56,6 +56,28 @@ trait CometExprShim extends CommonStringExprs {
       inputs: Seq[Attribute],
       binding: Boolean): Option[Expr] = {
     expr match {
+      case knc: KnownNotContainsNull =>
+        // On Spark 4.0, array_compact rewrites to KnownNotContainsNull(ArrayFilter(IsNotNull)).
+        // Strip the wrapper and serialize the inner ArrayFilter as spark_array_compact.
+        knc.child match {
+          case filter: ArrayFilter =>
+            filter.function.children.headOption match {
+              case Some(_: IsNotNull) =>
+                val arrayChild = filter.left
+                val elementType = arrayChild.dataType.asInstanceOf[ArrayType].elementType
+                val arrayExprProto = exprToProtoInternal(arrayChild, inputs, binding)
+                val returnType = ArrayType(elementType)
+                val scalarExpr = scalarFunctionExprToProtoWithReturnType(
+                  "spark_array_compact",
+                  returnType,
+                  false,
+                  arrayExprProto)
+                optExprWithInfo(scalarExpr, knc, arrayChild)
+              case _ => exprToProtoInternal(knc.child, inputs, binding)
+            }
+          case _ => exprToProtoInternal(knc.child, inputs, binding)
+        }
+
       case s: StaticInvoke
           if s.staticObject == classOf[StringDecode] &&
             s.dataType.isInstanceOf[StringType] &&
@@ -109,12 +131,6 @@ trait CometExprShim extends CommonStringExprs {
         val optExpr = scalarFunctionExprToProto("width_bucket", childExprs: _*)
         optExprWithInfo(optExpr, wb, wb.children: _*)
 
-      // KnownNotContainsNull is a TaggingExpression added in Spark 4.0 that only
-      // changes schema metadata (containsNull = false). It has no runtime effect,
-      // so we pass through to the child expression.
-      case k: KnownNotContainsNull =>
-        exprToProtoInternal(k.child, inputs, binding)
-
       // In Spark 4.0, StructsToJson is a RuntimeReplaceable whose replacement is
       // Invoke(Literal(StructsToJsonEvaluator), "evaluate", ...). Reconstruct the
       // original StructsToJson and recurse so support-level checks apply.

spark/src/test/resources/sql-tests/expressions/array/array_compact.sql

@@ -17,28 +17,48 @@
 
 
 statement
-CREATE TABLE test_array_compact(arr array<int>) USING parquet
+CREATE TABLE test_array_compact(
+  ints array<int>,
+  strs array<string>,
+  dbls array<double>,
+  nested array<array<int>>
+) USING parquet
 
 statement
-INSERT INTO test_array_compact VALUES (array(1, NULL, 2, NULL, 3)), (array()), (NULL), (array(NULL, NULL)), (array(1, 2, 3))
+INSERT INTO test_array_compact VALUES
+  (array(1, NULL, 2, NULL, 3), array('a', NULL, 'b', NULL, 'c'), array(1.0, NULL, 2.0), array(array(1, NULL, 3), NULL, array(4, NULL, 6))),
+  (array(), array(), array(), array()),
+  (NULL, NULL, NULL, NULL),
+  (array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL)),
+  (array(1, 2, 3), array('x', 'y', 'z'), array(1.5, 2.5), array(array(1, 2), array(3, 4)))
 
--- column argument
+-- integer column
 query
-SELECT array_compact(arr) FROM test_array_compact
+SELECT array_compact(ints) FROM test_array_compact
+
+-- string column
+query
+SELECT array_compact(strs) FROM test_array_compact
+
+-- double column
+query
+SELECT array_compact(dbls) FROM test_array_compact
+
+-- nested array column: outer nulls removed, inner nulls preserved
+query
+SELECT array_compact(nested) FROM test_array_compact
 
 -- literal arguments
 query
 SELECT array_compact(array(1, NULL, 2, NULL, 3))
 
--- string element type
-statement
-CREATE TABLE test_array_compact_str(arr array<string>) USING parquet
-
-statement
-INSERT INTO test_array_compact_str VALUES (array('a', NULL, 'b', NULL, 'c')), (array()), (NULL), (array(NULL, NULL)), (array('', NULL, '', NULL))
+-- literal string array
+query
+SELECT array_compact(array('a', NULL, 'b'))
 
+-- all-null literal array
 query
-SELECT array_compact(arr) FROM test_array_compact_str
+SELECT array_compact(array(NULL, NULL, NULL))
 
 -- double element type
 query