diff --git a/tests/ast/decoder.py b/tests/ast/decoder.py index 6989b9cefb..23b5c2c121 100644 --- a/tests/ast/decoder.py +++ b/tests/ast/decoder.py @@ -447,6 +447,11 @@ def decode_timezone_expr(self, tz_expr: proto.PythonTimeZone) -> Any: def binop(self, ast, fn): return fn(self.decode_expr(ast.lhs), self.decode_expr(ast.rhs)) + def bitop(self, ast, fn): + lhs = self.decode_expr(ast.lhs) + rhs = self.decode_expr(ast.rhs) + return getattr(lhs, fn)(rhs) + def get_statement_params(self, d: Dict): statement_params = {} statement_params_list = d.get("statementParams", []) @@ -732,6 +737,27 @@ def decode_expr(self, expr: proto.Expr) -> Any: # "or" is reserved keyword in python - so have to use getattr here. return self.binop(getattr(expr, "or"), lambda lhs, rhs: lhs | rhs) + # bit operations on columns + case "bit_and": + return self.bitop(expr.bit_and, "bitwiseAnd") + + case "bit_or": + return self.bitop(expr.bit_or, "bitwiseOR") + + case "bit_xor": + return self.bitop(expr.bit_xor, "bitwiseXOR") + + # Unary operations on columns: + case "neg": + col = self.decode_expr(expr.neg.operand) + return -col + + case "not": + # not is a reserved word in python. + col_expr = getattr(expr, "not").operand + col = self.decode_expr(col_expr) + return ~col + # DATAFRAME FUNCTIONS case "sp_create_dataframe": data = self.decode_dataframe_data_expr(expr.sp_create_dataframe.data)