From d4596e785f9189dd149eebca85b1c8f48a9c8485 Mon Sep 17 00:00:00 2001 From: May Liu Date: Mon, 20 Jul 2026 14:00:43 -0700 Subject: [PATCH 1/3] add support for to_polars() interop --- CHANGELOG.md | 4 + docs/source/snowpark/dataframe.rst | 1 + src/snowflake/snowpark/dataframe.py | 106 ++++++++ tests/integ/test_df_to_polars.py | 370 ++++++++++++++++++++++++++++ 4 files changed, 481 insertions(+) create mode 100644 tests/integ/test_df_to_polars.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3632e7adf3..04eac4f845 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ ### Snowpark Python API updates +#### New Features + +- Added `DataFrame.to_polars()` to convert a Snowpark DataFrame to a Polars DataFrame or LazyFrame. + #### Improvements - Removed the `experimental` tag from all AI SQL functions in `DataFrameAIFunctions` (`complete`, `filter`, `agg`, `classify`, `similarity`, `sentiment`, `embed`, `summarize_agg`, `transcribe`, `parse_document`, `extract`, `count_tokens`, `split_text_markdown_header`, `split_text_recursive_character`) and `RelationalGroupedDataFrame.ai_agg`. diff --git a/docs/source/snowpark/dataframe.rst b/docs/source/snowpark/dataframe.rst index a62bab7bc2..f6083cff84 100644 --- a/docs/source/snowpark/dataframe.rst +++ b/docs/source/snowpark/dataframe.rst @@ -92,6 +92,7 @@ DataFrame DataFrame.to_local_iterator DataFrame.to_pandas DataFrame.to_pandas_batches + DataFrame.to_polars DataFrame.to_snowpark_pandas DataFrame.union DataFrame.unionAll diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py index 6b0f59a1a0..442ae66085 100644 --- a/src/snowflake/snowpark/dataframe.py +++ b/src/snowflake/snowpark/dataframe.py @@ -238,6 +238,7 @@ if TYPE_CHECKING: import modin.pandas # pragma: no cover + import polars # pragma: no cover from table import Table # pragma: no cover _logger = getLogger(__name__) @@ -1371,6 +1372,111 @@ def to_arrow_batches( **kwargs, ) + @publicapi + @overload + def to_polars( + self, + *, + lazy: bool = False, + statement_params: Optional[Dict[str, str]] = None, + _emit_ast: bool = True, + **kwargs: Dict[str, Any], + ) -> "polars.DataFrame": + ... # pragma: no cover + + @publicapi + @overload + def to_polars( + self, + *, + lazy: bool = True, + statement_params: Optional[Dict[str, str]] = None, + _emit_ast: bool = True, + **kwargs: Dict[str, Any], + ) -> "polars.LazyFrame": + ... # pragma: no cover + + @experimental(version="1.54.0") + @df_collect_api_telemetry + @publicapi + def to_polars( + self, + *, + lazy: bool = False, + statement_params: Optional[Dict[str, str]] = None, + _emit_ast: bool = True, + **kwargs: Dict[str, Any], + ) -> Union["polars.DataFrame", "polars.LazyFrame"]: + """Executes the query representing this DataFrame and returns the result as a + `polars DataFrame `_ + or, when ``lazy=True``, a `polars LazyFrame `_. + + When the full result set is too large to materialize, use ``lazy=True`` and + push down column projections or row limits via polars expressions before + calling ``.collect()``. + + Example:: + + >>> df = session.create_dataframe([[1, 2], [3, 4]], schema=["a", "b"]) + >>> df.to_polars().shape + (2, 2) + >>> lf = df.to_polars(lazy=True) + >>> lf.collect().sort("A").to_dicts() + [{'A': 1, 'B': 2}, {'A': 3, 'B': 4}] + + Args: + lazy: If ``True``, the Snowpark plan is deferred until ``.collect()`` is called + on the returned :class:`polars.LazyFrame`. Defaults to ``False``. + statement_params: Dictionary of statement level parameters to be set while executing this action. + + Note: + Requires ``polars>=1.0``. + """ + import polars as pl + + if lazy: + schema = pl.from_arrow( + self.limit(1).to_arrow( + statement_params=statement_params, _emit_ast=False, **kwargs + ) + ).schema + + def _scan(with_columns, predicate, n_rows, batch_size): + # TODO(SNOW-3472759): push predicates down to Snowpark operations for Polars Exprs. + # Polars re-applies the predicate post-fetch, so correctness is preserved as-is. + scoped = self + if with_columns: + # Quote each name to preserve case for mixed-case / quoted + # Snowflake identifiers (unquoted names are always uppercase + # in Arrow, so quoting them is safe too). + scoped = scoped.select( + *[quote_name(c, keep_case=True) for c in with_columns] + ) + if n_rows is not None: + scoped = scoped.limit(n_rows) + for batch in scoped.to_arrow_batches( + statement_params=statement_params, _emit_ast=False, **kwargs + ): + yield pl.from_arrow(batch) + + return pl.io.plugins.register_io_source(_scan, schema=schema) + + parts = [ + pl.from_arrow(batch) + for batch in self.to_arrow_batches( + statement_params=statement_params, _emit_ast=False, **kwargs + ) + ] + if not parts: + # No batches returned: run a 0-row fetch to get the schema so the + # returned DataFrame has correct column names and types. + return pl.from_arrow( + self.limit(0).to_arrow( + statement_params=statement_params, _emit_ast=False, **kwargs + ) + ) + return pl.concat(parts) if len(parts) > 1 else parts[0] + @df_api_usage @publicapi def to_df( diff --git a/tests/integ/test_df_to_polars.py b/tests/integ/test_df_to_polars.py new file mode 100644 index 0000000000..a4f1bcab91 --- /dev/null +++ b/tests/integ/test_df_to_polars.py @@ -0,0 +1,370 @@ +#!/usr/bin/env python3 +# +# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved. +# + +from datetime import date, datetime +from decimal import Decimal +from unittest import mock + +import pytest + +from snowflake.snowpark.functions import col +from snowflake.snowpark.types import DecimalType + +from tests.utils import TestData + +try: + import polars as pl +except ImportError: + pytest.skip("polars not available", allow_module_level=True) + + +def polars_to_pydict(df: "pl.DataFrame") -> dict: + return df.to_arrow().to_pydict() + + +# --------------------------------------------------------------------------- +# Core type-family correctness +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +@pytest.mark.parametrize( + "example,expected", + [ + (TestData.integer1, {"A": [1, 2, 3]}), + ( + TestData.null_data1, + {"A": [None, Decimal("2"), Decimal("1"), Decimal("3"), None]}, + ), + ( + TestData.double1, + {"A": [Decimal("1.111"), Decimal("2.222"), Decimal("3.333")]}, + ), + (TestData.string1, {"A": ["test1", "test2", "test3"], "B": ["a", "b", "c"]}), + pytest.param( + TestData.array1, + { + "ARR1": ["[\n 1,\n 2,\n 3\n]", "[\n 6,\n 7,\n 8\n]"], + "ARR2": ["[\n 3,\n 4,\n 5\n]", "[\n 9,\n 0,\n 1\n]"], + }, + id="semi-structured array", + ), + pytest.param( + TestData.object2, + { + "OBJ": [ + '{\n "age": 21,\n "name": "Joe",\n "zip": 21021\n}', + '{\n "age": 26,\n "name": "Jay",\n "zip": 94021\n}', + ], + "K": ["age", "key"], + "V": [Decimal("0"), Decimal("0")], + "FLAG": [True, False], + }, + id="semi-structured object", + ), + ( + TestData.datetime_primitives2, + { + "TIMESTAMP": [ + datetime(9999, 12, 31, 0, 0, 0, 123456), + datetime(1583, 1, 1, 23, 59, 59, 567890), + ] + }, + ), + ( + TestData.date1, + { + "A": [date(2020, 8, 1), date(2010, 12, 1)], + "B": [Decimal("1"), Decimal("2")], + }, + ), + ], +) +def test_to_polars_type_correctness(session, example, expected): + df = example(session) + assert polars_to_pydict(df.to_polars()) == expected + + +# --------------------------------------------------------------------------- +# Decimal / NUMBER precision +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_decimal_precision(session): + data = [ + [1111111111111111111, 222222222222222222], + [3333333333333333333, 444444444444444444], + [5555555555555555555, 666666666666666666], + [7777777777777777777, 888888888888888888], + [9223372036854775807, 111111111111111111], + [2222222222222222222, 333333333333333333], + [4444444444444444444, 555555555555555555], + [6666666666666666666, 777777777777777777], + [-9223372036854775808, 999999999999999999], + ] + df = session.create_dataframe(data, schema=["A", "B"]).select( + col("A").cast(DecimalType(38, 0)).alias("A"), + col("B").cast(DecimalType(18, 0)).alias("B"), + ) + pl_df = df.to_polars() + pa_df = pl_df.to_arrow() + assert str(pa_df.schema[0].type) == "decimal128(38, 0)" + assert str(pa_df.schema[1].type) == "int64" + assert [[int(x) for x in row.values()] for row in pa_df.to_pylist()] == data + + +# --------------------------------------------------------------------------- +# NULL handling +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_null_integer_column(session): + df = session.create_dataframe([[0], [1], [None]], schema=["A"]) + col_values = df.to_polars()["A"].to_list() + assert col_values[0] == 0 + assert col_values[1] == 1 + assert col_values[2] is None + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_all_null_column(session): + df = session.create_dataframe([[None], [None], [None]], schema=["A"]) + pl_df = df.to_polars() + assert pl_df.height == 3 + assert all(v is None for v in pl_df["A"].to_list()) + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_empty_dataframe(session): + df = session.create_dataframe([[1, 2]], schema=["A", "B"]).filter(col("A") > 100) + pl_df = df.to_polars() + assert isinstance(pl_df, pl.DataFrame) + assert pl_df.height == 0 + assert set(pl_df.columns) == {"A", "B"} + + +# --------------------------------------------------------------------------- +# Eager path +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_eager_returns_dataframe(session): + df = session.create_dataframe([[1, "a"], [2, "b"]], schema=["A", "B"]) + result = df.to_polars() + assert isinstance(result, pl.DataFrame) + assert result.shape == (2, 2) + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_eager_multi_batch(session): + """Exercises the pl.concat(parts) path when the result spans multiple Arrow batches.""" + n = 100_000 + pl_df = session.range(n).to_polars() + assert isinstance(pl_df, pl.DataFrame) + assert pl_df.height == n + assert pl_df.columns == ["ID"] + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_timestamp_ltz_and_tz(session): + """TIMESTAMP_LTZ and TIMESTAMP_TZ survive the Arrow path (returned as tz-aware datetimes).""" + df = session.sql( + "SELECT " + "TO_TIMESTAMP_LTZ('2024-01-15 12:00:00 -0800') AS ts_ltz, " + "TO_TIMESTAMP_TZ('2024-01-15 20:00:00 +0530') AS ts_tz" + ) + pl_df = df.to_polars() + assert isinstance(pl_df, pl.DataFrame) + assert pl_df.height == 1 + assert pl_df["TS_LTZ"][0] is not None + assert pl_df["TS_TZ"][0] is not None + assert pl_df["TS_LTZ"].dtype.time_zone is not None + assert pl_df["TS_TZ"].dtype.time_zone is not None + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_eager_matches_to_arrow(session): + df = session.sql( + "SELECT 42::INT AS i, 3.14::FLOAT AS f, 'hello'::VARCHAR AS s, " + "TRUE AS b, DATE '2024-01-01' AS d, " + "TO_TIMESTAMP_NTZ('2024-01-01 12:00:00') AS t" + ) + assert polars_to_pydict(df.to_polars()) == df.to_arrow().to_pydict() + + +# --------------------------------------------------------------------------- +# Lazy path +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_lazy_returns_lazyframe(session): + df = session.create_dataframe([[1, "a"], [2, "b"]], schema=["A", "B"]) + lf = df.to_polars(lazy=True) + assert isinstance(lf, pl.LazyFrame) + assert set(lf.collect_schema().names()) == {"A", "B"} + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_lazy_collect_matches_eager(session): + df = session.create_dataframe( + [[1, "a", 1.5], [2, "b", 2.5], [3, "c", 3.5]], schema=["A", "B", "C"] + ) + assert df.to_polars().sort("A").equals(df.to_polars(lazy=True).collect().sort("A")) + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_lazy_type_correctness(session): + df = TestData.datetime_primitives2(session) + expected = { + "TIMESTAMP": [ + datetime(9999, 12, 31, 0, 0, 0, 123456), + datetime(1583, 1, 1, 23, 59, 59, 567890), + ] + } + assert polars_to_pydict(df.to_polars(lazy=True).collect()) == expected + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_lazy_projection_pushdown(session): + df = session.create_dataframe( + [[i, str(i), i * 1.5] for i in range(20)], schema=["A", "B", "C"] + ) + with mock.patch.object( + type(df), "select", autospec=True, side_effect=type(df).select + ) as spy: + # Single-column projection + result = df.to_polars(lazy=True).select("A").collect() + assert result.columns == ["A"] + assert result.height == 20 + called_col_sets = [ + {str(a).strip('"') for a in call.args[1:]} for call in spy.mock_calls + ] + assert any({"A"} == s for s in called_col_sets) + + with mock.patch.object( + type(df), "select", autospec=True, side_effect=type(df).select + ) as spy: + # Multi-column projection + result = df.to_polars(lazy=True).select("A", "B").collect() + assert set(result.columns) == {"A", "B"} + assert result.height == 20 + called_col_sets = [ + {str(a).strip('"') for a in call.args[1:]} for call in spy.mock_calls + ] + assert any({"A", "B"} == s for s in called_col_sets) + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_lazy_projection_pushdown_quoted_identifiers(session): + """Mixed-case quoted Snowflake identifiers survive projection pushdown without being uppercased.""" + df = session.sql('SELECT 1 AS "myInt", \'hello\' AS "myStr"') + result = df.to_polars(lazy=True).select("myInt").collect() + assert result.columns == ["myInt"] + assert result.height == 1 + assert result[0, 0] == 1 + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_lazy_limit_pushdown(session): + df = session.create_dataframe([[i, str(i)] for i in range(50)], schema=["A", "B"]) + assert df.to_polars(lazy=True).head(5).collect().height == 5 + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_lazy_construction_does_not_fetch_batches(session): + df = session.create_dataframe([[1, 2], [3, 4]], schema=["A", "B"]) + with mock.patch.object( + type(df), + "to_arrow_batches", + autospec=True, + side_effect=type(df).to_arrow_batches, + ) as batches_spy: + lf = df.to_polars(lazy=True) + assert isinstance(lf, pl.LazyFrame) + assert batches_spy.call_count == 0 + lf.collect() + assert batches_spy.call_count == 1 + + +# --------------------------------------------------------------------------- +# statement_params +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) +def test_to_polars_statement_params(session): + df = session.create_dataframe([[1]], schema=["A"]) + params = {"QUERY_TAG": "polars_integ_test"} + assert isinstance(df.to_polars(statement_params=params), pl.DataFrame) + assert isinstance( + df.to_polars(lazy=True, statement_params=params).collect(), pl.DataFrame + ) + + +# --------------------------------------------------------------------------- +# Missing dependency +# --------------------------------------------------------------------------- + + +def test_to_polars_raises_when_polars_missing(session): + df = session.create_dataframe([[1]], schema=["A"]) + with mock.patch.dict("sys.modules", {"polars": None}): + with pytest.raises(ModuleNotFoundError, match="polars"): + df.to_polars() From b293da514715bc615e881ec71d01426da3885903 Mon Sep 17 00:00:00 2001 From: May Liu Date: Mon, 20 Jul 2026 16:19:58 -0700 Subject: [PATCH 2/3] test coverage --- setup.py | 1 + src/snowflake/snowpark/dataframe.py | 6 +- tests/integ/test_df_to_polars.py | 390 ++++++++++++++++------------ 3 files changed, 222 insertions(+), 175 deletions(-) diff --git a/setup.py b/setup.py index 25d8590197..38471ff4bc 100644 --- a/setup.py +++ b/setup.py @@ -73,6 +73,7 @@ "psutil", # testing for telemetry "lxml", # used in XML reader unit tests "pyarrow", # used in dataframe reader tests + "polars>=1.0", # used in test_df_to_polars integration tests ] MODIN_DEVELOPMENT_REQUIREMENTS = [ # Snowpark pandas 3rd party library testing. Cap the scipy version because diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py index 442ae66085..bea143ffa1 100644 --- a/src/snowflake/snowpark/dataframe.py +++ b/src/snowflake/snowpark/dataframe.py @@ -1418,10 +1418,10 @@ def to_polars( Example:: >>> df = session.create_dataframe([[1, 2], [3, 4]], schema=["a", "b"]) - >>> df.to_polars().shape + >>> df.to_polars().shape # doctest: +SKIP (2, 2) - >>> lf = df.to_polars(lazy=True) - >>> lf.collect().sort("A").to_dicts() + >>> lf = df.to_polars(lazy=True) # doctest: +SKIP + >>> lf.collect().sort("A").to_dicts() # doctest: +SKIP [{'A': 1, 'B': 2}, {'A': 3, 'B': 4}] Args: diff --git a/tests/integ/test_df_to_polars.py b/tests/integ/test_df_to_polars.py index a4f1bcab91..acacfa5487 100644 --- a/tests/integ/test_df_to_polars.py +++ b/tests/integ/test_df_to_polars.py @@ -7,8 +7,10 @@ from decimal import Decimal from unittest import mock +import pyarrow as pa import pytest +from snowflake.snowpark import Session from snowflake.snowpark.functions import col from snowflake.snowpark.types import DecimalType @@ -16,8 +18,33 @@ try: import polars as pl + + _polars_available = True except ImportError: - pytest.skip("polars not available", allow_module_level=True) + pl = None # type: ignore[assignment] + _polars_available = False + +# Shorthand for tests that require a live Snowflake connection + Arrow support. +_skip_local = pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", + reason="arrow not fully supported by local testing.", +) + + +@pytest.fixture +def _no_polars_required(): + """Request this fixture to opt out of the polars-availability skip.""" + + +@pytest.fixture(autouse=True) +def _require_polars(request): + """Skip tests that need polars when it is not installed. + + Tests that mock polars or test the missing-dep path should request + the ``_no_polars_required`` fixture to opt out. + """ + if not _polars_available and "_no_polars_required" not in request.fixturenames: + pytest.skip("polars not available") def polars_to_pydict(df: "pl.DataFrame") -> dict: @@ -25,14 +52,38 @@ def polars_to_pydict(df: "pl.DataFrame") -> dict: # --------------------------------------------------------------------------- -# Core type-family correctness +# Fixtures / helpers # --------------------------------------------------------------------------- -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@pytest.fixture(scope="module") +def local_session(): + """Local-testing session — no Snowflake connection required.""" + with Session.builder.config("local_testing", True).create() as s: + yield s + + +def _mock_pl(): + """Minimal polars mock for unit tests that run without polars installed.""" + pl_mock = mock.MagicMock(name="polars") + mock_frame = mock.MagicMock(name="DataFrame") + mock_frame.schema = {"A": mock.MagicMock(), "B": mock.MagicMock()} + pl_mock.from_arrow.return_value = mock_frame + pl_mock.concat.return_value = mock_frame + mock_lf = mock.MagicMock(name="LazyFrame") + pl_mock.io.plugins.register_io_source.return_value = mock_lf + return pl_mock, mock_frame, mock_lf + + +_BATCH = pa.RecordBatch.from_pydict({"A": [1, 2], "B": ["a", "b"]}) + + +# --------------------------------------------------------------------------- +# Type correctness (eager) +# --------------------------------------------------------------------------- + + +@_skip_local @pytest.mark.parametrize( "example,expected", [ @@ -86,19 +137,10 @@ def polars_to_pydict(df: "pl.DataFrame") -> dict: ], ) def test_to_polars_type_correctness(session, example, expected): - df = example(session) - assert polars_to_pydict(df.to_polars()) == expected - + assert polars_to_pydict(example(session).to_polars()) == expected -# --------------------------------------------------------------------------- -# Decimal / NUMBER precision -# --------------------------------------------------------------------------- - -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@_skip_local def test_to_polars_decimal_precision(session): data = [ [1111111111111111111, 222222222222222222], @@ -115,8 +157,7 @@ def test_to_polars_decimal_precision(session): col("A").cast(DecimalType(38, 0)).alias("A"), col("B").cast(DecimalType(18, 0)).alias("B"), ) - pl_df = df.to_polars() - pa_df = pl_df.to_arrow() + pa_df = df.to_polars().to_arrow() assert str(pa_df.schema[0].type) == "decimal128(38, 0)" assert str(pa_df.schema[1].type) == "int64" assert [[int(x) for x in row.values()] for row in pa_df.to_pylist()] == data @@ -127,99 +168,67 @@ def test_to_polars_decimal_precision(session): # --------------------------------------------------------------------------- -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) -def test_to_polars_null_integer_column(session): - df = session.create_dataframe([[0], [1], [None]], schema=["A"]) - col_values = df.to_polars()["A"].to_list() - assert col_values[0] == 0 - assert col_values[1] == 1 - assert col_values[2] is None - +@_skip_local +def test_to_polars_null_handling(session): + # Nullable integer column + col_values = ( + session.create_dataframe([[0], [1], [None]], schema=["A"]) + .to_polars()["A"] + .to_list() + ) + assert col_values == [0, 1, None] -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) -def test_to_polars_all_null_column(session): - df = session.create_dataframe([[None], [None], [None]], schema=["A"]) - pl_df = df.to_polars() + # All-null column + pl_df = session.create_dataframe([[None], [None], [None]], schema=["A"]).to_polars() assert pl_df.height == 3 assert all(v is None for v in pl_df["A"].to_list()) -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) -def test_to_polars_empty_dataframe(session): - df = session.create_dataframe([[1, 2]], schema=["A", "B"]).filter(col("A") > 100) - pl_df = df.to_polars() - assert isinstance(pl_df, pl.DataFrame) - assert pl_df.height == 0 - assert set(pl_df.columns) == {"A", "B"} - - # --------------------------------------------------------------------------- # Eager path # --------------------------------------------------------------------------- -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) -def test_to_polars_eager_returns_dataframe(session): - df = session.create_dataframe([[1, "a"], [2, "b"]], schema=["A", "B"]) - result = df.to_polars() - assert isinstance(result, pl.DataFrame) - assert result.shape == (2, 2) - - -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@_skip_local def test_to_polars_eager_multi_batch(session): """Exercises the pl.concat(parts) path when the result spans multiple Arrow batches.""" - n = 100_000 - pl_df = session.range(n).to_polars() + pl_df = session.range(100_000).to_polars() assert isinstance(pl_df, pl.DataFrame) - assert pl_df.height == n + assert pl_df.height == 100_000 assert pl_df.columns == ["ID"] -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) -def test_to_polars_timestamp_ltz_and_tz(session): - """TIMESTAMP_LTZ and TIMESTAMP_TZ survive the Arrow path (returned as tz-aware datetimes).""" - df = session.sql( - "SELECT " - "TO_TIMESTAMP_LTZ('2024-01-15 12:00:00 -0800') AS ts_ltz, " - "TO_TIMESTAMP_TZ('2024-01-15 20:00:00 +0530') AS ts_tz" +@_skip_local +def test_to_polars_empty_dataframe(session): + pl_df = ( + session.create_dataframe([[1, 2]], schema=["A", "B"]) + .filter(col("A") > 100) + .to_polars() ) - pl_df = df.to_polars() assert isinstance(pl_df, pl.DataFrame) - assert pl_df.height == 1 - assert pl_df["TS_LTZ"][0] is not None - assert pl_df["TS_TZ"][0] is not None - assert pl_df["TS_LTZ"].dtype.time_zone is not None - assert pl_df["TS_TZ"].dtype.time_zone is not None + assert pl_df.height == 0 + assert set(pl_df.columns) == {"A", "B"} -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@_skip_local +def test_to_polars_timestamp_ltz_and_tz(session): + """TIMESTAMP_LTZ and TIMESTAMP_TZ survive the Arrow path as tz-aware datetimes.""" + pl_df = session.sql( + "SELECT TO_TIMESTAMP_LTZ('2024-01-15 12:00:00 -0800') AS ts_ltz," + " TO_TIMESTAMP_TZ('2024-01-15 20:00:00 +0530') AS ts_tz" + ).to_polars() + assert ( + pl_df["TS_LTZ"][0] is not None and pl_df["TS_LTZ"].dtype.time_zone is not None + ) + assert pl_df["TS_TZ"][0] is not None and pl_df["TS_TZ"].dtype.time_zone is not None + + +@_skip_local def test_to_polars_eager_matches_to_arrow(session): df = session.sql( - "SELECT 42::INT AS i, 3.14::FLOAT AS f, 'hello'::VARCHAR AS s, " - "TRUE AS b, DATE '2024-01-01' AS d, " - "TO_TIMESTAMP_NTZ('2024-01-01 12:00:00') AS t" + "SELECT 42::INT AS i, 3.14::FLOAT AS f, 'hello'::VARCHAR AS s," + " TRUE AS b, DATE '2024-01-01' AS d," + " TO_TIMESTAMP_NTZ('2024-01-01 12:00:00') AS t" ) assert polars_to_pydict(df.to_polars()) == df.to_arrow().to_pydict() @@ -229,47 +238,17 @@ def test_to_polars_eager_matches_to_arrow(session): # --------------------------------------------------------------------------- -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) -def test_to_polars_lazy_returns_lazyframe(session): - df = session.create_dataframe([[1, "a"], [2, "b"]], schema=["A", "B"]) - lf = df.to_polars(lazy=True) - assert isinstance(lf, pl.LazyFrame) - assert set(lf.collect_schema().names()) == {"A", "B"} - - -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@_skip_local def test_to_polars_lazy_collect_matches_eager(session): df = session.create_dataframe( [[1, "a", 1.5], [2, "b", 2.5], [3, "c", 3.5]], schema=["A", "B", "C"] ) - assert df.to_polars().sort("A").equals(df.to_polars(lazy=True).collect().sort("A")) + lf = df.to_polars(lazy=True) + assert isinstance(lf, pl.LazyFrame) + assert df.to_polars().sort("A").equals(lf.collect().sort("A")) -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) -def test_to_polars_lazy_type_correctness(session): - df = TestData.datetime_primitives2(session) - expected = { - "TIMESTAMP": [ - datetime(9999, 12, 31, 0, 0, 0, 123456), - datetime(1583, 1, 1, 23, 59, 59, 567890), - ] - } - assert polars_to_pydict(df.to_polars(lazy=True).collect()) == expected - - -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@_skip_local def test_to_polars_lazy_projection_pushdown(session): df = session.create_dataframe( [[i, str(i), i * 1.5] for i in range(20)], schema=["A", "B", "C"] @@ -277,54 +256,42 @@ def test_to_polars_lazy_projection_pushdown(session): with mock.patch.object( type(df), "select", autospec=True, side_effect=type(df).select ) as spy: - # Single-column projection result = df.to_polars(lazy=True).select("A").collect() - assert result.columns == ["A"] - assert result.height == 20 - called_col_sets = [ - {str(a).strip('"') for a in call.args[1:]} for call in spy.mock_calls - ] - assert any({"A"} == s for s in called_col_sets) + assert result.columns == ["A"] and result.height == 20 + assert any( + {"A"} == {str(a).strip('"') for a in c.args[1:]} for c in spy.mock_calls + ) with mock.patch.object( type(df), "select", autospec=True, side_effect=type(df).select ) as spy: - # Multi-column projection result = df.to_polars(lazy=True).select("A", "B").collect() - assert set(result.columns) == {"A", "B"} - assert result.height == 20 - called_col_sets = [ - {str(a).strip('"') for a in call.args[1:]} for call in spy.mock_calls - ] - assert any({"A", "B"} == s for s in called_col_sets) + assert set(result.columns) == {"A", "B"} and result.height == 20 + assert any( + {"A", "B"} == {str(a).strip('"') for a in c.args[1:]} + for c in spy.mock_calls + ) -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@_skip_local def test_to_polars_lazy_projection_pushdown_quoted_identifiers(session): - """Mixed-case quoted Snowflake identifiers survive projection pushdown without being uppercased.""" - df = session.sql('SELECT 1 AS "myInt", \'hello\' AS "myStr"') - result = df.to_polars(lazy=True).select("myInt").collect() - assert result.columns == ["myInt"] - assert result.height == 1 - assert result[0, 0] == 1 + """Mixed-case quoted identifiers survive projection pushdown without being uppercased.""" + result = ( + session.sql('SELECT 1 AS "myInt", \'hello\' AS "myStr"') + .to_polars(lazy=True) + .select("myInt") + .collect() + ) + assert result.columns == ["myInt"] and result.height == 1 and result[0, 0] == 1 -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@_skip_local def test_to_polars_lazy_limit_pushdown(session): df = session.create_dataframe([[i, str(i)] for i in range(50)], schema=["A", "B"]) assert df.to_polars(lazy=True).head(5).collect().height == 5 -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@_skip_local def test_to_polars_lazy_construction_does_not_fetch_batches(session): df = session.create_dataframe([[1, 2], [3, 4]], schema=["A", "B"]) with mock.patch.object( @@ -332,23 +299,14 @@ def test_to_polars_lazy_construction_does_not_fetch_batches(session): "to_arrow_batches", autospec=True, side_effect=type(df).to_arrow_batches, - ) as batches_spy: + ) as spy: lf = df.to_polars(lazy=True) - assert isinstance(lf, pl.LazyFrame) - assert batches_spy.call_count == 0 + assert isinstance(lf, pl.LazyFrame) and spy.call_count == 0 lf.collect() - assert batches_spy.call_count == 1 - - -# --------------------------------------------------------------------------- -# statement_params -# --------------------------------------------------------------------------- + assert spy.call_count == 1 -@pytest.mark.skipif( - "config.getoption('local_testing_mode', default=False)", - reason="arrow not fully supported by local testing.", -) +@_skip_local def test_to_polars_statement_params(session): df = session.create_dataframe([[1]], schema=["A"]) params = {"QUERY_TAG": "polars_integ_test"} @@ -358,13 +316,101 @@ def test_to_polars_statement_params(session): ) +# --------------------------------------------------------------------------- +# Unit coverage — no Snowflake, polars injected via mock +# These run regardless of whether polars is installed. +# --------------------------------------------------------------------------- + + +def test_to_polars_eager_unit_multi_batch_concat(local_session, _no_polars_required): + """pl.concat() is called when to_arrow_batches yields more than one batch.""" + pl_mock, mock_frame, _ = _mock_pl() + df = local_session.create_dataframe([[1, "a"], [2, "b"]], schema=["A", "B"]) + batch2 = pa.RecordBatch.from_pydict({"A": [3], "B": ["c"]}) + with mock.patch.dict("sys.modules", {"polars": pl_mock}): + with mock.patch.object( + type(df), "to_arrow_batches", return_value=[_BATCH, batch2] + ): + result = df.to_polars() + pl_mock.concat.assert_called_once() + assert result is mock_frame + + +def test_to_polars_eager_unit_no_batches_fallback(local_session, _no_polars_required): + """When to_arrow_batches yields nothing, falls back to limit(0).to_arrow() for schema.""" + pl_mock, _, _ = _mock_pl() + schema_table = pa.Table.from_pydict({"A": pa.array([], type=pa.int64())}) + df = local_session.create_dataframe([[1]], schema=["A"]) + with mock.patch.dict("sys.modules", {"polars": pl_mock}): + with mock.patch.object(type(df), "to_arrow_batches", return_value=[]): + with mock.patch.object(type(df), "to_arrow", return_value=schema_table): + df.to_polars() + pl_mock.from_arrow.assert_called_once_with(schema_table) + + +def test_to_polars_lazy_unit_scan_executes(local_session, _no_polars_required): + """_scan body yields pl.from_arrow(batch) for each batch returned by to_arrow_batches.""" + pl_mock, _, mock_lf = _mock_pl() + schema_table = pa.Table.from_pydict({"A": [1], "B": ["a"]}) + df = local_session.create_dataframe([[1, "a"]], schema=["A", "B"]) + + captured = {} + pl_mock.io.plugins.register_io_source.side_effect = ( + lambda fn, schema: captured.update(fn=fn) or mock_lf + ) + with mock.patch.dict("sys.modules", {"polars": pl_mock}): + with mock.patch.object(type(df), "to_arrow", return_value=schema_table): + df.to_polars(lazy=True) + + with mock.patch.object(type(df), "to_arrow_batches", return_value=[_BATCH]): + assert ( + len( + list( + captured["fn"]( + with_columns=None, predicate=None, n_rows=None, batch_size=None + ) + ) + ) + == 1 + ) + + +def test_to_polars_lazy_unit_scan_pushdowns(local_session, _no_polars_required): + """_scan calls select() for with_columns projection and limit() for n_rows.""" + pl_mock, _, mock_lf = _mock_pl() + schema_table = pa.Table.from_pydict({"A": [1], "B": ["a"]}) + df = local_session.create_dataframe([[1, "a"]], schema=["A", "B"]) + + captured = {} + pl_mock.io.plugins.register_io_source.side_effect = ( + lambda fn, schema: captured.update(fn=fn) or mock_lf + ) + with mock.patch.dict("sys.modules", {"polars": pl_mock}): + with mock.patch.object(type(df), "to_arrow", return_value=schema_table): + df.to_polars(lazy=True) + + col_batch = pa.RecordBatch.from_pydict({"A": [1]}) + with mock.patch.object(type(df), "select", wraps=df.select) as sel_spy: + with mock.patch.object(type(df), "limit", wraps=df.limit) as lim_spy: + with mock.patch.object( + type(df), "to_arrow_batches", return_value=[col_batch] + ): + list( + captured["fn"]( + with_columns=["A"], predicate=None, n_rows=5, batch_size=None + ) + ) + sel_spy.assert_called_once() + lim_spy.assert_called_once_with(5) + + # --------------------------------------------------------------------------- # Missing dependency # --------------------------------------------------------------------------- -def test_to_polars_raises_when_polars_missing(session): - df = session.create_dataframe([[1]], schema=["A"]) +def test_to_polars_raises_when_polars_missing(local_session, _no_polars_required): + df = local_session.create_dataframe([[1]], schema=["A"]) with mock.patch.dict("sys.modules", {"polars": None}): with pytest.raises(ModuleNotFoundError, match="polars"): df.to_polars() From 0e7f8fde88647205280c9d36b57eaad1544990c9 Mon Sep 17 00:00:00 2001 From: May Liu Date: Tue, 21 Jul 2026 11:52:05 -0700 Subject: [PATCH 3/3] AST for to_polars --- .../snowpark/_internal/proto/ast.proto | 440 +++++++++--------- src/snowflake/snowpark/dataframe.py | 16 +- src/snowflake/snowpark/mock/_connection.py | 19 + .../snowpark/mock/_nop_connection.py | 16 + tests/ast/data/DataFrame.to_polars.test | 304 ++++++++++++ 5 files changed, 578 insertions(+), 217 deletions(-) create mode 100644 tests/ast/data/DataFrame.to_polars.test diff --git a/src/snowflake/snowpark/_internal/proto/ast.proto b/src/snowflake/snowpark/_internal/proto/ast.proto index 03883767c3..a05533d59c 100644 --- a/src/snowflake/snowpark/_internal/proto/ast.proto +++ b/src/snowflake/snowpark/_internal/proto/ast.proto @@ -246,7 +246,7 @@ message FlattenMode { } } -// dataframe.ir:252 +// dataframe.ir:258 message JoinType { oneof variant { bool join_type__asof = 1; @@ -752,7 +752,7 @@ message CreateDataframe { SrcPosition src = 3; } -// dataframe.ir:183 +// dataframe.ir:189 message DataframeAgg { Expr df = 1; ExprArgList exprs = 2; @@ -892,7 +892,7 @@ message DataframeAiTranscribe { SrcPosition src = 5; } -// dataframe.ir:188 +// dataframe.ir:194 message DataframeAlias { Expr df = 1; string name = 2; @@ -955,7 +955,7 @@ message DataframeAnalyticsTimeSeriesAgg { repeated string windows = 8; } -// dataframe.ir:378 +// dataframe.ir:384 message DataframeCacheResult { Expr df = 1; NameRef object_name = 2; @@ -970,7 +970,7 @@ message DataframeCol { SrcPosition src = 3; } -// dataframe.ir:385 +// dataframe.ir:391 message DataframeColIlike { Expr df = 1; string pattern = 2; @@ -1042,7 +1042,7 @@ message DataframeCreateOrReplaceView { repeated Tuple_String_String statement_params = 7; } -// dataframe.ir:193 +// dataframe.ir:199 message DataframeCrossJoin { google.protobuf.BoolValue directed = 1; Expr lhs = 2; @@ -1059,7 +1059,7 @@ message DataframeCube { SrcPosition src = 3; } -// dataframe.ir:201 +// dataframe.ir:207 message DataframeDescribe { ExprArgList cols = 1; Expr df = 2; @@ -1067,41 +1067,41 @@ message DataframeDescribe { bool strings_include_math_stats = 4; } -// dataframe.ir:207 +// dataframe.ir:213 message DataframeDistinct { Expr df = 1; SrcPosition src = 2; } -// dataframe.ir:211 +// dataframe.ir:217 message DataframeDrop { ExprArgList cols = 1; Expr df = 2; SrcPosition src = 3; } -// dataframe.ir:216 +// dataframe.ir:222 message DataframeDropDuplicates { ExprArgList cols = 1; Expr df = 2; SrcPosition src = 3; } -// dataframe.ir:221 +// dataframe.ir:227 message DataframeExcept { Expr df = 1; Expr other = 2; SrcPosition src = 3; } -// dataframe.ir:226 +// dataframe.ir:232 message DataframeFilter { Expr condition = 1; Expr df = 2; SrcPosition src = 3; } -// dataframe.ir:240 +// dataframe.ir:246 message DataframeFirst { bool block = 1; Expr df = 2; @@ -1110,7 +1110,7 @@ message DataframeFirst { repeated Tuple_String_String statement_params = 5; } -// dataframe.ir:231 +// dataframe.ir:237 message DataframeFlatten { Expr df = 1; Expr input = 2; @@ -1135,14 +1135,14 @@ message DataframeGroupByGroupingSets { SrcPosition src = 3; } -// dataframe.ir:247 +// dataframe.ir:253 message DataframeIntersect { Expr df = 1; Expr other = 2; SrcPosition src = 3; } -// dataframe.ir:263 +// dataframe.ir:269 message DataframeJoin { google.protobuf.BoolValue directed = 1; Expr join_expr = 2; @@ -1155,14 +1155,14 @@ message DataframeJoin { SrcPosition src = 9; } -// dataframe.ir:282 +// dataframe.ir:288 message DataframeJoinTableFunction { Expr fn = 1; Expr lhs = 2; SrcPosition src = 3; } -// dataframe.ir:274 +// dataframe.ir:280 message DataframeLateralJoin { Expr join_expr = 1; Expr lhs = 2; @@ -1172,7 +1172,7 @@ message DataframeLateralJoin { SrcPosition src = 6; } -// dataframe.ir:287 +// dataframe.ir:293 message DataframeLimit { Expr df = 1; int64 n = 2; @@ -1180,7 +1180,7 @@ message DataframeLimit { SrcPosition src = 4; } -// dataframe.ir:153 +// dataframe.ir:159 message DataframeNaDrop_Python { Expr df = 1; string how = 2; @@ -1189,7 +1189,7 @@ message DataframeNaDrop_Python { google.protobuf.Int64Value thresh = 5; } -// dataframe.ir:147 +// dataframe.ir:153 message DataframeNaDrop_Scala { repeated string cols = 1; Expr df = 2; @@ -1197,7 +1197,7 @@ message DataframeNaDrop_Scala { SrcPosition src = 4; } -// dataframe.ir:160 +// dataframe.ir:166 message DataframeNaFill { Expr df = 1; bool include_decimal = 2; @@ -1207,7 +1207,7 @@ message DataframeNaFill { repeated Tuple_String_Expr value_map = 6; } -// dataframe.ir:168 +// dataframe.ir:174 message DataframeNaReplace { Expr df = 1; bool include_decimal = 2; @@ -1220,7 +1220,7 @@ message DataframeNaReplace { repeated Expr values = 9; } -// dataframe.ir:293 +// dataframe.ir:299 message DataframeNaturalJoin { google.protobuf.BoolValue directed = 1; JoinType join_type = 2; @@ -1238,7 +1238,7 @@ message DataframePivot { Expr values = 5; } -// dataframe.ir:308 +// dataframe.ir:314 message DataframeRandomSplit { Expr df = 1; google.protobuf.Int64Value seed = 2; @@ -1262,7 +1262,7 @@ message DataframeRef { SrcPosition src = 2; } -// dataframe.ir:315 +// dataframe.ir:321 message DataframeRename { Expr col_or_mapper = 1; Expr df = 2; @@ -1277,7 +1277,7 @@ message DataframeRollup { SrcPosition src = 3; } -// dataframe.ir:321 +// dataframe.ir:327 message DataframeSample { Expr df = 1; google.protobuf.Int64Value num = 2; @@ -1286,7 +1286,7 @@ message DataframeSample { SrcPosition src = 5; } -// dataframe.ir:328 +// dataframe.ir:334 message DataframeSelect { ExprArgList cols = 1; Expr df = 2; @@ -1301,7 +1301,7 @@ message DataframeShow { SrcPosition src = 3; } -// dataframe.ir:335 +// dataframe.ir:341 message DataframeSort { Expr ascending = 1; ExprArgList cols = 2; @@ -1353,7 +1353,7 @@ message DataframeStatSampleBy { SrcPosition src = 4; } -// dataframe.ir:138 +// dataframe.ir:144 message DataframeToDf { ExprArgList col_names = 1; Expr df = 2; @@ -1385,7 +1385,15 @@ message DataframeToPandasBatches { repeated Tuple_String_String statement_params = 4; } -// dataframe.ir:341 +// dataframe.ir:138 +message DataframeToPolars { + Expr df = 1; + bool is_lazy = 2; + SrcPosition src = 3; + repeated Tuple_String_String statement_params = 4; +} + +// dataframe.ir:347 message DataframeUnion { bool all = 1; bool allow_missing_columns = 2; @@ -1395,7 +1403,7 @@ message DataframeUnion { SrcPosition src = 6; } -// dataframe.ir:300 +// dataframe.ir:306 message DataframeUnpivot { repeated Expr column_list = 1; Expr df = 2; @@ -1405,7 +1413,7 @@ message DataframeUnpivot { string value_column = 6; } -// dataframe.ir:349 +// dataframe.ir:355 message DataframeWithColumn { Expr col = 1; string col_name = 2; @@ -1413,7 +1421,7 @@ message DataframeWithColumn { SrcPosition src = 4; } -// dataframe.ir:355 +// dataframe.ir:361 message DataframeWithColumnRenamed { Expr col = 1; Expr df = 2; @@ -1421,7 +1429,7 @@ message DataframeWithColumnRenamed { SrcPosition src = 4; } -// dataframe.ir:361 +// dataframe.ir:367 message DataframeWithColumns { repeated string col_names = 1; Expr df = 2; @@ -1625,92 +1633,93 @@ message Expr { DataframeToLocalIterator dataframe_to_local_iterator = 115; DataframeToPandas dataframe_to_pandas = 116; DataframeToPandasBatches dataframe_to_pandas_batches = 117; - DataframeUnion dataframe_union = 118; - DataframeUnpivot dataframe_unpivot = 119; - DataframeWithColumn dataframe_with_column = 120; - DataframeWithColumnRenamed dataframe_with_column_renamed = 121; - DataframeWithColumns dataframe_with_columns = 122; - DataframeWriter dataframe_writer = 123; - DatatypeVal datatype_val = 124; - Directory directory = 125; - Div div = 126; - Eq eq = 127; - Flatten flatten = 128; - Float64Val float64_val = 129; - FnRef fn_ref = 130; - Generator generator = 131; - Geq geq = 132; - GroupingSets grouping_sets = 133; - Gt gt = 134; - IndirectTableFnIdRef indirect_table_fn_id_ref = 135; - IndirectTableFnNameRef indirect_table_fn_name_ref = 136; - Int64Val int64_val = 137; - Leq leq = 138; - ListVal list_val = 139; - Lt lt = 140; - MergeDeleteWhenMatchedClause merge_delete_when_matched_clause = 141; - MergeInsertWhenNotMatchedClause merge_insert_when_not_matched_clause = 142; - MergeUpdateWhenMatchedClause merge_update_when_matched_clause = 143; - Mod mod = 144; - Mul mul = 145; - Neg neg = 146; - Neq neq = 147; - Not not = 148; - NullVal null_val = 149; - ObjectGetItem object_get_item = 150; - Or or = 151; - Pow pow = 152; - PythonDateVal python_date_val = 153; - PythonTimeVal python_time_val = 154; - PythonTimestampVal python_timestamp_val = 155; - Range range = 156; - ReadAvro read_avro = 157; - ReadCsv read_csv = 158; - ReadDirectory read_directory = 159; - ReadJson read_json = 160; - ReadLoad read_load = 161; - ReadOrc read_orc = 162; - ReadParquet read_parquet = 163; - ReadTable read_table = 164; - ReadXml read_xml = 165; - RedactedConst redacted_const = 166; - RelationalGroupedDataframeAgg relational_grouped_dataframe_agg = 167; - RelationalGroupedDataframeAiAgg relational_grouped_dataframe_ai_agg = 168; - RelationalGroupedDataframeApplyInPandas relational_grouped_dataframe_apply_in_pandas = 169; - RelationalGroupedDataframeBuiltin relational_grouped_dataframe_builtin = 170; - RelationalGroupedDataframePivot relational_grouped_dataframe_pivot = 171; - RelationalGroupedDataframeRef relational_grouped_dataframe_ref = 172; - Rollback rollback = 173; - Row row = 174; - SeqMapVal seq_map_val = 175; - SessionTableFunction session_table_function = 176; - Sql sql = 177; - SqlExpr sql_expr = 178; - StoredProcedure stored_procedure = 179; - StringVal string_val = 180; - Sub sub = 181; - Table table = 182; - TableDelete table_delete = 183; - TableDropTable table_drop_table = 184; - TableFnCallAlias table_fn_call_alias = 185; - TableFnCallOver table_fn_call_over = 186; - TableMerge table_merge = 187; - TableSample table_sample = 188; - TableUpdate table_update = 189; - ToSnowparkPandas to_snowpark_pandas = 190; - TruncatedExpr truncated_expr = 191; - TupleVal tuple_val = 192; - Udaf udaf = 193; - Udf udf = 194; - Udtf udtf = 195; - WriteCopyIntoLocation write_copy_into_location = 196; - WriteCsv write_csv = 197; - WriteInsertInto write_insert_into = 198; - WriteJson write_json = 199; - WritePandas write_pandas = 200; - WriteParquet write_parquet = 201; - WriteSave write_save = 202; - WriteTable write_table = 203; + DataframeToPolars dataframe_to_polars = 118; + DataframeUnion dataframe_union = 119; + DataframeUnpivot dataframe_unpivot = 120; + DataframeWithColumn dataframe_with_column = 121; + DataframeWithColumnRenamed dataframe_with_column_renamed = 122; + DataframeWithColumns dataframe_with_columns = 123; + DataframeWriter dataframe_writer = 124; + DatatypeVal datatype_val = 125; + Directory directory = 126; + Div div = 127; + Eq eq = 128; + Flatten flatten = 129; + Float64Val float64_val = 130; + FnRef fn_ref = 131; + Generator generator = 132; + Geq geq = 133; + GroupingSets grouping_sets = 134; + Gt gt = 135; + IndirectTableFnIdRef indirect_table_fn_id_ref = 136; + IndirectTableFnNameRef indirect_table_fn_name_ref = 137; + Int64Val int64_val = 138; + Leq leq = 139; + ListVal list_val = 140; + Lt lt = 141; + MergeDeleteWhenMatchedClause merge_delete_when_matched_clause = 142; + MergeInsertWhenNotMatchedClause merge_insert_when_not_matched_clause = 143; + MergeUpdateWhenMatchedClause merge_update_when_matched_clause = 144; + Mod mod = 145; + Mul mul = 146; + Neg neg = 147; + Neq neq = 148; + Not not = 149; + NullVal null_val = 150; + ObjectGetItem object_get_item = 151; + Or or = 152; + Pow pow = 153; + PythonDateVal python_date_val = 154; + PythonTimeVal python_time_val = 155; + PythonTimestampVal python_timestamp_val = 156; + Range range = 157; + ReadAvro read_avro = 158; + ReadCsv read_csv = 159; + ReadDirectory read_directory = 160; + ReadJson read_json = 161; + ReadLoad read_load = 162; + ReadOrc read_orc = 163; + ReadParquet read_parquet = 164; + ReadTable read_table = 165; + ReadXml read_xml = 166; + RedactedConst redacted_const = 167; + RelationalGroupedDataframeAgg relational_grouped_dataframe_agg = 168; + RelationalGroupedDataframeAiAgg relational_grouped_dataframe_ai_agg = 169; + RelationalGroupedDataframeApplyInPandas relational_grouped_dataframe_apply_in_pandas = 170; + RelationalGroupedDataframeBuiltin relational_grouped_dataframe_builtin = 171; + RelationalGroupedDataframePivot relational_grouped_dataframe_pivot = 172; + RelationalGroupedDataframeRef relational_grouped_dataframe_ref = 173; + Rollback rollback = 174; + Row row = 175; + SeqMapVal seq_map_val = 176; + SessionTableFunction session_table_function = 177; + Sql sql = 178; + SqlExpr sql_expr = 179; + StoredProcedure stored_procedure = 180; + StringVal string_val = 181; + Sub sub = 182; + Table table = 183; + TableDelete table_delete = 184; + TableDropTable table_drop_table = 185; + TableFnCallAlias table_fn_call_alias = 186; + TableFnCallOver table_fn_call_over = 187; + TableMerge table_merge = 188; + TableSample table_sample = 189; + TableUpdate table_update = 190; + ToSnowparkPandas to_snowpark_pandas = 191; + TruncatedExpr truncated_expr = 192; + TupleVal tuple_val = 193; + Udaf udaf = 194; + Udf udf = 195; + Udtf udtf = 196; + WriteCopyIntoLocation write_copy_into_location = 197; + WriteCsv write_csv = 198; + WriteInsertInto write_insert_into = 199; + WriteJson write_json = 200; + WritePandas write_pandas = 201; + WriteParquet write_parquet = 202; + WriteSave write_save = 203; + WriteTable write_table = 204; } } @@ -1802,7 +1811,7 @@ message Geq { SrcPosition src = 3; } -// dataframe.ir:367 +// dataframe.ir:373 message GroupingSets { ExprArgList sets = 1; SrcPosition src = 2; @@ -1937,98 +1946,99 @@ message HasSrcPosition { DataframeToLocalIterator dataframe_to_local_iterator = 118; DataframeToPandas dataframe_to_pandas = 119; DataframeToPandasBatches dataframe_to_pandas_batches = 120; - DataframeUnion dataframe_union = 121; - DataframeUnpivot dataframe_unpivot = 122; - DataframeWithColumn dataframe_with_column = 123; - DataframeWithColumnRenamed dataframe_with_column_renamed = 124; - DataframeWithColumns dataframe_with_columns = 125; - DataframeWriter dataframe_writer = 126; - DatatypeVal datatype_val = 127; - Directory directory = 128; - Div div = 129; - Eq eq = 130; - Flatten flatten = 131; - Float64Val float64_val = 132; - FnRef fn_ref = 133; - Generator generator = 134; - Geq geq = 135; - GroupingSets grouping_sets = 136; - Gt gt = 137; - IndirectTableFnIdRef indirect_table_fn_id_ref = 138; - IndirectTableFnNameRef indirect_table_fn_name_ref = 139; - Int64Val int64_val = 140; - Leq leq = 141; - ListVal list_val = 142; - Lt lt = 143; - MergeDeleteWhenMatchedClause merge_delete_when_matched_clause = 144; - MergeInsertWhenNotMatchedClause merge_insert_when_not_matched_clause = 145; - MergeUpdateWhenMatchedClause merge_update_when_matched_clause = 146; - Mod mod = 147; - Mul mul = 148; - NameRef name_ref = 149; - Neg neg = 150; - Neq neq = 151; - Not not = 152; - NullVal null_val = 153; - ObjectGetItem object_get_item = 154; - Or or = 155; - Pow pow = 156; - PythonDateVal python_date_val = 157; - PythonTimeVal python_time_val = 158; - PythonTimestampVal python_timestamp_val = 159; - Range range = 160; - ReadAvro read_avro = 161; - ReadCsv read_csv = 162; - ReadDirectory read_directory = 163; - ReadJson read_json = 164; - ReadLoad read_load = 165; - ReadOrc read_orc = 166; - ReadParquet read_parquet = 167; - ReadTable read_table = 168; - ReadXml read_xml = 169; - RedactedConst redacted_const = 170; - RelationalGroupedDataframeAgg relational_grouped_dataframe_agg = 171; - RelationalGroupedDataframeAiAgg relational_grouped_dataframe_ai_agg = 172; - RelationalGroupedDataframeApplyInPandas relational_grouped_dataframe_apply_in_pandas = 173; - RelationalGroupedDataframeBuiltin relational_grouped_dataframe_builtin = 174; - RelationalGroupedDataframePivot relational_grouped_dataframe_pivot = 175; - RelationalGroupedDataframeRef relational_grouped_dataframe_ref = 176; - Rollback rollback = 177; - Row row = 178; - SeqMapVal seq_map_val = 179; - SessionTableFunction session_table_function = 180; - Sql sql = 181; - SqlExpr sql_expr = 182; - StoredProcedure stored_procedure = 183; - StringVal string_val = 184; - Sub sub = 185; - Table table = 186; - TableDelete table_delete = 187; - TableDropTable table_drop_table = 188; - TableFnCallAlias table_fn_call_alias = 189; - TableFnCallOver table_fn_call_over = 190; - TableMerge table_merge = 191; - TableSample table_sample = 192; - TableUpdate table_update = 193; - ToSnowparkPandas to_snowpark_pandas = 194; - TruncatedExpr truncated_expr = 195; - TupleVal tuple_val = 196; - Udaf udaf = 197; - Udf udf = 198; - Udtf udtf = 199; - WindowSpecEmpty window_spec_empty = 200; - WindowSpecOrderBy window_spec_order_by = 201; - WindowSpecPartitionBy window_spec_partition_by = 202; - WindowSpecRangeBetween window_spec_range_between = 203; - WindowSpecRowsBetween window_spec_rows_between = 204; - WriteCopyIntoLocation write_copy_into_location = 205; - WriteCsv write_csv = 206; - WriteInsertInto write_insert_into = 207; - WriteJson write_json = 208; - WritePandas write_pandas = 209; - WriteParquet write_parquet = 210; - WriteSave write_save = 211; - WriteTable write_table = 212; + DataframeToPolars dataframe_to_polars = 121; + DataframeUnion dataframe_union = 122; + DataframeUnpivot dataframe_unpivot = 123; + DataframeWithColumn dataframe_with_column = 124; + DataframeWithColumnRenamed dataframe_with_column_renamed = 125; + DataframeWithColumns dataframe_with_columns = 126; + DataframeWriter dataframe_writer = 127; + DatatypeVal datatype_val = 128; + Directory directory = 129; + Div div = 130; + Eq eq = 131; + Flatten flatten = 132; + Float64Val float64_val = 133; + FnRef fn_ref = 134; + Generator generator = 135; + Geq geq = 136; + GroupingSets grouping_sets = 137; + Gt gt = 138; + IndirectTableFnIdRef indirect_table_fn_id_ref = 139; + IndirectTableFnNameRef indirect_table_fn_name_ref = 140; + Int64Val int64_val = 141; + Leq leq = 142; + ListVal list_val = 143; + Lt lt = 144; + MergeDeleteWhenMatchedClause merge_delete_when_matched_clause = 145; + MergeInsertWhenNotMatchedClause merge_insert_when_not_matched_clause = 146; + MergeUpdateWhenMatchedClause merge_update_when_matched_clause = 147; + Mod mod = 148; + Mul mul = 149; + NameRef name_ref = 150; + Neg neg = 151; + Neq neq = 152; + Not not = 153; + NullVal null_val = 154; + ObjectGetItem object_get_item = 155; + Or or = 156; + Pow pow = 157; + PythonDateVal python_date_val = 158; + PythonTimeVal python_time_val = 159; + PythonTimestampVal python_timestamp_val = 160; + Range range = 161; + ReadAvro read_avro = 162; + ReadCsv read_csv = 163; + ReadDirectory read_directory = 164; + ReadJson read_json = 165; + ReadLoad read_load = 166; + ReadOrc read_orc = 167; + ReadParquet read_parquet = 168; + ReadTable read_table = 169; + ReadXml read_xml = 170; + RedactedConst redacted_const = 171; + RelationalGroupedDataframeAgg relational_grouped_dataframe_agg = 172; + RelationalGroupedDataframeAiAgg relational_grouped_dataframe_ai_agg = 173; + RelationalGroupedDataframeApplyInPandas relational_grouped_dataframe_apply_in_pandas = 174; + RelationalGroupedDataframeBuiltin relational_grouped_dataframe_builtin = 175; + RelationalGroupedDataframePivot relational_grouped_dataframe_pivot = 176; + RelationalGroupedDataframeRef relational_grouped_dataframe_ref = 177; + Rollback rollback = 178; + Row row = 179; + SeqMapVal seq_map_val = 180; + SessionTableFunction session_table_function = 181; + Sql sql = 182; + SqlExpr sql_expr = 183; + StoredProcedure stored_procedure = 184; + StringVal string_val = 185; + Sub sub = 186; + Table table = 187; + TableDelete table_delete = 188; + TableDropTable table_drop_table = 189; + TableFnCallAlias table_fn_call_alias = 190; + TableFnCallOver table_fn_call_over = 191; + TableMerge table_merge = 192; + TableSample table_sample = 193; + TableUpdate table_update = 194; + ToSnowparkPandas to_snowpark_pandas = 195; + TruncatedExpr truncated_expr = 196; + TupleVal tuple_val = 197; + Udaf udaf = 198; + Udf udf = 199; + Udtf udtf = 200; + WindowSpecEmpty window_spec_empty = 201; + WindowSpecOrderBy window_spec_order_by = 202; + WindowSpecPartitionBy window_spec_partition_by = 203; + WindowSpecRangeBetween window_spec_range_between = 204; + WindowSpecRowsBetween window_spec_rows_between = 205; + WriteCopyIntoLocation write_copy_into_location = 206; + WriteCsv write_csv = 207; + WriteInsertInto write_insert_into = 208; + WriteJson write_json = 209; + WritePandas write_pandas = 210; + WriteParquet write_parquet = 211; + WriteSave write_save = 212; + WriteTable write_table = 213; } } @@ -2534,7 +2544,7 @@ message TableUpdate { repeated Tuple_String_String statement_params = 7; } -// dataframe.ir:371 +// dataframe.ir:377 message ToSnowparkPandas { repeated string columns = 1; Expr df = 2; diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py index bea143ffa1..a2016203ba 100644 --- a/src/snowflake/snowpark/dataframe.py +++ b/src/snowflake/snowpark/dataframe.py @@ -1432,14 +1432,26 @@ def to_polars( Note: Requires ``polars>=1.0``. """ + if _emit_ast: + stmt = self._session._ast_batch.bind() + ast = with_src_position(stmt.expr.dataframe_to_polars, stmt) + self._set_ast_ref(ast.df) + ast.is_lazy = lazy + if statement_params is not None: + build_expr_from_dict_str_str(ast.statement_params, statement_params) + self._session._ast_batch.eval(stmt) + _, kwargs[DATAFRAME_AST_PARAMETER] = self._session._ast_batch.flush(stmt) + import polars as pl if lazy: schema = pl.from_arrow( - self.limit(1).to_arrow( + self.limit(1, _emit_ast=False).to_arrow( statement_params=statement_params, _emit_ast=False, **kwargs ) ).schema + # Remove the AST parameter so _scan calls don't re-report the same batch. + kwargs.pop(DATAFRAME_AST_PARAMETER, None) def _scan(with_columns, predicate, n_rows, batch_size): # TODO(SNOW-3472759): push predicates down to Snowpark operations for Polars Exprs. @@ -1471,7 +1483,7 @@ def _scan(with_columns, predicate, n_rows, batch_size): # No batches returned: run a 0-row fetch to get the schema so the # returned DataFrame has correct column names and types. return pl.from_arrow( - self.limit(0).to_arrow( + self.limit(0, _emit_ast=False).to_arrow( statement_params=statement_params, _emit_ast=False, **kwargs ) ) diff --git a/src/snowflake/snowpark/mock/_connection.py b/src/snowflake/snowpark/mock/_connection.py index 92fea43002..bd52b15574 100644 --- a/src/snowflake/snowpark/mock/_connection.py +++ b/src/snowflake/snowpark/mock/_connection.py @@ -719,6 +719,25 @@ def execute( # we do not mock the splitting into data chunks behavior rows = [rows] if to_iter else rows + to_arrow = kwargs.get("to_arrow", False) + if to_arrow: + try: + import pyarrow as pa + + if isinstance(res, TableEmulator): + renamed = res.rename( + columns={c: unquote_if_quoted(c) for c in res.columns} + ) + arrow_table = pa.Table.from_pandas(renamed, preserve_index=False) + else: + arrow_table = pa.table({}) + except Exception: + import pyarrow as pa + + arrow_table = pa.table({}) + self.notify_mock_query_record_listener(**kwargs) + return iter([arrow_table]) if to_iter else arrow_table + # Notify query listeners. self.notify_mock_query_record_listener(**kwargs) diff --git a/src/snowflake/snowpark/mock/_nop_connection.py b/src/snowflake/snowpark/mock/_nop_connection.py index c7b98e1db9..e9d5a3f0ad 100644 --- a/src/snowflake/snowpark/mock/_nop_connection.py +++ b/src/snowflake/snowpark/mock/_nop_connection.py @@ -231,11 +231,27 @@ def execute( # Create a dummy single row DataFrame with the expected schema. Note that the schema # attributes is a based on best effort based on most common operators but won't work for # things like dynamic pivot and possibly other operators. + to_arrow_flag = kwargs.get("to_arrow", False) if to_pandas: result = pandas.DataFrame( [[v for v in result_row[0]]], columns=[rm.name for rm in result_meta], ) + elif to_arrow_flag: + try: + import pyarrow as pa + + cols = { + result_meta[i].name: [result_row[0][i]] + for i in range(len(result_meta)) + } + arrow_result = pa.table(cols) + except Exception: + import pyarrow as pa + + arrow_result = pa.table({}) + self.notify_mock_query_record_listener(**kwargs) + return iter([arrow_result]) if to_iter else arrow_result else: result = result_row diff --git a/tests/ast/data/DataFrame.to_polars.test b/tests/ast/data/DataFrame.to_polars.test new file mode 100644 index 0000000000..f9b003d52d --- /dev/null +++ b/tests/ast/data/DataFrame.to_polars.test @@ -0,0 +1,304 @@ +## TEST CASE + +df = session.table(tables.table1) + +df.to_polars() + +df.to_polars(lazy=True) + +df.to_polars(statement_params={"SF_PARTNER": "FAKE_PARTNER"}) + +df.to_polars(lazy=True, statement_params={"SF_PARTNER": "FAKE_PARTNER"}) + +## EXPECTED UNPARSER OUTPUT + +df = session.table("table1") + +df.to_polars() + +df = session.table("table1") + +df.to_polars(lazy=True) + +df = session.table("table1") + +df.to_polars(statement_params={"SF_PARTNER": "FAKE_PARTNER"}) + +df = session.table("table1") + +df.to_polars(statement_params={"SF_PARTNER": "FAKE_PARTNER"}, lazy=True) + +## EXPECTED ENCODED AST + +interned_value_table { + string_values { + key: -1 + } + string_values { + key: 2 + value: "SRC_POSITION_TEST_MODE" + } +} +body { + bind { + expr { + table { + name { + name { + name_flat { + name: "table1" + } + } + } + src { + end_column: 41 + end_line: 25 + file: 2 + start_column: 13 + start_line: 25 + } + variant { + session_table: true + } + } + } + first_request_id: "\003U\"\366q\366P\346\260\261?\234\303\254\316\353" + symbol { + value: "df" + } + uid: 1 + } +} +body { + bind { + expr { + dataframe_to_polars { + df { + dataframe_ref { + id: 1 + } + } + src { + end_column: 22 + end_line: 27 + file: 2 + start_column: 8 + start_line: 27 + } + } + } + first_request_id: "\003U\"\366q\366P\346\260\261?\234\303\254\316\353" + symbol { + } + uid: 2 + } +} +body { + eval { + bind_id: 2 + } +} +body { + bind { + expr { + table { + name { + name { + name_flat { + name: "table1" + } + } + } + src { + end_column: 41 + end_line: 25 + file: 2 + start_column: 13 + start_line: 25 + } + variant { + session_table: true + } + } + } + first_request_id: "\003U\"\366q\366P\346\260\261?\234\303\254\316\353" + symbol { + value: "df" + } + uid: 1 + } +} +body { + bind { + expr { + dataframe_to_polars { + df { + dataframe_ref { + id: 1 + } + } + is_lazy: true + src { + end_column: 31 + end_line: 29 + file: 2 + start_column: 8 + start_line: 29 + } + } + } + first_request_id: "\003U\"\366q\366P\346\260\261?\234\303\254\316\353" + symbol { + } + uid: 3 + } +} +body { + eval { + bind_id: 3 + } +} +body { + bind { + expr { + table { + name { + name { + name_flat { + name: "table1" + } + } + } + src { + end_column: 41 + end_line: 25 + file: 2 + start_column: 13 + start_line: 25 + } + variant { + session_table: true + } + } + } + first_request_id: "\003U\"\366q\366P\346\260\261?\234\303\254\316\353" + symbol { + value: "df" + } + uid: 1 + } +} +body { + bind { + expr { + dataframe_to_polars { + df { + dataframe_ref { + id: 1 + } + } + src { + end_column: 69 + end_line: 31 + file: 2 + start_column: 8 + start_line: 31 + } + statement_params { + _1: "SF_PARTNER" + _2: "FAKE_PARTNER" + } + } + } + first_request_id: "\003U\"\366q\366P\346\260\261?\234\303\254\316\353" + symbol { + } + uid: 4 + } +} +body { + eval { + bind_id: 4 + } +} +body { + bind { + expr { + table { + name { + name { + name_flat { + name: "table1" + } + } + } + src { + end_column: 41 + end_line: 25 + file: 2 + start_column: 13 + start_line: 25 + } + variant { + session_table: true + } + } + } + first_request_id: "\003U\"\366q\366P\346\260\261?\234\303\254\316\353" + symbol { + value: "df" + } + uid: 1 + } +} +body { + bind { + expr { + dataframe_to_polars { + df { + dataframe_ref { + id: 1 + } + } + is_lazy: true + src { + end_column: 80 + end_line: 33 + file: 2 + start_column: 8 + start_line: 33 + } + statement_params { + _1: "SF_PARTNER" + _2: "FAKE_PARTNER" + } + } + } + first_request_id: "\003U\"\366q\366P\346\260\261?\234\303\254\316\353" + symbol { + } + uid: 5 + } +} +body { + eval { + bind_id: 5 + } +} +client_ast_version: 1 +client_language { + python_language { + version { + label: "final" + major: 3 + minor: 9 + patch: 1 + } + } +} +client_version { + major: 1 + minor: 53 + patch: 1 +} +id: "\003U\"\366q\366P\346\260\261?\234\303\254\316\353"