Fix/sr ndit 1663

pyproject.toml

@@ -43,7 +43,7 @@ pandas = "2.3.3"
 polars = "0.20.31"
 pyarrow = "23.0.1"
 pydantic = "1.10.19"
-pyspark = ">=3.0.0,<=3.5.2"
+pyspark = ">=3.5.0,<=3.5.2"
 typing_extensions = "4.15.0"
 urllib3 = "2.7.0"  # dependency of boto3 & botocore
 

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

@@ -21,6 +21,7 @@
 
 from dve.common.error_utils import get_feedback_errors_uri
 from dve.core_engine.backends.base.utilities import _get_non_heterogenous_type
+from dve.core_engine.backends.utilities import DEFAULT_ISO_FORMATS, datetime_format_to_regex
 from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME
 from dve.core_engine.type_hints import URI, EntityName
 from dve.parser.file_handling.service import LocalFilesystemImplementation, _get_implementation
@@ -95,6 +96,8 @@ def __call__(self):
 """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())
@@ -361,7 +364,7 @@ def duckdb_record_index(cls):
 
 def _cast_as_ddb_type(field_expr: str, type_annotation: Any) -> str:
     """Cast to Duck DB type"""
-    return f"""try_cast({field_expr} as {get_duckdb_type_from_annotation(type_annotation)})"""
+    return f"""TRY_CAST({field_expr} as {get_duckdb_type_from_annotation(type_annotation)})"""
 
 
 def _ddb_safely_quote_name(field_name: str) -> str:
@@ -378,9 +381,6 @@ def get_duckdb_cast_statement_from_annotation(
     element_name: str,
     type_annotation: Any,
     parent_element: bool = True,
-    date_regex: str = r"^[0-9]{4}-[0-9]{2}-[0-9]{2}$",
-    timestamp_regex: str = r"^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}((\+|\-)[0-9]{2}:[0-9]{2})?$",  # pylint: disable=C0301
-    time_regex: str = r"^[0-9]{2}:[0-9]{2}:[0-9]{2}$",
 ) -> str:
     """Generate casting statements for duckdb relations from type annotations"""
     type_origin = get_origin(type_annotation)
@@ -391,19 +391,19 @@ def get_duckdb_cast_statement_from_annotation(
     if type_origin is Union:
         python_type = _get_non_heterogenous_type(get_args(type_annotation))
         return get_duckdb_cast_statement_from_annotation(
-            element_name, python_type, parent_element, date_regex, timestamp_regex
+            element_name, python_type, parent_element,
         )
 
     # 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))
-        stmt = f"list_transform({quoted_name}, x -> {get_duckdb_cast_statement_from_annotation('x',element_type, False, date_regex, timestamp_regex)})"  # pylint: disable=C0301
+        stmt = f"LIST_TRANSFORM({quoted_name}, x -> {get_duckdb_cast_statement_from_annotation('x',element_type, False,)})"  # pylint: disable=C0301
         return stmt if not parent_element else _cast_as_ddb_type(stmt, type_annotation)
 
     if type_origin is Annotated:
         python_type, *other_args = get_args(type_annotation)  # pylint: disable=unused-variable
         return get_duckdb_cast_statement_from_annotation(
-            element_name, python_type, parent_element, date_regex, timestamp_regex
+            element_name, python_type, parent_element,
         )  # add other expected params here
     # Ensure that we have a concrete type at this point.
     if not isinstance(type_annotation, type):
@@ -428,15 +428,15 @@ def get_duckdb_cast_statement_from_annotation(
                 continue
 
             fields[field_name] = get_duckdb_cast_statement_from_annotation(
-                f"{element_name}.{field_name}", field_annotation, False, date_regex, timestamp_regex
+                f"{element_name}.{field_name}", field_annotation, False,
             )
 
         if not fields:
             raise ValueError(
                 f"No type annotations in dict/dataclass type (got {type_annotation!r})"
             )
         cast_exprs = ",".join([f'"{nme}":= {stmt}' for nme, stmt in fields.items()])
-        stmt = f"struct_pack({cast_exprs})"
+        stmt = f"STRUCT_PACK({cast_exprs})"
         return stmt if not parent_element else _cast_as_ddb_type(stmt, type_annotation)
 
     if type_annotation is list:
@@ -447,18 +447,21 @@ def get_duckdb_cast_statement_from_annotation(
         raise ValueError(f"dict must be `typing.TypedDict` subclass, got {type_annotation!r}")
 
     for type_ in type_annotation.mro():
+        _date_format = getattr(type_, "DATE_FORMAT", DEFAULT_ISO_FORMATS.get(type_, DEFAULT_ISO_FORMATS.get(datetime)))
+        dt_cast_statement = fr"CASE WHEN REGEXP_FULL_MATCH(TRIM({quoted_name}), '{datetime_format_to_regex(_date_format)}') THEN TRY_STRPTIME(TRIM({quoted_name}), '{_date_format}') ELSE NULL END"
+
         # datetime is subclass of date, so needs to be handled first
         if issubclass(type_, datetime):
-            stmt = rf"CASE WHEN REGEXP_MATCHES(TRIM({quoted_name}), '{timestamp_regex}') THEN TRY_CAST(TRIM({quoted_name}) as TIMESTAMP) ELSE NULL END"  # pylint: disable=C0301
+            stmt = rf"TRY_CAST({dt_cast_statement} as TIMESTAMP)"
             return stmt
         if issubclass(type_, date):
-            stmt = rf"CASE WHEN REGEXP_MATCHES(TRIM({quoted_name}), '{date_regex}') THEN TRY_CAST(TRIM({quoted_name}) as DATE) ELSE NULL END"  # pylint: disable=C0301
+            stmt = rf"TRY_CAST({dt_cast_statement} as DATE)"
             return stmt
         if issubclass(type_, time):
-            stmt = rf"CASE WHEN REGEXP_MATCHES(TRIM({quoted_name}), '{time_regex}') THEN TRY_CAST(TRIM({quoted_name}) as TIME) ELSE NULL END"  # pylint: disable=C0301
+            stmt = rf"TRY_CAST({dt_cast_statement} as TIME)"
             return stmt
         duck_type = get_duckdb_type_from_annotation(type_)
         if duck_type:
-            stmt = f"trim({quoted_name})"
+            stmt = f"TRIM({quoted_name})"
             return _cast_as_ddb_type(stmt, type_) if parent_element else stmt
     raise ValueError(f"No equivalent DuckDB type for {type_annotation!r}")

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

@@ -7,6 +7,7 @@
 from uuid import uuid4
 
 from pydantic import BaseModel
+from pydantic.fields import ModelField
 from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql import functions as sf
 from pyspark.sql.functions import col, lit
@@ -30,6 +31,7 @@
     spark_read_parquet,
     spark_record_index,
     spark_write_parquet,
+    get_spark_cast_statement_from_annotation
 )
 from dve.core_engine.backends.implementations.spark.types import SparkEntities
 from dve.core_engine.backends.metadata.contract import DataContractMetadata
@@ -102,6 +104,9 @@ def apply_data_contract(
 
         successful = True
         for entity_name, record_df in entities.items():
+            entity_fields: dict[str, ModelField] = contract_metadata.schemas[
+                    entity_name
+                ].__fields__
             spark_schema = get_type_from_annotation(contract_metadata.schemas[entity_name])
             spark_schema.add(StructField(RECORD_INDEX_COLUMN_NAME, LongType()))
             if df_is_empty(record_df):
@@ -145,11 +150,13 @@ def apply_data_contract(
 
             try:
                 record_df = record_df.select(
-                    [
-                        col(column.name).cast(column.dataType)
-                        for column in spark_schema
-                        if column.name in record_df.columns
-                    ]
+                    *[
+                        get_spark_cast_statement_from_annotation(fld, mdl_field.annotation).alias(fld) 
+                        if fld in record_df.columns
+                        else lit(None).cast(get_type_from_annotation(mdl_field.annotation).alias(fld))
+                        for fld, mdl_field in entity_fields.items()
+                    ],
+                    col(RECORD_INDEX_COLUMN_NAME).cast(LongType()).alias(RECORD_INDEX_COLUMN_NAME)
                 )
             except Exception as err:  # pylint: disable=broad-except
                 successful = False

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

@@ -27,6 +27,7 @@
 from typing_extensions import Annotated, Protocol, TypedDict, get_args, get_origin, get_type_hints
 
 from dve.common.error_utils import get_feedback_errors_uri
+from dve.core_engine.backends.utilities import DEFAULT_ISO_FORMATS, datetime_format_to_regex
 from dve.core_engine.backends.base.utilities import _get_non_heterogenous_type
 from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME
 from dve.core_engine.type_hints import URI, EntityName
@@ -99,6 +100,21 @@ def __post_init__(self):
 }
 """A mapping of Python types to the equivalent Spark types."""
 
+PYTHON_TO_JAVA_DT_FORMAT_MAP: dict[str, str] = {
+    "%Y": "yyyy",
+    "%y": "yy",
+    "%m": "MM",
+    "%d": "dd",
+    "%H": "HH",
+    "%M": "mm",
+    "%S": "ss",
+    "%z": "XX",
+    "%Z": "z"
+}
+"""A mapping of python to java datetime component formats. Not exhaustive but will hopefully support
+   all requirements.
+"""
+
 
 PydanticModel = TypeVar("PydanticModel", bound=BaseModel)
 """An Pydantic model."""
@@ -257,6 +273,12 @@ def create_udf(function: FourArgWrappable) -> FourArgWrapped:  # pragma: no cove
     pass
 
 
+def python_to_java_datetime_format(datetime_fmt: str) -> str:
+    "Helper to convert python datetime formats to Java datetime formats"
+    for pt, jt in PYTHON_TO_JAVA_DT_FORMAT_MAP.items():
+        datetime_fmt = datetime_fmt.replace(pt, jt)
+    return datetime_fmt.replace("T", "'T'")
+
 def create_udf(function: Callable) -> Callable:
     """Get a UDF representing a specific function."""
     if not callable(function):
@@ -508,9 +530,7 @@ def _spark_safely_quote_name(field_name: str) -> str:
 def get_spark_cast_statement_from_annotation(
     element_name: str,
     type_annotation: Any,
-    parent_element: bool = True,
-    date_regex: str = r"^[0-9]{4}-[0-9]{2}-[0-9]{2}$",
-    timestamp_regex: str = r"^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}((\\+|\\-)[0-9]{2}:[0-9]{2})?$",  # pylint: disable=C0301
+    parent_element: bool = True
 ):
     """Generate casting statements for spark dataframes based on type annotations"""
     type_origin = get_origin(type_annotation)
@@ -521,19 +541,18 @@ def get_spark_cast_statement_from_annotation(
     if type_origin is Union:
         python_type = _get_non_heterogenous_type(get_args(type_annotation))
         return get_spark_cast_statement_from_annotation(
-            element_name, python_type, parent_element, date_regex, timestamp_regex
-        )
+            element_name, python_type, parent_element)
 
     # 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))
-        stmt = f"transform({quoted_name}, x -> {get_spark_cast_statement_from_annotation('x',element_type, False, date_regex, timestamp_regex)})"  # pylint: disable=C0301
+        stmt = f"TRANSFORM({quoted_name}, x -> {get_spark_cast_statement_from_annotation('x',element_type, False)})"  # pylint: disable=C0301
         return stmt if not parent_element else _cast_as_spark_type(stmt, type_annotation)
 
     if type_origin is Annotated:
         python_type, *_ = get_args(type_annotation)  # pylint: disable=unused-variable
         return get_spark_cast_statement_from_annotation(
-            element_name, python_type, parent_element, date_regex, timestamp_regex
+            element_name, python_type, parent_element
         )  # add other expected params here
     # Ensure that we have a concrete type at this point.
     if not isinstance(type_annotation, type):
@@ -558,15 +577,15 @@ def get_spark_cast_statement_from_annotation(
                 continue
 
             fields[field_name] = get_spark_cast_statement_from_annotation(
-                f"{element_name}.{field_name}", field_annotation, False, date_regex, timestamp_regex
+                f"{element_name}.{field_name}", field_annotation, False
             )
 
         if not fields:
             raise ValueError(
                 f"No type annotations in dict/dataclass type (got {type_annotation!r})"
             )
         cast_exprs = ",".join([f"{stmt} AS `{nme}`" for nme, stmt in fields.items()])
-        stmt = f"struct({cast_exprs})"
+        stmt = f"STRUCT({cast_exprs})"
         return stmt if not parent_element else _cast_as_spark_type(stmt, type_annotation)
     if type_annotation is list:
         raise ValueError(
@@ -576,15 +595,15 @@ def get_spark_cast_statement_from_annotation(
         raise ValueError(f"dict must be `typing.TypedDict` subclass, got {type_annotation!r}")
 
     for type_ in type_annotation.mro():
+        _date_format = getattr(type_, "DATE_FORMAT", DEFAULT_ISO_FORMATS.get(type_, DEFAULT_ISO_FORMATS.get(dt.datetime)))
+        dt_cast_statement = f"CASE WHEN REGEXP(TRIM({quoted_name}), '{datetime_format_to_regex(_date_format)}') THEN TRY_TO_TIMESTAMP(TRIM({quoted_name}), \"{python_to_java_datetime_format(_date_format)}\") ELSE NULL END"
         # datetime is subclass of date, so needs to be handled first
         if issubclass(type_, dt.datetime):
-            stmt = rf"CASE WHEN REGEXP(TRIM({quoted_name}), '{timestamp_regex}') THEN TRIM({quoted_name}) ELSE NULL END"  # pylint: disable=C0301
-            return _cast_as_spark_type(stmt, type_) if parent_element else stmt
+            return _cast_as_spark_type(dt_cast_statement, type_) if parent_element else dt_cast_statement
         if issubclass(type_, dt.date):
-            stmt = rf"CASE WHEN REGEXP(TRIM({quoted_name}), '{date_regex}') THEN TRIM({quoted_name}) ELSE NULL END"  # pylint: disable=C0301
-            return _cast_as_spark_type(stmt, type_) if parent_element else stmt
+            return _cast_as_spark_type(dt_cast_statement, type_) if parent_element else dt_cast_statement
         spark_type = get_type_from_annotation(type_)
         if spark_type:
-            stmt = f"trim({quoted_name})"
+            stmt = f"TRIM({quoted_name})"
             return _cast_as_spark_type(stmt, type_) if parent_element else stmt
     raise ValueError(f"No equivalent Spark type for {type_annotation!r}")

src/dve/core_engine/backends/utilities.py

@@ -23,6 +23,47 @@
 else:
     from typing import Annotated, get_args, get_origin, get_type_hints
 
+DEFAULT_ISO_FORMATS: dict[type, str] = {
+    date: "%Y-%m-%d",
+    datetime: "%Y-%m-%dT%H:%M:%S",
+    time: "%H:%M:%S"
+}
+"""Mapping of default ISO formats to use when date format not supplied"""
+
+PYTHON_DATE_FORMAT_REGEX_HELPER: dict[str, str] = {
+    "Y": r"[0-9]{4}",
+    "y": r"[0-9]{2}",
+    "m": r"[0-9]{2}",
+    "d": r"[0-9]{2}",
+    "H": r"[0-9]{2}",
+    "M": r"[0-9]{2}",
+    "S": r"[0-9]{2}",
+    "z": r"(\+|\-)?[0-9]+(\.[0-9]*)?",
+    "Z": r"[A-Z]{0,3}",
+}
+
+REGEXP_NEED_ESCAPE_CHARS: tuple[str] = ("+", "-", ".")
+"""Helper to map python date format to regexp expression. Not exhaustive, but aims to cover
+   all foreseen use cases."""
+
+def datetime_format_to_regex(format_str: str) -> str:
+    """
+    Helper function to convert python datetime formats to regexp string checks for casting
+    purposes.
+    """
+
+    tokens = list(format_str.replace("%", ""))
+
+    for idx, tkn in enumerate(tokens):
+        if rpl := PYTHON_DATE_FORMAT_REGEX_HELPER.get(tkn):
+            tokens[idx] = rpl
+        elif tkn in REGEXP_NEED_ESCAPE_CHARS:
+            tokens[idx] = rf"\{tkn}"
+        else:
+            continue
+        tokens = [PYTHON_DATE_FORMAT_REGEX_HELPER.get(tkn, tkn) for tkn in tokens]
+    return "".join(["^", *tokens, "$"])
+
 PYTHON_TYPE_TO_POLARS_TYPE: dict[type, PolarsType] = {
     # issue with decimal conversion at the moment...
     str: pl.Utf8,  # type: ignore

tests/features/books.feature

@@ -49,3 +49,30 @@ Feature: Pipeline tests using the books dataset
         Then the latest audit record for the submission is marked with processing status error_report
         When I run the error report phase
         Then An error report is produced
+
+    Scenario: Validate complex nested XML data (spark)
+        Given I submit the books file nested_books.XML for processing
+        And A spark pipeline is configured with schema file 'nested_books.dischema.json'
+        And I add initial audit entries for the submission
+        Then the latest audit record for the submission is marked with processing status file_transformation
+        When I run the file transformation phase
+        Then the header entity is stored as a parquet after the file_transformation phase
+        And the nested_books entity is stored as a parquet after the file_transformation phase
+        And the latest audit record for the submission is marked with processing status data_contract
+        When I run the data contract phase
+        Then there is 1 record rejection from the data_contract phase
+        And the header entity is stored as a parquet after the data_contract phase
+        And the nested_books entity is stored as a parquet after the data_contract phase
+        And the latest audit record for the submission is marked with processing status business_rules
+        When I run the business rules phase
+        Then The rules restrict "nested_books" to 3 qualifying records
+        And The entity "nested_books" contains an entry for "17.85" in column "total_value_of_books"
+        And the nested_books entity is stored as a parquet after the business_rules phase
+        And the latest audit record for the submission is marked with processing status error_report
+        When I run the error report phase
+        Then An error report is produced
+        And The statistics entry for the submission shows the following information
+            | parameter                | value |
+            | record_count             | 4     |
+            | number_record_rejections | 2     |
+            | number_warnings          | 0     |