palimpzest 0.8.7__py3-none-any.whl → 1.0.0__py3-none-any.whl

palimpzest/prompts/validator.py

@@ -22,17 +22,17 @@ OUTPUT FIELDS:
 - birth_year: the year the scientist was born
 
 CONTEXT:
-{{
+{
   "text": "Augusta Ada King, Countess of Lovelace, also known as Ada Lovelace, was an English mathematician and writer chiefly known for her work on Charles Babbage's proposed mechanical general-purpose computer, the Analytical Engine. She was the first to recognise that the machine had applications beyond pure calculation.",
   "birthday": "December 10, 1815"
-}}
+}
 
 OUTPUT:
 --------
-{{
+{
   "name": "Charles Babbage",
   "birth_year": 1815
-}}
+}
 
 EVALUATION: {"name": 0.0, "birth_year": 1.0}
 
@@ -66,18 +66,18 @@ OUTPUT FIELDS:
 - person_in_image: true if a person is in the image and false otherwise
 
 CONTEXT:
-{{
+{
   "image": <bytes>,
   "photographer": "CameraEnthusiast1"
-}}
+}
 <image content provided here; assume in this example the image shows a dog and a cat playing>
 
 OUTPUT:
 --------
-{{
+{
   "dog_in_image": true,
   "person_in_image": true
-}}
+}
 
 EVALUATION: {"dog_in_image": 1.0, "person_in_image": 0.0}
 
@@ -113,22 +113,22 @@ OUTPUT FIELDS:
 - birth_year: the year the scientist was born
 
 CONTEXT:
-{{
+{
   "text": "Augusta Ada King, Countess of Lovelace, also known as Ada Lovelace, was an English mathematician and writer chiefly known for her work on Charles Babbage's proposed mechanical general-purpose computer, the Analytical Engine. She was the first to recognise that the machine had applications beyond pure calculation.",
   "birthdays": "...Lovelace was born on December 10, 1815, almost exactly 24 years after Babbage's birth on 26 December 1791..."
-}}
+}
 
 OUTPUTS:
 --------
 [
-  {{
+  {
     "name": "Ada Lovelace",
     "birth_year": 1815
-  }},
-  {{
+  },
+  {
     "name": "Charles Babbage",
     "birth_year": 1790
-  }}
+  }
 ]
 
 EVALUATION: [{"name": 1.0, "birth_year": 1.0}, {"name": 1.0, "birth_year": 0.0}]
@@ -163,23 +163,23 @@ OUTPUT FIELDS:
 - animal_is_canine: true if the animal is a canine and false otherwise
 
 CONTEXT:
-{{
+{
   "image": <bytes>,
   "photographer": "CameraEnthusiast1"
-}}
+}
 <image content provided here; assume in this example the image shows a dog and a cat playing>
 
 OUTPUT:
 --------
 [
-  {{
+  {
     "animal": "dog",
     "animal_is_canine": true
-  }},
-  {{
+  },
+  {
     "animal": "cat",
     "animal_is_canine": true
-  }}
+  }
 ]
 
 EVALUATION: [{"animal": 1.0, "animal_is_canine": 1.0}, {"animal": 1.0, "animal_is_canine": 0.0}]
@@ -214,20 +214,20 @@ OUTPUT FIELDS:
 - related_scientists: list of scientists who perform similar work as the scientist described in the text
 
 CONTEXT:
-{{
+{
   "text": "Augusta Ada King, Countess of Lovelace, also known as Ada Lovelace, was an English mathematician and writer chiefly known for her work on Charles Babbage's proposed mechanical general-purpose computer, the Analytical Engine. She was the first to recognise that the machine had applications beyond pure calculation.",
-}}
+}
 
 OUTPUT:
 --------
-{{
+{
   "related_scientists": [
     "Charles Babbage",
     "Alan Turing",
     "Charles Darwin",
     "John von Neumann",
   ]
-}}
+}
 
 EVALUATION: {"related_scientists": 0.75}
 

palimpzest/query/execution/all_sample_execution_strategy.py

