From 3782249573990f8e8d025575a9da57aac405a17c Mon Sep 17 00:00:00 2001 From: Yuyang Wang Date: Thu, 6 Feb 2025 11:23:31 -0800 Subject: [PATCH] SNOW-1904191:[API Coverage] functions coverage (#2964) --- CHANGELOG.md | 1 + src/snowflake/snowpark/functions.py | 38 +++++++ tests/ast/data/functions2.test | 148 ++++++++++++++++++++++++++++ 3 files changed, 187 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 74537740c45..d18f1fb3dc5 100644 --- 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 index d64d21b3439..c14158e431c 100644 --- 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 index 5d8a584e68e..ff801027b67 100644 --- 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 {