chore: support Spark 4.2

.github/agents/spark-test-expert.agent.md

@@ -13,7 +13,7 @@ Run the tests inside the Hatch environments defined in `pyproject.toml`.
 The following test environments are currently available:
 
 - `test-spark.spark-3.5.7`
-- `test-spark.spark-4.1.1`
+- `test-spark.spark-4.2.0`
 - `test-ibis`
 
 The `test-spark.spark-*` environments support running PySpark unit tests and doctests. The `test-ibis` environment supports running Ibis unit tests with PySpark as the backend.

.github/workflows/build.yml

@@ -108,7 +108,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        spark-version: ["3.5.7", "4.1.1"]
+        spark-version: ["3.5.7", "4.2.0"]
     with:
       spark_version: ${{ matrix.spark-version }}
     needs:
@@ -134,7 +134,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        spark-version: ["3.5.7", "4.0.1", "4.1.1"]
+        spark-version: ["3.5.7", "4.0.1", "4.1.1", "4.2.0"]
     with:
       spark_version: ${{ matrix.spark-version }}
     needs:

.github/workflows/catalog-tests.yml

@@ -11,7 +11,7 @@ jobs:
     name: Test
     runs-on: ubuntu-latest
     env:
-      HATCH_ENV: "test.spark-4.1.1"
+      HATCH_ENV: "test.spark-4.2.0"
       # Coverage environment variables
       RUSTC_WORKSPACE_WRAPPER: ${{ github.workspace }}/.github/scripts/rustc-workspace-wrapper.sh
       LLVM_PROFILE_FILE: ${{ github.workspace }}/target/sail-catalog-it-%p-%m.profraw

.github/workflows/gold-data-script-validation.yml

@@ -26,7 +26,7 @@ jobs:
         with:
           repository: apache/spark
           path: opt/spark
-          ref: v4.1.1
+          ref: v4.2.0
           fetch-depth: 1
 
       - uses: actions/setup-java@v5

.github/workflows/report.yml

@@ -35,8 +35,8 @@ jobs:
         continue-on-error: true
         with:
           workflow_run_id: ${{ github.event.workflow_run.id }}
-          artifact_name: spark-4.1.1-test-report
-          report_name: spark-4.1.1-test
+          artifact_name: spark-4.2.0-test-report
+          report_name: spark-4.2.0-test
 
       - uses: ./.github/actions/pull-request-report
         if: github.event.workflow_run.name == 'Build'

.github/workflows/spark-package-artifacts.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        spark-version: ["3.5.7", "4.1.1"]
+        spark-version: ["3.5.7", "4.2.0"]
     steps:
       - uses: actions/checkout@v7
 

crates/sail-execution/proto/sail/plan/physical.proto

@@ -470,6 +470,7 @@ message PySparkUdfConfig {
   bool python_udtf_pandas_conversion_enabled = 7;
   bool python_udf_pandas_int_to_decimal_coercion_enabled = 8;
   bool binary_as_bytes = 9;
+  bool python_udf_pandas_prefer_int_extension_dtype = 10;
 }
 
 message EquivalenceProperties {

crates/sail-execution/src/proto/codec.rs

@@ -4203,6 +4203,8 @@ impl RemoteExecutionCodec {
             python_udf_pandas_int_to_decimal_coercion_enabled: config
                 .python_udf_pandas_int_to_decimal_coercion_enabled,
             binary_as_bytes: config.binary_as_bytes,
+            python_udf_pandas_prefer_int_extension_dtype: config
+                .python_udf_pandas_prefer_int_extension_dtype,
         };
         Ok(config)
     }
@@ -4223,6 +4225,8 @@ impl RemoteExecutionCodec {
             python_udf_pandas_int_to_decimal_coercion_enabled: config
                 .python_udf_pandas_int_to_decimal_coercion_enabled,
             binary_as_bytes: config.binary_as_bytes,
+            python_udf_pandas_prefer_int_extension_dtype: config
+                .python_udf_pandas_prefer_int_extension_dtype,
         };
         Ok(config)
     }

crates/sail-plan/src/resolver/query/udtf.rs

@@ -181,6 +181,7 @@ impl PlanResolver<'_> {
             function.eval_type,
             arguments_len,
             &input_types,
+            passthrough_columns,
             kwargs,
             &return_type,
             &self.config.pyspark_udf_config,

crates/sail-python-udf/src/cereal/mod.rs

@@ -1,5 +1,6 @@
 use datafusion::arrow::datatypes::DataType;
 use pyo3::prelude::PyAnyMethods;
