sqlalchemy-monetdb-adbc 0.1.0__py3-none-any.whl

sqlalchemy_monetdb_adbc/__init__.py

@@ -0,0 +1,37 @@
+from importlib.metadata import PackageNotFoundError, version
+
+from sqlalchemy_monetdb_adbc.arrow import fetch_arrow_table, fetch_record_batches, ingest_arrow
+from sqlalchemy_monetdb_adbc.connection import raw_adbc_connection
+from sqlalchemy_monetdb_adbc.dialect import MonetDBADBCDialect
+from sqlalchemy_monetdb_adbc.types import (
+    DOUBLE_PRECISION,
+    HUGEINT,
+    INET,
+    MONTH_INTERVAL,
+    SECOND_INTERVAL,
+    TINYINT,
+    URL,
+    PydanticJSON,
+)
+
+try:
+    __version__ = version("sqlalchemy-monetdb-adbc")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+__all__ = [
+    "DOUBLE_PRECISION",
+    "HUGEINT",
+    "INET",
+    "MONTH_INTERVAL",
+    "SECOND_INTERVAL",
+    "TINYINT",
+    "URL",
+    "MonetDBADBCDialect",
+    "PydanticJSON",
+    "__version__",
+    "fetch_arrow_table",
+    "fetch_record_batches",
+    "ingest_arrow",
+    "raw_adbc_connection",
+]

sqlalchemy_monetdb_adbc/_alembic.py

@@ -0,0 +1,74 @@
+"""Alembic support.
+
+Alembic resolves a migration implementation by dialect name from the subclasses
+of ``DefaultImpl`` that have been imported, so defining this class registers
+MonetDB. Alembic is not a runtime dependency, so ``dialect`` imports this module
+inside a ``try``.
+"""
+
+from typing import Any, Literal, cast
+
+from alembic.ddl.base import (
+    ColumnType,
+    alter_column,
+    alter_table,
+    format_type,  # pyright: ignore[reportUnknownVariableType]
+)
+from alembic.ddl.impl import DefaultImpl
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.compiler import DDLCompiler
+from sqlalchemy.sql.type_api import TypeEngine
+
+from sqlalchemy_monetdb_adbc.types import PydanticJSON
+
+
+def _unquote_string_default(value: str) -> str:
+    if len(value) >= 2 and value[0] == value[-1] == "'":
+        return value[1:-1].replace("''", "'")
+    return value
+
+
+class MonetDBImpl(DefaultImpl):
+    __dialect__ = "monetdb"
+
+    transactional_ddl = True
+
+    def render_type(self, type_obj: Any, autogen_context: Any) -> str | Literal[False]:
+        # A migration describes the database schema, so a PydanticJSON column
+        # renders as the JSON column it actually is. Rendering the model would
+        # make migrations import application code and break as soon as that
+        # model is renamed or removed.
+        if isinstance(type_obj, PydanticJSON):
+            autogen_context.imports.add("import sqlalchemy as sa")
+            return "sa.JSON()"
+        return False
+
+    def compare_server_default(
+        self,
+        inspector_column: Any,
+        metadata_column: Any,
+        rendered_metadata_default: str | None,
+        rendered_inspector_default: str | None,
+    ) -> bool:
+        # MonetDB echoes a string default back with its quotes, so compare the
+        # unquoted forms rather than reporting every such column as changed.
+        if rendered_inspector_default is not None:
+            rendered_inspector_default = _unquote_string_default(rendered_inspector_default)
+        if rendered_metadata_default is not None:
+            rendered_metadata_default = _unquote_string_default(rendered_metadata_default)
+        return rendered_inspector_default != rendered_metadata_default
+
+
+@compiles(ColumnType, "monetdb")
+def _monetdb_column_type(  # pyright: ignore[reportUnusedFunction]
+    element: ColumnType,
+    compiler: DDLCompiler,
+    **kw: Any,  # noqa: ARG001
+) -> str:
+    # MonetDB spells this "ALTER COLUMN <name> <type>", without the TYPE keyword
+    # the default rendering emits.
+    table = alter_table(compiler, element.table_name, element.schema)
+    column = alter_column(compiler, element.column_name)
+    column_type = cast(TypeEngine[Any], element.type_)  # pyright: ignore[reportUnknownMemberType]
+    rendered = format_type(compiler, column_type)
+    return f"{table} {column} {rendered}"

sqlalchemy_monetdb_adbc/_convert.py

