Skip to content

Commit

Permalink
feat(duckdb): add dict-like for DuckDB settings
Browse files Browse the repository at this point in the history
Adds a dict-like passthrough object for setting options in the DuckDB
backend.

The repr will show all values as stringified, but this is also how
they'll be passed in.

This avoids users having to pass in options via calls to `raw_sql` or similar.

Co-authored-by: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
  • Loading branch information
gforsyth and cpcloud committed Nov 9, 2023
1 parent 64dddbd commit ea2d317
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 4 deletions.
35 changes: 31 additions & 4 deletions ibis/backends/duckdb/__init__.py
Expand Up @@ -7,10 +7,7 @@
import os
import warnings
from pathlib import Path
from typing import (
TYPE_CHECKING,
Any,
)
from typing import TYPE_CHECKING, Any

import duckdb
import pyarrow as pa
Expand Down Expand Up @@ -69,11 +66,41 @@ def _format_kwargs(kwargs: Mapping[str, Any]):
}


class _Settings:
def __init__(self, con):
self.con = con

def __getitem__(self, key):
try:
with self.con.begin() as con:
return con.exec_driver_sql(
f"select value from duckdb_settings() where name = '{key}'"
).one()
except sa.exc.NoResultFound:
raise KeyError(key)

def __setitem__(self, key, value):
with self.con.begin() as con:
con.exec_driver_sql(f"SET {key}='{value}'")

def __repr__(self):
with self.con.begin() as con:
kv = con.exec_driver_sql(
"select map(array_agg(name), array_agg(value)) from duckdb_settings()"
).scalar()

return repr(dict(zip(kv["key"], kv["value"])))


class Backend(AlchemyCrossSchemaBackend, CanCreateSchema):
name = "duckdb"
compiler = DuckDBSQLCompiler
supports_create_or_replace = True

@property
def settings(self) -> _Settings:
return _Settings(self)

@property
def current_database(self) -> str:
return self._scalar_query(sa.select(sa.func.current_database()))
Expand Down
26 changes: 26 additions & 0 deletions ibis/backends/duckdb/tests/test_client.py
@@ -1,6 +1,7 @@
from __future__ import annotations

import duckdb
import pyarrow as pa
import pytest
import sqlalchemy as sa
from pytest import param
Expand Down Expand Up @@ -132,3 +133,28 @@ def test_create_table_with_timestamp_scales(con, scale):
schema = ibis.schema(dict(ts=dt.Timestamp(scale=scale)))
t = con.create_table(gen_name("duckdb_timestamp_scale"), schema=schema, temp=True)
assert t.schema() == schema


def test_config_options(con):
a_first = {"a": [None, 1]}
a_last = {"a": [1, None]}
nulls_first = pa.Table.from_pydict(a_first, schema=pa.schema([("a", pa.float64())]))
nulls_last = pa.Table.from_pydict(a_last, schema=pa.schema([("a", pa.float64())]))

t = ibis.memtable(a_last)

expr = t.order_by("a")

assert con.to_pyarrow(expr) == nulls_last

con.settings["null_order"] = "nulls_first"

assert con.to_pyarrow(expr) == nulls_first


def test_config_options_bad_option(con):
with pytest.raises(sa.exc.ProgrammingError):
con.settings["not_a_valid_option"] = "oopsie"

with pytest.raises(KeyError):
con.settings["i_didnt_set_this"]

0 comments on commit ea2d317

Please sign in to comment.