Create Initial Semantic Group By Operation

[mdr223]

The goal of this issue is to create a semantic group by operation such that users can execute an operation like the following:

# assume there is a directory of product-reviews/ laid out like:
# product-reviews/
# |
# |-- review1.txt
# |-- review2.txt
# | ...
# |-- review100.txt
#
# the following creates a dataset with default fields `contents` and `filename` for each review
ds = pz.TextFileDataset(id="reviews", dir="product-reviews/")

# this operation computes groups of complaints and then returns the count of reviews which belong to each complaint 
ds = ds.sem_groupby(gby_fields=['complaint'], agg_fields=['contents'], agg_funcs=['count'])

# run and print output
out = ds.run()
print(out.to_df()

# should output something like:
# 
#                       complaint       count
# 0               size not right            57
# 1    did not match description            33
# 2                not ergonomic            10

Here is a high-level overview of the changes which need to be made (although this may not be an exhaustive list)

1. Add sem_groupby() operation to pz.Dataset

In src/core/data/dataset.py you will need to add a function for sem_groupby() to the pz.Dataset class. (Classes like pz.TextFileDataset inherit from this base class, so defining the sem_groupby() function here will make this available to all Dataset classes).

The function will need to have the following signature:

def sem_groupby(gby_fields: list[str], agg_fields: list[str], agg_funcs: list[str]) --> Dataset

At this point, you should pause reading this issue and read the description of #252

This issue should similarly assume that pz.GroupBySig() no longer exists, and should pass these fields directly into the logical GroupByAggregate() operator. You can begin to implement the sem_groupby() function by copy-pasting the groupby() function, but be aware that you will likely need to pull the logic from pz.GroupBySig.output_schema into this function so that you can pass an output_schema to the GroupByAggregate() operator.

2. Update GroupByAggregate() to No Longer Accept pz.GroupyBySig

You will likely need to pass the gby_fields, agg_fields, and agg_funcs directly into the logical operator's constructor (in src/core/data/dataset.py::sem_groupby(), and update any internal logic to use these fields anywhere GroupByAggregate previously accessed the pz.GroupBySig object (in src/query/operators/logical.py).

For GroupByAggregate.get_logical_id_params() and GroupByAggregate.get_logical_op_params(), make sure that the param dictionaries which are returned by each function contain the new gby_fields, agg_fields, and agg_funcs parameters that you provide to __init__().

3. (The Fun Part!) Create Your Semantic Group By Implementation

In src/query/operators/aggregate.py, you should create a new class which contains your implementation of the semantic group by. I've pasted a template for you to get started with below:

class SemanticGroupByOp(AggregateOp):
    """
    Implementation of a GroupBy operator. This operator groups records by a set of fields
    and applies a function to each group. The group_by_sig object contains the fields to
    group by and the aggregation functions to apply to each group.
    """
    def __init__(self, gby_fields: list[str], agg_fields: list[str], agg_funcs: list[str], *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.gby_fields = gby_fields
        self.agg_fields = agg_fields
        self.agg_funcs = agg_funcs

    def __str__(self):
        op = super().__str__()
        op += f"    Group-by Fields: {self.gby_fields}\n"
        op += f"    Agg. Fields: {self.agg_fields}\n"
        op += f"    Agg. Funcs: {self.agg_funcs}\n"
        return op

    def get_id_params(self):
        id_params = super().get_id_params()
        return {"gby_fields": self.gby_fields, "agg_fields": self.agg_fields, "agg_funcs": self.agg_funcs, **id_params}

    def get_op_params(self):
        op_params = super().get_op_params()
        return {"gby_fields": self.gby_fields, "agg_fields": self.agg_fields, "agg_funcs": self.agg_funcs, **op_params}

    def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
        # for now, assume the groupby outputs NAIVE_EST_NUM_GROUPS 
        return OperatorCostEstimates(
            cardinality=NAIVE_EST_NUM_GROUPS,
            time_per_record=0,
            cost_per_record=0,
            quality=1.0,
        )

    def __call__(self, candidates: list[DataRecord]) -> DataRecordSet:
        # TODO: you will be given a list of `candidates` which are input `DataRecords`
        #             each `DataRecord` will have the values `dr.contents` and `dr.filename`.
        #
        #             1. Use some method (perhaps start with asking an LLM) to compute the unique group(s) across all `candidates` according to the `gby_fields`.
        #                (`gby_fields` may be a singleton, but it could also be a list of two or more fields (e.g. `["state", "city"]`), if it helps, just assume it's a singleton for now).
        #            2. Use some method (again perhaps asking an LLM) to assign each input `candidate` to one of the unique group(s) based on its `agg_fields`.
        #            3. Apply the `agg_funcs` over each group of records (e.g. if `agg_funcs` is just `["count"]`, then count the number of records in each group. For now let's assume that `agg_funcs` is a singleton and the only aggregation method is "count").
        #           4. For now, return a `list[dict]`, where each `dict` in the list contains the `gby_fields` and aggregate for one unique group.
        #                For example, you might return [{"complaint": "size not right", "count": 57}, {"complaint": "did not match description", "count": 33}, {"complaint": "not ergonomic", "count": 10}]

Final Notes

Once you have accomplished 1.-3., let's check-in to talk about the final steps (plugging this into the optimizer) as well as unit tests we can run to make sure this solution is working!

In order to test that your code works, I would suggest running the following script:

import pandas as pd
import palimpzest as pz
from palimpzest.query.operators.aggregate import SemanticGroupByOp

# create list of candidates from text file dataset
ds = pz.TextFileDataset(id="reviews", dir="product-reviews/")
output = ds.run()
candidates = [dr for dr in output]

# create instance of your physical operator
sem_group_by_op = SemanticGroupByOp(gby_fields=['complaint'], agg_fields=['contents'], agg_funcs=['count'])
grouped_output = sem_group_by_op(candidates)
print(pd.DataFrame(grouped_output))

[mdr223]

First and foremost -- fantastic job on the first portion of this issue! You also anticipated some next steps by using the Generator for LLM calls w/cost-tracking and marshaling output records and stats into a DataRecordSet. Great work!

Thanks to that effort, there are only a few more details which need to be ironed out before this MVP can work end-to-end.

4. Add an ImplementationRule for your SemanticGroupBy

In src/palimpzest/query/optimizer/rules.py, you will need to add an implementation rule for your new SemanticGroupBy operator.

In the Slack message I sent you, I shared a lecture which gives a high-level overview of a Cascades query optimizer. The lecture introduces the concept of an Implementation rule, which substitutes a logical expression in a logical query (sub)plan for one or more physical expressions. Implementing one of these rules requires implementing two functions:

1. def matches_pattern(cls, logical_expression: LogicalExpression) -> bool: which takes a logical expression and returns a T/F determination of whether this rule can be used to substitute in for the logical expression
2. def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]: which takes a logical expression (that passes matches_pattern()) and returns the unique set of physical expressions which can implement that logical expression.

Here is an example of the implementation rule for the SemanticAggregate operator:

class SemanticAggregateRule(ImplementationRule):
    """
    Substitute a logical expression for a SemanticAggregate with an llm physical implementation.
    """

    @classmethod
    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
        is_match = isinstance(logical_expression.operator, Aggregate) and logical_expression.operator.agg_str is not None
        logger.debug(f"SemanticAggregateRule matches_pattern: {is_match} for {logical_expression}")
        return is_match

    @classmethod
    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
        logger.debug(f"Substituting SemanticAggregateRule for {logical_expression}")

        # create variable physical operator kwargs for each model which can implement this logical_expression
        models = [model for model in runtime_kwargs["available_models"] if cls._model_matches_input(model, logical_expression) and not model.is_llama_model()]
        no_reasoning = runtime_kwargs["reasoning_effort"] in [None, "minimal", "low"]
        variable_op_kwargs = [
            {
                "model": model,
                "prompt_strategy": PromptStrategy.AGG_NO_REASONING if model.is_reasoning_model() and no_reasoning else PromptStrategy.AGG,
                "reasoning_effort": runtime_kwargs["reasoning_effort"]
            }
            for model in models
        ]

        return cls._perform_substitution(logical_expression, SemanticAggregate, runtime_kwargs, variable_op_kwargs)

In your case, your logical expression should:

1. only pattern match for logical expressions whose operator is a GroupByAggregate
2. substitute them for the physical SemanticGroupBy operator that you just created, with a variable_op_kwargs that looks similar to the one shown above.

5. Create Distinction Between Semantic and Non-Semantic GroupBys

There is one final detail which we will need to tackle before we can merge your PR to dev (and then main). Currently, there is no way to distinguish between semantic and non-semantic groupbys from just looking at the logical (and physical) operators.

groupby and sem_groupby both create logical GroupByAggregate operators, which are defined by their input / output schema, as well as the groupby fields, agg fields, and agg funcs. However, when we perform our pattern matching in the implementation rule above, we need to make sure that our SemanticGroupBy operator is only applied to groupbys which require an LLM to compute the group(s)!

If we only check:

is_match = isinstance(logical_expression.operator, GroupByAggregate)

Then this will match for both semantic and non-semantic group by operators. The solution is to create an additional piece of metadata in the logical GroupByAggregate operator which distinguishes between semantic and non-semantic groupbys. For now, it would suffice to add a boolean flag is_semantic which could be set to True when we create the GroupByAggregate in datasets.py. We can then update our pattern matching check to ensure that the logical_expression.operator is a GroupByAggregate and its is_semantic flag is set to True.

6. Create New Implementation Rule for Non-Semantic GroupBys

This brings us to our final sticking point. Before your code changes, we only had non-semantic groupbys. As a result, we used a default rule which performed deterministic substitutions between logical and physical operators in cases where there was a one-to-one relationship between the two (see the BasicSubstitutionRule defined at the bottom of src/palimpzest/query/optimizer/rules.py):

class BasicSubstitutionRule(ImplementationRule):
    """
    For logical operators with a single physical implementation, substitute the
    logical expression with its physical counterpart.
    """

    LOGICAL_OP_CLASS_TO_PHYSICAL_OP_CLASS_MAP = {
        BaseScan: MarshalAndScanDataOp,
        # ComputeOperator: SmolAgentsCompute,
        SearchOperator: SmolAgentsSearch, # SmolAgentsManagedSearch, # SmolAgentsCustomManagedSearch
        ContextScan: ContextScanOp,
        Distinct: DistinctOp,
        LimitScan: LimitScanOp,
        Project: ProjectOp,
        GroupByAggregate: ApplyGroupByOp,
    }

    @classmethod
    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
        logical_op_class = logical_expression.operator.__class__
        is_match = logical_op_class in cls.LOGICAL_OP_CLASS_TO_PHYSICAL_OP_CLASS_MAP
        logger.debug(f"BasicSubstitutionRule matches_pattern: {is_match} for {logical_expression}")
        return is_match

    @classmethod
    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
        logger.debug(f"Substituting BasicSubstitutionRule for {logical_expression}")
        physical_op_class = cls.LOGICAL_OP_CLASS_TO_PHYSICAL_OP_CLASS_MAP[logical_expression.operator.__class__]
        return cls._perform_substitution(logical_expression, physical_op_class, runtime_kwargs)

Now that our logical GroupByAggregate operator can be matched to a (non-semantic) ApplyGroupByOp or a SemanticGroupBy, we need to remove the mapping in the BasicSubstitutionRule which automatically implements GroupByAggregate with ApplyGroupByOp.

After removing this mapping from LOGICAL_OP_CLASS_TO_PHYSICAL_OP_CLASS_MAP we finally need to add a new implementation rule for non-semantic group bys. I will let you think about how to implement this rule (hint: it should be fairly trivial given the existing BasicSubstitutionRule and the creation of your is_semantic flag).

Final Notes

As always, I'm available on Slack if you have questions about this issue and / or the steps you need to take!