perf: Utilize ORDER BY LIMIT over ROW_NUMBER where possible #1077
Conversation
product-auto-label
bot
added
the
api: bigquery
| @@ -385,6 +385,44 @@ def common_selection_root( | |||
| return None | |||
|
|
|||
|
|
|||
| def pullup_limit_from_slice( | |||
There was a problem hiding this comment.
Are we able to have a test to verify this process?
product-auto-label
bot
added
size: l
bigframes/core/compile/compiled.py
Outdated
| sql += f"{order_by_clause}\n" | ||
| sql += f"\n{order_by_clause}" | ||
| if limit is not None: | ||
| assert isinstance(limit, int) |
There was a problem hiding this comment.
Should we raise a TypeError here instead of an assertion error?
I don't know how well our public APIs are validating the type of parameters. Not everyone is going to use a type checker.
There was a problem hiding this comment.
Yeah, in general, probably best to assume that wrong types can sometimes reach even the lower levels of the code base. In particular want to be safe about what we are putting into the sql strings. Switched to TypeError here.
bigframes/core/nodes.py
Outdated
| @property | ||
| def is_limit(self) -> bool: | ||
| """Returns whether this is equivalent to a ORDER BY ... LIMIT N.""" | ||
| return (not self.start) and (self.step == 1) and (self.stop is not None) |
There was a problem hiding this comment.
Maybe a TODO to handle the negative version (e.g. "tail") too? One would have to reverse the order twice, but doable.
There was a problem hiding this comment.
yeah, should be doable, if a bit tricky. adding a todo
| return None | ||
| return self.left_child.row_count * self.right_child.row_count | ||
|
|
||
| return None |
There was a problem hiding this comment.
TODO for left/right join with a unique join key?
There was a problem hiding this comment.
Need a bit more than just unique join keys to determine exact cardinality. The join keys on the outer side all need to have a match on the inner side as well. Most of these cases will be handled by implicit joiner.
I think upper bound will also be useful (and possible to generate for every node type except explode), so may add that as well later on.
bigframes/core/nodes.py
Outdated
| @@ -635,12 +681,24 @@ def relation_ops_created(self) -> int: | |||
| return 3 | |||
|
|
|||
| @property | |||
| def supports_fast_head(self) -> bool: | |||
| def fast_offsets(self) -> bool: | |||
| # Fast head is only supported when row offsets are available. | |||
There was a problem hiding this comment.
This comment is out of date now.
There was a problem hiding this comment.
Updated comment to mention clustering
bigframes/core/nodes.py
Outdated
| cast(ex.DerefOp, col.scalar_expression).id.name for col in order_cols | ||
| ) | ||
| cluster_col_ids = self.source.table.cluster_cols | ||
| return order_col_ids == cluster_col_ids |
There was a problem hiding this comment.
I suspect if cluster_col_ids is a prefix of order_col_ids, it'd be fast-ish too, at least if the clustering is selective enough.
There was a problem hiding this comment.
I think its best to be conservative on this. In theory, the important parts of the order ids could be the later parts, not in the clustering. Also, the query engine is binary about this optimization, either on or off. Can maybe run some tests and loosen the requirement later.
| """Attempts to simplify the slice.""" | ||
| row_count = traversals.row_count(node) | ||
| start, stop, step = node.start, node.stop, node.step | ||
| def rewrite_slice(node: nodes.SliceNode): |
There was a problem hiding this comment.
rewrite_slice seems like a good candidate for some unit tests if we don't have some already.
There was a problem hiding this comment.
ok, added some unit tests. Bit rough writing tree tests right now. Might need to setup a little framework if we want to do lots of these. Local engine could also be a way to ensure that tree rewrites do not change result value.
bigframes/session/executor.py
Outdated
| node, bigframes.core.identifiers.ColumnId(offsets_id) | ||
| ) | ||
| predicate = ops.lt_op.as_expr(ex.deref(offsets_id), ex.const(n)) | ||
| plan_w_head = nodes.FilterNode(plan_w_offsets, predicate) | ||
| # Finally, drop the offsets column |
There was a problem hiding this comment.
I think this comment is out of date now.
TrevorBergeron
force-pushed
the
slice_pullup
branch
from
5b27ed7 to
6c580a5
Compare
Probably most impact on notebook benchmarks initially. TPCH uses to_gbq which will not use this optimization until further changes (need to be careful not to use ORDER BY LIMIT when writing to clustered tables) |
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
Fixes #<issue_number_goes_here> 🦕