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
@@ -361,7 +362,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 +379,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 +389,23 @@ 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 +430,17 @@ 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 +451,23 @@ 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: str = getattr(  # type: ignore
+            type_, "DATE_FORMAT", DEFAULT_ISO_FORMATS.get(type_, DEFAULT_ISO_FORMATS.get(datetime))
+        )
+        dt_cast_statement = rf"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"  # pylint: disable=C0301
+
         # 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
@@ -26,6 +27,7 @@
 )
 from dve.core_engine.backends.implementations.spark.spark_helpers import (
     df_is_empty,
+    get_spark_cast_statement_from_annotation,
     get_type_from_annotation,
     spark_read_parquet,
     spark_record_index,
@@ -86,7 +88,7 @@ def create_entity_from_py_iterator(
             schema=get_type_from_annotation(schema),
         )
 
-    # pylint: disable=R0915
+    # pylint: disable=R0914,R0915
     def apply_data_contract(
         self,
         working_dir: URI,
@@ -102,6 +104,7 @@ 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 +148,19 @@ 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

@@ -28,6 +28,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
 
@@ -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."""
@@ -272,6 +288,13 @@ def create_udf(function: Callable) -> Callable:
     return udf(function, returnType=return_type)
 
 
+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'")
+
+
 SupportedBaseType = Union[str, int, bool, float, Decimal, dt.date, dt.datetime]
 """Supported base types for Spark literals."""
 SparkLiteralType = Union[  # type: ignore
@@ -506,11 +529,7 @@ def _spark_safely_quote_name(field_name: str) -> str:
 
 # pylint: disable=R0801
 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
+    element_name: str, type_annotation: Any, parent_element: bool = True
 ):
     """Generate casting statements for spark dataframes based on type annotations"""
     type_origin = get_origin(type_annotation)
@@ -520,20 +539,18 @@ def get_spark_cast_statement_from_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_spark_cast_statement_from_annotation(
-            element_name, python_type, parent_element, date_regex, timestamp_regex
-        )
+        return get_spark_cast_statement_from_annotation(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 +575,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 +593,29 @@ 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: str = getattr(  # type: ignore
+            type_,
+            "DATE_FORMAT",
+            DEFAULT_ISO_FORMATS.get(type_, DEFAULT_ISO_FORMATS.get(dt.datetime)),
+        )
+
+        # pylint: disable=C0301
+        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"  # pylint: disable=C0301
         # 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,49 @@
 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, str, 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     |