Skip to content

Commit a0c7e18

Browse files
bernhardschaefernicornk
bernhardschaefer
and
nicornk
authored
Implement query_foundry_sql(return_type='polars') (#115)
* query_foundry_sql(return_type='polars') * docs update * use assert_in_literal * do no extend query_foundry_sql_legacy * reduce diff * reduce diff v2 * v3 * add legacy query polars to follow same workflow as pandas * Update libs/foundry-dev-tools/src/foundry_dev_tools/clients/data_proxy.py --------- Co-authored-by: Nicolas Renkamp <nicornk@users.noreply.github.com>
1 parent 81330fc commit a0c7e18

9 files changed

Lines changed: 107 additions & 38 deletions

File tree

‎docs/dev/contribute.md‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,12 @@ cd foundry-dev-tools
1919
````{tab} Conda/Mamba
2020
First create an environment specifically for foundry-dev-tools with pdm and activate it
2121
```shell
22-
mamba create -n foundry-dev-tools pdm openjdk===17
22+
mamba create -n foundry-dev-tools python=3.12 pdm openjdk=17
2323
mamba activate foundry-dev-tools
2424
```
2525
````
26+
Note that python>3.12 throws an error when running `pdm install`: `configured Python interpreter version (3.14) is newer than PyO3's maximum supported version (3.12)`.
27+
2628
````{tab} Without Conda/Mamba
2729
If you don't want to use conda or mamba, you'll need to install pdm through other means.
2830
It is available in most Linux package managers, Homebrew, Scoop and also via pip.

‎docs/examples/dataset.md‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -260,9 +260,7 @@ import polars as pl
260260
261261
ctx = FoundryContext()
262262
ds = ctx.get_dataset_by_path("/path/to/test_dataset")
263-
arrow_table = ds.query_foundry_sql("SELECT *",return_type="arrow")
264-
265-
df = pl.from_arrow(arrow_table)
263+
df = ds.query_foundry_sql("SELECT *", return_type="polars")
266264
print(df)
267265
```
268266
````

‎libs/foundry-dev-tools/src/foundry_dev_tools/clients/data_proxy.py‎

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from pathlib import Path
2828

2929
import pandas as pd
30+
import polars as pl
3031
import pyarrow as pa
3132
import pyspark
3233
import requests
@@ -111,7 +112,17 @@ def query_foundry_sql_legacy(
111112
branch: Ref = ...,
112113
sql_dialect: SqlDialect = ...,
113114
timeout: int = ...,
114-
) -> pd.core.frame.DataFrame: ...
115+
) -> pd.DataFrame: ...
116+
117+
@overload
118+
def query_foundry_sql_legacy(
119+
self,
120+
query: str,
121+
return_type: Literal["polars"],
122+
branch: Ref = ...,
123+
sql_dialect: SqlDialect = ...,
124+
timeout: int = ...,
125+
) -> pl.DataFrame: ...
115126

116127
@overload
117128
def query_foundry_sql_legacy(
@@ -151,7 +162,7 @@ def query_foundry_sql_legacy(
151162
branch: Ref = ...,
152163
sql_dialect: SqlDialect = ...,
153164
timeout: int = ...,
154-
) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
165+
) -> tuple[dict, list[list]] | pd.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
155166

156167
def query_foundry_sql_legacy(
157168
self,
@@ -160,7 +171,7 @@ def query_foundry_sql_legacy(
160171
branch: Ref = "master",
161172
sql_dialect: SqlDialect = "SPARK",
162173
timeout: int = 600,
163-
) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pa.Table | pyspark.sql.DataFrame:
174+
) -> tuple[dict, list[list]] | pd.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame:
164175
"""Queries the dataproxy query API with spark SQL.
165176
166177
Example:
@@ -206,20 +217,29 @@ def query_foundry_sql_legacy(
206217
response_json = response.json()
207218
if return_type == "raw":
208219
return response_json["foundrySchema"], response_json["rows"]
220+
# return_type arrow, pandas and polars use the FakeModule implementation in
221+
# their _optional packages. The FakeModule throws an ImportError when trying
222+
# to access attributes of the module, so no need to explicitly catch ImportError.
209223
if return_type == "pandas":
210224
from foundry_dev_tools._optional.pandas import pd
211225

212226
return pd.DataFrame(
213227
data=response_json["rows"],
214228
columns=[e["name"] for e in response_json["foundrySchema"]["fieldSchemaList"]],
215229
)
216-
if return_type == "arrow":
230+
if return_type in {"arrow", "polars"}:
217231
from foundry_dev_tools._optional.pyarrow import pa
218232

219-
return pa.table(
233+
table = pa.table(
220234
data=response_json["rows"],
221235
names=[e["name"] for e in response_json["foundrySchema"]["fieldSchemaList"]],
222236
)
237+
if return_type == "arrow":
238+
return table
239+
240+
from foundry_dev_tools._optional.polars import pl
241+
242+
return pl.from_arrow(table)
223243
if return_type == "spark":
224244
from foundry_dev_tools.utils.converter.foundry_spark import (
225245
foundry_schema_to_spark_schema,

‎libs/foundry-dev-tools/src/foundry_dev_tools/clients/foundry_sql_server.py‎

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
if TYPE_CHECKING:
1919
import pandas as pd
20+
import polars as pl
2021
import pyarrow as pa
2122
import pyspark
2223
import requests
@@ -35,7 +36,7 @@ def query_foundry_sql(
3536
branch: Ref = ...,
3637
sql_dialect: SqlDialect = ...,
3738
timeout: int = ...,
38-
) -> pd.core.frame.DataFrame: ...
39+
) -> pd.DataFrame: ...
3940

4041
@overload
4142
def query_foundry_sql(
@@ -57,6 +58,16 @@ def query_foundry_sql(
5758
timeout: int = ...,
5859
) -> pa.Table: ...
5960

61+
@overload
62+
def query_foundry_sql(
63+
self,
64+
query: str,
65+
return_type: Literal["polars"],
66+
branch: Ref = ...,
67+
sql_dialect: SqlDialect = ...,
68+
timeout: int = ...,
69+
) -> pl.DataFrame: ...
70+
6071
@overload
6172
def query_foundry_sql(
6273
self,
@@ -75,16 +86,16 @@ def query_foundry_sql(
7586
branch: Ref = ...,
7687
sql_dialect: SqlDialect = ...,
7788
timeout: int = ...,
78-
) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
89+
) -> tuple[dict, list[list]] | pd.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
7990

80-
def query_foundry_sql(
91+
def query_foundry_sql( # noqa: C901
8192
self,
8293
query: str,
8394
return_type: SQLReturnType = "pandas",
8495
branch: Ref = "master",
8596
sql_dialect: SqlDialect = "SPARK",
8697
timeout: int = 600,
87-
) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pa.Table | pyspark.sql.DataFrame:
98+
) -> tuple[dict, list[list]] | pd.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame:
8899
"""Queries the Foundry SQL server with spark SQL dialect.
89100
90101
Uses Arrow IPC to communicate with the Foundry SQL Server Endpoint.
@@ -105,9 +116,9 @@ def query_foundry_sql(
105116
timeout: Query Timeout, default value is 600 seconds
106117
107118
Returns:
108-
:external+pandas:py:class:`~pandas.DataFrame` | :external+pyarrow:py:class:`~pyarrow.Table` | :external+spark:py:class:`~pyspark.sql.DataFrame`:
119+
:external+pandas:py:class:`~pandas.DataFrame` | :external+polars:py:class:`~polars.DataFrame` | :external+pyarrow:py:class:`~pyarrow.Table` | :external+spark:py:class:`~pyspark.sql.DataFrame`:
109120
110-
A pandas DataFrame, Spark DataFrame or pyarrow.Table with the result.
121+
A pandas, polars, Spark DataFrame or pyarrow.Table with the result.
111122
112123
Raises:
113124
ValueError: Only direct read eligible queries can be returned as arrow Table.
@@ -139,6 +150,15 @@ def query_foundry_sql(
139150
arrow_stream_reader = self.read_fsql_query_results_arrow(query_id=query_id)
140151
if return_type == "pandas":
141152
return arrow_stream_reader.read_pandas()
153+
if return_type == "polars":
154+
# The FakeModule implementation used in the _optional packages
155+
# throws an ImportError when trying to access attributes of the module.
156+
# This ImportError is caught below to fall back to query_foundry_sql_legacy
157+
# which will again raise an ImportError when polars is not installed.
158+
from foundry_dev_tools._optional.polars import pl
159+
160+
arrow_table = arrow_stream_reader.read_all()
161+
return pl.from_arrow(arrow_table)
142162

143163
if return_type == "spark":
144164
from foundry_dev_tools.utils.converter.foundry_spark import (
@@ -147,16 +167,12 @@ def query_foundry_sql(
147167

148168
return arrow_stream_to_spark_dataframe(arrow_stream_reader)
149169
return arrow_stream_reader.read_all()
150-
151-
return self._query_fsql(
152-
query=query,
153-
branch=branch,
154-
return_type=return_type,
155-
)
156170
except (
157171
FoundrySqlSerializationFormatNotImplementedError,
158172
ImportError,
159173
) as exc:
174+
# Swallow exception when return_type != 'arrow'
175+
# to fall back to query_foundry_sql_legacy
160176
if return_type == "arrow":
161177
msg = (
162178
"Only direct read eligible queries can be returned as arrow Table. Consider using setting"
@@ -166,6 +182,8 @@ def query_foundry_sql(
166182
msg,
167183
) from exc
168184

185+
# this fallback is not only used if return_type is 'raw', but also when one of
186+
# the above exceptions is caught and return_type != 'arrow'
169187
warnings.warn("Falling back to query_foundry_sql_legacy!")
170188
return self.context.data_proxy.query_foundry_sql_legacy(
171189
query=query,

‎libs/foundry-dev-tools/src/foundry_dev_tools/foundry_api_client.py‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
33
One of the gaols of this module is to be self-contained so that it can be
44
dropped into any python installation with minimal dependency to 'requests'
5-
Optional dependencies for the SQL functionality to work are pandas and pyarrow.
5+
Optional dependencies for the SQL functionality to work are pandas, polars and pyarrow.
66
77
"""
88

@@ -41,6 +41,7 @@
4141
from collections.abc import Iterator
4242

4343
import pandas as pd
44+
import polars as pl
4445
import pyarrow as pa
4546
import pyspark
4647
import requests
@@ -1000,7 +1001,7 @@ def query_foundry_sql(
10001001
branch: api_types.Ref = ...,
10011002
sql_dialect: api_types.SqlDialect = ...,
10021003
timeout: int = ...,
1003-
) -> pd.core.frame.DataFrame: ...
1004+
) -> pd.DataFrame: ...
10041005

10051006
@overload
10061007
def query_foundry_sql(
@@ -1022,6 +1023,16 @@ def query_foundry_sql(
10221023
timeout: int = ...,
10231024
) -> pa.Table: ...
10241025

1026+
@overload
1027+
def query_foundry_sql(
1028+
self,
1029+
query: str,
1030+
return_type: Literal["polars"],
1031+
branch: api_types.Ref = ...,
1032+
sql_dialect: api_types.SqlDialect = ...,
1033+
timeout: int = ...,
1034+
) -> pl.DataFrame: ...
1035+
10251036
@overload
10261037
def query_foundry_sql(
10271038
self,
@@ -1040,7 +1051,7 @@ def query_foundry_sql(
10401051
branch: api_types.Ref = ...,
10411052
sql_dialect: api_types.SqlDialect = ...,
10421053
timeout: int = ...,
1043-
) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
1054+
) -> tuple[dict, list[list]] | pd.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
10441055

10451056
def query_foundry_sql(
10461057
self,
@@ -1049,7 +1060,7 @@ def query_foundry_sql(
10491060
branch: api_types.Ref = "master",
10501061
sql_dialect: api_types.SqlDialect = "SPARK",
10511062
timeout: int = 600,
1052-
) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pa.Table | pyspark.sql.DataFrame:
1063+
) -> tuple[dict, list[list]] | pd.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame:
10531064
"""Queries the Foundry SQL server with spark SQL dialect.
10541065
10551066
Uses Arrow IPC to communicate with the Foundry SQL Server Endpoint.

‎libs/foundry-dev-tools/src/foundry_dev_tools/resources/dataset.py‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131

3232
import pandas as pd
3333
import pandas.core.frame
34-
import polars.dataframe.frame
34+
import polars as pl
3535
import pyarrow as pa
3636
import pyspark.sql
3737

@@ -539,7 +539,7 @@ def download_files_temporary(
539539

540540
def save_dataframe(
541541
self,
542-
df: pandas.core.frame.DataFrame | polars.dataframe.frame.DataFrame | pyspark.sql.DataFrame,
542+
df: pandas.core.frame.DataFrame | pl.DataFrame | pyspark.sql.DataFrame,
543543
transaction_type: api_types.FoundryTransaction = "SNAPSHOT",
544544
foundry_schema: api_types.FoundrySchema | None = None,
545545
) -> Self:
@@ -697,7 +697,7 @@ def query_foundry_sql(
697697
return_type: Literal["pandas"],
698698
sql_dialect: api_types.SqlDialect = ...,
699699
timeout: int = ...,
700-
) -> pd.core.frame.DataFrame: ...
700+
) -> pd.DataFrame: ...
701701

702702
@overload
703703
def query_foundry_sql(
@@ -717,6 +717,15 @@ def query_foundry_sql(
717717
timeout: int = ...,
718718
) -> pa.Table: ...
719719

720+
@overload
721+
def query_foundry_sql(
722+
self,
723+
query: str,
724+
return_type: Literal["polars"],
725+
sql_dialect: api_types.SqlDialect = ...,
726+
timeout: int = ...,
727+
) -> pl.DataFrame: ...
728+
720729
@overload
721730
def query_foundry_sql(
722731
self,
@@ -733,15 +742,15 @@ def query_foundry_sql(
733742
return_type: api_types.SQLReturnType = ...,
734743
sql_dialect: api_types.SqlDialect = ...,
735744
timeout: int = ...,
736-
) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
745+
) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
737746

738747
def query_foundry_sql(
739748
self,
740749
query: str,
741750
return_type: api_types.SQLReturnType = "pandas",
742751
sql_dialect: api_types.SqlDialect = "SPARK",
743752
timeout: int = 600,
744-
) -> tuple[dict, list[list]] | pd.DataFrame | pa.Table | pyspark.sql.DataFrame:
753+
) -> tuple[dict, list[list]] | pd.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame:
745754
"""Wrapper around :py:meth:`foundry_dev_tools.clients.foundry_sql_server.FoundrySqlServerClient.query_foundry_sql`.
746755
747756
But it automatically prepends the dataset location, so instead of:
@@ -783,17 +792,12 @@ def to_pandas(self) -> pandas.core.frame.DataFrame:
783792
"""
784793
return self.query_foundry_sql("SELECT *", return_type="pandas")
785794

786-
def to_polars(self) -> polars.dataframe.frame.DataFrame:
795+
def to_polars(self) -> pl.DataFrame:
787796
"""Get dataset as a :py:class:`polars.DataFrame`.
788797
789798
Via :py:meth:`foundry_dev_tools.resources.dataset.Dataset.query_foundry_sql`
790799
"""
791-
try:
792-
import polars as pl
793-
except ImportError as e:
794-
msg = "The optional 'polars' package is not installed. Please install it to use the 'to_polars' method"
795-
raise ImportError(msg) from e
796-
return pl.from_arrow(self.to_arrow())
800+
return self.query_foundry_sql("SELECT *", return_type="polars")
797801

798802
@contextmanager
799803
def transaction_context(

‎libs/foundry-dev-tools/src/foundry_dev_tools/utils/api_types.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,10 +95,11 @@ def assert_in_literal(option, literal, variable_name) -> None: # noqa: ANN001
9595
SqlDialect = Literal["ANSI", "SPARK"]
9696
"""The SQL Dialect for Foundry SQL queries."""
9797

98-
SQLReturnType = Literal["pandas", "spark", "arrow", "raw"]
98+
SQLReturnType = Literal["pandas", "polars", "spark", "arrow", "raw"]
9999
"""The return_types for sql queries.
100100
101101
pandas: :external+pandas:py:class:`pandas.DataFrame`
102+
polars: :external+polars:py:class:`polars.DataFrame`
102103
arrow: :external+pyarrow:py:class:`pyarrow.Table`
103104
spark: :external+spark:py:class:`~pyspark.sql.DataFrame`
104105
raw: Tuple of (foundry_schema, data) (can only be used in legacy)

‎tests/integration/clients/test_foundry_sql_server.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import polars as pl
12
import pytest
23

34
from foundry_dev_tools.errors.dataset import BranchNotFoundError, DatasetHasNoSchemaError, DatasetNotFoundError
@@ -12,6 +13,16 @@ def test_smoke():
1213
assert one_row_one_column.shape == (1, 1)
1314

1415

16+
def test_polars_return_type():
17+
polars_df = TEST_SINGLETON.ctx.foundry_sql_server.query_foundry_sql(
18+
f"SELECT sepal_length FROM `{TEST_SINGLETON.iris_new.rid}` LIMIT 2",
19+
return_type="polars",
20+
)
21+
assert isinstance(polars_df, pl.DataFrame)
22+
assert polars_df.height == 2
23+
assert polars_df.width == 1
24+
25+
1526
def test_exceptions():
1627
with pytest.raises(BranchNotFoundError) as exc:
1728
TEST_SINGLETON.ctx.foundry_sql_server.query_foundry_sql(

‎tests/integration/resources/test_dataset.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,10 @@ def test_crud_dataset(spark_session, tmp_path): # noqa: PLR0915
9797
ds_polars_branch = TEST_SINGLETON.ctx.get_dataset(ds.rid, branch="polars")
9898
ds_polars_branch.save_dataframe(polars_df)
9999
pl_assert_frame_equal(polars_df, ds_polars_branch.to_polars())
100+
pl_assert_frame_equal(
101+
polars_df,
102+
ds_polars_branch.query_foundry_sql("SELECT *", return_type="polars"),
103+
)
100104

101105
ds_spark_branch = TEST_SINGLETON.ctx.get_dataset(ds.rid, branch="spark")
102106
ds_spark_branch.save_dataframe(spark_df)

0 commit comments

Comments
 (0)