feat: complete Arrow type mapping for parametric and missing types

src/ibis_hotdata/types.py

@@ -2,53 +2,145 @@
 
 from __future__ import annotations
 
+import re
+
+import pyarrow as pa
 import ibis.expr.datatypes as dt
 from ibis.backends.sql.datatypes import PostgresType
+from ibis.formats.pyarrow import PyArrowType
 
-# Arrow-style type names returned by Hotdata's information_schema when tables are
-# loaded from Parquet/Arrow sources.  PostgresType.from_string() treats these as
-# USERDEFINED unknowns, so we resolve them explicitly before falling through.
-_ARROW_TYPE_MAP: dict[str, type[dt.DataType]] = {
+# Simple Arrow type strings → PyArrow instances.  Covers non-parametric types
+# that the Postgres dialect parser does not know (Arrow-specific names, unsigned
+# ints) or mis-maps (Arrow "int8" = 8-bit; Postgres "int8" = 8-byte / int64).
+# All scalar types that can appear as list/map element types must be listed here
+# because element type strings are resolved via this map, not the Postgres parser.
+_PA_TYPE_MAP: dict[str, pa.DataType] = {
     # dates
-    "date32": dt.Date,
-    "date64": dt.Date,
-    # floats
-    "float16": dt.Float16,
-    "float32": dt.Float32,
-    "float64": dt.Float64,
-    # unsigned ints
-    "uint8": dt.UInt8,
-    "uint16": dt.UInt16,
-    "uint32": dt.UInt32,
-    "uint64": dt.UInt64,
-    # strings
-    "utf8": dt.String,
-    "largeutf8": dt.String,
+    "date32": pa.date32(),
+    "date64": pa.date64(),
+    # floats — "halffloat" is PyArrow's str() name for float16
+    "float16": pa.float16(),
+    "float32": pa.float32(),
+    "float64": pa.float64(),
+    "halffloat": pa.float16(),
+    # signed ints — Arrow "int8" ≠ Postgres "int8" (8-byte/int64); all four
+    # listed here so they resolve correctly when used as list element types
+    "int8": pa.int8(),
+    "int16": pa.int16(),
+    "int32": pa.int32(),
+    "int64": pa.int64(),
+    # unsigned ints (Postgres parser returns Unknown for all of these)
+    "uint8": pa.uint8(),
+    "uint16": pa.uint16(),
+    "uint32": pa.uint32(),
+    "uint64": pa.uint64(),
+    # strings — large-offset variants not known to the Postgres parser
+    "utf8": pa.utf8(),
+    "largeutf8": pa.large_utf8(),
+    "large_string": pa.large_utf8(),
+    "string": pa.string(),
     # binary
-    "largebinary": dt.Binary,
-    # time
-    "time32": dt.Time,
-    "time64": dt.Time,
+    "binary": pa.binary(),
+    "largebinary": pa.large_binary(),
+    # boolean / null
+    "bool": pa.bool_(),
+    "boolean": pa.bool_(),
+    "null": pa.null(),
+    # time — unit is absent from these bare string forms; the unit does not
+    # affect the Ibis type (both time32 and time64 map to dt.Time)
+    "time32": pa.time32("ms"),
+    "time64": pa.time64("us"),
 }
 
+# Regex patterns for parametric Arrow types that embed parameters in the string.
+_TIMESTAMP_RE = re.compile(r"^timestamp\[(\w+)(?:,\s*tz=(.+))?\]$", re.IGNORECASE)
+_DURATION_RE = re.compile(r"^duration\[(\w+)\]$", re.IGNORECASE)
+_DECIMAL_RE = re.compile(r"^decimal(?:128|256)?\((\d+),\s*(\d+)\)$", re.IGNORECASE)
+_LIST_RE = re.compile(r"^(large_)?list<item:\s*(.+)>$", re.IGNORECASE)
+# PyArrow appends " not null" when a list's item field is non-nullable.
+_NOT_NULL_SUFFIX_RE = re.compile(r"\s+not\s+null$", re.IGNORECASE)
+
+
+def _pa_type_from_arrow_str(raw: str) -> pa.DataType | None:
+    """Best-effort: Arrow type string → PyArrow DataType, or ``None`` if not recognised.
+
+    Handles simple names (via ``_PA_TYPE_MAP``) and parametric forms
+    (timestamp, duration, decimal, list/large_list).  Returns ``None`` if the
+    string is not a known Arrow type, allowing the caller to fall through to the
+    Postgres dialect parser or String fallback.
+    """
+    s = raw.strip()
+
+    # Simple non-parametric types.
+    pa_type = _PA_TYPE_MAP.get(s.lower())
+    if pa_type is not None:
+        return pa_type
+
+    # timestamp[unit] or timestamp[unit, tz=…]
+    m = _TIMESTAMP_RE.match(s)
+    if m:
+        unit = m.group(1).lower()
+        tz: str | None = m.group(2).strip() if m.group(2) else None
+        try:
+            return pa.timestamp(unit, tz=tz)
+        except Exception:
+            return None
+
+    # duration[unit] — unknown units return None so the caller falls through
+    m = _DURATION_RE.match(s)
+    if m:
+        try:
+            return pa.duration(m.group(1).lower())
+        except Exception:
+            return None
+
+    # decimal / decimal128 / decimal256
+    m = _DECIMAL_RE.match(s)
+    if m:
+        precision, scale = int(m.group(1)), int(m.group(2))
+        try:
+            # decimal128 supports precision 1–38; fall back to decimal256 for wider values
+            return pa.decimal128(precision, scale) if precision <= 38 else pa.decimal256(precision, scale)
+        except Exception:
+            return None
+
+    # list<item: T> or large_list<item: T> (recursive for nested types)
+    m = _LIST_RE.match(s)
+    if m:
+        is_large = bool(m.group(1))
+        item_raw = m.group(2).strip()
+        item_not_null = bool(_NOT_NULL_SUFFIX_RE.search(item_raw))
+        item_str = _NOT_NULL_SUFFIX_RE.sub("", item_raw).strip()
+        item_pa_type = _pa_type_from_arrow_str(item_str)
+        if item_pa_type is None:
+            return None
+        item_field = pa.field("item", item_pa_type, nullable=not item_not_null)
+        return pa.large_list(item_field) if is_large else pa.list_(item_field)
+
+    return None
+
 
 def dtype_from_hotdata_sql_type(sql_type: str | None, *, nullable: bool) -> dt.DataType:
-    """Best-effort mapping from Hotdata `/information_schema` column `data_type` strings.
+    """Best-effort mapping from Hotdata ``/information_schema`` column ``data_type`` strings.
 
     Hotdata may return either SQL-style names (``INTEGER``, ``VARCHAR``, ``DOUBLE
     PRECISION``, …) or Arrow-style names (``Date32``, ``Float64``, ``Utf8``, …).
-    SQL-style names are delegated to the Postgres dialect parser; Arrow-style names
-    are resolved via an explicit lookup table before falling back to the parser.
+    Arrow-style names are resolved via PyArrow's type system and converted to Ibis
+    types using the Ibis–PyArrow bridge; SQL-style names fall through to the Postgres
+    dialect parser.
     """
     if not sql_type:
         return dt.String(nullable=nullable)
 
-    # Arrow-style names (case-insensitive lookup).
-    arrow_cls = _ARROW_TYPE_MAP.get(sql_type.strip().lower())
-    if arrow_cls is not None:
-        return arrow_cls(nullable=nullable)
+    raw = sql_type.strip()
+
+    # Try to parse as an Arrow type string (simple or parametric).
+    pa_type = _pa_type_from_arrow_str(raw)
+    if pa_type is not None:
+        return PyArrowType.to_ibis(pa_type).copy(nullable=nullable)
 
+    # Fall through to Postgres dialect parser for SQL-style type names.
     try:
-        return PostgresType.from_string(sql_type.strip(), nullable=nullable)
+        return PostgresType.from_string(raw, nullable=nullable)
     except Exception:  # ibis/sqlglot raise a variety of parse errors; fall back to String
         return dt.String(nullable=nullable)

tests/test_hotdata_types.py

@@ -55,6 +55,15 @@ def test_dtype_from_hotdata_malformed_fallback_string():
         ("LargeBinary", True, dt.Binary),
         ("Time32", True, dt.Time),
         ("Time64", False, dt.Time),
+        # Previously missing: signed int8 (Postgres "int8" means int64, not int8)
+        ("int8", True, dt.Int8),
+        ("Int8", False, dt.Int8),
+        # Previously missing: halffloat (PyArrow's str() name for float16)
+        ("halffloat", True, dt.Float16),
+        ("HALFFLOAT", False, dt.Float16),
+        # Previously missing: large_string (PyArrow large-offset string variant)
+        ("large_string", True, dt.String),
+        ("Large_String", False, dt.String),
         # Case-insensitive
         ("date32", True, dt.Date),
         ("FLOAT64", True, dt.Float64),
@@ -65,3 +74,95 @@ def test_dtype_from_hotdata_arrow_type_names(sql_type, nullable, expected_cls):
     out = dtype_from_hotdata_sql_type(sql_type, nullable=nullable)
     assert out.nullable is nullable
     assert isinstance(out, expected_cls)
+
+
+@pytest.mark.parametrize(
+    ("sql_type", "expected_tz", "expected_scale"),
+    [
+        ("timestamp[s]", None, 0),
+        ("timestamp[ms]", None, 3),
+        ("timestamp[us]", None, 6),
+        ("timestamp[ns]", None, 9),
+        ("timestamp[us, tz=UTC]", "UTC", 6),
+        ("timestamp[us, tz=America/New_York]", "America/New_York", 6),
+        ("TIMESTAMP[MS]", None, 3),
+    ],
+)
+def test_dtype_from_hotdata_arrow_timestamp(sql_type, expected_tz, expected_scale):
+    out = dtype_from_hotdata_sql_type(sql_type, nullable=True)
+    assert isinstance(out, dt.Timestamp)
+    assert out.timezone == expected_tz
+    assert out.scale == expected_scale
+    assert out.nullable is True
+
+
+@pytest.mark.parametrize(
+    ("sql_type", "expected_unit"),
+    [
+        ("duration[s]", "s"),
+        ("duration[ms]", "ms"),
+        ("duration[us]", "us"),
+        ("duration[ns]", "ns"),
+        ("DURATION[MS]", "ms"),
+    ],
+)
+def test_dtype_from_hotdata_arrow_duration(sql_type, expected_unit):
+    out = dtype_from_hotdata_sql_type(sql_type, nullable=False)
+    assert isinstance(out, dt.Interval)
+    assert out.unit.value == expected_unit
+    assert out.nullable is False
+
+
+def test_dtype_from_hotdata_arrow_duration_unknown_unit_falls_back():
+    # An unrecognised duration unit should not silently map to seconds;
+    # it falls through to the Postgres parser (which returns Unknown) or String fallback.
+    out = dtype_from_hotdata_sql_type("duration[foo]", nullable=True)
+    assert not isinstance(out, dt.Interval)  # must not produce a valid Interval
+
+
+@pytest.mark.parametrize(
+    ("sql_type", "expected_precision", "expected_scale"),
+    [
+        ("decimal128(10, 3)", 10, 3),
+        ("decimal128(38, 0)", 38, 0),
+        ("decimal256(76, 38)", 76, 38),
+        ("decimal(5, 2)", 5, 2),
+        ("DECIMAL128(18, 6)", 18, 6),
+        # decimal12 is NOT a valid form — should not be matched by the decimal regex
+    ],
+)
+def test_dtype_from_hotdata_arrow_decimal(sql_type, expected_precision, expected_scale):
+    out = dtype_from_hotdata_sql_type(sql_type, nullable=True)
+    assert isinstance(out, dt.Decimal)
+    assert out.precision == expected_precision
+    assert out.scale == expected_scale
+    assert out.nullable is True
+
+
+@pytest.mark.parametrize(
+    ("sql_type", "expected_value_cls", "expected_item_nullable"),
+    [
+        ("list<item: int32>", dt.Int32, True),
+        ("list<item: utf8>", dt.String, True),
+        ("list<item: float64>", dt.Float64, True),
+        ("large_list<item: int64>", dt.Int64, True),
+        ("LIST<ITEM: UINT8>", dt.UInt8, True),
+        # Non-nullable item fields — PyArrow appends " not null"
+        ("list<item: int32 not null>", dt.Int32, False),
+        ("list<item: utf8 not null>", dt.String, False),
+        ("large_list<item: float32 not null>", dt.Float32, False),
+    ],
+)
+def test_dtype_from_hotdata_arrow_list(sql_type, expected_value_cls, expected_item_nullable):
+    out = dtype_from_hotdata_sql_type(sql_type, nullable=True)
+    assert isinstance(out, dt.Array)
+    assert isinstance(out.value_type, expected_value_cls)
+    assert out.value_type.nullable is expected_item_nullable
+    assert out.nullable is True
+
+
+def test_dtype_from_hotdata_arrow_decimal12_not_matched():
+    # "decimal12" (only the trailing 8 made optional) must NOT match the decimal regex.
+    # The Postgres parser handles bare "decimal" forms; decimal12 is not a real type.
+    out = dtype_from_hotdata_sql_type("decimal12(10, 3)", nullable=True)
+    assert not isinstance(out, dt.Decimal)  # falls through to Unknown or String