expr: upgrade data_lib dependency to 0.1.21

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]