| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| """ |
| Backend-specific operations for the DB-API layer. |
| |
| These are mostly functions that convert Python types to/from Arrow types. |
| They are abstracted so that we can support multiple backends like PyArrow, |
| polars, and nanoarrow. |
| """ |
| |
| import abc |
| import typing |
| |
| from . import _lib |
| |
| if typing.TYPE_CHECKING: |
| from typing_extensions import CapsuleType |
| |
| |
| class DbapiBackend(abc.ABC): |
| """ |
| Python/Arrow type conversions that the DB-API layer needs. |
| |
| The return types can and should vary based on the backend. |
| """ |
| |
| @abc.abstractmethod |
| def convert_bind_parameters(self, parameters: typing.Any) -> typing.Any: |
| """Convert an arbitrary Python object into bind parameters. |
| |
| Parameters |
| ---------- |
| parameters |
| A sequence of bind parameters. For instance: a tuple, where each |
| item is a bind parameter in sequence. |
| |
| Returns |
| ------- |
| parameters : CapsuleType |
| This should be an Arrow stream capsule or an object implementing |
| the Arrow PyCapsule interface. |
| |
| See Also |
| -------- |
| https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html |
| |
| """ |
| ... |
| |
| @abc.abstractmethod |
| def convert_executemany_parameters( |
| self, parameters: typing.Any |
| ) -> typing.Tuple[typing.Any, bool]: |
| """Convert an arbitrary Python sequence into bind parameters. |
| |
| Parameters |
| ---------- |
| parameters |
| A sequence of bind parameters. For instance: an iterable of |
| tuples, where each tuple is a row of bind parameters. |
| |
| Returns |
| ------- |
| parameters : CapsuleType |
| This should be an Arrow stream capsule or an object implementing |
| the Arrow PyCapsule interface. |
| bind_by_name : bool |
| Whether the parameters should be bound by name (e.g. because they |
| contain a dictionary). |
| |
| See Also |
| -------- |
| https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html |
| |
| """ |
| ... |
| |
| @abc.abstractmethod |
| def convert_description(self, schema: "CapsuleType") -> typing.List[typing.Tuple]: |
| """Convert a schema capsule into a DB-API description. |
| |
| Parameters |
| ---------- |
| schema |
| A PyCapsule of type "arrow_schema". |
| |
| Returns |
| ------- |
| description : list[tuple] |
| A DB-API description, as a list of 7-item tuples. |
| |
| See Also |
| -------- |
| https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html |
| |
| """ |
| ... |
| |
| @abc.abstractmethod |
| def import_array_stream(self, handle: _lib.ArrowArrayStreamHandle) -> typing.Any: |
| """Import an Arrow stream.""" |
| ... |
| |
| @abc.abstractmethod |
| def import_schema(self, handle: _lib.ArrowSchemaHandle) -> typing.Any: |
| """Import an Arrow schema.""" |
| ... |
| |
| |
| _ALL_BACKENDS: list[DbapiBackend] = [] |
| |
| |
| def default_backend() -> DbapiBackend: |
| return _ALL_BACKENDS[-1] |
| |
| |
| class _NoOpBackend(DbapiBackend): |
| def convert_bind_parameters(self, parameters: typing.Any) -> typing.Any: |
| raise _lib.ProgrammingError( |
| "This API requires PyArrow or another suitable backend to be installed", |
| status_code=_lib.AdbcStatusCode.INVALID_STATE, |
| ) |
| |
| def convert_executemany_parameters( |
| self, parameters: typing.Any |
| ) -> typing.Tuple[typing.Any, bool]: |
| raise _lib.ProgrammingError( |
| "This API requires PyArrow or another suitable backend to be installed", |
| status_code=_lib.AdbcStatusCode.INVALID_STATE, |
| ) |
| |
| def convert_description(self, schema: "CapsuleType") -> typing.List[typing.Tuple]: |
| raise _lib.ProgrammingError( |
| "This API requires PyArrow or another suitable backend to be installed", |
| status_code=_lib.AdbcStatusCode.INVALID_STATE, |
| ) |
| |
| def import_array_stream( |
| self, handle: _lib.ArrowArrayStreamHandle |
| ) -> _lib.ArrowArrayStreamHandle: |
| return handle |
| |
| def import_schema(self, handle: _lib.ArrowSchemaHandle) -> _lib.ArrowSchemaHandle: |
| return handle |
| |
| |
| def param_iterable_to_dict(parameters: typing.Any) -> typing.Tuple[dict, bool]: |
| bind_by_name = False |
| cols = {} |
| for param in parameters: |
| if not cols: |
| # First iteration |
| if isinstance(param, dict): |
| bind_by_name = True |
| for k, v in param.items(): |
| cols[str(k)] = [] |
| else: |
| for col_idx, v in enumerate(param): |
| cols[str(col_idx)] = [] |
| |
| if isinstance(param, dict): |
| for k, v in param.items(): |
| cols[str(k)].append(v) |
| else: |
| for col_idx, v in enumerate(param): |
| cols[str(col_idx)].append(v) |
| return cols, bind_by_name |
| |
| |
| _ALL_BACKENDS.append(_NoOpBackend()) |
| |
| try: |
| import polars |
| |
| class _PolarsBackend(DbapiBackend): |
| def convert_bind_parameters(self, parameters: typing.Any) -> polars.DataFrame: |
| if isinstance(parameters, dict): |
| return polars.DataFrame( |
| {str(k): v for k, v in parameters.items()}, |
| ) |
| |
| return polars.DataFrame( |
| {str(col_idx): v for col_idx, v in enumerate(parameters)}, |
| ) |
| |
| def convert_executemany_parameters( |
| self, parameters: typing.Any |
| ) -> typing.Tuple[typing.Any, bool]: |
| cols, bind_by_name = param_iterable_to_dict(parameters) |
| return polars.DataFrame(cols), bind_by_name |
| |
| def convert_description( |
| self, schema: "CapsuleType" |
| ) -> typing.List[typing.Tuple]: |
| raise _lib.NotSupportedError("Polars does not support __arrow_c_schema__") |
| |
| def import_array_stream( |
| self, handle: _lib.ArrowArrayStreamHandle |
| ) -> typing.Any: |
| return polars.DataFrame(handle) # pyright: ignore[reportArgumentType] |
| |
| def import_schema(self, handle: _lib.ArrowSchemaHandle) -> typing.Any: |
| raise _lib.NotSupportedError("Polars does not support __arrow_c_schema__") |
| |
| _ALL_BACKENDS.append(_PolarsBackend()) |
| except ImportError: |
| pass |
| |
| # Keep PyArrow at the end so it stays default |
| try: |
| import pyarrow |
| |
| class _PyArrowBackend(DbapiBackend): |
| def convert_bind_parameters(self, parameters: typing.Any) -> typing.Any: |
| if isinstance(parameters, dict): |
| return pyarrow.record_batch( |
| {str(k): [v] for k, v in parameters.items()}, |
| ) |
| return pyarrow.record_batch( |
| [[param_value] for param_value in parameters], |
| names=[str(i) for i in range(len(parameters))], |
| ) |
| |
| def convert_executemany_parameters( |
| self, parameters: typing.Any |
| ) -> typing.Tuple[typing.Any, bool]: |
| cols, bind_by_name = param_iterable_to_dict(parameters) |
| return pyarrow.RecordBatch.from_pydict(cols), bind_by_name |
| |
| def convert_description( |
| self, schema: "CapsuleType" |
| ) -> typing.List[typing.Tuple]: |
| s = pyarrow.Schema._import_from_c_capsule(schema) |
| return [ |
| (field.name, field.type, None, None, None, None, None) for field in s |
| ] |
| |
| def import_array_stream( |
| self, handle: _lib.ArrowArrayStreamHandle |
| ) -> pyarrow.RecordBatchReader: |
| return pyarrow.RecordBatchReader._import_from_c(handle.address) |
| |
| def import_schema(self, handle: _lib.ArrowSchemaHandle) -> pyarrow.Schema: |
| return pyarrow.schema(handle) |
| |
| _ALL_BACKENDS.append(_PyArrowBackend()) |
| |
| except ImportError: |
| pass |
