Skip to content

Commit

Permalink
fix(duckdb): handle conversion from duckdb_engine unsigned int aliases
Browse files Browse the repository at this point in the history
  • Loading branch information
gforsyth authored and cpcloud committed Jul 18, 2023
1 parent 207587c commit e6fd0cc
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 0 deletions.
8 changes: 8 additions & 0 deletions ibis/backends/duckdb/datatypes.py
Expand Up @@ -101,9 +101,17 @@ def parse(text: str, default_decimal_parameters=(18, 3)) -> dt.DataType:
ducktypes.SmallInteger: dt.Int16,
ducktypes.Integer: dt.Int32,
ducktypes.BigInteger: dt.Int64,
ducktypes.HugeInteger: dt.Int64,
# The following dtypes are present in duckdb_engine and will show up
# in _some_ tables, but it isn't clear _why_.
# Please don't remove them.
ducktypes.UInt8: dt.UInt8,
ducktypes.UTinyInteger: dt.UInt8,
ducktypes.UInt16: dt.UInt16,
ducktypes.USmallInteger: dt.UInt16,
ducktypes.UInt32: dt.UInt32,
ducktypes.UInteger: dt.UInt32,
ducktypes.UInt64: dt.UInt64,
ducktypes.UBigInteger: dt.UInt64,
}

Expand Down
20 changes: 20 additions & 0 deletions ibis/backends/duckdb/tests/test_datatypes.py
Expand Up @@ -133,3 +133,23 @@ def test_generate_quoted_struct():
result = typ.compile(dialect=duckdb_engine.Dialect())
expected = 'STRUCT("in come" TEXT, "my count" BIGINT, thing INTEGER)'
assert result == expected


def test_read_uint8_from_parquet(tmp_path):
import ibis

con = ibis.connect("duckdb://:memory:")

t = ibis.memtable({"a": [1, 2, 3, 4]})
# There is an incorrect mapping in duckdb-engine from UInteger -> UInt8
# In order to get something that reads as a UInt8, we cast to UInt32 (UInteger)
t = t.mutate(a=t.a.cast("uint32"))

parqpath = tmp_path / "uint.parquet"

t.to_parquet(parqpath)

# If this doesn't fail, then things are working
t2 = con.read_parquet(parqpath)

assert t2.a.type() == dt.uint8

0 comments on commit e6fd0cc

Please sign in to comment.