data-validation-engine 0.6.2__py3-none-any.whl

dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py

@@ -0,0 +1,288 @@
+# pylint: disable=protected-access
+# ignore: type[attr-defined]
+
+"""Helper objects for duckdb data contract implementation"""
+from collections.abc import Generator, Iterator
+from dataclasses import is_dataclass
+from datetime import date, datetime, time
+from decimal import Decimal
+from pathlib import Path
+from typing import Any, ClassVar, Union
+from urllib.parse import urlparse
+
+import duckdb.typing as ddbtyp
+import numpy as np
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from duckdb.typing import DuckDBPyType
+from pandas import DataFrame
+from pydantic import BaseModel
+from typing_extensions import Annotated, get_args, get_origin, get_type_hints
+
+from dve.core_engine.backends.base.utilities import _get_non_heterogenous_type
+from dve.core_engine.type_hints import URI
+from dve.parser.file_handling.service import LocalFilesystemImplementation, _get_implementation
+
+
+class DDBDecimal:
+    """DuckDB Decimal type"""
+
+    TYPE_TEXT = "DECIMAL"
+
+    def __init__(self, width: int = 18, scale: int = 3):
+        self._width = width
+        self._scale = scale
+
+    def __str__(self):
+        return DDBDecimal.TYPE_TEXT + f"({self._width},{self._scale})"
+
+    def __call__(self):
+        return self.__str__()
+
+
+class DDBList:
+    """DuckDB List type"""
+
+    TYPE_TEXT = "[]"
+
+    def __init__(self, element_type: DuckDBPyType):
+        self._element_type = element_type
+
+    def __str__(self):
+        return str(self._element_type) + DDBList.TYPE_TEXT
+
+    def __call__(self):
+        return self.__str__()
+
+
+class DDBStruct:
+    """DuckDB StructType"""
+
+    TYPE_TEXT = "STRUCT"
+
+    def __init__(self, sub_elements: dict[str, DuckDBPyType]):
+        self._sub_elements = {**sub_elements}
+
+    def add_element(self, field_name: str, data_type: DuckDBPyType):
+        """Add another element to the struct"""
+        self._sub_elements.update({field_name: data_type})
+
+    def __str__(self):
+        return (
+            DDBStruct.TYPE_TEXT
+            + "("
+            + ", ".join(f"{fld} {dtype}" for fld, dtype in self._sub_elements.items())
+            + ")"
+        )
+
+    def __call__(self):
+        return self.__str__()
+
+
+PYTHON_TYPE_TO_DUCKDB_TYPE: dict[type, DuckDBPyType] = {
+    str: ddbtyp.VARCHAR,
+    int: ddbtyp.BIGINT,
+    bool: ddbtyp.BOOLEAN,
+    float: ddbtyp.FLOAT,
+    bytes: ddbtyp.BLOB,
+    date: ddbtyp.DATE,
+    datetime: ddbtyp.TIMESTAMP,
+    Decimal: DDBDecimal()(),
+    time: ddbtyp.TIME,
+}
+"""A mapping of Python types to the equivalent DuckDB types."""
+
+
+def table_exists(connection: DuckDBPyConnection, table_name: str) -> bool:
+    """check if a table exists in a given DuckDBPyConnection"""
+    return table_name in map(lambda x: x[0], connection.sql("SHOW TABLES").fetchall())
+
+
+def relation_is_empty(relation: DuckDBPyRelation) -> bool:
+    """Check if a duckdb relation is empty"""
+    if relation.limit(1).count("*"):
+        return False
+    return True
+
+
+def get_duckdb_type_from_annotation(type_annotation: Any) -> DuckDBPyType:
+    """Get a duckdb type from a Python type annotation.
+
+    Supported types are any of the following (this definition is recursive):
+    - Supported basic Python types. These are:
+      * `str`: VARCHAR
+      * `int`: BIGINT
+      * `bool`: BOOLEAN
+      * `float`: FLOAT
+      * `bytes`: BLOB
+      * `datetime.date`: DATE
+      * `datetime.datetime`: TIMESTAMP
+      * `decimal.Decimal`: DECIMAL with precision of 18 and scale of 3
+    - A list of supported types (e.g. `List[str]` or `typing.List[str]`).
+      This will return a duckdb LIST type (variable length)
+    - A `typing.Optional` type or a `typing.Union` of the type and `None` (e.g.
+      `typing.Optional[str]`, `typing.Union[List[str], None]`). This will remove the
+      'optional' wrapper and return the inner type
+    - A subclass of `typing.TypedDict` with values typed using supported types. This
+      will parse the value types as Polars types and return a duckdb STRUCT.
+    - A dataclass or `pydantic.main.ModelMetaClass` with values typed using supported types.
+      This will parse the field types as Polars types and return a duckdb STRUCT.
+    - Any supported type, with a `typing_extensions.Annotated` wrapper.
+
+    Any `ClassVar` types within `TypedDict`s, dataclasses, or `pydantic` models will be
+    ignored.
+
+    """
+    type_origin = get_origin(type_annotation)
+
+    # An `Optional` or `Union` type, check to ensure non-heterogenity.
+    if type_origin is Union:
+        python_type = _get_non_heterogenous_type(get_args(type_annotation))
+        return get_duckdb_type_from_annotation(python_type)
+
+    # Type hint is e.g. `List[str]`, check to ensure non-heterogenity.
+    if type_origin is list or (isinstance(type_origin, type) and issubclass(type_origin, list)):
+        element_type = _get_non_heterogenous_type(get_args(type_annotation))
+        return DDBList(get_duckdb_type_from_annotation(element_type))()
+
+    if type_origin is Annotated:
+        python_type, *other_args = get_args(type_annotation)  # pylint: disable=unused-variable
+        return get_duckdb_type_from_annotation(python_type)
+    # Ensure that we have a concrete type at this point.
+    if not isinstance(type_annotation, type):
+        raise ValueError(f"Unsupported type annotation {type_annotation!r}")
+
+    if (
+        # Type hint is a dict subclass, but not dict. Possibly a `TypedDict`.
+        (issubclass(type_annotation, dict) and type_annotation is not dict)
+        # Type hint is a dataclass.
+        or is_dataclass(type_annotation)
+        # Type hint is a `pydantic` model.
+        or (type_origin is None and issubclass(type_annotation, BaseModel))
+    ):
+        fields: dict[str, DuckDBPyType] = {}
+        for field_name, field_annotation in get_type_hints(type_annotation).items():
+            # Technically non-string keys are disallowed, but people are bad.
+            if not isinstance(field_name, str):
+                raise ValueError(
+                    f"Dictionary/Dataclass keys must be strings, got {type_annotation!r}"
+                )  # pragma: no cover
+            if get_origin(field_annotation) is ClassVar:
+                continue
+
+            fields[field_name] = get_duckdb_type_from_annotation(field_annotation)
+
+        if not fields:
+            raise ValueError(
+                f"No type annotations in dict/dataclass type (got {type_annotation!r})"
+            )
+
+        return DDBStruct(fields)()
+
+    if type_annotation is list:
+        raise ValueError(
+            f"List must have type annotation (e.g. `List[str]`), got {type_annotation!r}"
+        )
+    if type_annotation is dict or type_origin is dict:
+        raise ValueError(f"dict must be `typing.TypedDict` subclass, got {type_annotation!r}")
+
+    for type_ in type_annotation.mro():
+        duck_type = PYTHON_TYPE_TO_DUCKDB_TYPE.get(type_)
+        if duck_type:
+            return duck_type
+    raise ValueError(f"No equivalent DuckDB type for {type_annotation!r}")
+
+
+def coerce_inferred_numpy_array_to_list(pandas_df: DataFrame) -> DataFrame:
+    """Function to modify numpy inferred array when cnverting from duckdb relation to
+    pandas dataframe - these cause issues with pydantic models
+    (ie. numpy array can't be viewed as a list)
+
+    pandas_df (DataFrame): The dataframe to type check and covert where needed
+    """
+    pandas_df = pandas_df.replace({np.nan: None})
+    for col in pandas_df.columns:
+        if isinstance(pandas_df[col].iloc[0], (np.ndarray,)):
+            pandas_df[col] = pandas_df[col].apply(lambda x: x.tolist() if x is not None else x)
+    return pandas_df
+
+
+def _ddb_read_parquet(
+    self, path: URI, **kwargs  # pylint: disable=unused-argument
+) -> DuckDBPyRelation:
+    """Read entity from a parquet file. Due to different behaviours in writing parquet
+    files to single files/ mutiple files in a directory and duckdb being inflexible
+    checks whether a directory supplied (in which case points at parquet files within)
+    or a single file, in which case retains supplied path.
+    """
+    if isinstance(_get_implementation(path), LocalFilesystemImplementation):
+        path = urlparse(path).path
+
+    if Path(path).is_dir():
+        if not path.endswith("/"):
+            path += "/"
+        # is a directory - provide edit glob to include all parquet files
+        return self._connection.read_parquet(file_glob=f"{path}*.parquet")
+    return self._connection.read_parquet(file_glob=path)
+
+
+def _ddb_write_parquet(  # pylint: disable=unused-argument
+    self, entity: Union[Iterator[dict[str, Any]], DuckDBPyRelation], target_location: URI, **kwargs
+) -> URI:
+    """Method to write parquet files from type cast entities
+    following data contract application
+    """
+    if isinstance(_get_implementation(target_location), LocalFilesystemImplementation):
+        Path(target_location).parent.mkdir(parents=True, exist_ok=True)
+
+    if isinstance(entity, Generator):
+        entity = self._connection.query(
+            "select dta.* from (select unnest($data) as dta)", params={"data": list(entity)}
+        )
+
+    entity.to_parquet(file_name=target_location, compression="snappy", **kwargs)  # type: ignore
+    return target_location
+
+
+def duckdb_read_parquet(cls):
+    """Class decorator to add read_parquet method for duckdb implementations"""
+    cls.read_parquet = _ddb_read_parquet
+    return cls
+
+
+def duckdb_write_parquet(cls):
+    """Class decorator to add write_parquet method for duckdb implementations"""
+    cls.write_parquet = _ddb_write_parquet
+    return cls
+
+
+@staticmethod  # type: ignore
+def _duckdb_get_entity_count(entity: DuckDBPyRelation) -> int:
+    """Method to obtain entity count from a persisted parquet entity"""
+    return entity.shape[0]
+
+
+def duckdb_get_entity_count(cls):
+    """Class decorator to count records in an entity supplied"""
+    cls.get_entity_count = _duckdb_get_entity_count
+    return cls
+
+
+def get_all_registered_udfs(connection: DuckDBPyConnection) -> set[str]:
+    """Function to supply the names of a registered functions stored in the supplied
+    duckdb connection. Creates the temp table used to store registered functions (if not exists).
+    """
+    connection.sql("CREATE TEMP TABLE IF NOT EXISTS dve_udfs (function_name VARCHAR)")
+    return {rw[0] for rw in connection.sql("SELECT * FROM dve_udfs").fetchall()}
+
+
+def duckdb_rel_to_dictionaries(
+    entity: DuckDBPyRelation, batch_size=1000
+) -> Iterator[dict[str, Any]]:
+    """Iterator converting DuckDBPyRelation to lists of dictionaries.
+    Avoids issues where dates are getting converted to datetimes using polars as intermediate."""
+    # TODO - look into float conversion - floats that can't be stored exactly in binary
+    # TODO - are given to nearest approximation. Tried Decimal, causes issues in arrays
+    # TODO - with templating (as in complex fields, repr used when str called in jinja templating).
+    cols: tuple[str] = tuple(entity.columns)  # type: ignore
+    while rows := entity.fetchmany(batch_size):
+        yield from (dict(zip(cols, rw)) for rw in rows)

