[Rust] Expose Arrow batch encoded byte sizes - #790
Open
flaviofcruz wants to merge 1 commit into
Open
Conversation
Add ZerobusArrowStream::take_offset_details, returning Some(OffsetDetails) with the encoded wire byte size of the batch at an offset and the cumulative wire bytes sent through it. Callers report an accurate bytes-sent metric without re-serialising the RecordBatch (issue #779); None when no size is recorded. The batch channel now carries each batch's OffsetId alongside the RecordBatch, so the Flight encoder records byte sizes keyed by the durable OffsetId with no per-connection offset reconstruction. Per-offset and cumulative sizes accumulate every transmission, counting retransmits; a batch encoded as several IPC messages (e.g. a dictionary batch) sums against the same offset. State lives in a bounded BatchStatsTracker; reads are consume-once. Additive: wait_for_offset/flush and the FFI/JNI/PyO3 surfaces are unchanged. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Flavio Cruz <flavio.cruz@databricks.com>
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.
What changes are proposed in this pull request?
Add
ZerobusArrowStream::take_offset_details, returningSome(OffsetDetails)with the encoded wire byte size of the batch at an offset and the cumulative wire bytes sent through it. Callers report an accurate bytes-sent metric without re-serialising the RecordBatch (issue #779); None when no size is recorded.To make this happen, we carry the OffsetId alongside the RecordBatch so that we can encode the size. The stats recorder keeps only a certain amount of stats so if they are not consumed, the user won't be able to recover them, so they have to be consumed after waiting for the offset.
This change will be used to track the vector's
databricks_zerobusbytes sent which is helpful to debug what is happening.An alternative design that was considered was to introduce a separate
wait_for_offsetcalledwait_for_offset_with_detailswhich waits for the offset and then immediately returns the stats. The downside is that we end up with multiple similar functions but can do that if you feel that's better.How is this tested?
Added several unit tests.