test: Enable Spark 4.0 tests

.github/workflows/spark_sql_test_ansi.yml

@@ -22,17 +22,15 @@ concurrency:
   cancel-in-progress: true
 
 on:
-  # enable the following once Ansi support is completed
-  #  push:
-  #    paths-ignore:
-  #      - "doc/**"
-  #      - "**.md"
-  #  pull_request:
-  #    paths-ignore:
-  #      - "doc/**"
-  #      - "**.md"
-
-  # manual trigger ONLY
+  push:
+    paths-ignore:
+      - "docs/**"
+      - "**.md"
+  pull_request:
+    paths-ignore:
+      - "docs/**"
+      - "**.md"
+  # manual trigger
   # https://docs.github.com/en/actions/managing-workflow-runs/manually-running-a-workflow
   workflow_dispatch:
 
@@ -44,8 +42,8 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        java-version: [11]
-        spark-version: [{short: '3.4', full: '3.4.2'}]
+        java-version: [17]
+        spark-version: [{short: '4.0', full: '4.0.0-preview1'}]
         module:
           - {name: "catalyst", args1: "catalyst/test", args2: ""}
           - {name: "sql/core-1", args1: "", args2: sql/testOnly * -- -l org.apache.spark.tags.ExtendedSQLTest -l org.apache.spark.tags.SlowSQLTest}
@@ -75,7 +73,8 @@ jobs:
       - name: Run Spark tests
         run: |
           cd apache-spark
-          ENABLE_COMET=true ENABLE_COMET_ANSI_MODE=true build/sbt ${{ matrix.module.args1 }} "${{ matrix.module.args2 }}"
+          rm -rf /root/.m2/repository/org/apache/parquet # somehow parquet cache requires cleanups
+          RUST_BACKTRACE=1 ENABLE_COMET=true ENABLE_COMET_ANSI_MODE=true build/sbt ${{ matrix.module.args1 }} "${{ matrix.module.args2 }}"
         env:
           LC_ALL: "C.UTF-8"
 

common/src/main/java/org/apache/comet/parquet/BatchReader.java

