Skip to content

Commit

Permalink
feat(polars): allow seamless connection for DataFrame as well as Lazy…
Browse files Browse the repository at this point in the history
…Frame
  • Loading branch information
alexander-beedie authored and cpcloud committed Jun 20, 2023
1 parent 5015677 commit a2a3e45
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions ibis/backends/polars/__init__.py
Expand Up @@ -33,9 +33,9 @@ def __init__(self, *args, **kwargs):
self._context = pl.SQLContext()

def do_connect(
self, tables: MutableMapping[str, pl.LazyFrame] | None = None
self, tables: MutableMapping[str, pl.LazyFrame | pl.DataFrame] | None = None
) -> None:
"""Construct a client from a dictionary of `polars.LazyFrame`s.
"""Construct a client from a dictionary of polars `LazyFrame`s and/or `DataFrame`s.
Parameters
----------
Expand Down Expand Up @@ -128,6 +128,8 @@ def _register_failure(self):
)

def _add_table(self, name: str, obj: pl.LazyFrame | pl.DataFrame) -> None:
if isinstance(obj, pl.DataFrame):
obj = obj.lazy()
self._tables[name] = obj
self._context.register(name, obj)

Expand Down Expand Up @@ -173,7 +175,7 @@ def read_csv(
self._add_table(table_name, pl.scan_csv(path, **kwargs))
except pl.exceptions.ComputeError:
# handles compressed csvs
self._add_table(table_name, pl.read_csv(path, **kwargs).lazy())
self._add_table(table_name, pl.read_csv(path, **kwargs))
return self.table(table_name)

def read_pandas(
Expand Down

0 comments on commit a2a3e45

Please sign in to comment.