Merge streamed result set parts in execute_with_retries#862
Merged
Conversation
Stream parts that share a result_set_index are concatenated back into a single ResultSet, so execute_with_retries returns one result set per SELECT regardless of its size instead of one per stream part.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #862 +/- ##
==========================================
+ Coverage 82.18% 82.25% +0.07%
==========================================
Files 99 99
Lines 12683 12714 +31
Branches 1235 1241 +6
==========================================
+ Hits 10423 10458 +35
+ Misses 1804 1801 -3
+ Partials 456 455 -1
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
This PR fixes QuerySessionPool.execute_with_retries (sync + async) returning multiple ResultSet objects for a single logical SELECT when the Query service splits one result set across multiple streamed response parts sharing the same result_set_index. It introduces a utility to merge those parts back into one result set per index.
Changes:
- Aggregate streamed
ResultSetparts byresult_set_indexinexecute_with_retries(sync and async). - Add
aggregate_result_sets_by_indexinydb.convertand unit tests for its merging behavior (rows, schema, truncated, Arrow data). - Add integration tests validating large SELECT results are returned as a single merged result set; document the fix in
CHANGELOG.md.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| ydb/query/pool.py | Uses result-set aggregation in sync execute_with_retries. |
| ydb/aio/query/pool.py | Uses result-set aggregation in async execute_with_retries. |
| ydb/convert.py | Adds aggregate_result_sets_by_index to merge streamed parts sharing an index. |
| ydb/query/pool_test.py | Adds unit tests validating aggregation semantics (rows/schema/truncation/data). |
| tests/query/test_query_session_pool.py | Adds integration test ensuring large result sets are merged for sync pool. |
| tests/aio/query/test_query_session_pool.py | Adds integration test ensuring large result sets are merged for async pool. |
| CHANGELOG.md | Adds user-facing changelog entry for the fix. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
alex268
approved these changes
Jul 20, 2026
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
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.
execute_with_retriesmaterializes the query stream, but the query service splits one logical result set into several stream parts that share aresult_set_index. Those parts were returned as separateResultSetobjects, so a large SELECT came back as multiple result sets instead of one.Row-format (
VALUE) parts that share an index are now concatenated back into a single result set, so a large SELECT is returned as one result set instead of one per stream chunk.ARROWparts are left untouched: each already carries its own schema and record-batchdataand is independently decodable, so they stay as separate result sets rather than splicing opaque bytes together.The stream is aggregated in a single pass (sync and async) without an intermediate list.
Closes: #863