@@ -285,6 +285,7 @@ public void init() throws URISyntaxException, IOException {
     missingColumns = new boolean[columns.size()];
     List<String[]> paths = requestedSchema.getPaths();
     StructField[] nonPartitionFields = sparkSchema.fields();
+    ShimFileFormat.findRowIndexColumnIndexInSchema(sparkSchema);
     for (int i = 0; i < requestedSchema.getFieldCount(); i++) {
       Type t = requestedSchema.getFields().get(i);
       Preconditions.checkState(

common/src/main/java/org/apache/comet/parquet/TypeUtil.java

@@ -27,6 +27,7 @@
 import org.apache.parquet.schema.PrimitiveType;
 import org.apache.parquet.schema.Type;
 import org.apache.parquet.schema.Types;
+import org.apache.spark.package$;
 import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException;
 import org.apache.spark.sql.types.*;
 
@@ -169,6 +170,7 @@ && isUnsignedIntTypeMatched(logicalTypeAnnotation, 64)) {
         break;
       case INT96:
         if (sparkType == TimestampNTZType$.MODULE$) {
+          if (isSpark40Plus()) return; // Spark 4.0+ supports Timestamp NTZ with INT96
           convertErrorForTimestampNTZ(typeName.name());
         } else if (sparkType == DataTypes.TimestampType) {
           return;
@@ -218,7 +220,8 @@ private static void validateTimestampType(
     // Throw an exception if the Parquet type is TimestampLTZ and the Catalyst type is TimestampNTZ.
     // This is to avoid mistakes in reading the timestamp values.
     if (((TimestampLogicalTypeAnnotation) logicalTypeAnnotation).isAdjustedToUTC()
-        && sparkType == TimestampNTZType$.MODULE$) {
+        && sparkType == TimestampNTZType$.MODULE$
+        && !isSpark40Plus()) {
       convertErrorForTimestampNTZ("int64 time(" + logicalTypeAnnotation + ")");
     }
   }
@@ -278,4 +281,8 @@ private static boolean isUnsignedIntTypeMatched(
         && !((IntLogicalTypeAnnotation) logicalTypeAnnotation).isSigned()
         && ((IntLogicalTypeAnnotation) logicalTypeAnnotation).getBitWidth() == bitWidth;
   }
+
+  private static boolean isSpark40Plus() {
+    return package$.MODULE$.SPARK_VERSION().compareTo("4.0") >= 0;
+  }
 }

common/src/main/spark-3.x/org/apache/comet/shims/ShimFileFormat.scala

@@ -19,6 +19,8 @@
 
 package org.apache.comet.shims
 
+import org.apache.spark.sql.types.{LongType, StructField, StructType}
+
 object ShimFileFormat {
 
   // TODO: remove after dropping Spark 3.2 & 3.3 support and directly use FileFormat.ROW_INDEX
@@ -32,4 +34,20 @@ object ShimFileFormat {
 
   // TODO: remove after dropping Spark 3.2 support and use FileFormat.OPTION_RETURNING_BATCH
   val OPTION_RETURNING_BATCH = "returning_batch"
+
+  // TODO: remove after dropping Spark 3.2 & 3.3 support and directly use
+  //       RowIndexUtil.findRowIndexColumnIndexInSchema
+  def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = {
+    sparkSchema.fields.zipWithIndex.find { case (field: StructField, _: Int) =>
+      field.name == ShimFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME
+    } match {
+      case Some((field: StructField, idx: Int)) =>
+        if (field.dataType != LongType) {
+          throw new RuntimeException(
+            s"${ShimFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME} must be of LongType")
+        }
+        idx
+      case _ => -1
+    }
+  }
 }

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

@@ -21,11 +21,16 @@ package org.apache.comet.shims
 
 import org.apache.spark.sql.execution.datasources.FileFormat
 import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
+import org.apache.spark.sql.execution.datasources.parquet.ParquetRowIndexUtil
+import org.apache.spark.sql.types.StructType
 
 object ShimFileFormat {
   // A name for a temporary column that holds row indexes computed by the file format reader
   // until they can be placed in the _metadata struct.
   val ROW_INDEX_TEMPORARY_COLUMN_NAME = ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME
 
   val OPTION_RETURNING_BATCH = FileFormat.OPTION_RETURNING_BATCH
+
+  def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int =
+    ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema)
 }

core/src/errors.rs

@@ -21,6 +21,7 @@ use arrow::error::ArrowError;
 use datafusion_common::DataFusionError;
 use jni::errors::{Exception, ToException};
 use regex::Regex;
+
 use std::{
     any::Any,
     convert,
@@ -37,6 +38,7 @@ use std::{
 use jni::sys::{jboolean, jbyte, jchar, jdouble, jfloat, jint, jlong, jobject, jshort};
 
 use crate::execution::operators::ExecutionError;
+use jni::objects::{GlobalRef, JThrowable};
 use jni::JNIEnv;
 use lazy_static::lazy_static;
 use parquet::errors::ParquetError;
@@ -159,7 +161,11 @@ pub enum CometError {
     },
 
     #[error("{class}: {msg}")]
-    JavaException { class: String, msg: String },
+    JavaException {
+        class: String,
+        msg: String,
+        throwable: GlobalRef,
+    },
 }
 
 pub fn init() {
@@ -198,6 +204,15 @@ impl From<CometError> for ExecutionError {
     fn from(value: CometError) -> Self {
         match value {
             CometError::Execution { source } => source,
+            CometError::JavaException {
+                class,
+                msg,
+                throwable,
+            } => ExecutionError::JavaException {
+                class,
+                msg,
+                throwable,
+            },
             _ => ExecutionError::GeneralError(value.to_string()),
         }
     }
@@ -369,17 +384,34 @@ pub fn unwrap_or_throw_default<T: JNIDefault>(
     }
 }
 
-fn throw_exception<E: ToException>(env: &mut JNIEnv, error: &E, backtrace: Option<String>) {
+fn throw_exception(env: &mut JNIEnv, error: &CometError, backtrace: Option<String>) {
     // If there isn't already an exception?
     if env.exception_check().is_ok() {
         // ... then throw new exception
-        let exception = error.to_exception();
-        match backtrace {
-            Some(backtrace_string) => env.throw_new(
-                exception.class,
-                to_stacktrace_string(exception.msg, backtrace_string).unwrap(),
-            ),
-            _ => env.throw_new(exception.class, exception.msg),
+        match error {
+            CometError::JavaException {
+                class: _,
+                msg: _,
+                throwable,
+            } => env.throw(<&JThrowable>::from(throwable.as_obj())),
+            CometError::Execution {
+                source:
+                    ExecutionError::JavaException {
+                        class: _,
+                        msg: _,
+                        throwable,
+                    },
+            } => env.throw(<&JThrowable>::from(throwable.as_obj())),
+            _ => {
+                let exception = error.to_exception();
+                match backtrace {
+                    Some(backtrace_string) => env.throw_new(
+                        exception.class,
+                        to_stacktrace_string(exception.msg, backtrace_string).unwrap(),
+                    ),
+                    _ => env.throw_new(exception.class, exception.msg),
+                }
+            }
         }
         .expect("Thrown exception")
     }

core/src/execution/operators/mod.rs

@@ -25,6 +25,7 @@ use arrow::{
 
 use arrow::compute::{cast_with_options, CastOptions};
 use arrow_schema::ArrowError;
+use jni::objects::GlobalRef;
 use std::{fmt::Debug, sync::Arc};
 
 mod scan;
@@ -52,6 +53,13 @@ pub enum ExecutionError {
     /// DataFusion error
     #[error("Error from DataFusion {0}.")]
     DataFusionError(String),
+
+    #[error("{class}: {msg}")]
+    JavaException {
+        class: String,
+        msg: String,
+        throwable: GlobalRef,
+    },
 }
 
 /// Copy an Arrow Array

core/src/jvm_bridge/mod.rs

@@ -385,5 +385,6 @@ pub(crate) fn convert_exception(
     Ok(CometError::JavaException {
         class: exception_class_name_str,
         msg: message_str,
+        throwable: env.new_global_ref(throwable)?,
     })
 }