Expose logical rules via package exports

Signed-off-by: 400Ping <jiekai.chang326@gmail.com>

Author: 400Ping

python/ray/data/_internal/logical/optimizers.py

@@ -8,19 +8,17 @@
     Plan,
     Rule,
 )
-from ray.data._internal.logical.rules.combine_shuffles import CombineShuffles
-from ray.data._internal.logical.rules.configure_map_task_memory import (
+from ray.data._internal.logical.rules import (
+    CombineShuffles,
     ConfigureMapTaskMemoryUsingOutputSize,
-)
-from ray.data._internal.logical.rules.inherit_batch_format import InheritBatchFormatRule
-from ray.data._internal.logical.rules.inherit_target_max_block_size import (
+    FuseOperators,
+    InheritBatchFormatRule,
     InheritTargetMaxBlockSizeRule,
+    LimitPushdownRule,
+    PredicatePushdown,
+    ProjectionPushdown,
+    SetReadParallelismRule,
 )
-from ray.data._internal.logical.rules.limit_pushdown import LimitPushdownRule
-from ray.data._internal.logical.rules.operator_fusion import FuseOperators
-from ray.data._internal.logical.rules.predicate_pushdown import PredicatePushdown
-from ray.data._internal.logical.rules.projection_pushdown import ProjectionPushdown
-from ray.data._internal.logical.rules.set_read_parallelism import SetReadParallelismRule
 from ray.util.annotations import DeveloperAPI
 
 _LOGICAL_RULESET = Ruleset(

python/ray/data/_internal/logical/rules/__init__.py

@@ -1,3 +1,22 @@
-from ray.data._internal.logical.rules.operator_fusion import FuseOperators
+"""Expose rule classes in ray.data._internal.logical.rules.
 
-__all__ = ["FuseOperators"]
+This module dynamically imports each submodule in the package and re-exports
+names listed in the submodule's __all__ attribute. This allows callers to
+import rules directly from ray.data._internal.logical.rules.
+"""
+
+from importlib import import_module
+from pkgutil import iter_modules
+
+__all__ = []
+
+for _loader, _module_name, _is_pkg in iter_modules(__path__):
+    if _module_name.startswith("_"):
+        continue
+    _module = import_module(f"{__name__}.{_module_name}")
+    _exported_names = getattr(_module, "__all__", [])
+    for _name in _exported_names:
+        globals()[_name] = getattr(_module, _name)
+    __all__.extend(_exported_names)
+
+__all__.sort()

python/ray/data/_internal/logical/rules/combine_shuffles.py

@@ -4,12 +4,16 @@
     Plan,
     Rule,
 )
-from ray.data._internal.logical.operators.all_to_all_operator import (
+from ray.data._internal.logical.operators import (
     Aggregate,
     Repartition,
     Sort,
+    StreamingRepartition,
 )
-from ray.data._internal.logical.operators.map_operator import StreamingRepartition
+
+__all__ = [
+    "CombineShuffles",
+]
 
 
 class CombineShuffles(Rule):

python/ray/data/_internal/logical/rules/configure_map_task_memory.py

@@ -7,6 +7,11 @@
 from ray.data._internal.logical.interfaces.physical_plan import PhysicalPlan
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
+__all__ = [
+    "ConfigureMapTaskMemoryRule",
+    "ConfigureMapTaskMemoryUsingOutputSize",
+]
+
 
 class ConfigureMapTaskMemoryRule(Rule, abc.ABC):
     def apply(self, plan: PhysicalPlan) -> PhysicalPlan:

python/ray/data/_internal/logical/rules/inherit_batch_format.py

@@ -2,8 +2,11 @@
 from typing import Iterable
 
 from ray.data._internal.logical.interfaces import LogicalOperator, LogicalPlan, Rule
-from ray.data._internal.logical.operators.all_to_all_operator import AbstractAllToAll
-from ray.data._internal.logical.operators.map_operator import MapBatches
+from ray.data._internal.logical.operators import AbstractAllToAll, MapBatches
+
+__all__ = [
+    "InheritBatchFormatRule",
+]
 
 
 class InheritBatchFormatRule(Rule):

python/ray/data/_internal/logical/rules/inherit_target_max_block_size.py

@@ -3,6 +3,10 @@
 from ray.data._internal.execution.interfaces import PhysicalOperator
 from ray.data._internal.logical.interfaces import PhysicalPlan, Rule
 
+__all__ = [
+    "InheritTargetMaxBlockSizeRule",
+]
+
 
 class InheritTargetMaxBlockSizeRule(Rule):
     """For each op that has overridden the default target max block size,

python/ray/data/_internal/logical/rules/limit_pushdown.py

@@ -3,13 +3,18 @@
 from typing import List
 
 from ray.data._internal.logical.interfaces import LogicalOperator, LogicalPlan, Rule
-from ray.data._internal.logical.operators.map_operator import AbstractMap
-from ray.data._internal.logical.operators.n_ary_operator import Union
-from ray.data._internal.logical.operators.one_to_one_operator import (
+from ray.data._internal.logical.operators import (
+    AbstractMap,
     AbstractOneToOne,
     Limit,
+    Union,
 )
 
+__all__ = [
+    "LimitPushdownRule",
+]
+
+
 logger = logging.getLogger(__name__)
 
 

python/ray/data/_internal/logical/rules/operator_fusion.py

@@ -26,20 +26,24 @@
     TaskPoolMapOperator,
 )
 from ray.data._internal.logical.interfaces import PhysicalPlan, Rule
-from ray.data._internal.logical.operators.all_to_all_operator import (
+from ray.data._internal.logical.operators import (
     AbstractAllToAll,
-    RandomShuffle,
-    Repartition,
-)
-from ray.data._internal.logical.operators.map_operator import (
     AbstractMap,
     AbstractUDFMap,
     MapBatches,
+    RandomShuffle,
+    Repartition,
     StreamingRepartition,
 )
 from ray.data._internal.streaming_repartition import StreamingRepartitionRefBundler
 from ray.util.annotations import DeveloperAPI
 
+__all__ = [
+    "FuseOperators",
+    "are_remote_args_compatible",
+]
+
+
 # Scheduling strategy can be inherited from upstream operator if not specified.
 INHERITABLE_REMOTE_ARGS = ["scheduling_strategy"]
 

python/ray/data/_internal/logical/rules/predicate_pushdown.py

@@ -9,12 +9,16 @@
     PredicatePassThroughBehavior,
     Rule,
 )
-from ray.data._internal.logical.operators.map_operator import Filter, Project
+from ray.data._internal.logical.operators import Filter, Project
 from ray.data._internal.planner.plan_expression.expression_visitors import (
     _ColumnSubstitutionVisitor,
 )
 from ray.data.expressions import Expr, col
 
+__all__ = [
+    "PredicatePushdown",
+]
+
 
 class PredicatePushdown(Rule):
     """Pushes down predicates across the graph.

python/ray/data/_internal/logical/rules/projection_pushdown.py

@@ -6,7 +6,7 @@
     LogicalPlan,
     Rule,
 )
-from ray.data._internal.logical.operators.map_operator import Project
+from ray.data._internal.logical.operators import Project
 from ray.data._internal.planner.plan_expression.expression_visitors import (
     _ColumnReferenceCollector,
     _ColumnSubstitutionVisitor,
@@ -19,6 +19,10 @@
     StarExpr,
 )
 
+__all__ = [
+    "ProjectionPushdown",
+]
+
 
 def _collect_referenced_columns(exprs: List[Expr]) -> Optional[List[str]]:
     """