feat: Add read of arrow ipc files to reference data loaders

Author: stevenhsd

src/dve/core_engine/backends/base/reference_data.py

@@ -8,9 +8,23 @@
 from typing_extensions import Annotated, Literal
 
 from dve.core_engine.backends.base.core import get_entity_type
-from dve.core_engine.backends.exceptions import MissingRefDataEntity
+from dve.core_engine.backends.exceptions import MissingRefDataEntity, RefdataLacksFileExtensionSupport
 from dve.core_engine.backends.types import EntityType
-from dve.core_engine.type_hints import EntityName
+from dve.core_engine.type_hints import URI, EntityName
+import dve.parser.file_handling as fh
+from dve.parser.file_handling.implementations.file import LocalFilesystemImplementation
+from dve.parser.file_handling.service import _get_implementation
+
+_FILE_EXTENSION_NAME: str = "_REFDATA_FILE_EXTENSION"
+"""Name of attribute added to methods where they relate
+   to loading a particular reference file type."""
+
+def mark_refdata_file_extension(file_extension):
+    """Mark a method for loading a particular file extension"""
+    def wrapper(func: Callable):
+        setattr(func, _FILE_EXTENSION_NAME, file_extension)
+        return func
+    return wrapper
 
 
 class ReferenceTable(BaseModel, frozen=True):
@@ -37,7 +51,10 @@ class ReferenceFile(BaseModel, frozen=True):
     type: Literal["filename"]
     """The object type."""
     filename: str
-    """The path to the reference data (as Parquet) relative to the contract."""
+    """The path to the reference data relative to the contract."""
+    @property
+    def file_extension(self) -> str:
+        return fh.get_file_suffix(self.filename)
 
 
 class ReferenceURI(BaseModel, frozen=True):
@@ -47,6 +64,9 @@ class ReferenceURI(BaseModel, frozen=True):
     """The object type."""
     uri: str
     """The absolute URI of the reference data (as Parquet)."""
+    @property
+    def file_extension(self) -> str:
+        return fh.get_file_suffix(self.uri)
 
 
 ReferenceConfig = Union[ReferenceFile, ReferenceTable, ReferenceURI]
