[SPARK-60698][CONNECT] Batch Arrow input normalization config reads

crates/pyspark-rs/src/lib.rs

@@ -38,6 +38,7 @@ mod observation;
 mod pipelines;
 mod profiler;
 mod pyarrow_export;
+mod pyarrow_import;
 mod readwriter;
 mod resource;
 mod row;

crates/pyspark-rs/src/pyarrow_import.rs

@@ -0,0 +1,94 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! PyArrow Table input through arrow-rs's owned C Stream importer.
+
+use arrow::pyarrow::{FromPyArrow, Table};
+use pyo3::prelude::*;
+use spark_connect::session::SparkSession;
+
+use crate::dataframe::PyDataFrame;
+use crate::errors::ResultExt;
+
+pub(crate) fn create_dataframe(
+    py: Python<'_>,
+    session: &SparkSession,
+    table: &Bound<'_, PyAny>,
+    schema: Option<&Bound<'_, PyAny>>,
+) -> PyResult<PyDataFrame> {
+    let helper = py.import("pyspark.sql._arrow_input")?;
+    // Resolve DDL with the same converter as the existing local-data path.
+    let schema = match schema {
+        Some(s) if s.extract::<String>().is_ok() => Some(crate::types::data_type_to_py(
+            py,
+            &crate::types::py_to_data_type(s)?,
+        )?),
+        Some(s) => Some(s.clone().unbind()),
+        None => None,
+    };
+    let schema = schema.unwrap_or_else(|| py.None());
+    // Empty input needs only the explicit schema, and no server configuration.
+    let empty: bool = table.call_method0("__len__")?.extract::<usize>()? == 0;
+    let (timestamp_type, timezone, large_types) = if empty {
+        (None, None, None)
+    } else {
+        py.detach(|| -> PyResult<_> {
+            let values = session
+                .conf()
+                .get_many(&[
+                    "spark.sql.timestampType",
+                    "spark.sql.session.timeZone",
+                    "spark.sql.execution.arrow.useLargeVarTypes",
+                ])
+                .to_pyerr()?;
+            let mut values = values.into_iter();
+            Ok((
+                values.next().flatten(),
+                values.next().flatten(),
+                values.next().flatten(),
+            ))
+        })?
+    };
+    let result = helper.getattr("prepare_table")?.call1((
+        table,
+        schema,
+        timestamp_type,
+        timezone,
+        large_types,
+    ))?;
+    let normalized = result.get_item(0)?;
+    let spark_schema = crate::types::py_to_data_type(&result.get_item(1)?)?;
+    let names: Option<Vec<String>> = result.get_item(2)?.extract()?;
+    let batches = if empty {
+        Vec::new()
+    } else {
+        // PyArrow 10 readers support _export_to_c, even without PyCapsules.
+        Table::from_pyarrow_bound(&normalized.call_method0("to_reader")?)?
+            .into_inner()
+            .0
+    };
+    let dataframe = py.detach(|| {
+        session
+            .create_dataframe_from_arrow(batches, spark_schema)
+            .to_pyerr()
+    })?;
+    let dataframe = match names {
+        Some(names) => dataframe.to_df(names.iter().map(String::as_str).collect()),
+        None => dataframe,
+    };
+    Ok(PyDataFrame::new(dataframe))
+}

crates/pyspark-rs/src/session.rs

@@ -348,6 +348,12 @@ impl PySparkSession {
         // signature parity. See PARITY_GAP_ANALYSIS.md.
         let _ = (samplingRatio, verifySchema);
 
+        if let Ok(pa) = py.import("pyarrow") {
+            if data.is_instance(&pa.getattr("Table")?)? {
+                return crate::pyarrow_import::create_dataframe(py, &self.session, data, schema);
+            }
+        }
+
         // Accept a pandas DataFrame (as pyspark does): convert it to a list of row lists
         // with NaN/NaT mapped to None, and take its column names as the default schema.
         // pandas-on-Spark always passes an explicit `schema`, so this mainly needs to

crates/spark-connect-core/Cargo.toml

@@ -26,6 +26,7 @@ serde_json = "1"
 regex = "1"
 prost-types.workspace = true
 crc32fast = "1"
+sha2 = "0.10"
 
 [dev-dependencies]
 tempfile = "3"

crates/spark-connect-core/src/client.rs

@@ -800,6 +800,73 @@ impl SparkConnectClient {
         Ok(())
     }
 
+    /// Upload session-scoped cache blobs, returning SHA-256 hashes in input order.
+    /// Existing artifacts and duplicate blobs are uploaded only once.
+    pub async fn cache_artifacts(&self, blobs: Vec<Vec<u8>>) -> Result<Vec<String>> {
+        use crate::artifact::InMemoryArtifact;
+        use sha2::{Digest, Sha256};
+        use std::collections::HashSet;
+
+        let hashes: Vec<String> = blobs
+            .iter()
+            .map(|blob| format!("{:x}", Sha256::digest(blob)))
+            .collect();
+        let names: Vec<String> = hashes.iter().map(|hash| format!("cache/{hash}")).collect();
+        let unique: HashSet<&str> = names.iter().map(String::as_str).collect();
+        if unique.is_empty() {
+            return Ok(hashes);
+        }
+        let cached = self
+            .artifact_status(&unique.into_iter().collect::<Vec<_>>())
+            .await?;
+        let mut seen = HashSet::new();
+        let mut artifacts = Vec::new();
+        for (blob, name) in blobs.into_iter().zip(names) {
+            if !cached.get(&name).copied().unwrap_or(false) && seen.insert(name.clone()) {
+                artifacts.push(NamedArtifact::new(
+                    name,
+                    Box::new(InMemoryArtifact::new(blob)),
+                ));
+            }
+        }
+        if artifacts.is_empty() {
+            return Ok(hashes);
+        }
+        let expected: HashSet<String> = artifacts.iter().map(|a| a.name.clone()).collect();
+        let mut requests = build_artifact_request_stream(self.session_id.clone(), artifacts)?;
+        for request in &mut requests {
+            request.user_context = Some(UserContext {
+                user_id: self.user_id.clone().unwrap_or_default(),
+                ..Default::default()
+            });
+        }
+        let mut req = Request::new(futures::stream::iter(requests));
+        self._attach_metadata(&mut req);
+        let response = self
+            .stub
+            .clone()
+            .add_artifacts(req)
+            .await
+            .map_err(SparkError::from_grpc_status)?
+            .into_inner();
+        let mut received = HashSet::new();
+        for summary in response.artifacts {
+            if !summary.is_crc_successful {
+                return Err(SparkError::connect_msg(format!(
+                    "CRC check failed for artifact: {}",
+                    summary.name
+                )));
+            }
+            received.insert(summary.name);
+        }
+        if !expected.is_subset(&received) {
+            return Err(SparkError::connect_msg(
+                "Missing cache artifact upload acknowledgement",
+            ));
+        }
+        Ok(hashes)
+    }
+
     /// Check artifact status on the server.
     ///
     /// Mirrors functionality from `pyspark.sql.connect.client.artifact.ArtifactManager.is_cached_artifact`.