dve/core_engine/backends/implementations/duckdb/readers/__init__.py

@@ -0,0 +1,13 @@
+"""Readers for use with duckdb backend"""
+
+from .csv import DuckDBCSVReader, DuckDBCSVRepeatingHeaderReader, PolarsToDuckDBCSVReader
+from .json import DuckDBJSONReader
+from .xml import DuckDBXMLStreamReader
+
+__all__ = [
+    "DuckDBCSVReader",
+    "DuckDBCSVRepeatingHeaderReader",
+    "DuckDBJSONReader",
+    "DuckDBXMLStreamReader",
+    "PolarsToDuckDBCSVReader",
+]

dve/core_engine/backends/implementations/duckdb/readers/csv.py

@@ -0,0 +1,222 @@
+"""A csv reader to create duckdb relations"""
+
+# pylint: disable=arguments-differ
+from collections.abc import Iterator
+from typing import Any, Optional
+
+import duckdb as ddb
+import polars as pl
+from duckdb import DuckDBPyConnection, DuckDBPyRelation, default_connection, read_csv
+from pydantic import BaseModel
+
+from dve.core_engine.backends.base.reader import BaseFileReader, read_function
+from dve.core_engine.backends.exceptions import EmptyFileError, MessageBearingError
+from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import (
+    duckdb_write_parquet,
+    get_duckdb_type_from_annotation,
+)
+from dve.core_engine.backends.implementations.duckdb.types import SQLType
+from dve.core_engine.backends.readers.utilities import check_csv_header_expected
+from dve.core_engine.backends.utilities import get_polars_type_from_annotation
+from dve.core_engine.message import FeedbackMessage
+from dve.core_engine.type_hints import URI, EntityName
+from dve.parser.file_handling import get_content_length
+
+
+@duckdb_write_parquet
+class DuckDBCSVReader(BaseFileReader):
+    """A reader for CSV files including the ability to compare the passed model
+    to the file header, if it exists.
+
+    field_check: flag to compare submitted file header to the accompanying pydantic model
+    field_check_error_code: The error code to provide if the file header doesn't contain
+                            the expected fields
+    field_check_error_message: The error message to provide if the file header doesn't contain
+                               the expected fields"""
+
+    # TODO - the read_to_relation should include the schema and determine whether to
+    # TODO - stringify or not
+    def __init__(
+        self,
+        *,
+        header: bool = True,
+        delim: str = ",",
+        quotechar: str = '"',
+        connection: Optional[DuckDBPyConnection] = None,
+        field_check: bool = False,
+        field_check_error_code: Optional[str] = "ExpectedVsActualFieldMismatch",
+        field_check_error_message: Optional[str] = "The submitted header is missing fields",
+        **_,
+    ):
+        self.header = header
+        self.delim = delim
+        self.quotechar = quotechar
+        self._connection = connection if connection else default_connection
+        self.field_check = field_check
+        self.field_check_error_code = field_check_error_code
+        self.field_check_error_message = field_check_error_message
+
+        super().__init__()
+
+    def perform_field_check(
+        self, resource: URI, entity_name: str, expected_schema: type[BaseModel]
+    ):
+        """Check that the header of the CSV aligns with the provided model"""
+        if not self.header:
+            raise ValueError("Cannot perform field check without a CSV header")
+
+        if missing := check_csv_header_expected(resource, expected_schema, self.delim):
+            raise MessageBearingError(
+                "The CSV header doesn't match what is expected",
+                messages=[
+                    FeedbackMessage(
+                        entity=entity_name,
+                        record=None,
+                        failure_type="submission",
+                        error_location="Whole File",
+                        error_code=self.field_check_error_code,
+                        error_message=f"{self.field_check_error_message} - missing fields: {missing}",  # pylint: disable=line-too-long
+                    )
+                ],
+            )
+
+    def read_to_py_iterator(
+        self, resource: URI, entity_name: EntityName, schema: type[BaseModel]
+    ) -> Iterator[dict[str, Any]]:
+        """Creates an iterable object of rows as dictionaries"""
+        yield from self.read_to_relation(resource, entity_name, schema).pl().iter_rows(named=True)
+
+    @read_function(DuckDBPyRelation)
+    def read_to_relation(  # pylint: disable=unused-argument
+        self, resource: URI, entity_name: EntityName, schema: type[BaseModel]
+    ) -> DuckDBPyRelation:
+        """Returns a relation object from the source csv"""
+        if get_content_length(resource) == 0:
+            raise EmptyFileError(f"File at {resource} is empty.")
+
+        if self.field_check:
+            self.perform_field_check(resource, entity_name, schema)
+
+        reader_options: dict[str, Any] = {
+            "header": self.header,
+            "delimiter": self.delim,
+            "quotechar": self.quotechar,
+        }
+
+        ddb_schema: dict[str, SQLType] = {
+            fld.name: str(get_duckdb_type_from_annotation(fld.annotation))  # type: ignore
+            for fld in schema.__fields__.values()
+        }
+
+        reader_options["columns"] = ddb_schema
+        return read_csv(resource, **reader_options)
+
+
+class PolarsToDuckDBCSVReader(DuckDBCSVReader):
+    """
+    Utilises the polars lazy csv reader which is then converted into a DuckDBPyRelation object.
+
+    The primary reason this reader exists is due to the limitation within duckdb csv reader and
+    it not being able to read partial content from a csv (i.e. select a, b NOT y).
+    """
+
+    @read_function(DuckDBPyRelation)
+    def read_to_relation(  # pylint: disable=unused-argument
+        self, resource: URI, entity_name: EntityName, schema: type[BaseModel]
+    ) -> DuckDBPyRelation:
+        """Returns a relation object from the source csv"""
+        if get_content_length(resource) == 0:
+            raise EmptyFileError(f"File at {resource} is empty.")
+
+        if self.field_check:
+            self.perform_field_check(resource, entity_name, schema)
+
+        reader_options: dict[str, Any] = {
+            "has_header": self.header,
+            "separator": self.delim,
+            "quote_char": self.quotechar,
+        }
+
+        polars_types = {
+            fld.name: get_polars_type_from_annotation(fld.annotation)  # type: ignore
+            for fld in schema.__fields__.values()
+        }
+        reader_options["dtypes"] = polars_types
+
+        # there is a raise_if_empty arg for 0.18+. Future reference when upgrading. Makes L85
+        # redundant
+        df = pl.scan_csv(resource, **reader_options).select(list(polars_types.keys()))  # type: ignore  # pylint: disable=W0612
+
+        return ddb.sql("SELECT * FROM df")
+
+
+class DuckDBCSVRepeatingHeaderReader(PolarsToDuckDBCSVReader):
+    """A Reader for files with a `.csv` extension and where there are repeating "header" values
+    within the file. Header in this case is not the column names at the top of a csv, rather a
+    collection of unique records that would usually be structured in another entity. However, due
+    to the fact that `csv` is a semi-structured data format, you cannot define complex entities,
+    hence the values are then repeated on all rows.
+
+    Example of a repeating header data may look like this...
+
+    | headerCol1 | headerCol2 | headerCol3 | nonHeaderCol1 | nonHeaderCol2 |
+    | ---------- | ---------- | ---------- | ------------- | ------------- |
+    | shop 1     | clothes    | 2025-01-01 | jeans         | 20.39         |
+    | shop 1     | clothes    | 2025-01-01 | shirt         | 14.99         |
+
+    This reader will just pull out the distinct values from the header column. Where there are
+    more/less than one distinct value per column, the reader will produce a
+    `NonDistinctHeaderError`.
+
+    So using the example above, the expected entity would look like this...
+    | headerCol1 | headerCol2 | headerCol3 |
+    | ---------- | ---------- | ---------- |
+    | shop1      | clothes    | 2025-01-01 |
+    """
+
+    def __init__(
+        self,
+        *args,
+        non_unique_header_error_code: Optional[str] = "NonUniqueHeader",
+        non_unique_header_error_message: Optional[str] = None,
+        **kwargs,
+    ):
+        self._non_unique_header_code = non_unique_header_error_code
+        self._non_unique_header_message = non_unique_header_error_message
+        super().__init__(*args, **kwargs)
+
+    @read_function(DuckDBPyRelation)
+    def read_to_relation(  # pylint: disable=unused-argument
+        self, resource: URI, entity_name: EntityName, schema: type[BaseModel]
+    ) -> DuckDBPyRelation:
+        entity = super().read_to_relation(resource=resource, entity_name=entity_name, schema=schema)
+        entity = entity.distinct()
+        no_records = entity.shape[0]
+
+        if no_records != 1:
+            rows = entity.pl().to_dicts()
+            differing_values = [
+                f"{key}: {', '.join(sorted(str(val) for val in values))}"
+                for key, *values in zip(rows[0], *map(dict.values, rows))  # type: ignore
+                if len(set(values)) > 1
+            ]
+            raise MessageBearingError(
+                "More than one set of Headers found in CSV file",
+                messages=[
+                    FeedbackMessage(
+                        record={entity_name: differing_values},
+                        entity="Pre-validation",
+                        failure_type="submission",
+                        error_message=(
+                            f"Found {no_records} distinct combination of header values."
+                            if not self._non_unique_header_message
+                            else self._non_unique_header_message
+                        ),
+                        error_location=entity_name,
+                        category="Bad file",
+                        error_code=self._non_unique_header_code,
+                    )
+                ],
+            )
+
+        return entity

