This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 38b465de5d00 [SPARK-58838][PYTHON][TEST] Add tests for
pa.RecordBatch.from_pandas
38b465de5d00 is described below
commit 38b465de5d00f1557032ac093d980d4cb81625a7
Author: Spenser Sun <[email protected]>
AuthorDate: Thu Aug 20 16:10:55 2026 +0800
[SPARK-58838][PYTHON][TEST] Add tests for pa.RecordBatch.from_pandas
### What changes were proposed in this pull request?
Adds a golden-file test for `pa.RecordBatch.from_pandas` in a new file
`test_pyarrow_dataframe_from_pandas.py`. Per-column type inference is already
covered by `test_pyarrow_array_from_pandas_*`, so this pins the DataFrame-level
behavior:
- **`num_rows` for a 0-column DataFrame** — the headline case. A 0-column
frame still has N rows, and `RecordBatch.from_pandas` (with the default
`preserve_index=None`) is what carries that count, via the pandas index
metadata. The golden records `num_rows` explicitly because a 0-column batch has
no column to imply it.
- **The pandas index under `preserve_index` (`None` / `False` / `True`)** —
a RangeIndex is dropped while a named or other index becomes a column, and on a
0-column frame `False` collapses `num_rows` to 0.
- **Column assembly** over a small dtype sample, including nulls and empty
frames.
Structure:
- A shared base class (`_PyArrowFromPandasFrameTestBase`, no tests of its
own) holds the pandas-DataFrame inputs; the `RecordBatch` class subclasses it.
The `Table` / `Schema` `from_pandas` tests will land in this file and reuse the
base.
- `goldenutils.py`: `repr_value` can now renders a `pa.RecordBatch` with
the new `repr_arrow_record_batch_value` branch, mirroring the existing
`pa.Table` support.
### Why are the changes needed?
A drift canary under SPARK-54936: `createDataFrame` uses
`RecordBatch.from_pandas` for the 0-column case (`pandas/conversion.py:1026`,
`connect/session.py:632`) and relies on the behavior above, so an upstream
change would silently drop rows.
### Does this PR introduce _any_ user-facing change?
No. Test-only.
### How was this patch tested?
New golden-file test (no Spark session needed). Swept pyarrow 18–25 ×
pandas 2/3 — all 16 combinations pass, no version overrides needed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #58073 from Spenserrrr/recordbatch-from-pandas-tests.
Authored-by: Spenser Sun <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit afce937016fb3ae8e227a6d525cc07cbf325db8a)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
dev/sparktestsupport/modules.py | 1 +
python/pyspark/testing/goldenutils.py | 59 +++++--
.../golden_pyarrow_record_batch_from_pandas.csv | 11 ++
.../golden_pyarrow_record_batch_from_pandas.md | 12 ++
.../pyarrow/test_pyarrow_dataframe_from_pandas.py | 179 +++++++++++++++++++++
5 files changed, 247 insertions(+), 15 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index d097ad556531..425879ea48e8 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -542,6 +542,7 @@ pyspark_core = Module(
"pyspark.tests.upstream.pyarrow.test_pyarrow_array_type_inference",
"pyspark.tests.upstream.pyarrow.test_pyarrow_arrow_to_pandas_default",
"pyspark.tests.upstream.pyarrow.test_pyarrow_arrow_to_pandas_non_default",
+ "pyspark.tests.upstream.pyarrow.test_pyarrow_dataframe_from_pandas",
"pyspark.tests.upstream.pyarrow.test_pyarrow_ignore_timezone",
"pyspark.tests.upstream.pyarrow.test_pyarrow_scalar_type_coercion",
"pyspark.tests.upstream.pyarrow.test_pyarrow_scalar_type_inference",
diff --git a/python/pyspark/testing/goldenutils.py
b/python/pyspark/testing/goldenutils.py
index 0499bbe31857..8374d20eb422 100644
--- a/python/pyspark/testing/goldenutils.py
+++ b/python/pyspark/testing/goldenutils.py
@@ -15,7 +15,7 @@
# limitations under the License.
#
-from typing import Any, Callable, List, Optional
+from typing import Any, Callable, List, Optional, Union
import inspect
import os
import time
@@ -332,7 +332,9 @@ class GoldenFileTestMixin:
raise
@classmethod
- def repr_arrow_value(cls, value: Any, max_len: int = 32) -> str:
+ def repr_arrow_value(
+ cls, value: Union["pa.Array", "pa.ChunkedArray"], max_len: int = 32
+ ) -> str:
"""
Format a PyArrow Array/ChunkedArray for golden file.
@@ -359,7 +361,23 @@ class GoldenFileTestMixin:
return f"{v_str}@{cls.repr_type(value.type)}"
@classmethod
- def repr_arrow_table_value(cls, value: Any, max_len: int = 32) -> str:
+ def _repr_arrow_columns(
+ cls, value: Union["pa.Table", "pa.RecordBatch"], max_len: int
+ ) -> "tuple[str, str]":
+ """Render a Table/RecordBatch as a "{name: [scalars], ...}" body and
schema string."""
+ columns = []
+ for name, column in zip(value.schema.names, value.columns):
+ # Escape NULL bytes so the value can be safely stored in CSV files.
+ elements = [cls._scalar_str(scalar) for scalar in column]
+ columns.append(f"{name}: [" + ", ".join(elements) + "]")
+ v_str = "{" + ", ".join(columns) + "}"
+ if max_len > 0:
+ v_str = v_str[:max_len]
+ schema = ", ".join(f"{f.name}: {cls.repr_type(f.type)}" for f in
value.schema)
+ return v_str, schema
+
+ @classmethod
+ def repr_arrow_table_value(cls, value: "pa.Table", max_len: int = 32) ->
str:
"""
Format a PyArrow Table for golden file.
@@ -371,19 +389,27 @@ class GoldenFileTestMixin:
str
"{col: [val1, val2, None], ...}@Table[name: type, ...]"
"""
- columns = []
- for name, column in zip(value.column_names, value.columns):
- # Escape NULL bytes so the value can be safely stored in CSV files.
- elements = [cls._scalar_str(scalar) for scalar in column]
- columns.append(f"{name}: [" + ", ".join(elements) + "]")
- v_str = "{" + ", ".join(columns) + "}"
- if max_len > 0:
- v_str = v_str[:max_len]
- schema = ", ".join(f"{f.name}: {cls.repr_type(f.type)}" for f in
value.schema)
+ v_str, schema = cls._repr_arrow_columns(value, max_len)
return f"{v_str}@Table[{schema}]"
@classmethod
- def repr_pandas_value(cls, value: Any, max_len: int = 32) -> str:
+ def repr_arrow_record_batch_value(cls, value: "pa.RecordBatch", max_len:
int = 32) -> str:
+ """
+ Format a PyArrow RecordBatch for golden file.
+
+ Same shape as ``repr_arrow_table_value`` (a RecordBatch is a single
batch of
+ columns), keyed by column name, plus the Arrow schema.
+
+ Returns
+ -------
+ str
+ "{col: [val1, val2, None], ...}@RecordBatch[name: type, ...]"
+ """
+ v_str, schema = cls._repr_arrow_columns(value, max_len)
+ return f"{v_str}@RecordBatch[{schema}]"
+
+ @classmethod
+ def repr_pandas_value(cls, value: "pd.DataFrame", max_len: int = 32) ->
str:
"""
Format a pandas DataFrame for golden file.
@@ -412,7 +438,7 @@ class GoldenFileTestMixin:
return f"{v_str}@Dataframe[{simple_schema}]"
@classmethod
- def repr_numpy_value(cls, value: Any, max_len: int = 32) -> str:
+ def repr_numpy_value(cls, value: "np.ndarray", max_len: int = 32) -> str:
"""
Format a numpy ndarray for golden file.
@@ -456,6 +482,7 @@ class GoldenFileTestMixin:
- PyArrow Array/ChunkedArray -> repr_arrow_value
- PyArrow Table -> repr_arrow_table_value
+ - PyArrow RecordBatch -> repr_arrow_record_batch_value
- pandas DataFrame -> repr_pandas_value
- numpy ndarray -> repr_numpy_value
- Everything else -> repr_python_value
@@ -476,6 +503,8 @@ class GoldenFileTestMixin:
return cls.repr_arrow_value(value, max_len)
if have_pyarrow and isinstance(value, pa.Table):
return cls.repr_arrow_table_value(value, max_len)
+ if have_pyarrow and isinstance(value, pa.RecordBatch):
+ return cls.repr_arrow_record_batch_value(value, max_len)
if have_pandas and isinstance(value, pd.DataFrame):
return cls.repr_pandas_value(value, max_len)
@@ -487,7 +516,7 @@ class GoldenFileTestMixin:
return cls.repr_python_value(value, max_len)
@classmethod
- def repr_pandas_series_value(cls, value: Any, max_len: int = 32) -> str:
+ def repr_pandas_series_value(cls, value: "pd.Series", max_len: int = 32)
-> str:
"""
Format a pandas Series for golden file.
diff --git
a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.csv
b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.csv
new file mode 100644
index 000000000000..b918763436da
--- /dev/null
+++
b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.csv
@@ -0,0 +1,11 @@
+test case pandas dataframe preserve_index=None
preserve_index=False preserve_index=True
+0-columns:range-index {}@Dataframe[][index=RangeIndex[0:3:1]]
{}@RecordBatch[][num_rows=3] {}@RecordBatch[][num_rows=0]
{__index_level_0__: [0, 1, 2]}@RecordBatch[__index_level_0__: int64][num_rows=3]
+0-columns:named-index {}@Dataframe[][index='idx':[100, 200, 300]] {idx:
[100, 200, 300]}@RecordBatch[idx: int64][num_rows=3]
{}@RecordBatch[][num_rows=0] {idx: [100, 200, 300]}@RecordBatch[idx:
int64][num_rows=3]
+0-columns:unnamed-index {}@Dataframe[][index=None:[10, 20, 30]]
{__index_level_0__: [10, 20, 30]}@RecordBatch[__index_level_0__:
int64][num_rows=3] {}@RecordBatch[][num_rows=0] {__index_level_0__: [10,
20, 30]}@RecordBatch[__index_level_0__: int64][num_rows=3]
+0-columns:empty {}@Dataframe[][index=RangeIndex[0:0:1]]
{}@RecordBatch[][num_rows=0] {}@RecordBatch[][num_rows=0]
{__index_level_0__: []}@RecordBatch[__index_level_0__: int64][num_rows=0]
+single-column:range-index {'a': [1, 2, 3]}@Dataframe[a
int64][index=RangeIndex[0:3:1]] {a: [1, 2, 3]}@RecordBatch[a:
int64][num_rows=3] {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3]
{a: [1, 2, 3], __index_level_0__: [0, 1, 2]}@RecordBatch[a: int64,
__index_level_0__: int64][num_rows=3]
+single-column:named-index {'a': [1, 2, 3]}@Dataframe[a
int64][index='idx':[100, 200, 300]] {a: [1, 2, 3], idx: [100, 200,
300]}@RecordBatch[a: int64, idx: int64][num_rows=3] {a: [1, 2,
3]}@RecordBatch[a: int64][num_rows=3] {a: [1, 2, 3], idx: [100, 200,
300]}@RecordBatch[a: int64, idx: int64][num_rows=3]
+single-column:unnamed-index {'a': [1, 2, 3]}@Dataframe[a
int64][index=None:[10, 20, 30]] {a: [1, 2, 3], __index_level_0__: [10, 20,
30]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] {a: [1, 2,
3]}@RecordBatch[a: int64][num_rows=3] {a: [1, 2, 3], __index_level_0__:
[10, 20, 30]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3]
+multi-column:standard {'i': [1, 2, 3], 'f': [1.5, 2.5, 3.5], 'b': [True,
False, True], 's': ['a', 'b', 'c'], 't': [Timestamp('2020-01-01 05:30:00'),
Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01
05:30:00')]}@Dataframe[i int64, f float64, b bool, s object, t
datetime64[ns]][index=RangeIndex[0:3:1]] {i: [1, 2, 3], f: [1.5, 2.5, 3.5], b:
[True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01
05:30:00, 2020-01-01 05:30:00]}@RecordBatch[i: int64, f: float64, b: bool, [...]
+multi-column:nullable {'f': [1.5, nan, 3.5], 'b': [True, None, False], 's':
['a', None, 'c'], 't': [Timestamp('2020-01-01 05:30:00'), NaT,
Timestamp('2020-01-01 05:30:00')]}@Dataframe[f float64, b object, s object, t
datetime64[ns]][index=RangeIndex[0:3:1]] {f: [1.5, None, 3.5], b: [True,
None, False], s: [a, None, c], t: [2020-01-01 05:30:00, None, 2020-01-01
05:30:00]}@RecordBatch[f: float64, b: bool, s: string, t:
timestamp[ns]][num_rows=3] {f: [1.5, None, 3.5], b: [True, None, False] [...]
+multi-column:no-rows {'i': [], 'f': [], 'b': [], 't': []}@Dataframe[i int64,
f float64, b bool, t datetime64[ns]][index=RangeIndex[0:0:1]] {i: [], f: [],
b: [], t: []}@RecordBatch[i: int64, f: float64, b: bool, t:
timestamp[ns]][num_rows=0] {i: [], f: [], b: [], t: []}@RecordBatch[i: int64,
f: float64, b: bool, t: timestamp[ns]][num_rows=0] {i: [], f: [], b: [], t:
[], __index_level_0__: []}@RecordBatch[i: int64, f: float64, b: bool, t:
timestamp[ns], __index_level_0__: int64][num_rows=0]
diff --git
a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.md
b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.md
new file mode 100644
index 000000000000..bc1bc210c536
--- /dev/null
+++
b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.md
@@ -0,0 +1,12 @@
+| test case | pandas dataframe
| preserve_index=None
[...]
+|-----------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------
[...]
+| 0-columns:range-index | {}@Dataframe[][index=RangeIndex[0:3:1]]
| {}@RecordBatch[][num_rows=3]
[...]
+| 0-columns:named-index | {}@Dataframe[][index='idx':[100, 200, 300]]
| {idx: [100, 200, 300]}@RecordBatch[idx: int64][num_rows=3]
[...]
+| 0-columns:unnamed-index | {}@Dataframe[][index=None:[10, 20, 30]]
| {__index_level_0__: [10, 20, 30]}@RecordBatch[__index_level_0__:
int64][num_rows=3]
[...]
+| 0-columns:empty | {}@Dataframe[][index=RangeIndex[0:0:1]]
| {}@RecordBatch[][num_rows=0]
[...]
+| single-column:range-index | {'a': [1, 2, 3]}@Dataframe[a
int64][index=RangeIndex[0:3:1]]
| {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3]
[...]
+| single-column:named-index | {'a': [1, 2, 3]}@Dataframe[a
int64][index='idx':[100, 200, 300]]
| {a: [1, 2, 3], idx: [100, 200, 300]}@RecordBatch[a: int64,
idx: int64][num_rows=3]
[...]
+| single-column:unnamed-index | {'a': [1, 2, 3]}@Dataframe[a
int64][index=None:[10, 20, 30]]
| {a: [1, 2, 3], __index_level_0__: [10, 20,
30]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3]
[...]
+| multi-column:standard | {'i': [1, 2, 3], 'f': [1.5, 2.5, 3.5], 'b':
[True, False, True], 's': ['a', 'b', 'c'], 't': [Timestamp('2020-01-01
05:30:00'), Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01
05:30:00')]}@Dataframe[i int64, f float64, b bool, s object, t
datetime64[ns]][index=RangeIndex[0:3:1]] | {i: [1, 2, 3], f: [1.5, 2.5, 3.5],
b: [True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01
05:30:00, 2020-01-01 05:30:00]}@RecordBatch[i: int64, f: float [...]
+| multi-column:nullable | {'f': [1.5, nan, 3.5], 'b': [True, None,
False], 's': ['a', None, 'c'], 't': [Timestamp('2020-01-01 05:30:00'), NaT,
Timestamp('2020-01-01 05:30:00')]}@Dataframe[f float64, b object, s object, t
datetime64[ns]][index=RangeIndex[0:3:1]]
| {f: [1.5, None, 3.5], b: [True, None, False], s: [a, None, c], t:
[2020-01-01 05:30:00, None, 2020-01-01 05:30:00]}@RecordBatch[f: float64, b:
bool, s: string, t: timestam [...]
+| multi-column:no-rows | {'i': [], 'f': [], 'b': [], 't':
[]}@Dataframe[i int64, f float64, b bool, t
datetime64[ns]][index=RangeIndex[0:0:1]]
| {i: [], f: [], b: [], t:
[]}@RecordBatch[i: int64, f: float64, b: bool, t: timestamp[ns]][num_rows=0]
[...]
\ No newline at end of file
diff --git
a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_dataframe_from_pandas.py
b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_dataframe_from_pandas.py
new file mode 100644
index 000000000000..0b7af403fa07
--- /dev/null
+++
b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_dataframe_from_pandas.py
@@ -0,0 +1,179 @@
+#
+# 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.
+#
+
+"""
+Golden-file tests for the PyArrow ``from_pandas`` constructors that take a
whole pandas
+DataFrame: ``pa.RecordBatch.from_pandas`` (with ``pa.Table.from_pandas`` /
+``pa.Schema.from_pandas`` to follow in this file). These take a DataFrame,
unlike
+``pa.Array.from_pandas`` which takes a Series (covered by
test_pyarrow_array_from_pandas_*).
+
+Per-column type inference matches the Array tests, so these pin the
DataFrame-level
+behavior instead: whole-frame assembly, the pandas index under
``preserve_index``, and --
+for RecordBatch -- num_rows preservation for a 0-column DataFrame. Spark calls
+``RecordBatch.from_pandas`` bare at pandas/conversion.py:1026 and
connect/session.py:632
+(the createDataFrame 0-column branch) and
stateful_processor_api_client.py:557, relying on
+the default ``preserve_index=None`` to carry num_rows via the index metadata
-- otherwise a
+0-column relation loses its rows.
+
+Regenerate with SPARK_GENERATE_GOLDEN_FILES=1.
+"""
+
+import datetime
+import unittest
+
+from pyspark.testing.utils import (
+ have_pyarrow,
+ have_pandas,
+ pyarrow_requirement_message,
+ pandas_requirement_message,
+)
+from pyspark.testing.goldenutils import GoldenFileTestMixin
+
+if have_pandas:
+ import pandas as pd
+if have_pyarrow:
+ import pyarrow as pa
+
+
+class _PyArrowFromPandasFrameTestBase(GoldenFileTestMixin, unittest.TestCase):
+ """
+ Shared machinery for the DataFrame-input ``from_pandas`` constructors
(RecordBatch
+ here; Table/Schema as followups). Owns the source-frame inventory and
defines no
+ ``test_*`` of its own.
+ """
+
+ def _build_source_frames(self):
+ """Named pandas DataFrames covering shape x index-kind, plus a dtype
sample."""
+ dt = datetime.datetime(2020, 1, 1, 5, 30)
+ named = pd.Index([100, 200, 300], name="idx")
+ unnamed = pd.Index([10, 20, 30])
+ frames = {}
+
+ # =====================================================================
+ # 0-column frames -- only the index carries the row count
+ # =====================================================================
+ frames["0-columns:range-index"] = pd.DataFrame(index=range(3))
+ frames["0-columns:named-index"] = pd.DataFrame(index=named)
+ frames["0-columns:unnamed-index"] = pd.DataFrame(index=unnamed)
+ frames["0-columns:empty"] = pd.DataFrame(index=range(0))
+
+ # =====================================================================
+ # Single column -- a non-RangeIndex becomes an extra column
+ # =====================================================================
+ frames["single-column:range-index"] = pd.DataFrame({"a": [1, 2, 3]})
+ frames["single-column:named-index"] = pd.DataFrame({"a": [1, 2, 3]},
index=named)
+ frames["single-column:unnamed-index"] = pd.DataFrame({"a": [1, 2, 3]},
index=unnamed)
+
+ # =====================================================================
+ # Multi-column assembly. Drift-prone columns (object strings,
datetime64[ns]) are
+ # pinned so the Arrow output is stable across pandas 2/3; per-dtype
inference itself
+ # is already covered by test_pyarrow_array_from_pandas_default.
+ # =====================================================================
+ frames["multi-column:standard"] = pd.DataFrame(
+ {
+ "i": pd.Series([1, 2, 3], dtype="int64"),
+ "f": pd.Series([1.5, 2.5, 3.5], dtype="float64"),
+ "b": pd.Series([True, False, True], dtype=bool),
+ "s": pd.Series(["a", "b", "c"], dtype=object),
+ "t": pd.Series([dt, dt, dt], dtype="datetime64[ns]"),
+ }
+ )
+ frames["multi-column:nullable"] = pd.DataFrame(
+ {
+ "f": pd.Series([1.5, None, 3.5], dtype="float64"),
+ "b": pd.Series([True, None, False], dtype=object),
+ "s": pd.Series(["a", None, "c"], dtype=object),
+ "t": pd.Series([dt, None, dt], dtype="datetime64[ns]"),
+ }
+ )
+ # Multiple columns but zero rows (an empty object column is omitted --
it would
+ # infer to Arrow ``null`` rather than a concrete type).
+ frames["multi-column:no-rows"] = pd.DataFrame(
+ {
+ "i": pd.Series([], dtype="int64"),
+ "f": pd.Series([], dtype="float64"),
+ "b": pd.Series([], dtype=bool),
+ "t": pd.Series([], dtype="datetime64[ns]"),
+ }
+ )
+ return frames
+
+
[email protected](
+ not have_pyarrow or not have_pandas,
+ pyarrow_requirement_message or pandas_requirement_message,
+)
+class PyArrowRecordBatchFromPandasTests(_PyArrowFromPandasFrameTestBase):
+ """Tests pa.RecordBatch.from_pandas() across preserve_index via golden
file comparison."""
+
+ @staticmethod
+ def _index_desc(index) -> str:
+ """Compact, deterministic description of a pandas index for the input
cell."""
+ if isinstance(index, pd.RangeIndex):
+ return f"RangeIndex[{index.start}:{index.stop}:{index.step}]"
+ return f"{index.name!r}:{index.tolist()}"
+
+ def _input_cell(self, df) -> str:
+ """Input DataFrame repr, extended with its index (repr_value drops
it)."""
+ return f"{self.repr_value(df,
max_len=0)}[index={self._index_desc(df.index)}]"
+
+ def _from_pandas_cell(self, df, **kwargs) -> str:
+ """
+ Convert ``df`` via RecordBatch.from_pandas(**kwargs) and append
num_rows -- the
+ property this test pins, which a 0-column batch has no column to
imply. Returns
+ ERR@<ExceptionClass> if the conversion raises; a formatting error is a
test bug.
+ """
+ try:
+ batch = pa.RecordBatch.from_pandas(df, **kwargs)
+ except Exception as e:
+ return f"ERR@{type(e).__name__}"
+ return f"{self.repr_value(batch,
max_len=0)}[num_rows={batch.num_rows}]"
+
+ def test_from_pandas(self):
+ """Test pa.RecordBatch.from_pandas() across preserve_index against
golden file."""
+ sources = self._build_source_frames()
+ row_names = list(sources.keys())
+ preserve = {
+ "preserve_index=None": None,
+ "preserve_index=False": False,
+ "preserve_index=True": True,
+ }
+ col_names = ["pandas dataframe", *preserve.keys()]
+
+ # Version-specific expected values go here, keyed by (row, col), for
known drift.
+ overrides: dict[tuple[str, str], str] = {}
+
+ def compute_cell(row_name, col_name):
+ df = sources[row_name]
+ if col_name == "pandas dataframe":
+ return self._input_cell(df)
+ return self._from_pandas_cell(df,
preserve_index=preserve[col_name])
+
+ self.compare_or_generate_golden_matrix(
+ row_names=row_names,
+ col_names=col_names,
+ compute_cell=compute_cell,
+ golden_file_prefix="golden_pyarrow_record_batch_from_pandas",
+ index_name="test case",
+ overrides=overrides,
+ )
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
