feat: add some initial optimizer functions to run atop query plans #2669
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,97 @@ | ||
| package query | ||
|
|
||
| // TypedOptimizerFunc is a function that transforms an iterator of a specific type T | ||
| // into a potentially optimized iterator. It returns the optimized iterator, a boolean | ||
| // indicating whether any optimization was performed, and an error if the optimization failed. | ||
| // | ||
| // The type parameter T constrains the function to operate only on specific iterator types, | ||
| // providing compile-time type safety when creating typed optimizers. | ||
| type TypedOptimizerFunc[T Iterator] func(it T) (Iterator, bool, error) | ||
|
|
||
| // OptimizerFunc is a type-erased wrapper around TypedOptimizerFunc[T] that can be | ||
| // stored in a homogeneous list while maintaining type safety at runtime. | ||
| type OptimizerFunc func(it Iterator) (Iterator, bool, error) | ||
|
|
||
| // WrapOptimizer wraps a typed TypedOptimizerFunc[T] into a type-erased OptimizerFunc. | ||
| // This allows optimizer functions for different concrete iterator types to be stored | ||
| // together in a heterogeneous list. | ||
| func WrapOptimizer[T Iterator](fn TypedOptimizerFunc[T]) OptimizerFunc { | ||
| return func(it Iterator) (Iterator, bool, error) { | ||
| if v, ok := it.(T); ok { | ||
| return fn(v) | ||
| } | ||
| return it, false, nil | ||
| } | ||
| } | ||
|
|
||
| // StaticOptimizations is a list of optimization functions that can be safely applied | ||
| // to any iterator tree without needing runtime information or context. | ||
| var StaticOptimizations = []OptimizerFunc{ | ||
| RemoveNullIterators, | ||
| ElideSingletonUnionAndIntersection, | ||
|
Contributor
There was a problem hiding this comment. Had to google what "elide" meant 😅 Now that I looked at the code I would have called this "collapse". But i'm not going to die on this hill |
||
| WrapOptimizer(PushdownCaveatEvaluation), | ||
| } | ||
|
|
||
| // ApplyOptimizations recursively applies a list of optimizer functions to an iterator | ||
| // tree, transforming it into an optimized form. | ||
| // | ||
| // The function operates bottom-up, optimizing leafs and subiterators first, and replacing the | ||
| // subtrees up to the top, which it then returns. | ||
| // | ||
| // Parameters: | ||
| // - it: The iterator tree to optimize | ||
| // - fns: A list of optimizer functions to apply | ||
| // | ||
| // Returns: | ||
| // - The optimized iterator (which may be the same as the input if no optimizations applied) | ||
| // - A boolean indicating whether any changes were made | ||
| // - An error if any optimization failed | ||
| func ApplyOptimizations(it Iterator, fns []OptimizerFunc) (Iterator, bool, error) { | ||
| var err error | ||
| origSubs := it.Subiterators() | ||
| changed := false | ||
| if len(origSubs) != 0 { | ||
| // Make a copy of the subiterators slice to avoid mutating the original iterator | ||
| subs := make([]Iterator, len(origSubs)) | ||
| copy(subs, origSubs) | ||
|
|
||
| subChanged := false | ||
| for i, subit := range subs { | ||
| newit, ok, err := ApplyOptimizations(subit, fns) | ||
| if err != nil { | ||
| return nil, false, err | ||
| } | ||
| if ok { | ||
| subs[i] = newit | ||
| subChanged = true | ||
| } | ||
| } | ||
| if subChanged { | ||
| changed = true | ||
| it, err = it.ReplaceSubiterators(subs) | ||
| if err != nil { | ||
| return nil, false, err | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Apply each optimizer to the current iterator | ||
| // If any optimizer transforms the iterator, recursively optimize the new tree | ||
| for _, fn := range fns { | ||
| newit, fnChanged, err := fn(it) | ||
| if err != nil { | ||
| return nil, false, err | ||
| } | ||
| if fnChanged { | ||
| // The iterator was transformed - recursively optimize the new tree | ||
| // to ensure all optimizations are fully applied | ||
| optimizedIt, _, err := ApplyOptimizations(newit, fns) | ||
| if err != nil { | ||
| return nil, false, err | ||
| } | ||
| // Return true for changed since we did transform the iterator | ||
| return optimizedIt, true, nil | ||
| } | ||
| } | ||
| return it, changed, nil | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,93 @@ | ||
| package query | ||
|
|
||
| import ( | ||
| core "github.com/authzed/spicedb/pkg/proto/core/v1" | ||
| "github.com/authzed/spicedb/pkg/spiceerrors" | ||
| ) | ||
|
|
||
| // PushdownCaveatEvaluation pushes caveat evaluation down through certain composite iterators | ||
| // to allow earlier filtering and better performance. | ||
| // | ||
| // This optimization transforms: | ||
| // | ||
| // Caveat(Union[A, B]) -> Union[Caveat(A), B] (if only A contains the caveat) | ||
| // Caveat(Union[A, B]) -> Union[Caveat(A), Caveat(B)] (if both contain the caveat) | ||
| // | ||
| // The pushdown does NOT occur through IntersectionArrow iterators, as they have special | ||
| // semantics that require caveat evaluation to happen after the intersection. | ||
| func PushdownCaveatEvaluation(c *CaveatIterator) (Iterator, bool, error) { | ||
| // Don't push through IntersectionArrow | ||
| if _, ok := c.subiterator.(*IntersectionArrow); ok { | ||
| return c, false, nil | ||
| } | ||
|
|
||
| // Don't push down if the subiterator is already a CaveatIterator | ||
| // This prevents infinite recursion | ||
| if _, ok := c.subiterator.(*CaveatIterator); ok { | ||
| return c, false, nil | ||
| } | ||
|
|
||
| // Get the subiterators of the child | ||
| subs := c.subiterator.Subiterators() | ||
| if len(subs) == 0 { | ||
| // No subiterators to push down into (e.g., leaf iterator) | ||
| return c, false, nil | ||
| } | ||
|
|
||
| // Find which subiterators contain relations with this caveat | ||
| newSubs := make([]Iterator, len(subs)) | ||
| changed := false | ||
| for i, sub := range subs { | ||
| if containsCaveat(sub, c.caveat) { | ||
| // Wrap this subiterator with the caveat | ||
| newSubs[i] = NewCaveatIterator(sub, c.caveat) | ||
| changed = true | ||
| } else { | ||
| // Leave unchanged | ||
| newSubs[i] = sub | ||
| } | ||
| } | ||
|
|
||
| if !changed { | ||
| return c, false, nil | ||
| } | ||
|
|
||
| // Replace the subiterators in the child iterator | ||
| newChild, err := c.subiterator.ReplaceSubiterators(newSubs) | ||
| if err != nil { | ||
| return nil, false, err | ||
| } | ||
|
|
||
| // Return the child without the caveat wrapper | ||
| return newChild, true, nil | ||
| } | ||
|
|
||
| // containsCaveat checks if an iterator tree contains a RelationIterator | ||
| // that references the given caveat. | ||
| func containsCaveat(it Iterator, caveat *core.ContextualizedCaveat) bool { | ||
| found := false | ||
| _, err := Walk(it, func(node Iterator) (Iterator, error) { | ||
| if rel, ok := node.(*RelationIterator); ok { | ||
| if relationContainsCaveat(rel, caveat) { | ||
| found = true | ||
| } | ||
| } | ||
| return node, nil | ||
| }) | ||
| if err != nil { | ||
| spiceerrors.MustPanicf("should never error -- callback contains no errors, but linters must always check") | ||
| } | ||
|
|
||
| return found | ||
| } | ||
|
|
||
| // relationContainsCaveat checks if a RelationIterator's base relation | ||
| // has a caveat that matches the given caveat name. | ||
| func relationContainsCaveat(rel *RelationIterator, caveat *core.ContextualizedCaveat) bool { | ||
| if rel.base == nil || caveat == nil { | ||
| return false | ||
|
Contributor
There was a problem hiding this comment. this code path isn't covered with unit tests |
||
| } | ||
|
|
||
| // Check if the relation has this caveat | ||
| return rel.base.Caveat() == caveat.CaveatName | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
FYI, having to copy the variable hasn't been necessary since go 1.22 https://go.dev/wiki/LoopvarExperiment