SNOW-3472759: support for to_polars() for snowpark -> polars interoperability - #4288
SNOW-3472759: support for to_polars() for snowpark -> polars interoperability#4288sfc-gh-mayliu wants to merge 8 commits into
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #4288 +/- ##
==========================================
- Coverage 95.46% 95.27% -0.20%
==========================================
Files 171 172 +1
Lines 44720 44804 +84
Branches 7676 7686 +10
==========================================
- Hits 42694 42689 -5
- Misses 1253 1299 +46
- Partials 773 816 +43 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![graphite-app[bot]](https://avatars.githubusercontent.com/in/158384?s=60&v=4){kind=small}
…lelize file open and read for parquet, to_arrow for plans with 1 query
…row_batches; 3. add usage notes in docstring
![snowflake-security-bot[bot]](https://avatars.githubusercontent.com/u/140554735?s=60&v=4){kind=small}
| ) | ||
| # The stream objects are owned by the returned LazyFrame; Polars closes them | ||
| # when the scan is materialized (or if the LazyFrame is discarded). | ||
| streams = _open_stage_files_parallel( |
There was a problem hiding this comment.
i have a question about this ingestion path.
Since this path is ingesting with a stream of parquet file,
if I call collect() two times, will the second collect() still be able to fetch anything since the stream is already consumed at that time?
There was a problem hiding this comment.
Both session.file.get_stream and SnowflakeFile.open eagerly downloads/buffers, so the multiple Polars .collect() calls are safe to get the predicate/projection pruned data on CPU without needing to re-download. However, it's a Snowflake infra-limitation that network I/O pruning costs can't be saved due to the lack of HTTP Range GET support in session.file and SnowflakeFile
Tested with 1M rows: first .collect() takes 0.04s, second .collect() takes 0.03s — effectively the same cpu-decode cost with no additional network I/O.
| rid = uuid.uuid4().hex[:12] | ||
| stage = df._session.get_session_stage().rstrip("/") | ||
| sub = f"{stage}/{sub_prefix}/{rid}/" | ||
| df.write.copy_into_location( |
There was a problem hiding this comment.
it appears like the uploaded parquet file is never removed, do we need a garbage cleaning mechanism to make sure the stage is not filled when user invoke this ingestion route multiple times?
There was a problem hiding this comment.
get_session_stage() returns a temporary stage that Snowflake auto-drops on session close, so no need to recycle here
There was a problem hiding this comment.
I am thinking about a scenario in which user have a long running pipeline that call to_polars() in this ingestion route, in that scenario, the stage will keep accumulating uploaded parquet files. Do we need to worry about this scenario or is this not a expected usage?
There was a problem hiding this comment.
Good question, discussed offline. TLDR: long-lived sessions exist but are rare, and there's no storage limit on temp stages. sproc also rejects REMOVE {file} calls, so a clean up would result in inconsistency between local and server environments. We will rely on temp stages' property to clean up after itself, similar to existing Snowpark patterns
| overwrite=True, | ||
| statement_params=statement_params, | ||
| ) | ||
| rows = df._session.sql(f"LIST '{sub}'").collect(statement_params=statement_params) |
There was a problem hiding this comment.
I remember that List would not return fully qualified stage name, should we make it fully qualified before return?
There was a problem hiding this comment.
Good point, LS returns stage-relative paths and we should make it fully qualified. We already have the fully-qualified sub prefix at hand, so the fix is to extract just the filename from the LIST result and prepend sub instead. Fix incoming.
| ) -> list: | ||
| """Open stage files concurrently and return file-like objects. | ||
|
|
||
| Uses ``SnowflakeFile.open`` inside a stored procedure (no local /tmp | ||
| materialization) and ``session.file.get_stream`` on the client. | ||
| """ | ||
| if is_sproc: | ||
| from snowflake.snowpark.files import SnowflakeFile | ||
|
|
||
| def _open(p: str): | ||
| return SnowflakeFile.open(p, "rb", require_scoped_url=False) | ||
|
|
||
| else: | ||
|
|
||
| def _open(p: str): | ||
| return session.file.get_stream(p) | ||
|
|
||
| if not paths: | ||
| return [] | ||
| with ThreadPoolExecutor(max_workers=max_workers) as ex: | ||
| return list(ex.map(_open, paths)) | ||
|
|
||
|
|
||
| def _open_and_read_parquet_parallel( | ||
| session, | ||
| paths: list[str], | ||
| is_sproc: bool, | ||
| max_workers: int | None = None, | ||
| ) -> list[pl.DataFrame]: | ||
| """Open and read each stage Parquet file in a thread pool. | ||
|
|
||
| Opens and decodes run together per file so the read step is not serialized | ||
| after a parallel open — the decode releases the GIL for I/O and native | ||
| Arrow work. | ||
| """ | ||
| import polars as pl | ||
|
|
||
| if is_sproc: | ||
| from snowflake.snowpark.files import SnowflakeFile | ||
|
|
||
| def _open_read(p: str): | ||
| f = SnowflakeFile.open(p, "rb", require_scoped_url=False) | ||
| try: | ||
| return pl.read_parquet(f) | ||
| finally: | ||
| try: | ||
| f.close() | ||
| except Exception as e: | ||
| _logger.warning("failed to close stage file %s: %s", p, e) | ||
|
|
||
| else: | ||
|
|
||
| def _open_read(p: str): | ||
| f = session.file.get_stream(p) | ||
| try: | ||
| return pl.read_parquet(f) | ||
| finally: | ||
| try: | ||
| f.close() | ||
| except Exception as e: | ||
| _logger.warning("failed to close stage file %s: %s", p, e) | ||
|
|
||
| if not paths: | ||
| return [] | ||
| with ThreadPoolExecutor(max_workers=max_workers) as ex: | ||
| return list(ex.map(_open_read, paths)) |
There was a problem hiding this comment.
this is a styling issue -- I'm seeing repeated code structures here in 2 ways:
- within the function, two similar _open_read, one for sproc, one for non-sproc
- between _open_and_read_parquet_parallel and _open_stage_files_parallel, one is eager one is lazy
is it possible to improve reusability? There is a Python built-in functools.partial which might help.
I think the 4 (2x2) can be consolidated into one function without sacrificing too much readability, but would improve maintainability
def _open_stage_files_parallel(session, path, is_sproc, is_lazy) -> list[pl.Dafarame] | list
| ) -> "polars.LazyFrame": | ||
| ... # pragma: no cover | ||
|
|
||
| @experimental(version="1.54.0") |
There was a problem hiding this comment.
nit:
| @experimental(version="1.54.0") | |
| @experimental(version="1.55.0") |
| @publicapi | ||
| @overload | ||
| def to_polars( | ||
| self, | ||
| *, | ||
| lazy: bool = False, | ||
| use_parquet: bool = False, | ||
| max_workers: Optional[int] = None, | ||
| statement_params: Optional[Dict[str, str]] = None, | ||
| _emit_ast: bool = False, | ||
| ) -> "polars.DataFrame": | ||
| ... # pragma: no cover | ||
|
|
||
| @publicapi | ||
| @overload | ||
| def to_polars( | ||
| self, | ||
| *, | ||
| lazy: bool = True, | ||
| use_parquet: bool = False, | ||
| max_workers: Optional[int] = None, | ||
| statement_params: Optional[Dict[str, str]] = None, | ||
| _emit_ast: bool = False, | ||
| ) -> "polars.LazyFrame": | ||
| ... # pragma: no cover |
There was a problem hiding this comment.
nit: I remember overload function requires the parameter value to be defined in the following format
lazy: Literal[True]
lazy: Literal[False]
4 participants
Which Jira issue is this PR addressing? Make sure that there is an accompanying issue to your PR.
Fixes SNOW-3472759
Fill out the following pre-review checklist:
Please describe how your code solves the related issue.
Please write a short description of how your code change solves the related issue.
This PR adds
DataFrame.to_polars()to enable direct conversion from a Snowpark DataFrame to a Polars DataFrame/LazyFrame