zhengruifeng commented on code in PR #50382:
URL: https://github.com/apache/spark/pull/50382#discussion_r2011371226
##########
python/pyspark/pandas/namespace.py:
##########
@@ -1157,89 +1161,99 @@ def pd_read_excel(
**kwds,
)
- if isinstance(io, str):
- # 'binaryFile' format is available since Spark 3.0.0.
- binaries =
default_session().read.format("binaryFile").load(io).select("content").head(2)
- io_or_bin = binaries[0][0]
- single_file = len(binaries) == 1
- else:
- io_or_bin = io
- single_file = True
-
- pdf_or_psers = pd_read_excel(io_or_bin, sn=sheet_name)
-
- if single_file:
+ if not isinstance(io, str):
+ # When io is not a path, always need to load all data to python side
+ pdf_or_psers = pd_read_excel(io, sn=sheet_name, nr=nrows)
if isinstance(pdf_or_psers, dict):
return {
sn: cast(Union[DataFrame, Series], from_pandas(pdf_or_pser))
for sn, pdf_or_pser in pdf_or_psers.items()
}
else:
return cast(Union[DataFrame, Series], from_pandas(pdf_or_psers))
- else:
- def read_excel_on_spark(
- pdf_or_pser: Union[pd.DataFrame, pd.Series],
- sn: Union[str, int, List[Union[str, int]], None],
- ) -> Union[DataFrame, Series]:
- if isinstance(pdf_or_pser, pd.Series):
- pdf = pdf_or_pser.to_frame()
- else:
- pdf = pdf_or_pser
+ spark = default_session()
+
+ # Collect the first #nr rows from the first file
+ nr = get_option("compute.max_rows", 1000)
+ if nrows is not None and nrows < nr:
+ nr = nrows
+
+ def sample_data(pdf: pd.DataFrame) -> pd.DataFrame:
+ raw_data = BytesIO(pdf.content[0])
+ pdf_or_dict = pd_read_excel(raw_data, sn=sheet_name, nr=nr)
+ return pd.DataFrame({"sampled": [pickle.dumps(pdf_or_dict)]})
Review Comment:
pickle the sampled rows
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
