palimpzest 0.7.20__py3-none-any.whl → 0.8.0__py3-none-any.whl

palimpzest/query/processor/sentinel_processor.py

@@ -1,88 +0,0 @@
-import logging
-
-from palimpzest.core.data.dataclasses import ExecutionStats, SentinelPlanStats
-from palimpzest.core.elements.records import DataRecordCollection
-from palimpzest.query.optimizer.cost_model import SampleBasedCostModel
-from palimpzest.query.optimizer.optimizer_strategy_type import OptimizationStrategyType
-from palimpzest.query.optimizer.plan import SentinelPlan
-from palimpzest.query.processor.query_processor import QueryProcessor
-
-logger = logging.getLogger(__name__)
-
-class SentinelQueryProcessor(QueryProcessor):
-
-    def _generate_sample_observations(self, sentinel_plan: SentinelPlan) -> SentinelPlanStats:
-        """
-        This function is responsible for generating sample observation data which can be
-        consumed by the CostModel.
-
-        To accomplish this, we construct a special sentinel plan using the Optimizer which is
-        capable of executing any valid physical implementation of a Filter or Convert operator
-        on each record.
-        """
-        # if we're using validation data, get the set of expected output records
-        expected_outputs = {}
-        for source_idx in range(len(self.val_datasource)):
-            expected_output = self.val_datasource[source_idx]
-            expected_outputs[source_idx] = expected_output
-
-        # execute sentinel plan; returns sentinel_plan_stats
-        return self.sentinel_execution_strategy.execute_sentinel_plan(sentinel_plan, expected_outputs)
-
-    def _create_sentinel_plan(self) -> SentinelPlan:
-        """
-        Generates and returns a SentinelPlan for the given dataset.
-        """
-        # create a new optimizer and update its strategy to SENTINEL
-        optimizer = self.optimizer.deepcopy_clean()
-        optimizer.update_strategy(OptimizationStrategyType.SENTINEL)
-
-        # create copy of dataset, but change its data source to the validation data source
-        dataset = self.dataset.copy()
-        dataset._set_data_source(self.val_datasource)
-
-        # get the sentinel plan for the given dataset
-        sentinel_plans = optimizer.optimize(dataset)
-        sentinel_plan = sentinel_plans[0]
-
-        return sentinel_plan
-
-    def execute(self) -> DataRecordCollection:
-        # for now, enforce that we are using validation data; we can relax this after paper submission
-        if self.val_datasource is None:
-            raise Exception("Make sure you are using validation data with SentinelQueryProcessor")
-        logger.info(f"Executing {self.__class__.__name__}")
-
-        # create execution stats
-        execution_stats = ExecutionStats(execution_id=self.execution_id())
-        execution_stats.start()
-
-        # create sentinel plan
-        sentinel_plan = self._create_sentinel_plan()
-
-        # generate sample execution data
-        sentinel_plan_stats = self._generate_sample_observations(sentinel_plan)
-
-        # update the execution stats to account for the work done in optimization
-        execution_stats.add_plan_stats(sentinel_plan_stats)
-        execution_stats.finish_optimization()
-
-        # (re-)initialize the optimizer
-        optimizer = self.optimizer.deepcopy_clean()
-
-        # construct the CostModel with any sample execution data we've gathered
-        cost_model = SampleBasedCostModel(sentinel_plan_stats, self.verbose)
-        optimizer.update_cost_model(cost_model)
-
-        # execute plan(s) according to the optimization strategy
-        records, plan_stats = self._execute_best_plan(self.dataset, optimizer)
-
-        # update the execution stats to account for the work to execute the final plan
-        execution_stats.add_plan_stats(plan_stats)
-        execution_stats.finish()
-
-        # construct and return the DataRecordCollection
-        result = DataRecordCollection(records, execution_stats=execution_stats)
-        logger.info("Done executing SentinelQueryProcessor")
-
-        return result

