fdolce commented on code in PR #29146:
URL: https://github.com/apache/flink/pull/29146#discussion_r3979008005
##########
flink-python/pyflink/dataframe/sql.py:
##########
@@ -90,48 +103,84 @@ def sql(query: str, *, auto_bind: bool = True, **bindings:
DataFrame) -> DataFra
... auto_bind=False,
... src=df1,
... )
+ >>> # UDFs are bound the same way, under their variable or keyword name
+ >>> @pf.udf
+ ... def add_one(value: int) -> int:
+ ... return value + 1
+ >>> pf.sql("SELECT add_one(a) AS a1 FROM df1")
+ >>> pf.sql("SELECT inc(a) FROM src", auto_bind=False, src=df1,
inc=add_one)
>>> # Mix SQL and the DataFrame API
>>> pf.sql("SELECT a, b FROM df1").filter(pf.col("a") > 1).to_pandas()
.. versionadded:: 2.4.0
"""
if not isinstance(query, str):
raise TypeError("query must be a string")
- auto_bindings: Dict[str, DataFrame] = {}
+
+ # Collect auto bindings first.
+ variables = {}
+ # Gather the variables in the namespace
if auto_bind:
- frame = inspect.currentframe()
- caller = frame.f_back if frame is not None else None
- try:
- if caller is not None:
- # Locals take precedence over globals.
- namespace = {**caller.f_globals, **caller.f_locals}
- auto_bindings = {
- name: value
- for name, value in namespace.items()
- if isinstance(value, DataFrame)
- }
- finally:
- del frame, caller
- t_env = _resolve_table_environment(bindings, auto_bindings)
- registered: List[str] = []
+ if frame := inspect.currentframe():
+ if outer_frame := frame.f_back:
+ variables = {**outer_frame.f_globals, **outer_frame.f_locals}
+ # Suggested by python docs
+ del outer_frame
+ del frame
+ auto_frames = _get_dataframes(variables)
+ auto_udfs = _get_udfs(variables)
+
+ # Check that explicit bindings are the correct type first
+ for name, value in bindings.items():
+ if not isinstance(value, _BINDABLE_TYPES):
+ raise TypeError(
+ f"sql() binding '{name}' must be a DataFrame or a UDF created
with "
+ f"pyflink.dataframe.udf, got {type(value).__name__}"
+ )
+ explicit_frames = _get_dataframes(bindings)
+ explicit_udfs = _get_udfs(bindings)
+
+ t_env = _resolve_table_environment(explicit_frames, auto_frames)
+ # Each registration step is all-or-nothing: it rolls back its own partial
work on
+ # failure and only returns names on success, so a step that raises leaves
nothing
+ # of its own behind and the finally block only drops what earlier steps
returned.
+ views: List[str] = []
+ functions: List[str] = []
try:
- _register_bindings(t_env, bindings, auto_bindings, registered)
+ views = _register_views(t_env, explicit_frames, auto_frames)
Review Comment:
It's true that the outer `finally` doesn't catch that, but both
`_register_views` and `_register_functions` have their own internal try/catch
to drop partially registered ones. If you prefer, I can move the responsibility
of the drop up here, but as it is it should work correctly, or am I missing
anything else?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