+use pyo3::sync::PyOnceLock;
 use pyo3::types::{PyList, PyModule};
 use pyo3::{PyResult, Python, intern};
 use sail_common::spec;
@@ -10,36 +11,34 @@ use crate::error::{PyUdfError, PyUdfResult};
 pub mod pyspark_udf;
 pub mod pyspark_udtf;
 
+#[derive(Clone, Copy, PartialEq, Eq)]
 enum PySparkVersion {
     V3,
     V4_0,
     V4_1,
-}
-
-impl PySparkVersion {
-    fn is_v4(&self) -> bool {
-        matches!(self, PySparkVersion::V4_0 | PySparkVersion::V4_1)
-    }
+    V4_2,
 }
 
 fn get_pyspark_version() -> PyUdfResult<PySparkVersion> {
-    use pyo3::prelude::PyAnyMethods;
-    use pyo3::types::PyModule;
+    static PYSPARK_VERSION: PyOnceLock<PySparkVersion> = PyOnceLock::new();
 
     Python::attach(|py| {
-        let module = PyModule::import(py, "pyspark")?;
-        let version: String = module.getattr("__version__")?.extract()?;
-        if version.starts_with("3.") {
-            Ok(PySparkVersion::V3)
-        } else if version.starts_with("4.0.") {
-            Ok(PySparkVersion::V4_0)
-        } else if version.starts_with("4.") {
-            Ok(PySparkVersion::V4_1)
-        } else {
-            Err(PyUdfError::invalid(format!(
-                "unsupported PySpark version: {version}"
-            )))
-        }
+        PYSPARK_VERSION
+            .get_or_try_init(py, || {
+                let module = PyModule::import(py, "pyspark")?;
+                let version: String = module.getattr("__version__")?.extract()?;
+                let parts: Vec<_> = version.split('.').collect();
+                match parts.as_slice() {
+                    ["3", ..] => Ok(PySparkVersion::V3),
+                    ["4", "0", ..] => Ok(PySparkVersion::V4_0),
+                    ["4", "1", ..] => Ok(PySparkVersion::V4_1),
+                    ["4", "2", ..] => Ok(PySparkVersion::V4_2),
+                    _ => Err(PyUdfError::invalid(format!(
+                        "unsupported PySpark version: {version}"
+                    ))),
+                }
+            })
+            .copied()
     })
 }
 
@@ -104,6 +103,16 @@ pub(crate) fn write_kwarg(data: &mut Vec<u8>, kwargs: &[Option<String>], index:
     }
 }
 
