datajunction-server 0.0.1a47__tar.gz → 0.0.1a47.dev0__tar.gz

{datajunction_server-0.0.1a47 → datajunction_server-0.0.1a47.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.1
 Name: datajunction-server
-Version: 0.0.1a47
+Version: 0.0.1a47.dev0
 Summary: DataJunction server library for running to a DataJunction server
 Project-URL: repository, https://github.com/DataJunction/dj
 Author-email: DataJunction Authors <yian.shang@gmail.com>

datajunction_server-0.0.1a47.dev0/datajunction_server/__about__.py

@@ -0,0 +1,4 @@
+"""
+Version for Hatch
+"""
+__version__ = "0.0.1a47.dev0"

{datajunction_server-0.0.1a47 → datajunction_server-0.0.1a47.dev0}/datajunction_server/api/nodes.py

@@ -622,6 +622,7 @@ async def create_cube(
     """
     Create a cube node.
     """
+    print("data", data)
     await raise_if_node_exists(session, data.name)
 
     # if the node previously existed and now is inactive

{datajunction_server-0.0.1a47 → datajunction_server-0.0.1a47.dev0}/datajunction_server/internal/nodes.py

@@ -732,6 +732,8 @@ async def update_cube_node(  # pylint: disable=too-many-locals
     """
     Update cube node based on changes
     """
+    node = await Node.get_cube_by_name(session, node_revision.name)
+    node_revision = node.current  # type: ignore
     minor_changes = has_minor_changes(node_revision, data)
     old_metrics = [m.name for m in node_revision.cube_metrics()]
     old_dimensions = node_revision.cube_dimensions()

{datajunction_server-0.0.1a47 → datajunction_server-0.0.1a47.dev0}/datajunction_server/sql/functions.py

@@ -278,28 +278,48 @@ class Aggregate(Function):
             ast,
         )
 
-        expr, start, merge = args
+        if len(args) == 4:
+            expr, start, merge, finish = args
+        else:
+            expr, start, merge = args
+            finish = None
+
         available_identifiers = {
             identifier.name: idx for idx, identifier in enumerate(merge.identifiers)
         }
-        columns = list(
+        merge_columns = list(
             merge.expr.filter(
                 lambda x: isinstance(x, ast.Column)
                 and x.alias_or_name.name in available_identifiers,
             ),
         )
-        for col in columns:
+        finish_columns = (
+            list(
+                finish.expr.filter(
+                    lambda x: isinstance(x, ast.Column)
+                    and x.alias_or_name.name in available_identifiers,
+                ),
+            )
+            if finish
+            else []
+        )
+        for col in merge_columns:
             if available_identifiers.get(col.alias_or_name.name) == 0:
                 col.add_type(start.type)
             if available_identifiers.get(col.alias_or_name.name) == 1:
                 col.add_type(expr.type.element.type)
+        for col in finish_columns:
+            if (
+                available_identifiers.get(col.alias_or_name.name) == 0
+            ):  # pragma: no cover
+                col.add_type(start.type)
 
 
 @Aggregate.register  # type: ignore
 def infer_type(
     expr: ct.ListType,
-    start: ct.PrimitiveType,
-    merge: ct.PrimitiveType,
+    start: ct.ColumnType,
+    merge: ct.ColumnType,
 ) -> ct.ColumnType:
     return merge.expr.type
 
@@ -1028,6 +1048,13 @@ def infer_type(
     return arrays[0].type
 
 
+@Concat.register  # type: ignore
+def infer_type(
+    *maps: ct.MapType,
+) -> ct.MapType:
+    return maps[0].type
+
+
 class ConcatWs(Function):
     """
     concat_ws(separator, [str | array(str)]+) - Returns the concatenation of the
@@ -3130,19 +3157,48 @@ def infer_type(map_: ct.MapType) -> ct.ColumnType:
     return ct.ListType(element_type=map_.type.value.type)
 
 
-# TODO  # pylint: disable=fixme
-# class MapZipWith(Function):
-#     """
-#     map_zip_with(map1, map2, function) - Returns a merged map of two given maps by
-#     applying function to the pair of values with the same key.
-#     """
-#
-#
-# @MapZipWith.register  # type: ignore
-# def infer_type(
-#     map1: ct.MapType, map2: ct.MapType, function: ct.ColumnType
-# ) -> ct.ColumnType:
-#     return ct.MapType()
+class MapZipWith(Function):
+    """
+    map_zip_with(map1, map2, function) - Returns a merged map of two given maps by
+    applying function to the pair of values with the same key.
+    """
+
+    @staticmethod
+    def compile_lambda(*args):
+        """
+        Compiles the lambda function used by the `map_zip_with` Spark function so that
+        the lambda's expression can be evaluated to determine the result's type.
+        """
+        from datajunction_server.sql.parsing import (  # pylint: disable=import-outside-toplevel
+            ast,
+        )
+
+        map1, map2, func = args
+        available_identifiers = {
+            identifier.name: idx for idx, identifier in enumerate(func.identifiers)
+        }
+        columns = list(
+            func.expr.filter(
+                lambda x: isinstance(x, ast.Column)
+                and x.alias_or_name.name in available_identifiers,
+            ),
+        )
+        for col in columns:
+            if available_identifiers.get(col.alias_or_name.name) == 0:
+                col.add_type(map1.type.key.type)  # pragma: no cover
+            if available_identifiers.get(col.alias_or_name.name) == 1:
+                col.add_type(map1.type)
+            if available_identifiers.get(col.alias_or_name.name) == 2:
+                col.add_type(map2.type)
+
+
+@MapZipWith.register  # type: ignore
+def infer_type(
+    map1: ct.MapType,
+    map2: ct.MapType,
+    function: ct.ColumnType,
+) -> ct.ColumnType:
+    return map1.type
 
 
 class Mask(Function):
@@ -4152,6 +4208,92 @@ def infer_type(
     return ct.ListType(element_type=func.expr.type)
 
 
+class TransformKeys(Function):
+    """
+    transform_keys(expr, func) - Transforms keys in a map using the function
+    """
+
+    @staticmethod
+    def compile_lambda(*args):
+        """
+        Compiles the lambda function used by the `transform_keys` Spark function so that
+        the lambda's expression can be evaluated to determine the result's type.
+        """
+        from datajunction_server.sql.parsing import (  # pylint: disable=import-outside-toplevel
+            ast,
+        )
+
+        expr, func = args
+        available_identifiers = {
+            identifier.name: idx for idx, identifier in enumerate(func.identifiers)
+        }
+        columns = list(
+            func.expr.filter(
+                lambda x: isinstance(x, ast.Column)
+                and x.alias_or_name.name in available_identifiers,
+            ),
+        )
+        for col in columns:
+            # The map key arg
+            if available_identifiers.get(col.alias_or_name.name) == 0:
+                col.add_type(expr.type.key.type)
+
+            # The map value arg
+            if available_identifiers.get(col.alias_or_name.name) == 1:
+                col.add_type(expr.type.value.type)
+
+
+@TransformKeys.register  # type: ignore
+def infer_type(
+    expr: ct.MapType,
+    func: ct.ColumnType,
+) -> ct.MapType:
+    return ct.MapType(key_type=func.expr.type, value_type=expr.type.value.type)
+
+
+class TransformValues(Function):
+    """
+    transform_values(expr, func) - Transforms values in a map using the function
+    """
+
+    @staticmethod
+    def compile_lambda(*args):
+        """
+        Compiles the lambda function used by the `transform_values` Spark function so that
+        the lambda's expression can be evaluated to determine the result's type.
+        """
+        from datajunction_server.sql.parsing import (  # pylint: disable=import-outside-toplevel
+            ast,
+        )
+
+        expr, func = args
+        available_identifiers = {
+            identifier.name: idx for idx, identifier in enumerate(func.identifiers)
+        }
+        columns = list(
+            func.expr.filter(
+                lambda x: isinstance(x, ast.Column)
+                and x.alias_or_name.name in available_identifiers,
+            ),
+        )
+        for col in columns:
+            # The map key arg
+            if available_identifiers.get(col.alias_or_name.name) == 0:
+                col.add_type(expr.type.key.type)
+
+            # The map value arg
+            if available_identifiers.get(col.alias_or_name.name) == 1:
+                col.add_type(expr.type.value.type)
+
+
+@TransformValues.register  # type: ignore
+def infer_type(
+    expr: ct.MapType,
+    func: ct.ColumnType,
+) -> ct.MapType:
+    return ct.MapType(key_type=expr.type.key.type, value_type=func.expr.type)
+
+
 class Trim(Function):
     """
     Removes leading and trailing whitespace from a string value.

{datajunction_server-0.0.1a47 → datajunction_server-0.0.1a47.dev0}/datajunction_server/sql/parsing/ast.py

@@ -3,6 +3,7 @@
 # mypy: ignore-errors
 import collections
 import decimal
+import logging
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from dataclasses import dataclass, field, fields
@@ -63,6 +64,7 @@ from datajunction_server.sql.parsing.types import (
 )
 
 PRIMITIVES = {int, float, str, bool, type(None)}
+logger = logging.getLogger(__name__)
 
 
 def flatten(maybe_iterables: Any) -> Iterator:
@@ -1661,6 +1663,7 @@ class Function(Named, Operation):
     args: List[Expression] = field(default_factory=list)
     quantifier: str = ""
     over: Optional[Over] = None
+    args_compiled: bool = False
 
     def __new__(
         cls,
@@ -1724,7 +1727,14 @@ class Function(Named, Operation):
                 await arg.compile(ctx)
                 arg._is_compiled = True
 
-        self.function().compile_lambda(*self.args)
+        # FIXME: We currently catch this exception because we are unable  # pylint: disable=fixme
+        # to infer types for nested lambda functions. For the time being, an easy workaround is to
+        # add a CAST(...) wrapper around the nested lambda function so that the type hard-coded by
+        # the argument to CAST
+        try:
+            self.function().compile_lambda(*self.args)
+        except DJParseException as parse_exc:
+            logger.warning(parse_exc)
 
         for child in self.children:
             if not child.is_compiled():

{datajunction_server-0.0.1a47 → datajunction_server-0.0.1a47.dev0}/tests/sql/functions_test.py

@@ -140,6 +140,69 @@ async def test_aggregate(session: AsyncSession):
     assert query.select.projection[0].type == StringType()  # type: ignore
 
 
+@pytest.mark.asyncio
+async def test_complex_nested_functions(session: AsyncSession):
+    """
+    Test that the CAST(...) workaround works for complex nested lambda functions
+    """
+    query = parse(
+        """
+    SELECT
+      CAST(
+        AGGREGATE(
+          MAP_VALUES(
+            AGGREGATE(
+              COLLECT_LIST(
+                MAP(
+                  x,
+                  NAMED_STRUCT(
+                    'a1',
+                    CASE
+                      WHEN y = 'apples'
+                      THEN z ELSE 0
+                    END,
+                    'b1',
+                    COALESCE(y, 0)
+                  )
+                )
+              ),
+              CAST(MAP() AS MAP<STRING, STRUCT<z: BIGINT, y: BIGINT>>),
+              (acc, x) -> MAP_ZIP_WITH(
+                acc, x,
+                (k, v1, v2) -> NAMED_STRUCT(
+                        'z',
+                        COALESCE(v1['z'], 0) + COALESCE(v2['z'], 0),
+                        'y',
+                        COALESCE(v1['y'], 0) + COALESCE(v2['y'], 0)
+
+                    )
+                ),
+                acc -> TRANSFORM_VALUES(acc, (_, v) -> v['z']/(20.0*v['y']))
+             )
+          ),
+          NAMED_STRUCT(
+              'z1',
+              CAST(0 AS DOUBLE),
+              'y1',
+              CAST(0 AS DOUBLE)
+          ),
+          (acc, x) -> NAMED_STRUCT('y1', acc['z1'] + x, 'y1', acc['y1'] + 1),
+          acc -> acc['y1']/acc['z1']
+      ) AS STRING
+    ) AS etc
+FROM (
+  SELECT '124345' as x,
+  'something' as y,
+  29109 as z
+)
+        """,
+    )
+    exc = DJException()
+    ctx = ast.CompileContext(session=session, exception=exc)
+    await query.compile(ctx)
+    assert query.select.projection[0].type == StringType()  # type: ignore
+
+
 @pytest.mark.asyncio
 async def test_any_value(session: AsyncSession):
     """
@@ -826,7 +889,9 @@ async def test_concat_func(session: AsyncSession):
     Test the `concat` function
     """
     query = parse(
-        "SELECT concat('hello', '+', 'world'), concat(array(1, 2), array(3))",
+        "SELECT concat('hello', '+', 'world'), "
+        "concat(array(1, 2), array(3)), "
+        "concat(map(1, 'a'), map(2, 'b'))",
     )
     exc = DJException()
     ctx = ast.CompileContext(session=session, exception=exc)
@@ -834,6 +899,10 @@ async def test_concat_func(session: AsyncSession):
     assert not exc.errors
     assert query.select.projection[0].type == ct.StringType()  # type: ignore
     assert query.select.projection[1].type == ct.ListType(ct.IntegerType())  # type: ignore
+    assert query.select.projection[2].type == ct.MapType(  # type: ignore
+        key_type=ct.IntegerType(),
+        value_type=ct.StringType(),
+    )
 
 
 @pytest.mark.asyncio
@@ -2634,20 +2703,23 @@ async def test_max() -> None:
     )
 
 
-# TODO
-# @pytest.mark.asyncio
-# async def test_map_zip_with_func(session: AsyncSession):
-#     """
-#     Test the `map_zip_with` function
-#     """
-#     # The third argument to map_zip_with is a function, which needs special handling
-#     # Assuming that we have a function "func" defined elsewhere in the code
-#     query = parse("SELECT map_zip_with(map_col1, map_col2, func) FROM table")
-#     exc = DJException()
-#     ctx = ast.CompileContext(session=session, exception=exc)
-#     await query.compile(ctx)
-#     assert not exc.errors
-#     assert isinstance(query.select.projection[0].type, ct.MapType)  # type: ignore
+@pytest.mark.asyncio
+async def test_map_zip_with_func(session: AsyncSession):
+    """
+    Test the `map_zip_with` function
+    """
+    # The third argument to map_zip_with is a function, which needs special handling
+    # Assuming that we have a function "func" defined elsewhere in the code
+    query = parse(
+        "SELECT map_zip_with(map(1, 'a', 2, 'b'), map(1, 'x', 2, 'y'), (k, v1, v2) -> concat(v1, v2))",
+    )
+    exc = DJException()
+    ctx = ast.CompileContext(session=session, exception=exc)
+    await query.compile(ctx)
+    assert query.select.projection[0].type == ct.MapType(  # type: ignore
+        key_type=ct.IntegerType(),
+        value_type=ct.StringType(),
+    )
 
 
 @pytest.mark.asyncio
@@ -3402,6 +3474,66 @@ async def test_transform(session: AsyncSession):
     assert query.select.projection[0].type == ct.ListType(element_type=ct.IntegerType())  # type: ignore
 
 
+@pytest.mark.asyncio
+async def test_transform_keys(session: AsyncSession):
+    """
+    Test the `transform_keys` Spark function
+    """
+    query = parse(
+        """
+        SELECT transform_keys(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), (k, v) -> k + 1.1);
+        """,
+    )
+    ctx = ast.CompileContext(session=session, exception=DJException())
+    await query.compile(ctx)
+    assert query.select.projection[0].type == ct.MapType(  # type: ignore
+        key_type=ct.FloatType(),
+        value_type=ct.IntegerType(),
+    )
+
+    query = parse(
+        """
+        SELECT transform_keys(map_from_arrays(array('1', '2', '3'), array('1', '2', '3')), (k, v) -> k + v)
+        """,
+    )
+    ctx = ast.CompileContext(session=session, exception=DJException())
+    await query.compile(ctx)
+    assert query.select.projection[0].type == ct.MapType(  # type: ignore
+        key_type=ct.StringType(),
+        value_type=ct.StringType(),
+    )
+
+
+@pytest.mark.asyncio
+async def test_transform_values(session: AsyncSession):
+    """
+    Test the `transform_values` Spark function
+    """
+    query = parse(
+        """
+        SELECT transform_values(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), (k, v) -> v*1.0)
+        """,
+    )
+    ctx = ast.CompileContext(session=session, exception=DJException())
+    await query.compile(ctx)
+    assert query.select.projection[0].type == ct.MapType(  # type: ignore
+        key_type=ct.IntegerType(),
+        value_type=ct.FloatType(),
+    )
+
+    query = parse(
+        """
+        SELECT transform_values(map_from_arrays(array('1', '2', '3'), array('1', '2', '3')), (k, v) -> k + v)
+        """,
+    )
+    ctx = ast.CompileContext(session=session, exception=DJException())
+    await query.compile(ctx)
+    assert query.select.projection[0].type == ct.MapType(  # type: ignore
+        key_type=ct.StringType(),
+        value_type=ct.StringType(),
+    )
+
+
 @pytest.mark.asyncio
 async def test_to_date() -> None:
     """

datajunction_server-0.0.1a47/datajunction_server/__about__.py

@@ -1,4 +0,0 @@
-"""
-Version for Hatch
-"""
-__version__ = "0.0.1a47"