Skip to content

Commit

Permalink
Snow-1830500: Bit and unary operations on columns. (#2741)
Browse files Browse the repository at this point in the history
<!---
Please answer these questions before creating your pull request. Thanks!
--->

1. Which Jira issue is this PR addressing? Make sure that there is an
accompanying issue to your PR.

   <!---
   In this section, please add a Snowflake Jira issue number.

Note that if a corresponding GitHub issue exists, you should still
include
   the Snowflake Jira issue number. For example, for GitHub issue
#1400, you should
   add "SNOW-1335071" here.
    --->

   Fixes SNOW-1830500

2. Fill out the following pre-review checklist:

- [ ] I am adding a new automated test(s) to verify correctness of my
new code
- [ ] If this test skips Local Testing mode, I'm requesting review from
@snowflakedb/local-testing
   - [ ] I am adding new logging messages
   - [ ] I am adding a new telemetry message
   - [ ] I am adding new credentials
   - [ ] I am adding a new dependency
- [ ] If this is a new feature/behavior, I'm adding the Local Testing
parity changes.
- [ ] I acknowledge that I have ensured my changes to be thread-safe.
Follow the link for more information: [Thread-safe Developer
Guidelines](https://github.com/snowflakedb/snowpark-python/blob/main/CONTRIBUTING.md#thread-safe-development)

3. Please describe how your code solves the related issue.

Please write a short description of how your code change solves the
related issue.
  • Loading branch information
sfc-gh-batur authored Dec 10, 2024
1 parent 402fbde commit 8ccf5e0
Showing 1 changed file with 26 additions and 0 deletions.
26 changes: 26 additions & 0 deletions tests/ast/decoder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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", [])
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 8ccf5e0

Please sign in to comment.