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 9e76ac48dab [FLINK-40197][python] Add sql() to the DataFrame API
9e76ac48dab is described below
commit 9e76ac48dab13733dce130de0da5b017f4c2dae2
Author: Federico Dolce <[email protected]>
AuthorDate: Fri Aug 28 12:39:36 2026 +0200
[FLINK-40197][python] Add sql() to the DataFrame API
Add pyflink.dataframe.sql(), which executes a SQL SELECT query and
returns the result as a DataFrame. Referenced DataFrames are registered
as temporary views for the duration of the call and dropped afterwards:
auto-binding registers DataFrames from the caller's scope under their
variable names (best-effort, with warnings on invalid identifiers and
collisions, never shadowing permanent catalog objects), while explicit
keyword bindings are strict, take precedence, and may intentionally
shadow permanent catalog tables.
This closes #29036.
---
.../docs/reference/pyflink.dataframe/index.rst | 1 +
.../pyflink.dataframe/{index.rst => sql.rst} | 33 +-
flink-python/pyflink/dataframe/__init__.py | 2 +
flink-python/pyflink/dataframe/sql.py | 266 +++++++++++++++
flink-python/pyflink/dataframe/tests/test_sql.py | 370 +++++++++++++++++++++
5 files changed, 661 insertions(+), 11 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/index.rst
b/flink-python/docs/reference/pyflink.dataframe/index.rst
index e07a631d1fe..5f0645c2baa 100644
--- a/flink-python/docs/reference/pyflink.dataframe/index.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/index.rst
@@ -28,5 +28,6 @@ This page gives an overview of all public PyFlink DataFrame
APIs.
dataframe
creation
io
+ sql
datatype
environment
diff --git a/flink-python/docs/reference/pyflink.dataframe/index.rst
b/flink-python/docs/reference/pyflink.dataframe/sql.rst
similarity index 61%
copy from flink-python/docs/reference/pyflink.dataframe/index.rst
copy to flink-python/docs/reference/pyflink.dataframe/sql.rst
index e07a631d1fe..fdfb26ab796 100644
--- a/flink-python/docs/reference/pyflink.dataframe/index.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/sql.rst
@@ -16,17 +16,28 @@
limitations under the License.
################################################################################
-==================
-PyFlink DataFrame
-==================
+===
+SQL
+===
-This page gives an overview of all public PyFlink DataFrame APIs.
+Execute SQL SELECT queries against DataFrames.
-.. toctree::
- :maxdepth: 1
+Example::
- dataframe
- creation
- io
- datatype
- environment
+ >>> import pyflink.dataframe as pf
+ >>> df1 = pf.from_dict({"a": [1, 2, 3], "b": ["x", "y", "z"]})
+ >>> df2 = pf.from_dict({"a": [1, 2, 3], "c": ["p", "q", "r"]})
+ >>> joined = pf.sql("SELECT df1.a, b, c FROM df1 JOIN df2 ON df1.a =
df2.a")
+ >>> result = pf.sql(
+ ... "SELECT * FROM src WHERE a > 1",
+ ... auto_bind=False,
+ ... src=df1,
+ ... )
+ >>> pf.sql("SELECT a, b FROM df1").filter(pf.col("a") > 1).to_pandas()
+
+.. currentmodule:: pyflink.dataframe
+
+.. autosummary::
+ :toctree: api/
+
+ sql
diff --git a/flink-python/pyflink/dataframe/__init__.py
b/flink-python/pyflink/dataframe/__init__.py
index 88e7ca2aad5..326b139efc9 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -54,6 +54,7 @@ from pyflink.dataframe.context import (
from pyflink.dataframe.dataframe import DataFrame, GroupedDataFrame, col, lit
from pyflink.dataframe.datatype import DataType
from pyflink.dataframe.io import read_generic
+from pyflink.dataframe.sql import sql
__all__ = [
"DataFrame",
@@ -68,6 +69,7 @@ __all__ = [
"from_table",
"range",
"read_generic",
+ "sql",
"set_table_environment",
"get_table_environment",
"get_or_create_table_environment",
diff --git a/flink-python/pyflink/dataframe/sql.py
b/flink-python/pyflink/dataframe/sql.py
new file mode 100644
index 00000000000..f244a81897b
--- /dev/null
+++ b/flink-python/pyflink/dataframe/sql.py
@@ -0,0 +1,266 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+import inspect
+import logging
+import warnings
+from typing import Any, Dict, List
+
+from py4j.protocol import Py4JJavaError
+
+from pyflink.dataframe.context import get_or_create_table_environment
+from pyflink.dataframe.dataframe import DataFrame
+from pyflink.table import Table, TableEnvironment
+from pyflink.util.api_stability_decorators import PublicEvolving
+from pyflink.util.java_utils import is_instance_of
+
+__all__ = ["sql"]
+
+_LOG = logging.getLogger(__name__)
+
+
+@PublicEvolving()
+def sql(query: str, *, auto_bind: bool = True, **bindings: DataFrame) ->
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.
+
+ 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.
+
+ 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.
+
+ The query runs in the :class:`TableEnvironment` of the bound DataFrames:
the
+ environment shared by the explicit ``bindings`` when given, otherwise the
+ environment shared by all valid auto-bound candidates. Explicit bindings
from
+ 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
+ environment is used only for this call and never replaces the global one.
+
+ :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.
+ :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`.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df1 = pf.from_dict({"a": [1, 2, 3], "b": ["x", "y", "z"]})
+ >>> df2 = pf.from_dict({"a": [1, 2, 3], "c": ["p", "q", "r"]})
+ >>> # Auto-bind: df1 / df2 are registered under their variable names
+ >>> joined = pf.sql("SELECT df1.a, b, c FROM df1 JOIN df2 ON df1.a =
df2.a")
+ >>> # Explicit bindings: pick the SQL names, turn off scanning
+ >>> result = pf.sql(
+ ... "SELECT * FROM src WHERE a > 1",
+ ... auto_bind=False,
+ ... src=df1,
+ ... )
+ >>> # 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] = {}
+ 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] = []
+ try:
+ _register_bindings(t_env, bindings, auto_bindings, registered)
+ 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
+ )
+
+
+def _resolve_table_environment(
+ explicit: Dict[str, Any], auto: Dict[str, DataFrame]
+) -> TableEnvironment:
+ """
+ Pick the environment to run the query in: the environment shared by the
explicit
+ bindings when given, otherwise the environment shared by all auto-bound
+ candidates. Explicit bindings from different environments are an error the
caller
+ must resolve. Auto-bound candidates with invalid identifiers do not
participate in
+ 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()}
+ 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
+ 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"
+ )
+ return get_or_create_table_environment()
+
+
+def _execute_query(t_env: TableEnvironment, query: str) -> Table:
+ """
+ Run ``query`` through :meth:`TableEnvironment.sql_query`, which parses the
statement
+ and rejects anything that is not a single query returning a result.
Translate that
+ rejection into a plain :class:`ValueError`.
+ """
+ try:
+ return t_env.sql_query(query)
+ except Py4JJavaError as e:
+ if "Unsupported SQL query!" in str(e.java_exception):
+ raise ValueError(
+ "sql() only supports queries that return a result, such as
SELECT "
+ "or VALUES (no INSERT / DDL); use
TableEnvironment.execute_sql() "
+ "for other statements."
+ ) from e
+ raise
+
+
+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
+ 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.
+ """
+ try:
+ identifier = t_env._j_tenv.getParser().parseIdentifier(name)
+ except Py4JJavaError as e:
+ if not is_instance_of(
+ e.java_exception, "org.apache.flink.table.api.SqlParserException"
+ ):
+ raise
+ return False
+ return (
+ not identifier.getCatalogName().isPresent()
+ and not identifier.getDatabaseName().isPresent()
+ and identifier.getObjectName() == name
+ )
+
+
+def _register_bindings(
+ t_env: TableEnvironment,
+ explicit: Dict[str, DataFrame],
+ auto: Dict[str, DataFrame],
+ registered: List[str],
+) -> None:
+ """
+ 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`).
+ """
+ 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())
+
+ 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:
+ t_env.create_temporary_view(name, value.to_table())
+ except Exception as e:
+ _warn_skipped(name, f"registration failed: {e}")
+ continue
+ registered.append(name)
+
+
+def _warn_skipped(name: str, reason: str) -> None:
+ warnings.warn(
+ f"sql() auto-bind skipped '{name}': {reason}.",
+ UserWarning,
+ )
diff --git a/flink-python/pyflink/dataframe/tests/test_sql.py
b/flink-python/pyflink/dataframe/tests/test_sql.py
new file mode 100644
index 00000000000..bed0eecae5b
--- /dev/null
+++ b/flink-python/pyflink/dataframe/tests/test_sql.py
@@ -0,0 +1,370 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+import unittest
+
+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.testing.test_case_utils import PyFlinkDataFrameUTTestCase
+
+
+class SqlValidationTests(unittest.TestCase):
+ def setUp(self):
+ previous_environment = pf.get_table_environment()
+ self.addCleanup(pf.set_table_environment, previous_environment)
+ pf.set_table_environment(None)
+
+ def test_query_must_be_a_string_checked_before_environment_creation(self):
+ with self.assertRaisesRegex(TypeError, "query must be a string"):
+ pf.sql(42)
+
+ self.assertIsNone(pf.get_table_environment())
+
+
+class SqlTests(PyFlinkDataFrameUTTestCase):
+ def test_non_query_statements_are_rejected(self):
+ self.t_env.execute_sql(
+ "CREATE TABLE sink (a BIGINT) WITH ('connector' = 'blackhole')"
+ )
+ self.addCleanup(self.t_env.execute_sql, "DROP TABLE sink")
+
+ for statement in [
+ "INSERT INTO sink VALUES (1)",
+ "CREATE TABLE t (a INT)",
+ "DROP TABLE t",
+ "EXPLAIN SELECT 1",
+ ]:
+ with self.subTest(statement=statement):
+ with self.assertRaisesRegex(
+ ValueError, "only supports queries that return a result"
+ ):
+ pf.sql(statement)
+
+ def test_unparsable_statements_surface_the_java_error(self):
+ for statement, error in [
+ ("", "SQL parse failed"),
+ ("-- only a comment", "only single statement supported"),
+ ("SELECT 1; SELECT 2", "only single statement supported"),
+ ]:
+ with self.subTest(statement=statement):
+ with self.assertRaisesRegex(Py4JJavaError, error):
+ pf.sql(statement)
+
+ def test_auto_bind_joins_dataframes_by_variable_name(self):
+ df1 = pf.from_dict({"a": [1, 2, 3], "b": ["x", "y", "z"]}) # noqa:
F841
+ df2 = pf.from_dict({"a": [1, 2, 3], "c": ["p", "q", "r"]}) # noqa:
F841
+
+ joined = pf.sql(
+ "SELECT df1.a, b, c FROM df1 JOIN df2 ON df1.a = df2.a"
+ )
+
+ self.assertEqual(
+ sorted(joined.collect(), key=lambda row: row[0]),
+ [Row(1, "x", "p"), Row(2, "y", "q"), Row(3, "z", "r")],
+ )
+
+ def test_select_variants_are_accepted(self):
+ df = pf.from_dict({"a": [1, 2]}) # noqa: F841
+
+ for query in [
+ "SELECT a FROM df",
+ "-- leading comment\nSELECT a FROM df",
+ "/* block comment */ SELECT a FROM df",
+ "WITH cte AS (SELECT a FROM df) SELECT a FROM cte",
+ "(SELECT a FROM df)",
+ ]:
+ with self.subTest(query=query):
+ self.assertEqual(
+ sorted(pf.sql(query).collect(), key=lambda row: row[0]),
+ [Row(1), Row(2)],
+ )
+
+ self.assertEqual(pf.sql("VALUES (1)").collect(), [Row(1)])
+
+ def test_explicit_bindings_with_auto_bind_disabled(self):
+ src = pf.from_dict({"a": [1, 2]}) # noqa: F841
+ other = pf.from_dict({"a": [10, 20]})
+
+ result = pf.sql(
+ "SELECT a FROM src WHERE a > 1", auto_bind=False, src=other
+ )
+
+ self.assertEqual(
+ sorted(result.collect(), key=lambda row: row[0]),
+ [Row(10), Row(20)],
+ )
+
+ def test_auto_bind_disabled_ignores_caller_variables(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+
+ with self.assertRaisesRegex(Py4JJavaError, "Object 'df' not found"):
+ pf.sql("SELECT a FROM df", auto_bind=False)
+
+ self.assertNotIn("df", self.t_env.list_temporary_views())
+
+ def test_explicit_bindings_take_precedence_over_auto_bind(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+ other = pf.from_dict({"a": [42]})
+
+ result = pf.sql("SELECT a FROM df", df=other)
+
+ self.assertEqual(result.collect(), [Row(42)])
+
+ def test_auto_bind_finds_module_level_globals(self):
+ globals()["global_test_df"] = pf.from_dict({"a": [7]})
+ self.addCleanup(globals().pop, "global_test_df", None)
+
+ result = pf.sql("SELECT a FROM global_test_df")
+
+ self.assertEqual(result.collect(), [Row(7)])
+
+ def test_auto_bind_locals_take_precedence_over_globals(self):
+ globals()["shadow_df"] = pf.from_dict({"a": [1]})
+ self.addCleanup(globals().pop, "shadow_df", None)
+ shadow_df = pf.from_dict({"a": [2]}) # noqa: F841
+
+ self.assertEqual(pf.sql("SELECT a FROM shadow_df").collect(), [Row(2)])
+
+ def test_auto_bind_warns_and_skips_on_collision_with_existing_view(self):
+ self.t_env.create_temporary_view(
+ "df", pf.from_dict({"a": [100]}).to_table()
+ )
+ self.addCleanup(self.t_env.drop_temporary_view, "df")
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+
+ with self.assertWarnsRegex(UserWarning, "skipped 'df'"):
+ result = pf.sql("SELECT a FROM df")
+
+ # The pre-existing view wins and survives the call.
+ self.assertEqual(result.collect(), [Row(100)])
+ self.assertIn("df", self.t_env.list_temporary_views())
+
+ def test_explicit_binding_collision_with_temporary_view_raises(self):
+ self.t_env.create_temporary_view(
+ "src", pf.from_dict({"a": [100]}).to_table()
+ )
+ self.addCleanup(self.t_env.drop_temporary_view, "src")
+
+ with self.assertRaisesRegex(ValueError, "'src'.*already exists"):
+ pf.sql(
+ "SELECT a FROM src",
+ auto_bind=False,
+ src=pf.from_dict({"a": [1]}),
+ )
+
+ def
test_partial_registrations_are_dropped_when_a_later_binding_fails(self):
+ self.t_env.create_temporary_view(
+ "taken", pf.from_dict({"a": [100]}).to_table()
+ )
+ self.addCleanup(self.t_env.drop_temporary_view, "taken")
+
+ with self.assertRaisesRegex(ValueError, "'taken'.*already exists"):
+ pf.sql(
+ "SELECT a FROM fresh",
+ auto_bind=False,
+ fresh=pf.from_dict({"a": [1]}),
+ taken=pf.from_dict({"a": [2]}),
+ )
+
+ # The binding registered before the failure is cleaned up.
+ self.assertNotIn("fresh", self.t_env.list_temporary_views())
+ self.assertIn("taken", self.t_env.list_temporary_views())
+
+ def test_explicit_binding_shadows_permanent_table(self):
+ self.t_env.execute_sql(
+ "CREATE TABLE perm (a BIGINT) "
+ "WITH ('connector' = 'datagen', 'number-of-rows' = '1')"
+ )
+ self.addCleanup(self.t_env.execute_sql, "DROP TABLE perm")
+
+ result = pf.sql(
+ "SELECT a FROM perm",
+ auto_bind=False,
+ perm=pf.from_dict({"a": [42]}),
+ )
+
+ self.assertEqual(result.collect(), [Row(42)])
+ # The permanent table is intact after the call.
+ self.assertIn("perm", self.t_env.list_tables())
+ self.assertNotIn("perm", self.t_env.list_temporary_views())
+
+ def test_auto_bind_warns_and_skips_on_collision_with_permanent_table(self):
+ self.t_env.execute_sql(
+ "CREATE TABLE perm (a BIGINT) WITH ("
+ "'connector' = 'datagen', 'fields.a.kind' = 'sequence', "
+ "'fields.a.start' = '100', 'fields.a.end' = '100')"
+ )
+ self.addCleanup(self.t_env.execute_sql, "DROP TABLE perm")
+ perm = pf.from_dict({"a": [1]}) # noqa: F841
+
+ with self.assertWarnsRegex(UserWarning, "skipped 'perm'"):
+ result = pf.sql("SELECT a FROM perm")
+
+ # The permanent table wins and is never shadowed.
+ self.assertEqual(result.collect(), [Row(100)])
+ self.assertNotIn("perm", self.t_env.list_temporary_views())
+
+ def test_auto_bind_skips_invalid_sql_identifiers_with_warning(self):
+ globals()["my df"] = pf.from_dict({"a": [1]})
+ self.addCleanup(globals().pop, "my df", None)
+ df = pf.from_dict({"a": [2]}) # noqa: F841
+
+ with self.assertWarnsRegex(
+ UserWarning, "skipped 'my df'.*not a valid SQL identifier"
+ ):
+ result = pf.sql("SELECT a FROM df")
+
+ self.assertEqual(result.collect(), [Row(2)])
+
+ def test_auto_bind_supports_unicode_identifiers(self):
+ globals()["dfé"] = pf.from_dict({"a": [1]})
+ self.addCleanup(globals().pop, "dfé", None)
+
+ self.assertEqual(pf.sql("SELECT a FROM dfé").collect(), [Row(1)])
+
+ def test_auto_bind_supports_keyword_names_via_quoting(self):
+ order = pf.from_dict({"a": [1]}) # noqa: F841
+
+ self.assertEqual(pf.sql("SELECT a FROM `order`").collect(), [Row(1)])
+
+ def test_bindings_are_dropped_after_success(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+
+ pf.sql("SELECT a FROM df")
+
+ self.assertNotIn("df", self.t_env.list_temporary_views())
+
+ def test_bindings_are_dropped_after_failure(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+
+ with self.assertRaises(Py4JJavaError):
+ pf.sql("SELECT nonexistent_column FROM df")
+
+ self.assertNotIn("df", self.t_env.list_temporary_views())
+
+ def test_result_composes_with_dataframe_api(self):
+ df1 = pf.from_dict({"a": [1, 2, 3], "b": ["x", "y", "z"]}) # noqa:
F841
+
+ result = (
+ pf.sql("SELECT a, b FROM df1")
+ .filter(pf.col("a") > 1)
+ .to_pandas()
+ )
+
+ self.assertEqual(sorted(result["a"].tolist()), [2, 3])
+
+ def test_explicit_binding_of_unsupported_type_raises(self):
+ with self.assertRaisesRegex(TypeError, "'x' must be a DataFrame"):
+ pf.sql("SELECT * FROM x", auto_bind=False, x=42)
+
+ def test_explicit_binding_of_raw_table_raises(self):
+ table = pf.from_dict({"a": [1]}).to_table()
+
+ 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())
+
+ with self.assertRaisesRegex(TypeError, "'add_one' must be a
DataFrame"):
+ pf.sql("SELECT add_one(a) FROM df", auto_bind=False,
add_one=add_one)
+
+ def test_auto_bind_ignores_udfs(self):
+ df = pf.from_dict({"a": [1]}) # noqa: F841
+ add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT()) #
noqa: F841
+
+ with self.assertRaisesRegex(Py4JJavaError, "No match found for
function"):
+ pf.sql("SELECT add_one(a) FROM df")
+
+ 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"]))
+
+ 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)]
+ )
+ # The environment is resolved per call; the global one is untouched.
+ self.assertIs(pf.get_table_environment(), self.t_env)
+ self.assertNotIn("src", other_env.list_temporary_views())
+
+ def test_explicit_bindings_from_different_environments_raise(self):
+ other_env =
TableEnvironment.create(EnvironmentSettings.in_batch_mode())
+ foreign = pf.DataFrame(other_env.from_elements([(1,)], ["a"]))
+ local = pf.from_dict({"b": [2]})
+
+ with self.assertRaisesRegex(ValueError, "different TableEnvironments"):
+ pf.sql(
+ "SELECT * FROM one JOIN two ON TRUE",
+ auto_bind=False,
+ one=foreign,
+ two=local,
+ )
+
+ self.assertNotIn("one", other_env.list_temporary_views())
+ self.assertNotIn("two", self.t_env.list_temporary_views())
+
+ def test_auto_bound_dataframes_sharing_an_environment_resolve_it(self):
+ other_env =
TableEnvironment.create(EnvironmentSettings.in_batch_mode())
+ remote_df = pf.DataFrame(other_env.from_elements([(1,)], ["a"])) #
noqa: F841
+
+ self.assertEqual(pf.sql("SELECT a FROM remote_df").collect(), [Row(1)])
+ self.assertIs(pf.get_table_environment(), self.t_env)
+
+ def test_environment_resolved_from_bindings_does_not_become_global(self):
+ pf.set_table_environment(None)
+ self.addCleanup(pf.set_table_environment, self.t_env)
+ other_env =
TableEnvironment.create(EnvironmentSettings.in_batch_mode())
+ source = pf.DataFrame(other_env.from_elements([(1,)], ["a"]))
+
+ self.assertEqual(
+ pf.sql("SELECT a FROM src", auto_bind=False, src=source).collect(),
+ [Row(1)],
+ )
+ self.assertIsNone(pf.get_table_environment())
+
+ def test_explicit_binding_with_invalid_sql_identifier_raises(self):
+ df = pf.from_dict({"a": [1]})
+
+ with self.assertRaisesRegex(
+ ValueError, "'my df'.*not a valid SQL identifier"
+ ):
+ pf.sql("SELECT a FROM `my df`", auto_bind=False, **{"my df": df})
+
+ self.assertNotIn("my df", self.t_env.list_temporary_views())
+
+ def test_auto_bound_dataframe_from_other_environment_raises(self):
+ other_env =
TableEnvironment.create(EnvironmentSettings.in_batch_mode())
+ foreign = pf.DataFrame(other_env.from_elements([(1,)], ["a"])) #
noqa: F841
+ df = pf.from_dict({"a": [2]}) # noqa: F841
+
+ with self.assertRaisesRegex(
+ ValueError, "auto-bound DataFrames belong to different
TableEnvironments"
+ ):
+ pf.sql("SELECT a FROM df")
+
+
+if __name__ == "__main__":
+ unittest.main()
