From 013e616ffdf8eebe7ade80be7203b348a34149bc Mon Sep 17 00:00:00 2001 From: Henry Harbeck Date: Sat, 11 Oct 2025 20:41:11 +1000 Subject: [PATCH 1/4] add support for adbc_get_table_schema with Polars backend --- .../adbc_driver_manager/_dbapi_backend.py | 17 +++++++++++++++-- .../adbc_driver_manager/_lib.pyi | 5 +++++ .../tests/test_dbapi_polars_nopyarrow.py | 4 ++-- 3 files changed, 22 insertions(+), 4 deletions(-) diff --git a/python/adbc_driver_manager/adbc_driver_manager/_dbapi_backend.py b/python/adbc_driver_manager/adbc_driver_manager/_dbapi_backend.py index 9457a74761..b513364a0d 100644 --- a/python/adbc_driver_manager/adbc_driver_manager/_dbapi_backend.py +++ b/python/adbc_driver_manager/adbc_driver_manager/_dbapi_backend.py @@ -179,8 +179,21 @@ def import_array_stream( ) -> typing.Any: return polars.from_arrow(handle) - def import_schema(self, handle: _lib.ArrowSchemaHandle) -> typing.Any: - raise _lib.NotSupportedError("Polars does not support __arrow_c_schema__") + def import_schema(self, handle: _lib.ArrowSchemaHandle) -> polars.Schema: + def parse_version(version: str) -> tuple[int, ...]: + return tuple(int(v) for v in version.split(".")) + + # The version that Polars added support to initialize a schema via the + # __arrow_c_schema__ interface + required_version = "1.32.2" + polars_version = polars.__version__ + if parse_version(polars_version) >= parse_version(required_version): + return polars.Schema(handle) + msg = ( + "Initializing Polars Schema from __arrow_c_schema__ interface requires " + f"version {required_version} or higher. Found {polars_version!r}" + ) + raise _lib.NotSupportedError(msg) _ALL_BACKENDS.append(_PolarsBackend()) except ImportError: diff --git a/python/adbc_driver_manager/adbc_driver_manager/_lib.pyi b/python/adbc_driver_manager/adbc_driver_manager/_lib.pyi index 33d234e825..9d11c0f75d 100644 --- a/python/adbc_driver_manager/adbc_driver_manager/_lib.pyi +++ b/python/adbc_driver_manager/adbc_driver_manager/_lib.pyi @@ -145,16 +145,21 @@ class ArrowArrayHandle: address: int is_valid: bool def release(self) -> None: ... + def __arrow_c_array__( + self, requested_schema: object | None + ) -> tuple[object, object]: ... class ArrowArrayStreamHandle: address: int is_valid: bool def release(self) -> None: ... + def __arrow_c_stream__(self, requested_schema: object | None) -> object: ... class ArrowSchemaHandle: address: int is_valid: bool def release(self) -> None: ... + def __arrow_c_schema__(self) -> object: ... class DataError(DatabaseError): ... class DatabaseError(Error): ... diff --git a/python/adbc_driver_manager/tests/test_dbapi_polars_nopyarrow.py b/python/adbc_driver_manager/tests/test_dbapi_polars_nopyarrow.py index da86fd3a12..3c0ca903e7 100644 --- a/python/adbc_driver_manager/tests/test_dbapi_polars_nopyarrow.py +++ b/python/adbc_driver_manager/tests/test_dbapi_polars_nopyarrow.py @@ -295,15 +295,15 @@ def test_query_double_capsule(sqlite: dbapi.Connection) -> None: polars.from_arrow(capsule) -@pytest.mark.xfail(raises=dbapi.NotSupportedError) def test_get_table_schema(sqlite: dbapi.Connection) -> None: with sqlite.cursor() as cursor: cursor.execute("CREATE TABLE test_table_schema (a INT, b STRING)") + cursor.execute("INSERT INTO test_table_schema VALUES (1, 'hello')") schema = sqlite.adbc_get_table_schema("test_table_schema") assert schema == polars.Schema( [ - ("a", polars.Int32), + ("a", polars.Int64), ("b", polars.String), ] ) From 5c5867d76547bfbf87dcc31d91bcb4f4c76e54e2 Mon Sep 17 00:00:00 2001 From: Henry Harbeck Date: Sat, 11 Oct 2025 20:41:25 +1000 Subject: [PATCH 2/4] drive by typo --- python/adbc_driver_snowflake/adbc_driver_snowflake/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/adbc_driver_snowflake/adbc_driver_snowflake/__init__.py b/python/adbc_driver_snowflake/adbc_driver_snowflake/__init__.py index 49b4c891e0..14de86bc95 100644 --- a/python/adbc_driver_snowflake/adbc_driver_snowflake/__init__.py +++ b/python/adbc_driver_snowflake/adbc_driver_snowflake/__init__.py @@ -29,7 +29,7 @@ class DatabaseOptions(enum.Enum): - """Database options specific to the Flight SQL driver.""" + """Database options specific to the Snowflake driver.""" ACCOUNT = "adbc.snowflake.sql.account" APPLICATION_NAME = "adbc.snowflake.sql.client_option.app_name" From b969d6e95bbf82b861bc8f5c6c416f35fb543253 Mon Sep 17 00:00:00 2001 From: Henry Harbeck Date: Sun, 12 Oct 2025 16:49:55 +1000 Subject: [PATCH 3/4] update return type hints --- .../adbc_driver_manager/dbapi.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/python/adbc_driver_manager/adbc_driver_manager/dbapi.py b/python/adbc_driver_manager/adbc_driver_manager/dbapi.py index 4c3641dca8..e55bc210ce 100644 --- a/python/adbc_driver_manager/adbc_driver_manager/dbapi.py +++ b/python/adbc_driver_manager/adbc_driver_manager/dbapi.py @@ -491,7 +491,7 @@ def adbc_get_table_schema( *, catalog_filter: Optional[str] = None, db_schema_filter: Optional[str] = None, - ) -> "pyarrow.Schema": + ) -> "pyarrow.Schema | polars.Schema": """ Get the Arrow schema of a table by name. @@ -1027,16 +1027,16 @@ def adbc_execute_partitions( self, operation, parameters=None, - ) -> Tuple[List[bytes], "pyarrow.Schema"]: + ) -> Tuple[List[bytes], "pyarrow.Schema | polars.Schema"]: """ Execute a query and get the partitions of a distributed result set. Returns ------- - partitions : list of byte + partitions : list of bytes A list of partition descriptors, which can be read with read_partition. - schema : pyarrow.Schema or None + schema : pyarrow.Schema, polars.Schema or None The schema of the result set. May be None if incremental query execution is enabled and the server has not returned a schema. @@ -1055,13 +1055,15 @@ def adbc_execute_partitions( schema = None return partitions, schema - def adbc_execute_schema(self, operation, parameters=None) -> "pyarrow.Schema": + def adbc_execute_schema( + self, operation, parameters=None + ) -> "pyarrow.Schema | polars.Schema": """ Get the schema of the result set of a query without executing it. Returns ------- - pyarrow.Schema + pyarrow.Schema or polars.Schema The schema of the result set. Notes @@ -1073,7 +1075,9 @@ def adbc_execute_schema(self, operation, parameters=None) -> "pyarrow.Schema": schema = _blocking_call(self._stmt.execute_schema, (), {}, self._stmt.cancel) return self._conn._backend.import_schema(schema) - def adbc_prepare(self, operation: Union[bytes, str]) -> Optional["pyarrow.Schema"]: + def adbc_prepare( + self, operation: Union[bytes, str] + ) -> Optional["pyarrow.Schema | polars.Schema"]: """ Prepare a query without executing it. @@ -1083,7 +1087,7 @@ def adbc_prepare(self, operation: Union[bytes, str]) -> Optional["pyarrow.Schema Returns ------- - pyarrow.Schema or None + pyarrow.Schema, polars.Schema or None The schema of the bind parameters, or None if the schema could not be determined. From 40030accc0a4ed3f848f40292e9cdb88a2b60223 Mon Sep 17 00:00:00 2001 From: Henry Harbeck Date: Sun, 12 Oct 2025 17:09:39 +1000 Subject: [PATCH 4/4] reduce work in happy path --- .../adbc_driver_manager/_dbapi_backend.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/python/adbc_driver_manager/adbc_driver_manager/_dbapi_backend.py b/python/adbc_driver_manager/adbc_driver_manager/_dbapi_backend.py index b513364a0d..b340681870 100644 --- a/python/adbc_driver_manager/adbc_driver_manager/_dbapi_backend.py +++ b/python/adbc_driver_manager/adbc_driver_manager/_dbapi_backend.py @@ -180,18 +180,16 @@ def import_array_stream( return polars.from_arrow(handle) def import_schema(self, handle: _lib.ArrowSchemaHandle) -> polars.Schema: - def parse_version(version: str) -> tuple[int, ...]: - return tuple(int(v) for v in version.split(".")) - # The version that Polars added support to initialize a schema via the # __arrow_c_schema__ interface - required_version = "1.32.2" + required_version = (1, 32, 2) polars_version = polars.__version__ - if parse_version(polars_version) >= parse_version(required_version): + if tuple(int(v) for v in polars_version.split(".")) >= required_version: return polars.Schema(handle) msg = ( "Initializing Polars Schema from __arrow_c_schema__ interface requires " - f"version {required_version} or higher. Found {polars_version!r}" + f"version {'.'.join(str(m) for m in required_version)} or higher. " + f"Found {polars_version!r}" ) raise _lib.NotSupportedError(msg)