feat: bug fixes with readers and base pipeline. Added spark csv reader. Added new dischema files for tests

Author: stevenhsd

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

@@ -1,3 +1,4 @@
+from dve.core_engine.backends.implementations.duckdb.readers.json import DuckDBJSONReader
 from dve.core_engine.backends.readers import register_reader
 
 from .contract import DuckDBDataContract
@@ -6,6 +7,7 @@
 from .rules import DuckDBStepImplementations
 
 register_reader(DuckDBCSVReader)
+register_reader(DuckDBJSONReader)
 register_reader(DuckDBXMLStreamReader)
 
 __all__ = [

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

@@ -6,7 +6,7 @@
 from datetime import date, datetime
 from decimal import Decimal
 from pathlib import Path
-from typing import Any, ClassVar, Dict, Set, Union
+from typing import Any, ClassVar, Dict, Generator, Iterator, Set, Union
 from urllib.parse import urlparse
 
 import duckdb.typing as ddbtyp
@@ -224,13 +224,21 @@ def _ddb_read_parquet(
 
 
 def _ddb_write_parquet(  # pylint: disable=unused-argument
-    self, entity: DuckDBPyRelation, target_location: URI, **kwargs
+    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)
     return target_location

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

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

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

@@ -3,42 +3,17 @@
 # pylint: disable=arguments-differ
 from typing import Any, Dict, Iterator, Type
 
-from duckdb import DuckDBPyConnection, DuckDBPyRelation, read_csv
+from duckdb import DuckDBPyConnection, DuckDBPyRelation, read_csv, default_connection
 from pydantic import BaseModel
-from typing_extensions import Literal
 
 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
 
-SQLType = Literal[
-    "BIGINT",
-    "BIT",
-    "BLOB",
-    "BOOLEAN",
-    "DATE",
-    "DECIMAL",
-    "DOUBLE",
-    "HUGEINT",
-    "INTEGER",
-    "INTERVAL",
-    "REAL",
-    "SMALLINT",
-    "TIME",
-    "UBIGINT",
-    "UHUGEINT",
-    "UINTEGER",
-    "USMALLINT",
-    "UTINYINT",
-    "UUID",
-    "VARCHAR",
-]
-"""SQL types recognised in duckdb"""
-
-
 @duckdb_write_parquet
 class DuckDBCSVReader(BaseFileReader):
     """A reader for CSV files"""
@@ -47,21 +22,21 @@ class DuckDBCSVReader(BaseFileReader):
     # TODO - stringify or not
     def __init__(
         self,
-        header: bool,
-        delim: str,
-        connection: DuckDBPyConnection,
+        header: bool = True,
+        delim: str = ",",
+        connection: DuckDBPyConnection = None,
     ):
         self.header = header
         self.delim = delim
-        self._connection = connection
+        self._connection = connection if connection else default_connection
 
         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)
+        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

src/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 typing import Any, Dict, Iterator, Optional, Type
+
+from duckdb import DuckDBPyRelation, read_json
+from pydantic import BaseModel
+from typing_extensions import Literal
+
+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,
+        format: Optional[str] = "array"
+    ):
+        self._format = 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._format)

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

@@ -1,10 +1,11 @@
 # mypy: disable-error-code="attr-defined"
 """An xml reader to create duckdb relations"""
 
-from typing import Dict, Type
+from typing import Dict, Optional, Type
 
+import pandas as pd
 import polars as pl
-from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from duckdb import DuckDBPyConnection, DuckDBPyRelation, default_connection
 from pydantic import BaseModel
 
 from dve.core_engine.backends.base.reader import read_function
@@ -20,15 +21,17 @@
 class DuckDBXMLStreamReader(XMLStreamReader):
     """A reader for XML files"""
 
