dianfu commented on code in PR #29146:
URL: https://github.com/apache/flink/pull/29146#discussion_r3979327529
##########
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:
The internal cleanup only covers the explicit-binding loops. The
auto-binding loops are outside those `try/except` blocks.
The same issue exists in `_register_views()`. Moving the registration
tracker to the outer scope, or wrapping both explicit and auto registration in
the helper cleanup block, would cover this path.
--
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]
