fix: streaming

Author: astandrik

src/services/api/streaming.ts

@@ -1,7 +1,8 @@
-import {isSessionChunk} from '../../store/reducers/query/utils';
+import {isQueryResponseChunk, isSessionChunk} from '../../store/reducers/query/utils';
 import type {Actions, TracingLevel} from '../../types/api/query';
 import type {QuerySyntax, StatisticsMode} from '../../types/store/query';
 import type {StreamingChunk} from '../../types/store/streaming';
+import {StreamDataBatcher} from '../../utils/StreamDataBatcher';
 import {BINARY_DATA_IN_PLAIN_TEXT_DISPLAY} from '../../utils/constants';
 import {MultipartStreamParser} from '../parsers/parseMultipartStream';
 import {settingsManager} from '../settings';
@@ -35,13 +36,25 @@ export class StreamingAPI extends BaseYdbAPI {
         );
 
         let traceId: string | undefined = '';
+        const batcher = new StreamDataBatcher();
 
         const streamParser = new MultipartStreamParser((chunk: StreamingChunk) => {
             if (isSessionChunk(chunk)) {
                 const sessionChunk = chunk;
                 sessionChunk.meta.trace_id = traceId;
+                options.onChunk(chunk);
+            } else if (isQueryResponseChunk(chunk)) {
+                // Flush any batched chunks before sending response
+                const batchedChunks = batcher.flush();
+                batchedChunks.forEach((batchedChunk) => options.onChunk(batchedChunk));
+                options.onChunk(chunk);
+            } else {
+                // For stream data chunks, try to batch them
+                const batchedChunk = batcher.addChunk(chunk);
+                if (batchedChunk) {
+                    options.onChunk(batchedChunk);
+                }
             }
-            options.onChunk(chunk);
         });
 
         const queryParams = new URLSearchParams();
@@ -88,6 +101,9 @@ export class StreamingAPI extends BaseYdbAPI {
                 streamParser.processBuffer(value);
             }
         } finally {
+            // Flush any remaining batched chunks
+            const remainingChunks = batcher.flush();
+            remainingChunks.forEach((chunk) => options.onChunk(chunk));
             reader.releaseLock();
         }
     }

src/services/parsers/parseMultipartStream.ts

@@ -96,11 +96,11 @@ export class MultipartStreamParser {
             if (boundaryPos === -1) {
                 // Keep last boundary length bytes in case boundary is split between chunks
                 if (this.buffer.length > this.boundaryBytes.length) {
-                    const newBuffer = new Uint8Array(this.boundaryBytes.length);
-                    newBuffer.set(
+                    const partialBoundaryBuffer = new Uint8Array(this.boundaryBytes.length);
+                    partialBoundaryBuffer.set(
                         this.buffer.slice(this.buffer.length - this.boundaryBytes.length),
                     );
-                    this.buffer = newBuffer;
+                    this.buffer = partialBoundaryBuffer;
                 }
                 this.bufferPos = 0;
                 break;
@@ -134,9 +134,9 @@ export class MultipartStreamParser {
 
         // Compact buffer if we've processed some data
         if (this.bufferPos > 0) {
-            const newBuffer = new Uint8Array(this.buffer.length - this.bufferPos);
-            newBuffer.set(this.buffer.slice(this.bufferPos));
-            this.buffer = newBuffer;
+            const compactedBuffer = new Uint8Array(this.buffer.length - this.bufferPos);
+            compactedBuffer.set(this.buffer.slice(this.bufferPos));
+            this.buffer = compactedBuffer;
             this.bufferPos = 0;
         }
     }

src/store/reducers/query/query.ts

@@ -168,27 +168,27 @@ const slice = createSlice({
                         };
                     }
                 }
-                // Convert and append new rows
+                // Process rows in batches
                 if (
                     rows &&
                     rows.length > 0 &&
                     state.result.data.resultSets?.[resultIndex]?.columns
                 ) {
-                    const currentRows = state.result.data.resultSets[resultIndex].result;
-                    let currentIndex = currentRows?.length || 1;
+                    const currentRows = state.result.data.resultSets[resultIndex].result || [];
+                    let currentIndex = currentRows.length || 1;
 
+                    // Process all rows at once instead of one by one
                     const indexedRows = rows.map((row) => [currentIndex++, ...row]);
-
                     const formattedRows = parseResult(
                         indexedRows,
                         state.result.data.resultSets[resultIndex].columns,
                     );
 
-                    if (state.result.data.resultSets[resultIndex].result) {
-                        state.result.data.resultSets[resultIndex].result.push(...formattedRows);
-                    } else {
-                        state.result.data.resultSets[resultIndex].result = formattedRows;
-                    }
+                    // Update result set with all formatted rows at once
+                    state.result.data.resultSets[resultIndex].result = [
+                        ...(state.result.data.resultSets[resultIndex].result || []),
+                        ...formattedRows,
+                    ];
                 }
             } else if (isQueryResponseChunk(chunk)) {
                 state.result.isLoading = false;
@@ -259,7 +259,7 @@ interface QueryStats {
     endTime?: string | number;
 }
 
-const DEFAULT_STREAM_CHUNK_SIZE = 2;
+const DEFAULT_STREAM_CHUNK_SIZE = 10000;
 const DEFAULT_CONCURRENT_RESULTS = false;
 
 export const queryApi = api.injectEndpoints({

src/utils/StreamDataBatcher.ts

@@ -0,0 +1,61 @@
+import {isStreamDataChunk} from '../store/reducers/query/utils';
+import type {StreamDataChunk, StreamingChunk} from '../types/store/streaming';
+
+const DEFAULT_BATCH_SIZE = 1000;
+
+export class StreamDataBatcher {
+    private mergedChunks: Record<number, StreamDataChunk> = {}; // resultIndex -> mergedChunk
+    private readonly batchSize: number;
+
+    constructor(options = {batchSize: DEFAULT_BATCH_SIZE}) {
+        this.batchSize = options.batchSize;
+    }
+
+    addChunk(chunk: StreamingChunk): StreamingChunk | null {
+        if (!isStreamDataChunk(chunk)) {
+            return chunk;
+        }
+
+        const {
+            meta: {result_index: resultIndex},
+        } = chunk;
+
+        if (!this.mergedChunks[resultIndex]) {
+            this.mergedChunks[resultIndex] = {
+                ...chunk,
+                result: {
+                    rows: [],
+                    columns: chunk.result.columns,
+                },
+            };
+        }
+
+        const mergedChunk = this.mergedChunks[resultIndex];
+
+        const rows = mergedChunk.result.rows || [];
+        if (chunk.result.rows) {
+            rows.push(...chunk.result.rows);
+        }
+
+        // If we have enough rows, return the merged chunk and reset
+        if (rows.length >= this.batchSize) {
+            const result = {
+                ...mergedChunk,
+                result: {
+                    rows: mergedChunk.result.rows?.slice() || [],
+                    columns: mergedChunk.result.columns,
+                },
+            };
+            mergedChunk.result.rows = [];
+            return result;
+        }
+
+        return null; // Not enough rows yet
+    }
+
+    flush(): StreamingChunk[] {
+        const results = Object.values(this.mergedChunks);
+        this.mergedChunks = {};
+        return results;
+    }
+}