Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feedback on public api #2735

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
+12 −3
Draft

Feedback on public api #2735

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2fee257
add feedback
sfc-gh-aalam Dec 10, 2024
9155aa1
add feedback
sfc-gh-aalam Dec 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions src/snowflake/snowpark/_internal/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -846,8 +846,9 @@ def infer_ast_enabled_from_global_sessions(func: Callable) -> bool: # pragma: n
f"Could not retrieve default session "
f"for function {func.__qualname__}, capturing AST by default."
)
# session has not been created yet. To not lose information, always encode AST.
return True # noqa: B012
# session has not been created yet, do not encode AST be default.
# TODO: flip this to True when we go GA.
return False # noqa: B012
Copy link
Contributor

@sfc-gh-oplaton sfc-gh-oplaton Dec 11, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we're moving away from hard-coded True, the value should come from an environment variable, or some other suitable global configuration source. We don't want to be in a place where the feature is untestable while under development/opt-in mode.

Copy link
Contributor Author

@sfc-gh-aalam sfc-gh-aalam Dec 11, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why would it not be testable if user has a session and they have oped-in such that session.ast_enabled returns true.

We reach this case when a valid session is not found. Now, default to assuming that ast is enabled but I suggest we default to assuming it is disabled.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Much of the Snowpark public API surface is session-agnostic. Hard-coding this return value to False breaks the following valid code:

  avg_a = call_function("avg", col("a"))
  with Session.builder.config(...).create() as session:
    df = session.create_dataframe([1, 2, 3, 4], schema=["a"])
    df.select(avg_a).show()

else:
return session.ast_enabled # noqa: B012

Expand All @@ -868,6 +869,10 @@ def func_call_wrapper(*args, **kwargs): # pragma: no cover
if "_emit_ast" in func.__code__.co_varnames and "_emit_ast" not in kwargs:
# No arguments, or single argument with function.
if len(args) == 0 or (len(args) == 1 and isinstance(args[0], Callable)):
# this can be refactored to
# - extract session
# - check isinstance(session, snowflake.snowpark.Session)
# - update kwargs["_emit_ast"] = value
Comment on lines +874 to +875
Copy link
Contributor Author

@sfc-gh-aalam sfc-gh-aalam Dec 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed we are not checking isinstance(session, Session) in all branches so this would make it more consistent

if func.__name__ in {
"udf",
"udtf",
Expand Down
6 changes: 5 additions & 1 deletion src/snowflake/snowpark/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -787,13 +787,17 @@ def ast_enabled(self, value: bool) -> None:
# )
# except Exception:
# pass
warn_session_config_update_in_multithreaded_mode(
"ast_enabled", self._conn._thread_safe_session_enabled
)

self._ast_enabled = value

# Auto temp cleaner has bad interactions with AST at the moment, disable when enabling AST.
# This feature should get moved server-side anyways.
if self._ast_enabled:
# TODO SNOW-1770278: Ensure auto temp table cleaner works with AST.
_logger.warning(
"TODO SNOW-1770278: Ensure auto temp table cleaner works with AST."
" Disabling auto temp cleaner for full test suite due to buggy behavior."
)
self.auto_clean_up_temp_table_enabled = False
Expand Down
Loading