Yicong-Huang commented on code in PR #58729: URL: https://github.com/apache/spark/pull/58729#discussion_r4017060817
########## python/pyspark/sql/eval_handlers/__init__.py: ########## @@ -0,0 +1,46 @@ +# +# 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. +# + +"""Handlers for the Arrow/Pandas UDF eval types. + +Each eval type is an ``EvalTypeHandler`` subclass that declares its ``eval_type`` +and self-registers in ``EVAL_TYPE_HANDLERS`` via ``__init_subclass__``; the +worker's ``read_udfs`` dispatches on the registry. Base classes live in +``_base``; concrete handlers live in private per-family submodules (``_arrow``), +imported below so importing this package registers them. +""" + +# Re-export the concrete handlers from their private submodules so callers reach +# them from this package (or via EVAL_TYPE_HANDLERS), never from the ``_`` +# submodules. The import also registers them. +from pyspark.sql.eval_handlers._arrow import ArrowScalarUDFHandler +from pyspark.sql.eval_handlers._base import ( + EVAL_TYPE_HANDLERS, + BatchEvalTypeHandler, + CoGroupedEvalTypeHandler, + EvalTypeHandler, + GroupedEvalTypeHandler, +) + +__all__ = [ Review Comment: Agreed, this isn't a real cross-process extension point yet. I removed `__all__` and marked the package internal in the docstring so it isn't advertised as a supported API, and consumers now import from the private submodules. A supported extension API can come later with a proper driver-to-worker registration mechanism. ########## python/pyspark/sql/eval_handlers/__init__.py: ########## @@ -0,0 +1,46 @@ +# +# 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. +# + +"""Handlers for the Arrow/Pandas UDF eval types. + +Each eval type is an ``EvalTypeHandler`` subclass that declares its ``eval_type`` Review Comment: Done. Updated the docstring to scope it to migrated eval types and spell out the eval-type-to-handler mapping. ########## python/pyspark/sql/eval_handlers/_base.py: ########## @@ -0,0 +1,115 @@ +# +# 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. +# + +"""Base classes and registry for the eval type handlers. + +Kept in a leaf module (importing only ``_typing`` and serializers) so both the +package ``__init__`` and the concrete-handler submodules can import it without a +circular import. +""" + +from abc import ABCMeta, abstractmethod +from collections.abc import Iterator +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional + +from pyspark.serializers import Serializer +from pyspark.sql.eval_handlers._typing import ( + CoGroupedBatch, + GroupedBatch, + InputBatch, + OutputBatch, +) +from pyspark.sql.pandas.serializers import ( + ArrowStreamCoGroupSerializer, + ArrowStreamGroupSerializer, + ArrowStreamSerializer, +) + +if TYPE_CHECKING: + import pyarrow as pa # noqa: F401 # only in the batch category's forward-ref subscript + + from pyspark.worker_util import EvalConf, RunnerConf + +# eval type -> handler class, populated by EvalTypeHandler.__init_subclass__. +EVAL_TYPE_HANDLERS: "dict[int, type[EvalTypeHandler]]" = {} + + +class EvalTypeHandler(Generic[InputBatch, OutputBatch], metaclass=ABCMeta): + """Base class for the Arrow/Pandas UDF execution model. + + A handler declares the ``serializer`` for the input/output streams and + implements ``run``, which consumes the input stream and yields the output + stream. Concrete handlers subclass a typed category base + (``BatchEvalTypeHandler``, ``GroupedEvalTypeHandler``, + ``CoGroupedEvalTypeHandler``), which fixes the input type and serializer. + """ + + # PythonEvalType this handler serves; None on the abstract bases. + eval_type: ClassVar[Optional[int]] = None + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + # Register concrete handlers, i.e. those that declare an eval type. + eval_type = cls.__dict__.get("eval_type") Review Comment: Good catch. Registration now happens in a metaclass that runs after `ABCMeta` computes `__abstractmethods__`, so a class declaring an `eval_type` while leaving `run`/`serializer` abstract is rejected at definition time. Added a focused test. ########## python/pyspark/sql/eval_handlers/_typing.py: ########## @@ -0,0 +1,37 @@ +# +# 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. +# + +"""Type aliases and variables for the eval type handlers. + +The Arrow element types are forward refs so pyarrow stays a type-checking-only +import at runtime, as elsewhere in the worker. +""" + +from collections.abc import Iterator +from typing import TYPE_CHECKING, TypeVar + +if TYPE_CHECKING: + import pyarrow as pa + +# Input stream element type for the grouped categories (the batch category's +# element is a plain ``pa.RecordBatch``). +GroupedBatch = Iterator["pa.RecordBatch"] # one group of batches +CoGroupedBatch = tuple[Iterator["pa.RecordBatch"], Iterator["pa.RecordBatch"]] # a co-group pair Review Comment: Right, the co-group serializer materializes each side as a list. Fixed the alias to `tuple[list[...], list[...]]` and added a category test that deserializes one co-group and asserts both sides are lists. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
