support polars in v2 as well.

Author: nicornk

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

@@ -322,6 +322,16 @@ def query_foundry_sql(
         timeout: int = ...,
     ) -> pd.core.frame.DataFrame: ...
 
+    @overload
+    def query_foundry_sql(
+        self,
+        query: str,
+        return_type: Literal["polars"],
+        branch: Ref = ...,
+        sql_dialect: SqlDialect = ...,
+        timeout: int = ...,
+    ) -> pl.DataFrame: ...
+
     @overload
     def query_foundry_sql(
         self,
@@ -350,15 +360,15 @@ def query_foundry_sql(
         return_type: SQLReturnType = ...,
         disable_arrow_compression: bool = ...,
         timeout: int = ...,
-    ) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
+    ) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
 
     def query_foundry_sql(
         self,
         query: str,
         return_type: SQLReturnType = "pandas",
         disable_arrow_compression: bool = False,
         application_id: str | None = None,
-    ) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pa.Table | pyspark.sql.DataFrame:
+    ) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame:
         """Queries the Foundry SQL server using the V2 API.
 
         Uses Arrow IPC to communicate with the Foundry SQL Server Endpoint.
@@ -376,9 +386,9 @@ def query_foundry_sql(
             application_id: The application/dataset RID, defaults to foundry-dev-tools User-Agent
 
         Returns:
-            :external+pandas:py:class:`~pandas.DataFrame` | :external+pyarrow:py:class:`~pyarrow.Table` | :external+spark:py:class:`~pyspark.sql.DataFrame`:
+            :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`:
 
-            A pandas DataFrame, Spark DataFrame or pyarrow.Table with the result.
+            A pandas DataFrame, polars, Spark DataFrame or pyarrow.Table with the result.
 
         Raises:
             FoundrySqlQueryFailedError: If the query fails
@@ -414,6 +424,16 @@ def query_foundry_sql(
         if return_type == "pandas":
             return arrow_stream_reader.read_pandas()
 
+        if return_type == "polars":
+            # The FakeModule implementation used in the _optional packages
+            # throws an ImportError when trying to access attributes of the module.
+            # This ImportError is caught below to fall back to query_foundry_sql_legacy
+            # which will again raise an ImportError when polars is not installed.
+            from foundry_dev_tools._optional.polars import pl
+
+            arrow_table = arrow_stream_reader.read_all()
+            return pl.from_arrow(arrow_table)
+
         if return_type == "spark":
             from foundry_dev_tools.utils.converter.foundry_spark import (
                 arrow_stream_to_spark_dataframe,

tests/integration/clients/test_foundry_sql_server.py

@@ -176,3 +176,13 @@ def test_v2_with_where_clause():
     # Verify all returned rows have is_setosa = 'setosa'
     if result.shape[0] > 0:
         assert all(result["is_setosa"] == "setosa")
+
+
+def test_v2_polars_return_type():
+    polars_df = TEST_SINGLETON.ctx.foundry_sql_server_v2.query_foundry_sql(
+        f"SELECT sepal_length FROM `{TEST_SINGLETON.iris_new.rid}` LIMIT 2",
+        return_type="polars",
+    )
+    assert isinstance(polars_df, pl.DataFrame)
+    assert polars_df.height == 2
+    assert polars_df.width == 1