zhengruifeng commented on code in PR #58729: URL: https://github.com/apache/spark/pull/58729#discussion_r4013925817
########## 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: This exports an implementor-facing API from a non-underscore package, but external subclasses cannot populate the executor-local registry before `read_udfs` consults it: user code is deserialized only afterward. Please keep the package and symbols internal for this refactor; a supported extension point would need an explicit 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: This should describe the key-to-handler relationship and limit the claim to the incrementally migrated eval types. ```suggestion Each migrated eval type has an ``EvalTypeHandler`` subclass that declares its ``eval_type`` ``` ########## 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: This registers a class solely because it declares `eval_type`, even when `ABCMeta` still considers `run` abstract. `read_udfs` then fails when it instantiates that handler. Please reject abstract declarations here and add a focused case so an invalid handler fails during class definition instead of task startup. ########## 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: `ArrowStreamCoGroupSerializer.load_stream` eagerly materializes both sides as lists, so this annotation gives future handlers the wrong contract. Please cover one deserialized co-group in the category tests as well. ```suggestion CoGroupedBatch = tuple[list["pa.RecordBatch"], list["pa.RecordBatch"]] ``` -- 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]
