fix: enhance duckdb casting to be less permissive of poorly formatted dates and trim whitespace

Author: stevenhsd

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

@@ -313,3 +313,85 @@ def duckdb_record_index(cls):
     setattr(cls, "add_record_index", _add_duckdb_record_index)
     setattr(cls, "drop_record_index", _drop_duckdb_record_index)
     return cls
+
+def _cast_as_ddb_type(field_expr:str, type_annotation:Any) -> str:
+    return f"try_cast({field_expr} as {get_duckdb_type_from_annotation(type_annotation)})"
+
+
+def get_duckdb_cast_statement_from_annotation(element_name:str,
+                                              type_annotation: Any,
+                                              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} [0-9]{2}:[0-9]{2}:[0-9]{2}$",
+                                              include_cast: bool = True) -> DuckDBPyType:
+    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_cast_statement_from_annotation(element_name, python_type, date_regex, timestamp_regex, include_cast)
+
+    # 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({element_name}, x -> {get_duckdb_cast_statement_from_annotation('x',element_type, date_regex, timestamp_regex, False)})"
+        return stmt if not include_cast 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, date_regex, timestamp_regex, include_cast) # add other expected params here
+    # 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, str] = {}
+        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_cast_statement_from_annotation(
+                f"{element_name}.{field_name}",
+                field_annotation,
+                date_regex,
+                timestamp_regex,
+                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})"
+        return stmt if not include_cast else _cast_as_ddb_type(stmt, type_annotation)
+
+    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():
+        if issubclass(type_, datetime):
+            stmt = f"CASE WHEN REGEXP_MATCHES(TRIM({element_name}), '{timestamp_regex}') THEN TRY_CAST(TRIM({element_name}) as TIMESTAMP) ELSE NULL END"
+            return stmt
+        if issubclass(type_, date):
+            stmt = f"CASE WHEN REGEXP_MATCHES(TRIM({element_name}), '{date_regex}') THEN TRY_CAST(TRIM({element_name}) as DATE) ELSE NULL END"
+            return stmt
+        duck_type = get_duckdb_type_from_annotation(type_)   
+        if duck_type:
+            stmt =  f"trim({element_name})"
+            return _cast_as_ddb_type(stmt, type_) if include_cast else stmt
+    raise ValueError(f"No equivalent DuckDB type for {type_annotation!r}")

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

@@ -3,17 +3,73 @@
 import datetime
 import tempfile
 from pathlib import Path
-from typing import Any
+from typing import Any, List
 
 import pytest
 import pyspark.sql.types as pst
 from duckdb import DuckDBPyRelation, DuckDBPyConnection
+from pydantic import BaseModel
 from pyspark.sql import Row, SparkSession
 
 from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import (
     _ddb_read_parquet,
-    duckdb_rel_to_dictionaries)
+    duckdb_rel_to_dictionaries,
+    get_duckdb_cast_statement_from_annotation,
+    get_duckdb_type_from_annotation)
 