dve/core_engine/backends/implementations/duckdb/readers/json.py

@@ -0,0 +1,50 @@
+"""A csv reader to create duckdb relations"""
+
+# pylint: disable=arguments-differ
+from collections.abc import Iterator
+from typing import Any, Optional
+
+from duckdb import DuckDBPyRelation, read_json
+from pydantic import BaseModel
+
+from dve.core_engine.backends.base.reader import BaseFileReader, read_function
+from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import (
+    duckdb_write_parquet,
+    get_duckdb_type_from_annotation,
+)
+from dve.core_engine.backends.implementations.duckdb.types import SQLType
+from dve.core_engine.type_hints import URI, EntityName
+
+
+@duckdb_write_parquet
+class DuckDBJSONReader(BaseFileReader):
+    """A reader for JSON files"""
+
+    def __init__(
+        self,
+        *,
+        json_format: Optional[str] = "array",
+        **_,
+    ):
+        self._json_format = json_format
+
+        super().__init__()
+
+    def read_to_py_iterator(
+        self, resource: URI, entity_name: EntityName, schema: type[BaseModel]
+    ) -> Iterator[dict[str, Any]]:
+        """Creates an iterable object of rows as dictionaries"""
+        return self.read_to_relation(resource, entity_name, schema).pl().iter_rows(named=True)
+
+    @read_function(DuckDBPyRelation)
+    def read_to_relation(  # pylint: disable=unused-argument
+        self, resource: URI, entity_name: EntityName, schema: type[BaseModel]
+    ) -> DuckDBPyRelation:
+        """Returns a relation object from the source json"""
+
+        ddb_schema: dict[str, SQLType] = {
+            fld.name: str(get_duckdb_type_from_annotation(fld.annotation))  # type: ignore
+            for fld in schema.__fields__.values()
+        }
+
+        return read_json(resource, columns=ddb_schema, format=self._json_format)  # type: ignore

