From 8ccf5e0ff71d0fd1f4d55376cf069ba6b1425fd7 Mon Sep 17 00:00:00 2001
From: Balachandran Atur <bala.atur@snowflake.com>
Date: Tue, 10 Dec 2024 15:33:08 -0800
Subject: [PATCH] Snow-1830500: Bit and unary operations on columns. (#2741)

<!---
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
https://github.com/snowflakedb/snowpark-python/issues/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.
---
 tests/ast/decoder.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

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)