From f194bb244e328a825e4067f526db150002c2c80c Mon Sep 17 00:00:00 2001 From: Matthew Russo Date: Wed, 22 Jan 2025 20:19:30 -0500 Subject: [PATCH] rename optimize --> get_optimal_plan to clarify purpose a bit --- src/palimpzest/query/optimizer/optimizer.py | 3 +- .../query/optimizer/optimizer_strategy.py | 34 ++++++++++--------- src/palimpzest/query/optimizer/tasks.py | 2 +- 3 files changed, 20 insertions(+), 19 deletions(-) diff --git a/src/palimpzest/query/optimizer/optimizer.py b/src/palimpzest/query/optimizer/optimizer.py index 54c89855..ac95fbf0 100644 --- a/src/palimpzest/query/optimizer/optimizer.py +++ b/src/palimpzest/query/optimizer/optimizer.py @@ -452,5 +452,4 @@ def optimize(self, query_plan: Dataset, policy: Policy | None = None) -> list[Ph # search the optimization space by applying logical and physical transformations to the initial group tree self.search_optimization_space(final_group_id) - return self.strategy.optimize(self.groups, final_group_id, policy) - \ No newline at end of file + return self.strategy.get_optimal_plans(self.groups, final_group_id, policy) diff --git a/src/palimpzest/query/optimizer/optimizer_strategy.py b/src/palimpzest/query/optimizer/optimizer_strategy.py index 9c013abf..a1db6c66 100644 --- a/src/palimpzest/query/optimizer/optimizer_strategy.py +++ b/src/palimpzest/query/optimizer/optimizer_strategy.py @@ -1,6 +1,7 @@ +from __future__ import annotations + from abc import ABC, abstractmethod from enum import Enum -from typing import Type from palimpzest.policy import Policy from palimpzest.query.optimizer.plan import PhysicalPlan, SentinelPlan @@ -21,12 +22,12 @@ class OptimizationStrategyType(str, Enum): class OptimizationStrategy(ABC): @abstractmethod - def optimize(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: + def get_optimal_plan(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: """Strategy decides how to search through the groups for optimal plan(s)""" pass @classmethod - def get_strategy(cls, strategy_type: str) -> 'OptimizationStrategy': + def get_strategy(cls, strategy_type: str) -> OptimizationStrategy: """Factory method to create strategy instances""" return OptimizerStrategyRegistry.get_strategy(strategy_type) @@ -51,7 +52,7 @@ def _get_greedy_physical_plan(self, groups: dict, group_id: int) -> PhysicalPlan # add this operator to best physical plan and return return PhysicalPlan.from_ops_and_sub_plan([best_phys_expr.operator], input_best_phys_plan, best_phys_expr.plan_cost) - def optimize(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: + def get_optimal_plan(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: return [self._get_greedy_physical_plan(groups, final_group_id)] @@ -101,7 +102,7 @@ def _get_candidate_pareto_physical_plans(self, groups: dict, group_id: int, poli return pareto_optimal_plans - def optimize(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: + def get_optimal_plan(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: # compute all of the pareto optimal physical plans plans = self._get_candidate_pareto_physical_plans(groups, final_group_id, policy) @@ -150,12 +151,12 @@ def _get_sentinel_plan(self, groups: dict, group_id: int) -> SentinelPlan: # add this operator set to best physical plan and return return SentinelPlan.from_ops_and_sub_plan([phys_op_set], best_phys_subplan) - def optimize(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: + def get_optimal_plan(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: return [self._get_sentinel_plan(groups, final_group_id)] class NoOptimizationStrategy(OptimizationStrategy): - def optimize(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: + def get_optimal_plan(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: raise NotImplementedError("No optimization strategy selected") @@ -201,20 +202,21 @@ def _get_confidence_interval_optimal_plans(self, groups: dict, group_id: int) -> return best_plans - def optimize(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: + def get_optimal_plan(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: # TODO: fix this to properly handle multiple potential plans raise Exception("NotImplementedError") - plans = self.get_confidence_interval_optimal_plans(final_group_id) - + # plans = self._get_confidence_interval_optimal_plans(final_group_id) + + class AutoOptimizationStrategy(OptimizationStrategy): - def optimize(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: + def get_optimal_plan(self, groups: dict, final_group_id: int, policy: Policy) -> list[PhysicalPlan]: raise NotImplementedError("Auto optimization strategy not implemented") - + class OptimizerStrategyRegistry: """Registry to map strategy types to their implementations""" - - _strategies: dict[str, Type[OptimizationStrategy]] = { + + _strategies: dict[str, type[OptimizationStrategy]] = { OptimizationStrategyType.GREEDY.value: GreedyStrategy, OptimizationStrategyType.CONFIDENCE_INTERVAL.value: ConfidenceIntervalStrategy, OptimizationStrategyType.PARETO.value: ParetoStrategy, @@ -222,11 +224,11 @@ class OptimizerStrategyRegistry: OptimizationStrategyType.NONE.value: NoOptimizationStrategy, OptimizationStrategyType.AUTO.value: AutoOptimizationStrategy, } - + @classmethod def get_strategy(cls, strategy_type: str) -> OptimizationStrategy: """Get strategy instance by type""" strategy_class = cls._strategies.get(strategy_type) if not strategy_class: raise ValueError(f"Unknown optimization strategy: {strategy_type}") - return strategy_class() \ No newline at end of file + return strategy_class() diff --git a/src/palimpzest/query/optimizer/tasks.py b/src/palimpzest/query/optimizer/tasks.py index c6ce7289..0b96b574 100644 --- a/src/palimpzest/query/optimizer/tasks.py +++ b/src/palimpzest/query/optimizer/tasks.py @@ -2,10 +2,10 @@ from typing import Any -from palimpzest.query.optimizer.optimizer_strategy import OptimizationStrategyType from palimpzest.core.data.dataclasses import PlanCost from palimpzest.policy import Policy from palimpzest.query.optimizer.cost_model import BaseCostModel +from palimpzest.query.optimizer.optimizer_strategy import OptimizationStrategyType from palimpzest.query.optimizer.primitives import Expression, Group from palimpzest.query.optimizer.rules import ImplementationRule, Rule, TransformationRule