diff --git a/crates/pyspark-rs/src/lib.rs b/crates/pyspark-rs/src/lib.rs index 15796178..807fce24 100644 --- a/crates/pyspark-rs/src/lib.rs +++ b/crates/pyspark-rs/src/lib.rs @@ -36,6 +36,7 @@ mod group; mod infer; mod ml; mod observation; +mod pandas_input; mod pipelines; mod profiler; mod pyarrow_export; @@ -130,6 +131,8 @@ fn _pyspark(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(types::parse_datatype_json, m)?)?; m.add_function(wrap_pyfunction!(column::column_reference_name, m)?)?; + m.add_function(wrap_pyfunction!(pandas_input::_pandas_string_ipc, m)?)?; + // Spark Declarative Pipelines (SDP) command execution m.add_class::()?; m.add_function(wrap_pyfunction!( diff --git a/crates/pyspark-rs/src/pandas_input.rs b/crates/pyspark-rs/src/pandas_input.rs new file mode 100644 index 00000000..ab6eb8fc --- /dev/null +++ b/crates/pyspark-rs/src/pandas_input.rs @@ -0,0 +1,113 @@ +// 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. + +//! Fused validation and inline encoding of guarded pandas string columns. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, StringBuilder}; +use arrow::datatypes::{DataType as ArrowType, Field, Schema}; +use arrow::ipc::writer::StreamWriter; +use arrow::record_batch::RecordBatch; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyString}; + +use crate::types::PyStructType; + +/// The Python eligibility helper authenticates pandas storage and native schema. +/// Validate each object before accepting it: extraction must never call a custom +/// numeric/string conversion hook or turn a pandas missing sentinel into null. +#[pyfunction] +pub(crate) fn _pandas_string_ipc<'py>( + py: Python<'py>, + columns: Vec>, + schema: PyRef<'py, PyStructType>, +) -> PyResult>> { + if columns.len() != schema.fields.len() { + return Ok(None); + } + // Reject unsupported objects before copying any string data. Otherwise a + // late subclass would pay for most of an Arrow build before the row path. + // Keep this pass to exact type/null checks; UTF-8 extraction stays below. + for (column, field) in columns.iter().zip(&schema.fields) { + if !matches!(&field.data_type, spark_connect::types::DataType::String { collation } + if collation == "UTF8_BINARY") + { + return Ok(None); + } + for value in column.try_iter()? { + let value = value?; + if !(value.is_none() && field.nullable || value.is_exact_instance_of::()) { + return Ok(None); + } + } + } + // A lone surrogate must fall back before allocating partial Arrow buffers. + // Drop each temporary UTF-8 bytes object immediately rather than retaining + // an additional full copy of every column. Keep the cheap type pass above + // so a late subclass does not pay for this UTF-8 traversal. + for (column, field) in columns.iter().zip(&schema.fields) { + for value in column.try_iter()? { + let value = value?; + if value.is_none() && field.nullable { + continue; + } + let Ok(value) = value.cast_exact::() else { + return Ok(None); + }; + if value.encode_utf8().is_err() { + return Ok(None); + } + } + } + let mut arrays: Vec = Vec::with_capacity(columns.len()); + let mut fields = Vec::with_capacity(columns.len()); + for (column, field) in columns.iter().zip(&schema.fields) { + let mut builder = StringBuilder::with_capacity(column.len()?, column.len()? * 8); + for value in column.try_iter()? { + let value = value?; + if value.is_none() && field.nullable { + builder.append_null(); + } else if let Ok(value) = value.cast_exact::() { + let Ok(bytes) = value.encode_utf8() else { + return Ok(None); + }; + let Ok(text) = std::str::from_utf8(bytes.as_bytes()) else { + return Ok(None); + }; + // append_value copies while the bytes object and GIL are live. + // This avoids to_cow's owned String on the abi3-py39 build. + builder.append_value(text); + } else { + return Ok(None); + } + } + arrays.push(Arc::new(builder.finish())); + // Like the row encoder, Spark metadata lives in the logical schema; + // Arrow fields carry the name, standard string type and nullability. + fields.push(Field::new(&field.name, ArrowType::Utf8, field.nullable)); + } + let bytes = py + .detach(move || -> Result, arrow::error::ArrowError> { + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays)?; + let mut writer = StreamWriter::try_new(Vec::new(), &batch.schema())?; + writer.write(&batch)?; + writer.into_inner() + }) + .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))?; + Ok(Some(PyBytes::new(py, &bytes))) +} diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 4b9ce5c5..8345a595 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -765,6 +765,54 @@ impl PySparkSession { } let spec = schema_spec(py, schema)?; if is_loaded_instance(py, data, "pandas", "DataFrame")? { + // Guarded integer columns share the existing Arrow importer. Schema + // validation above and unsupported pandas conversions keep master's + // normal ordering and row fallback. + let table = py + .import("pyspark.sql._pandas_input")? + .getattr("prepare_int64_table")? + .call1((data, schema))?; + if !table.is_none() { + let frame = + crate::pyarrow_import::create_dataframe(py, &self.session, &table, schema)?; + // pandas applies explicit field names with toDF even though the + // local relation already carries them. Empty input stays local. + return Ok(match &spec { + Spec::Struct(fields) if !fields.is_empty() && table.len()? > 0 => { + PyDataFrame::new( + frame + .dataframe + .to_df(fields.iter().map(|field| field.name.as_str()).collect()), + ) + } + _ => frame, + }); + } + let ipc = py + .import("pyspark.sql._pandas_input")? + .getattr("prepare_string_input")? + .call1((data, schema))?; + if !ipc.is_none() { + let fields = match &spec { + Spec::Struct(fields) => fields, + _ => { + return Err(pyo3::exceptions::PyValueError::new_err( + "string input requires a struct schema", + )); + } + }; + let schema = DataType::Struct { + fields: fields.clone(), + }; + let frame = self + .session + .create_dataframe_from_inline_arrow_ipc(ipc.extract()?, schema); + // Match pandas row input: explicit nonempty field names travel + // through toDF as well as the local relation's logical schema. + return Ok(PyDataFrame::new( + frame.to_df(fields.iter().map(|field| field.name.as_str()).collect()), + )); + } return self.create_from_pandas(py, data, spec); } self.create_from_local(py, data, spec) diff --git a/crates/spark-connect/src/session.rs b/crates/spark-connect/src/session.rs index 96c3b1f6..e7290132 100644 --- a/crates/spark-connect/src/session.rs +++ b/crates/spark-connect/src/session.rs @@ -360,6 +360,23 @@ impl SparkSession { Ok(DataFrame::new(self.clone(), plan)) } + /// Internal binding bridge for Arrow IPC already validated and encoded from + /// local rows or columns. Retains the row path's inline transport policy. + #[doc(hidden)] + pub fn create_dataframe_from_inline_arrow_ipc( + &self, + data: Vec, + schema: DataType, + ) -> DataFrame { + DataFrame::new( + self.clone(), + LogicalPlan::LocalRelation { + schema: Some(schema), + data: Some(data), + }, + ) + } + /// Create a local DataFrame from Arrow batches without materializing rows. /// /// `schema` optionally supplies the Spark schema applied by the server. With diff --git a/docs/architecture.md b/docs/architecture.md index 5cce8782..089232b3 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -142,6 +142,16 @@ timestamp conversion helpers normalize the table, then arrow-rs imports its 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. +Pandas inputs with only NumPy `int64` or pandas nullable `Int64` columns also +use this importer, preserving integer precision and null masks without Python +rows. Small frames (fewer than 1,000 rows) with only NumPy columns retain the +row path to avoid configuration RPC overhead. Nullable `Int64` columns always +use Arrow so schema inference does not depend on the first value being non-null. +This initial path accepts inferred schemas, complete string column names, +and explicit `StructType` schemas whose fields are all `LongType`. Other dtypes, +DDL/atomic schemas, schema conversions, NUL names, and duplicate/non-string input names +retain the existing row conversion path. `verifySchema` remains ignored for +pandas Arrow input, matching Spark 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 diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py new file mode 100644 index 00000000..78d1eaf2 --- /dev/null +++ b/python/pyspark/sql/_pandas_input.py @@ -0,0 +1,397 @@ +# +# 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. +# + + +"""Conservative columnar pandas input; unsupported conversions use the row path.""" + + +def _has_canonical_columns(columns): + """Decline label containers whose row-path iteration can run custom hooks.""" + import sys + from types import FunctionType + + import numpy as np + import pandas as pd + from pandas.core import base as base_module + from pandas.core.indexes import base as index_module + + Index = index_module.Index + if (pd.Index is not Index or type(columns) is not Index + or Index.__getattribute__ is not object.__getattribute__ + or "__iter__" in vars(columns)): + return False + + def ordinary(method, name, module): + return (type(method) is FunctionType and method.__name__ == name + and method.__module__ == module.__name__ + and method.__globals__ is vars(module)) + + # The native row path iterates columns before normalizing any cells. + # Decline overrides without invoking them, so fallback preserves that order. + if (Index.__iter__ is not base_module.IndexOpsMixin.__iter__ + or not ordinary(Index.__iter__, "__iter__", base_module)): + return False + values = vars(Index).get("_values") + if (type(values) is not property + or not ordinary(values.fget, "_values", index_module)): + return False + if any("_data" in vars(base) for base in Index.__mro__): + return False + backing = vars(columns).get("_data") + if type(backing) is np.ndarray: + return True + + # Arrow-backed labels already hold PyArrow storage. Do not import optional + # Arrow modules just to decide whether a fallback input is eligible. + pa = sys.modules.get("pyarrow") + if pa is None: + return False + + try: + from pandas.core.arrays import numpy_ as numpy_module, string_ as string_module + from pandas.core.arrays import base as array_base_module, _mixins as mixins_module + from pandas._libs.arrays import NDArrayBacked + from pandas.core.arrays import string_arrow as string_arrow_module + from pandas.core.arrays.arrow import array as arrow_module + except ImportError: + # Older pandas releases may lack these concrete string storage types. + return False + + python_strings = tuple(vars(string_module).get(name) for name in + ("StringArray", "StringArrayNumpySemantics")) + arrow_strings = tuple(vars(string_arrow_module).get(name) for name in + ("ArrowStringArray", "ArrowStringArrayNumpySemantics")) + cls = type(backing) + if not any(cls is candidate for candidate in python_strings + arrow_strings): + return False + if (cls.__getattribute__ is not object.__getattribute__ + or any(name in vars(backing) for name in + ("__iter__", "__getitem__", "__len__", "_box_func", + "tolist", "to_numpy", "isna", "fillna", "_hasna")) + or getattr(cls, "tolist", None) is not getattr( + vars(string_module).get("BaseStringArray"), "tolist", None) + or not ordinary(getattr(cls, "tolist", None), "tolist", string_module)): + return False + if any(cls is candidate for candidate in python_strings): + owner, module = vars(numpy_module).get("NumpyExtensionArray"), numpy_module + string_array = vars(string_module).get("StringArray") + if owner is None or string_array is None: + return False + try: + storage = object.__getattribute__(backing, "_ndarray") + except AttributeError: + return False + if type(storage) is not np.ndarray: + return False + # ExtensionArray iteration dispatches through len and indexed boxing. + for name, hook_module in (("__iter__", array_base_module), + ("__getitem__", mixins_module), ("_box_func", mixins_module)): + if not ordinary(getattr(cls, name, None), name, hook_module): + return False + if (getattr(cls, "__len__", None) is not NDArrayBacked.__len__ + or getattr(cls, "ndim", None) is not NDArrayBacked.ndim + or any("_ndarray" in vars(base) for base in cls.__mro__ + if base is not NDArrayBacked)): + return False + isna = vars(string_array).get("isna", getattr(owner, "isna", None)) + isna_module = string_module if "isna" in vars(string_array) else module + else: + owner, module = vars(arrow_module).get("ArrowExtensionArray"), arrow_module + if owner is None: + return False + try: + storage = object.__getattribute__(backing, "_pa_array") + except AttributeError: + return False + if (type(storage) is not pa.ChunkedArray + or not ordinary(getattr(cls, "__iter__", None), "__iter__", module)): + return False + isna, isna_module = getattr(owner, "isna", None), module + hasna = getattr(cls, "_hasna", None) + if (type(hasna) is not property + or not ordinary(hasna.fget, "_hasna", module) + or getattr(cls, "fillna", None) is not getattr(owner, "fillna", None) + or not ordinary(getattr(cls, "fillna", None), "fillna", module)): + return False + # Keep the existing storage/conversion authentication as well as the + # iterator hooks: these objects also supply labels to pandas normalization. + return (getattr(cls, "to_numpy", None) is getattr(owner, "to_numpy", None) + and ordinary(getattr(cls, "to_numpy", None), "to_numpy", module) + and getattr(cls, "isna", None) is isna + and ordinary(getattr(cls, "isna", None), "isna", isna_module)) + + +def prepare_int64_table(frame, schema): + """Return an Arrow table only when each column is already a signed 64-bit integer. + + In particular, never coerce a mixed pandas frame to a common NumPy dtype: + doing so can round integers above 2**53. Other dtypes and schema conversions + deliberately retain the existing row path until their semantics are covered. + ``verifySchema`` is ignored for pandas Arrow input, as in Spark Connect. + """ + import sys + from types import GetSetDescriptorType, WrapperDescriptorType + + import numpy as np + import pandas as pd + try: + from pandas.core.arrays.integer import Int64Dtype, IntegerArray + from pandas.core.arrays.numpy_ import NumpyExtensionArray + from pandas.core.frame import DataFrame + from pandas.core.series import Series + except ImportError: + # Older supported pandas versions call the NumPy wrapper PandasArray. + # Keep their existing row conversion when these concrete types are absent. + return None + + from pyspark._pyspark import ( + AtomicType, BooleanType, ByteType, ShortType, IntegerType, LongType, StructField, StructType, + ) + from pyspark.sql import types as sql_types + + # A subclass or custom extension array can change the values exposed to Arrow. + # Check concrete identities before consulting dtype equality or conversion hooks. + if pd.DataFrame is not DataFrame or type(frame) is not DataFrame: + return None + + if not _has_canonical_columns(frame.columns): + return None + width = len(frame.columns) + if not width or not frame.columns.is_unique or not all( + type(name) is str and "\x00" not in name for name in frame.columns + ): + return None + dtypes = frame.dtypes + numpy_dtype = np.dtype("int64") + if not all((type(dtype) is type(numpy_dtype) and dtype == numpy_dtype) + or type(dtype) is Int64Dtype + for dtype in dtypes): + return None + # The importer's configuration RPCs outweigh row conversion for small NumPy + # frames. Nullable extension columns still use Arrow for reliable null typing. + if len(frame) < 1000 and not any(type(dtype) is Int64Dtype for dtype in dtypes): + return None + + # The native row path reads native schema fields directly. The Arrow converter + # instead uses Python descriptors and imported class aliases, including aliases + # captured before this helper was loaded. Customized schemas retain the row path. + if sys.modules.get("pyspark.sql.types") is not sql_types: + return None + importer = sys.modules.get("pyspark.sql._arrow_input") + if vars(sql_types).get("AtomicType") is not AtomicType: + return None + if importer is not None: + if getattr(importer, "StructType", None) is not StructType: + return None + # The importer uses this alias while validating schema widths. + if vars(importer).get("AtomicType") is not AtomicType: + return None + converter = sys.modules.get("pyspark.sql.pandas.types") + namespaces = [vars(sql_types)] + if converter is not None: + namespaces.append(vars(converter)) + if importer is not None: + # An importer may still hold functions from a removed converter module. + for name in ("from_arrow_schema", "to_arrow_schema"): + # Later importers resolve these converters locally inside prepare_table. + if name not in vars(importer): + continue + namespace = getattr(getattr(importer, name, None), "__globals__", None) + if type(namespace) is not dict: + return None + namespaces.append(namespace) + for name, cls in ( + ("BooleanType", BooleanType), ("ByteType", ByteType), ("ShortType", ShortType), + ("IntegerType", IntegerType), ("LongType", LongType), + ("StructField", StructField), ("StructType", StructType), + ): + if any(namespace.get(name) is not cls for namespace in namespaces): + return None + for cls in (LongType, StructField, StructType): + if (cls.__getattribute__ is not object.__getattribute__ + or any("__class__" in vars(base) for base in cls.__mro__ if base is not object)): + return None + for cls, names, descriptor_type in ( + (StructType, ("fields", "names"), GetSetDescriptorType), + (StructType, ("__iter__",), WrapperDescriptorType), + (StructField, ("name", "dataType", "nullable", "metadata"), GetSetDescriptorType), + ): + for name in names: + descriptor = vars(cls).get(name) + if (type(descriptor) is not descriptor_type + or descriptor.__objclass__ is not cls or descriptor.__name__ != name): + return None + + names = list(frame.columns) + fields = None + if type(schema) is StructType: + fields = schema.fields + if len(fields) != width or not all( + type(field) is StructField and type(field.dataType) is LongType + and type(field.name) is str and "\x00" not in field.name + for field in fields + ): + return None + names = schema.names + if len(set(names)) != width: + return None + elif schema is not None: + if type(schema) not in (list, tuple) or len(schema) != width or not all( + type(name) is str and "\x00" not in name for name in schema + ): + return None + + columns = [frame.iloc[:, i] for i in range(width)] + for i, (column, dtype) in enumerate(zip(columns, dtypes)): + if type(column) is not Series: + return None + expected = IntegerArray if type(dtype) is Int64Dtype else NumpyExtensionArray + if type(column.array) is not expected: + return None + # The native row path rejects invalid nullability before any Config RPC. + if (fields is not None and not fields[i].nullable + and type(dtype) is Int64Dtype and column.hasnans): + return None + + import pyarrow as pa + + # Arrow consumes each Series' integer buffer and nullable mask independently. + return pa.Table.from_arrays( + [pa.array(column, type=pa.int64(), from_pandas=True) for column in columns], + names=names, + ) + + +def prepare_string_input(frame, schema): + """Prepare inline IPC for explicit columns of exact strings or None. + + Inspect every object before Arrow conversion: str subclasses can override + conversion hooks, and pandas treats several other objects as missing values. + Those cases must keep the existing row conversion and its error ordering. + """ + import sys + + import numpy as np + import pandas as pd + from pandas.core import frame as frame_module, generic as generic_module + from pandas.core.frame import DataFrame + from pandas.core.series import Series + from types import FunctionType, GetSetDescriptorType, MethodDescriptorType, WrapperDescriptorType + + from pyspark._pyspark import ( + BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, + DoubleType, DecimalType, StringType, StructField, StructType as NativeStructType, + ) + from pyspark.sql import types as sql_types + + # Keep small frames on the established row path. + if (pd.DataFrame is not DataFrame or type(frame) is not DataFrame + or type(schema) is not NativeStructType + or len(frame) < 10000 + or any(name in vars(frame) for name in ("astype", "notna", "where"))): + return None + # Preserve the direct normalization hooks used by the existing row path. + for name, module in (("astype", generic_module), ("notna", frame_module), + ("where", generic_module)): + method = getattr(DataFrame, name) + if (type(method) is not FunctionType or method.__name__ != name + or method.__module__ != module.__name__ or method.__globals__ is not vars(module)): + return None + if (DataFrame.astype is not generic_module.NDFrame.astype + or DataFrame.where is not generic_module.NDFrame.where): + return None + # Earlier converter branches must not recognize a string as another type. + # Check a loaded converter too: it keeps its own imported type aliases. + if sys.modules.get("pyspark.sql.types") is not sql_types: + return None + importer = sys.modules.get("pyspark.sql._arrow_input") + if importer is not None and getattr(importer, "StructType", None) is not NativeStructType: + return None + converter = sys.modules.get("pyspark.sql.pandas.types") + for name, cls in ( + ("BooleanType", BooleanType), ("ByteType", ByteType), ("ShortType", ShortType), + ("IntegerType", IntegerType), ("LongType", LongType), ("FloatType", FloatType), + ("DoubleType", DoubleType), ("DecimalType", DecimalType), ("StringType", StringType), + ("StructField", StructField), ("StructType", NativeStructType), + ): + if (getattr(sql_types, name, None) is not cls + or converter is not None and getattr(converter, name, None) is not cls): + return None + # The row path reads native fields directly. Class attributes are mutable, + # so customized Python schema access must not select a different Arrow path. + for name, cls in (("StringType", StringType), ("StructField", StructField), + ("StructType", NativeStructType)): + if (getattr(sql_types, name, None) is not cls + or cls.__getattribute__ is not object.__getattribute__ + or any("__class__" in vars(base) for base in cls.__mro__ if base is not object)): + return None + for cls, names, descriptor_type in ( + (NativeStructType, ("fields", "names"), GetSetDescriptorType), + (NativeStructType, ("__iter__",), WrapperDescriptorType), + (StructField, ("name", "dataType", "nullable", "metadata"), GetSetDescriptorType), + (StringType, ("isUTF8BinaryCollation",), MethodDescriptorType), + ): + for name in names: + descriptor = vars(cls).get(name) + if (type(descriptor) is not descriptor_type + or descriptor.__objclass__ is not cls or descriptor.__name__ != name): + return None + if not _has_canonical_columns(frame.columns): + return None + width = len(frame.columns) + if not width or len(schema.fields) != width or not frame.columns.is_unique: + return None + if not all(type(name) is str and "\x00" not in name for name in frame.columns): + return None + if not all( + type(field) is StructField and type(field.dataType) is StringType + and field.dataType.isUTF8BinaryCollation() + and type(field.name) is str and "\x00" not in field.name + for field in schema.fields + ) or len(set(schema.names)) != width: + return None + if not all(type(dtype) is type(np.dtype("object")) and dtype == np.dtype("object") + for dtype in frame.dtypes): + return None + # A custom ExtensionArray may report a genuine NumPy dtype while its + # Arrow hook emits different values from to_numpy(). Trust only pandas' + # concrete NumPy backing array, including its name in older pandas releases. + try: + from pandas.core.arrays.numpy_ import NumpyExtensionArray + except ImportError: + try: + from pandas.core.arrays.numpy_ import PandasArray as NumpyExtensionArray + except ImportError: + return None + columns = [frame.iloc[:, i] for i in range(width)] + if any(type(column) is not Series or type(column.array) is not NumpyExtensionArray + for column in columns): + return None + + from pyspark._pyspark import _pandas_string_ipc + + # Preflight exact object types before the native Arrow buffer traversal. + # Unsupported values, non-nullable None, and surrogates return to the row + # path before serialization. Keep its inline transport and configuration + # behavior rather than selecting the public Arrow import/upload policy. + arrays = [column.array._ndarray for column in columns] + if any(type(array) is not np.ndarray or array.ndim != 1 + or type(array.dtype) is not type(np.dtype("object")) + or array.dtype != np.dtype("object") or len(array) != len(frame) + for array in arrays): + return None + return _pandas_string_ipc(arrays, schema) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 459a65e1..f6f056f7 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -26,17 +26,19 @@ import subprocess from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from datetime import datetime, timezone +from datetime import date, datetime, timezone from decimal import Decimal from pathlib import Path from urllib.parse import quote import grpc +import numpy as np import pandas as pd import pyarrow as pa import pytest from google.protobuf import descriptor_pb2, descriptor_pool, message_factory -from pyspark.sql import SparkSession +from pyspark.sql import SparkSession, types as T +from pyspark.sql._pandas_input import prepare_int64_table @pytest.fixture(scope="module") @@ -403,6 +405,7 @@ def _arrow_input_service( protocol, *, upload_error=None, + config_error=False, user_id=None, user_agent=None, allow_file_artifacts=False, @@ -424,7 +427,8 @@ def _arrow_input_service( } config.update(overrides) state = {"cache": {}, "uploads": 0, "relations": [], "message_sizes": [], - "config_requests": [], "upload_contexts": [], "status_contexts": []} + "config_requests": [], "input_tables": [], + "upload_contexts": [], "status_contexts": []} server_id = "00000000-0000-0000-0000-000000000001" def attribution(request, context): @@ -446,6 +450,8 @@ def configure(raw, context): operation = request.operation.WhichOneof("op_type") keys = getattr(request.operation, operation).keys state["config_requests"].append((operation, list(keys))) + if config_error: + context.abort(grpc.StatusCode.UNAVAILABLE, "configuration unavailable") return protocol("ConfigResponse", session_id=request.session_id, pairs=[dict(key=k, value=config[k]) for k in keys if config.get(k) is not None] @@ -509,6 +515,9 @@ def execute(raw, context): state["column_names"] = list(relation.to_df.column_names) relation = relation.to_df.input state["relations"].append(relation) + # Column-name schemas wrap the local input with a rename relation. + if relation.HasField("to_df"): + relation = relation.to_df.input if relation.HasField("chunked_cached_local_relation"): cached = relation.chunked_cached_local_relation payloads = [state["cache"]["cache/" + h] for h in cached.dataHashes] @@ -523,6 +532,7 @@ def execute(raw, context): if relation.local_relation.HasField("schema") else None ) rows, first, last = 0, None, None + batches = [] state["arrow_schemas"], state["batch_rows"] = [], [] for payload in payloads: reader = pa.ipc.open_stream(payload) @@ -531,12 +541,14 @@ def execute(raw, context): state["batch_rows"].append(stream_rows) for batch in reader: stream_rows.append(batch.num_rows) + batches.append(batch) 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) + state["input_tables"].append(pa.Table.from_batches(batches)) result = pa.table({"rows": [rows]}) sink = pa.BufferOutputStream() with pa.ipc.new_stream(sink, result.schema) as writer: @@ -889,3 +901,1765 @@ def import_module(name, *args, **kwargs): spark.createDataFrame(data) assert imported == [] assert state["config_requests"] == [] + + +# Signed-int64 pandas input uses the same Arrow importer and mock service. + +@pytest.mark.parametrize("storage", ["object", "default", "python", "pyarrow", "python_nan", "pyarrow_nan"]) +@pytest.mark.parametrize("dtype", ["int64", "Int64"]) +def test_integer_canonical_columns_storage_stays_columnar(protocol, monkeypatch, storage, dtype): + from pyspark.sql._pandas_input import _has_canonical_columns + + values = [2**53 + 1, -(2**63), 2**63 - 1, 0] * 250 + if dtype == "Int64": + values[1] = None + frame = pd.DataFrame({"original": pd.Series(values, dtype=dtype)}) + if storage == "object": + frame.columns = pd.Index(["original"], dtype=object) + elif storage != "default": + kind = storage.removesuffix("_nan") + label_dtype = (pd.StringDtype(storage=kind, na_value=np.nan) + if storage.endswith("_nan") else pd.StringDtype(storage=kind)) + frame.columns = pd.Index(pd.array(["original"], dtype=label_dtype)) + assert _has_canonical_columns(frame.columns) + assert prepare_int64_table(frame, None).column(0).to_pylist() == values + + def no_rows(*args, **kwargs): + raise AssertionError("ordinary columns must retain integer columnar construction") + + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame).toArrow() + assert state["input_tables"][0].column(0).to_pylist() == values + assert state["config_requests"] + + +@pytest.mark.parametrize("dtype", ["int64", "Int64"]) +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +@pytest.mark.parametrize("location", [ + "subclass", "class", "owner", "backing_subclass_item", "backing_class_iter", + "backing_class_getitem", "backing_instance_box", +]) +@pytest.mark.parametrize("behavior", ["raise", "mutate"]) +def test_integer_columns_iterator_hooks_keep_native_order( + protocol, monkeypatch, dtype, schema_kind, location, behavior +): + from pandas.core.base import IndexOpsMixin + + values = [2**53 + 1, -(2**63), 2**63 - 1, 0] * 250 + if dtype == "Int64": + values[1] = None + frame = pd.DataFrame({"original": pd.Series(values, dtype=dtype)}) + schema = None + expected_name = "original" + nullable = True + metadata = {} + if schema_kind == "names": + schema = ["renamed"] + expected_name = "renamed" + elif schema_kind == "explicit": + nullable, metadata = dtype == "Int64", {"source": "native"} + schema = T.StructType([T.StructField("renamed", T.LongType(), nullable, metadata)]) + expected_name = "renamed" + original_iter = pd.Index.__iter__ + original_notna, original_astype = frame.notna, frame.astype + calls = [] + invoked = False + + class ColumnsError(RuntimeError): + pass + + def notna(*args, **kwargs): + calls.append("notna") + return original_notna(*args, **kwargs) + + def astype(*args, **kwargs): + calls.append("astype") + return original_astype(*args, **kwargs) + + def before_hook(): + nonlocal invoked + if invoked: + return + invoked = True + calls.append("columns iterator") + if behavior == "raise": + raise ColumnsError("columns hook precedes pandas normalization") + # Install spies only when native iteration reaches the columns hook. + # Before eligibility, normalization methods are the canonical methods. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + frame.iat[0, 0] = 2**53 + 7 + frame.iat[-1, 0] = None if dtype == "Int64" else -(2**63) + 1 + + def iterate(index): + if index is columns: + before_hook() + return original_iter(index) + + if location == "subclass": + class Columns(pd.Index): + pass + monkeypatch.setattr(Columns, "__iter__", iterate) + labels = pd.Index(["original"], dtype=object) + frame.columns = Columns._simple_new(labels._values, name=None) + assert type(frame.columns) is Columns + elif location == "backing_subclass_item": + # IndexOpsMixin iteration calls ndarray.item, rather than ndarray.__iter__. + class Names(np.ndarray): + def item(self, *args): + before_hook() + return super().item(*args) + backing = np.asarray(["original"], dtype=object).view(Names) + frame.columns = pd.Index._simple_new(backing, name=None) + assert type(frame.columns._values) is Names + elif location.startswith("backing_"): + frame.columns = pd.Index(pd.array(["original"], dtype=pd.StringDtype(storage="python"))) + backing = frame.columns._values + if location == "backing_class_iter": + original = type(backing).__iter__ + + def backing_iter(array): + if array is backing: + before_hook() + return original(array) + + monkeypatch.setattr(type(backing), "__iter__", backing_iter) + elif location == "backing_class_getitem": + original = type(backing).__getitem__ + + def backing_getitem(array, key): + if array is backing: + before_hook() + return original(array, key) + + monkeypatch.setattr(type(backing), "__getitem__", backing_getitem) + else: + original = backing._box_func + + def box(value): + before_hook() + return original(value) + + monkeypatch.setitem(vars(backing), "_box_func", box) + else: + monkeypatch.setattr(pd.Index if location == "class" else IndexOpsMixin, + "__iter__", iterate) + columns = frame.columns + + with _arrow_input_service(protocol) as (spark, state): + if behavior == "raise": + with pytest.raises(ColumnsError, + match="columns hook precedes pandas normalization"): + spark.createDataFrame(frame, schema) + assert calls == ["columns iterator"] + assert state["config_requests"] == [] + assert state["input_tables"] == [] + else: + spark.createDataFrame(frame, schema).toArrow() + expected = values.copy() + expected[0], expected[-1] = 2**53 + 7, (None if dtype == "Int64" else -(2**63) + 1) + table = state["input_tables"][0] + assert table.num_rows == 1000 + assert table.column_names == [expected_name] + assert table.column(0).type == pa.int64() + assert table.column(0).to_pylist() == expected + assert table.column(0).null_count == sum(value is None for value in expected) + assert table.schema.field(0).nullable is nullable + assert table.schema.field(0).metadata is None + import json + logical = json.loads(state["relations"][0].local_relation.schema) + assert logical == {"type": "struct", "fields": [{ + "name": expected_name, "type": "long", "nullable": nullable, + "metadata": metadata, + }]} + assert state["config_requests"] == [] + assert calls == ["columns iterator", "notna", "astype"] + + +def test_integer_rebound_index_values_keeps_native_error(protocol, monkeypatch): + frame = pd.DataFrame({"original": np.arange(1000, dtype=np.int64)}) + schema = T.StructType([T.StructField("renamed", T.LongType())]) + columns = frame.columns + original = pd.Index._values.fget + calls = [] + + def values(index): + if index is columns: + calls.append("_values") + raise RuntimeError("columns backing hook precedes normalization") + return original(index) + + monkeypatch.setattr(pd.Index, "_values", property(values)) + with _arrow_input_service(protocol) as (spark, state): + with pytest.raises(RuntimeError, + match="columns backing hook precedes normalization"): + spark.createDataFrame(frame, schema) + assert calls == ["_values"] + assert state["config_requests"] == [] + assert state["input_tables"] == [] + + +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +@pytest.mark.parametrize("verify", [True, False]) +def test_integer_columns_preserve_values_without_python_rows(protocol, monkeypatch, schema_kind, verify): + frame = pd.DataFrame({ + "a": pd.Series([2**53 + 1, None, -(2**63)], dtype="Int64"), + "b": np.array([2**63 - 1, -1, 0], dtype=np.int64), + }) + schema = None + names = ["a", "b"] + if schema_kind == "names": + schema = names = ["renamed_a", "renamed_b"] + elif schema_kind == "explicit": + schema = T.StructType([T.StructField("a", T.LongType()), + T.StructField("b", T.LongType(), False, {"source": "test"})]) + import pyspark.sql._arrow_input as helper + original = helper.prepare_table + observed = [] + + def capture(*args): + result = original(*args) + observed.append(result) + return result + + def no_rows(*args, **kwargs): + raise AssertionError("supported pandas input must not materialize Python rows") + + monkeypatch.setattr(helper, "prepare_table", capture) + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + df = spark.createDataFrame(frame, schema=schema, verifySchema=verify) + # Every nonempty pandas input applies its names through toDF, including + # explicit struct schemas, just as the current row fallback does. + assert not df.isLocal() + assert df.toArrow().column(0).to_pylist() == [3] + assert state["column_names"] == names + assert state["result"] == (3, 2**53 + 1, -(2**63)) + table, resolved, renamed = observed[0] + assert table.to_pydict() == {"a": [2**53 + 1, None, -(2**63)], + "b": [2**63 - 1, -1, 0]} + assert (renamed or resolved.names) == names + if schema_kind == "explicit": + assert resolved == schema + assert not table.schema.field(1).nullable + assert resolved.fields[1].metadata == {"source": "test"} + + +@pytest.mark.parametrize("dtype", ["int64", "Int64"]) +def test_empty_integer_frame_requires_explicit_schema(protocol, dtype): + frame = pd.DataFrame({"a": pd.Series([], dtype=dtype)}) + with _arrow_input_service(protocol) as (spark, state): + with pytest.raises(ValueError): + spark.createDataFrame(frame) + df = spark.createDataFrame(frame, T.StructType([T.StructField("a", T.LongType())])) + assert df.isLocal() + + +@pytest.mark.parametrize("values", [ + [1.0, float("nan")], [True, False], ["x", None], + [datetime(2024, 1, 1), None], + [datetime(2024, 1, 1, tzinfo=timezone.utc), None], + [date(2024, 1, 1), None], [Decimal("1.20"), None], + [[1, 2], None], [{"a": 1}, None], +]) +def test_other_dtypes_keep_row_fallback(values): + frame = pd.DataFrame({"a": values}) + assert prepare_int64_table(frame, None) is None + assert prepare_int64_table(frame, T.StructType([T.StructField("a", T.LongType())])) is None + + +@pytest.mark.parametrize("schema", [ + "a long", T.LongType(), ["a", "b"], [1], + T.StructType([T.StructField("a", T.IntegerType())]), + T.StructType([T.StructField("a", T.DoubleType())]), + T.StructType([]), +]) +def test_schema_conversions_keep_row_fallback(schema): + assert prepare_int64_table(pd.DataFrame({"a": [1, 2]}), schema) is None + + +def test_duplicate_names_fallback_and_nullable_masks(): + frame = pd.DataFrame([[2**53 + 1, 2**53 + 3]], columns=["x", "x"], dtype="int64") + assert prepare_int64_table(frame, None) is None + assert prepare_int64_table(pd.DataFrame({0: [1]}), None) is None + frame = pd.DataFrame({"a": pd.Series([None, None], dtype="Int64")}) + assert prepare_int64_table(frame, None).column(0).to_pylist() == [None, None] + + +@pytest.mark.parametrize("dtype", ["int32", "uint64", "Int32", "UInt64", "int64[pyarrow]"]) +def test_other_integer_dtypes_keep_row_fallback(dtype): + assert prepare_int64_table(pd.DataFrame({"a": pd.Series([1, 2], dtype=dtype)}), None) is None + + +@pytest.mark.parametrize("rows", [100, 1000, 100000, 1000000]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_integer_scales_reach_mock_service(protocol, monkeypatch, rows, explicit): + values = np.arange(rows, dtype=np.int64) + (2**53 + 1) + frame = pd.DataFrame({name: values for name in ("a", "b", "c")}) + schema = T.StructType([T.StructField(name, T.LongType()) for name in frame.columns]) + + def no_rows(*args, **kwargs): + raise AssertionError("integer input must stay columnar at every scale") + + if rows >= 1000: + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema if explicit else None).toArrow() + assert state["result"] == (rows, 2**53 + 1, 2**53 + rows) + if frame.memory_usage(index=False).sum() > 1048576: + assert state["relations"][0].HasField("chunked_cached_local_relation") + + +def test_nullable_inference_does_not_depend_on_first_value(protocol): + frame = pd.DataFrame({"a": pd.Series([None, 2**53 + 1], dtype="Int64")}) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame).toArrow() + assert state["result"] == (2, None, 2**53 + 1) + + +@pytest.mark.parametrize("verify", [True, False]) +@pytest.mark.parametrize("config_error", [False, True]) +def test_non_nullable_schema_is_preserved(protocol, verify, config_error): + frame = pd.DataFrame({"a": pd.Series([1, None], dtype="Int64")}) + schema = T.StructType([T.StructField("a", T.LongType(), False)]) + assert prepare_int64_table(frame, schema) is None + with _arrow_input_service(protocol, config_error=config_error) as (spark, state): + with pytest.raises(RuntimeError, match="null"): + spark.createDataFrame(frame, schema, verifySchema=verify) + assert not state["relations"] + assert not state["uploads"] + assert state["config_requests"] == [] + + +def test_small_numpy_frames_avoid_importer_configuration_overhead(): + assert prepare_int64_table(pd.DataFrame({"a": np.arange(999, dtype=np.int64)}), None) is None + assert prepare_int64_table(pd.DataFrame({"a": np.arange(1000, dtype=np.int64)}), None) is not None + assert prepare_int64_table(pd.DataFrame({"a": pd.Series([None], dtype="Int64")}), None) is not None + + +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +def test_nul_names_keep_lossless_ipc_row_fallback(protocol, schema_kind): + import pyarrow as pa + + name = "a\x00b" + frame = pd.DataFrame({name if schema_kind == "inferred" else "original": + np.arange(1000, dtype=np.int64)}) + schema = None + if schema_kind == "names": + schema = [name] + elif schema_kind == "explicit": + schema = T.StructType([T.StructField(name, T.LongType())]) + # C Stream field names are NUL-terminated; IPC preserves the full string. + assert prepare_int64_table(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["result"] == (1000, 0, 999) + payload = state["relations"][0].local_relation.data + assert pa.ipc.open_stream(payload).schema.names == [name] + + +@pytest.mark.parametrize("nullable", [False, True]) +@pytest.mark.parametrize("rows", [999, 1000]) +def test_integer_explicit_renamed_schema_preserves_all_ipc_cells(protocol, monkeypatch, nullable, rows): + values = ([2**53 + 1, -(2**63), 2**63 - 1, 0] * 250)[:rows] + other = [None if nullable and i % 7 == 0 else value + for i, value in enumerate(reversed(values))] + frame = pd.DataFrame({"original": np.asarray(values, dtype="int64"), + "second": pd.Series(other, dtype="Int64" if nullable else "int64")}) + schema = T.StructType([ + T.StructField("renamed", T.LongType(), False, {"source": "native", "id": 3}), + T.StructField("nullable", T.LongType(), True, {"tag": [1, 2]}), + ]) + + def no_rows(*args, **kwargs): + raise AssertionError("trusted integer columns must stay columnar") + + if rows >= 1000 or nullable: + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + table = state["input_tables"][0] + assert table.column_names == ["renamed", "nullable"] + assert [column.to_pylist() for column in table.columns] == [values, other] + assert table.schema.field(0).type == pa.int64() + assert table.schema.field(0).nullable is False + assert table.schema.field(1).nullable is True + if rows >= 1000 or nullable: + from pyspark.sql.pandas.types import from_arrow_schema + assert from_arrow_schema(table.schema).json() == schema.json(), ( + from_arrow_schema(table.schema).json(), schema.json(), table.schema) + assert state["config_requests"] + else: + assert state["config_requests"] == [] + + +def test_integer_dataframe_and_name_subclasses_keep_fallback(monkeypatch): + class Frame(pd.DataFrame): + pass + + class Name(str): + pass + + class Names(list): + pass + + frame = pd.DataFrame({"a": np.arange(1000, dtype="int64")}) + assert prepare_int64_table(Frame(frame), None) is None + assert prepare_int64_table(frame, Names(["a"])) is None + assert prepare_int64_table(frame, [Name("a")]) is None + # Pandas 3 normalizes constructor labels to plain strings by default. + # An object Index keeps the subclass this fallback regression exercises. + named_frame = frame.copy() + named_frame.columns = pd.Index([Name("a")], dtype=object) + assert type(named_frame.columns[0]) is Name + assert prepare_int64_table(named_frame, None) is None + monkeypatch.setattr(pd, "DataFrame", Frame) + assert prepare_int64_table(Frame(frame), None) is None + + +def test_integer_missing_internal_pandas_array_class_keeps_fallback(): + import sys + + script = r''' +import sys +import types +import numpy as np +import pandas as pd +from pyspark.sql._pandas_input import prepare_int64_table +frame = pd.DataFrame({"a": np.arange(1000, dtype="int64")}) +nullable = pd.DataFrame({"a": pd.Series([None, 1], dtype="Int64")}) +# Older pandas exposes PandasArray rather than NumpyExtensionArray. Missing +# optional internals must disable the optimization without breaking pandas input. +sys.modules["pandas.core.arrays.numpy_"] = types.ModuleType("pandas.core.arrays.numpy_") +assert prepare_int64_table(frame, None) is None +assert prepare_int64_table(nullable, None) is None +''' + result = subprocess.run([sys.executable, "-c", script], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize("class_spoof", [False, True]) +def test_integer_custom_extension_dtype_cannot_select_arrow(protocol, class_spoof): + from pandas.api.extensions import ExtensionArray, ExtensionDtype, take + + arrow_calls = [] + + class Int64EqualDtype(ExtensionDtype): + name = "int64_equal" + type = int + kind = "i" + na_value = None + + @classmethod + def construct_array_type(cls): + return DifferentArrowArray + + def __eq__(self, other): + return (isinstance(other, Int64EqualDtype) + or isinstance(other, np.dtype) and other == np.dtype("int64")) + + def __hash__(self): + return hash(self.name) + + class DifferentArrowArray(ExtensionArray): + def __init__(self, values): + self._values = np.asarray(values, dtype=object) + + @property + def dtype(self): + return Int64EqualDtype() + + @property + def nbytes(self): + return self._values.nbytes + + def __len__(self): + return len(self._values) + + def __getitem__(self, key): + result = self._values[key] + return type(self)(result) if isinstance(result, np.ndarray) else result + + def isna(self): + return pd.isna(self._values) + + def take(self, indices, allow_fill=False, fill_value=None): + return type(self)(take(self._values, indices, allow_fill=allow_fill, + fill_value=fill_value)) + + def copy(self): + return type(self)(self._values.copy()) + + @classmethod + def _from_sequence(cls, scalars, *, dtype=None, copy=False): + values = np.asarray(scalars, dtype=object) + return cls(values.copy() if copy else values) + + def to_numpy(self, dtype=None, copy=False, na_value=None): + if dtype is not None: + return self._values.astype(dtype, copy=copy) + return self._values.copy() if copy else self._values + + def __arrow_array__(self, type=None): + arrow_calls.append("arrow") + return pa.array([-7] * len(self), type=type) + + values = [2**53 + 1, None, -(2**63), 0] * 250 + frame = pd.DataFrame({"a": pd.Series(DifferentArrowArray(values), copy=False)}) + if class_spoof: + Int64EqualDtype.__getattribute__ = lambda self, name: ( + type(np.dtype("int64")) if name == "__class__" + else object.__getattribute__(self, name)) + assert frame.dtypes.iloc[0] == np.dtype("int64") + assert isinstance(frame.dtypes.iloc[0], np.dtype) is class_spoof + schema = T.StructType([T.StructField("renamed", T.LongType())]) + assert prepare_int64_table(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + assert state["input_tables"][0].column(0).to_pylist() == values + assert arrow_calls == [] + + +@pytest.mark.parametrize("custom", ["dtype", "array"]) +def test_integer_nullable_dtype_and_array_subclasses_keep_fallback(protocol, custom): + from pandas.core.arrays.integer import IntegerArray, Int64Dtype + + arrow_calls = [] + + class CustomDtype(Int64Dtype): + pass + + class CustomArray(IntegerArray): + @property + def dtype(self): + if custom == "dtype": + return CustomDtype() + return Int64Dtype() + + def __arrow_array__(self, type=None): + arrow_calls.append("arrow") + return pa.array([-7] * len(self), type=type) + + values = [2**53 + 1, 0, -(2**63), 7] * 250 + array = CustomArray(np.asarray(values, dtype="int64"), np.zeros(1000, dtype=bool)) + frame = pd.DataFrame({"a": pd.Series(array, copy=False)}) + assert type(frame["a"].array) is CustomArray + schema = T.StructType([T.StructField("a", T.LongType())]) + assert prepare_int64_table(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + assert state["input_tables"][0].column(0).to_pylist() == values + assert arrow_calls == [] + + +def _integer_guard_process(mode, argument, before_import=False): + """Class mutation and converter imports must be isolated from other tests.""" + import sys + + if mode == "captured_alias": + import pyspark.sql._arrow_input as importer + if not hasattr(importer, "to_arrow_schema"): + pytest.skip("this importer resolves converters locally and holds no captured aliases") + + script = r''' +import importlib +import importlib.util +import sys +import tempfile +from pathlib import Path +import numpy as np +import pandas as pd +from pyspark.sql import types as T +from pyspark import _pyspark as native + +spec = importlib.util.spec_from_file_location("integer_guard_capture", sys.argv[4]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +values = [2**53 + 1, -(2**63), 2**63 - 1, 0] * 250 +frame = pd.DataFrame({"renamed": np.asarray(values, dtype="int64")}) +schema = T.StructType([T.StructField("renamed", T.LongType(), False, {"source": "native"})]) +mode, argument, before_import = sys.argv[1:4] +if before_import == "True": + sys.modules.pop("pyspark.sql._pandas_input", None) +if mode == "descriptor": + class_name, attribute = argument.split(":") + cls = getattr(native, class_name) + if attribute == "alias": + setattr(T, class_name, object) + elif attribute == "__getattribute__": + setattr(cls, attribute, lambda self, name: object.__getattribute__(self, name)) + elif attribute == "__class__": + setattr(cls, "__getattribute__", lambda self, name: ( + object if name == "__class__" else object.__getattribute__(self, name))) + elif attribute == "__iter__": + setattr(cls, attribute, lambda self: iter([T.StructField("changed", T.BooleanType())])) + else: + value = {"names": ["changed"], "fields": [], "name": "changed", + "dataType": T.BooleanType(), "nullable": True, + "metadata": {"changed": True}}[attribute] + setattr(cls, attribute, property(lambda self: value)) +elif mode in ("public_alias", "converter_alias", "captured_alias"): + replacement = native.BooleanType if argument == "LongType" else native.LongType + if mode == "public_alias": + sys.modules.pop("pyspark.sql.pandas.types", None) + assert "pyspark.sql.pandas.types" not in sys.modules + setattr(T, argument, replacement) + else: + converter = importlib.import_module("pyspark.sql.pandas.types") + if mode == "captured_alias": + importer = importlib.import_module("pyspark.sql._arrow_input") + setattr(converter, argument, replacement) + if mode == "captured_alias": + sys.modules.pop("pyspark.sql.pandas.types") + current = importlib.import_module("pyspark.sql.pandas.types") + assert current is not converter + assert importer.to_arrow_schema.__globals__ is vars(converter) +elif mode == "atomic_alias": + scope, replacement_name = argument.split(":") + replacement = None if replacement_name == "none" else native.LongType + if scope == "public_before": + sys.modules.pop("pyspark.sql._arrow_input", None) + T.AtomicType = replacement + elif scope == "public_after": + importlib.import_module("pyspark.sql._arrow_input") + T.AtomicType = replacement + else: + importer = importlib.import_module("pyspark.sql._arrow_input") + if replacement_name == "missing": + del importer.AtomicType + else: + importer.AtomicType = replacement +elif mode == "importer_alias": + importer = importlib.import_module("pyspark.sql._arrow_input") + importer.StructType = None if argument == "none" else native.LongType +else: + raise AssertionError(mode) +class Factory: + def mktemp(self, name): + return Path(tempfile.mkdtemp(prefix=name)) +with capture._arrow_input_service(capture.protocol.__wrapped__(Factory())) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [], state["config_requests"] + table = state["input_tables"][0] + assert table.schema.names == ["renamed"], table.schema + assert table.column(0).type == __import__("pyarrow").int64() + assert table.column(0).to_pylist() == values +''' + result = subprocess.run([sys.executable, "-c", script, mode, argument, + str(before_import), __file__], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize("class_name,attribute", [ + ("StructType", "fields"), ("StructType", "names"), ("StructType", "__iter__"), + ("StructField", "name"), ("StructField", "dataType"), ("StructField", "nullable"), + ("StructField", "metadata"), ("StructType", "__getattribute__"), + ("StructField", "__getattribute__"), ("LongType", "__getattribute__"), + ("StructType", "__class__"), ("StructField", "__class__"), ("LongType", "__class__"), + ("StructType", "alias"), ("StructField", "alias"), ("LongType", "alias"), +]) +@pytest.mark.parametrize("before_import", [False, True]) +def test_integer_custom_schema_descriptors_keep_native_row_path(class_name, attribute, before_import): + _integer_guard_process("descriptor", f"{class_name}:{attribute}", before_import) + + +@pytest.mark.parametrize("class_name", ["BooleanType", "ByteType", "ShortType", "IntegerType", "LongType"]) +@pytest.mark.parametrize("mode", ["public_alias", "converter_alias", "captured_alias"]) +def test_integer_rebound_converter_aliases_keep_native_row_path(class_name, mode): + _integer_guard_process(mode, class_name) + + +@pytest.mark.parametrize("replacement", ["none", "long"]) +def test_integer_rebound_importer_schema_alias_keeps_native_row_path(replacement): + _integer_guard_process("importer_alias", replacement) + + +@pytest.mark.parametrize("scope", ["public_before", "public_after", "importer"]) +@pytest.mark.parametrize("replacement", ["none", "long"]) +def test_integer_rebound_atomic_alias_keeps_native_row_path(scope, replacement): + _integer_guard_process("atomic_alias", f"{scope}:{replacement}") + + +def test_integer_deleted_importer_atomic_alias_keeps_native_row_path(): + _integer_guard_process("atomic_alias", "importer:missing") + + +def test_integer_warm_importer_without_global_converter_aliases(protocol, monkeypatch): + import pyspark.sql._arrow_input as importer + + original = importer.prepare_table + + def local_imports(table, schema, *configs): + # Newer importers import converters within prepare_table. Exercise that + # module shape even on the pinned PR118 dependency used by this branch. + from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema + table, resolved, names = original(table, schema, *configs) + assert from_arrow_schema(table.schema).fields + assert to_arrow_schema(resolved).names == table.column_names + return table, resolved, names + + frame = pd.DataFrame({"a": np.arange(1000, dtype="int64")}) + schema = T.StructType([T.StructField("renamed", T.LongType())]) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(pa.table({"warm": [1, 2]})).toArrow() + # Keep original's module globals available in the dependency importer; + # the full-stack validation separately uses the actual local-import helper. + import types + substitute = types.ModuleType(importer.__name__) + substitute.StructType = T.StructType + substitute.AtomicType = T.AtomicType + substitute.prepare_table = local_imports + monkeypatch.setitem(__import__("sys").modules, importer.__name__, substitute) + + def no_rows(*args, **kwargs): + raise AssertionError("warm integer input must stay columnar") + + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + for _ in range(3): + spark.createDataFrame(frame, schema).toArrow() + for table in state["input_tables"][1:]: + assert table.column_names == ["renamed"] + assert table.column(0).to_pylist() == list(range(1000)) + + + + + + + +@pytest.mark.parametrize("dtype", ["int64", "Int64"]) +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +def test_integer_rebound_index_data_keeps_native_read_order( + protocol, monkeypatch, dtype, schema_kind +): + import json + + values = [2**53 + 1, -(2**63), 2**63 - 1, 0] * 250 + if dtype == "Int64": + values[1] = None + frame = pd.DataFrame({"original": pd.Series(values, dtype=dtype)}) + schema, expected_name, nullable, metadata = None, "original", True, {} + if schema_kind == "names": + schema, expected_name = ["renamed"], "renamed" + elif schema_kind == "explicit": + expected_name, nullable, metadata = "renamed", dtype == "Int64", {"source": "native"} + schema = T.StructType([T.StructField(expected_name, T.LongType(), nullable, metadata)]) + columns, backing = frame.columns, vars(frame.columns)["_data"] + original_notna, original_astype = frame.notna, frame.astype + calls = [] + normalizing = False + allowed_reads = None + + class ColumnsDataError(RuntimeError): + pass + + def notna(*args, **kwargs): + nonlocal normalizing + normalizing = True + calls.append("notna") + return original_notna(*args, **kwargs) + + def astype(*args, **kwargs): + calls.append("astype") + return original_astype(*args, **kwargs) + + def data(index): + if index is columns: + if not normalizing: + calls.append("_data") + if allowed_reads is not None and calls.count("_data") > allowed_reads: + raise ColumnsDataError("eligibility added a columns storage read") + # Install spies only when the native columns storage hook runs. + # Canonical frame methods remain intact before eligibility. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + return backing + return vars(index)["_data"] + + def set_data(index, value): + # Preserve ordinary storage initialization for indexes pandas creates + # while normalizing values; only reads of the input columns are hooked. + vars(index)["_data"] = value + + monkeypatch.setattr(pd.Index, "_data", property(data, set_data), raising=False) + import pyspark.sql._pandas_input as helper + + with _arrow_input_service(protocol) as (spark, state): + # Record current master's fallback rather than assuming the old tolist + # path's single read. Index iteration reads NumPy-backed _values three + # times, and extension-backed _values twice, before normalization. + with monkeypatch.context() as baseline: + baseline.setattr(helper, "prepare_int64_table", lambda *_: None) + spark.createDataFrame(frame, schema).toArrow() + baseline_calls = calls.copy() + allowed_reads = baseline_calls.index("notna") + assert allowed_reads > 0 + assert baseline_calls == ["_data"] * allowed_reads + ["notna", "astype"] + calls.clear() + normalizing = False + # Restore canonical methods before eligibility, just as for the first + # call. The storage hook installs them when native iteration starts. + monkeypatch.delitem(vars(frame), "notna") + monkeypatch.delitem(vars(frame), "astype") + spark.createDataFrame(frame, schema).toArrow() + assert calls == baseline_calls + assert not state["config_requests"] + assert len(state["input_tables"]) == 2 + table = state["input_tables"][1] + assert table.equals(state["input_tables"][0], check_metadata=True) + assert table.column(0).to_pylist() == values + assert table.column(0).null_count == sum(value is None for value in values) + assert table.schema.equals( + pa.schema([pa.field(expected_name, pa.int64(), nullable=nullable)]), + check_metadata=True, + ) + assert json.loads(state["relations"][0].local_relation.schema) == { + "type": "struct", "fields": [ + {"name": expected_name, "type": "long", "nullable": nullable, + "metadata": metadata} + ], + } + + +@pytest.mark.parametrize("location", [ + "subclass", "class", "owner", "backing_subclass_item", "backing_class_iter", + "backing_class_getitem", "backing_instance_box", "arrow_class_iter", +]) +@pytest.mark.parametrize("behavior", ["raise", "mutate"]) +def test_string_columns_iterator_hooks_keep_native_order( + protocol, monkeypatch, location, behavior +): + from pandas.core.base import IndexOpsMixin + + values = ["plain"] * 10000 + frame = pd.DataFrame({"original": values}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.StringType())]) + original_iter = pd.Index.__iter__ + original_notna, original_astype = frame.notna, frame.astype + calls = [] + invoked = False + + class ColumnsError(RuntimeError): + pass + + def notna(*args, **kwargs): + calls.append("notna") + return original_notna(*args, **kwargs) + + def astype(*args, **kwargs): + calls.append("astype") + return original_astype(*args, **kwargs) + + def before_hook(): + nonlocal invoked + if invoked: + return + invoked = True + calls.append("columns iterator") + if behavior == "raise": + raise ColumnsError("columns hook precedes pandas normalization") + # Install spies only when native iteration reaches the columns hook. + # Before eligibility, normalization methods are the canonical methods. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + frame.iat[0, 0], frame.iat[-1, 0] = "changed", None + + def iterate(index): + if index is columns: + before_hook() + return original_iter(index) + + if location == "subclass": + class Columns(pd.Index): + pass + monkeypatch.setattr(Columns, "__iter__", iterate) + labels = pd.Index(["original"], dtype=object) + frame.columns = Columns._simple_new(labels._values, name=None) + assert type(frame.columns) is Columns + elif location == "backing_subclass_item": + # IndexOpsMixin iteration uses ndarray.item, not ndarray.__iter__. + class Names(np.ndarray): + def item(self, *args): + before_hook() + return super().item(*args) + backing = np.asarray(["original"], dtype=object).view(Names) + frame.columns = pd.Index._simple_new(backing, name=None) + assert type(frame.columns._values) is Names + elif location.startswith(("backing_", "arrow_")): + storage = "pyarrow" if location.startswith("arrow_") else "python" + frame.columns = pd.Index(pd.array(["original"], dtype=pd.StringDtype(storage=storage))) + backing = frame.columns._values + if location.endswith("class_iter"): + original = type(backing).__iter__ + + def backing_iter(array): + if array is backing: + before_hook() + return original(array) + + monkeypatch.setattr(type(backing), "__iter__", backing_iter) + elif location.endswith("class_getitem"): + original = type(backing).__getitem__ + + def backing_getitem(array, key): + if array is backing: + before_hook() + return original(array, key) + + monkeypatch.setattr(type(backing), "__getitem__", backing_getitem) + else: + original = backing._box_func + + def box(value): + before_hook() + return original(value) + + monkeypatch.setitem(vars(backing), "_box_func", box) + else: + monkeypatch.setattr(pd.Index if location == "class" else IndexOpsMixin, + "__iter__", iterate) + columns = frame.columns + + with _arrow_input_service(protocol) as (spark, state): + if behavior == "raise": + with pytest.raises(ColumnsError, + match="columns hook precedes pandas normalization"): + spark.createDataFrame(frame, schema) + assert calls == ["columns iterator"] + assert state["config_requests"] == [] + assert state["input_tables"] == [] + assert state["uploads"] == 0 + else: + spark.createDataFrame(frame, schema).toArrow() + expected = values.copy() + expected[0], expected[-1] = "changed", None + import json + + table = state["input_tables"][0] + assert table.schema.names == ["renamed"] + assert table.column(0).type == pa.string() + assert table.column(0).to_pylist() == expected + assert table.column(0).is_null().to_pylist() == [False] * 9999 + [True] + assert table.schema.equals(pa.schema([pa.field("renamed", pa.string())]), + check_metadata=True) + assert json.loads(state["relations"][0].local_relation.schema) == json.loads(schema.json()) + assert state["config_requests"] == [] + assert state["uploads"] == 0 + assert calls == ["columns iterator", "notna", "astype"] + + + +def test_string_rebound_index_values_keeps_native_error(protocol, monkeypatch): + + frame = pd.DataFrame({"original": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.StringType())]) + columns = frame.columns + original = pd.Index._values.fget + calls = [] + + def values(index): + if index is columns: + calls.append("_values") + raise RuntimeError("columns backing hook precedes normalization") + return original(index) + + monkeypatch.setattr(pd.Index, "_values", property(values)) + assert calls == [] + with _arrow_input_service(protocol) as (spark, state): + with pytest.raises(RuntimeError, + match="columns backing hook precedes normalization"): + spark.createDataFrame(frame, schema) + assert calls == ["_values"] + assert state["config_requests"] == [] + assert state["input_tables"] == [] + + + +@pytest.mark.parametrize("nullable", [False, True]) +def test_string_rebound_index_data_keeps_native_read_order(protocol, monkeypatch, nullable): + import json + + values = ["plain", "é", "漢字", ""] * 2500 + if nullable: + values[1] = None + frame = pd.DataFrame({"original": values}, dtype=object) + schema = T.StructType([ + T.StructField("renamed", T.StringType(), nullable, {"source": "native"}) + ]) + columns, backing = frame.columns, vars(frame.columns)["_data"] + original_notna, original_astype = frame.notna, frame.astype + calls = [] + normalizing = False + allowed_reads = None + + class ColumnsDataError(RuntimeError): + pass + + def notna(*args, **kwargs): + nonlocal normalizing + normalizing = True + calls.append("notna") + return original_notna(*args, **kwargs) + + def astype(*args, **kwargs): + calls.append("astype") + return original_astype(*args, **kwargs) + + def data(index): + if index is columns: + if not normalizing: + calls.append("_data") + if allowed_reads is not None and calls.count("_data") > allowed_reads: + raise ColumnsDataError("eligibility added a columns storage read") + # Canonical normalization remains intact before eligibility. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + return backing + return vars(index)["_data"] + + def set_data(index, value): + # Keep ordinary storage initialization for indexes created by pandas. + vars(index)["_data"] = value + + monkeypatch.setattr(pd.Index, "_data", property(data, set_data), raising=False) + import pyspark.sql._pandas_input as helper + + with _arrow_input_service(protocol) as (spark, state): + # Measure the current native fallback's storage reads, rather than the + # former tolist path's single-read contract. + with monkeypatch.context() as baseline: + baseline.setattr(helper, "prepare_string_input", lambda *_: None) + spark.createDataFrame(frame, schema).toArrow() + baseline_calls = calls.copy() + allowed_reads = baseline_calls.index("notna") + assert allowed_reads > 0 + assert baseline_calls == ["_data"] * allowed_reads + ["notna", "astype"] + calls.clear() + normalizing = False + # Restore canonical methods before the candidate eligibility check; + # native label iteration will install the spies through the storage hook. + monkeypatch.delitem(vars(frame), "notna") + monkeypatch.delitem(vars(frame), "astype") + spark.createDataFrame(frame, schema).toArrow() + assert calls == baseline_calls + assert not state["config_requests"] + assert state["uploads"] == 0 + assert len(state["input_tables"]) == 2 + assert len(state["relations"]) == 2 + table = state["input_tables"][1] + assert table.equals(state["input_tables"][0], check_metadata=True) + assert table.column(0).to_pylist() == values + assert table.column(0).is_null().to_pylist() == [value is None for value in values] + assert table.schema.equals( + pa.schema([pa.field("renamed", pa.string(), nullable=nullable)]), + check_metadata=True, + ) + assert json.loads(state["relations"][1].local_relation.schema) == { + "type": "struct", "fields": [ + {"name": "renamed", "type": "string", "nullable": nullable, + "metadata": {"source": "native"}} + ], + } + + + +@pytest.mark.parametrize("kind", [ + "object", "default", "python", "pyarrow", "python_numpy", "pyarrow_numpy", +]) +def test_string_canonical_column_storage_stays_eligible(kind): + from pyspark.sql._pandas_input import _has_canonical_columns, prepare_string_input + + frame = pd.DataFrame({"original": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.StringType())]) + if kind == "object": + frame.columns = pd.Index(np.asarray(["original"], dtype=object), dtype=object) + elif kind == "default": + frame.columns = pd.Index(["original"]) + elif kind.endswith("_numpy"): + from pandas.core.arrays import string_ as strings, string_arrow as arrow_strings + + module = strings if kind.startswith("python") else arrow_strings + name = ("StringArrayNumpySemantics" if kind.startswith("python") + else "ArrowStringArrayNumpySemantics") + cls = getattr(module, name, None) + if cls is None: + pytest.skip("older infer-string storage class is absent") + frame.columns = pd.Index._simple_new(cls._from_sequence(["original"]), name=None) + else: + frame.columns = pd.Index(pd.array(["original"], dtype=pd.StringDtype(storage=kind))) + assert _has_canonical_columns(frame.columns) + assert type(prepare_string_input(frame, schema)) is bytes + + + +@pytest.mark.parametrize("missing", ["module", "owner", "base"]) +def test_string_absent_private_columns_storage_keeps_row_fallback(monkeypatch, missing): + import builtins + from pandas.core.arrays import numpy_ as numpy_module, string_ as string_module + from pyspark.sql._pandas_input import _has_canonical_columns + + columns = pd.Index(pd.array(["original"], dtype=pd.StringDtype(storage="python"))) + if missing == "module": + original = builtins.__import__ + + def without_arrow(name, *args, **kwargs): + if name == "pandas.core.arrays.arrow": + raise ImportError("private storage module is absent") + return original(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", without_arrow) + elif missing == "owner": + monkeypatch.delattr(numpy_module, "NumpyExtensionArray") + else: + monkeypatch.delattr(string_module, "BaseStringArray") + assert not _has_canonical_columns(columns) + + + +@pytest.mark.parametrize("cached", [False, True]) +@pytest.mark.parametrize("verify", [False, True]) +@pytest.mark.parametrize("large", ["false", "true"]) +def test_string_columns_preserve_every_ipc_cell(protocol, monkeypatch, cached, verify, large): + import json + + from pyspark.sql.pandas.types import from_arrow_metadata + + strings = ["", "plain", "é", "e\u0301", "\x00embedded\x00", "漢字", "🙂", + "\U0010ffff", "\uffff", "\u202eright-to-left"] * 1000 + nullable = [None if i % 7 == 0 else value for i, value in enumerate(reversed(strings))] + expected = [strings, nullable, [None] * len(strings)] + frame = pd.DataFrame(dict(zip(["original", "b", "c"], expected)), dtype=object) + schema = T.StructType([ + T.StructField("renamed.`text", T.StringType(), False, {"source": "string", "id": 3}), + T.StructField("nullable", T.StringType(), True, {"tag": [1, 2]}), + T.StructField("all_null", T.StringType()), + ]) + + def no_rows(*args, **kwargs): + raise AssertionError("string fast path must not materialize Python rows") + + for method in ("to_numpy", "iterrows", "itertuples"): + monkeypatch.setattr(pd.DataFrame, method, no_rows) + monkeypatch.setattr(pd.DataFrame, "values", property(no_rows)) + monkeypatch.setattr(pd.Series, "tolist", no_rows) + overrides = { + "spark.sql.session.localRelationCacheThreshold": "1" if cached else "999999999", + "spark.sql.execution.arrow.useLargeVarTypes": large, + "spark.sql.session.localRelationChunkSizeRows": "997", + } + with _arrow_input_service(protocol, **overrides) as (spark, state): + result = spark.createDataFrame(frame, schema, verifySchema=verify) + assert not result.isLocal() + result.toArrow() + assert state["column_names"] == schema.names + table = state["input_tables"][0] + assert table.num_rows == len(strings) + assert table.schema.names == schema.names + for column, values, field in zip(table.columns, expected, schema.fields): + assert column.type == pa.string() + assert column.to_pylist() == values + actual_field = table.schema.field(field.name) + assert actual_field.nullable == field.nullable + assert from_arrow_metadata(actual_field.metadata) is None + spark_field = json.loads(state["relations"][0].local_relation.schema)["fields"][ + schema.names.index(field.name)] + assert spark_field["metadata"] == field.metadata + assert state["config_requests"] == [] + assert state["uploads"] == 0 + + + +@pytest.mark.parametrize("reason", [ + "small", "empty", "inferred", "names", "ddl", "cast", "atomic", "mixed_schema", + "width", "duplicate_input", "duplicate_schema", "input_nul", "schema_nul", + "input_number", "datetime", "date", "nan", "pd_na", "nat", "mixed", "bytes", + "nonnullable_null", "frame_subclass", "string_subclass", "string_dtype", + "arrow_string_dtype", "late_unsupported", "surrogate", "late_surrogate", "collation", +]) +def test_string_column_guard_preserves_row_fallback(reason): + from pyspark.sql._pandas_input import prepare_string_input + + class Frame(pd.DataFrame): + pass + + class Text(str): + def __str__(self): + raise AssertionError("guard must not call conversion hooks") + + class Unsupported: + def __str__(self): + raise AssertionError("guard must not stringify unsupported objects") + + values = ["plain"] * 10000 + frame = pd.DataFrame({"text": values}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + if reason == "small": + frame = frame.iloc[:9999] + elif reason == "empty": + frame = frame.iloc[:0] + elif reason in ("inferred", "names", "ddl", "cast", "atomic"): + schema = {"inferred": None, "names": ["text"], "ddl": "text STRING", + "cast": T.StructType([T.StructField("text", T.LongType())]), + "atomic": T.StringType()}[reason] + elif reason == "mixed_schema": + frame["long"] = np.arange(10000) + schema.add("long", T.LongType()) + elif reason == "width": + schema.add("extra", T.StringType()) + elif reason in ("duplicate_input", "duplicate_schema"): + frame["other"] = values + schema.add("other", T.StringType()) + if reason == "duplicate_input": + frame.columns = ["text", "text"] + else: + schema = T.StructType([T.StructField("text", T.StringType())] * 2) + elif reason in ("input_nul", "input_number"): + frame.columns = ["text\x00" if reason == "input_nul" else 1] + elif reason == "schema_nul": + schema = T.StructType([T.StructField("text\x00", T.StringType())]) + elif reason == "frame_subclass": + frame = Frame(frame) + elif reason in ("string_dtype", "arrow_string_dtype"): + frame = pd.DataFrame({"text": pd.Series( + values, dtype=pd.StringDtype(storage="pyarrow" if reason == "arrow_string_dtype" + else "python"))}) + elif reason == "collation": + schema = T.StructType([T.StructField("text", T.StringType("UTF8_LCASE"))]) + else: + value = { + "datetime": datetime(2020, 1, 1), "date": date(2020, 1, 1), + "nan": float("nan"), "pd_na": pd.NA, "nat": pd.NaT, "mixed": 42, + "bytes": b"plain", "nonnullable_null": None, "string_subclass": Text("plain"), + "late_unsupported": Unsupported(), "surrogate": "\ud800", "late_surrogate": "\udfff", + }[reason] + if reason == "nonnullable_null": + schema = T.StructType([T.StructField("text", T.StringType(), False)]) + if reason in ("late_unsupported", "late_surrogate"): + # pandas 3 infers StringDtype for a newly assigned list column. + # Keep these unsupported-object fixtures on concrete object storage. + frame["other"] = pd.Series(values, dtype=object) + schema.add("other", T.StringType()) + backing = frame["other"].array._ndarray + assert type(backing) is np.ndarray and backing.dtype == np.dtype("object") + frame.iat[-1, -1] = value + if reason in ("late_unsupported", "late_surrogate"): + assert frame.iat[-1, -1] is value + assert prepare_string_input(frame, schema) is None + + + +@pytest.mark.parametrize("reason", [ + "string_subclass", "frame_subclass", "nan", "pd_na", "nat", "string_dtype", + "nonnullable_null", "inferred", "names", "ddl", "small", "collation", "cast", "mixed", +]) +def test_string_public_fallback_keeps_row_semantics(protocol, monkeypatch, reason): + calls = [] + + class Text(str): + def __str__(self): + calls.append("str") + return "changed" + + class Frame(pd.DataFrame): + pass + + values = ["plain"] * (9999 if reason == "small" else 10000) + if reason in ("nan", "pd_na", "nat", "nonnullable_null", "mixed"): + values[-1] = {"nan": float("nan"), "pd_na": pd.NA, "nat": pd.NaT, + "nonnullable_null": None, "mixed": 42}[reason] + elif reason == "string_subclass": + values[-1] = Text("plain") + frame_type = Frame if reason == "frame_subclass" else pd.DataFrame + frame = frame_type({"text": values}, dtype="string" if reason == "string_dtype" else object) + schema = T.StructType([T.StructField("text", T.StringType(), reason != "nonnullable_null")]) + if reason == "inferred": + schema = None + elif reason == "names": + schema = ["renamed"] + elif reason == "ddl": + schema = "text STRING" + elif reason == "collation": + schema = T.StructType([T.StructField("text", T.StringType("UTF8_LCASE"))]) + elif reason == "cast": + schema = T.StructType([T.StructField("text", T.LongType())]) + astype = pd.DataFrame.astype + + def record_rows(self, *args, **kwargs): + calls.append("astype") + return astype(self, *args, **kwargs) + + monkeypatch.setattr(pd.DataFrame, "astype", record_rows) + with _arrow_input_service(protocol) as (spark, state): + if reason in ("nonnullable_null", "cast", "mixed"): + from pyspark.errors import PySparkRuntimeError + + message = ("Column 'text' is declared as non-nullable but contains null values" + if reason == "nonnullable_null" else "Type mismatch in row data") + with pytest.raises(PySparkRuntimeError, match=message): + spark.createDataFrame(frame, schema) + assert calls == ["astype"] + assert state["config_requests"] == [] + assert state["relations"] == [] + return + spark.createDataFrame(frame, schema).toArrow() + assert calls == ["astype"] # Native extraction keeps str.__str__ overrides uncalled. + table = state["input_tables"][0] + assert table.schema.names == ["renamed" if reason == "names" else "text"] + expected = ["plain"] * len(values) + if reason in ("nan", "pd_na", "nat"): + expected[-1] = None + assert table.column(0).to_pylist() == expected + assert state["config_requests"] == [] + + + +@pytest.mark.parametrize("surrogate", ["\ud800", "\udfff"]) +def test_string_surrogate_fallback_precedes_importer_configuration(protocol, monkeypatch, surrogate): + values = ["plain"] * 10000 + values[-1] = surrogate + frame = pd.DataFrame({"text": values}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + astype = pd.DataFrame.astype + calls = [] + + def record_rows(self, *args, **kwargs): + calls.append("astype") + return astype(self, *args, **kwargs) + + monkeypatch.setattr(pd.DataFrame, "astype", record_rows) + with _arrow_input_service(protocol) as (spark, state): + state["config_error"] = grpc.StatusCode.INVALID_ARGUMENT + with pytest.raises(TypeError, match="unsupported value type for createDataFrame: str"): + spark.createDataFrame(frame, schema) + assert calls == ["astype"] + assert state["config_requests"] == [] + assert state["relations"] == [] + assert state["uploads"] == 0 + + + +@pytest.mark.parametrize("class_name,attribute", [ + ("StructType", "names"), ("StructType", "fields"), ("StructType", "__iter__"), + ("StructField", "name"), ("StructField", "dataType"), ("StructField", "nullable"), + ("StructField", "metadata"), ("StringType", "isUTF8BinaryCollation"), + ("StructType", "__getattribute__"), ("StructField", "__getattribute__"), + ("StringType", "__getattribute__"), ("StructType", "__class__"), + ("StructField", "__class__"), ("StringType", "__class__"), + ("StructType", "alias"), ("StructField", "alias"), ("StringType", "alias"), +]) +@pytest.mark.parametrize("before_import", [False, True]) +def test_string_custom_schema_descriptors_keep_native_row_path( + protocol, class_name, attribute, before_import +): + import sys + + script = r''' +import sys +import tempfile +from pathlib import Path +import pandas as pd +from pyspark.sql import types as T + +import importlib.util +spec = importlib.util.spec_from_file_location("string_schema_capture", sys.argv[4]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +_arrow_input_service, protocol = capture._arrow_input_service, capture.protocol +frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) +schema = T.StructType([T.StructField("text", T.StringType(), True, {"source": "native"})]) +class_name, attribute, before_import = sys.argv[1:4] +if before_import == "True": + # The capture module imports prepare_int64_table; force the public call to + # import a new helper after the class or public alias has been customized. + sys.modules.pop("pyspark.sql._pandas_input", None) +assert ("pyspark.sql._pandas_input" in sys.modules) == (before_import == "False") +cls = getattr(T, class_name) +if attribute == "alias": + setattr(T, class_name, object) +else: + if attribute == "__getattribute__": + replacement = lambda self, name: object.__getattribute__(self, name) + elif attribute == "__iter__": + replacement = lambda self: iter([T.StructField("changed", T.StringType())]) + elif attribute == "__class__": + attribute = "__getattribute__" + replacement = lambda self, name: object if name == "__class__" else ( + object.__getattribute__(self, name)) + elif attribute == "isUTF8BinaryCollation": + replacement = lambda self: False + else: + value = {"names": ["changed"], "fields": [], "name": "changed", + "dataType": T.StringType(), "nullable": False, "metadata": {"changed": True}}[attribute] + replacement = property(lambda self: value) + setattr(cls, attribute, replacement) +class Factory: + def mktemp(self, name): + return Path(tempfile.mkdtemp(prefix=name)) +with _arrow_input_service(protocol.__wrapped__(Factory())) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [], state["config_requests"] + table = state["input_tables"][0] + assert table.schema.names == ["text"], table.schema + assert table.column(0).to_pylist() == ["plain"] * 10000 +''' + result = subprocess.run([sys.executable, "-c", script, class_name, attribute, + str(before_import), __file__], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + + +def test_string_rebound_pandas_alias_keeps_subclass_fallback(monkeypatch): + from pyspark.sql._pandas_input import prepare_string_input + + class Frame(pd.DataFrame): + pass + + frame = Frame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + monkeypatch.setattr(pd, "DataFrame", Frame) + assert prepare_string_input(frame, schema) is None + + + +@pytest.mark.parametrize("class_name", [ + "BooleanType", "ByteType", "ShortType", "IntegerType", "LongType", + "FloatType", "DoubleType", "DecimalType", +]) +@pytest.mark.parametrize("converter_loaded", [False, True]) +def test_string_rebound_preceding_converter_alias_keeps_native_row_path( + protocol, class_name, converter_loaded +): + import sys + + script = r''' +import importlib.util +import sys +import tempfile +from pathlib import Path +import pandas as pd +from pyspark.sql import types as T + +spec = importlib.util.spec_from_file_location("string_alias_capture", sys.argv[3]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +_arrow_input_service, protocol = capture._arrow_input_service, capture.protocol +values = ["plain", "", "\x00", "é", "漢字", "🙂", None, "tail"] * 1250 +frame = pd.DataFrame({"text": values}, dtype=object) +schema = T.StructType([T.StructField("text", T.StringType())]) +class_name, converter_loaded = sys.argv[1:3] +if converter_loaded == "True": + import pyspark.sql.pandas.types as converter + setattr(converter, class_name, T.StringType) +else: + sys.modules.pop("pyspark.sql.pandas.types", None) + setattr(T, class_name, T.StringType) + assert "pyspark.sql.pandas.types" not in sys.modules +class Factory: + def mktemp(self, name): + return Path(tempfile.mkdtemp(prefix=name)) +with _arrow_input_service(protocol.__wrapped__(Factory())) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [], state["config_requests"] + table = state["input_tables"][0] + assert table.schema.names == ["text"], table.schema + assert table.column(0).to_pylist() == values +''' + result = subprocess.run([sys.executable, "-c", script, class_name, + str(converter_loaded), __file__], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + + +@pytest.mark.parametrize("replacement", ["string", "none"]) +def test_string_rebound_importer_schema_alias_keeps_native_row_path(protocol, replacement): + import sys + + script = r''' +import importlib.util +import json +import sys +import tempfile +from pathlib import Path +import pandas as pd +from pyspark.sql import types as T + +spec = importlib.util.spec_from_file_location("string_importer_capture", sys.argv[2]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +_arrow_input_service, protocol = capture._arrow_input_service, capture.protocol +strings = ["plain", "", "\x00", "é", "漢字", "🙂", "e\u0301", "tail"] * 1250 +nullable = [None if i % 7 == 0 else value for i, value in enumerate(reversed(strings))] +expected = [strings, nullable, [None] * len(strings)] +frame = pd.DataFrame(dict(zip(["original", "b", "c"], expected)), dtype=object) +schema = T.StructType([ + T.StructField("text", T.StringType(), False, {"source": "native", "id": 3}), + T.StructField("nullable", T.StringType(), True, {"tag": [1, 2]}), + T.StructField("all_null", T.StringType()), +]) +import pyspark.sql._arrow_input as importer +importer.StructType = T.StringType if sys.argv[1] == "string" else None +class Factory: + def mktemp(self, name): + return Path(tempfile.mkdtemp(prefix=name)) +with _arrow_input_service(protocol.__wrapped__(Factory())) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [], state["config_requests"] + table = state["input_tables"][0] + assert table.schema.names == schema.names, table.schema + spark_fields = json.loads(state["relations"][0].local_relation.schema)["fields"] + for column, values, field, spark_field in zip(table.columns, expected, schema.fields, spark_fields): + assert column.to_pylist() == values + actual_field = table.schema.field(field.name) + assert actual_field.nullable == field.nullable + assert actual_field.metadata is None + assert spark_field["metadata"] == field.metadata +''' + result = subprocess.run([sys.executable, "-c", script, replacement, __file__], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + + +@pytest.mark.parametrize("dtype_mode", ["equal", "class_spoof", "exact_numpy"]) +def test_string_custom_extension_dtype_cannot_impersonate_numpy_object(protocol, dtype_mode): + from pandas.api.extensions import ExtensionArray, ExtensionDtype, take + from pyspark.sql._pandas_input import prepare_string_input + + arrow_calls = [] + + class ObjectEqualDtype(ExtensionDtype): + name = "object_equal_string" + type = str + kind = "O" + na_value = None + + @classmethod + def construct_array_type(cls): + return DifferentArrowArray + + def __eq__(self, other): + return (isinstance(other, ObjectEqualDtype) + or isinstance(other, np.dtype) and other.kind == "O") + + def __hash__(self): + return hash(self.name) + + class DifferentArrowArray(ExtensionArray): + def __init__(self, values): + self._values = np.asarray(values, dtype=object) + + @property + def dtype(self): + return ObjectEqualDtype() + + @property + def nbytes(self): + return self._values.nbytes + + def __len__(self): + return len(self._values) + + def __getitem__(self, key): + result = self._values[key] + return type(self)(result) if isinstance(result, np.ndarray) else result + + def isna(self): + return pd.isna(self._values) + + def take(self, indices, allow_fill=False, fill_value=None): + return type(self)(take(self._values, indices, allow_fill=allow_fill, + fill_value=fill_value)) + + def copy(self): + return type(self)(self._values.copy()) + + @classmethod + def _from_sequence(cls, scalars, *, dtype=None, copy=False): + values = np.asarray(scalars, dtype=object) + return cls(values.copy() if copy else values) + + def to_numpy(self, dtype=None, copy=False, na_value=None): + if dtype is not None: + return self._values.astype(dtype, copy=copy) + return self._values.copy() if copy else self._values + + def __arrow_array__(self, type=None): + arrow_calls.append("arrow") + # A scan of to_numpy() cannot establish the values this hook emits. + return pa.array(["different"] * len(self), type=type) + + array = DifferentArrowArray(["plain"] * 10000) + frame = pd.DataFrame({"text": pd.Series(array, copy=False)}) + schema = T.StructType([T.StructField("text", T.StringType())]) + if dtype_mode == "exact_numpy": + # An ExtensionArray may change its advertised dtype after construction. + # The dtype and NumPy scan both pass while its Arrow hook changes cells. + DifferentArrowArray.dtype = property(lambda self: np.dtype("object")) + DifferentArrowArray.__array__ = lambda self, dtype=None, copy=None: self._values + elif dtype_mode == "class_spoof": + def spoof_class(self, name): + if name == "__class__": + return type(np.dtype("object")) + return object.__getattribute__(self, name) + + ObjectEqualDtype.__getattribute__ = spoof_class + dtype = frame.dtypes.iloc[0] + assert type(frame) is pd.DataFrame + assert type(frame.iloc[:, 0].array) is DifferentArrowArray + assert type(dtype) is (type(np.dtype("object")) if dtype_mode == "exact_numpy" + else ObjectEqualDtype) + assert isinstance(dtype, np.dtype) is (dtype_mode != "equal") + assert dtype == np.dtype("object") + assert prepare_string_input(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + if dtype_mode == "exact_numpy": + # This deliberately inconsistent EA already fails pandas' row-path + # astype(object). Preserve that error before importer configuration. + with pytest.raises(TypeError, match="cannot set astype"): + spark.createDataFrame(frame, schema) + assert state["input_tables"] == [] + else: + spark.createDataFrame(frame, schema).toArrow() + assert state["input_tables"][0].column(0).to_pylist() == ["plain"] * 10000 + assert state["config_requests"] == [] + assert arrow_calls == [] + + + +@pytest.mark.parametrize("method", ["astype", "notna"]) +def test_string_instance_conversion_hooks_keep_row_fallback(protocol, method): + frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + original = getattr(frame, method) + calls = [] + expected = ["plain"] * 10000 + if method == "astype": + expected[0] = "changed" + + def customized(*args, **kwargs): + calls.append(method) + result = original(*args, **kwargs) + result.iat[0, 0] = "changed" + return result + else: + expected[0] = None + + def customized(*args, **kwargs): + calls.append(method) + result = original(*args, **kwargs) + result.iat[0, 0] = False + return result + setattr(frame, method, customized) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert calls == [method] + assert state["config_requests"] == [] + assert state["uploads"] == 0 + assert state["input_tables"][0].column(0).to_pylist() == expected + + + +@pytest.mark.parametrize("method", ["astype", "notna", "where"]) +@pytest.mark.parametrize("before_import", [False, True]) +def test_string_class_conversion_hooks_keep_row_fallback( + protocol, monkeypatch, method, before_import +): + import importlib + import sys + + name = "pyspark.sql._pandas_input" + if before_import: + monkeypatch.delitem(sys.modules, name, raising=False) + else: + importlib.import_module(name) + frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + original = getattr(pd.DataFrame, method) + calls = [] + expected = ["plain"] * 10000 + expected[0] = None if method == "notna" else "changed" + + def customized(self, *args, **kwargs): + calls.append(method) + result = original(self, *args, **kwargs) + result.iat[0, 0] = False if method == "notna" else "changed" + return result + + monkeypatch.setattr(pd.DataFrame, method, customized) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert calls == [method] + assert state["config_requests"] == [] + assert state["uploads"] == 0 + assert state["input_tables"][0].column(0).to_pylist() == expected + + + +def test_string_instance_where_override_keeps_existing_row_behavior(protocol): + from pyspark.sql._pandas_input import prepare_string_input + + frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + + def customized(*args, **kwargs): + raise AssertionError("row normalization calls where on the astype result") + + frame.where = customized + assert prepare_string_input(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + assert state["input_tables"][0].column(0).to_pylist() == ["plain"] * 10000 + + + +def test_string_native_backing_does_not_add_series_conversion_hooks(protocol, monkeypatch): + frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + + def customized(*args, **kwargs): + raise AssertionError("native string input must not call Series.to_numpy") + + monkeypatch.setattr(pd.Series, "to_numpy", customized) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + assert state["input_tables"][0].column(0).to_pylist() == ["plain"] * 10000 + + + +def test_string_large_inline_input_keeps_row_transport_and_schema(protocol): + values = ["漢字-🙂-" + str(i) for i in range(100000)] + frame = pd.DataFrame({"text": values}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.StringType(), False, + {"source": "inline-string", "size": len(values)})]) + with _arrow_input_service(protocol, **{ + "spark.sql.session.localRelationCacheThreshold": "1", + "spark.sql.session.localRelationSizeLimit": "1", + "spark.sql.execution.arrow.useLargeVarTypes": "true", + }) as (spark, state): + result = spark.createDataFrame(frame, schema) + assert not result.isLocal() + result.toArrow() + assert state["column_names"] == schema.names + assert state["config_requests"] == [] + assert state["uploads"] == 0 + table = state["input_tables"][0] + assert table.schema.names == ["renamed"] + assert table.column(0).type == pa.string() + assert table.schema.field(0).nullable is False + assert table.schema.field(0).metadata is None + assert table.column(0).to_pylist() == values + assert __import__("json").loads(state["relations"][0].local_relation.schema) == ( + __import__("json").loads(schema.json()))