SNOW-1904191:[API Coverage] functions coverage (#2964)

snowflakedb · Feb 6, 2025 · 3782249 · 3782249
1 parent 7f6d091
commit 3782249
Show file tree

Hide file tree

Showing 3 changed files with 187 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,6 +21,7 @@
   - `bitmap_bit_position`
   - `bitmap_bucket_number`
   - `bitmap_construct_agg`
+  - `bitshiftright_unsigned`
   - `cbrt`
   - `equal_null`
   - `from_json`

diff --git a/src/snowflake/snowpark/functions.py b/src/snowflake/snowpark/functions.py
@@ -684,6 +684,44 @@ def bitshiftleft(
     return call_builtin("bitshiftleft", c, n, _emit_ast=_emit_ast)
 
 
+@publicapi
+def bitshiftright_unsigned(
+    to_shift_column: ColumnOrName, n: Union[Column, int], _emit_ast: bool = True
+) -> Column:
+    """Returns the bitwise negation of a numeric expression.
+
+    Example:
+        >>> df = session.createDataFrame([(-1999)], ['a'])
+        >>> df.select(bitshiftright_unsigned('a', 1)).collect()[0][0]
+        9223372036854774808
+
+        >>> df = session.createDataFrame([(42)], ['a'])
+        >>> df.select(bitshiftright_unsigned('a', 1)).collect()[0][0]
+        21
+
+        >>> df = session.createDataFrame([(-21)], ['a'])
+        >>> df.select(bitshiftright_unsigned('a', 1)).collect()[0][0]
+        9223372036854775797
+    """
+    # AST.
+    ast = None
+    if _emit_ast:
+        ast = proto.Expr()
+        build_builtin_fn_apply(ast, "bitshiftright_unsigned", to_shift_column, n)
+
+    c = _to_col_if_str(to_shift_column, "bitshiftright_unsigned")
+    max_bit = bitshiftleft(lit(1, _emit_ast=False), 64, _emit_ast=False)
+    unsigned_c = iff(
+        c < 0,
+        bitshiftright(c + max_bit, n, _emit_ast=False),
+        bitshiftright(c, n, _emit_ast=False),
+        _emit_ast=False,
+    )
+    col = call_builtin("bitand", unsigned_c, max_bit - 1, _emit_ast=False)
+    col._ast = ast
+    return col
+
+
 @publicapi
 def bitshiftright(
     to_shift_column: ColumnOrName, n: Union[Column, int], _emit_ast: bool = True

diff --git a/tests/ast/data/functions2.test b/tests/ast/data/functions2.test
@@ -322,6 +322,8 @@ df314 = df.select(instr("A", "test_str"))
 
 df315 = df.select(nth_value("A", 2), nth_value("A", 2, True), nth_value(col("B"), 2, False))
 
+df316 = df.select(bitshiftright_unsigned("A", 2), bitshiftright_unsigned("A", col("B")))
+
 ## EXPECTED UNPARSER OUTPUT
 
 df = session.table("table1")
@@ -644,6 +646,8 @@ df314 = df.select(instr("A", "test_str"))
 
 df315 = df.select(nth_value("A", 2, False), nth_value("A", 2, True), nth_value(col("B"), 2, False))
 
+df316 = df.select(bitshiftright_unsigned("A", 2), bitshiftright_unsigned("A", col("B")))
+
 ## EXPECTED ENCODED AST
 
 interned_value_table {
@@ -26271,6 +26275,150 @@ body {
     }
   }
 }
+body {
+  assign {
+    expr {
+      sp_dataframe_select__columns {
+        cols {
+          apply_expr {
+            fn {
+              builtin_fn {
+                name {
+                  name {
+                    sp_name_flat {
+                      name: "bitshiftright_unsigned"
+                    }
+                  }
+                }
+              }
+            }
+            pos_args {
+              string_val {
+                src {
+                  end_column: 56
+                  end_line: 347
+                  file: 2
+                  start_column: 26
+                  start_line: 347
+                }
+                v: "A"
+              }
+            }
+            pos_args {
+              int64_val {
+                src {
+                  end_column: 56
+                  end_line: 347
+                  file: 2
+                  start_column: 26
+                  start_line: 347
+                }
+                v: 2
+              }
+            }
+            src {
+              end_column: 56
+              end_line: 347
+              file: 2
+              start_column: 26
+              start_line: 347
+            }
+          }
+        }
+        cols {
+          apply_expr {
+            fn {
+              builtin_fn {
+                name {
+                  name {
+                    sp_name_flat {
+                      name: "bitshiftright_unsigned"
+                    }
+                  }
+                }
+              }
+            }
+            pos_args {
+              string_val {
+                src {
+                  end_column: 95
+                  end_line: 347
+                  file: 2
+                  start_column: 58
+                  start_line: 347
+                }
+                v: "A"
+              }
+            }
+            pos_args {
+              apply_expr {
+                fn {
+                  builtin_fn {
+                    name {
+                      name {
+                        sp_name_flat {
+                          name: "col"
+                        }
+                      }
+                    }
+                  }
+                }
+                pos_args {
+                  string_val {
+                    src {
+                      end_column: 94
+                      end_line: 347
+                      file: 2
+                      start_column: 86
+                      start_line: 347
+                    }
+                    v: "B"
+                  }
+                }
+                src {
+                  end_column: 94
+                  end_line: 347
+                  file: 2
+                  start_column: 86
+                  start_line: 347
+                }
+              }
+            }
+            src {
+              end_column: 95
+              end_line: 347
+              file: 2
+              start_column: 58
+              start_line: 347
+            }
+          }
+        }
+        df {
+          sp_dataframe_ref {
+            id {
+              bitfield1: 1
+            }
+          }
+        }
+        src {
+          end_column: 96
+          end_line: 347
+          file: 2
+          start_column: 16
+          start_line: 347
+        }
+        variadic: true
+      }
+    }
+    symbol {
+      value: "df316"
+    }
+    uid: 161
+    var_id {
+      bitfield1: 161
+    }
+  }
+}
 client_ast_version: 1
 client_language {
   python_language {