+pub(crate) fn write_conf(data: &mut Vec<u8>, conf: Vec<(String, String)>) {
+    data.extend((conf.len() as i32).to_be_bytes());
+    for (key, value) in conf {
+        data.extend((key.len() as i32).to_be_bytes());
+        data.extend(key.as_bytes());
+        data.extend((value.len() as i32).to_be_bytes());
+        data.extend(value.as_bytes());
+    }
+}
+
 fn should_write_config(eval_type: spec::PySparkUdfType) -> bool {
     use spec::PySparkUdfType;
 

crates/sail-python-udf/src/cereal/pyspark_udf.rs

@@ -7,7 +7,7 @@ use sail_common::spec;
 
 use crate::cereal::{
     PySparkVersion, build_input_types_json, check_python_udf_version, get_pyspark_version,
-    should_write_config, supports_kwargs, write_kwarg,
+    should_write_config, supports_kwargs, write_conf, write_kwarg,
 };
 use crate::config::PySparkUdfConfig;
 use crate::error::{PyUdfError, PyUdfResult};
@@ -29,9 +29,20 @@ impl PySparkUdfPayload {
         let serializer = PyModule::import(py, intern!(py, "pyspark.serializers"))?
             .getattr(intern!(py, "CPickleSerializer"))?
             .call0()?;
-        let tuple = PyModule::import(py, intern!(py, "pyspark.worker"))?
-            .getattr(intern!(py, "read_udfs"))?
-            .call1((serializer, infile, eval_type))?;
+        let worker = PyModule::import(py, intern!(py, "pyspark.worker"))?;
+        let read_udfs = worker.getattr(intern!(py, "read_udfs"))?;
+        let tuple = match get_pyspark_version()? {
+            PySparkVersion::V4_2 => {
+                let runner_conf = worker
+                    .getattr(intern!(py, "RunnerConf"))?
+                    .call1((&infile,))?;
+                let eval_conf = worker.getattr(intern!(py, "EvalConf"))?.call1((&infile,))?;
+                read_udfs.call1((serializer, infile, eval_type, runner_conf, eval_conf))?
+            }
+            PySparkVersion::V3 | PySparkVersion::V4_0 | PySparkVersion::V4_1 => {
+                read_udfs.call1((serializer, infile, eval_type))?
+            }
+        };
         tuple
             .get_item(0)?
             .into_pyobject(py)
@@ -54,29 +65,41 @@ impl PySparkUdfPayload {
 
         data.extend(i32::from(eval_type).to_be_bytes());
 
-        if should_write_config(eval_type) {
-            let config = config.to_key_value_pairs();
-            data.extend((config.len() as i32).to_be_bytes()); // number of configuration options
-            for (key, value) in config {
-                data.extend((key.len() as i32).to_be_bytes()); // length of the key
-                data.extend(key.as_bytes());
-                data.extend((value.len() as i32).to_be_bytes()); // length of the value
-                data.extend(value.as_bytes());
+        match pyspark_version {
+            PySparkVersion::V4_2 => {
+                // Spark 4.2 reads both maps before dispatching to read_udfs.
+                write_conf(&mut data, config.to_key_value_pairs());
+                let mut eval_conf = vec![];
+                if eval_type == spec::PySparkUdfType::ArrowBatched {
+                    eval_conf.push((
+                        "input_type".to_string(),
+                        build_input_types_json(input_types)?,
+                    ))
+                }
+                write_conf(&mut data, eval_conf);
             }
-        }
+            PySparkVersion::V3 | PySparkVersion::V4_0 | PySparkVersion::V4_1 => {
+                if should_write_config(eval_type) {
+                    write_conf(&mut data, config.to_key_value_pairs());
+                }
 
-        // PySpark 4.1+ reads input types for ArrowBatched UDFs.
-        // PySpark 4.0.x does not read input types and would misparse the stream.
-        if matches!(pyspark_version, PySparkVersion::V4_1)
-            && matches!(eval_type, spec::PySparkUdfType::ArrowBatched)
-        {
-            let schema_json = build_input_types_json(input_types)?;
-            data.extend((schema_json.len() as i32).to_be_bytes());
-            data.extend(schema_json.as_bytes());
+                // PySpark 4.1 reads input types for ArrowBatched UDFs.
+                // PySpark 4.0.x does not read input types and would misparse the stream.
+                if pyspark_version == PySparkVersion::V4_1
+                    && eval_type == spec::PySparkUdfType::ArrowBatched
+                {
+                    let schema_json = build_input_types_json(input_types)?;
+                    data.extend((schema_json.len() as i32).to_be_bytes());
+                    data.extend(schema_json.as_bytes());
+                }
+            }
         }
 
-        if pyspark_version.is_v4() {
-            data.extend(0u8.to_be_bytes()); // profiling is not enabled
+        match pyspark_version {
+            PySparkVersion::V4_0 | PySparkVersion::V4_1 => {
+                data.extend(0u8.to_be_bytes()); // profiling is not enabled
+            }
+            PySparkVersion::V3 | PySparkVersion::V4_2 => {}
         }
 
         data.extend(1i32.to_be_bytes()); // number of UDFs
@@ -87,7 +110,11 @@ impl PySparkUdfPayload {
             .map_err(|e| PyUdfError::invalid(format!("num args: {e}")))?;
         data.extend(num_arg_offsets.to_be_bytes()); // number of argument offsets
 
-        let allow_kwargs = pyspark_version.is_v4() && supports_kwargs(eval_type);
+        let allow_kwargs = match pyspark_version {
+            PySparkVersion::V4_2 => true,
+            PySparkVersion::V4_0 | PySparkVersion::V4_1 => supports_kwargs(eval_type),
+            PySparkVersion::V3 => false,
+        };
 
         for (i, offset) in arg_offsets.iter().enumerate() {
             let offset: i32 = (*offset)
@@ -103,6 +130,14 @@ impl PySparkUdfPayload {
         data.extend((command.len() as i32).to_be_bytes()); // length of the function
         data.extend_from_slice(command);
 
+        match pyspark_version {
+            PySparkVersion::V4_2 => {
+                // Spark 4.2 always sends this field, even with profiling disabled.
+                data.extend(0i64.to_be_bytes());
+            }
+            PySparkVersion::V3 | PySparkVersion::V4_0 | PySparkVersion::V4_1 => {}
+        }
+
         Ok(data)
     }
 }