@@ -0,0 +1,107 @@
+"""Fast Arrow column to Python list conversion.
+
+``pyarrow.Array.to_pylist`` boxes every element through a generic scalar
+conversion. For bulk reads that cost can dominate pulling the result off the
+wire. ``to_numpy(zero_copy_only=False).tolist()`` does the same work one column
+at a time in compiled code.
+
+This is why numpy is a dependency rather than an optional accelerator: without
+it the fastest available path is reading the Arrow values buffer through
+``memoryview.cast``, which matches numpy on fixed-width numerics but does
+nothing for the string and temporal columns where most of the cost actually is.
+
+The fast path is guarded, and anything not provably identical to ``to_pylist``
+falls back to it. Each exclusion is a case where numpy returns a different
+Python type or value, and every one is asserted in ``tests/test_convert.py``:
+
+* Integer and floating-point columns containing nulls. numpy has no integer
+  null and represents both cases with ``nan``.
+* Nanosecond timestamps and times, which come back as a raw ``int`` count
+  because numpy cannot represent a nanosecond instant as a ``datetime``.
+* ``date64``, which widens ``datetime.date`` to ``datetime.datetime``.
+* Nested and extension types, which are simply not whitelisted.
+
+For a single non-temporal value, the generic scalar conversion is faster than
+setting up numpy, so point-query batches take that path directly.
+"""
+
+
+# pyarrow ships no type information, so its data types, arrays, and predicates
+# all read as unknown under strict mode. As in ``arrow.py``, Arrow objects are
+# annotated ``Any`` and their concrete types given per function.
+# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false
+
+import datetime
+from typing import Any, Final, cast
+
+import pyarrow as pa
+
+# Temporal units numpy converts to the same Python object ``to_pylist`` returns.
+_SAFE_TEMPORAL_UNITS: Final = frozenset({"s", "ms", "us"})
+_ARROW_EXTENSION_NAME: Final = b"ARROW:extension:name"
+_MONETDB_HUGEINT_EXTENSION: Final = b"monetdb.hugeint"
+
+
+def _numpy_safe(data_type: Any) -> bool:
+    """Whether ``to_numpy().tolist()`` round-trips this ``pyarrow.DataType`` exactly.
+
+    A whitelist rather than a blacklist: an unrecognized type falls back to
+    ``to_pylist`` instead of silently converting through an unverified path.
+    """
+    if pa.types.is_timestamp(data_type):
+        return data_type.unit in _SAFE_TEMPORAL_UNITS
+
+    if pa.types.is_time(data_type):
+        return data_type.unit in _SAFE_TEMPORAL_UNITS
+
+    if pa.types.is_date(data_type):
+        return bool(pa.types.is_date32(data_type))
+
+    return bool(
+        pa.types.is_integer(data_type)
+        or pa.types.is_floating(data_type)
+        or pa.types.is_boolean(data_type)
+        or pa.types.is_string(data_type)
+        or pa.types.is_large_string(data_type)
+        or pa.types.is_binary(data_type)
+        or pa.types.is_large_binary(data_type)
+        or pa.types.is_decimal(data_type)
+    )
+
+
+def _timestamp_to_pylist(column: Any) -> list[datetime.datetime]:
+    values = cast(list[datetime.datetime], column.to_numpy(zero_copy_only=False).tolist())
+    timezone = cast(datetime.tzinfo, cast(Any, pa).lib.string_to_tzinfo(column.type.tz))
+    if timezone.utcoffset(None) == datetime.timedelta(0):
+        return [value.replace(tzinfo=timezone) for value in values]
+
+    return [value.replace(tzinfo=datetime.UTC).astimezone(timezone) for value in values]
+
+
+def column_to_pylist(column: Any) -> list[Any]:
+    """Convert one :class:`pyarrow.Array` to Python objects.
+
+    Always equal to ``column.to_pylist()``, down to the exact Python type of
+    every element.
+    """
+    nulls_change_values = column.null_count and (pa.types.is_integer(column.type) or pa.types.is_floating(column.type))
+    tiny_scalar_batch = len(column) == 1 and not pa.types.is_temporal(column.type)
+    timestamp_with_timezone = bool(pa.types.is_timestamp(column.type) and column.type.tz is not None)
+    if timestamp_with_timezone and not column.null_count and _numpy_safe(column.type):
+        return cast(list[Any], _timestamp_to_pylist(column))
+
+    if tiny_scalar_batch or nulls_change_values or timestamp_with_timezone or not _numpy_safe(column.type):
+        return cast(list[Any], column.to_pylist())
+
+    return cast(list[Any], column.to_numpy(zero_copy_only=False).tolist())
+
+
+def batch_to_rows(batch: Any) -> list[tuple[Any, ...]]:
+    """Convert a :class:`pyarrow.RecordBatch` to row tuples, one column at a time."""
+    columns: list[list[Any]] = []
+    for field, column in zip(batch.schema, batch.columns, strict=True):
+        values = column_to_pylist(column)
+        if field.metadata and field.metadata.get(_ARROW_EXTENSION_NAME) == _MONETDB_HUGEINT_EXTENSION:
+            values = [None if value is None else int(value) for value in values]
+        columns.append(values)
+    return list(zip(*columns, strict=True))

