test: added small tests for duckdb and spark utils

Author: stevenhsd

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

@@ -1,4 +1,5 @@
 """Implementation of duckdb backend"""
+
 from dve.core_engine.backends.implementations.duckdb.readers.json import DuckDBJSONReader
 from dve.core_engine.backends.readers import register_reader
 

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

@@ -1,6 +1,7 @@
 """Utility objects for use with duckdb backend"""
 
 import itertools
+
 from dve.core_engine.backends.base.utilities import _split_multiexpr_string
 
 
@@ -27,10 +28,9 @@ def expr_array_to_columns(expressions: list[str]) -> list[str]:
     """Create list of duckdb expressions from list of expressions"""
     return list(
         itertools.chain.from_iterable(
-            _split_multiexpr_string(expression) 
-            for expression in expressions
-            )
+            _split_multiexpr_string(expression) for expression in expressions
         )
+    )
 
 
 def multiexpr_string_to_columns(expressions: str) -> list[str]:

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

@@ -12,14 +12,7 @@
 from dataclasses import dataclass, is_dataclass
 from decimal import Decimal
 from functools import wraps
-from typing import (
-    Any,
-    ClassVar,
-    Optional,
-    TypeVar,
-    Union,
-    overload,
-)
+from typing import Any, ClassVar, Optional, TypeVar, Union, overload
 
 from delta.exceptions import ConcurrentAppendException, DeltaConcurrentModificationException
 from pydantic import BaseModel

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

@@ -1,6 +1,7 @@
 """Some utilities which are useful for implementing Spark transformations."""
 
 import datetime as dt
+import itertools
 from collections.abc import Callable
 from json import JSONEncoder
 from operator import and_, or_
@@ -70,7 +71,13 @@ def expr_mapping_to_columns(expressions: ExpressionMapping) -> list[Column]:
 
 def expr_array_to_columns(expressions: ExpressionArray) -> list[Column]:
     """Convert an array of expressions to a list of columns."""
-    return list(map(sf.expr, expressions))
+
+    _expr_list = list(
+        itertools.chain.from_iterable(
+            _split_multiexpr_string(expression) for expression in expressions
+        )
+    )
+    return list(map(sf.expr, _expr_list))
 
 
 def multiexpr_string_to_columns(expressions: MultiExpression) -> list[Column]:

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

@@ -3,14 +3,7 @@
 
 import re
 from collections.abc import Collection, Iterator
-from typing import (
-    IO,
-    Any,
-    GenericAlias,  # type: ignore
-    Optional,
-    Union,
-    overload
-)
+from typing import IO, Any, GenericAlias, Optional, Union, overload  # type: ignore
 
 import polars as pl
 from lxml import etree  # type: ignore

src/dve/core_engine/backends/utilities.py

@@ -4,8 +4,8 @@
 from dataclasses import is_dataclass
 from datetime import date, datetime
 from decimal import Decimal
-from typing import Any, ClassVar, Union
 from typing import GenericAlias  # type: ignore
+from typing import Any, ClassVar, Union
 
 import polars as pl  # type: ignore
 from polars.datatypes.classes import DataTypeClass as PolarsType
@@ -39,7 +39,9 @@
 
 def stringify_type(type_: Union[type, GenericAlias]) -> type:
     """Stringify an individual type."""
-    if isinstance(type_, type) and not isinstance(type_, GenericAlias):  # A model, return the contents.  # pylint: disable=C0301
+    if isinstance(type_, type) and not isinstance(
+        type_, GenericAlias
+    ):  # A model, return the contents.  # pylint: disable=C0301
         if issubclass(type_, BaseModel):
             return stringify_model(type_)
 

src/dve/core_engine/message.py

@@ -2,8 +2,8 @@
 
 import copy
 import datetime as dt
-import operator
 import json
+import operator
 from collections.abc import Callable
 from decimal import Decimal
 from functools import reduce

