Nikhil/complex lit

fennel/CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## [1.5.29] - 2024-09-22
+- Enable support for complex literals (e.g. list literals or struct literals)
+
 ## [1.5.28] - 2024-09-21
 - Add FirstK aggregation
 

fennel/expr/expr.py

@@ -619,6 +619,12 @@ def __init__(self, expr: Expr, op: DictOp):
         self.expr = expr
         super(_Dict, self).__init__()
 
+    # Gets the value of the key in the dict
+    # When default is None and the key is not present, the result is None (and
+    # consequently type of the `get` expression is Optional[T]). However, when
+    # the default is not None and the key is not present, the result is the
+    # default value (and consequently type of the `get` expression is the common
+    # type of the default value and the value of the key).
     def get(self, key: str, default: Optional[Expr] = None) -> Expr:
         key = make_expr(key)
         default = make_expr(default) if default is not None else None  # type: ignore

fennel/expr/test_expr.py

@@ -321,9 +321,7 @@ def test_dict_op():
         "a"
     ).dict.len()
     printer = ExprPrinter()
-    expected = (
-        """(CEIL((col("a").get("x") + col("a").get("y"))) + LEN(col("a")))"""
-    )
+    expected = """(CEIL((col("a").dict.get("x") + col("a").dict.get("y"))) + LEN(col("a")))"""
     ref_extractor = FetchReferences()
     ref_extractor.visit(expr.root)
     assert ref_extractor.refs == {"a"}
@@ -378,6 +376,10 @@ def test_dict_op():
     }
     expected_expr = ParseDict(d, Expr())
     assert expected_expr == proto_expr, error_message(proto_expr, expected_expr)
+    assert expr.typeof({"a": Dict[str, int]}) == Optional[int]
+
+    # but types change when default is not None
+    expr = col("a").dict.get("x", 10)
     assert expr.typeof({"a": Dict[str, int]}) == int
 
 
@@ -1268,29 +1270,28 @@ def test_fillnull():
             expected_dtype=datetime,
             proto_json=None,
         ),
-        # TODO(Nikhil): Add support for filling nulls for complex types
         # Fill null for struct
-        # ExprTestCase(
-        #     expr=(col("a").fillnull(lit(A(1, 2, "a"), A))),
-        #     df=pd.DataFrame({"a": [A(1, 2, "a"), None, A(3, 4, "b")]}),
-        #     schema={"a": Optional[A]},
-        #     display="""FILL_NULL(col('a'), {"x": 1, "y": 2, "z": "a"})""",
-        #     refs={"a"},
-        #     eval_result=[A(1, 2, "a"), A(1, 2, "a"), A(3, 4, "b")],
-        #     expected_dtype=A,
-        #     proto_json=None,
-        # ),
+        ExprTestCase(
+            expr=(col("a").fillnull(lit(A(1, 2, "a"), A))),
+            df=pd.DataFrame({"a": [A(1, 2, "a"), None, A(3, 4, "b")]}),
+            schema={"a": Optional[A]},
+            display="""FILL_NULL(col("a"), {"x": 1, "y": 2, "z": "a"})""",
+            refs={"a"},
+            eval_result=[A(1, 2, "a"), A(1, 2, "a"), A(3, 4, "b")],
+            expected_dtype=A,
+            proto_json=None,
+        ),
         # Fill null for Optional[List[int]]
-        # ExprTestCase(
-        #     expr=(col("a").fillnull(lit([1, 2, 3], List[int]))),
-        #     df=pd.DataFrame({"a": [[1, 2, 3], None, [4, 5, 6]]}),
-        #     schema={"a": Optional[List[int]]},
-        #     display="FILL_NULL(col('a'), [1, 2, 3])",
-        #     refs={"a"},
-        #     eval_result=[[1, 2, 3], [1, 2, 3], [4, 5, 6]],
-        #     expected_dtype=List[int],
-        #     proto_json=None,
-        # ),
+        ExprTestCase(
+            expr=(col("a").fillnull(lit([1, 2, 3], List[int]))),
+            df=pd.DataFrame({"a": [[1, 2, 3], None, [4, 5, 6]]}),
+            schema={"a": Optional[List[int]]},
+            display='FILL_NULL(col("a"), [1, 2, 3])',
+            refs={"a"},
+            eval_result=[[1, 2, 3], [1, 2, 3], [4, 5, 6]],
+            expected_dtype=List[int],
+            proto_json=None,
+        ),
     ]
     for case in cases:
         check_test_case(case)

fennel/expr/visitor.py

@@ -305,9 +305,11 @@ def visitDict(self, obj):
             )
         elif isinstance(obj.op, DictGet):
             if obj.op.default is None:
-                return f"{self.visit(obj.expr)}.get({self.visit(obj.op.key)})"
+                return (
+                    f"{self.visit(obj.expr)}.dict.get({self.visit(obj.op.key)})"
+                )
             else:
-                return f"{self.visit(obj.expr)}.get('{self.visit(obj.op.key)}', {self.visit(obj.op.default)})"
+                return f"{self.visit(obj.expr)}.dict.get('{self.visit(obj.op.key)}', {self.visit(obj.op.default)})"
         elif isinstance(obj.op, DictLen):
             return f"LEN({self.visit(obj.expr)})"
 

pyproject.toml

@@ -20,7 +20,7 @@ pytest = "7.1.3"
 pytest-rerunfailures = "^13.0"
 sortedcontainers = "^2.4.0"
 typing-extensions = "^4.12.0"
-fennel-data-lib = "0.1.20"
+fennel-data-lib = "0.1.21"
 pyarrow = "^14.0.2"
 
 [tool.poetry.dev-dependencies]