diff --git a/Cargo.lock b/Cargo.lock index 29e7d3fe..6238cb96 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -87,6 +87,7 @@ dependencies = [ "prost-types", "regex", "serde_json", + "sha2", "tempfile", "thiserror", "tokio", diff --git a/crates/pyspark-rs/src/lib.rs b/crates/pyspark-rs/src/lib.rs index fec1a7d0..68507bd5 100644 --- a/crates/pyspark-rs/src/lib.rs +++ b/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; diff --git a/crates/pyspark-rs/src/pyarrow_import.rs b/crates/pyspark-rs/src/pyarrow_import.rs new file mode 100644 index 00000000..058fb99f --- /dev/null +++ b/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 { + 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::().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::()? == 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> = 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)) +} diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 5f6302e3..0e271dbd 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/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 diff --git a/crates/spark-connect-core/Cargo.toml b/crates/spark-connect-core/Cargo.toml index c6a0c088..6cdf0e01 100644 --- a/crates/spark-connect-core/Cargo.toml +++ b/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" diff --git a/crates/spark-connect-core/src/client.rs b/crates/spark-connect-core/src/client.rs index 31820ed5..8f658bdc 100644 --- a/crates/spark-connect-core/src/client.rs +++ b/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>) -> Result> { + use crate::artifact::InMemoryArtifact; + use sha2::{Digest, Sha256}; + use std::collections::HashSet; + + let hashes: Vec = blobs + .iter() + .map(|blob| format!("{:x}", Sha256::digest(blob))) + .collect(); + let names: Vec = 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::>()) + .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 = 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`. diff --git a/crates/spark-connect/src/arrow_input.rs b/crates/spark-connect/src/arrow_input.rs new file mode 100644 index 00000000..8ef091e3 --- /dev/null +++ b/crates/spark-connect/src/arrow_input.rs @@ -0,0 +1,235 @@ +//! Arrow local relations and session-scoped chunked uploads. + +use arrow::datatypes::SchemaRef; +use arrow::ipc::writer::StreamWriter; +use arrow::record_batch::RecordBatch; +use spark_connect_core::client::SparkConnectClient; +use spark_connect_core::error::{Result, SparkError, SparkErrorKind}; +use spark_connect_core::runtime::block_on; + +use crate::plan::LogicalPlan; +use crate::types::DataType; + +struct Options { + cache_threshold: usize, + size_limit: usize, + chunk_rows: usize, + chunk_bytes: usize, + upload_batch_bytes: usize, +} + +impl Options { + fn from_client(client: &SparkConnectClient) -> Result { + let keys = [ + "spark.sql.session.localRelationCacheThreshold", + "spark.sql.session.localRelationSizeLimit", + "spark.sql.session.localRelationChunkSizeRows", + "spark.sql.session.localRelationChunkSizeBytes", + "spark.sql.session.localRelationBatchOfChunksSizeBytes", + ]; + let values = block_on(client.get_configs(&keys))?; + let mut parsed = Vec::new(); + for (i, key) in keys.iter().enumerate() { + let value = values + .get(i) + .and_then(Option::as_deref) + .and_then(|v| v.parse::().ok()) + .filter(|v| i == 0 || *v > 0) + .ok_or_else(|| { + SparkError::connect_msg(format!("Invalid or missing configuration: {key}")) + })?; + parsed.push(value); + } + Ok(Self { + cache_threshold: parsed[0], + size_limit: parsed[1], + chunk_rows: parsed[2], + chunk_bytes: parsed[3].min(parsed[4]), + upload_batch_bytes: parsed[4], + }) + } +} + +fn arrow_error(error: arrow::error::ArrowError) -> SparkError { + SparkError::connect_msg(error.to_string()) +} + +fn validate(batches: &[RecordBatch]) -> Result<()> { + if let Some(first) = batches.first() { + if batches.iter().any(|batch| batch.schema() != first.schema()) { + return Err(SparkError::connect_msg("Arrow batch schemas must match")); + } + } + Ok(()) +} + +pub(crate) fn inline_relation(batches: &[RecordBatch], schema: DataType) -> Result { + validate(batches)?; + let data = if let Some(first) = batches.first() { + let mut writer = StreamWriter::try_new(Vec::new(), &first.schema()).map_err(arrow_error)?; + for batch in batches { + writer.write(batch).map_err(arrow_error)?; + } + Some(writer.into_inner().map_err(arrow_error)?) + } else { + None + }; + Ok(LogicalPlan::LocalRelation { schema, data }) +} + +pub(crate) fn create_relation( + client: &SparkConnectClient, + batches: Vec, + schema: DataType, +) -> Result { + // Validate before any RPC, and preserve the offline empty-input path. + validate(&batches)?; + if batches.iter().all(|batch| batch.num_rows() == 0) { + return inline_relation(&batches, schema); + } + let options = Options::from_client(client)?; + // Count the logical slices rather than their potentially much larger backing buffers. + let bytes: usize = batches + .iter() + .flat_map(|b| b.columns()) + .map(|a| { + a.to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| a.get_buffer_memory_size()) + }) + .sum(); + if bytes < options.cache_threshold { + return inline_relation(&batches, schema); + } + + let mut pending = Vec::new(); + let mut pending_size = 0; + let mut total_size = 0; + let mut hashes = Vec::new(); + let mut enqueue = |chunk: Vec| -> Result<()> { + total_size += chunk.len(); + if total_size > options.size_limit { + return Err(SparkError::classed( + SparkErrorKind::Analysis, + "LOCAL_RELATION_SIZE_LIMIT_EXCEEDED", + &[ + ("actualSize", &total_size.to_string()), + ("sizeLimit", &options.size_limit.to_string()), + ], + )); + } + if !pending.is_empty() && pending_size + chunk.len() > options.upload_batch_bytes { + hashes.extend(block_on( + client.cache_artifacts(std::mem::take(&mut pending)), + )?); + pending_size = 0; + } + pending_size += chunk.len(); + pending.push(chunk); + Ok(()) + }; + enqueue(schema.json().into_bytes())?; + serialize_chunks(&batches, &options, &mut enqueue)?; + hashes.extend(block_on(client.cache_artifacts(pending))?); + let schema_hash = hashes.remove(0); + Ok(LogicalPlan::ChunkedCachedLocalRelation { + data_hashes: hashes, + schema_hash: Some(schema_hash), + }) +} + +fn new_writer(schema: &SchemaRef) -> Result>> { + StreamWriter::try_new(Vec::new(), schema).map_err(arrow_error) +} + +/// Each chunk is an independently readable IPC stream. Check limits at most every +/// 1024 rows, as reference PySpark does; a single oversized row remains indivisible. +fn serialize_chunks( + batches: &[RecordBatch], + options: &Options, + emit: &mut impl FnMut(Vec) -> Result<()>, +) -> Result<()> { + let schema = batches[0].schema(); + let mut writer = new_writer(&schema)?; + let mut rows = 0; + for batch in batches { + let mut offset = 0; + while offset < batch.num_rows() { + let length = (batch.num_rows() - offset) + .min(1024) + .min(options.chunk_rows - rows); + writer + .write(&batch.slice(offset, length)) + .map_err(arrow_error)?; + offset += length; + rows += length; + if rows >= options.chunk_rows || writer.get_ref().len() >= options.chunk_bytes { + emit(writer.into_inner().map_err(arrow_error)?)?; + writer = new_writer(&schema)?; + rows = 0; + } + } + } + if rows > 0 { + emit(writer.into_inner().map_err(arrow_error)?)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow::datatypes::{DataType as ArrowType, Field, Schema}; + use arrow::ipc::reader::StreamReader; + use std::io::Cursor; + use std::sync::Arc; + + #[test] + fn chunks_preserve_rows_schema_and_order() { + let schema = Arc::new(Schema::new(vec![Field::new("id", ArrowType::Int64, true)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from_iter_values(0..2500))], + ) + .unwrap(); + for (chunk_rows, chunk_bytes) in [(700, usize::MAX), (10000, 1024)] { + let options = Options { + cache_threshold: 0, + size_limit: usize::MAX, + chunk_rows, + chunk_bytes, + upload_batch_bytes: usize::MAX, + }; + let mut chunks = Vec::new(); + serialize_chunks(std::slice::from_ref(&batch), &options, &mut |c| { + chunks.push(c); + Ok(()) + }) + .unwrap(); + assert!(chunks.len() > 1); + let mut actual = Vec::new(); + for chunk in chunks { + let reader = StreamReader::try_new(Cursor::new(chunk), None).unwrap(); + assert_eq!(reader.schema(), schema); + let mut rows = 0; + for batch in reader { + let batch = batch.unwrap(); + rows += batch.num_rows(); + actual.extend( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied(), + ); + } + assert!(rows <= chunk_rows); + } + assert_eq!(actual, (0..2500).collect::>()); + } + } +} diff --git a/crates/spark-connect/src/conf.rs b/crates/spark-connect/src/conf.rs index d43d611f..17a16ec6 100644 --- a/crates/spark-connect/src/conf.rs +++ b/crates/spark-connect/src/conf.rs @@ -58,10 +58,18 @@ impl RuntimeConf { /// println!("Shuffle partitions: {:?}", value); /// ``` pub fn get(&self, key: &str) -> Result> { - let results = block_on(self.client.get_configs(&[key]))?; + let results = self.get_many(&[key])?; Ok(results.into_iter().next().flatten()) } + /// Get configuration values in one request, in the same order as `keys`. + /// + /// Missing keys return `None`, as with [`Self::get`]. Values are read from + /// the server on every call; they are not cached locally. + pub fn get_many(&self, keys: &[&str]) -> Result>> { + block_on(self.client.get_configs(keys)) + } + /// Get a configuration value, returning `default` when the key is unset — instead of /// erroring as `get` does for an unknown key. Uses the server's `GetWithDefault` /// operation, mirroring `RuntimeConfig.get(key, default)`. diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index ea0e2c62..d66deb6a 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -1609,7 +1609,10 @@ impl DataFrame { /// Check if this DataFrame is local (collected). pub fn is_local(&self) -> bool { - matches!(self.plan.as_ref(), LogicalPlan::LocalRelation { .. }) + matches!( + self.plan.as_ref(), + LogicalPlan::LocalRelation { .. } | LogicalPlan::ChunkedCachedLocalRelation { .. } + ) } /// Check if this DataFrame is streaming. diff --git a/crates/spark-connect/src/lib.rs b/crates/spark-connect/src/lib.rs index 5f1a8faf..180bf089 100644 --- a/crates/spark-connect/src/lib.rs +++ b/crates/spark-connect/src/lib.rs @@ -5,6 +5,7 @@ #![allow(clippy::result_large_err)] // SparkError is large; changing the type would break the API +mod arrow_input; pub mod catalog; pub mod column; pub mod conf; diff --git a/crates/spark-connect/src/plan.rs b/crates/spark-connect/src/plan.rs index b24309bf..86dacf27 100644 --- a/crates/spark-connect/src/plan.rs +++ b/crates/spark-connect/src/plan.rs @@ -228,6 +228,11 @@ pub enum LogicalPlan { schema: DataType, data: Option>, }, + /// Local Arrow data uploaded as session-scoped cache artifacts. + ChunkedCachedLocalRelation { + data_hashes: Vec, + schema_hash: Option, + }, /// CachedRemoteRelation: a reference to a cached remote relation. CachedRemoteRelation { relation_id: String }, /// Read: `spark.read.format(...).load(...)` or `spark.read.table(...)`. @@ -944,6 +949,17 @@ impl LogicalPlan { Some(proto::relation::RelType::SubqueryAlias(Box::new(sq_alias))); } + LogicalPlan::ChunkedCachedLocalRelation { + data_hashes, + schema_hash, + } => { + relation.rel_type = Some(proto::relation::RelType::ChunkedCachedLocalRelation( + proto::ChunkedCachedLocalRelation { + data_hashes: data_hashes.clone(), + schema_hash: schema_hash.clone(), + }, + )); + } LogicalPlan::LocalRelation { schema, data } => { let mut local = proto::LocalRelation::default(); if let Some(d) = data { diff --git a/crates/spark-connect/src/session.rs b/crates/spark-connect/src/session.rs index 7eb29ace..ce92949e 100644 --- a/crates/spark-connect/src/session.rs +++ b/crates/spark-connect/src/session.rs @@ -290,6 +290,21 @@ impl SparkSession { Ok(DataFrame::new(self.clone(), plan)) } + /// Create a local DataFrame from Arrow batches without materializing rows. + /// + /// `schema` is the Spark schema applied by the server. All batches must share + /// the same Arrow schema; IPC serialization is still required for transport. + /// Non-empty inputs read the server's local-relation configuration. Large inputs + /// are uploaded as cache artifacts and referenced by hash in the resulting plan. + pub fn create_dataframe_from_arrow( + &self, + batches: Vec, + schema: DataType, + ) -> Result { + let plan = crate::arrow_input::create_relation(&self.client, batches, schema)?; + Ok(DataFrame::new(self.clone(), plan)) + } + /// Create a DataFrameReader for reading data from various sources. /// /// Mirrors `pyspark.sql.SparkSession.read`. @@ -1224,6 +1239,53 @@ mod tests { .expect("failed to build session") } + #[test] + fn arrow_local_relation_preserves_batches_and_rejects_schema_mismatch() { + use arrow::array::Array; + use arrow::datatypes::{DataType as ArrowType, Field, Schema}; + let arrow_schema = Arc::new(Schema::new(vec![Field::new("x", ArrowType::Int32, true)])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![Arc::new(Int32Array::from(vec![Some(1), None]))], + ) + .unwrap(); + let schema = DataType::Struct { + fields: vec![struct_field("x", DataType::Integer)], + }; + let spark = session(); + let plan = + crate::arrow_input::inline_relation(&[batch.clone(), batch.clone()], schema.clone()) + .unwrap(); + let LogicalPlan::LocalRelation { + data: Some(data), .. + } = &plan + else { + panic!("expected Arrow local relation") + }; + let batches = StreamReader::try_new(Cursor::new(data), None) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.len(), 2); + assert_eq!(batches[0], batch); + assert!(batches[1].column(0).is_null(1)); + let other = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("y", ArrowType::Int32, true)])), + batch.columns().to_vec(), + ) + .unwrap(); + assert!(spark + .create_dataframe_from_arrow(vec![batch, other], schema.clone()) + .is_err()); + assert!(matches!( + spark + .create_dataframe_from_arrow(vec![], schema) + .unwrap() + .plan(), + LogicalPlan::LocalRelation { data: None, .. } + )); + } + #[test] fn session_is_not_stopped_initially() { let spark = session(); diff --git a/docs/architecture.md b/docs/architecture.md index 17559ed9..93107c3e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -111,6 +111,19 @@ Batch schema consistency is validated before export. Field names or timezones containing NUL use IPC over the collected batches, without another RPC. Empty batch vectors retain the empty-schema behavior. +`createDataFrame(pyarrow_table)` uses the reverse boundary: Spark's schema and +timestamp conversion helpers normalize the table, then arrow-rs imports its +`RecordBatchReader` through the C Stream interface. The Rust core serializes +those batches to IPC for the local relation sent to Spark. No Python rows are +materialized; casts, buffer alignment, and IPC serialization may still copy. +Empty tables require an explicit schema, matching PySpark Connect. +Inputs at or above `spark.sql.session.localRelationCacheThreshold` are serialized +into independent IPC streams and uploaded as session-scoped cache artifacts. +The plan carries `ChunkedCachedLocalRelation` hashes instead of embedding the +entire payload. Chunk sizes, upload batch sizes, and the total relation size +limit follow the server's local-relation settings; artifact RPC messages use +32 KiB pieces, with SHA-256 deduplication and CRC checks. + ## Two front-ends, one core === "Rust (native)" diff --git a/python/pyspark/sql/_arrow_input.py b/python/pyspark/sql/_arrow_input.py new file mode 100644 index 00000000..f77540c8 --- /dev/null +++ b/python/pyspark/sql/_arrow_input.py @@ -0,0 +1,66 @@ +# +# 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. +# + + +"""Schema normalization for the native Arrow Table input path.""" + +from pyspark.errors import PySparkValueError +from pyspark.sql.pandas.types import ( + _check_arrow_table_timestamps_localize, + from_arrow_schema, + to_arrow_schema, +) +from pyspark.sql.types import StructType + + +def prepare_table(table, schema, timestamp_type, timezone, large_types): + """Match Connect's Arrow casts without converting any column to Python rows.""" + names = list(schema) if isinstance(schema, (list, tuple)) else None + explicit = schema is not None and names is None + if len(table) == 0: + if not explicit: + raise PySparkValueError(errorClass="CANNOT_INFER_EMPTY_SCHEMA", messageParameters={}) + return table, schema, None + + if names is not None: + names.extend(f"_{i + 1}" for i in range(len(names), table.num_columns)) + inferred = ( + schema + if isinstance(schema, StructType) + else from_arrow_schema(table.schema, prefer_timestamp_ntz=timestamp_type == "TIMESTAMP_NTZ") + ) + normalized = ( + _check_arrow_table_timestamps_localize(table, inferred, True, timezone) + .cast( + to_arrow_schema( + inferred, + error_on_duplicated_field_names_in_struct=True, + timezone="UTC", + prefers_large_types=large_types == "true", + ) + ) + .rename_columns(inferred.names) + ) + if names is not None and len(names) != normalized.num_columns: + raise PySparkValueError( + errorClass="AXIS_LENGTH_MISMATCH", + messageParameters={ + "expected_length": str(len(names)), + "actual_length": str(normalized.num_columns), + }, + ) + return normalized, schema if explicit else inferred, names diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 214a09a0..bba16caa 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -73,18 +73,6 @@ def message(name, **fields): return message -@pytest.mark.parametrize("in_ci", [False, True]) -def test_missing_protoc_fails_in_ci_and_skips_locally(monkeypatch, tmp_path_factory, in_ci): - monkeypatch.setattr(shutil, "which", lambda _: None) - if in_ci: - monkeypatch.setenv("CI", "true") - else: - monkeypatch.delenv("CI", raising=False) - error = pytest.fail.Exception if in_ci else pytest.skip.Exception - with pytest.raises(error, match="require protoc"): - protocol.__wrapped__(tmp_path_factory) - - @contextmanager def _dataframe(batches, *, protocol, payloads=None, calls=None): """Serve Arrow results using the repository's protobuf message definitions.""" @@ -219,16 +207,6 @@ def test_handoff_preserves_values_metadata_and_chunks(protocol): assert calls == {"execute": 1} -def test_normal_handoff_does_not_reopen_ipc(protocol, monkeypatch): - def unexpected_ipc(*args, **kwargs): - pytest.fail("Normal toArrow handoff must use the C Stream importer") - - monkeypatch.setattr(pa.ipc, "open_file", unexpected_ipc) - batch = _mixed_batch() - with _dataframe(protocol=protocol, batches=[batch]) as (frame, _): - assert frame.toArrow().equals(pa.Table.from_batches([batch]), check_metadata=True) - - def test_handoff_buffers_outlive_dataframe_session_and_table(protocol): batch = _mixed_batch() expected = batch.to_pydict() @@ -245,10 +223,13 @@ def collect(): assert table.to_pydict() == expected numeric = table.column(0).chunk(0).slice(1, 3) strings = table.column(2).chunk(0).slice(1, 3) + addresses = [buffer.address for buffer in (numeric.buffers()[1], strings.buffers()[2])] del table gc.collect() assert numeric.to_pylist() == [None, -3, 4] assert strings.to_pylist() == [None, "", "omega"] + assert numeric.buffers()[1].address == addresses[0] + assert strings.buffers()[2].address == addresses[1] assert calls == {"execute": 1} @@ -395,3 +376,367 @@ def test_malformed_ipc_is_reported(protocol): with pytest.raises(RuntimeError, match="Arrow"): frame.toArrow() assert calls == {"execute": 1} + + +@contextmanager +def _arrow_input_service(protocol, *, upload_error=None, **overrides): + """Resolve cache artifacts at a mock service with a 4 MiB receive limit.""" + import hashlib + import zlib + + config = { + "spark.sql.timestampType": "TIMESTAMP_LTZ", + "spark.sql.session.timeZone": "UTC", + "spark.sql.execution.arrow.useLargeVarTypes": "false", + "spark.sql.session.localRelationCacheThreshold": "1048576", + "spark.sql.session.localRelationSizeLimit": str(3 * 1024**3), + "spark.sql.session.localRelationChunkSizeRows": "10000", + "spark.sql.session.localRelationChunkSizeBytes": str(16 * 1024**2), + "spark.sql.session.localRelationBatchOfChunksSizeBytes": str(1024**3), + } + config.update(overrides) + state = {"cache": {}, "uploads": 0, "relations": [], "message_sizes": [], + "config": config, "config_requests": [], "input_tables": []} + server_id = "00000000-0000-0000-0000-000000000001" + + def configure(raw, context): + request = protocol("ConfigRequest").FromString(raw) + keys = list(request.operation.get.keys) + state["config_requests"].append(keys) + if error := state.pop("config_error", None): + context.abort(error, "configuration read failed") + # Deliberately reverse responses and omit missing keys to test key mapping. + return protocol("ConfigResponse", session_id=request.session_id, + pairs=[dict(key=k, value=config[k]) for k in reversed(keys) + if k in config] + ).SerializeToString() + + def statuses(raw, context): + request = protocol("ArtifactStatusesRequest").FromString(raw) + return protocol("ArtifactStatusesResponse", session_id=request.session_id, + statuses={name: dict(exists=name in state["cache"]) + for name in request.names}).SerializeToString() + + def upload(requests, context): + summaries = [] + name, expected, buffer = None, 0, bytearray() + session_id = None + + def save(name, data): + assert name == "cache/" + hashlib.sha256(data).hexdigest() + assert name not in state["cache"], "cached artifacts should not be uploaded again" + state["cache"][name] = bytes(data) + summaries.append(dict(name=name, is_crc_successful=True)) + + def checked(chunk): + assert zlib.crc32(chunk.data) == chunk.crc + return chunk.data + + for raw in requests: + state["message_sizes"].append(len(raw)) + request = protocol("AddArtifactsRequest").FromString(raw) + session_id = request.session_id + if request.HasField("batch"): + for artifact in request.batch.artifacts: + save(artifact.name, checked(artifact.data)) + elif request.HasField("begin_chunk"): + begin = request.begin_chunk + name, expected = begin.name, begin.total_bytes + buffer = bytearray(checked(begin.initial_chunk)) + else: + buffer.extend(checked(request.chunk)) + if name is not None and len(buffer) == expected: + save(name, buffer) + name, buffer = None, bytearray() + assert name is None + state["uploads"] += 1 + if upload_error == "crc": + summaries[0]["is_crc_successful"] = False + elif upload_error == "missing": + summaries = [] + return protocol("AddArtifactsResponse", session_id=session_id, + artifacts=summaries).SerializeToString() + + def execute(raw, context): + request = protocol("ExecutePlanRequest").FromString(raw) + relation = request.plan.root + state["relations"].append(relation) + if relation.HasField("chunked_cached_local_relation"): + cached = relation.chunked_cached_local_relation + payloads = [state["cache"]["cache/" + h] for h in cached.dataHashes] + state["spark_schema"] = state["cache"]["cache/" + cached.schemaHash].decode() + else: + payloads = [relation.local_relation.data] + rows, first, last = 0, None, None + tables = [pa.ipc.open_stream(payload).read_all() for payload in payloads] + state["input_tables"].append(pa.concat_tables(tables)) + for table in tables: + for batch in table.to_batches(): + rows += batch.num_rows + if batch.num_rows: + if first is None: + first = batch.column(0)[0].as_py() + last = batch.column(0)[-1].as_py() + state["result"] = (rows, first, last) + result = pa.table({"rows": [rows]}) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, result.schema) as writer: + writer.write_table(result) + header = dict(session_id=request.session_id, operation_id=request.operation_id, + server_side_session_id=server_id) + yield protocol("ExecutePlanResponse", **header, response_id="0", + arrow_batch=dict(row_count=1, data=sink.getvalue().to_pybytes()) + ).SerializeToString() + yield protocol("ExecutePlanResponse", **header, response_id="1", + result_complete={}).SerializeToString() + + def release(raw, context): + request = protocol("ReleaseExecuteRequest").FromString(raw) + return protocol("ReleaseExecuteResponse", session_id=request.session_id, + operation_id=request.operation_id, + server_side_session_id=server_id).SerializeToString() + + def stop(raw, context): + request = protocol("ReleaseSessionRequest").FromString(raw) + return protocol("ReleaseSessionResponse", session_id=request.session_id, + server_side_session_id=server_id).SerializeToString() + + with ThreadPoolExecutor(max_workers=2) as executor: + server = grpc.server(executor, options=[("grpc.max_receive_message_length", 4 * 1024**2)]) + server.add_generic_rpc_handlers([grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", { + "Config": grpc.unary_unary_rpc_method_handler(configure), + "ArtifactStatus": grpc.unary_unary_rpc_method_handler(statuses), + "AddArtifacts": grpc.stream_unary_rpc_method_handler(upload), + "ExecutePlan": grpc.unary_stream_rpc_method_handler(execute), + "ReleaseExecute": grpc.unary_unary_rpc_method_handler(release), + "ReleaseSession": grpc.unary_unary_rpc_method_handler(stop), + })]) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + session = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() + try: + yield session, state + finally: + session.stop() + server.stop(0).wait(timeout=5) + + +@pytest.mark.parametrize("rows", [2, 17000000]) +def test_arrow_input_upload_above_grpc_limit(protocol, rows): + import numpy as np + import json + + table = pa.table({"id": np.arange(rows, dtype=np.int64)}) + with _arrow_input_service(protocol) as (spark, state): + df = spark.createDataFrame(table) + assert df.isLocal() + assert df.toArrow().column(0).to_pylist() == [rows] + assert state["result"] == (rows, 0, rows - 1) + if rows == 2: + assert state["uploads"] == 0 + assert state["relations"][0].HasField("local_relation") + else: + assert table.nbytes > 128 * 1024**2 + assert state["relations"][0].HasField("chunked_cached_local_relation") + assert state["relations"][0].ByteSize() < 1024**2 + assert max(state["message_sizes"]) < 33 * 1024 + assert json.loads(state["spark_schema"])["fields"][0]["name"] == "id" + uploads = state["uploads"] + spark.createDataFrame(table).toArrow() + assert state["uploads"] == uploads + + +def test_arrow_input_upload_limits_and_chunk_order(protocol): + import numpy as np + from pyspark.errors import AnalysisException + table = pa.table({"id": np.arange(100, dtype=np.int64)}) + options = { + "spark.sql.session.localRelationCacheThreshold": "0", + "spark.sql.session.localRelationChunkSizeRows": "7", + "spark.sql.session.localRelationChunkSizeBytes": "1000", + "spark.sql.session.localRelationBatchOfChunksSizeBytes": "1000", + } + with _arrow_input_service(protocol, **options) as (spark, state): + df = spark.createDataFrame(table) + df.toArrow() + assert state["result"] == (100, 0, 99) + assert state["uploads"] > 1 + cached = state["relations"][0].chunked_cached_local_relation + assert len(cached.dataHashes) > 1 + for hash in cached.dataHashes: + assert pa.ipc.open_stream(state["cache"]["cache/" + hash]).read_all().num_rows <= 7 + + options["spark.sql.session.localRelationSizeLimit"] = "10" + with _arrow_input_service(protocol, **options) as (spark, state): + with pytest.raises(AnalysisException, match="LOCAL_RELATION_SIZE_LIMIT_EXCEEDED"): + spark.createDataFrame(table) + assert not state["relations"] + assert state["uploads"] == 0 + + +@pytest.mark.parametrize("failure,match", [("crc", "CRC check failed"), + ("missing", "Missing cache artifact")]) +def test_arrow_input_upload_failure(protocol, failure, match): + with _arrow_input_service(protocol, upload_error=failure, + **{"spark.sql.session.localRelationCacheThreshold": "0"}) as (spark, state): + with pytest.raises(Exception, match=match): + spark.createDataFrame(pa.table({"id": [1, 2]})) + assert not state["relations"] + + +def test_arrow_input_duplicate_chunks_preserve_multiplicity(protocol): + options = {"spark.sql.session.localRelationCacheThreshold": "0", + "spark.sql.session.localRelationChunkSizeRows": "5"} + with _arrow_input_service(protocol, **options) as (spark, state): + spark.createDataFrame(pa.table({"id": [1] * 20})).toArrow() + cached = state["relations"][0].chunked_cached_local_relation + assert len(cached.dataHashes) == 4 + assert len(set(cached.dataHashes)) == 1 + assert len(state["cache"]) == 2 # One schema and one unique data chunk. + assert state["result"] == (20, 1, 1) + + +_NORMALIZATION_KEYS = [ + "spark.sql.timestampType", + "spark.sql.session.timeZone", + "spark.sql.execution.arrow.useLargeVarTypes", +] +_UPLOAD_KEYS = [ + "spark.sql.session.localRelationCacheThreshold", + "spark.sql.session.localRelationSizeLimit", + "spark.sql.session.localRelationChunkSizeRows", + "spark.sql.session.localRelationChunkSizeBytes", + "spark.sql.session.localRelationBatchOfChunksSizeBytes", +] + + +@pytest.mark.parametrize("timestamp_type", ["TIMESTAMP_LTZ", "TIMESTAMP_NTZ"]) +@pytest.mark.parametrize("zone", ["UTC", "America/Los_Angeles"]) +@pytest.mark.parametrize("large", ["false", "true"]) +def test_arrow_input_batched_normalization_configs(protocol, timestamp_type, zone, large): + import json + from zoneinfo import ZoneInfo + + value = datetime(2024, 1, 15, 12, 34, 56, 123456) + table = pa.table({"t": [value, None], "s": ["hello", None], "b": [b"xyz", None]}) + with _arrow_input_service(protocol, **dict(zip( + _NORMALIZATION_KEYS, [timestamp_type, zone, large] + ))) as (spark, state): + frame = spark.createDataFrame(table) + assert state["config_requests"] == [_NORMALIZATION_KEYS, _UPLOAD_KEYS] + frame.collect() + expected_time = (value if timestamp_type == "TIMESTAMP_NTZ" else + value.replace(tzinfo=ZoneInfo(zone)).astimezone(timezone.utc)) + expected = pa.table({ + "t": pa.array([expected_time, None], type=pa.timestamp( + "us", tz=None if timestamp_type == "TIMESTAMP_NTZ" else "UTC")), + "s": pa.array(["hello", None], type=pa.large_string() if large == "true" + else pa.string()), + "b": pa.array([b"xyz", None], type=pa.large_binary() if large == "true" + else pa.binary()), + }) + actual = state["input_tables"][0] + assert actual.schema == expected.schema + assert actual.to_pydict() == expected.to_pydict() + fields = json.loads(state["relations"][0].local_relation.schema)["fields"] + assert [f["type"] for f in fields] == [ + "timestamp_ntz" if timestamp_type == "TIMESTAMP_NTZ" else "timestamp", + "string", "binary", + ] + + +@pytest.mark.parametrize("missing", _NORMALIZATION_KEYS) +@pytest.mark.parametrize("present_null", [False, True]) +def test_arrow_input_missing_normalization_config(protocol, missing, present_null): + with _arrow_input_service(protocol) as (spark, state): + if present_null: + state["config"][missing] = None + else: + del state["config"][missing] + spark.createDataFrame(pa.table({"x": [1, None], "s": ["hello", None]})).collect() + assert state["config_requests"] == [_NORMALIZATION_KEYS, _UPLOAD_KEYS] + expected = pa.table({"x": pa.array([1, None], type=pa.int64()), + "s": ["hello", None]}) + assert state["input_tables"][0].equals(expected) + + +def test_arrow_input_normalization_config_failure_then_fresh_read(protocol): + with _arrow_input_service(protocol) as (spark, state): + state["config_error"] = grpc.StatusCode.INVALID_ARGUMENT + table = pa.table({"x": [1]}) + with pytest.raises(Exception, match="configuration read failed"): + spark.createDataFrame(table) + assert state["config_requests"] == [_NORMALIZATION_KEYS] + assert state["uploads"] == 0 + spark.createDataFrame(table).collect() + assert state["config_requests"] == [ + _NORMALIZATION_KEYS, _NORMALIZATION_KEYS, _UPLOAD_KEYS] + assert state["input_tables"][0].to_pydict() == {"x": [1]} + + +def test_arrow_input_normalization_config_transient_retry(protocol): + with _arrow_input_service(protocol) as (spark, state): + state["config_error"] = grpc.StatusCode.UNAVAILABLE + spark.createDataFrame(pa.table({"x": [1]})).collect() + assert state["config_requests"] == [ + _NORMALIZATION_KEYS, _NORMALIZATION_KEYS, _UPLOAD_KEYS] + assert state["input_tables"][0].to_pydict() == {"x": [1]} + + +def test_arrow_input_normalization_configs_are_fresh(protocol): + import json + + value = datetime(2024, 1, 15, 12) + table = pa.table({"t": [value], "s": ["hello"]}) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(table).collect() + state["config"].update(dict(zip(_NORMALIZATION_KEYS, [ + "TIMESTAMP_NTZ", "America/Los_Angeles", "true"]))) + spark.createDataFrame(table).collect() + state["config"].update(dict(zip(_NORMALIZATION_KEYS, [ + "TIMESTAMP_LTZ", "America/Los_Angeles", "false"]))) + spark.createDataFrame(table).collect() + assert state["config_requests"] == [_NORMALIZATION_KEYS, _UPLOAD_KEYS] * 3 + tables = state["input_tables"] + assert [t.column("t")[0].as_py() for t in tables] == [ + value.replace(tzinfo=timezone.utc), value, + datetime(2024, 1, 15, 20, tzinfo=timezone.utc)] + assert [t.schema.field("s").type for t in tables] == [ + pa.string(), pa.large_string(), pa.string()] + assert [json.loads(r.local_relation.schema)["fields"][0]["type"] + for r in state["relations"]] == ["timestamp", "timestamp_ntz", "timestamp"] + + +def test_arrow_input_empty_skips_all_configs(protocol): + from pyspark.sql.types import LongType, StructField, StructType + + with _arrow_input_service(protocol) as (spark, state): + table = pa.table({"x": pa.array([], type=pa.int64())}) + spark.createDataFrame(table, StructType([StructField("x", LongType())])) + with pytest.raises(Exception, match="CANNOT_INFER_EMPTY_SCHEMA"): + spark.createDataFrame(table) + assert state["config_requests"] == [] + assert state["uploads"] == 0 + + +@pytest.mark.parametrize("in_ci", [False, True]) +def test_missing_protoc_fails_in_ci_and_skips_locally(monkeypatch, tmp_path_factory, in_ci): + monkeypatch.setattr(shutil, "which", lambda _: None) + if in_ci: + monkeypatch.setenv("CI", "true") + else: + monkeypatch.delenv("CI", raising=False) + error = pytest.fail.Exception if in_ci else pytest.skip.Exception + with pytest.raises(error, match="require protoc"): + protocol.__wrapped__(tmp_path_factory) + + +def test_normal_handoff_does_not_reopen_ipc(protocol, monkeypatch): + def unexpected_ipc(*args, **kwargs): + pytest.fail("Normal toArrow handoff must use the C Stream importer") + + monkeypatch.setattr(pa.ipc, "open_file", unexpected_ipc) + batch = _mixed_batch() + with _dataframe(protocol=protocol, batches=[batch]) as (frame, _): + assert frame.toArrow().equals(pa.Table.from_batches([batch]), check_metadata=True) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 67ccaa4d..168ce7b7 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -2755,3 +2755,132 @@ def test_create_dataframe_width_mismatch_is_a_normal_exception(): else: raise AssertionError("expected a width validation error") """], check=True) + + +@pytest.mark.parametrize("names,expected", [(None, None), (["a"], ["a", "_2"]), ([], ["_1", "_2"])]) +def test_arrow_input_normalization(names, expected): + pa = pytest.importorskip("pyarrow") + from pyspark.sql._arrow_input import prepare_table + + table = pa.table({"x": pa.chunked_array([[1], [2]]), "y": ["a", None]}) + normalized, schema, columns = prepare_table(table, names, "TIMESTAMP_LTZ", "UTC", "false") + assert normalized.equals(table) + assert schema.names == ["x", "y"] + assert columns == expected + + +def test_arrow_input_empty_and_width_errors(): + pa = pytest.importorskip("pyarrow") + from pyspark.errors import PySparkValueError + from pyspark.sql._arrow_input import prepare_table + + empty = pa.table({"x": pa.array([], type=pa.int64())}) + with pytest.raises(PySparkValueError, match="CANNOT_INFER_EMPTY_SCHEMA"): + prepare_table(empty, None, None, None, None) + schema = T.StructType([T.StructField("x", T.LongType())]) + assert prepare_table(empty, schema, None, None, None)[1] is schema + with pytest.raises(PySparkValueError, match="AXIS_LENGTH_MISMATCH"): + prepare_table(pa.table({"x": [1]}), ["a", "b"], "TIMESTAMP_LTZ", "UTC", "false") + + +@pytest.mark.parametrize("timestamp_type,expected", [("TIMESTAMP_NTZ", T.TimestampNTZType), ("TIMESTAMP_LTZ", T.TimestampType)]) +def test_arrow_input_timestamps(timestamp_type, expected): + pa = pytest.importorskip("pyarrow") + from pyspark.sql._arrow_input import prepare_table + + table = pa.table({"t": pa.array([1001], type=pa.timestamp("ns"))}) + normalized, schema, _ = prepare_table(table, None, timestamp_type, "UTC", "false") + assert isinstance(schema.fields[0].dataType, expected) + assert normalized.column(0).type.unit == "us" + assert normalized.column(0).cast(pa.int64()).to_pylist() == [1] + + +def test_arrow_input_explicit_cast_and_nested_values(): + pa = pytest.importorskip("pyarrow") + from pyspark.sql._arrow_input import prepare_table + + table = pa.table({"x": [1, 2], "nested": [[1, None], []]}) + schema = T.StructType([T.StructField("x", T.IntegerType()), + T.StructField("nested", T.ArrayType(T.LongType()))]) + normalized, result, _ = prepare_table(table, schema, "TIMESTAMP_LTZ", "UTC", "false") + assert result is schema + assert normalized.column(0).type == pa.int32() + assert normalized.to_pylist() == table.to_pylist() + with pytest.raises(pa.ArrowInvalid): + prepare_table(pa.table({"x": [2**40]}), T.StructType([T.StructField("x", T.IntegerType())]), + "TIMESTAMP_LTZ", "UTC", "false") + + +def test_arrow_input_roundtrip_live_server(): + """Run explicitly with SPARK_REMOTE; also checks imported buffers own their lifetime.""" + import gc + import os + pa = pytest.importorskip("pyarrow") + remote = os.environ.get("SPARK_REMOTE") + if not remote: + pytest.skip("requires SPARK_REMOTE") + from pyspark.sql import SparkSession + + spark = SparkSession.builder.remote(remote).getOrCreate() + try: + table = pa.table({"x": pa.chunked_array([[1, 2], [3]]), "nested": [[1], None, []]}) + df = spark.createDataFrame(table, ["id", "items"]) + del table + gc.collect() + assert df.columns == ["id", "items"] + assert [(r[0], r[1]) for r in df.collect()] == [(1, [1]), (2, None), (3, [])] + empty = spark.createDataFrame(pa.table({"x": pa.array([], type=pa.int64())}), + T.StructType([T.StructField("x", T.LongType())])) + assert empty.count() == 0 + assert empty.columns == ["x"] + finally: + spark.stop() + + +@pytest.mark.parametrize("kind", ["slice", "dictionary", "decimal", "null", "struct", "nul"]) +def test_arrow_input_types_live_server(kind): + import os + from decimal import Decimal + pa = pytest.importorskip("pyarrow") + remote = os.environ.get("SPARK_REMOTE") + if not remote: + pytest.skip("requires SPARK_REMOTE") + from pyspark.sql import SparkSession + + tables = { + "slice": pa.table({"x": pa.array([0, 1, 2, 3]).slice(1, 2)}), + "dictionary": pa.table({"x": pa.array(["a", None, "b", "a"]).dictionary_encode()}), + "decimal": pa.table({"x": pa.array([Decimal("1.23"), None], type=pa.decimal128(10, 2))}), + "null": pa.table({"x": pa.nulls(2)}), + "struct": pa.table({"x": [{"a": [1, None]}, None]}), + "nul": pa.table({"a\0b": [1, 2]}), + } + spark = SparkSession.builder.remote(remote).getOrCreate() + try: + table = tables[kind] + df = spark.createDataFrame(table) + assert df.columns == table.column_names + assert df.toArrow().to_pylist() == table.to_pylist() + finally: + spark.stop() + + +def test_arrow_input_large_upload_live_server(): + """A table larger than one gRPC message must use uploaded cache artifacts.""" + import os + pa = pytest.importorskip("pyarrow") + remote = os.environ.get("SPARK_REMOTE") + if not remote: + pytest.skip("requires SPARK_REMOTE") + from pyspark.sql import SparkSession + + rows = 140000 + table = pa.table({"payload": [b"x" * 1024] * rows}) + assert table.nbytes > 128 * 1024**2 + spark = SparkSession.builder.remote(remote).getOrCreate() + try: + df = spark.createDataFrame(table) + assert df.count() == rows + assert df.selectExpr("sum(length(payload)) AS total").collect()[0][0] == rows * 1024 + finally: + spark.stop()