Skip to content

feat: track rows processed during model evaluation #5162

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 25 commits into
base: main
Choose a base branch
from
Open

feat: track rows processed during model evaluation #5162

wants to merge 25 commits into from

Conversation

treysp
Copy link
Contributor

@treysp treysp commented Aug 14, 2025 •
edited
Loading

Captures the number of rows processed (and potentially other information) during a model evaluation for display to the user.

Example running sushi on Bigquery (includes bytes processed for each model batch):

Screenshot 2025-08-18 at 4 21 44 PM

Context

  • Python client libraries conformant to the DB-API 2.0 spec provide a cursor.row_count property that contains the number of rows affected by a DML command execution (e.g., INSERT, CTAS)
  • Variations
    • Some clients provide the property, but they don't actually track rows so always return -1
    • Snowflake conforms other than CTAS
    • Bigquery python API necessary to handle all scenarios

Implementation context

  • SQLMesh makes many execute calls that should not be tracked
  • In most cases, SQLMesh DML actions take place via the engine adapter _execute method, so we read the row_count property there immediately after cursor.execute
  • Snapshot evaluation may take place concurrently, where a single snapshot evaluation occurs within one thread

Determining what to track

  • We should only track specific operations, so execute does NOT track rows by default
  • Some engine adapter methods should always track, so they pass track_row_count=True in the execute call within them
  • Other engine methods vary in their need to track so expose a track_row_count arg
    • Example: engine_adapter.insert_append
      • Should almost always be tracked during snapshot evaluation
      • Should never be tracked when called from state sync

Tracker implementation

  • Row count values are collected in the QueryExecutionContext dataclass
    • Accumulates row counts if called multiple times
    • Ignores row counts of -1 (default value when cursor.row_count not populated)
  • Trackers implemented as context manager
  • Engine adapter _execute calls single record_execution method, which determines which tracker is active and should be used (if any)

@treysp treysp changed the title feat: add rows tracking feat: track rows processed during model evaluation Aug 14, 2025
@treysp treysp force-pushed the trey/rows-tracker branch 2 times, most recently from d722c67 to c3a8760 Compare August 14, 2025 20:34
erindru
erindru reviewed Aug 15, 2025
View reviewed changes
Copy link
Collaborator

@erindru erindru left a comment •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice start, I think I follow the logic. I've left some comments for your consideration.

I'm not sure about the track_row_count parameter on the engine adapter methods. I'd like to understand more about the undesirable queries that were getting picked up - perhaps that indicates we were being too coarse with the statistics gathering / wrapping in the wrong place

@eakmanrq
Copy link
Contributor

Building on @erindru's feedback, I wonder if in general this tracking logic could get pushed to the SnapshotEvaluator layer. Right now the engine adapter methods return nothing but if they returned row count then it seems like the SnapshotEvaluator could then determine what to do with that information.

I don't think though I deeply understand this problem and this comes from a quick review and my general intuition.

@treysp treysp force-pushed the trey/rows-tracker branch 5 times, most recently from 1e405ce to 9377241 Compare August 15, 2025 23:23
erindru
erindru reviewed Aug 17, 2025
View reviewed changes
@erindru
Copy link
Collaborator

erindru commented Aug 18, 2025 •
edited
Loading

Right now the engine adapter methods return nothing but if they returned row count then it seems like the SnapshotEvaluator could then determine what to do with that information.

If we pass row counts around through the entire call stack between the scheduler and the engine adapter, I believe it starts getting very messy with the number of method signature changes.

The scheduler is where we want to display the information, but it's typically on a different thread and reasonably far removed from EngineAdapter._execute where the relevant queries are executed. In addition, the EngineAdapter layer is a driver that deliberately has no idea what Snapshot it's executing because that's a higher level concept.

This design that Trey is proposing means that we don't have to change the return type of the engine adapter methods or have to worry about passing things back up the stack / rely on every component in the chain between the Scheduler and EngineAdapter._execute to pass things back correctly.

It essentially uses a combination of thread-local + global state to bypass the call stack. It emits row counts in EngineAdapter._execute, assigns meaning to them in SnapshotEvaluator and allows the information to be queried back in Scheduler so it can be outputted to the console for the user to see - without making every method in between return some kind of Stats object.

imo this is a much simpler / less invasive implementation

@treysp treysp force-pushed the trey/rows-tracker branch from 9377241 to 6957330 Compare August 18, 2025 21:11
izeigerman
izeigerman reviewed Aug 18, 2025
View reviewed changes
@treysp treysp force-pushed the trey/rows-tracker branch from 6957330 to d4d5afa Compare August 18, 2025 22:35
izeigerman
izeigerman reviewed Aug 19, 2025
View reviewed changes
@treysp treysp force-pushed the trey/rows-tracker branch 14 times, most recently from d00e0b9 to a2bdfdc Compare August 20, 2025 23:18
@treysp treysp force-pushed the trey/rows-tracker branch 2 times, most recently from f072873 to 97db434 Compare August 22, 2025 20:58
treysp added 24 commits August 22, 2025 16:05
@treysp
Add rows tracking
8549479
@treysp
Add flag for supporting row tracking to engine adapters
66827cd
@treysp
Run cloud tests in CI
c66b001
@treysp
Add engine support flags
5b38dd9
@treysp
Use threading.local() instead of locks
8872f6c
@treysp
Move all tracking into snapshot evaluator, remove seed tracker class
335feb5
@treysp
Remove 'processed' and 'inserted' from console output
1a659ca
@treysp
Add BQ support and track bytes processed
cc80ba5
@treysp
Remove seed tracking, have snapshot evaluator own tracker instance
f866937
@treysp
Move tracker class into snapshot module
4d92f31
@treysp
Fix circular import
144a947
@treysp
Handle snowflake lack of CTAS tracking
233a09d
@treysp
Fix tests and snowflake regex
7939e6b
@treysp
Change tracking arg name to track_rows_processed
2648a47
@treysp
Report 0 rows correctly
c2e2dc3
@treysp
Add databricks support
ac38f0f
@treysp
Remove time travel test for cloud engines, handle pyspark DFs in dbx
0075b55
@treysp
Fix rebase
263c68c
@treysp
Seeds are now handled in evaluator
6a17974
@treysp
Fix rebase
173650a
@treysp
Handle snowflake table already exists
347836d
@treysp
Query info schema for snowflake CTAS num rows
74ae584
@treysp
Remove databricks, snowflake metadata calls
7dca8da
@treysp
Add snowflake test
7ac3118
@treysp treysp force-pushed the trey/rows-tracker branch from 97db434 to b66fc6a Compare August 22, 2025 21:24
@treysp treysp force-pushed the trey/rows-tracker branch from b66fc6a to eebb26e Compare August 22, 2025 21:38
@treysp treysp marked this pull request as ready for review August 22, 2025 21:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@erindru erindru erindru left review comments

@izeigerman izeigerman izeigerman left review comments

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@treysp @eakmanrq @erindru @izeigerman