src/dve/core_engine/type_hints.py

@@ -6,12 +6,13 @@
 from pathlib import Path
 from queue import Queue as ThreadQueue
 from typing import TYPE_CHECKING, Any, List, Optional, TypeVar, Union  # pylint: disable=W1901
-# TODO - cannot remove List from Typing. See L60 for details.
 
 from pyspark.sql import DataFrame
 from pyspark.sql.types import StructType
 from typing_extensions import Literal, ParamSpec, get_args
 
+# TODO - cannot remove List from Typing. See L60 for details.
+
 
 if TYPE_CHECKING:  # pragma: no cover
     from dve.core_engine.message import FeedbackMessage

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

@@ -0,0 +1,56 @@
+from typing import Dict, List
+import pytest
+
+from dve.core_engine.backends.implementations.duckdb.utilities import (
+    expr_mapping_to_columns,
+    expr_array_to_columns,
+)
+
+
+@pytest.mark.parametrize(
+    ["expressions", "expected"],
+    [
+        (
+            {"size(array_field)": "field_length", "another_field": "rename_another_field"},
+            ["size(array_field) as field_length", "another_field as rename_another_field"],
+        ),
+    ],
+)
+def test_expr_mapping_to_columns(expressions: Dict[str, str], expected: list[str]):
+    observed = expr_mapping_to_columns(expressions)
+    assert observed == expected
+
+
+@pytest.mark.parametrize(
+    ["expressions", "expected"],
+    [
+        (
+            [
+                "a_field",
+                "another_field as renamed",
+                "struct(a_field, another_field) as struct_field",
+            ],
+            [
+                "a_field",
+                "another_field as renamed",
+                "struct(a_field, another_field) as struct_field",
+            ],
+        ),
+        (
+            [
+                "size(array_field)",
+                "another_field as rename_another_field",
+                "a_dynamic_field, another_dynamic_field",
+            ],
+            [
+                "size(array_field)",
+                "another_field as rename_another_field",
+                "a_dynamic_field",
+                "another_dynamic_field",
+            ],
+        ),
+    ],
+)
+def test_expr_array_to_columns(expressions: Dict[str, str], expected: list[str]):
+    observed = expr_array_to_columns(expressions)
+    assert observed == expected

tests/test_core_engine/test_backends/test_implementations/test_spark/test_spark_utils.py

@@ -0,0 +1,38 @@
+import pytest
+from pyspark.sql.functions import expr
+
+from dve.core_engine.backends.implementations.spark.utilities import (
+    expr_mapping_to_columns,
+    expr_array_to_columns,
+)
+
+
+@pytest.mark.parametrize(
+    ["expressions"],
+    [
+        (
+            {"size(array_field)": "field_length", "another_field": "rename_another_field"},
+        ),
+    ]
+)
+def test_expr_mapping_to_columns(spark, expressions: dict[str, str]):
+    observed = expr_mapping_to_columns(expressions)
+    assert [cl._jc.toString() for cl in observed] == [expr(expression).alias(rename)._jc.toString() for expression, rename in expressions.items()]
+
+
+@pytest.mark.parametrize(
+    ["expressions", "expected"],
+    [   
+        (
+            ["a_field", "another_field as renamed", "struct(a_field, another_field) as struct_field"],
+            ["a_field", "another_field as renamed", "struct(a_field, another_field) as struct_field"]
+            ),
+        (
+            ["size(array_field)", "another_field as rename_another_field", "a_dynamic_field, another_dynamic_field"],
+            ["size(array_field)", "another_field as rename_another_field", "a_dynamic_field", "another_dynamic_field"],
+        ),
+    ],
+)
+def test_expr_array_to_columns(spark, expressions: dict[str, str], expected: list[str]):
+    observed = expr_array_to_columns(expressions)
+    assert [cl._jc.toString() for cl in observed] == [expr(expression)._jc.toString() for expression in expected]