palimpzest/query/processor/streaming_processor.py

@@ -1,149 +0,0 @@
-import logging
-import time
-
-from palimpzest.core.data.dataclasses import PlanStats
-from palimpzest.core.elements.records import DataRecordCollection
-from palimpzest.query.operators.aggregate import AggregateOp
-from palimpzest.query.operators.filter import FilterOp
-from palimpzest.query.operators.limit import LimitScanOp
-from palimpzest.query.operators.scan import ScanPhysicalOp
-from palimpzest.query.optimizer.plan import PhysicalPlan
-from palimpzest.query.processor.query_processor import QueryProcessor
-from palimpzest.sets import Dataset
-
-logger = logging.getLogger(__name__)
-
-class StreamingQueryProcessor(QueryProcessor):
-    """This class can be used for a streaming, record-based execution.
-    Results are returned as an iterable that can be consumed by the caller."""
-
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-        self._plan: PhysicalPlan | None = None
-        self._plan_stats: PlanStats | None = None
-        self.last_record = False
-        self.current_scan_idx: int = 0
-        self.plan_generated: bool = False
-        self.records_count: int = 0
-        logger.info("Initialized StreamingQueryProcessor")
-
-    @property
-    def plan(self) -> PhysicalPlan:
-        if self._plan is None:
-            raise Exception("Plan has not been generated yet.")
-        return self._plan
-
-    @plan.setter
-    def plan(self, plan: PhysicalPlan):
-        self._plan = plan
-
-    @property
-    def plan_stats(self) -> PlanStats:
-        if self._plan_stats is None:
-            raise Exception("Plan stats have not been generated yet.")
-        return self._plan_stats
-
-    @plan_stats.setter
-    def plan_stats(self, plan_stats: PlanStats):
-        self._plan_stats = plan_stats
-
-    def generate_plan(self, dataset: Dataset):
-        # self.clear_cached_examples()
-        start_time = time.time()
-
-        # check that the plan does not contain any aggregation operators
-        for op in self.plan.operators:
-            if isinstance(op, AggregateOp):
-                raise Exception("You cannot have a Streaming Execution if there is an Aggregation Operator")
-
-        # TODO: Do we need to re-initialize the optimizer here? 
-        # Effectively always use the optimal strategy   
-        optimizer = self.optimizer.deepcopy_clean()
-        plans = optimizer.optimize(dataset)
-        self.plan = plans[0]
-        self.plan_stats = PlanStats.from_plan(self.plan)
-        self.plan_stats.start()
-        logger.info(f"Time for planning: {time.time() - start_time:.2f} seconds")
-        self.plan_generated = True
-        logger.info(f"Generated plan:\n{self.plan}")
-        return self.plan
-
-    def execute(self):
-        logger.info("Executing StreamingQueryProcessor")
-        # Always delete cache
-        if not self.plan_generated:
-            self.generate_plan(self.dataset)
-
-        # if dry_run:
-        #     yield [], self.plan, self.plan_stats
-        #     return
-
-        input_records = self.get_input_records()
-        for idx, record in enumerate(input_records):
-            # print("Iteration number: ", idx+1, "out of", len(input_records))
-            output_records = self.execute_opstream(self.plan, record)
-            if idx == len(input_records) - 1:
-                # finalize plan stats
-                self.plan_stats.finish()
-            self.plan_stats.plan_str = str(self.plan)
-            yield DataRecordCollection(output_records, plan_stats=self.plan_stats)
-
-        logger.info("Done executing StreamingQueryProcessor")
-
-
-    def get_input_records(self):
-        scan_operator = self.plan.operators[0]
-        assert isinstance(scan_operator, ScanPhysicalOp), "First operator in physical plan must be a ScanPhysicalOp"
-        datareader = scan_operator.datareader
-        if not datareader:
-            raise Exception("DataReader not found")
-        datareader_len = len(datareader)
-
-        input_records = []
-        record_op_stats = []
-        for source_idx in range(datareader_len):
-            record_set = scan_operator(source_idx)
-            input_records += record_set.data_records
-            record_op_stats += record_set.record_op_stats
-
-        self.plan_stats.add_record_op_stats(record_op_stats)
-
-        return input_records
-
-    def execute_opstream(self, plan, record):
-        # initialize list of output records and intermediate variables
-        input_records = [record]
-        record_op_stats_lst = []
-
-        for operator in plan.operators:
-            # TODO: this being defined in the for loop potentially makes the return
-            # unbounded if plan.operators is empty. This should be defined outside the loop
-            # and the loop refactored to account for not redeclaring this for each operator
-            output_records = []
-
-            if isinstance(operator, ScanPhysicalOp):
-                continue
-            # only invoke aggregate operator(s) once there are no more source records and all
-            # upstream operators' processing queues are empty
-            # elif isinstance(operator, AggregateOp):
-            # output_records, record_op_stats_lst = operator(candidates=input_records)
-            elif isinstance(operator, LimitScanOp):
-                if self.records_count >= operator.limit:
-                    break
-            else:
-                for r in input_records:
-                    record_set = operator(r)
-                    output_records += record_set.data_records
-                    record_op_stats_lst += record_set.record_op_stats
-
-                if isinstance(operator, FilterOp):
-                    # delete all records that did not pass the filter
-                    output_records = [r for r in output_records if r.passed_operator]
-                    if not output_records:
-                        break
-
-            self.plan_stats.add_record_op_stats(record_op_stats_lst)
-            input_records = output_records
-            self.records_count += len(output_records)
-
-        return output_records

