auroflow commented on code in PR #29029: URL: https://github.com/apache/flink/pull/29029#discussion_r3881556852
########## flink-python/pyflink/dataframe/udf.py: ########## @@ -0,0 +1,1036 @@ +################################################################################ +# 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. +################################################################################ + +"""User-defined scalar functions for the DataFrame API.""" + +import functools +import inspect +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Optional, + Tuple, + Type, + Union, + cast, + get_type_hints, + overload, +) + +from pyflink.common import Row +from pyflink.dataframe.datatype import DataType +from pyflink.table.expression import Expression +from pyflink.table.expressions import call as table_call +from pyflink.table.types import ArrayType, MapType, RowType +from pyflink.table.udf import ( + AsyncScalarFunction, + ScalarFunction, + UserDefinedFunction, + UserDefinedFunctionWrapper, + udf as table_udf, +) +from pyflink.util.api_stability_decorators import PublicEvolving + +__all__ = ["DataFrameUDFWrapper", "udf"] + +_UDFInput = Union[Callable[..., Any], ScalarFunction, AsyncScalarFunction, Type] +_DataTypeLike = Union[DataType, Type, str] + + +class _UDFUsage(Enum): + EXPRESSION = "expression" + MAP = "map" + MAP_BATCHES = "map_batches" + + +class _UDFSourceKind(Enum): + """How a resolved UDF source is initialized and invoked on a worker.""" + + DIRECT_CALLABLE = "direct_callable" + CALLABLE_INSTANCE = "callable_instance" + CALLABLE_CLASS = "callable_class" + SCALAR_FUNCTION_INSTANCE = "scalar_function_instance" + SCALAR_FUNCTION_CLASS = "scalar_function_class" + + +@dataclass(frozen=True) +class _ResolvedUDFSource: + """Callable metadata resolved once on the client and reused on workers.""" + + source: _UDFInput + kind: _UDFSourceKind + is_async: bool + ignored_hint_names: FrozenSet[str] = frozenset() + + @property + def inspection_target(self) -> Callable[..., Any]: + if self.kind is _UDFSourceKind.DIRECT_CALLABLE: + return _get_callable_inspection_target( + cast(Callable[..., Any], self.source) + ) + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: + return cast( + Union[ScalarFunction, AsyncScalarFunction], self.source + ).eval + if self.kind in ( + _UDFSourceKind.CALLABLE_CLASS, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + ): + hint_method, _ = _get_callable_class_hint_method( + cast(Type, self.source), + "eval" + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS + else "__call__", + ) + if hint_method is None: + raise RuntimeError("Resolved UDF class has no inspection target.") + return hint_method + return cast(Callable[..., Any], getattr(self.source, "__call__")) + + @property + def default_name(self) -> str: + return _default_udf_name(self.source) + + @property + def is_scalar_function(self) -> bool: + return self.kind in ( + _UDFSourceKind.SCALAR_FUNCTION_INSTANCE, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + ) + + @property + def constructs_on_worker(self) -> bool: + return self.kind in ( + _UDFSourceKind.CALLABLE_CLASS, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + ) + + def create_worker_source(self) -> _UDFInput: + if not self.constructs_on_worker: + return self.source + source_class = cast(Type, self.source) + source = source_class() + if self.is_scalar_function: + if not isinstance(source, (ScalarFunction, AsyncScalarFunction)): + raise TypeError( + f"Scalar UDF class '{source_class.__name__}' constructed an " + f"unsupported object of type '{type(source).__name__}'." + ) + elif not callable(source): + raise TypeError( + f"Callable class '{source_class.__name__}' constructed a non-callable " + f"object of type '{type(source).__name__}'." + ) + return cast(_UDFInput, source) + + def validate_deterministic( + self, declared: bool, worker_source: Optional[_UDFInput] = None + ) -> None: + source: Optional[_UDFInput] + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: + source = self.source + elif self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS: + source = worker_source + else: + source = None + if source is not None: + _validate_deterministic( + declared, + cast( + Union[ScalarFunction, AsyncScalarFunction], source + ).is_deterministic(), + ) + + def open_worker_source( + self, worker_source: _UDFInput, function_context: Any + ) -> None: + if self.is_scalar_function: + cast( + Union[ScalarFunction, AsyncScalarFunction], worker_source + ).open(function_context) + + def worker_invocation( + self, worker_source: _UDFInput + ) -> Callable[..., Any]: + if self.kind is _UDFSourceKind.DIRECT_CALLABLE: + return cast(Callable[..., Any], worker_source) + if self.is_scalar_function: + return cast( + Union[ScalarFunction, AsyncScalarFunction], worker_source + ).eval + return cast(Callable[..., Any], getattr(worker_source, "__call__")) + + def close_worker_source(self, worker_source: Optional[_UDFInput]) -> None: + if self.is_scalar_function and worker_source is not None: + cast( + Union[ScalarFunction, AsyncScalarFunction], worker_source + ).close() + + +@PublicEvolving() +class DataFrameUDFWrapper: + """ + A callable DataFrame scalar UDF declaration. + + Instances are created with :func:`udf` and can be called with DataFrame + expressions or Python literals to produce an expression. + + Example:: + + >>> import pyflink.dataframe as pf + >>> @pf.udf + ... def add_one(value: int) -> int: + ... return value + 1 + >>> expression = add_one(pf.col("value")) + + .. versionadded:: 2.4.0 + """ + + _source: _ResolvedUDFSource + _return_dtype: DataType + _deterministic: bool + _func_type: str + _cached_table_udf_wrapper: Optional[UserDefinedFunctionWrapper] + _frozen: bool + __name__: str + + def __init__( + self, + source: _ResolvedUDFSource, + return_dtype: DataType, + deterministic: bool, + name: str, + func_type: str, + ) -> None: + object.__setattr__(self, "_source", source) + object.__setattr__(self, "_return_dtype", return_dtype) + object.__setattr__(self, "_deterministic", deterministic) + object.__setattr__(self, "_func_type", func_type) + object.__setattr__(self, "_cached_table_udf_wrapper", None) + + declaration_metadata = _unwrap_partial(source.source) + functools.update_wrapper(self, declaration_metadata, updated=()) + object.__setattr__(self, "__name__", name) + object.__setattr__(self, "__wrapped__", source.source) Review Comment: Good catch! I fixed this issue by setting an explicit `__signature__` from the correct invocation target, so the result in your example is correct. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
