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

palimpzest/query/operators/physical.py

@@ -2,9 +2,10 @@ from __future__ import annotations
 
 import json
 
-from palimpzest.core.data.dataclasses import OperatorCostEstimates
+from pydantic import BaseModel
+
 from palimpzest.core.elements.records import DataRecord, DataRecordSet
-from palimpzest.core.lib.schemas import Schema
+from palimpzest.core.models import OperatorCostEstimates
 from palimpzest.utils.hash_helpers import hash_for_id
 
 
@@ -17,12 +18,13 @@ class PhysicalOperator:
 
     def __init__(
         self,
-        output_schema: Schema,
-        input_schema: Schema | None = None,
+        output_schema: BaseModel,
+        input_schema: BaseModel | None = None,
         depends_on: list[str] | None = None,
         logical_op_id: str | None = None,
+        unique_logical_op_id: str | None = None,
         logical_op_name: str | None = None,
-        target_cache_id: str | None = None,
+        api_base: str | None = None,
         verbose: bool = False,
         *args,
         **kwargs,
@@ -31,16 +33,17 @@ class PhysicalOperator:
         self.input_schema = input_schema
         self.depends_on = depends_on if depends_on is None else sorted(depends_on)
         self.logical_op_id = logical_op_id
+        self.unique_logical_op_id = unique_logical_op_id
         self.logical_op_name = logical_op_name
-        self.target_cache_id = target_cache_id
+        self.api_base = api_base
         self.verbose = verbose
         self.op_id = None
 
         # compute the fields generated by this physical operator
-        input_field_names = self.input_schema.field_names() if self.input_schema is not None else []
+        input_field_names = list(self.input_schema.model_fields) if self.input_schema is not None else []
         self.generated_fields = sorted([
             field_name
-            for field_name in self.output_schema.field_names()
+            for field_name in self.output_schema.model_fields
             if field_name not in input_field_names
         ])
 
@@ -50,16 +53,18 @@ class PhysicalOperator:
         self.__class__.__hash__ = PhysicalOperator.__hash__
 
     def __str__(self):
-        op = f"{self.input_schema.class_name()} -> {self.op_name()} -> {self.output_schema.class_name()}\n"
-        op += f"    ({', '.join(self.input_schema.field_names())[:30]}) "
-        op += f"-> ({', '.join(self.output_schema.field_names())[:30]})\n"
+        op = f"{self.input_schema.__name__} -> {self.op_name()} -> {self.output_schema.__name__}\n"
+        op += f"    ({', '.join(sorted(self.input_schema.model_fields))[:30]}) "
+        op += f"-> ({', '.join(sorted(self.output_schema.model_fields))[:30]})\n"
         if getattr(self, "model", None):
             op += f"    Model: {self.model}\n"
         return op
 
+    # def __eq__(self, other) -> bool:
+    #     all_op_params_match = all(value == getattr(other, key) for key, value in self.get_op_params().items())
+    #     return isinstance(other, self.__class__) and all_op_params_match
     def __eq__(self, other) -> bool:
-        all_op_params_match = all(value == getattr(other, key) for key, value in self.get_op_params().items())
-        return isinstance(other, self.__class__) and all_op_params_match
+        return isinstance(other, self.__class__) and self.get_full_op_id() == other.get_full_op_id()
 
     def copy(self) -> PhysicalOperator:
         return self.__class__(**self.get_op_params())
@@ -94,8 +99,9 @@ class PhysicalOperator:
             "input_schema": self.input_schema,
             "depends_on": self.depends_on,
             "logical_op_id": self.logical_op_id,
+            "unique_logical_op_id": self.unique_logical_op_id,
             "logical_op_name": self.logical_op_name,
-            "target_cache_id": self.target_cache_id,
+            "api_base": self.api_base,
             "verbose": self.verbose,
         }
 
@@ -116,10 +122,7 @@ class PhysicalOperator:
         # get op name and op parameters which are relevant for computing the id
         op_name = self.op_name()
         id_params = self.get_id_params()