@@ -71,6 +91,12 @@ class BaseRefDataLoader(Generic[EntityType], Mapping[EntityName, EntityType], AB
     A mapping between refdata config types and functions to call to load these configs
     into reference data entities
     """
+    
+    __reader_functions__: ClassVar[dict[str, Callable]] = {}
+    """
+    A mapping between file extensions and functions to load the file uris
+    into reference data entities
+    """
     prefix: str = "refdata_"
 
     def __init_subclass__(cls, *_, **__) -> None:
@@ -91,20 +117,29 @@ class variable for the subclass.
             method = getattr(cls, method_name, None)
             if method is None or not callable(method):
                 continue
+            
+            if ext := getattr(method, _FILE_EXTENSION_NAME, None):
+                cls.__reader_functions__[ext] = method
+                continue
 
             type_hints = get_type_hints(method)
             if set(type_hints.keys()) != {"config", "return"}:
                 continue
             config_type = type_hints["config"]
             if not issubclass(config_type, BaseModel):
                 continue
+
             cls.__step_functions__[config_type] = method  # type: ignore
 
     # pylint: disable=unused-argument
     def __init__(
-        self, reference_entity_config: dict[EntityName, ReferenceConfig], **kwargs
+        self,
+        reference_entity_config: dict[EntityName, ReferenceConfig],
+        dataset_config_uri: Optional[URI] = None,
+        **kwargs
     ) -> None:
         self.reference_entity_config = reference_entity_config
+        self.dataset_config_uri = dataset_config_uri
         """
         Configuration options for the reference data. This is likely to vary
         from backend to backend (e.g. might be locations and file types for
@@ -119,15 +154,30 @@ def load_table(self, config: ReferenceTable) -> EntityType:
         """Load reference entity from a database table"""
         raise NotImplementedError()
 
-    @abstractmethod
     def load_file(self, config: ReferenceFile) -> EntityType:
         "Load reference entity from a relative file path"
-        raise NotImplementedError()
+        if not self.dataset_config_uri:
+            raise AttributeError("dataset_config_uri must be specified if using relative paths")
+        target_location = fh.build_relative_uri(self.dataset_config_uri, config.filename)
+        if isinstance(_get_implementation(self.dataset_config_uri), LocalFilesystemImplementation):
+            target_location = fh.file_uri_to_local_path(target_location).as_posix()
+        try:
+            impl = self.__reader_functions__[config.file_extension]
+            return impl(self, target_location)
+        except KeyError:
+            raise RefdataLacksFileExtensionSupport(file_extension=config.file_extension)
 
-    @abstractmethod
     def load_uri(self, config: ReferenceURI) -> EntityType:
         "Load reference entity from an absolute URI"
-        raise NotImplementedError()
+        if isinstance(_get_implementation(config.uri), LocalFilesystemImplementation):
+            target_location = fh.file_uri_to_local_path(config.uri).as_posix()
+        else:
+            target_location = config.uri
+        try:
+            impl = self.__reader_functions__[config.file_extension]
+            return impl(self, target_location)
+        except KeyError:
+            raise RefdataLacksFileExtensionSupport(file_extension=config.file_extension)
 
     def load_entity(self, entity_name: EntityName, config: ReferenceConfig) -> EntityType:
         """Load a reference entity given the reference config"""

src/dve/core_engine/backends/exceptions.py

@@ -165,6 +165,21 @@ def __init__(self, *args: object, entity_type: Any) -> None:
     def get_message_preamble(self) -> EntityName:
         return f"Reader does not support reading directly to entity type {self.entity_type!r}"
 
+class RefdataLacksFileExtensionSupport(BackendError):
+    """An error raised when trying to load a refdata file where the loader
+       lacks support for the given file type
+
+    """
+
+    def __init__(self, *args: object, file_extension: str) -> None:
+        super().__init__(*args)
+        self.file_extension = file_extension
+        """The file extension that is not supported directly by the
+           refdata loader"""
+
+    def get_message_preamble(self) -> EntityName:
+        return f"Refdata loader does not support reading refdata from {self.file_extension} files"
+
 
 class EmptyFileError(ReaderErrorMixin, ValueError):
     """The read file was empty."""

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

@@ -3,6 +3,7 @@
 from typing import Optional
 
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
+import pyarrow.ipc as ipc
 
 import dve.parser.file_handling as fh
 from dve.core_engine.backends.base.reference_data import (
@@ -11,6 +12,7 @@
     ReferenceFile,
     ReferenceTable,
     ReferenceURI,
+    mark_refdata_file_extension
 )
 from dve.core_engine.type_hints import EntityName
 from dve.parser.file_handling.implementations.file import (
@@ -35,28 +37,20 @@ def __init__(
         reference_entity_config: dict[EntityName, ReferenceConfigUnion],
         **kwargs,
     ) -> None:
-        super().__init__(reference_entity_config, **kwargs)
+        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}")
-
-    def load_file(self, config: ReferenceFile) -> DuckDBPyRelation:
-        "Load reference entity from a relative file path"
-        if not self.dataset_config_uri:
-            raise AttributeError("dataset_config_uri must be specified if using relative paths")
-        target_location = fh.build_relative_uri(self.dataset_config_uri, config.filename)
-        if isinstance(_get_implementation(self.dataset_config_uri), LocalFilesystemImplementation):
-            target_location = file_uri_to_local_path(target_location).as_posix()
-        return self.connection.read_parquet(target_location)
-
-    def load_uri(self, config: ReferenceURI) -> DuckDBPyRelation:
-        "Load reference entity from an absolute URI"
-        if isinstance(_get_implementation(config.uri), LocalFilesystemImplementation):
-            target_location = file_uri_to_local_path(config.uri).as_posix()
-        else:
-            target_location = config.uri
-        return self.connection.read_parquet(target_location)
+        return self.connection.sql(f"select * from {config.fq_table_name}")         
+    
+    @mark_refdata_file_extension("parquet")
+    def load_parquet_file(self, uri: str) -> DuckDBPyRelation:
+        return self.connection.read_parquet(uri)
+    
+    @mark_refdata_file_extension("arrow")
+    def load_arrow_file(self, uri: str) -> DuckDBPyRelation:
+        return self.connection.from_arrow(ipc.open_file(uri).read_all())
+            

src/dve/core_engine/backends/implementations/spark/reference_data.py

@@ -12,6 +12,7 @@
     ReferenceFile,
     ReferenceTable,
     ReferenceURI,
+    mark_refdata_file_extension,
 )
 from dve.core_engine.type_hints import EntityName
 from dve.parser.type_hints import URI
@@ -31,18 +32,13 @@ def __init__(
         reference_entity_config: dict[EntityName, ReferenceConfig],
         **kwargs,
     ) -> None:
-        super().__init__(reference_entity_config, **kwargs)
+        super().__init__(reference_entity_config, self.dataset_config_uri, **kwargs)
         if not self.spark:
             raise AttributeError("Spark session must be provided")
 
     def load_table(self, config: ReferenceTable) -> DataFrame:
         return self.spark.table(f"{config.fq_table_name}")
 
-    def load_file(self, config: ReferenceFile) -> DataFrame:
-        if not self.dataset_config_uri:
-            raise AttributeError("dataset_config_uri must be specified if using relative paths")
-        target_location = fh.build_relative_uri(self.dataset_config_uri, config.filename)
-        return self.spark.read.parquet(target_location)
-
-    def load_uri(self, config: ReferenceURI) -> DataFrame:
-        return self.spark.read.parquet(config.uri)
+    @mark_refdata_file_extension("parquet")
+    def load_parquet_file(self, uri:str) -> DataFrame:
+        return self.spark.read.parquet(uri)

tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_ddb_refdata.py

@@ -0,0 +1,33 @@
+import shutil
+
+import duckdb
+import pytest
+from dve.core_engine.backends.implementations.duckdb.reference_data import DuckDBRefDataLoader
+from dve.core_engine.backends.base.reference_data import ReferenceFile
+from tempfile import TemporaryDirectory
+
+from tests.conftest import get_test_file_path
+
+@pytest.fixture(scope="module")
+def temp_working_dir():
+    with TemporaryDirectory(prefix="refdata_test") as tmp:
+        refdata_path = get_test_file_path("movies/refdata/movies_sequels.arrow")
+        shutil.copy(refdata_path.as_posix(), tmp)
+        yield tmp
+
+@pytest.fixture(scope="module")
+def ddb_refdata_loader(temp_working_dir):
+    DuckDBRefDataLoader.connection = duckdb.connect()
+    DuckDBRefDataLoader.dataset_config_uri = temp_working_dir
+    yield DuckDBRefDataLoader
+
+def test_load_arrow_file(ddb_refdata_loader):
+    config = {
+        "test_refdata": ReferenceFile(type="filename",
+                                      filename="./movies_sequels.arrow")
+    }
+    refdata_loader: DuckDBRefDataLoader = ddb_refdata_loader(config)
+    
+    test = refdata_loader.load_file(config.get("test_refdata"))
+    
+    assert test.shape == (2, 3)