Skip to content

Commit

Permalink
SNOW-1857250 Fix AST generation in functions.py (#2765)
Browse files Browse the repository at this point in the history
1. Which Jira issue is this PR addressing? Make sure that there is an
accompanying issue to your PR.

   Fixes SNOW-1857250

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.
- [x] 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.

Fixed the AST generation in multiple functions in `functions.py` by
encoding them outside of the `builtin` method. Also added
`_emit_ast=False` in several places to prevent recording public API used
internally when generating AST.
  • Loading branch information
sfc-gh-vbudati authored Dec 19, 2024
1 parent eacec5c commit a0ac492
Show file tree
Hide file tree
Showing 4 changed files with 482 additions and 1,796 deletions.
28 changes: 27 additions & 1 deletion src/snowflake/snowpark/_internal/ast/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -380,6 +380,33 @@ def build_sp_table_name( # type: ignore[no-untyped-def] # TODO(SNOW-1491199) #
raise ValueError(f"Invalid name type {type(name)} for SpTableName entity.")


def build_function_expr(
builtin_name: str,
args: List[Any],
ignore_null_args: bool = False,
) -> proto.Expr:
"""
Creates AST encoding for the methods in function.py.
Args:
builtin_name: Name of the builtin function to call.
args: Positional arguments to pass to function, in the form of a list.
ignore_null_args: If True, null arguments will be ignored.
Returns:
The AST encoding of the function.
"""
ast = proto.Expr()
args_list = [arg for arg in args if arg is not None] if ignore_null_args else args
build_builtin_fn_apply(
ast,
builtin_name,
*tuple(
snowpark_expression_to_ast(arg) if isinstance(arg, Expression) else arg
for arg in args_list
),
)
return ast


# TODO(SNOW-1491199) - This method is not covered by tests until the end of phase 0. Drop the pragma when it is covered.
def build_builtin_fn_apply(
ast: proto.Expr,
Expand All @@ -395,7 +422,6 @@ def build_builtin_fn_apply(
builtin_name: Name of the builtin function to call.
*args: Positional arguments to pass to function.
**kwargs: Keyword arguments to pass to function.
"""
expr = with_src_position(ast.apply_expr) # type: ignore[arg-type] # TODO(SNOW-1491199) # Argument 1 to "with_src_position" has incompatible type "ApplyExpr"; expected "Expr"
_set_fn_name(builtin_name, expr.fn.builtin_fn) # type: ignore[attr-defined] # TODO(SNOW-1491199) # "Expr" has no attribute "fn"
Expand Down
Loading

0 comments on commit a0ac492

Please sign in to comment.