perf: cache row count on VirtualDataFrame - #1903
Open
Charisn wants to merge 1 commit into
Open
Conversation
rows_count ran a SELECT COUNT(*) against the source on every read, and the dataframe serializer reads it on every prompt build. That is one full-table count per chat turn, repeated on every retry: a turn that exhausted max_retries=3 issued four of them. The count only fills in the dimensions attribute of the serialized table, and that attribute was already wrong. columns_count measured the empty backing frame, so tables serialized as dimensions="98765432x0". Cache the count in _rows_count, following the _head pattern already in the class, and read columns_count from the schema, falling back to the cached head when the schema declares no columns. Neither path adds a query, since serialize calls head() regardless.
dosubot
Bot
added
the
size:S
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Problem
VirtualDataFrame.rows_countruns aSELECT COUNT(*)against the source on every read, andDataframeSerializer.serializereads it every time a prompt is built. Thehead()method three lines above it caches into_head;rows_countdid not cache at all.So every chat turn pays a full-table count, and so does every retry. Tracing turns against a stubbed Postgres source:
COUNT(*)max_retries=3: 4COUNT(*)At a 250 ms round trip, those counts were 41% of wall time. On a local view over a 2M-row parquet the count took 26.9 ms, 29% of the turn, since it re-reads both parquet files and redoes the join on each call. Against a large remote table it costs considerably more than that, and it repeats per turn.
The count exists only to fill in the
dimensionsattribute of the serialized table, and that attribute was already half broken.columns_countmeasures the empty backing frame, so a 98,765,432-row table serialized as:Fix
Cache the count in
_rows_count, following the_headpattern already in the class. Overridecolumns_countto read the schema, falling back to the cached head when the schema declares no columns. Neither path adds a query, becauseserializecallshead()regardless.Follow-up turns now issue no
COUNT(*), anddimensionsreports both axes.Cache lifetime
_rows_countlives as long as_headdoes, so a long-livedVirtualDataFramekeeps reporting the row count from its first read. That is the same staleness the head cache already accepts, which is why the count is cached the same way rather than given its own invalidation rule.Tests
Six tests in
tests/unit_tests/dataframe/test_virtual_dataframe.py. Five of them fail without this change; the sixth checks that two dataframes from one loader keep separate counts, guarding against a shared cache later.Full unit suite: 539 passed, 40 skipped.