Skip to content

[SPARK-40353][PS][CONNECT] Fix index nullable mismatch in ps.read_excel #50323

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

[SPARK-40353][PS][CONNECT] Fix index nullable mismatch in ps.read_excel #50323

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
069247a
fix
zhengruifeng Mar 19, 2025
c9d7ff4
fix
zhengruifeng Mar 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 25 additions & 5 deletions python/pyspark/pandas/namespace.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@
StringType,
DateType,
StructType,
StructField,
DataType,
)
from pyspark.sql.dataframe import DataFrame as PySparkDataFrame
Expand All @@ -85,6 +86,7 @@
from pyspark.pandas.frame import DataFrame, _reduce_spark_multi
from pyspark.pandas.internal import (
InternalFrame,
InternalField,
DEFAULT_SERIES_NAME,
HIDDEN_COLUMNS,
SPARK_INDEX_NAME_FORMAT,
Expand Down Expand Up @@ -1186,9 +1188,28 @@ def read_excel_on_spark(
pdf = pdf_or_pser

psdf = cast(DataFrame, from_pandas(pdf))
return_schema = force_decimal_precision_scale(
as_nullable_spark_type(psdf._internal.spark_frame.drop(*HIDDEN_COLUMNS).schema)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

as_nullable_spark_type convert all fields (both index and data) to nullable=true, while in InternalFrame it assert the index field should have nullable=false

)

raw_schema = psdf._internal.spark_frame.drop(*HIDDEN_COLUMNS).schema
index_scol_names = psdf._internal.index_spark_column_names
nullable_fields = []
for field in raw_schema.fields:
if field.name in index_scol_names:
nullable_fields.append(field)
else:
nullable_fields.append(
StructField(
field.name,
as_nullable_spark_type(field.dataType),
nullable=True,
metadata=field.metadata,
)
)
nullable_schema = StructType(nullable_fields)
return_schema = force_decimal_precision_scale(nullable_schema)

return_data_fields: Optional[List[InternalField]] = None
if psdf._internal.data_fields is not None:
return_data_fields = [f.normalize_spark_type() for f in psdf._internal.data_fields]

def output_func(pdf: pd.DataFrame) -> pd.DataFrame:
pdf = pd.concat([pd_read_excel(bin, sn=sn) for bin in pdf[pdf.columns[0]]])
Expand All @@ -1211,8 +1232,7 @@ def output_func(pdf: pd.DataFrame) -> pd.DataFrame:
.select("content")
.mapInPandas(lambda iterator: map(output_func, iterator), schema=return_schema)
)

return DataFrame(psdf._internal.with_new_sdf(sdf))
return DataFrame(psdf._internal.with_new_sdf(sdf, data_fields=return_data_fields))

if isinstance(pdf_or_psers, dict):
return {
Expand Down
53 changes: 8 additions & 45 deletions python/pyspark/pandas/tests/io/test_dataframe_spark_io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -253,8 +253,6 @@ def test_spark_io(self):
expected_idx.sort_values(by="f").to_spark().toPandas(),
)

# TODO(SPARK-40353): re-enabling the `test_read_excel`.
@unittest.skip("openpyxl")
def test_read_excel(self):
with self.temp_dir() as tmp:
path1 = "{}/file1.xlsx".format(tmp)
Expand All @@ -266,33 +264,34 @@ def test_read_excel(self):
pd.read_excel(open(path1, "rb"), index_col=0),
)
self.assert_eq(
ps.read_excel(open(path1, "rb"), index_col=0, squeeze=True),
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

squeeze is dropped at pandas 2.0, so we need to remove it

pd.read_excel(open(path1, "rb"), index_col=0, squeeze=True),
ps.read_excel(open(path1, "rb"), index_col=0),
pd.read_excel(open(path1, "rb"), index_col=0),
)

self.assert_eq(ps.read_excel(path1), pd.read_excel(path1))
self.assert_eq(ps.read_excel(path1, index_col=0), pd.read_excel(path1, index_col=0))
self.assert_eq(
ps.read_excel(path1, index_col=0, squeeze=True),
pd.read_excel(path1, index_col=0, squeeze=True),
ps.read_excel(path1, index_col=0),
pd.read_excel(path1, index_col=0),
)

self.assert_eq(ps.read_excel(tmp), pd.read_excel(path1))

