diff --git a/flink-python/docs/reference/pyflink.dataframe/index.rst b/flink-python/docs/reference/pyflink.dataframe/index.rst index e07a631d1fe4f..5f0645c2baa38 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/sql.rst b/flink-python/docs/reference/pyflink.dataframe/sql.rst new file mode 100644 index 0000000000000..fdfb26ab7961b --- /dev/null +++ b/flink-python/docs/reference/pyflink.dataframe/sql.rst @@ -0,0 +1,43 @@ +.. ################################################################################ + 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. + ################################################################################ + +=== +SQL +=== + +Execute SQL SELECT queries against DataFrames. + +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"]}) + >>> 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 88e7ca2aad55a..326b139efc9da 100644 --- a/flink-python/pyflink/dataframe/__init__.py +++ b/flink-python/pyflink/dataframe/__init__.py @@ -54,6 +54,7 @@ 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 @@ "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 0000000000000..a7137570f9526 --- /dev/null +++ b/flink-python/pyflink/dataframe/sql.py @@ -0,0 +1,253 @@ +################################################################################ +# 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 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"] + + +@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 auto-bound candidates, falling back to the global + environment (see :func:`get_or_create_table_environment`). The resolved + environment is used only for this call and never replaces the global one. + Explicit bindings from different environments raise :class:`ValueError`; + auto-bound candidates that do not match the resolved environment are skipped + with a warning. + + :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. + :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 as e: + warnings.warn( + f"sql() failed to drop temporary view '{name}': {e}", + UserWarning, + ) + + +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, falling back to the global environment. Explicit bindings from + different environments are an error the caller must resolve; auto-bind is + best-effort, so mixed auto-bound candidates fall back to the global environment + (the non-matching ones are skipped with a warning during registration). + """ + 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(v._table._t_env): v._table._t_env for v in auto.values()} + if len(auto_envs) == 1: + return next(iter(auto_envs.values())) + 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}. Pass it as an explicit binding " + "to override.", + 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 0000000000000..0bbf458b221d1 --- /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_is_skipped(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.assertWarnsRegex(UserWarning, "skipped 'foreign'"): + result = pf.sql("SELECT a FROM df") + + self.assertEqual(result.collect(), [Row(2)]) + + +if __name__ == "__main__": + unittest.main()