################################################################################
# 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 entry point for the DataFrame API.
Allows users to write a SQL SELECT statement that references DataFrames in
the calling scope (or passed explicitly), returning a new DataFrame.
"""
import inspect
import logging
import warnings
from types import FrameType
from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Set, Tuple, Union
from pyflink.dataframe.context import (
get_or_create_table_environment,
get_table_environment,
)
from pyflink.dataframe.dataframe import DataFrame
from pyflink.dataframe.udf import DataFrameFunctionWrapper
if TYPE_CHECKING:
from pyflink.table import StreamTableEnvironment
__all__ = ["sql"]
_logger = logging.getLogger(__name__)
SqlBinding = Union[DataFrame, DataFrameFunctionWrapper]
def _parse_identifier(t_env: "StreamTableEnvironment", name: str):
return t_env._j_tenv.getParser().parseIdentifier(name)
def _is_valid_ident(t_env: "StreamTableEnvironment", name: str) -> bool:
try:
_parse_identifier(t_env, name)
return True
except Exception: # noqa: BLE001
return False
def _validate_identifier(
t_env: "StreamTableEnvironment", name: str, kind: str
) -> None:
try:
_parse_identifier(t_env, name)
except Exception as e: # noqa: BLE001
raise ValueError(f"'{name}' is not a valid SQL {kind} name.") from e
def _name_exists(name: str, existing_names: Set[str]) -> bool:
return name in existing_names
def _normalize_function_name(name: str) -> str:
# FunctionCatalog uses FunctionIdentifier.normalizeName(name), which is
# lower-case normalization. Python lower() matches that for the ASCII SQL
# identifiers expected here; unlike functions, table/view names remain
# case-sensitive and must not be normalized.
return name.lower()
def _normalize_function_names(names: List[str]) -> Set[str]:
return {_normalize_function_name(name) for name in names}
def _function_name_exists(name: str, normalized_existing_names: Set[str]) -> bool:
return _normalize_function_name(name) in normalized_existing_names
def _table_or_view_exists(name: str, table_or_view_names: Set[str]) -> bool:
# Table/view identifiers are case-sensitive in Flink catalog resolution.
return _name_exists(name, table_or_view_names)
def _temp_table_or_view_exists(name: str, temp_table_or_view_names: Set[str]) -> bool:
# Table/view identifiers are case-sensitive in Flink catalog resolution.
return _name_exists(name, temp_table_or_view_names)
def _function_exists(name: str, normalized_function_names: Set[str]) -> bool:
# Function identifiers are normalized to lower case by FunctionCatalog.
return _function_name_exists(name, normalized_function_names)
def _user_defined_function_exists(
name: str, normalized_user_defined_function_names: Set[str]
) -> bool:
# Includes temporary system, temporary catalog, and catalog functions.
return _function_name_exists(name, normalized_user_defined_function_names)
def _parse_function_binding_name(
t_env: "StreamTableEnvironment", name: str
) -> str:
"""Parse a function binding name into a temporary system function name.
Temporary system functions are global simple names. A quoted simple SQL
identifier is accepted, but catalog/database-qualified names are not.
"""
try:
identifier = _parse_identifier(t_env, name)
except Exception as e: # noqa: BLE001
raise ValueError(f"'{name}' is not a valid SQL function name.") from e
if (
identifier.getCatalogName().isPresent()
or identifier.getDatabaseName().isPresent()
):
raise ValueError(
f"binding '{name}' is not a valid simple SQL function name. "
f"Temporary system function bindings do not support "
f"catalog/database-qualified names."
)
return identifier.getObjectName()
def _resolve_explicit_bindings(
bindings: Mapping[str, SqlBinding],
) -> Tuple[Dict[str, DataFrame], Dict[str, DataFrameFunctionWrapper]]:
"""Validate explicit binding values and split views from functions."""
views: Dict[str, DataFrame] = {}
functions: Dict[str, DataFrameFunctionWrapper] = {}
for name, value in bindings.items():
if isinstance(value, DataFrame):
views[name] = value
elif isinstance(value, DataFrameFunctionWrapper):
functions[name] = value
else:
raise ValueError(
f"binding '{name}' must be a DataFrame or "
f"DataFrameFunctionWrapper, "
f"got {type(value).__name__}"
)
return views, functions
def _normalize_explicit_function_bindings(
t_env: "StreamTableEnvironment",
functions: Mapping[str, DataFrameFunctionWrapper],
) -> Dict[str, DataFrameFunctionWrapper]:
"""Normalize explicit function aliases and reject case-insensitive collisions."""
result: Dict[str, DataFrameFunctionWrapper] = {}
seen: Dict[str, str] = {}
for raw_name, value in functions.items():
name = _parse_function_binding_name(t_env, raw_name)
normalized_name = _normalize_function_name(name)
existing_name = seen.get(normalized_name)
if existing_name is not None:
raise ValueError(
f"binding '{raw_name}' conflicts with binding "
f"'{existing_name}' "
f"because SQL function names are case-insensitive. Use only "
f"one explicit binding for this function name."
)
seen[normalized_name] = raw_name
result[name] = value
return result
def _resolve_table_environment(
binding_views: Mapping[str, DataFrame],
auto_candidates: Mapping[str, DataFrame],
) -> "StreamTableEnvironment":
"""Resolve t_env by priority: explicit bindings, auto-bind candidates,
configured or newly created global environment.
"""
if binding_views:
first_name, first_df = next(iter(binding_views.items()))
t_env = first_df._table._t_env
_logger.debug(
"resolved TableEnvironment from explicit binding '%s'", first_name
)
for name, df in binding_views.items():
if df._table._t_env is not t_env:
raise ValueError(
f"binding '{name}' belongs to a different "
f"TableEnvironment; all explicitly bound DataFrames "
f"must share the same TableEnvironment."
)
return t_env
for name, df in auto_candidates.items():
if _is_valid_ident(df._table._t_env, name):
_logger.debug(
"resolved TableEnvironment from auto-bound DataFrame '%s'",
name,
)
return df._table._t_env
t_env = get_table_environment()
resolved_existing = t_env is not None
t_env = get_or_create_table_environment()
if resolved_existing:
_logger.debug("resolved TableEnvironment from configured global environment")
else:
_logger.debug("created new TableEnvironment for SQL query")
return t_env
def _resolve_auto_bind_dataframes(
candidates: Mapping[str, DataFrame],
bindings: Mapping[str, DataFrame],
t_env: "StreamTableEnvironment",
) -> Dict[str, DataFrame]:
"""Filter caller-scope DataFrames and decide which to register.
Auto-bind is best-effort: it never raises. Candidates that can't be
registered emit a warning so the user knows their DataFrame won't be
visible in SQL:
- the name is not a valid SQL view name;
- the name collides with an existing table/view.
If the same ``DataFrame`` is already reachable from SQL via a binding alias,
these registration failures are skipped silently instead.
"""
bound_df_ids = {id(df) for df in bindings.values()}
table_or_view_names = set(t_env.list_tables())
result: Dict[str, DataFrame] = {}
for name, value in candidates.items():
# Binding aliases make the DataFrame reachable under another name. If
# this auto candidate then fails registration, skip it without warning.
df_is_aliased = id(value) in bound_df_ids
if not _is_valid_ident(t_env, name):
if df_is_aliased:
continue
warnings.warn(
f"Skipping auto-bind of '{name}': not a valid SQL view "
f"name. To use this DataFrame in your query, pass "
f"it under a valid view name via bindings "
f"(e.g. `pf.sql(..., my_df={name})`).",
stacklevel=3, # warn -> _resolve_auto_bind_dataframes -> sql -> user
)
continue
if value._table._t_env is not t_env:
_logger.info(
"skipping auto-bind of '%s': DataFrame belongs to a "
"different TableEnvironment",
name,
)
continue
if _table_or_view_exists(name, table_or_view_names):
if df_is_aliased:
continue
warnings.warn(
f"Skipping auto-bind of '{name}': a table/view "
f"with this name already exists. To use this DataFrame in "
f"your query, either rename your Python variable, or pass "
f"it under a different name via bindings "
f"(e.g. `pf.sql(..., my_df={name})`).",
stacklevel=3, # warn -> _resolve_auto_bind_dataframes -> sql -> user
)
continue
result[name] = value
return result
def _resolve_auto_bind_functions(
candidates: Mapping[str, DataFrameFunctionWrapper],
bindings: Mapping[str, DataFrameFunctionWrapper],
t_env: "StreamTableEnvironment",
) -> Dict[str, DataFrameFunctionWrapper]:
"""Filter caller-scope DataFrame functions and decide which to register.
Auto-bind is best-effort: it never raises. Candidates that can't be
registered emit a warning so the user knows their function won't be
visible in SQL:
- the name is not a valid SQL function name;
- the name collides with an existing function;
- the name collides case-insensitively with another auto-bound function.
If the same ``DataFrameFunctionWrapper`` is already reachable from SQL via
a binding alias, these registration failures are skipped silently instead.
DataFrameFunctionWrapper does not carry a TableEnvironment, so this helper
does not perform the t_env mismatch check used for DataFrames.
"""
bound_function_ids = {id(function) for function in bindings.values()}
normalized_function_names = _normalize_function_names(t_env.list_functions())
normalized_user_defined_function_names = _normalize_function_names(
t_env.list_user_defined_functions()
)
result: Dict[str, DataFrameFunctionWrapper] = {}
# Track normalized names only for conflict detection. Functions are still
# registered under the original candidate name.
# Case variants follow the merged scope iteration order. Exact same Python
# names are local-over-global because locals update globals in sql(), but
# differently cased names are distinct Python keys and remain first-wins.
registered_normalized_function_names = set()
for name, value in candidates.items():
function_is_aliased = id(value) in bound_function_ids
if not _is_valid_ident(t_env, name):
if function_is_aliased:
continue
warnings.warn(
f"Skipping auto-bind of '{name}': not a valid SQL function "
f"name. To use this function in your query, pass it under a "
f"valid function name via bindings "
f"(e.g. `pf.sql(..., my_function={name})`).",
# warn -> _resolve_auto_bind_functions -> sql -> user
stacklevel=3,
)
continue
if _function_exists(name, normalized_function_names):
if function_is_aliased:
continue
if _user_defined_function_exists(
name, normalized_user_defined_function_names
):
guidance = (
f"To use this function in your query, either rename your "
f"Python variable, or pass it under a different name via "
f"bindings (e.g. `pf.sql(..., my_function={name})`)."
)
else:
guidance = (
f"To use this function under a different name, pass it via "
f"bindings (e.g. `pf.sql(..., my_function={name})`). If you "
f"intend to override the built-in/system function, pass "
f"it explicitly with the same SQL name "
f"(e.g. `pf.sql(..., {name}={name})`)."
)
warnings.warn(
f"Skipping auto-bind of '{name}': a function "
f"with this name already exists. {guidance}",
stacklevel=3, # warn -> _resolve_auto_bind_functions -> sql -> user
)
continue
normalized_name = _normalize_function_name(name)
# Python variables can differ only by case, but SQL function names are
# case-insensitive. Register only the first auto-bound variant.
if normalized_name in registered_normalized_function_names:
if function_is_aliased:
continue
warnings.warn(
f"Skipping auto-bind of '{name}': a function "
f"with this name is already auto-bound. Function names are "
f"case-insensitive in SQL. To use this function in your query, "
f"either rename your Python variable, or pass it under a "
f"different name via bindings "
f"(e.g. `pf.sql(..., my_function={name})`).",
# warn -> _resolve_auto_bind_functions -> sql -> user
stacklevel=3,
)
continue
result[name] = value
registered_normalized_function_names.add(normalized_name)
return result
def _execute_with_temp_objects(
t_env: "StreamTableEnvironment",
auto_views_to_register: Mapping[str, DataFrame],
binding_views_to_register: Mapping[str, DataFrame],
auto_functions_to_register: Mapping[str, DataFrameFunctionWrapper],
binding_functions_to_register: Mapping[str, DataFrameFunctionWrapper],
query: str,
) -> DataFrame:
"""Register temp views/functions, run the query, and drop them in finally.
Auto-bound objects are best-effort: registration failures warn and skip.
Explicit bindings are strict and propagate registration failures.
Drop failures are logged but never propagated — letting a cleanup
error mask the actual query outcome would be a debugging trap.
"""
registered_views: List[str] = []
registered_functions: List[str] = []
try:
for name, df in auto_views_to_register.items():
try:
t_env.create_temporary_view(name, df._table)
except Exception as e: # noqa: BLE001
warnings.warn(
f"Skipping auto-bind of '{name}': DataFrame could not "
f"be registered as a temporary view: {e}",
stacklevel=3, # warn -> _execute_with_temp_objects -> sql -> user
)
continue
registered_views.append(name)
_logger.debug("registered temp view '%s'", name)
for name, df in binding_views_to_register.items():
t_env.create_temporary_view(name, df._table)
registered_views.append(name)
_logger.debug("registered temp view '%s'", name)
for name, function in auto_functions_to_register.items():
try:
t_env.create_temporary_system_function(
name, function._table_wrapper
)
except Exception as e: # noqa: BLE001
warnings.warn(
f"Skipping auto-bind of '{name}': function could not be "
f"registered as a temporary system function: {e}",
stacklevel=3, # warn -> _execute_with_temp_objects -> sql -> user
)
continue
registered_functions.append(name)
_logger.debug("registered temp system function '%s'", name)
for name, function in binding_functions_to_register.items():
t_env.create_temporary_system_function(name, function._table_wrapper)
registered_functions.append(name)
_logger.debug("registered temp system function '%s'", name)
result = t_env.sql_query(query)
return DataFrame(result)
finally:
for name in registered_functions:
try:
t_env.drop_temporary_system_function(name)
_logger.debug("dropped temp system function '%s'", name)
except Exception as e: # noqa: BLE001
_logger.warning(
"drop_temporary_system_function(%s) failed: %s", name, e
)
for name in registered_views:
try:
t_env.drop_temporary_view(name)
_logger.debug("dropped temp view '%s'", name)
except Exception as e: # noqa: BLE001
_logger.warning(
"drop_temporary_view(%s) failed: %s", name, e
)
[docs]def sql(
query: str,
*,
auto_bind: bool = True,
**bindings: SqlBinding,
) -> DataFrame:
"""
Run a SQL SELECT query, returning a DataFrame.
In the SELECT query, ``DataFrame`` values can be used as temporary views,
and ``DataFrameFunctionWrapper`` values can be used as temporary system
functions. These temporary objects are registered only while this function
runs and are dropped before the result is returned.
You can make them visible to SQL in two ways:
- **Auto-bind**: with ``auto_bind=True`` (the default), ``sql`` scans the
caller's local and global variables for ``DataFrame`` and
``DataFrameFunctionWrapper`` instances. It registers each one under its
Python variable name.
- **Explicit bindings**: pass keyword arguments to choose the SQL name
yourself, for example ``pf.sql("SELECT * FROM source", source=df)`` or
``pf.sql("SELECT f(a) FROM source", source=df, f=my_udf)``.
Notes on the differences between auto-bind and explicit bindings:
- Auto-bind is best-effort. It warns and skips invalid SQL names, existing
tables/views, and existing functions. It logs and skips DataFrames from
a different ``TableEnvironment``. If an auto-bind candidate cannot be
registered but is already reachable through an explicit alias, the failed
auto-bind attempt is skipped silently.
- Explicit bindings are strict. Invalid SQL names, unsupported binding
values, temporary table/view conflicts, non-built-in function conflicts,
or explicitly bound DataFrames from different ``TableEnvironment``
instances raise ``ValueError``.
- Auto-bind never shadows existing SQL objects, including permanent catalog
tables/views and built-in/system functions. Use explicit bindings when you
intentionally want to shadow a permanent catalog table/view or a
built-in/system function.
Explicit bindings take precedence over auto-bind when they use the same
name.
Args:
query: A SQL SELECT statement. Other statement kinds (INSERT, DDL)
are not supported.
auto_bind: If True (default), auto-detect ``DataFrame`` and
``DataFrameFunctionWrapper`` instances in the caller's
globals and locals. See above.
**bindings: Explicit alias -> DataFrame or DataFrameFunctionWrapper
mappings. Always registered, even when ``auto_bind`` is
False. Take precedence over auto-bind on name collision.
Returns:
A new DataFrame backed by the query result.
Note:
This function is **not safe** under concurrent calls against the
same TableEnvironment: register / sql_query / drop is not an
atomic sequence. Don't call it from multiple threads sharing one
t_env.
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 * FROM df1 JOIN df2 ON df1.a = df2.a")
>>> @pf.udf
... def add_one(x: int) -> int:
... return x + 1
>>> df = pf.from_dict({"a": [1, 2, 3]})
>>> incremented = pf.sql("SELECT add_one(a) FROM df")
>>> plus_one = pf.udf(lambda x: x + 1, return_dtype=pf.DataType.int64())
>>> incremented = pf.sql(
... "SELECT f(a) FROM source",
... auto_bind=False,
... source=df,
... f=plus_one,
... )
"""
binding_views, binding_functions = _resolve_explicit_bindings(bindings)
auto_view_candidates: Dict[str, DataFrame] = {}
auto_function_candidates: Dict[str, DataFrameFunctionWrapper] = {}
if auto_bind:
# `sql` is a plain function with no decorators, so the user's
# frame is exactly one f_back. If a decorator is ever added,
# this walk must be updated in lockstep — test_frame_back_level
# pins the invariant.
frame = inspect.currentframe()
caller: Optional[FrameType] = frame.f_back if frame is not None else None
try:
if caller is not None:
scope: Dict[str, object] = {}
scope.update(caller.f_globals)
scope.update(caller.f_locals)
for name, value in scope.items():
if isinstance(value, DataFrame):
auto_view_candidates[name] = value
if isinstance(value, DataFrameFunctionWrapper):
auto_function_candidates[name] = value
finally:
del frame, caller
# Tables/views and functions are separate SQL namespaces. Explicit
# bindings suppress auto-bind candidates only within the same lane.
explicit_view_binding_names = set(binding_views)
auto_view_candidates = {
name: value
for name, value in auto_view_candidates.items()
if name not in explicit_view_binding_names
}
# DataFrameFunctionWrapper does not carry a TableEnvironment. Function-only
# calls therefore use the configured/global environment, or create one if
# needed.
t_env = _resolve_table_environment(binding_views, auto_view_candidates)
if binding_functions:
binding_functions = _normalize_explicit_function_bindings(
t_env, binding_functions
)
explicit_function_binding_names = {
_normalize_function_name(name) for name in binding_functions
}
auto_function_candidates = {
name: value
for name, value in auto_function_candidates.items()
if (
_normalize_function_name(name)
not in explicit_function_binding_names
)
}
if binding_views:
temp_table_or_view_names = (
set(t_env.list_temporary_tables()) | set(t_env.list_temporary_views())
)
for name in binding_views:
_validate_identifier(t_env, name, kind="view")
if _temp_table_or_view_exists(name, temp_table_or_view_names):
raise ValueError(
f"binding '{name}' conflicts with an existing temporary "
f"table/view of the same name. Explicit DataFrame bindings "
f"can override permanent catalog tables/views, but not "
f"temporary tables/views. Use a different binding name, or "
f"drop the existing '{name}' first."
)
if binding_functions:
normalized_user_defined_function_names = _normalize_function_names(
t_env.list_user_defined_functions()
)
for name in binding_functions:
if _user_defined_function_exists(
name, normalized_user_defined_function_names
):
raise ValueError(
f"binding '{name}' conflicts with an existing user-defined "
f"function of the same name. Explicit function bindings can "
f"override built-in/system functions, but not temporary or "
f"catalog user-defined functions. Use a different binding "
f"name, or drop the existing '{name}' first."
)
auto_views: Dict[str, DataFrame] = {}
auto_functions: Dict[str, DataFrameFunctionWrapper] = {}
if auto_bind:
auto_views = _resolve_auto_bind_dataframes(
auto_view_candidates, binding_views, t_env
)
auto_functions = _resolve_auto_bind_functions(
auto_function_candidates, binding_functions, t_env
)
return _execute_with_temp_objects(
t_env,
auto_views,
binding_views,
auto_functions,
binding_functions,
query,
)