path2 = "{}/file2.xlsx".format(tmp)
self.test_pdf[["i32"]].to_excel(path2)
print(ps.read_excel(tmp, index_col=0).sort_index())
self.assert_eq(
ps.read_excel(tmp, index_col=0).sort_index(),
pd.concat(
[pd.read_excel(path1, index_col=0), pd.read_excel(path2, index_col=0)]
).sort_index(),
)
self.assert_eq(
ps.read_excel(tmp, index_col=0, squeeze=True).sort_index(),
ps.read_excel(tmp, index_col=0).sort_index(),
pd.concat(
[
pd.read_excel(path1, index_col=0, squeeze=True),
pd.read_excel(path2, index_col=0, squeeze=True),
pd.read_excel(path1, index_col=0),
pd.read_excel(path2, index_col=0),
]
).sort_index(),
)
Expand All @@ -306,21 +305,12 @@ def test_read_excel(self):
sheet_names = [["Sheet_name_1", "Sheet_name_2"], None]

pdfs1 = pd.read_excel(open(path1, "rb"), sheet_name=None, index_col=0)
pdfs1_squeezed = pd.read_excel(
open(path1, "rb"), sheet_name=None, index_col=0, squeeze=True
)

for sheet_name in sheet_names:
psdfs = ps.read_excel(open(path1, "rb"), sheet_name=sheet_name, index_col=0)
self.assert_eq(psdfs["Sheet_name_1"], pdfs1["Sheet_name_1"])
self.assert_eq(psdfs["Sheet_name_2"], pdfs1["Sheet_name_2"])

psdfs = ps.read_excel(
open(path1, "rb"), sheet_name=sheet_name, index_col=0, squeeze=True
)
self.assert_eq(psdfs["Sheet_name_1"], pdfs1_squeezed["Sheet_name_1"])
self.assert_eq(psdfs["Sheet_name_2"], pdfs1_squeezed["Sheet_name_2"])

self.assert_eq(
ps.read_excel(tmp, index_col=0, sheet_name="Sheet_name_2"),
pdfs1["Sheet_name_2"],
Expand All @@ -331,30 +321,17 @@ def test_read_excel(self):
self.assert_eq(psdfs["Sheet_name_1"], pdfs1["Sheet_name_1"])
self.assert_eq(psdfs["Sheet_name_2"], pdfs1["Sheet_name_2"])

psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0, squeeze=True)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

such tests make no sense since squeeze=True is not allowed any more

self.assert_eq(psdfs["Sheet_name_1"], pdfs1_squeezed["Sheet_name_1"])
self.assert_eq(psdfs["Sheet_name_2"], pdfs1_squeezed["Sheet_name_2"])

path2 = "{}/file2.xlsx".format(tmp)
with pd.ExcelWriter(path2) as writer:
self.test_pdf.to_excel(writer, sheet_name="Sheet_name_1")
self.test_pdf[["i32"]].to_excel(writer, sheet_name="Sheet_name_2")

pdfs2 = pd.read_excel(path2, sheet_name=None, index_col=0)
pdfs2_squeezed = pd.read_excel(path2, sheet_name=None, index_col=0, squeeze=True)

self.assert_eq(
ps.read_excel(tmp, sheet_name="Sheet_name_2", index_col=0).sort_index(),
pd.concat([pdfs1["Sheet_name_2"], pdfs2["Sheet_name_2"]]).sort_index(),
)
self.assert_eq(
ps.read_excel(
tmp, sheet_name="Sheet_name_2", index_col=0, squeeze=True
).sort_index(),
pd.concat(
[pdfs1_squeezed["Sheet_name_2"], pdfs2_squeezed["Sheet_name_2"]]
).sort_index(),
)

for sheet_name in sheet_names:
psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0)
Expand All @@ -367,20 +344,6 @@ def test_read_excel(self):
pd.concat([pdfs1["Sheet_name_2"], pdfs2["Sheet_name_2"]]).sort_index(),
)

psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0, squeeze=True)
self.assert_eq(
psdfs["Sheet_name_1"].sort_index(),
pd.concat(
[pdfs1_squeezed["Sheet_name_1"], pdfs2_squeezed["Sheet_name_1"]]
).sort_index(),
)
self.assert_eq(
psdfs["Sheet_name_2"].sort_index(),
pd.concat(
[pdfs1_squeezed["Sheet_name_2"], pdfs2_squeezed["Sheet_name_2"]]
).sort_index(),
)

def test_read_orc(self):
with self.temp_dir() as tmp:
path = "{}/file1.orc".format(tmp)
Expand Down