sqlalchemy_monetdb_adbc/arrow.py

@@ -0,0 +1,99 @@
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any, Literal, cast
+
+from sqlalchemy.engine import Connection
+from sqlalchemy.sql.compiler import SQLCompiler
+from sqlalchemy.sql.elements import ClauseElement
+from sqlalchemy.sql.expression import Executable
+
+from sqlalchemy_monetdb_adbc.connection import raw_adbc_connection
+
+# pyarrow ships no type information, so the Arrow objects these helpers return
+# and accept are annotated as Any. Their concrete types are given per function.
+
+
+def compile_arrow_statement(connection: Connection, statement: Executable | str) -> tuple[str, list[Any]]:
+    if isinstance(statement, str):
+        return statement, []
+
+    schema_translate_map = connection.get_execution_options().get("schema_translate_map")
+    compiled = cast(
+        SQLCompiler,
+        cast(ClauseElement, statement).compile(
+            dialect=connection.dialect,
+            schema_translate_map=schema_translate_map,
+            render_schema_translate=bool(schema_translate_map),
+            compile_kwargs={"render_postcompile": True},
+        ),
+    )
+    parameters = compiled.params
+    # The dialect uses qmark, so bind values go positionally in the order the
+    # compiler emitted them.
+    ordered = [parameters[name] for name in compiled.positiontup or ()]
+    return str(compiled), ordered
+
+
+def fetch_arrow_table(
+    connection: Connection,
+    statement: Executable | str,
+    parameters: list[Any] | None = None,
+) -> Any:
+    """Run a statement on ``connection`` and return a ``pyarrow.Table``.
+
+    The query runs on the ADBC session backing ``connection``, so it sees that
+    connection's uncommitted work and takes part in its transaction. Rows are
+    never converted to Python objects.
+    """
+    sql, bound = compile_arrow_statement(connection, statement)
+    adbc_connection = raw_adbc_connection(connection)
+
+    with adbc_connection.cursor() as cursor:
+        adbc_cursor = cast(Any, cursor)
+        adbc_cursor.execute(sql, parameters if parameters is not None else bound)
+        return adbc_cursor.fetch_arrow_table()
+
+
+@contextmanager  # pyright: ignore[reportDeprecated]
+def fetch_record_batches(
+    connection: Connection,
+    statement: Executable | str,
+    parameters: list[Any] | None = None,
+) -> Iterator[Any]:
+    """Stream a statement's result as a ``pyarrow.RecordBatchReader``.
+
+    Like :func:`fetch_arrow_table`, but the result is not materialized at once.
+    A context manager, because the reader stays valid only while the cursor
+    behind it is open. MonetDB carries one result channel per session, so
+    finish the stream before using ``connection`` again.
+    """
+    sql, bound = compile_arrow_statement(connection, statement)
+    adbc_connection = raw_adbc_connection(connection)
+
+    with adbc_connection.cursor() as cursor:
+        adbc_cursor = cast(Any, cursor)
+        adbc_cursor.execute(sql, parameters if parameters is not None else bound)
+        yield adbc_cursor.fetch_record_batch()
+
+
+def ingest_arrow(
+    connection: Connection,
+    table_name: str,
+    data: Any,
+    *,
+    mode: Literal["append", "create", "replace", "create_append"] = "append",
+    schema_name: str | None = None,
+) -> int:
+    """Bulk-load Arrow data into ``table_name`` on ``connection``'s transaction.
+
+    ``data`` is anything ADBC accepts: a ``pyarrow`` table, record batch, or
+    reader, or any object exposing the Arrow PyCapsule interface. Returns the
+    number of rows written.
+    """
+    adbc_connection = raw_adbc_connection(connection)
+
+    with adbc_connection.cursor() as cursor:
+        return cast(int, cast(Any, cursor).adbc_ingest(table_name, data, mode=mode, db_schema_name=schema_name))
+
+
+__all__ = ["fetch_arrow_table", "fetch_record_batches", "ingest_arrow"]