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

sfc-gh-yuwang · 2025-01-30T19:56:15Z

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

Fixes SNOW-1904191
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
Please describe how your code solves the related issue.

Please write a short description of how your code change solves the related issue.

src/snowflake/snowpark/functions.py

sfc-gh-aalam · 2025-01-31T18:21:41Z

src/snowflake/snowpark/dataframe.py

+            elif isinstance(e, (list, tuple, set)):
+                for sub_e in e:
+                    names.append(sub_e._named())
+                    if _emit_ast and _ast_stmt is None:
+                        ast_cols.append(sub_e._ast)


how does this change affect this new funcs?

json_tuple will have to return a list of columns, while our select function currently does not support this form of call:
df.select(col1, [col2,col3]), this is meant to support it

sfc-gh-aalam · 2025-02-01T00:55:56Z

CHANGELOG.md

+  - `array_slice`
  - `try_to_binary`


did this change get removed?

+ need to update functions.rst?

I found out that we already support it, so I removed it from this PR

sfc-gh-aalam

minor fixes about _emit_ast. Please also take a review from IR team.

sfc-gh-aalam · 2025-02-04T18:23:28Z

src/snowflake/snowpark/functions.py

+    max_bit = bitshiftleft(lit(1), 64)
+    unsigned_c = iff(c < 0, bitshiftright(c + max_bit, n), bitshiftright(c, n))
+    return call_builtin("bitand", unsigned_c, max_bit - 1, _emit_ast=_emit_ast)


Suggested change

max_bit = bitshiftleft(lit(1), 64)

unsigned_c = iff(c < 0, bitshiftright(c + max_bit, n), bitshiftright(c, n))

return call_builtin("bitand", unsigned_c, max_bit - 1, _emit_ast=_emit_ast)

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)

return call_builtin("bitand", unsigned_c, max_bit - 1, _emit_ast=_emit_ast)

let's make sure that for all internal calls to our functions, we set _emit_ast=False

sfc-gh-aalam · 2025-02-04T18:23:44Z

src/snowflake/snowpark/functions.py

