Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(polars): use drop_table when cleaning the cache and remove duplicated _remove_table method #9922

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 2 additions & 5 deletions ibis/backends/polars/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,10 +148,6 @@ def _add_table(self, name: str, obj: pl.LazyFrame | pl.DataFrame) -> None:
self._tables[name] = obj
self._context.register(name, obj)

def _remove_table(self, name: str) -> None:
del self._tables[name]
self._context.unregister(name)

def sql(
self, query: str, schema: sch.Schema | None = None, dialect: str | None = None
) -> ir.Table:
Expand Down Expand Up @@ -407,6 +403,7 @@ def create_view(
def drop_table(self, name: str, *, force: bool = False) -> None:
if name in self._tables:
del self._tables[name]
self._context.unregister(name)
elif not force:
raise com.IbisError(f"Table {name!r} does not exist")

Expand Down Expand Up @@ -559,7 +556,7 @@ def _load_into_cache(self, name, expr):
self.create_table(name, self.compile(expr).cache())

def _clean_up_cached_table(self, name):
self._remove_table(name)
self.drop_table(name, force=True)


@lazy_singledispatch
Expand Down
24 changes: 24 additions & 0 deletions ibis/backends/polars/tests/test_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
from __future__ import annotations

import pytest

from ibis.backends.tests.errors import PolarsSQLInterfaceError
from ibis.util import gen_name


def test_cannot_run_sql_after_drop(con):
t = con.table("functional_alltypes")
n = t.count().execute()

name = gen_name("polars_dot_sql")
con.create_table(name, t)

sql = f"SELECT COUNT(*) FROM {name}"

expr = con.sql(sql)
result = expr.execute()
assert result.iat[0, 0] == n

con.drop_table(name)
with pytest.raises(PolarsSQLInterfaceError):
con.sql(sql)
3 changes: 2 additions & 1 deletion ibis/backends/tests/errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,10 +65,11 @@
from polars.exceptions import InvalidOperationError as PolarsInvalidOperationError
from polars.exceptions import PanicException as PolarsPanicException
from polars.exceptions import SchemaError as PolarsSchemaError
from polars.exceptions import SQLInterfaceError as PolarsSQLInterfaceError
except ImportError:
PolarsComputeError = PolarsPanicException = PolarsInvalidOperationError = (
PolarsSchemaError
) = PolarsColumnNotFoundError = None
) = PolarsColumnNotFoundError = PolarsSQLInterfaceError = None
jcrist marked this conversation as resolved.
Show resolved Hide resolved

try:
from pyarrow import ArrowInvalid, ArrowNotImplementedError
Expand Down