-        id_params = {
-            k: str(v) if k != "output_schema" else sorted(v.field_names())
-            for k, v in id_params.items()
-        }
+        id_params = {k: str(v) for k, v in id_params.items()}
 
         # compute, set, and return the op_id
         hash_str = json.dumps({"op_name": op_name, **id_params}, sort_keys=True)
@@ -127,9 +130,12 @@ class PhysicalOperator:
 
         return self.op_id
     
-    def get_logical_op_id(self) -> str | None:
+    def get_logical_op_id(self) -> str:
         return self.logical_op_id
 
+    def get_unique_logical_op_id(self) -> str:
+        return self.unique_logical_op_id
+
     def get_full_op_id(self):
         return f"{self.get_logical_op_id()}-{self.get_op_id()}"
 
@@ -148,9 +154,9 @@ class PhysicalOperator:
             else None
         )
         input_fields = (
-            self.input_schema.field_names()
+            list(self.input_schema.model_fields)
             if depends_on_fields is None
-            else [field for field in self.input_schema.field_names() if field in depends_on_fields]
+            else [field for field in self.input_schema.model_fields if field in depends_on_fields]
         )
 
         return input_fields

palimpzest/query/operators/project.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from palimpzest.core.data.dataclasses import OperatorCostEstimates, RecordOpStats
 from palimpzest.core.elements.records import DataRecord, DataRecordSet
+from palimpzest.core.models import OperatorCostEstimates, RecordOpStats
 from palimpzest.query.operators.physical import PhysicalOperator
 
 