@@ -225,7 +225,7 @@ class AllSamplingExecutionStrategy(SentinelExecutionStrategy):
         dataset_id_to_source_indices = {}
         for dataset_id, dataset in train_dataset.items():
             total_num_samples = len(dataset)
-            source_indices = [f"{dataset_id}-{int(idx)}" for idx in np.arange(total_num_samples)]
+            source_indices = [f"{dataset_id}---{int(idx)}" for idx in np.arange(total_num_samples)]
             dataset_id_to_source_indices[dataset_id] = source_indices
 
         # initialize set of physical operators for each logical operator

palimpzest/query/execution/execution_strategy.py

@@ -14,8 +14,8 @@ from palimpzest.query.operators.convert import LLMConvert
 from palimpzest.query.operators.filter import LLMFilter
 from palimpzest.query.operators.join import JoinOp
 from palimpzest.query.operators.physical import PhysicalOperator
-from palimpzest.query.operators.retrieve import RetrieveOp
 from palimpzest.query.operators.scan import ContextScanOp, ScanPhysicalOp
+from palimpzest.query.operators.topk import TopKOp
 from palimpzest.query.optimizer.plan import PhysicalPlan, SentinelPlan
 from palimpzest.utils.progress import PZSentinelProgressManager
 from palimpzest.validator.validator import Validator
@@ -123,7 +123,7 @@ class SentinelExecutionStrategy(BaseExecutionStrategy, ABC):
             return (
                 not isinstance(op, LLMConvert)
                 and not isinstance(op, LLMFilter)
-                and not isinstance(op, RetrieveOp)
+                and not isinstance(op, TopKOp)
                 and not isinstance(op, JoinOp)
             )
 
@@ -167,8 +167,8 @@ class SentinelExecutionStrategy(BaseExecutionStrategy, ABC):
                             full_hashes.add(full_hash)
                             futures.append(executor.submit(validator._score_flat_map, op, fields, input_record, output, full_hash))
 
-                    # create future for retrieve
-                    elif isinstance(op, RetrieveOp):
+                    # create future for top-k
+                    elif isinstance(op, TopKOp):
                         fields = op.generated_fields
                         input_record: DataRecord = record_set.input
                         output = record_set.data_records[0].to_dict(project_cols=fields)
@@ -176,7 +176,7 @@ class SentinelExecutionStrategy(BaseExecutionStrategy, ABC):
                         full_hash = f"{hash(input_record)}{hash(output_str)}"
                         if full_hash not in full_hashes:
                             full_hashes.add(full_hash)
-                            futures.append(executor.submit(validator._score_retrieve, op, fields, input_record, output, full_hash))
+                            futures.append(executor.submit(validator._score_topk, op, fields, input_record, output, full_hash))
 
                     # create future for filter
                     elif isinstance(op, LLMFilter):
@@ -235,7 +235,7 @@ class SentinelExecutionStrategy(BaseExecutionStrategy, ABC):
 
                 # TODO: this scoring function will (likely) bias towards small values of k since it
                 # measures precision and not recall / F1; will need to revisit this in the future
-                elif isinstance(op, RetrieveOp):
+                elif isinstance(op, TopKOp):
                     fields = op.generated_fields
                     input_record: DataRecord = record_set.input
                     output_str = record_set.data_records[0].to_json_str(project_cols=fields, bytes_to_str=True, sorted=True)
@@ -341,9 +341,9 @@ class SentinelExecutionStrategy(BaseExecutionStrategy, ABC):
     def _is_llm_op(self, physical_op: PhysicalOperator) -> bool:
         is_llm_convert = isinstance(physical_op, LLMConvert)
         is_llm_filter = isinstance(physical_op, LLMFilter)
-        is_llm_retrieve = isinstance(physical_op, RetrieveOp) and isinstance(physical_op.index, Collection)
+        is_llm_topk = isinstance(physical_op, TopKOp) and isinstance(physical_op.index, Collection)
         is_llm_join = isinstance(physical_op, JoinOp)