-    def __init__(self, ddb_connection: DuckDBPyConnection, **kwargs):
-        self.ddb_connection = ddb_connection
+    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"""
         polars_schema: Dict[str, pl.DataType] = {  # type: ignore
-            fld.name: get_polars_type_from_annotation(fld.type_)
+            fld.name: get_polars_type_from_annotation(fld.annotation)
             for fld in stringify_model(schema).__fields__.values()
         }
 

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

@@ -2,11 +2,36 @@
 
 # pylint: disable=C0103
 from typing import MutableMapping
+from typing_extensions import Literal
 
 from duckdb import DuckDBPyRelation
 
 from dve.core_engine.type_hints import EntityName
 
+SQLType = Literal[
+    "BIGINT",
+    "BIT",
+    "BLOB",
+    "BOOLEAN",
+    "DATE",
+    "DECIMAL",
+    "DOUBLE",
+    "HUGEINT",
+    "INTEGER",
+    "INTERVAL",
+    "REAL",
+    "SMALLINT",
+    "TIME",
+    "UBIGINT",
+    "UHUGEINT",
+    "UINTEGER",
+    "USMALLINT",
+    "UTINYINT",
+    "UUID",
+    "VARCHAR",
+]
+"""SQL types recognised in duckdb"""
+
 Source = DuckDBPyRelation
 """The source entity for a join. This will be aliased to the source entity name."""
 Target = DuckDBPyRelation

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

@@ -4,11 +4,20 @@
 
 from .backend import SparkBackend
 from .contract import SparkDataContract
-from .readers import SparkXMLReader
+from .readers import (
+    SparkCSVReader,
+    SparkJSONReader,
+    SparkXMLReader,
+    SparkXMLStreamReader
+)
 from .reference_data import SparkRefDataLoader
 from .rules import SparkStepImplementations
 
+register_reader(SparkCSVReader)
+register_reader(SparkJSONReader)
 register_reader(SparkXMLReader)
+register_reader(SparkXMLStreamReader)
+
 
 __all__ = [
     "SparkBackend",

src/dve/core_engine/backends/implementations/spark/readers/__init__.py

@@ -1,11 +1,22 @@
 """Spark-specific readers."""
 
+from dve.core_engine.backends.implementations.spark.readers.csv import (
+    SparkCSVReader
+)
+from dve.core_engine.backends.implementations.spark.readers.json import (
+    SparkJSONReader
+)
+
 from dve.core_engine.backends.implementations.spark.readers.xml import (
     SparkXMLReader,
     SparkXMLStreamReader,
 )
 
+
+
 __all__ = [
+    "SparkCSVReader"
+    "SparkJSONReader",
     "SparkXMLReader",
     "SparkXMLStreamReader",
 ]

src/dve/core_engine/backends/implementations/spark/readers/csv.py

@@ -0,0 +1,81 @@
+"""A reader implementation using the Databricks Spark XML reader."""
+
+
+from typing import Any, Dict, Iterator, Optional, Type
+
+from pydantic import BaseModel
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.types import StructType
+
+
+from dve.core_engine.backends.base.reader import BaseFileReader, read_function
+from dve.core_engine.backends.exceptions import EmptyFileError
+from dve.core_engine.backends.implementations.spark.spark_helpers import (
+    get_type_from_annotation,
+    spark_write_parquet,
+)
+from dve.core_engine.type_hints import URI, EntityName
+from dve.parser.file_handling import get_content_length
+
+
+@spark_write_parquet
+class SparkCSVReader(BaseFileReader):
+    """A Spark reader for JSON files."""
+
+    def __init__(
+        self,
+        *,
+        delimiter: str = ",",
+        escape_char: str = "\\",
+        quote_char: str = '"',
+        header: bool = True,
+        multi_line: bool = False,
+        encoding: str = "utf-8-sig",
+        spark_session: SparkSession = None
+    ) -> None:
+        
+        self.delimiter = delimiter
+        self.escape_char = escape_char
+        self.encoding = encoding
+        self.quote_char = quote_char
+        self.header = header
+        self.multi_line = multi_line
+        self.spark_session = spark_session if spark_session else SparkSession.builder.getOrCreate()
+        
+        super().__init__()
+
+    def read_to_py_iterator(
+        self, resource: URI, entity_name: EntityName, schema: Type[BaseModel]
+    ) -> Iterator[Dict[URI, Any]]:
+        df = self.read_to_dataframe(resource, entity_name, schema)
+        yield from (record.asDict(True) for record in df.toLocalIterator())
+
+    @read_function(DataFrame)
+    def read_to_dataframe(
+        self,
+        resource: URI,
+        entity_name: EntityName,  # pylint: disable=unused-argument
+        schema: Type[BaseModel],
+    ) -> DataFrame:
+        """Read an JSON file directly to a Spark DataFrame.
+
+        """
+        if get_content_length(resource) == 0:
+            raise EmptyFileError(f"File at {resource} is empty.")
+
+        spark_schema: StructType = get_type_from_annotation(schema)
+        kwargs = {
+            "sep": self.delimiter,
+            "header": self.header,
+            "escape": self.escape_char,
+            "quote": self.quote_char, 
+            "multiLine": self.multi_line,
+            
+        }
+        
+        return (
+            self.spark_session.read.format("csv")
+            .options(**kwargs)  # type: ignore
+            .load(resource, schema=spark_schema)
+        )
+