Skip to content

Commit

Permalink
expr: upgrade data_lib dependency to 0.1.21
Browse files Browse the repository at this point in the history
  • Loading branch information
nikhilgarg28 committed Sep 23, 2024
1 parent 2c9df25 commit bbb4729
Show file tree
Hide file tree
Showing 6 changed files with 78 additions and 66 deletions.
3 changes: 3 additions & 0 deletions fennel/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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

Expand Down
6 changes: 6 additions & 0 deletions fennel/expr/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
49 changes: 25 additions & 24 deletions fennel/expr/test_expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"}
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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)
Expand Down
6 changes: 4 additions & 2 deletions fennel/expr/visitor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)})"

Expand Down
78 changes: 39 additions & 39 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down

0 comments on commit bbb4729

Please sign in to comment.