-        return is_llm_convert or is_llm_filter or is_llm_retrieve or is_llm_join
+        return is_llm_convert or is_llm_filter or is_llm_topk or is_llm_join
 
     @abstractmethod
     def execute_sentinel_plan(self, sentinel_plan: SentinelPlan, train_dataset: dict[str, Dataset], validator: Validator) -> SentinelPlanStats:

palimpzest/query/execution/mab_execution_strategy.py

@@ -14,8 +14,8 @@ from palimpzest.query.operators.convert import LLMConvert
 from palimpzest.query.operators.filter import FilterOp, LLMFilter, NonLLMFilter
 from palimpzest.query.operators.join import JoinOp
 from palimpzest.query.operators.physical import PhysicalOperator
-from palimpzest.query.operators.retrieve import RetrieveOp
 from palimpzest.query.operators.scan import ContextScanOp, ScanPhysicalOp
+from palimpzest.query.operators.topk import TopKOp
 from palimpzest.query.optimizer.plan import SentinelPlan
 from palimpzest.utils.progress import create_progress_manager
 from palimpzest.validator.validator import Validator
@@ -66,8 +66,8 @@ class OpFrontier:
         self.is_llm_join = isinstance(sample_op, JoinOp)
         is_llm_convert = isinstance(sample_op, LLMConvert)
         is_llm_filter = isinstance(sample_op, LLMFilter)
-        is_llm_retrieve = isinstance(sample_op, RetrieveOp) and isinstance(sample_op.index, Collection)
-        self.is_llm_op = is_llm_convert or is_llm_filter or is_llm_retrieve or self.is_llm_join
+        is_llm_topk = isinstance(sample_op, TopKOp) and isinstance(sample_op.index, Collection)
+        self.is_llm_op = is_llm_convert or is_llm_filter or is_llm_topk or self.is_llm_join
 
         # get order in which we will sample physical operators for this logical operator
         sample_op_indices = self._get_op_index_order(op_set, seed)
@@ -96,6 +96,12 @@ class OpFrontier:
         """
         return self.frontier_ops
 
+    def get_off_frontier_ops(self) -> list[PhysicalOperator]:
+        """
+        Returns the set of off-frontier operators for this OpFrontier.
+        """
+        return self.off_frontier_ops
+
     def _compute_op_id_to_pareto_distance(self, priors: dict[str, dict[str, float]]) -> dict[str, float]:
         """
         Return l2-distance for each operator from the pareto frontier.
@@ -298,7 +304,7 @@ class OpFrontier:
         def remove_unavailable_root_datasets(source_indices: str | tuple) -> str | tuple | None:
             # base case: source_indices is a string
             if isinstance(source_indices, str):
-                return source_indices if source_indices.split("-")[0] in self.root_dataset_ids else None
+                return source_indices if source_indices.split("---")[0] in self.root_dataset_ids else None
 
             # recursive case: source_indices is a tuple
             left_indices = source_indices[0]
@@ -383,6 +389,12 @@ class OpFrontier:
             # compute final list of record op stats
             full_op_id_to_record_op_stats[full_op_id] = list(record_id_to_max_quality_record_op_stats.values())
 
+        # NOTE: it is possible for the full_op_id_to_record_op_stats to be empty if there is a duplicate operator
+        # (e.g. a scan of the same dataset) which has all of its results cached and no new_record_op_stats;
+        # in this case, we do not update the frontier
+        if full_op_id_to_record_op_stats == {}:
+            return
+
         # update the set of source indices processed by each physical operator
         for full_op_id, source_indices_processed in full_op_id_to_source_indices_processed.items():
             # update the set of source indices processed
@@ -641,8 +653,8 @@ class MABExecutionStrategy(SentinelExecutionStrategy):
         """
         Returns the operator in the frontier with the highest (estimated) quality.
         """
-        # get the operators in the frontier set for this logical_op_id
-        frontier_ops = op_frontiers[unique_logical_op_id].get_frontier_ops()
+        # get the (off) frontier operators for this logical_op_id
+        frontier_ops = op_frontiers[unique_logical_op_id].get_frontier_ops() + op_frontiers[unique_logical_op_id].get_off_frontier_ops()
 
         # get a mapping from full_op_id --> list[RecordOpStats]
         full_op_id_to_op_stats: dict[str, OperatorStats] = plan_stats.operator_stats.get(unique_logical_op_id, {})