dve/core_engine/backends/implementations/duckdb/readers/xml.py

@@ -0,0 +1,45 @@
+# mypy: disable-error-code="attr-defined"
+"""An xml reader to create duckdb relations"""
+
+from typing import Optional
+
+import polars as pl
+from duckdb import DuckDBPyConnection, DuckDBPyRelation, default_connection
+from pydantic import BaseModel
+
+from dve.core_engine.backends.base.reader import read_function
+from dve.core_engine.backends.exceptions import MessageBearingError
+from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import duckdb_write_parquet
+from dve.core_engine.backends.readers.xml import XMLStreamReader
+from dve.core_engine.backends.utilities import get_polars_type_from_annotation, stringify_model
+from dve.core_engine.type_hints import URI
+
+
+@duckdb_write_parquet
+class DuckDBXMLStreamReader(XMLStreamReader):
+    """A reader for XML files"""
+
+    def __init__(self, *, ddb_connection: Optional[DuckDBPyConnection] = None, **kwargs):
+        self.ddb_connection = ddb_connection if ddb_connection else default_connection
+        super().__init__(**kwargs)
+
+    @read_function(DuckDBPyRelation)
+    def read_to_relation(self, resource: URI, entity_name: str, schema: type[BaseModel]):
+        """Returns a relation object from the source xml"""
+        if self.xsd_location:
+            msg = self._run_xmllint(file_uri=resource)
+            if msg:
+                raise MessageBearingError(
+                    "Submitted file failed XSD validation.",
+                    messages=[msg],
+                )
+
+        polars_schema: dict[str, pl.DataType] = {  # type: ignore
+            fld.name: get_polars_type_from_annotation(fld.annotation)
+            for fld in stringify_model(schema).__fields__.values()
+        }
+
+        _lazy_frame = pl.LazyFrame(
+            data=self.read_to_py_iterator(resource, entity_name, schema), schema=polars_schema
+        )
+        return self.ddb_connection.sql("select * from _lazy_frame")

