fix(polars): avoid using unnecessary subquery for schema inference

ibis-project · Jan 27, 2024 · 0f43667 · 0f43667
1 parent e14107c
commit 0f43667
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 21 deletions.
diff --git a/ibis/backends/polars/__init__.py b/ibis/backends/polars/__init__.py
@@ -395,9 +395,7 @@ def compile(
         return translate(node, ctx=ctx)
 
     def _get_schema_using_query(self, query: str) -> sch.Schema:
-        return schema_from_polars(
-            self._context.execute(f"SELECT * FROM ({query}) LIMIT 0").schema
-        )
+        return schema_from_polars(self._context.execute(query).schema)
 
     def execute(
         self,

diff --git a/ibis/backends/tests/test_client.py b/ibis/backends/tests/test_client.py
@@ -189,7 +189,7 @@ def test_query_schema(ddl_backend, expr_fn, expected):
 }
 
 
-@pytest.mark.notimpl(["datafusion", "polars", "mssql"])
+@pytest.mark.notimpl(["datafusion", "mssql"])
 @pytest.mark.never(["dask", "pandas"], reason="dask and pandas do not support SQL")
 @pytest.mark.notimpl(
     ["flink"], raises=AttributeError, reason="'Backend' object has no attribute 'sql'"

diff --git a/ibis/backends/tests/test_dot_sql.py b/ibis/backends/tests/test_dot_sql.py
@@ -33,14 +33,7 @@
 @pytest.mark.parametrize(
     "schema",
     [
-        param(
-            None,
-            id="implicit_schema",
-            marks=[
-                pytest.mark.notimpl(["druid"]),
-                pytest.mark.notyet(["polars"], raises=PolarsComputeError),
-            ],
-        ),
+        param(None, id="implicit_schema", marks=[pytest.mark.notimpl(["druid"])]),
         param({"s": "string", "new_col": "double"}, id="explicit_schema"),
     ],
 )
@@ -82,7 +75,6 @@ def test_con_dot_sql(backend, con, schema):
 @dot_sql_notimpl
 @dot_sql_notyet
 @dot_sql_never
-@pytest.mark.notyet(["polars"], raises=PolarsComputeError)
 def test_table_dot_sql(backend, con):
     alltypes = con.table("functional_alltypes")
     t = (
@@ -120,7 +112,6 @@ def test_table_dot_sql(backend, con):
 @dot_sql_notimpl
 @dot_sql_notyet
 @dot_sql_never
-@pytest.mark.notyet(["polars"], raises=PolarsComputeError)
 def test_table_dot_sql_with_join(backend, con):
     alltypes = con.table("functional_alltypes")
     t = (
@@ -168,7 +159,6 @@ def test_table_dot_sql_with_join(backend, con):
 @dot_sql_notimpl
 @dot_sql_notyet
 @dot_sql_never
-@pytest.mark.notyet(["polars"], raises=PolarsComputeError)
 def test_table_dot_sql_repr(con):
     alltypes = con.table("functional_alltypes")
     t = (
@@ -235,12 +225,8 @@ def test_dot_sql_reuse_alias_with_different_types(backend, alltypes, df):
 
 @pytest.mark.parametrize(
     "dialect",
-    [
-        *sorted(_get_backend_names() - _NO_SQLGLOT_DIALECT),
-        *no_sqlglot_dialect,
-    ],
+    [*sorted(_get_backend_names() - _NO_SQLGLOT_DIALECT), *no_sqlglot_dialect],
 )
-@pytest.mark.notyet(["polars"], raises=PolarsComputeError)
 @table_dot_sql_notimpl
 @dot_sql_notimpl
 @dot_sql_notyet
@@ -264,7 +250,6 @@ def test_table_dot_sql_transpile(backend, alltypes, dialect, df):
         *no_sqlglot_dialect,
     ],
 )
-@pytest.mark.notyet(["polars"], raises=PolarsComputeError)
 @pytest.mark.notyet(["druid"], raises=ValueError)
 @pytest.mark.notyet(["snowflake", "bigquery"])
 @pytest.mark.notyet(