Skip to content

perf(spanner): optimize chunk processing and row decoding in PartialResultStream #9336

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

Open
olavloite wants to merge 1 commit into
base: main
Choose a base branch
from
+372 −42
Open

perf(spanner): optimize chunk processing and row decoding in PartialResultStream #9336

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 92 additions & 42 deletions handwritten/spanner/src/partial-result-stream.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,15 +367,17 @@ export class PartialResultStream extends Transform implements ResultEvents {
private _options: RowOptions;
private _pendingValue?: p.IValue;
private _pendingValueForResume?: p.IValue;
private _values: p.IValue[];
private _rowValues?: Value[];
private _valueIndex: number;
private _valueIndexForResume?: number;
private _numPushFailed = 0;
private _isFirstChunk = true;
constructor(options = {}) {
super({objectMode: true});

this._destroyed = false;
this._options = Object.assign({maxResumeRetries: 20}, options);
this._values = [];
this._valueIndex = 0;
this._isFirstChunk = true;
}
/**
Expand All @@ -389,6 +391,7 @@ export class PartialResultStream extends Transform implements ResultEvents {
}

this._destroyed = true;
this._rowValues = undefined;

process.nextTick(() => {
if (err) {
Expand Down Expand Up @@ -486,70 +489,117 @@ export class PartialResultStream extends Transform implements ResultEvents {
}

_resetPendingValues() {
if (this._pendingValueForResume) {
this._pendingValue = this._pendingValueForResume;
this._pendingValue = this._pendingValueForResume;
if (this._valueIndexForResume !== undefined) {
this._valueIndex = this._valueIndexForResume;
} else {
delete this._pendingValue;
this._valueIndex = 0;
}
this._rowValues?.fill(undefined, this._valueIndex);
}
Comment thread
olavloite marked this conversation as resolved.

/**
* Manages any chunked values.
* Manages stream chunks, chunked value merging across chunk boundaries,
* row assembly, and resume token checkpoints.
*
* @private
* Processing follows 5 distinct stages:
* 1. Single-chunk fast path: If the entire stream response is in this first chunk,
* decode rows directly and bypass incremental chunk buffering.
* 2. Pending value merge: If the previous chunk ended with an incomplete chunked
* value, merge it with the incoming continuation value at chunkValues[0].
* 3. Chunked value hold: If this chunk ends with a chunked value, hold the tail
* value (chunkValues[numValues - 1]) in `this._pendingValue` for the next chunk.
* 4. Complete value decoding: Iterate from `startIndex` to `endIndex`, decoding
* and adding values into row buffers via `_addValue`.
* 5. Checkpoint tracking: If this chunk contains a resume token, save the pending
* chunked value and the current column write cursor (`_valueIndex`) for retries.
*
* @param {object} chunk The partial result set.
* @private
* @param {google.spanner.v1.PartialResultSet} chunk The partial result set.
* @returns {boolean} Whether downstream can accept more data.
*/
private _addChunk(chunk: google.spanner.v1.PartialResultSet): boolean {
const isFirstChunk = this._isFirstChunk;
this._isFirstChunk = false;

// Fast path for single-chunk stream responses:
// Stage 1: Fast path for single-chunk stream responses:
if (isFirstChunk && chunk.last && !chunk.chunkedValue) {
return this._addSingleChunk(chunk);
}

const chunkValues = chunk.values;
const numValues = chunkValues.length;
const values: Value[] = new Array(numValues);
for (let i = 0; i < numValues; i++) {
values[i] = GrpcService.decodeValue_(chunkValues[i]);
}
let startIndex = 0;
let endIndex = numValues;
let canAcceptMore = true;

// If we have a chunk to merge, merge the values now.
if (this._pendingValue) {
const currentField = this._values.length % this._fields.length;
// Stage 2: Merge pending chunked value from the previous chunk with the
// incoming continuation value at chunkValues[0].
if (this._pendingValue && numValues > 0) {
const currentField = this._valueIndex;
const field = this._fields[currentField];
const headValue = this._pendingValue;
const continuationValue = GrpcService.decodeValue_(chunkValues[0]);
const merged = PartialResultStream.merge(
field.type as google.spanner.v1.Type,
this._pendingValue,
values.shift(),
headValue,
continuationValue,
);

values.unshift(...merged);
delete this._pendingValue;
// We consumed chunkValues[0] as the continuation of the pending value.
startIndex = 1;

// If this chunk only had 1 value and is still chunked, the last element
// of merged remains pending for the next chunk.
let mergedCount = merged.length;
if (numValues === 1 && chunk.chunkedValue) {
mergedCount--;
this._pendingValue = merged[mergedCount];
} else {
this._pendingValue = undefined;
}

for (let i = 0; i < mergedCount; i++) {
if (!this._addValue(merged[i]) && canAcceptMore) {
canAcceptMore = false;
this.emit('paused');
}
}
}

// If the chunk is chunked, store the last value for merging with the next
// chunk to be processed.
if (chunk.chunkedValue) {
this._pendingValue = values.pop();
if (_hasResumeToken(chunk)) {
this._pendingValueForResume = this._pendingValue;
// Stage 3: If this chunk ends with a chunked value, hold the tail value for
// merging with the next chunk instead of decoding it into the current row now.
if (chunk.chunkedValue && numValues > 0) {
if (numValues > 1 || !startIndex) {
endIndex = numValues - 1;
this._pendingValue = GrpcService.decodeValue_(
chunkValues[numValues - 1],
);
}
} else if (_hasResumeToken(chunk)) {
delete this._pendingValueForResume;
}

let res = true;
const len = values.length;
for (let i = 0; i < len; i++) {
res = this._addValue(values[i]) && res;
if (!res) {
// Stage 4: Decode in-place and push complete values into row buffers.
for (let i = startIndex; i < endIndex; i++) {
const value = GrpcService.decodeValue_(chunkValues[i]);
if (!this._addValue(value) && canAcceptMore) {
canAcceptMore = false;
this.emit('paused');
}
}
return res;

// Stage 5: If this chunk contains a resume token, record the checkpoint state.
// Both the pending chunked value and the row column index are preserved so that
// a retry can resume at the exact column position.
if (_hasResumeToken(chunk)) {
if (chunk.chunkedValue) {
this._pendingValueForResume = this._pendingValue;
} else {
this._pendingValueForResume = undefined;
}
this._valueIndexForResume = this._valueIndex;
}

return canAcceptMore;
}

/**
Expand Down Expand Up @@ -587,18 +637,20 @@ export class PartialResultStream extends Transform implements ResultEvents {
* @param {*} value The complete value.
*/
private _addValue(value: Value): boolean {
const values = this._values;
if (!this._rowValues) {
this._rowValues = new Array(this._fields.length);
}
Comment thread
olavloite marked this conversation as resolved.

values.push(value);
this._rowValues[this._valueIndex++] = value;

if (values.length !== this._fields.length) {
if (this._valueIndex !== this._fields.length) {
return true;
}

this._values = [];
this._valueIndex = 0;

return this.push(
formatRow(this._fields, this._decoders, values, this._options),
formatRow(this._fields, this._decoders, this._rowValues, this._options),
);
}

Expand Down Expand Up @@ -898,9 +950,7 @@ export function partialResultStream(
};

const makeRequest = (): void => {
if (isDefined(lastResumeToken) && lastResumeToken.length > 0) {
partialRSStream._resetPendingValues();
}
partialRSStream._resetPendingValues();
lastRequestStream = requestFn(lastResumeToken);
lastRequestStream.on('end', endListener);
errorListener = (err: grpc.ServiceError) => {
Expand Down
Loading
Loading