This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new a9d5b293f6a [FLINK-40197][python] Support UDF bindings in
DataFrame.sql (#29146)
a9d5b293f6a is described below
commit a9d5b293f6ab152bd80d46433a7b934b07620149
Author: Federico Dolce <[email protected]>
AuthorDate: Fri Sep 11 17:01:23 2026 +0200
[FLINK-40197][python] Support UDF bindings in DataFrame.sql (#29146)
---
.../docs/reference/pyflink.dataframe/sql.rst | 7 +-
flink-python/pyflink/dataframe/sql.py | 343 +++++++++++++------
flink-python/pyflink/dataframe/tests/test_sql.py | 364 ++++++++++++++++++++-
3 files changed, 603 insertions(+), 111 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/sql.rst
b/flink-python/docs/reference/pyflink.dataframe/sql.rst
index fdfb26ab796..b56e784031e 100644
--- a/flink-python/docs/reference/pyflink.dataframe/sql.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/sql.rst
@@ -20,7 +20,7 @@
SQL
===
-Execute SQL SELECT queries against DataFrames.
+Execute SQL SELECT queries against DataFrames and DataFrame UDFs.
Example::
@@ -33,6 +33,11 @@ Example::
... auto_bind=False,
... src=df1,
... )
+ >>> @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)
>>> pf.sql("SELECT a, b FROM df1").filter(pf.col("a") > 1).to_pandas()
.. currentmodule:: pyflink.dataframe
diff --git a/flink-python/pyflink/dataframe/sql.py
b/flink-python/pyflink/dataframe/sql.py
index f244a81897b..879a060b29c 100644
--- a/flink-python/pyflink/dataframe/sql.py
+++ b/flink-python/pyflink/dataframe/sql.py
@@ -19,13 +19,16 @@
import inspect
import logging
import warnings
-from typing import Any, Dict, List
+from typing import Any, Callable, Dict, Iterable, List, Set, Union
from py4j.protocol import Py4JJavaError
from pyflink.dataframe.context import get_or_create_table_environment
from pyflink.dataframe.dataframe import DataFrame
+from pyflink.dataframe.udf import _DataFrameUDFWrapper
+from pyflink.java_gateway import get_gateway
from pyflink.table import Table, TableEnvironment
+from pyflink.table.expression import Expression
from pyflink.util.api_stability_decorators import PublicEvolving
from pyflink.util.java_utils import is_instance_of
@@ -33,28 +36,36 @@ __all__ = ["sql"]
_LOG = logging.getLogger(__name__)
+# UDFs are annotated with the declared return type of
:func:`pyflink.dataframe.udf`
+# so that its result type-checks as a binding; at runtime a binding must be an
actual
+# UDF object, which is what _BINDABLE_TYPES enforces.
+_Binding = Union[DataFrame, Callable[..., Expression]]
+_BINDABLE_TYPES = (DataFrame, _DataFrameUDFWrapper)
+
@PublicEvolving()
-def sql(query: str, *, auto_bind: bool = True, **bindings: DataFrame) ->
DataFrame:
+def sql(query: str, *, auto_bind: bool = True, **bindings: _Binding) ->
DataFrame:
"""
Execute a SQL query and return the result as a :class:`DataFrame`.
The query must be a single statement that returns a result, such as SELECT
or
VALUES (no INSERT / DDL; use :meth:`TableEnvironment.execute_sql` for
those).
- The referenced DataFrames are registered as temporary views for the
duration of
- the call and dropped afterwards. The result can be further transformed
with the
- DataFrame API.
+ The referenced DataFrames are registered as temporary views and the
referenced
+ UDFs (created with :func:`~pyflink.dataframe.udf`) as temporary system
functions
+ for the duration of the call, and both are dropped afterwards. The result
can be
+ further transformed with the DataFrame API.
When ``auto_bind`` is ``True`` (the default), the caller's local and
global variables
- are scanned for :class:`DataFrame` objects and each is registered under
its Python
- variable name. Auto-binding is best-effort: it warns and skips names that
are not
- valid SQL identifiers or that collide with an existing table or view, and
it never
- shadows permanent catalog objects.
+ are scanned for :class:`DataFrame` and UDF objects and each is registered
under its
+ Python variable name. Auto-binding is best-effort: it warns and skips
names that
+ are not valid SQL identifiers or that collide with an existing table, view
or
+ function, and it never shadows permanent catalog objects or built-in
functions.
Explicit keyword ``bindings`` define the SQL names directly. They are
strict
- (invalid names and conflicts with existing temporary views raise
- :class:`ValueError`), take precedence over auto-bind on name collisions,
and are
- required to intentionally shadow a permanent catalog table or view.
+ (invalid names and conflicts with existing temporary views or temporary
functions
+ raise :class:`ValueError`), take precedence over auto-bind on name
collisions, and
+ are required to intentionally shadow a permanent catalog table, view or
function
+ or a built-in function.
The query runs in the :class:`TableEnvironment` of the bound DataFrames:
the
environment shared by the explicit ``bindings`` when given, otherwise the
@@ -62,20 +73,25 @@ def sql(query: str, *, auto_bind: bool = True, **bindings:
DataFrame) -> DataFra
different environments raise :class:`ValueError`. Without explicit
bindings,
valid auto-bound candidates from different environments also raise
:class:`ValueError`; when explicit bindings determine the environment,
auto-bound
- candidates from other environments are skipped with a warning. The resolved
+ candidates from other environments are skipped with a warning. UDF
bindings are
+ not tied to an environment and do not take part in this resolution. The
resolved
environment is used only for this call and never replaces the global one.
+ Function names are case-insensitive in SQL, so a UDF bound as ``addOne`` is
+ referenced as ``addOne(...)`` or ``addone(...)`` alike.
+
:param query: The query to execute.
- :param auto_bind: Whether to scan the caller's variables for DataFrames.
- :param bindings: Explicit name to :class:`DataFrame` bindings.
+ :param auto_bind: Whether to scan the caller's variables for DataFrames
and UDFs.
+ :param bindings: Explicit name to :class:`DataFrame` or UDF bindings.
:return: The query result.
:raises ValueError: If the query is not a query statement, if an explicit
binding
is not a valid SQL identifier or conflicts with an
existing
- temporary view, or if explicit bindings belong to
different
- TableEnvironments, or if auto-bound candidates belong
to
- different TableEnvironments when there are no explicit
- bindings.
- :raises TypeError: If an explicit binding is not a :class:`DataFrame`.
+ temporary view or temporary function, or if explicit
+ bindings belong to different TableEnvironments, or if
+ auto-bound candidates belong to different
TableEnvironments
+ when there are no explicit bindings.
+ :raises TypeError: If an explicit binding is neither a :class:`DataFrame`
nor a
+ UDF created with :func:`~pyflink.dataframe.udf`.
Example::
@@ -90,6 +106,12 @@ 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()
@@ -97,41 +119,71 @@ def sql(query: str, *, auto_bind: bool = True, **bindings:
DataFrame) -> DataFra
"""
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)
+ functions = _register_functions(t_env, explicit_udfs, auto_udfs)
return DataFrame(_execute_query(t_env, query))
finally:
- for name in registered:
- # Best-effort cleanup: dropping must not mask an exception raised
by the
- # query itself, but only names registered by this call are
dropped, so a
- # failure is an anomaly the user should hear about.
- try:
- t_env.drop_temporary_view(name)
- except Exception:
- _LOG.warning(
- "sql() failed to drop temporary view '%s'", name,
exc_info=True
- )
+ _drop_functions(t_env, functions)
+ _drop_views(t_env, views)
+
+
+def _get_dataframes(namespace: Dict[str, Any]) -> Dict[str, DataFrame]:
+ return {k: v for k, v in namespace.items() if isinstance(v, DataFrame)}
+
+
+def _get_udfs(namespace: Dict[str, Any]) -> Dict[str, _DataFrameUDFWrapper]:
+ return {k: v for k, v in namespace.items() if isinstance(v,
_DataFrameUDFWrapper)}
+
+
+def _drop_views(t_env: TableEnvironment, names: List[str]) -> None:
+ for name in names:
+ try:
+ t_env.drop_temporary_view(name)
+ except Exception:
+ _LOG.warning("sql() failed to drop view '%s'", name, exc_info=True)
+
+
+def _drop_functions(t_env: TableEnvironment, names: List[str]) -> None:
+ for name in names:
+ try:
+ t_env.drop_temporary_system_function(name)
+ except Exception:
+ _LOG.warning("sql() failed to drop function '%s'", name,
exc_info=True)
def _resolve_table_environment(
- explicit: Dict[str, Any], auto: Dict[str, DataFrame]
+ explicit: Dict[str, DataFrame], auto: Dict[str, DataFrame]
) -> TableEnvironment:
"""
Pick the environment to run the query in: the environment shared by the
explicit
@@ -141,36 +193,36 @@ def _resolve_table_environment(
environment resolution. Mixed valid auto-bound candidates are also
ambiguous and
require explicit bindings, regardless of the configured global environment.
"""
- for name, value in explicit.items():
- if not isinstance(value, DataFrame):
- raise TypeError(
- f"sql() binding '{name}' must be a DataFrame, got
{type(value).__name__}"
- )
- # Deduplicate by identity: environments are not comparable by value.
- explicit_envs = {id(v._table._t_env): v._table._t_env for v in
explicit.values()}
+ explicit_envs = _distinct_environments(explicit.values())
if len(explicit_envs) > 1:
raise ValueError(
"sql() explicit bindings belong to different TableEnvironments; "
"bind DataFrames from a single environment"
)
if explicit_envs:
- return next(iter(explicit_envs.values()))
- auto_envs = {
- id(value._table._t_env): value._table._t_env
+ return explicit_envs[0]
+
+ auto_envs = _distinct_environments(
+ value
for name, value in auto.items()
if _is_simple_sql_identifier(value._table._t_env, name)
- }
- if len(auto_envs) == 1:
- return next(iter(auto_envs.values()))
+ )
if len(auto_envs) > 1:
raise ValueError(
"sql() auto-bound DataFrames belong to different
TableEnvironments; "
"set auto_bind=False and pass explicit bindings from a single "
"TableEnvironment"
)
+ if auto_envs:
+ return auto_envs[0]
return get_or_create_table_environment()
+def _distinct_environments(frames: Iterable[DataFrame]) ->
List[TableEnvironment]:
+ # Deduplicate by identity: environments are not comparable by value.
+ return list({id(f._table._t_env): f._table._t_env for f in
frames}.values())
+
+
def _execute_query(t_env: TableEnvironment, query: str) -> Table:
"""
Run ``query`` through :meth:`TableEnvironment.sql_query`, which parses the
statement
@@ -192,8 +244,8 @@ def _execute_query(t_env: TableEnvironment, query: str) ->
Table:
def _is_simple_sql_identifier(t_env: TableEnvironment, name: str) -> bool:
"""
Whether ``name`` is accepted verbatim as a single-part identifier by the
SQL parser,
- i.e. whether registering a temporary view under it can succeed. This is
the same
- validation :meth:`TableEnvironment.create_temporary_view` applies to its
path, so
+ i.e. whether registering a temporary view or function under it can
succeed. This is
+ the same validation :meth:`TableEnvironment.create_temporary_view`
applies, so
keywords like ``order`` pass (queries reference them with backticks) while
names
that would need quoting or resolve to a different or multi-part path do
not.
"""
@@ -212,51 +264,146 @@ def _is_simple_sql_identifier(t_env: TableEnvironment,
name: str) -> bool:
)
-def _register_bindings(
+def _register_views(
t_env: TableEnvironment,
explicit: Dict[str, DataFrame],
auto: Dict[str, DataFrame],
- registered: List[str],
-) -> None:
+) -> List[str]:
"""
- Register explicit and auto-collected bindings as temporary views,
appending each
- successful registration to ``registered``. The explicit bindings have
already been
- type-checked and share ``t_env`` (see :func:`_resolve_table_environment`).
+ Register explicit and auto-collected DataFrames as temporary views and
return the
+ registered names. Registration is all-or-nothing: if anything raises, the
views
+ registered so far are dropped again before re-raising. The explicit
bindings have
+ already been type-checked and share ``t_env`` (see
+ :func:`_resolve_table_environment`).
"""
temporary_tables = set(t_env.list_temporary_tables())
# list_tables() covers both permanent and temporary tables and views.
all_tables = set(t_env.list_tables())
+ registered: List[str] = []
- for name, value in explicit.items():
- if not _is_simple_sql_identifier(t_env, name):
- raise ValueError(f"cannot bind '{name}': it is not a valid SQL
identifier")
- if name in temporary_tables:
- raise ValueError(
- f"cannot bind '{name}': a temporary table or view with this
name "
- "already exists"
- )
- t_env.create_temporary_view(name, value.to_table())
- registered.append(name)
-
- for name, value in auto.items():
- if name in explicit:
- # Explicit bindings take precedence on name collisions.
- continue
- if not _is_simple_sql_identifier(t_env, name):
- _warn_skipped(name, "it is not a valid SQL identifier")
- continue
- if value._table._t_env is not t_env:
- _warn_skipped(name, "it belongs to a different TableEnvironment")
- continue
- if name in all_tables:
- _warn_skipped(name, "a table or view with this name already
exists")
- continue
- try:
+ try:
+ for name, value in explicit.items():
+ if not _is_simple_sql_identifier(t_env, name):
+ raise ValueError(
+ f"cannot bind '{name}': it is not a valid SQL identifier"
+ )
+ if name in temporary_tables:
+ raise ValueError(
+ f"cannot bind '{name}': a temporary table or view with
this name "
+ "already exists"
+ )
t_env.create_temporary_view(name, value.to_table())
- except Exception as e:
- _warn_skipped(name, f"registration failed: {e}")
- continue
- registered.append(name)
+ registered.append(name)
+
+ # Auto-bound candidates are not expected to raise: problems are
reported as
+ # warnings. Anything unexpected still rolls back the registered views.
+ for name, value in auto.items():
+ if name in explicit:
+ # Explicit bindings take precedence on name collisions.
+ continue
+ if not _is_simple_sql_identifier(t_env, name):
+ _warn_skipped(name, "it is not a valid SQL identifier")
+ continue
+ if value._table._t_env is not t_env:
+ _warn_skipped(name, "it belongs to a different
TableEnvironment")
+ continue
+ if name in all_tables:
+ _warn_skipped(name, "a table or view with this name already
exists")
+ continue
+ try:
+ t_env.create_temporary_view(name, value.to_table())
+ except Exception as e:
+ _warn_skipped(name, f"registration failed: {e}")
+ continue
+ registered.append(name)
+ except Exception:
+ _drop_views(t_env, registered)
+ raise
+ return registered
+
+
+def _register_functions(
+ t_env: TableEnvironment,
+ explicit: Dict[str, _DataFrameUDFWrapper],
+ auto: Dict[str, _DataFrameUDFWrapper],
+) -> List[str]:
+ """
+ Register explicit and auto-collected UDFs as temporary system functions
and return
+ the registered names. Registration is all-or-nothing: if anything raises,
the
+ functions registered so far are dropped again before re-raising.
+
+ System functions are looked up by bare name independently of the current
catalog
+ and database, which matches how a Python name is referenced in the query.
Function
+ names are case-insensitive: the catalog normalizes them to lower case, so
+ collisions are checked case-insensitively. Mirroring views, explicit
bindings may
+ shadow built-in and permanent catalog functions but never an existing
temporary
+ function; auto-bind shadows nothing.
+ """
+ if not explicit and not auto:
+ return []
+ # Explicit bindings take precedence on name collisions, so only the
remaining
+ # auto-bound candidates need the built-in function names. list_functions()
covers
+ # those as well, but is comparatively expensive, so skip it when nothing
needs it.
+ explicit_names = {name.lower() for name in explicit}
+ auto_candidates = {
+ name: value for name, value in auto.items() if name.lower() not in
explicit_names
+ }
+ all_functions: Set[str] = set()
+ if auto_candidates:
+ all_functions = {f.lower() for f in t_env.list_functions()}
+ registered: List[str] = []
+
+ try:
+ for name, value in explicit.items():
+ if not _is_simple_sql_identifier(t_env, name):
+ raise ValueError(
+ f"cannot bind '{name}': it is not a valid SQL identifier"
+ )
+ if _has_temporary_function(t_env, name):
+ raise ValueError(
+ f"cannot bind '{name}': a temporary function with this
name "
+ "already exists"
+ )
+ t_env.create_temporary_system_function(name,
value._table_udf_wrapper)
+ registered.append(name)
+
+ # Auto-bound candidates are not expected to raise: problems are
reported as
+ # warnings. Anything unexpected still rolls back the registered
functions.
+ for name, value in auto_candidates.items():
+ if not _is_simple_sql_identifier(t_env, name):
+ _warn_skipped(name, "it is not a valid SQL identifier")
+ continue
+ if name.lower() in all_functions or _has_temporary_function(t_env,
name):
+ _warn_skipped(name, "a function with this name already exists")
+ continue
+ try:
+ t_env.create_temporary_system_function(name,
value._table_udf_wrapper)
+ except Exception as e:
+ _warn_skipped(name, f"registration failed: {e}")
+ continue
+ registered.append(name)
+ except Exception:
+ _drop_functions(t_env, registered)
+ raise
+ return registered
+
+
+def _has_temporary_function(t_env: TableEnvironment, name: str) -> bool:
+ """
+ Whether a temporary system function or a temporary catalog function in the
current
+ catalog and database is registered under ``name``. The environment's
function
+ listings merge temporary and permanent functions, so a temporary function
that
+ shares its name with a permanent one cannot be told apart from them; the
function
+ catalog keeps them separate.
+ """
+ function_catalog =
t_env._j_tenv.getPlanner().getFlinkContext().getFunctionCatalog()
+ if function_catalog.hasTemporarySystemFunction(name.lower()):
+ return True
+ gateway = get_gateway()
+ identifier =
gateway.jvm.org.apache.flink.table.catalog.ObjectIdentifier.of(
+ t_env.get_current_catalog(), t_env.get_current_database(), name
+ )
+ return function_catalog.hasTemporaryCatalogFunction(identifier)
def _warn_skipped(name: str, reason: str) -> None:
diff --git a/flink-python/pyflink/dataframe/tests/test_sql.py
b/flink-python/pyflink/dataframe/tests/test_sql.py
index bed0eecae5b..bcc7b1190b7 100644
--- a/flink-python/pyflink/dataframe/tests/test_sql.py
+++ b/flink-python/pyflink/dataframe/tests/test_sql.py
@@ -17,16 +17,27 @@
################################################################################
import unittest
+import warnings
from py4j.protocol import Py4JJavaError
import pyflink.dataframe as pf
from pyflink.common import Row
from pyflink.table import DataTypes, EnvironmentSettings, TableEnvironment
-from pyflink.table.udf import udf
+from pyflink.table.udf import udf as table_udf
from pyflink.testing.test_case_utils import PyFlinkDataFrameUTTestCase
[email protected]
+def module_level_add_one(value: int) -> int:
+ return value + 1
+
+
+def _collect_sorted(df: pf.DataFrame) -> list:
+ """Collect rows in a deterministic order: query results are unordered."""
+ return sorted(df.collect(), key=lambda row: row[0])
+
+
class SqlValidationTests(unittest.TestCase):
def setUp(self):
previous_environment = pf.get_table_environment()
@@ -78,7 +89,7 @@ class SqlTests(PyFlinkDataFrameUTTestCase):
)
self.assertEqual(
- sorted(joined.collect(), key=lambda row: row[0]),
+ _collect_sorted(joined),
[Row(1, "x", "p"), Row(2, "y", "q"), Row(3, "z", "r")],
)
@@ -94,7 +105,7 @@ class SqlTests(PyFlinkDataFrameUTTestCase):
]:
with self.subTest(query=query):
self.assertEqual(
- sorted(pf.sql(query).collect(), key=lambda row: row[0]),
+ _collect_sorted(pf.sql(query)),
[Row(1), Row(2)],
)
@@ -109,7 +120,7 @@ class SqlTests(PyFlinkDataFrameUTTestCase):
)
self.assertEqual(
- sorted(result.collect(), key=lambda row: row[0]),
+ _collect_sorted(result),
[Row(10), Row(20)],
)
@@ -282,21 +293,350 @@ class SqlTests(PyFlinkDataFrameUTTestCase):
with self.assertRaisesRegex(TypeError, "'x' must be a DataFrame"):
pf.sql("SELECT * FROM x", auto_bind=False, x=table)
- def test_udfs_are_not_bindable(self):
- # UDF support will come in a separate change once the DataFrame API
grows
- # UDF support in general: sql() must reject them rather than
half-support them.
- add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT())
+ def test_table_udfs_are_not_bindable(self):
+ # Only DataFrame UDFs (pf.udf) are supported, mirroring the rejection
of raw
+ # Tables in favour of DataFrames.
+ add_one = table_udf(lambda i: i + 1, result_type=DataTypes.BIGINT())
- with self.assertRaisesRegex(TypeError, "'add_one' must be a
DataFrame"):
+ with self.assertRaisesRegex(
+ TypeError, "'add_one' must be a DataFrame or a UDF created with"
+ ):
pf.sql("SELECT add_one(a) FROM df", auto_bind=False,
add_one=add_one)
- def test_auto_bind_ignores_udfs(self):
+ def test_auto_bind_ignores_table_udfs(self):
df = pf.from_dict({"a": [1]}) # noqa: F841
- add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT()) #
noqa: F841
+ add_one = table_udf( # noqa: F841
+ lambda i: i + 1, result_type=DataTypes.BIGINT()
+ )
with self.assertRaisesRegex(Py4JJavaError, "No match found for
function"):
pf.sql("SELECT add_one(a) FROM df")
+ def test_auto_bind_registers_udfs_by_variable_name(self):
+ df = pf.from_dict({"a": [1, 2]}) # noqa: F841
+
+ @pf.udf
+ def add_one(value: int) -> int:
+ return value + 1
+
+ result = pf.sql("SELECT add_one(a) AS b FROM df")
+
+ self.assertEqual(
+ _collect_sorted(result), [Row(2), Row(3)]
+ )
+ self.assertNotIn("add_one", self.t_env.list_user_defined_functions())
+
+ def test_auto_bind_finds_module_level_udfs(self):
+ df = pf.from_dict({"a": [1, 2]}) # noqa: F841
+
+ result = pf.sql("SELECT module_level_add_one(a) AS b FROM df")
+
+ self.assertEqual(
+ _collect_sorted(result), [Row(2), Row(3)]
+ )
+
+ def test_explicit_udf_binding_picks_the_sql_name(self):
+ src = pf.from_dict({"a": [1, 2]})
+
+ @pf.udf
+ def add_one(value: int) -> int:
+ return value + 1
+
+ result = pf.sql(
+ "SELECT inc(a) AS b FROM src", auto_bind=False, src=src,
inc=add_one
+ )
+
+ self.assertEqual(
+ _collect_sorted(result), [Row(2), Row(3)]
+ )
+ self.assertNotIn("inc", self.t_env.list_user_defined_functions())
+
+ def test_auto_bind_disabled_ignores_caller_udfs(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+ add_one = pf.udf(lambda value: value + 1, return_dtype=int) # noqa:
F841
+
+ with self.assertRaisesRegex(Py4JJavaError, "No match found for
function"):
+ pf.sql("SELECT add_one(a) FROM df", auto_bind=False, df=df)
+
+ def test_explicit_udf_binding_takes_precedence_over_auto_bind(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+ add_one = pf.udf(lambda value: value + 1, return_dtype=int) # noqa:
F841
+ add_ten = pf.udf(lambda value: value + 10, return_dtype=int)
+
+ result = pf.sql("SELECT add_one(a) FROM df", add_one=add_ten)
+
+ self.assertEqual(result.collect(), [Row(11)])
+
+ def test_udf_names_are_case_insensitive(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+ addOne = pf.udf(lambda value: value + 1, return_dtype=int) # noqa:
F841
+
+ for query in ["SELECT addOne(a) FROM df", "SELECT ADDONE(a) FROM df"]:
+ with self.subTest(query=query):
+ self.assertEqual(pf.sql(query).collect(), [Row(2)])
+ self.assertNotIn("addone", self.t_env.list_user_defined_functions())
+
+ def test_pandas_udfs_are_bindable(self):
+ df = pf.from_dict({"a": [1, 2]}) # noqa: F841
+ add_one = pf.udf( # noqa: F841
+ lambda values: values + 1, return_dtype=int, func_type="pandas"
+ )
+
+ result = pf.sql("SELECT add_one(a) AS b FROM df")
+
+ self.assertEqual(
+ _collect_sorted(result), [Row(2), Row(3)]
+ )
+
+ def test_udf_bindings_compose_with_the_dataframe_api(self):
+ df = pf.from_dict({"a": [1, 2, 3]}) # noqa: F841
+ add_one = pf.udf(lambda value: value + 1, return_dtype=int) # noqa:
F841
+
+ result = (
+ pf.sql("SELECT add_one(a) AS b FROM df")
+ .filter(pf.col("b") > 2)
+ .with_columns(c=add_one(pf.col("b")))
+ )
+
+ self.assertEqual(
+ _collect_sorted(result),
+ [Row(3, 4), Row(4, 5)],
+ )
+
+ def test_udf_bindings_are_dropped_after_failure(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+ add_one = pf.udf(lambda value: value + 1, return_dtype=int) # noqa:
F841
+
+ with self.assertRaises(Py4JJavaError):
+ pf.sql("SELECT add_one(nonexistent_column) FROM df")
+
+ self.assertNotIn("add_one", self.t_env.list_user_defined_functions())
+
+ def
test_auto_bind_warns_and_skips_udf_colliding_with_existing_function(self):
+ self.t_env.create_temporary_system_function(
+ "add_one", table_udf(lambda i: i + 100,
result_type=DataTypes.BIGINT())
+ )
+ self.addCleanup(self.t_env.drop_temporary_system_function, "add_one")
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+ add_one = pf.udf(lambda value: value + 1, return_dtype=int) # noqa:
F841
+
+ with self.assertWarnsRegex(UserWarning, "skipped 'add_one'.*already
exists"):
+ result = pf.sql("SELECT add_one(a) FROM df")
+
+ # The pre-existing function wins and survives the call.
+ self.assertEqual(result.collect(), [Row(101)])
+ self.assertIn("add_one", self.t_env.list_user_defined_functions())
+
+ def
test_auto_bind_warns_and_skips_udf_colliding_with_builtin_function(self):
+ df = pf.from_dict({"a": [-1]}) # noqa: F841
+ abs = pf.udf(lambda value: value + 100, return_dtype=int) # noqa: F841
+
+ with self.assertWarnsRegex(UserWarning, "skipped 'abs'.*already
exists"):
+ result = pf.sql("SELECT abs(a) FROM df")
+
+ # The built-in function is never shadowed.
+ self.assertEqual(result.collect(), [Row(1)])
+
+ def test_explicit_udf_binding_shadows_builtin_function(self):
+ df = pf.from_dict({"a": [-1]})
+ my_abs = pf.udf(lambda value: value + 100, return_dtype=int)
+
+ result = pf.sql("SELECT abs(a) FROM df", auto_bind=False, df=df,
abs=my_abs)
+
+ self.assertEqual(result.collect(), [Row(99)])
+ # The built-in function is restored after the call.
+ self.assertEqual(pf.sql("SELECT abs(a) FROM df", auto_bind=False,
df=df)
+ .collect(), [Row(1)])
+
+ def test_explicit_udf_binding_shadows_permanent_function(self):
+ # RichFunc0 adds one to its argument.
+ self.t_env.create_java_function(
+ "perm", "org.apache.flink.table.utils.TestingFunctions$RichFunc0"
+ )
+ self.addCleanup(self.t_env.drop_function, "perm")
+ df = pf.from_dict({"a": [1]})
+ my_perm = pf.udf(lambda value: value + 100, return_dtype=int)
+
+ result = pf.sql("SELECT perm(CAST(a AS INT)) FROM df",
auto_bind=False, df=df, perm=my_perm)
+
+ self.assertEqual(result.collect(), [Row(101)])
+ # The permanent function is restored after the call.
+ self.assertEqual(
+ pf.sql("SELECT perm(CAST(a AS INT)) FROM df", auto_bind=False,
df=df).collect(),
+ [Row(2)],
+ )
+ self.assertIn("perm", self.t_env.list_user_defined_functions())
+
+ def
test_auto_bind_warns_and_skips_udf_colliding_with_permanent_function(self):
+ self.t_env.create_java_function(
+ "perm", "org.apache.flink.table.utils.TestingFunctions$RichFunc0"
+ )
+ self.addCleanup(self.t_env.drop_function, "perm")
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+ perm = pf.udf(lambda value: value + 100, return_dtype=int) # noqa:
F841
+
+ with self.assertWarnsRegex(UserWarning, "skipped 'perm'.*already
exists"):
+ result = pf.sql("SELECT perm(CAST(a AS INT)) FROM df")
+
+ # The permanent function wins.
+ self.assertEqual(result.collect(), [Row(2)])
+
+ def
test_explicit_udf_binding_collision_with_temporary_function_raises(self):
+ self.t_env.create_temporary_system_function(
+ "taken", table_udf(lambda i: i + 100,
result_type=DataTypes.BIGINT())
+ )
+ self.addCleanup(self.t_env.drop_temporary_system_function, "taken")
+ df = pf.from_dict({"a": [1]})
+
+ with self.assertRaisesRegex(
+ ValueError, "'taken'.*temporary function.*already exists"
+ ):
+ pf.sql(
+ "SELECT taken(a) FROM df",
+ auto_bind=False,
+ df=df,
+ taken=pf.udf(lambda value: value + 1, return_dtype=int),
+ )
+
+ # The view registered before the failure is cleaned up.
+ self.assertNotIn("df", self.t_env.list_temporary_views())
+
+ def
test_explicit_udf_binding_collision_with_temporary_catalog_function_raises(self):
+ self.t_env.create_temporary_function(
+ "taken", table_udf(lambda i: i + 100,
result_type=DataTypes.BIGINT())
+ )
+ self.addCleanup(self.t_env.drop_temporary_function, "taken")
+
+ with self.assertRaisesRegex(
+ ValueError, "'taken'.*temporary function.*already exists"
+ ):
+ pf.sql(
+ "SELECT 1",
+ auto_bind=False,
+ taken=pf.udf(lambda value: value + 1, return_dtype=int),
+ )
+
+ def
test_explicit_udf_binding_collision_with_temporary_function_shadowing_permanent_raises(
+ self,
+ ):
+ # A temporary function with the same name as a permanent one is not
visible
+ # in the environment's merged function listings; it must still block
the binding.
+ self.t_env.create_java_function(
+ "taken", "org.apache.flink.table.utils.TestingFunctions$RichFunc0"
+ )
+ self.addCleanup(self.t_env.drop_function, "taken")
+ self.t_env.create_temporary_system_function(
+ "taken", table_udf(lambda i: i + 100,
result_type=DataTypes.BIGINT())
+ )
+ self.addCleanup(self.t_env.drop_temporary_system_function, "taken")
+
+ with self.assertRaisesRegex(
+ ValueError, "'taken'.*temporary function.*already exists"
+ ):
+ pf.sql(
+ "SELECT 1",
+ auto_bind=False,
+ taken=pf.udf(lambda value: value + 1, return_dtype=int),
+ )
+
+ def
test_explicit_udf_binding_collision_with_temporary_catalog_function_shadowing_permanent_raises(
# noqa: E501
+ self,
+ ):
+ self.t_env.create_java_function(
+ "taken", "org.apache.flink.table.utils.TestingFunctions$RichFunc0"
+ )
+ self.addCleanup(self.t_env.drop_function, "taken")
+ self.t_env.create_temporary_function(
+ "taken", table_udf(lambda i: i + 100,
result_type=DataTypes.BIGINT())
+ )
+ self.addCleanup(self.t_env.drop_temporary_function, "taken")
+
+ with self.assertRaisesRegex(
+ ValueError, "'taken'.*temporary function.*already exists"
+ ):
+ pf.sql(
+ "SELECT 1",
+ auto_bind=False,
+ taken=pf.udf(lambda value: value + 1, return_dtype=int),
+ )
+
+ def
test_auto_bound_registrations_are_rolled_back_when_auto_bind_raises(self):
+ # With warnings turned into errors, the skipped auto-bound candidate
raises
+ # after earlier candidates were already registered; nothing must be
left behind.
+ df = pf.from_dict({"a": [-1]}) # noqa: F841
+ add_one = pf.udf(lambda value: value + 1, return_dtype=int) # noqa:
F841
+ abs = pf.udf(lambda value: value + 100, return_dtype=int) # noqa: F841
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("error")
+ with self.assertRaisesRegex(UserWarning, "skipped 'abs'"):
+ pf.sql("SELECT add_one(a), abs(a) FROM df")
+
+ self.assertNotIn("df", self.t_env.list_temporary_views())
+ self.assertNotIn("add_one", self.t_env.list_user_defined_functions())
+
+ def
test_explicit_udf_binding_precedence_over_auto_bind_is_case_insensitive(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+ add_one = pf.udf(lambda value: value + 1, return_dtype=int) # noqa:
F841
+ override = pf.udf(lambda value: value + 100, return_dtype=int)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("error")
+ result = pf.sql("SELECT ADD_ONE(a) FROM df", ADD_ONE=override)
+
+ self.assertEqual(result.collect(), [Row(101)])
+
+ def
test_explicit_udf_bindings_are_rolled_back_when_a_later_one_fails(self):
+ self.t_env.create_temporary_system_function(
+ "taken", table_udf(lambda i: i + 100,
result_type=DataTypes.BIGINT())
+ )
+ self.addCleanup(self.t_env.drop_temporary_system_function, "taken")
+
+ with self.assertRaisesRegex(ValueError, "'taken'.*already exists"):
+ pf.sql(
+ "SELECT 1",
+ auto_bind=False,
+ first=pf.udf(lambda value: value + 1, return_dtype=int),
+ taken=pf.udf(lambda value: value + 1, return_dtype=int),
+ )
+
+ self.assertNotIn("first", self.t_env.list_user_defined_functions())
+
+ def test_explicit_udf_binding_collision_is_case_insensitive(self):
+ self.t_env.create_temporary_system_function(
+ "taken", table_udf(lambda i: i + 100,
result_type=DataTypes.BIGINT())
+ )
+ self.addCleanup(self.t_env.drop_temporary_system_function, "taken")
+
+ with self.assertRaisesRegex(ValueError, "'TAKEN'.*already exists"):
+ pf.sql(
+ "SELECT 1",
+ auto_bind=False,
+ TAKEN=pf.udf(lambda value: value + 1, return_dtype=int),
+ )
+
+ def test_explicit_udf_binding_with_invalid_sql_identifier_raises(self):
+ with self.assertRaisesRegex(
+ ValueError, "'my udf'.*not a valid SQL identifier"
+ ):
+ pf.sql(
+ "SELECT 1",
+ auto_bind=False,
+ **{"my udf": pf.udf(lambda value: value + 1,
return_dtype=int)},
+ )
+
+ def test_udf_bindings_do_not_take_part_in_environment_resolution(self):
+ other_env =
TableEnvironment.create(EnvironmentSettings.in_batch_mode())
+ source = pf.DataFrame(other_env.from_elements([(1,)], ["a"]))
+ add_one = pf.udf(lambda value: value + 1, return_dtype=int) # noqa:
F841
+
+ # The explicit binding selects other_env; the auto-bound UDF follows
it.
+ result = pf.sql("SELECT add_one(a) FROM src", src=source)
+
+ self.assertEqual(result.collect(), [Row(2)])
+ self.assertNotIn("add_one", other_env.list_user_defined_functions())
+ self.assertNotIn("add_one", self.t_env.list_user_defined_functions())
+
def test_explicit_bindings_resolve_the_environment(self):
other_env =
TableEnvironment.create(EnvironmentSettings.in_batch_mode())
source = pf.DataFrame(other_env.from_elements([(1,), (2,)], ["a"]))
@@ -304,7 +644,7 @@ class SqlTests(PyFlinkDataFrameUTTestCase):
result = pf.sql("SELECT a FROM src", auto_bind=False, src=source)
self.assertEqual(
- sorted(result.collect(), key=lambda row: row[0]), [Row(1), Row(2)]
+ _collect_sorted(result), [Row(1), Row(2)]
)
# The environment is resolved per call; the global one is untouched.
self.assertIs(pf.get_table_environment(), self.t_env)
