Skip to content

Commit

Permalink
feat(duckdb): support sqlalchemy 2
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud authored and gforsyth committed Mar 16, 2023
1 parent 7d29c63 commit 679bb52
Show file tree
Hide file tree
Showing 5 changed files with 26 additions and 12 deletions.
6 changes: 5 additions & 1 deletion .github/workflows/ibis-backends.yml
Original file line number Diff line number Diff line change
Expand Up @@ -454,7 +454,7 @@ jobs:
- run: python -m pip install --upgrade pip 'poetry<1.4'

- name: remove deps that are not compatible with sqlalchemy 2
run: poetry remove duckdb-engine snowflake-sqlalchemy
run: poetry remove snowflake-sqlalchemy

- name: add sqlalchemy 2
run: poetry add --lock --optional 'sqlalchemy>=2,<3'
Expand Down Expand Up @@ -522,6 +522,10 @@ jobs:
extras:
- trino
- postgres
- name: duckdb
title: DuckDB
extras:
- duckdb
steps:
- name: checkout
uses: actions/checkout@v3
Expand Down
11 changes: 10 additions & 1 deletion ibis/backends/duckdb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -525,7 +525,16 @@ def to_pyarrow_batches(
query_ast = self.compiler.to_ast_ensure_limit(expr, limit, params=params)
sql = query_ast.compile()

cursor = self.con.connect().execute(sql)
con = self.con.connect()

# end the current transaction started by sqlalchemy; without this
# duckdb-engine raises an exception disallowing nested transactions
#
# not clear if the value of returning a RecordBatchReader versus an
# iterator of record batches is worth the cursor leakage here
con.exec_driver_sql("COMMIT")

cursor = con.execute(sql)

reader = cursor.cursor.fetch_record_batch(chunk_size=chunk_size)
return IbisRecordBatchReader(reader, cursor)
Expand Down
11 changes: 6 additions & 5 deletions ibis/backends/duckdb/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,23 +53,24 @@ def _round(t, op):
}


def _generic_log(arg, base):
return sa.func.ln(arg) / sa.func.ln(base)
def _generic_log(arg, base, *, type_):
return sa.func.ln(arg, type_=type_) / sa.func.ln(base, type_=type_)


def _log(t, op):
arg, base = op.args
sqla_type = t.get_sqla_type(op.output_dtype)
sa_arg = t.translate(arg)
if base is not None:
sa_base = t.translate(base)
try:
base_value = sa_base.value
except AttributeError:
return _generic_log(sa_arg, sa_base)
return _generic_log(sa_arg, sa_base, type_=sqla_type)
else:
func = _LOG_BASE_FUNCS.get(base_value, _generic_log)
return func(sa_arg)
return sa.func.ln(sa_arg)
return func(sa_arg, type_=sqla_type)
return sa.func.ln(sa_arg, type_=sqla_type)


def _timestamp_from_unix(t, op):
Expand Down
8 changes: 4 additions & 4 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion requirements.txt

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 679bb52

Please sign in to comment.