palimpzest/sets.py

@@ -1,405 +0,0 @@
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Callable
-
-import pandas as pd
-from chromadb.api.models.Collection import Collection
-
-from palimpzest.constants import AggFunc, Cardinality
-from palimpzest.core.data.datareaders import DataReader
-from palimpzest.core.elements.filters import Filter
-from palimpzest.core.elements.groupbysig import GroupBySig
-from palimpzest.core.lib.schemas import Number, Schema
-from palimpzest.policy import construct_policy_from_kwargs
-from palimpzest.query.processor.config import QueryProcessorConfig
-from palimpzest.utils.datareader_helpers import get_local_datareader
-from palimpzest.utils.hash_helpers import hash_for_serialized_dict
-
-
-#####################################################
-#
-#####################################################
-class Set:
-    """
-    """
-
-    def __init__(
-        self,
-        source: Set | DataReader,
-        schema: Schema,
-        desc: str | None = None,
-        filter: Filter | None = None,
-        udf: Callable | None = None,
-        agg_func: AggFunc | None = None,
-        group_by: GroupBySig | None = None,
-        project_cols: list[str] | None = None,
-        index: Collection | None = None,
-        search_func: Callable | None = None,
-        search_attr: str | None = None,
-        output_attrs: list[dict] | None = None,
-        k: int | None = None,  # TODO: disambiguate `k` to be something like `retrieve_k`
-        limit: int | None = None,
-        cardinality: Cardinality = Cardinality.ONE_TO_ONE,
-        depends_on: list[str] | None = None,
-        cache: bool = False,
-    ):
-        self._schema = schema
-        self._source = source
-        self._desc = desc
-        self._filter = filter
-        self._udf = udf
-        self._agg_func = agg_func
-        self._group_by = group_by
-        self._project_cols = None if project_cols is None else sorted(project_cols)
-        self._index = index
-        self._search_func = search_func
-        self._search_attr = search_attr
-        self._output_attrs = output_attrs
-        self._k = k
-        self._limit = limit
-        self._cardinality = cardinality
-        self._depends_on = [] if depends_on is None else sorted(depends_on)
-        self._cache = cache
-
-    @property
-    def schema(self) -> Schema:
-        return self._schema
-
-    def _set_data_source(self, source: DataReader):
-        if isinstance(self._source, Set):
-            self._source._set_data_source(source)
-        else:
-            self._source = source
-
-    def serialize(self):
-        # NOTE: I needed to remove depends_on from the serialization dictionary because
-        # the optimizer changes the name of the depends_on fields to be their "full" name.
-        # This created an issue with the node.universal_identifier() not being consistent
-        # after changing the field to its full name.
-        d = {
-            "schema": self.schema.json_schema(),
-            "source": self._source.serialize(),
-            "desc": repr(self._desc),
-            "filter": None if self._filter is None else self._filter.serialize(),
-            "udf": None if self._udf is None else self._udf.__name__,
-            "agg_func": None if self._agg_func is None else self._agg_func.value,
-            "cardinality": self._cardinality,
-            "limit": self._limit,
-            "group_by": None if self._group_by is None else self._group_by.serialize(),
-            "project_cols": None if self._project_cols is None else self._project_cols,
-            "index": None if self._index is None else self._index.__class__.__name__,
-            "search_func": None if self._search_func is None else self._search_func.__name__,
-            "search_attr": self._search_attr,
-            "output_attrs": None if self._output_attrs is None else str(self._output_attrs),
-            "k": self._k,
-        }
-
-        return d
-
-    def universal_identifier(self):
-        """Return a unique identifier for this Set."""
-        return hash_for_serialized_dict(self.serialize())
-
-    def json_schema(self):
-        """Return the JSON schema for this Set."""
-        return self.schema.json_schema()
-
-
-class Dataset(Set):
-    """
-    A Dataset is the intended abstraction for programmers to interact with when writing PZ programs.
-
-    Users instantiate a Dataset by specifying a `source` that either points to a DataReader
-    or an existing Dataset. Users can then perform computations on the Dataset in a lazy fashion
-    by leveraging functions such as `filter`, `sem_filter`, `sem_add_columns`, `aggregate`, etc.
-    Underneath the hood, each of these operations creates a new Dataset. As a result, the Dataset
-    defines a lineage of computation.
-    """
-
-    def __init__(
-        self,
-        source: str | Path | list | pd.DataFrame | DataReader | Dataset,
-        schema: Schema | None = None,
-        *args,
-        **kwargs,
-    ) -> None:
-        # NOTE: this function currently assumes that DataReader will always be provided with a schema;
-        #       we will relax this assumption in a subsequent PR
-        # convert source into a DataReader
-        updated_source = get_local_datareader(source, **kwargs) if isinstance(source, (str, Path, list, pd.DataFrame)) else source
-
-        # get the schema
-        schema = updated_source.schema if schema is None else schema
-
-        # intialize class
-        super().__init__(updated_source, schema, *args, **kwargs)
-
-    def copy(self):
-        return Dataset(
-            source=self._source.copy() if isinstance(self._source, Set) else self._source,
-            schema=self._schema,
-            desc=self._desc,
-            filter=self._filter,
-            udf=self._udf,
-            agg_func=self._agg_func,
-            group_by=self._group_by,
-            project_cols=self._project_cols,
-            index=self._index,
-            search_func=self._search_func,
-            search_attr=self._search_attr,
-            output_attrs=self._output_attrs,
-            k=self._k,
-            limit=self._limit,
-            cardinality=self._cardinality,
-            depends_on=self._depends_on,
-            cache=self._cache,
-        )
-
-    def filter(
-        self,
-        _filter: Callable,
-        depends_on: str | list[str] | None = None,
-    ) -> Dataset:
-        """Add a user defined function as a filter to the Set. This filter will possibly restrict the items that are returned later."""
-        f = None
-        if callable(_filter):
-            f = Filter(filter_fn=_filter)
-        else:
-            error_str = f"Only support callable for filter, currently got {type(_filter)}"
-            if isinstance(_filter, str):
-                error_str += ". Consider using sem_filter() for semantic filters."
-            raise Exception(error_str)
-
-        if isinstance(depends_on, str):
-            depends_on = [depends_on]
-
-        return Dataset(
-            source=self,
-            schema=self.schema,
-            filter=f,
-            depends_on=depends_on,
-            cache=self._cache,
-        )
-
-    def sem_filter(
-        self,
-        _filter: str,
-        depends_on: str | list[str] | None = None,
-    ) -> Dataset:
-        """Add a natural language description of a filter to the Set. This filter will possibly restrict the items that are returned later."""
-        f = None
-        if isinstance(_filter, str):
-            f = Filter(_filter)
-        else:
-            raise Exception("sem_filter() only supports `str` input for _filter.", type(_filter))
-
-        if isinstance(depends_on, str):
-            depends_on = [depends_on]
-
-        return Dataset(
-            source=self,
-            schema=self.schema,
-            filter=f,
-            depends_on=depends_on,
-            cache=self._cache,
-        )
-
-    def sem_add_columns(self, cols: list[dict] | type[Schema],
-                        cardinality: Cardinality = Cardinality.ONE_TO_ONE,
-                        depends_on: str | list[str] | None = None,
-                        desc: str = "Add new columns via semantic reasoning") -> Dataset:
-        """
-        Add new columns by specifying the column names, descriptions, and types.
-        The column will be computed during the execution of the Dataset.
-        Example:
-            sem_add_columns(
-                [{'name': 'greeting', 'desc': 'The greeting message', 'type': str},
-                 {'name': 'age', 'desc': 'The age of the person', 'type': int},
-                 {'name': 'full_name', 'desc': 'The name of the person', 'type': str}]
-            )
-        """
-        if isinstance(depends_on, str):
-            depends_on = [depends_on]
-
-        new_output_schema = None
-        if isinstance(cols, list):
-            new_output_schema = self.schema.add_fields(cols)
-        elif issubclass(cols, Schema):
-            new_output_schema = self.schema.union(cols)
-        else:
-            raise ValueError("`cols` must be a list of dictionaries or a Schema.")
-
-        return Dataset(
-            source=self,
-            schema=new_output_schema,
-            udf=None,
-            cardinality=cardinality,
-            depends_on=depends_on,
-            desc=desc,
-            cache=self._cache,
-        )
-
-    def add_columns(self, udf: Callable,
-                    cols: list[dict] | type[Schema],
-                    cardinality: Cardinality = Cardinality.ONE_TO_ONE,
-                    depends_on: str | list[str] | None = None,
-                    desc: str = "Add new columns via UDF") -> Dataset:
-        """
-        Add new columns by specifying UDFs.
-
-        Examples:
-            add_columns(
-                udf=compute_personal_greeting,
-                cols=[
-                    {'name': 'greeting', 'desc': 'The greeting message', 'type': str},
-                    {'name': 'age', 'desc': 'The age of the person', 'type': int},
-                    {'name': 'full_name', 'desc': 'The name of the person', 'type': str},
-                ]
-            )
-        """
-        if udf is None or cols is None:
-            raise ValueError("`udf` and `cols` must be provided for add_columns.")
-
-        if isinstance(depends_on, str):
-            depends_on = [depends_on]
-
-        new_output_schema = None
-        if isinstance(cols, list):
-            updated_cols = []
-            for col_dict in cols:
-                assert isinstance(col_dict, dict), "each entry in `cols` must be a dictionary"
-                assert "name" in col_dict, "each type must contain a 'name' key specifying the column name"
-                assert "type" in col_dict, "each type must contain a 'type' key specifying the column type"
-                col_dict["desc"] = col_dict.get("desc", "New column: " + col_dict["name"])
-                updated_cols.append(col_dict)
-            new_output_schema = self.schema.add_fields(updated_cols)
-
-        elif issubclass(cols, Schema):
-            new_output_schema = self.schema.union(cols)
-
-        else:
-            raise ValueError("`cols` must be a list of dictionaries or a Schema.")
-
-        return Dataset(
-            source=self,
-            schema=new_output_schema,
-            udf=udf,
-            cardinality=cardinality,
-            desc=desc,
-            depends_on=depends_on,
-            cache=self._cache,
-        )
-
-    def map(self, udf: Callable) -> Dataset:
-        """
-        Apply a UDF map function.
-
-        Examples:
-            map(udf=clean_column_values)
-        """
-        if udf is None:
-            raise ValueError("`udf` must be provided for map.")
-
-        return Dataset(
-            source=self,
-            schema=self.schema,
-            udf=udf,
-            cache=self._cache,
-        )
-
-    def count(self) -> Dataset:
-        """Apply a count aggregation to this set"""
-        return Dataset(
-            source=self,
-            schema=Number,
-            desc="Count results",
-            agg_func=AggFunc.COUNT,
-            cache=self._cache,
-        )
-
-    def average(self) -> Dataset:
-        """Apply an average aggregation to this set"""
-        return Dataset(
-            source=self,
-            schema=Number,
-            desc="Average results",
-            agg_func=AggFunc.AVERAGE,
-            cache=self._cache,
-        )
-
-    def groupby(self, groupby: GroupBySig) -> Dataset:
-        return Dataset(
-            source=self,
-            schema=groupby.output_schema(),
-            desc="Group By",
-            group_by=groupby,
-            cache=self._cache,
-        )
-
-    def retrieve(
-        self,
-        index: Collection,
-        search_attr: str,
-        output_attrs: list[dict] | type[Schema],
-        search_func: Callable | None = None,
-        k: int = -1,
-    ) -> Dataset:
-        """
-        Retrieve the top-k nearest neighbors of the value of the `search_attr` from the `index` and
-        use these results to construct the `output_attrs` field(s).
-        """
-        new_output_schema = None
-        if isinstance(output_attrs, list):
-            new_output_schema = self.schema.add_fields(output_attrs)
-        elif issubclass(output_attrs, Schema):
-            new_output_schema = self.schema.union(output_attrs)
-        else:
-            raise ValueError("`cols` must be a list of dictionaries or a Schema.")
-
-        # TODO: revisit once we can think through abstraction(s)
-        # # construct the PZIndex from the user-provided index
-        # index = index_factory(index)
-
-        return Dataset(
-            source=self,
-            schema=new_output_schema,
-            desc="Retrieve",
-            index=index,
-            search_func=search_func,
-            search_attr=search_attr,
-            output_attrs=output_attrs,
-            k=k,
-            cache=self._cache,
-        )
-
-    def limit(self, n: int) -> Dataset:
-        """Limit the set size to no more than n rows"""
-        return Dataset(
-            source=self,
-            schema=self.schema,
-            desc="LIMIT " + str(n),
-            limit=n,
-            cache=self._cache,
-        )
-
-    def project(self, project_cols: list[str] | str) -> Dataset:
-        """Project the Set to only include the specified columns."""
-        return Dataset(
-            source=self,
-            schema=self.schema.project(project_cols),
-            project_cols=project_cols if isinstance(project_cols, list) else [project_cols],
-            cache=self._cache,
-        )
-
-    def run(self, config: QueryProcessorConfig | None = None, **kwargs):
-        """Invoke the QueryProcessor to execute the query. `kwargs` will be applied to the QueryProcessorConfig."""
-        # TODO: this import currently needs to be here to avoid a circular import; we should fix this in a subsequent PR
-        from palimpzest.query.processor.query_processor_factory import QueryProcessorFactory
-
-        # as syntactic sugar, we will allow some keyword arguments to parameterize our policies
-        policy = construct_policy_from_kwargs(**kwargs)
-        if policy is not None:
-            kwargs["policy"] = policy
-
-        return QueryProcessorFactory.create_and_run_processor(self, config, **kwargs)