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..b340681870 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,19 @@ 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: + # 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 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 {'.'.join(str(m) for m in required_version)} or higher. " + f"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/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. 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), ] ) 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"