@@ -39,8 +39,8 @@ class ProjectOp(PhysicalOperator):
         # create RecordOpStats object
         record_op_stats = RecordOpStats(
             record_id=dr.id,
-            record_parent_id=dr.parent_id,
-            record_source_idx=dr.source_idx,
+            record_parent_ids=dr.parent_ids,
+            record_source_indices=dr.source_indices,
             record_state=dr.to_dict(include_bytes=False),
             full_op_id=self.get_full_op_id(),
             logical_op_id=self.logical_op_id,

palimpzest/query/operators/rag_convert.py

@@ -5,15 +5,15 @@ import time
 from numpy import dot
 from numpy.linalg import norm
 from openai import OpenAI
+from pydantic.fields import FieldInfo
 
 from palimpzest.constants import (
     MODEL_CARDS,
     NAIVE_EST_NUM_OUTPUT_TOKENS,
     Model,
 )
-from palimpzest.core.data.dataclasses import GenerationStats, OperatorCostEstimates
 from palimpzest.core.elements.records import DataRecord
-from palimpzest.core.lib.fields import Field, StringField
+from palimpzest.core.models import GenerationStats, OperatorCostEstimates
 from palimpzest.query.operators.convert import LLMConvert
 
 
@@ -143,9 +143,9 @@ class RAGConvert(LLMConvert):
 
         # compute embedding for output fields
         output_fields_desc = ""
-        field_desc_map = self.output_schema.field_desc_map()
         for field_name in output_fields:
-            output_fields_desc += f"- {field_name}: {field_desc_map[field_name]}\n"
+            desc = self.output_schema.model_fields[field_name].description
+            output_fields_desc += f"- {field_name}: {'no description available' if desc is None else desc}\n"
         query_embedding, query_embed_stats = self.compute_embedding(output_fields_desc)
 
         # add cost of embedding the query to embed_stats
@@ -156,8 +156,8 @@ class RAGConvert(LLMConvert):
             field = candidate.get_field_type(field_name)
 
             # skip this field if it is not a string or a list of strings
-            is_string_field = isinstance(field, StringField)
-            is_list_string_field = hasattr(field, "element_type") and isinstance(field.element_type, StringField)
+            is_string_field = field.annotation in [str, str | None]
+            is_list_string_field = field.annotation in [list[str], list[str] | None]
             if not (is_string_field or is_list_string_field):
                 continue
 
@@ -193,7 +193,7 @@ class RAGConvert(LLMConvert):
 
         return candidate, embed_stats
 
-    def convert(self, candidate: DataRecord, fields: dict[str, Field]) -> tuple[dict[str, list], GenerationStats]:
+    def convert(self, candidate: DataRecord, fields: dict[str, FieldInfo]) -> tuple[dict[str, list], GenerationStats]:
         # set client
         self.client = OpenAI() if self.client is None else self.client
 

palimpzest/query/operators/retrieve.py

@@ -8,12 +8,12 @@ from chromadb.api.models.Collection import Collection
 from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
 from chromadb.utils.embedding_functions.openai_embedding_function import OpenAIEmbeddingFunction
 from openai import OpenAI
+from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
 
 from palimpzest.constants import MODEL_CARDS, Model
-from palimpzest.core.data.dataclasses import GenerationStats, OperatorCostEstimates, RecordOpStats
 from palimpzest.core.elements.records import DataRecord, DataRecordSet
-from palimpzest.core.lib.schemas import Schema
+from palimpzest.core.models import GenerationStats, OperatorCostEstimates, RecordOpStats
 from palimpzest.query.operators.physical import PhysicalOperator
 
 
@@ -22,7 +22,7 @@ class RetrieveOp(PhysicalOperator):
         self,
         index: Collection,
         search_attr: str,
-        output_attrs: list[dict] | type[Schema],
+        output_attrs: list[dict] | type[BaseModel],
         search_func: Callable | None,
         k: int,
         *args,
@@ -41,12 +41,12 @@ class RetrieveOp(PhysicalOperator):
         super().__init__(*args, **kwargs)
 
         # extract the field names from the output_attrs
-        if isinstance(output_attrs, Schema):
-            self.output_field_names = output_attrs.field_names()
+        if issubclass(output_attrs, BaseModel):
+            self.output_field_names = list(output_attrs.model_fields)
         elif isinstance(output_attrs, list):
             self.output_field_names = [attr["name"] for attr in output_attrs]
         else:
-            raise ValueError("`output_attrs` must be a list of dicts or a Schema object.")
+            raise ValueError("`output_attrs` must be a list of dicts or a `pydantic.BaseModel` object.")
 
         if len(self.output_field_names) != 1 and search_func is None:
             raise ValueError("If `search_func` is None, `output_attrs` must have a single field.")
@@ -160,8 +160,8 @@ class RetrieveOp(PhysicalOperator):
         # construct the RecordOpStats object
         record_op_stats = RecordOpStats(
             record_id=output_dr.id,
-            record_parent_id=output_dr.parent_id,
-            record_source_idx=output_dr.source_idx,
+            record_parent_ids=output_dr.parent_ids,
+            record_source_indices=output_dr.source_indices,
             record_state=record_state,
             full_op_id=self.get_full_op_id(),
             logical_op_id=self.logical_op_id,
@@ -169,7 +169,7 @@ class RetrieveOp(PhysicalOperator):
             time_per_record=total_time,
             cost_per_record=generation_stats.cost_per_record,
             answer=answer,
-            input_fields=self.input_schema.field_names(),
+            input_fields=list(self.input_schema.model_fields),
             generated_fields=generated_fields,
             fn_call_duration_secs=total_time - generation_stats.llm_call_duration_secs,
             llm_call_duration_secs=generation_stats.llm_call_duration_secs,

palimpzest/query/operators/scan.py

@@ -2,46 +2,43 @@ from __future__ import annotations
 
 import time
 from abc import ABC, abstractmethod
+from typing import Any
 
-from palimpzest.constants import (
-    LOCAL_SCAN_TIME_PER_KB,
-    MEMORY_SCAN_TIME_PER_KB,
-    Cardinality,
-)
-from palimpzest.core.data.dataclasses import OperatorCostEstimates, RecordOpStats
-from palimpzest.core.data.datareaders import DataReader, DirectoryReader, FileReader
+from palimpzest.constants import LOCAL_SCAN_TIME_PER_KB
+from palimpzest.core.data import context
 from palimpzest.core.elements.records import DataRecord, DataRecordSet
+from palimpzest.core.models import OperatorCostEstimates, RecordOpStats
 from palimpzest.query.operators.physical import PhysicalOperator
 
 
 class ScanPhysicalOp(PhysicalOperator, ABC):
     """
-    Physical operators which implement DataReaders require slightly more information
+    Physical operators which implement root Datasets require slightly more information
     in order to accurately compute naive cost estimates. Thus, we use a slightly
     modified abstract base class for these operators.
     """
-
-    def __init__(self, datareader: DataReader, *args, **kwargs):
+    # datasource: IterDataset
+    def __init__(self, datasource: Any, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.datareader = datareader
+        self.datasource = datasource
 
     def __str__(self):
-        op = f"{self.op_name()}({self.datareader}) -> {self.output_schema}\n"
-        op += f"    ({', '.join(self.output_schema.field_names())[:30]})\n"
+        op = f"{self.op_name()}({self.datasource}) -> {self.output_schema}\n"
+        op += f"    ({', '.join(list(self.output_schema.model_fields))[:30]})\n"
         return op
 
     def get_id_params(self):
-        return super().get_id_params()
+        id_params = super().get_id_params()
+        return {"datasource_id": self.datasource.id, **id_params}
 
     def get_op_params(self):
         op_params = super().get_op_params()
-        return {"datareader": self.datareader, **op_params}
+        return {"datasource": self.datasource, **op_params}
 
     @abstractmethod
     def naive_cost_estimates(
         self,
         source_op_cost_estimates: OperatorCostEstimates,
-        input_cardinality: Cardinality,
         input_record_size_in_bytes: int | float,
     ) -> OperatorCostEstimates:
         """
@@ -62,30 +59,27 @@ class ScanPhysicalOp(PhysicalOperator, ABC):
 
     def __call__(self, idx: int) -> DataRecordSet:
         """
-        This function invokes `self.datareader.__getitem__` on the given `idx` to retrieve the next data item.
+        This function invokes `self.datasource.__getitem__` on the given `idx` to retrieve the next data item.
         It then returns this item as a DataRecord wrapped in a DataRecordSet.
         """
         start_time = time.time()
-        item = self.datareader[idx]
+        item = self.datasource[idx]
         end_time = time.time()
 
-        # TODO: remove once validation data is refactored
-        item_field_dict = item.get("fields", item)
-
         # check that item covers fields in output schema
-        output_field_names = self.output_schema.field_names()
-        assert all([field in item_field_dict for field in output_field_names]), f"Some fields in DataReader schema not present in item!\n - DataReader fields: {output_field_names}\n - Item fields: {list(item.keys())}"
+        output_field_names = list(self.output_schema.model_fields)
+        assert all([field in item for field in output_field_names]), f"Some fields in Dataset schema not present in item!\n - Dataset fields: {output_field_names}\n - Item fields: {list(item.keys())}"
 
         # construct a DataRecord from the item
-        dr = DataRecord(self.output_schema, source_idx=idx)
+        dr = DataRecord(self.output_schema, source_indices=[f"{self.datasource.id}-{idx}"])
         for field in output_field_names:
-            setattr(dr, field, item_field_dict[field])
+            setattr(dr, field, item[field])
 
         # create RecordOpStats objects
         record_op_stats = RecordOpStats(
             record_id=dr.id,
-            record_parent_id=dr.parent_id,
-            record_source_idx=dr.source_idx,
+            record_parent_ids=dr.parent_ids,
+            record_source_indices=dr.source_indices,
             record_state=dr.to_dict(include_bytes=False),
             full_op_id=self.get_full_op_id(),
             logical_op_id=self.logical_op_id,
@@ -97,7 +91,7 @@ class ScanPhysicalOp(PhysicalOperator, ABC):
  
         # construct and return DataRecordSet object
         return DataRecordSet([dr], [record_op_stats])
-        
+
 
 class MarshalAndScanDataOp(ScanPhysicalOp):
     def naive_cost_estimates(
@@ -110,11 +104,14 @@ class MarshalAndScanDataOp(ScanPhysicalOp):
 
         # estimate time spent reading each record
         per_record_size_kb = input_record_size_in_bytes / 1024.0
-        time_per_record = (
-            LOCAL_SCAN_TIME_PER_KB * per_record_size_kb
-            if isinstance(self.datareader, (DirectoryReader, FileReader))
-            else MEMORY_SCAN_TIME_PER_KB * per_record_size_kb
-        )
+
+        # TODO: cannot do the first computation b/c we cannot import iter_dataset; possibly revisit
+        # time_per_record = (
+        #     MEMORY_SCAN_TIME_PER_KB * per_record_size_kb
+        #     if isinstance(self.datasource, (iter_dataset.MemoryDataset))
+        #     else LOCAL_SCAN_TIME_PER_KB * per_record_size_kb
+        # )
+        time_per_record = LOCAL_SCAN_TIME_PER_KB * per_record_size_kb
 
         # estimate output cardinality
         cardinality = source_op_cost_estimates.cardinality
@@ -128,26 +125,68 @@ class MarshalAndScanDataOp(ScanPhysicalOp):
         )
 
 
-class CacheScanDataOp(ScanPhysicalOp):
+class ContextScanOp(PhysicalOperator):
+    """
+    Physical operator which facillitates the loading of a Context for processing.
+    """
+
+    def __init__(self, context: context.Context, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.context = context
+
+    def __str__(self):
+        op = f"{self.op_name()}({self.context}) -> {self.output_schema}\n"
+        op += f"    ({', '.join(list(self.output_schema.model_fields))[:30]})\n"
+        return op
+
+    def get_id_params(self):
+        return super().get_id_params()
+
+    def get_op_params(self):
+        op_params = super().get_op_params()
+        return {"context": self.context, **op_params}
+
     def naive_cost_estimates(
         self,
         source_op_cost_estimates: OperatorCostEstimates,
-        input_record_size_in_bytes: int | float,
     ):
         # get inputs needed for naive cost estimation
         # TODO: we should rename cardinality --> "multiplier" or "selectivity" one-to-one / one-to-many
 
         # estimate time spent reading each record
-        per_record_size_kb = input_record_size_in_bytes / 1024.0
-        time_per_record = LOCAL_SCAN_TIME_PER_KB * per_record_size_kb
+        time_per_record = LOCAL_SCAN_TIME_PER_KB * 1.0
 
-        # estimate output cardinality
-        cardinality = source_op_cost_estimates.cardinality
-
-        # for now, assume no cost per record for reading from cache
+        # for now, assume no cost per record for reading data
         return OperatorCostEstimates(
-            cardinality=cardinality,
+            cardinality=1.0,
             time_per_record=time_per_record,
             cost_per_record=0,
             quality=1.0,
         )
+
+    def __call__(self, *args, **kwargs) -> DataRecordSet:
+        """
+        This function returns the context as a DataRecord wrapped in a DataRecordSet.
+        """
+        # construct a DataRecord from the context
+        start_time = time.time()
+        dr = DataRecord(self.output_schema, source_indices=[f"{self.context.id}-{0}"])
+        dr.context = self.context
+        end_time = time.time()
+
+        # create RecordOpStats objects
+        record_op_stats = RecordOpStats(
+            record_id=dr.id,
+            record_parent_ids=dr.parent_ids,
+            record_source_indices=dr.source_indices,
+            record_state=dr.to_dict(include_bytes=False),
+            full_op_id=self.get_full_op_id(),
+            logical_op_id=self.logical_op_id,
+            op_name=self.op_name(),
+            time_per_record=(end_time - start_time),
+            cost_per_record=0.0,
+            op_details={k: str(v) for k, v in self.get_id_params().items()},
+        )
+ 
+        # construct and return DataRecordSet object
+        return DataRecordSet([dr], [record_op_stats])