+@pytest.fixture
+def casting_test_table(temp_ddb_conn):
+    _, conn = temp_ddb_conn
+    conn.sql("""CREATE TABLE test_casting (
+    str_test VARCHAR,
+    int_test VARCHAR,
+    date_test VARCHAR,
+    timestamp_test VARCHAR,
+    list_int_field VARCHAR[],
+    basic_model STRUCT(str_field VARCHAR, date_field VARCHAR),
+    another_model STRUCT(unique_id VARCHAR, basic_models STRUCT(str_field VARCHAR, date_field VARCHAR)[]))""")
+    
+    conn.sql("""INSERT INTO test_casting
+              VALUES(
+              'good_one',
+              '1',
+              '2024-11-13',
+              '2024-04-15 12:25:36',
+              ['1', '2', '3'],
+              {'str_field': 'test', 'date_field': '2024-12-11'},
+              {'unique_id': '1', "basic_models": [{'str_field': 'test_nest', 'date_field': '2020-01-04'}, {'str_field': 'test_nest2', 'date_field': '2020-01-05'}]}),
+              (
+              'dodgy_dates',
+              '2',
+              '24-11-13',
+              '2024-4-15 12:25:36',
+              ['4', '5', '6'],
+              {'str_field': 'test', 'date_field': '202-1-11'},
+              {'unique_id': '2', "basic_models": [{'str_field': 'test_dd', 'date_field': '20-01-04'}, {'str_field': 'test_dd2', 'date_field': '2020-1-5'}]})""")
+    
+    
+    yield temp_ddb_conn
+    
+    conn.sql("DROP TABLE IF EXISTS test_casting")
+    
+    
+
+class BasicModel(BaseModel):
+    str_field: str
+    date_field: datetime.date
+    
+class AnotherModel(BaseModel):
+    unique_id: int
+    basic_models: List[BasicModel]
+
+class CastingRecord(BaseModel):
+    str_test: str
+    int_test: int
+    date_test: datetime.date
+    timestamp_test: datetime.datetime
+    list_int_field: list[int]
+    basic_model: BasicModel
+    another_model: AnotherModel
 
 class TempConnection:
     """
@@ -25,6 +81,7 @@ def __init__(self, connection: DuckDBPyConnection) -> None:
         self._connection = connection
 
 
+
 @pytest.mark.parametrize(
     "outpath",
     [
@@ -94,4 +151,26 @@ def test_duckdb_rel_to_dictionaries(temp_ddb_conn: DuckDBPyConnection,
         res.append(chunk)
 
     assert res == data
+
+# add decimal check
+@pytest.mark.parametrize("field_name,field_type,cast_statement",
+                         [("str_test", str, "try_cast(trim(str_test) as VARCHAR)"),
+                          ("int_test", int, "try_cast(trim(int_test) as BIGINT)"),
+                          ("date_test", datetime.date,"CASE WHEN REGEXP_MATCHES(TRIM(date_test), '^[0-9]{4}-[0-9]{2}-[0-9]{2}$') THEN TRY_CAST(TRIM(date_test) as DATE) ELSE NULL END"),
+                          ("timestamp_test", datetime.datetime,"CASE WHEN REGEXP_MATCHES(TRIM(timestamp_test), '^[0-9]{4}-[0-9]{2}-[0-9]{2} [0-9]{2}:[0-9]{2}:[0-9]{2}$') THEN TRY_CAST(TRIM(timestamp_test) as TIMESTAMP) ELSE NULL END"),
+                          ("list_int_field", list[int], "try_cast(list_transform(list_int_field, x -> trim(x)) as BIGINT[])"),
+                          ("basic_model", BasicModel, "try_cast(struct_pack(str_field:= trim(basic_model.str_field),date_field:= CASE WHEN REGEXP_MATCHES(TRIM(basic_model.date_field), '^[0-9]{4}-[0-9]{2}-[0-9]{2}$') THEN TRY_CAST(TRIM(basic_model.date_field) as DATE) ELSE NULL END) as STRUCT(str_field VARCHAR, date_field DATE))"),
+                          ("another_model", AnotherModel, "try_cast(struct_pack(unique_id:= trim(another_model.unique_id),basic_models:= list_transform(another_model.basic_models, x -> struct_pack(str_field:= trim(x.str_field),date_field:= CASE WHEN REGEXP_MATCHES(TRIM(x.date_field), '^[0-9]{4}-[0-9]{2}-[0-9]{2}$') THEN TRY_CAST(TRIM(x.date_field) as DATE) ELSE NULL END))) as STRUCT(unique_id BIGINT, basic_models STRUCT(str_field VARCHAR, date_field DATE)[]))")])
+def test_get_duckdb_cast_statement_from_annotation(field_name, field_type, cast_statement):
+    assert get_duckdb_cast_statement_from_annotation(field_name, field_type) == cast_statement
+
+
+def test_use_cast_statements(casting_test_table):
+    _, conn = casting_test_table
+    test_rel = conn.sql("SELECT * from test_casting")
+    casting_statements = [ f"{get_duckdb_cast_statement_from_annotation(fld.name, fld.annotation)} as {fld.name}" for fld in CastingRecord.__fields__.values()]
+    test_rel = test_rel.project(",".join(casting_statements))
+    assert dict(zip(test_rel.columns, test_rel.dtypes)) == {fld.name: get_duckdb_type_from_annotation(fld.annotation) for fld in CastingRecord.__fields__.values()}
+    dodgy_date_rec = test_rel.pl()[1].to_dicts()[0]
+    assert not dodgy_date_rec.get("date_test") and not dodgy_date_rec.get("basic_model",{}).get("date_field") and all(not val.get("date_field") for val in dodgy_date_rec.get("another_model",{}).get("basic_models",[]))