+    """
+    c = _to_col_if_str(col, "json_tuple")
+    return [
+        json_extract_path_text(parse_json(c), lit(field)).as_(f"c{i}")


same here, _emit_ast=False

sfc-gh-heshah

Please add expectation tests in the mentioned files, use the suggested helper functions from ast.utils, and see notes about json_tuple.

sfc-gh-heshah · 2025-02-04T22:03:12Z

src/snowflake/snowpark/dataframe.py

@@ -1414,6 +1418,11 @@ def select(
                names.append(e._named())
                if _emit_ast and _ast_stmt is None:
                    ast_cols.append(e._ast)
+            elif isinstance(e, (list, tuple)):


Instead of adding the column ASTs manually for this case, please use build_expr_from_python_val(ast_cols.add(), e) (from ast.utils) to capture when an input provided is a list or tuple.

It would also be very helpful if you could add an expectation test for this in tests/ast/data/select.test

sfc-gh-heshah · 2025-02-04T22:18:36Z

src/snowflake/snowpark/functions.py

+        bitshiftright(c, n, _emit_ast=False),
+        _emit_ast=False,
+    )
+    return call_builtin("bitand", unsigned_c, max_bit - 1, _emit_ast=_emit_ast)


Since the arguments to the public API and the internal call_builtin method are different the AST needs to be constructed manually via build_builtin_fn_apply and assigned to the _ast field of the Column instance returned by call_builtin before returning. See functions.bround for a relevant example.

An expectation test in tests/ast/data/functions2.test would also be very helpful.

sfc-gh-heshah · 2025-02-04T22:32:42Z

src/snowflake/snowpark/functions.py

+            _emit_ast=False,
+        ).as_(f"c{i}", _emit_ast=False)
+        for i, field in enumerate(fields)
+    ]


I don't think this API would emit any AST currently.

Modeling this API with the AST would still be more involved though, as this is the first Snowpark function that returns a List[Column] AFAIK. Usually we rely on the returned Column object storing its own AST in the ._ast field.

Is there any way the return type here can be a Column object? This would allow the AST to be constructed via build_builtin_fn_apply and assigned to the Column instance to be returned.

sfc-gh-heshah · 2025-02-05T18:43:43Z

src/snowflake/snowpark/functions.py

+    ast = None
+    if _emit_ast:
+        ast = proto.Expr()
+        build_builtin_fn_apply(ast, "bround", to_shift_column, n)


Nit: Change builtin_name arg to "bitshiftright_unsigned" instead of "bround"

sfc-gh-heshah · 2025-02-05T22:08:04Z

src/snowflake/snowpark/functions.py

+        bitshiftright(c, n, _emit_ast=False),
+        _emit_ast=False,
+    )
+    col = call_builtin("bitand", unsigned_c, max_bit - 1, _emit_ast=_emit_ast)


One more nit: _emit_ast should be False here as well now

sfc-gh-yuwang added 5 commits January 30, 2025 11:55

add function

22c3b30

add function

1126e20

add bit shift right unsigned

78278c2

Merge branch 'main' into SNOW-1904191

6d8ec25

update changelog

9d42921

sfc-gh-yuwang marked this pull request as ready for review January 31, 2025 18:16

sfc-gh-yuwang requested review from a team as code owners January 31, 2025 18:16

sfc-gh-yuwang requested review from sfc-gh-smirzaei, sfc-gh-jdu and sfc-gh-jrose January 31, 2025 18:16

sfc-gh-aalam reviewed Jan 31, 2025

View reviewed changes

change function name

fd7e143

sfc-gh-yuwang requested a review from sfc-gh-aalam January 31, 2025 21:07

sfc-gh-aalam reviewed Feb 1, 2025

View reviewed changes

sfc-gh-yuwang added 2 commits February 3, 2025 10:47

update changelog

0583ea8

add change log and test for change in dataframe

7e3bc33

sfc-gh-yuwang requested a review from sfc-gh-aalam February 3, 2025 20:53

sfc-gh-jrose approved these changes Feb 3, 2025

View reviewed changes

sfc-gh-yuwang added 2 commits February 3, 2025 14:14

Merge branch 'main' into SNOW-1904191

657e00b

fix docstring test

a38ad96

sfc-gh-aalam approved these changes Feb 4, 2025

View reviewed changes

sfc-gh-yuwang added 3 commits February 4, 2025 13:41

add ast option

90e17b5

add ast option

08b2f7f

add ast option

765c3b6

sfc-gh-heshah requested changes Feb 4, 2025

View reviewed changes

remove json tuple

bb9ceac

sfc-gh-heshah reviewed Feb 5, 2025

View reviewed changes

sfc-gh-yuwang added 2 commits February 5, 2025 10:46

fix nit

ca77239

remove doc string

24c6a2a

sfc-gh-heshah reviewed Feb 5, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

sfc-gh-yuwang commented Jan 30, 2025 •

edited

Loading

sfc-gh-aalam Jan 31, 2025

sfc-gh-yuwang Jan 31, 2025

sfc-gh-aalam Feb 1, 2025

sfc-gh-aalam Feb 1, 2025 •

edited

Loading

sfc-gh-yuwang Feb 3, 2025

sfc-gh-aalam left a comment

sfc-gh-aalam Feb 4, 2025

sfc-gh-aalam Feb 4, 2025

sfc-gh-heshah left a comment

sfc-gh-heshah Feb 4, 2025

sfc-gh-heshah Feb 4, 2025 •

edited

Loading

sfc-gh-heshah Feb 4, 2025

sfc-gh-heshah Feb 5, 2025 •

edited

Loading

sfc-gh-heshah Feb 5, 2025

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

Are you sure you want to change the base?

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

Conversation

sfc-gh-yuwang commented Jan 30, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-aalam Feb 1, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-aalam left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-heshah left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-heshah Feb 4, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-heshah Feb 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-yuwang commented Jan 30, 2025 •

edited

Loading

sfc-gh-aalam Feb 1, 2025 •

edited

Loading

sfc-gh-heshah Feb 4, 2025 •

edited

Loading

sfc-gh-heshah Feb 5, 2025 •

edited

Loading