@@ -693,14 +705,21 @@ class MABExecutionStrategy(SentinelExecutionStrategy):
                 max_quality_op = self._get_max_quality_op(unique_logical_op_id, op_frontiers, plan_stats)
 
                 # get frontier ops and their next input
-                def is_filtered_out(tup: tuple) -> bool:
-                    return tup[-1] is None or isinstance(tup[-1], list) and all([record is None for record in tup[-1]])
+                def filter_and_clean_inputs(frontier_op_inputs: list[tuple]) -> bool:
+                    cleaned_inputs = []
+                    for tup in frontier_op_inputs:
+                        input = tup[-1]
+                        if isinstance(input, list):
+                            input = [record for record in input if record is not None]
+                        if input is not None and input != []:
+                            cleaned_inputs.append((tup[0], tup[1], input))
+                    return cleaned_inputs
                 frontier_op_inputs = op_frontiers[unique_logical_op_id].get_frontier_op_inputs(source_indices_to_sample, max_quality_op)
-                frontier_op_inputs = list(filter(lambda tup: not is_filtered_out(tup), frontier_op_inputs))
+                frontier_op_inputs = filter_and_clean_inputs(frontier_op_inputs)
 
                 # break out of the loop if frontier_op_inputs is empty, as this means all records have been filtered out
                 if len(frontier_op_inputs) == 0:
-                    break
+                    continue
 
                 # run sampled operators on sampled inputs and update the number of samples drawn
                 source_indices_to_record_set_tuples, num_llm_ops = self._execute_op_set(unique_logical_op_id, frontier_op_inputs)
@@ -764,7 +783,7 @@ class MABExecutionStrategy(SentinelExecutionStrategy):
         dataset_id_to_shuffled_source_indices = {}
         for dataset_id, dataset in train_dataset.items():
             total_num_samples = len(dataset)
-            shuffled_source_indices = [f"{dataset_id}-{int(idx)}" for idx in np.arange(total_num_samples)]
+            shuffled_source_indices = [f"{dataset_id}---{int(idx)}" for idx in np.arange(total_num_samples)]
             self.rng.shuffle(shuffled_source_indices)
             dataset_id_to_shuffled_source_indices[dataset_id] = shuffled_source_indices
 

palimpzest/query/execution/parallel_execution_strategy.py

@@ -9,7 +9,6 @@ from palimpzest.query.operators.aggregate import AggregateOp
 from palimpzest.query.operators.distinct import DistinctOp
 from palimpzest.query.operators.join import JoinOp
 from palimpzest.query.operators.limit import LimitScanOp
-from palimpzest.query.operators.physical import PhysicalOperator
 from palimpzest.query.operators.scan import ContextScanOp, ScanPhysicalOp
 from palimpzest.query.optimizer.plan import PhysicalPlan
 from palimpzest.utils.progress import create_progress_manager
@@ -35,14 +34,27 @@ class ParallelExecutionStrategy(ExecutionStrategy):
                 return True
         return False
 
-    def _upstream_ops_finished(self, plan: PhysicalPlan, topo_idx: int, operator: PhysicalOperator, input_queues: dict[str, dict[str, list]], future_queues: dict[str, list]) -> bool:
+    def _upstream_ops_finished(self, plan: PhysicalPlan, unique_full_op_id: str, input_queues: dict[str, dict[str, list]], future_queues: dict[str, list]) -> bool:
         """Helper function to check if agg / join operator is ready to process its inputs."""
-        # for agg / join operator, we can only process it when all upstream operators have finished processing their inputs
-        upstream_unique_full_op_ids = plan.get_upstream_unique_full_op_ids(topo_idx, operator)
+        upstream_unique_full_op_ids = plan.get_upstream_unique_full_op_ids(unique_full_op_id)
         upstream_input_queues = {upstream_unique_full_op_id: input_queues[upstream_unique_full_op_id] for upstream_unique_full_op_id in upstream_unique_full_op_ids}
         upstream_future_queues = {upstream_unique_full_op_id: future_queues[upstream_unique_full_op_id] for upstream_unique_full_op_id in upstream_unique_full_op_ids}
         return not (self._any_queue_not_empty(upstream_input_queues) or self._any_queue_not_empty(upstream_future_queues))
 
