Skip to content

Commit

Permalink
Fix SQL vs. Polars backend detection for pydiverse.transform
Browse files Browse the repository at this point in the history
  • Loading branch information
windiana42 committed Feb 18, 2025
1 parent 77d0b71 commit f04b072
Showing 1 changed file with 8 additions and 11 deletions.
19 changes: 8 additions & 11 deletions src/pydiverse/pipedag/backend/table/sql/hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -940,22 +940,19 @@ def materialize(cls, store, table: Table[pdt.Table], stage_name):

t = table.obj
table = table.copy_without_obj()
try:
# detect SQL by checking whether build_query() succeeds
query = t >> build_query()
except Exception:
query = t >> build_query()
# detect SQL by checking whether build_query() succeeds
if query is not None:
# continue with SQL case handling
table.obj = sa.text(str(query))
hook = store.get_hook_subclass(SQLAlchemyTableHook)
return hook.materialize(store, table, stage_name)
else:
# use Polars for dataframe handling
table.obj = t >> export(Polars())
hook = store.get_hook_subclass(PolarsTableHook)
return hook.materialize(store, table, stage_name)

# continue with SQL case handling
table.obj = sa.text(query)
hook = store.get_hook_subclass(SQLAlchemyTableHook)
return hook.materialize(store, table, stage_name)

raise NotImplementedError

@classmethod
def retrieve(
cls,
Expand Down

0 comments on commit f04b072

Please sign in to comment.