dve/core_engine/backends/implementations/duckdb/reference_data.py

@@ -0,0 +1,49 @@
+"""A reference data loader for duckdb."""
+
+from typing import Optional
+
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from pyarrow import ipc  # type: ignore
+
+from dve.core_engine.backends.base.reference_data import (
+    BaseRefDataLoader,
+    ReferenceConfigUnion,
+    ReferenceTable,
+    mark_refdata_file_extension,
+)
+from dve.core_engine.type_hints import EntityName
+from dve.parser.type_hints import URI
+
+
+# pylint: disable=too-few-public-methods
+class DuckDBRefDataLoader(BaseRefDataLoader[DuckDBPyRelation]):
+    """A reference data loader using already existing DuckDB tables."""
+
+    connection: DuckDBPyConnection
+    """The DuckDB connection for the backend."""
+    dataset_config_uri: Optional[URI] = None
+    """The location of the dischema file"""
+
+    def __init__(
+        self,
+        reference_entity_config: dict[EntityName, ReferenceConfigUnion],
+        **kwargs,
+    ) -> None:
+        super().__init__(reference_entity_config, self.dataset_config_uri, **kwargs)
+
+        if not self.connection:
+            raise AttributeError("DuckDBConnection must be specified")
+
+    def load_table(self, config: ReferenceTable) -> DuckDBPyRelation:
+        """Load reference entity from a database table"""
+        return self.connection.sql(f"select * from {config.fq_table_name}")
+
+    @mark_refdata_file_extension("parquet")
+    def load_parquet_file(self, uri: str) -> DuckDBPyRelation:
+        """Load a parquet file into a duckdb relation"""
+        return self.connection.read_parquet(uri)
+
+    @mark_refdata_file_extension("arrow")
+    def load_arrow_file(self, uri: str) -> DuckDBPyRelation:
+        """Load an arrow ipc file into a duckdb relation"""
+        return self.connection.from_arrow(ipc.open_stream(uri).read_all())  # type:ignore