+    def _finish_outer_join(self, executor: ThreadPoolExecutor, plan: PhysicalPlan, unique_full_op_id: str, input_queues: dict[str, dict[str, list]], future_queues: dict[str, list]) -> None:
+        join_op_upstream_finished = self._upstream_ops_finished(plan, unique_full_op_id, input_queues, future_queues)
+        join_input_queues_empty = all(len(inputs) == 0 for inputs in input_queues[unique_full_op_id].values())
+        join_future_queue_empty = len(future_queues[unique_full_op_id]) == 0
+        if join_op_upstream_finished and join_input_queues_empty and join_future_queue_empty:
+            # process the join one last time with final=True to handle any left/right/outer join logic
+            operator = self.unique_full_op_id_to_operator[unique_full_op_id]
+            if not operator.finished:
+                def finalize_op(operator):
+                    return operator([], [], final=True)
+                future = executor.submit(finalize_op, operator)
+                future_queues[unique_full_op_id].append(future)
+                operator.set_finished()
+
     def _process_future_results(self, unique_full_op_id: str, future_queues: dict[str, list], plan_stats: PlanStats) -> list[DataRecord]:
         """
         Helper function which takes a full operator id, the future queues, and plan stats, and performs
@@ -117,15 +129,23 @@ class ParallelExecutionStrategy(ExecutionStrategy):
                             records = self._process_future_results(source_unique_full_op_id, future_queues, plan_stats)
                             input_queues[unique_full_op_id][source_unique_full_op_id].extend(records)
 
+                            # if the source is a left/right/outer join operator with no more inputs to process, then finish it
+                            if self.is_outer_join_op[source_unique_full_op_id]:
+                                self._finish_outer_join(executor, plan, source_unique_full_op_id, input_queues, future_queues)
+
                     # for the final operator, add any finished futures to the output_records
                     if unique_full_op_id == f"{topo_idx}-{final_op.get_full_op_id()}":
                         records = self._process_future_results(unique_full_op_id, future_queues, plan_stats)
                         output_records.extend(records)
 
+                        # if this is a left/right/outer join operator with no more inputs to process, then finish it
+                        if self.is_outer_join_op[unique_full_op_id]:
+                            self._finish_outer_join(executor, plan, unique_full_op_id, input_queues, future_queues)
+
                     # if this operator does not have enough inputs to execute, then skip it
                     num_inputs = sum(len(inputs) for inputs in input_queues[unique_full_op_id].values())
-                    agg_op_not_ready = isinstance(operator, AggregateOp) and not self._upstream_ops_finished(plan, topo_idx, operator, input_queues, future_queues)
-                    join_op_not_ready = isinstance(operator, JoinOp) and not self.join_has_downstream_limit_op[unique_full_op_id] and not self._upstream_ops_finished(plan, topo_idx, operator, input_queues, future_queues)
+                    agg_op_not_ready = isinstance(operator, AggregateOp) and not self._upstream_ops_finished(plan, unique_full_op_id, input_queues, future_queues)
+                    join_op_not_ready = isinstance(operator, JoinOp) and not self.join_has_downstream_limit_op[unique_full_op_id] and not self._upstream_ops_finished(plan, unique_full_op_id, input_queues, future_queues)
                     if num_inputs == 0 or agg_op_not_ready or join_op_not_ready:
                         continue
 
@@ -225,8 +245,9 @@ class ParallelExecutionStrategy(ExecutionStrategy):
         input_queues = self._create_input_queues(plan)
         future_queues = {f"{topo_idx}-{op.get_full_op_id()}": [] for topo_idx, op in enumerate(plan)}
 
-        # precompute which operators are joins and which joins have downstream limit ops
+        # precompute which operators are (outer) joins and which joins have downstream limit ops
         self.is_join_op = {f"{topo_idx}-{op.get_full_op_id()}": isinstance(op, JoinOp) for topo_idx, op in enumerate(plan)}
+        self.is_outer_join_op = {f"{topo_idx}-{op.get_full_op_id()}": isinstance(op, JoinOp) and op.how in ("left", "right", "outer") for topo_idx, op in enumerate(plan)}
         self.join_has_downstream_limit_op = {}
         for topo_idx, op in enumerate(plan):
             if isinstance(op, JoinOp):
@@ -240,6 +261,9 @@ class ParallelExecutionStrategy(ExecutionStrategy):
                         break
                 self.join_has_downstream_limit_op[unique_full_op_id] = has_downstream_limit_op
 
+        # precompute mapping from unique_full_op_id to operator instance
+        self.unique_full_op_id_to_operator = {f"{topo_idx}-{op.get_full_op_id()}": op for topo_idx, op in enumerate(plan)}
+
         # initialize and start the progress manager
         self.progress_manager = create_progress_manager(plan, num_samples=self.num_samples, progress=self.progress)
         self.progress_manager.start()

palimpzest/query/execution/single_threaded_execution_strategy.py

@@ -6,7 +6,6 @@ from palimpzest.query.execution.execution_strategy import ExecutionStrategy
 from palimpzest.query.operators.aggregate import AggregateOp
 from palimpzest.query.operators.join import JoinOp
 from palimpzest.query.operators.limit import LimitScanOp
-from palimpzest.query.operators.physical import PhysicalOperator
 from palimpzest.query.operators.scan import ContextScanOp, ScanPhysicalOp
 from palimpzest.query.optimizer.plan import PhysicalPlan
 from palimpzest.utils.progress import create_progress_manager
@@ -70,6 +69,13 @@ class SequentialSingleThreadExecutionStrategy(ExecutionStrategy):
                 record_set, num_inputs_processed = operator(left_input_records, right_input_records)
                 records = record_set.data_records
                 record_op_stats = record_set.record_op_stats
+
+                # process the join one last time with final=True to handle any left/right/outer join logic
+                if operator.how in ("left", "right", "outer"):
+                    record_set, num_inputs_processed = operator([], [], final=True)
+                    records.extend(record_set.data_records)
+                    record_op_stats.extend(record_set.record_op_stats)
+      
                 num_outputs = sum(record._passed_operator for record in records)
 
                 # update the progress manager
@@ -168,10 +174,9 @@ class PipelinedSingleThreadExecutionStrategy(ExecutionStrategy):
                 return True
         return False
 
-    def _upstream_ops_finished(self, plan: PhysicalPlan, topo_idx: int, operator: PhysicalOperator, input_queues: dict[str, dict[str, list]]) -> bool:
+    def _upstream_ops_finished(self, plan: PhysicalPlan, unique_full_op_id: str, input_queues: dict[str, dict[str, list]]) -> bool:
         """Helper function to check if agg / join operator is ready to process its inputs."""
-        # for agg / join operator, we can only process it when all upstream operators have finished processing their inputs
-        upstream_unique_full_op_ids = plan.get_upstream_unique_full_op_ids(topo_idx, operator)
+        upstream_unique_full_op_ids = plan.get_upstream_unique_full_op_ids(unique_full_op_id)
         upstream_input_queues = {upstream_unique_full_op_id: input_queues[upstream_unique_full_op_id] for upstream_unique_full_op_id in upstream_unique_full_op_ids}
         return not self._any_queue_not_empty(upstream_input_queues)
 
@@ -192,8 +197,8 @@ class PipelinedSingleThreadExecutionStrategy(ExecutionStrategy):
                 unique_full_op_id = f"{topo_idx}-{operator.get_full_op_id()}"
 
                 num_inputs = sum(len(input_queues[unique_full_op_id][source_unique_full_op_id]) for source_unique_full_op_id in source_unique_full_op_ids)
-                agg_op_not_ready = isinstance(operator, AggregateOp) and not self._upstream_ops_finished(plan, topo_idx, operator, input_queues)
-                join_op_not_ready = isinstance(operator, JoinOp) and not self._upstream_ops_finished(plan, topo_idx, operator, input_queues)
+                agg_op_not_ready = isinstance(operator, AggregateOp) and not self._upstream_ops_finished(plan, unique_full_op_id, input_queues)
+                join_op_not_ready = isinstance(operator, JoinOp) and not self._upstream_ops_finished(plan, unique_full_op_id, input_queues)
                 if num_inputs == 0 or agg_op_not_ready or join_op_not_ready:
                     continue
 
@@ -242,6 +247,18 @@ class PipelinedSingleThreadExecutionStrategy(ExecutionStrategy):
                     # update the progress manager
                     self.progress_manager.incr(unique_full_op_id, num_inputs=1, num_outputs=num_outputs, total_cost=record_set.get_total_cost())
 
+                # if this is a join operator with no more inputs to process, then finish it
+                if isinstance(operator, JoinOp) and operator.how in ("left", "right", "outer"):
+                    join_op_upstream_finished = self._upstream_ops_finished(plan, unique_full_op_id, input_queues)
+                    join_input_queues_empty = all(len(inputs) == 0 for inputs in input_queues[unique_full_op_id].values())
+                    if join_op_upstream_finished and join_input_queues_empty and not operator.finished:
+                        # process the join one last time with final=True to handle any left/right/outer join logic
+                        record_set, num_inputs_processed = operator([], [], final=True)
+                        records.extend(record_set.data_records)
+                        record_op_stats.extend(record_set.record_op_stats)
+                        num_outputs += sum(record._passed_operator for record in record_set.data_records)
+                        operator.set_finished()
+
                 # update plan stats
                 plan_stats.add_record_op_stats(unique_full_op_id, record_op_stats)
 

palimpzest/query/generators/generators.py

@@ -296,9 +296,9 @@ class Generator(Generic[ContextType, InputType]):
 
         return field_answers
 
-    def __call__(self, candidate: DataRecord, fields: dict[str, FieldInfo] | None, right_candidate: DataRecord | None = None, json_output: bool=True, **kwargs) -> GenerationOutput:
-        """Take the input record (`candidate`), generate the output `fields`, and return the generated output."""
-        logger.debug(f"Generating for candidate {candidate} with fields {fields}")
+    def __call__(self, candidate: DataRecord | list[DataRecord], fields: dict[str, FieldInfo] | None, right_candidate: DataRecord | None = None, json_output: bool=True, **kwargs) -> GenerationOutput:
+        """Take the input record(s) (`candidate`), generate the output `fields`, and return the generated output."""
+        logger.debug(f"Generating for candidate(s) {candidate} with fields {fields}")
 
         # fields can only be None if the user provides an answer parser
         fields_check = fields is not None or "parse_answer" in kwargs
@@ -338,7 +338,7 @@ class Generator(Generic[ContextType, InputType]):
                     reasoning_effort = "minimal" if self.reasoning_effort is None else self.reasoning_effort
                     completion_kwargs = {"reasoning_effort": reasoning_effort, **completion_kwargs}
             if self.model.is_vllm_model():
-                completion_kwargs = {"api_base": self.api_base, **completion_kwargs}
+                completion_kwargs = {"api_base": self.api_base, "api_key": os.environ.get("VLLM_API_KEY", "fake-api-key") **completion_kwargs}
             completion = litellm.completion(model=self.model_name, messages=messages, **completion_kwargs)
             end_time = time.time()
             logger.debug(f"Generated completion in {end_time - start_time:.2f} seconds")
@@ -405,15 +405,17 @@ class Generator(Generic[ContextType, InputType]):
 
         # pretty print prompt + full completion output for debugging
         completion_text = completion.choices[0].message.content
-        prompt = ""
+        prompt, system_prompt = "", ""
         for message in messages:
+            if message["role"] == "system":
+                system_prompt += message["content"] + "\n"
             if message["role"] == "user":
                 if message["type"] == "text":
                     prompt += message["content"] + "\n"
                 elif message["type"] == "image":
-                    prompt += "<image>\n"
+                    prompt += "<image>\n" * len(message["content"])
                 elif message["type"] == "input_audio":
-                    prompt += "<audio>\n"
+                    prompt += "<audio>\n" * len(message["content"])
         logger.debug(f"PROMPT:\n{prompt}")
         logger.debug(Fore.GREEN + f"{completion_text}\n" + Style.RESET_ALL)
 

palimpzest/query/operators/__init__.py

@@ -2,6 +2,10 @@ from palimpzest.query.operators.aggregate import AggregateOp as _AggregateOp
 from palimpzest.query.operators.aggregate import ApplyGroupByOp as _ApplyGroupByOp
 from palimpzest.query.operators.aggregate import AverageAggregateOp as _AverageAggregateOp
 from palimpzest.query.operators.aggregate import CountAggregateOp as _CountAggregateOp
+from palimpzest.query.operators.aggregate import MaxAggregateOp as _MaxAggregateOp
+from palimpzest.query.operators.aggregate import MinAggregateOp as _MinAggregateOp
+from palimpzest.query.operators.aggregate import SemanticAggregate as _SemanticAggregate
+from palimpzest.query.operators.aggregate import SumAggregateOp as _SumAggregateOp
 from palimpzest.query.operators.convert import ConvertOp as _ConvertOp
 from palimpzest.query.operators.convert import LLMConvert as _LLMConvert
 from palimpzest.query.operators.convert import LLMConvertBonded as _LLMConvertBonded
@@ -47,7 +51,7 @@ from palimpzest.query.operators.logical import (
     Project as _Project,
 )
 from palimpzest.query.operators.logical import (
-    RetrieveScan as _RetrieveScan,
+    TopKScan as _TopKScan,
 )
 from palimpzest.query.operators.mixture_of_agents import MixtureOfAgentsConvert as _MixtureOfAgentsConvert
 from palimpzest.query.operators.mixture_of_agents import MixtureOfAgentsFilter as _MixtureOfAgentsFilter
@@ -55,11 +59,11 @@ from palimpzest.query.operators.physical import PhysicalOperator as _PhysicalOpe
 from palimpzest.query.operators.project import ProjectOp as _ProjectOp
 from palimpzest.query.operators.rag import RAGConvert as _RAGConvert
 from palimpzest.query.operators.rag import RAGFilter as _RAGFilter
-from palimpzest.query.operators.retrieve import RetrieveOp as _RetrieveOp
 from palimpzest.query.operators.scan import MarshalAndScanDataOp as _MarshalAndScanDataOp
 from palimpzest.query.operators.scan import ScanPhysicalOp as _ScanPhysicalOp
 from palimpzest.query.operators.split import SplitConvert as _SplitConvert
 from palimpzest.query.operators.split import SplitFilter as _SplitFilter
+from palimpzest.query.operators.topk import TopKOp as _TopKOp
 
 LOGICAL_OPERATORS = [
     _LogicalOperator,
@@ -72,12 +76,12 @@ LOGICAL_OPERATORS = [
     _LogicalJoinOp,
     _LimitScan,
     _Project,
-    _RetrieveScan,
+    _TopKScan,
 ]
 
 PHYSICAL_OPERATORS = (
     # aggregate
-    [_AggregateOp, _ApplyGroupByOp, _AverageAggregateOp, _CountAggregateOp]
+    [_AggregateOp, _ApplyGroupByOp, _AverageAggregateOp, _CountAggregateOp, _MaxAggregateOp, _MinAggregateOp, _SemanticAggregate, _SumAggregateOp]
     # convert
     + [_ConvertOp, _NonLLMConvert, _LLMConvert, _LLMConvertBonded]
     # critique and refine
@@ -100,8 +104,8 @@ PHYSICAL_OPERATORS = (
     + [_ProjectOp]
     # rag
     + [_RAGConvert, _RAGFilter]
-    # retrieve
-    + [_RetrieveOp]
+    # top-k
+    + [_TopKOp]
     # split
     + [_SplitConvert, _SplitFilter]
 )