From cd0023bb329179136e1461e1c3c6f4cf3c3c7ff8 Mon Sep 17 00:00:00 2001 From: margaretjgu Date: Fri, 31 Jul 2026 15:45:47 -0400 Subject: [PATCH 1/2] feat: streaming bulk-ingest bounded by flush_bytes * concurrency --- src/es/helpers/bulk-ingest.ts | 322 ++++++++++++++++++++-------------- 1 file changed, 195 insertions(+), 127 deletions(-) diff --git a/src/es/helpers/bulk-ingest.ts b/src/es/helpers/bulk-ingest.ts index dc54a5d2..1dd75007 100644 --- a/src/es/helpers/bulk-ingest.ts +++ b/src/es/helpers/bulk-ingest.ts @@ -4,20 +4,18 @@ */ import { z } from 'zod' -import { readFileSync } from 'node:fs' +import { createReadStream } from 'node:fs' +import { createInterface } from 'node:readline' +import type { Readable } from 'node:stream' +import { parse as parseCsvStream } from 'csv-parse' import type { EsClient } from '../../lib/es-client.ts' import { defineCommand } from '../../factory.ts' import type { OpaqueCommandHandle, JsonValue } from '../../factory.ts' import { getEsClient } from '../../lib/es-client.ts' import { missingConfigError, transportError } from '../errors.ts' import { - parseInput, - parseCsvInput, - readRawInput, globFiles, - buildBulkNdjsonBody, retryWithBackoff, - runWithConcurrency, ProgressReporter } from './shared.ts' @@ -51,44 +49,23 @@ const inputSchema = z.object({ type BulkIngestInput = z.infer -/** - * Splits an array of documents into batches where each batch's serialized - * size does not exceed the byte threshold. - */ -function splitIntoBatches (docs: unknown[], flushBytes: number): unknown[][] { - const batches: unknown[][] = [] - let currentBatch: unknown[] = [] - let currentSize = 0 - - for (const doc of docs) { - const docSize = JSON.stringify(doc).length + 1 // +1 for newline - if (currentBatch.length > 0 && currentSize + docSize > flushBytes) { - batches.push(currentBatch) - currentBatch = [] - currentSize = 0 - } - currentBatch.push(doc) - currentSize += docSize - } - if (currentBatch.length > 0) { - batches.push(currentBatch) +/** Bounded concurrency via a counting semaphore. */ +class Semaphore { + private count: number + private readonly waiters: Array<() => void> = [] + + constructor (n: number) { this.count = n } + + async acquire (): Promise { + if (this.count > 0) { this.count--; return } + await new Promise(r => this.waiters.push(r)) } - return batches -} -/** Parses raw file content according to the selected source format. */ -function parseByFormat (raw: string, opts: BulkIngestInput): unknown[] { - if (opts.source_format === 'csv') { - const csvColumns = opts.csv_columns != null - ? opts.csv_columns.split(',').map((c) => c.trim()).filter(Boolean) - : undefined - return parseCsvInput(raw, { - ...(opts.csv_delimiter != null && { delimiter: opts.csv_delimiter }), - ...(csvColumns != null && { columns: csvColumns }), - ...(opts.skip_header != null && { skipHeader: opts.skip_header }), - }) + release (): void { + const next = this.waiters.shift() + if (next != null) next() + else this.count++ } - return parseInput(raw) } /** Returns the default glob pattern for the given source format. */ @@ -97,46 +74,6 @@ function defaultGlob (format: SourceFormat): string { return '**/*.{json,ndjson,jsonl}' } -/** Collects documents from the resolved input source. */ -function collectDocuments (opts: BulkIngestInput): { docs: unknown[], filesProcessed: number } { - const { data_file, data_dir } = opts - - if (data_file != null && data_dir != null) { - throw new Error('Provide only one input source: --data-file or --data-dir (not both)') - } - - if (data_dir != null) { - const pattern = opts.glob ?? defaultGlob(opts.source_format) - const recursive = opts.no_recursive !== true - const resolvedPattern = recursive ? pattern : pattern.replace(/^\*\*\//, '') - const files = globFiles(data_dir, resolvedPattern) - if (files.length === 0) { - throw new Error(`No files matched pattern "${resolvedPattern}" in ${data_dir}`) - } - const allDocs: unknown[] = [] - for (const file of files) { - const raw = readFileSync(file, 'utf-8') - allDocs.push(...parseByFormat(raw, opts)) - } - return { docs: allDocs, filesProcessed: files.length } - } - - if (data_file != null) { - const raw = readRawInput(data_file) - if (raw == null || raw.trim().length === 0) { - throw new Error('No input data received from file') - } - return { docs: parseByFormat(raw, opts), filesProcessed: 1 } - } - - // Fall back to stdin - const raw = readRawInput() - if (raw == null || raw.trim().length === 0) { - throw new Error('No input provided. Use --data-file, --data-dir, or pipe data to stdin') - } - return { docs: parseByFormat(raw, opts), filesProcessed: 0 } -} - /** Sends a single bulk batch to Elasticsearch. Returns the count of errors. */ async function sendBatch ( transport: EsClient, @@ -161,6 +98,167 @@ async function sendBatch ( return { errors: errorCount, total } } +/** + * Streams documents from all input sources and sends them as bulk batches. + * + * Peak memory is bounded by flush_bytes * concurrency regardless of input size. + * The semaphore provides backpressure: the producer blocks once concurrency + * slots are exhausted, preventing unbounded batch accumulation. + */ +async function streamBulkIngest ( + opts: BulkIngestInput, + transport: EsClient, + reporter: ProgressReporter +): Promise { + const { flush_bytes, concurrency, retries, retry_delay, index, pipeline, routing } = opts + + const actionLine = JSON.stringify({ + index: { + ...(index != null && { _index: index }), + ...(pipeline != null && { pipeline }), + ...(routing != null && { routing }), + } + }) + + const sem = new Semaphore(concurrency) + const errors: unknown[] = [] + + let buf = '' + let bufBytes = 0 + + const submitBatch = async (body: string): Promise => { + await sem.acquire() + // Fire-and-forget: producer continues reading while this batch is in flight. + retryWithBackoff( + async () => { + const res = await sendBatch(transport, body, index) + if (res.errors > 0 && res.errors === res.total) { + throw new Error(`Bulk batch failed: ${res.errors}/${res.total} errors`) + } + return res + }, + { retries, delay: retry_delay } + ).then(res => { + reporter.report(res.total, res.errors) + }).catch(err => { + errors.push(err) + }).finally(() => sem.release()) + } + + const flush = async (): Promise => { + if (bufBytes === 0) return + const body = buf + buf = '' + bufBytes = 0 + await submitBatch(body) + } + + const addDoc = async (docJson: string): Promise => { + const pair = actionLine + '\n' + docJson + '\n' + buf += pair + bufBytes += pair.length + if (bufBytes >= flush_bytes) await flush() + } + + // Resolve file list + const { data_file, data_dir, source_format } = opts + + if (data_file != null && data_dir != null) { + throw Object.assign(new Error('Provide only one input source: --data-file or --data-dir (not both)'), { code: 'input_error' }) + } + + let filePaths: Array + + if (data_dir != null) { + const pattern = opts.glob ?? defaultGlob(source_format) + const recursive = opts.no_recursive !== true + const resolvedPattern = recursive ? pattern : pattern.replace(/^\*\*\//, '') + const found = globFiles(data_dir, resolvedPattern) + if (found.length === 0) { + throw Object.assign(new Error(`No files matched pattern "${resolvedPattern}" in ${data_dir}`), { code: 'input_error' }) + } + reporter.filesProcessed = found.length + filePaths = found + } else if (data_file != null) { + reporter.filesProcessed = 1 + filePaths = [data_file] + } else { + if (process.stdin.isTTY === true) { + throw Object.assign(new Error('No input provided. Use --data-file, --data-dir, or pipe data to stdin'), { code: 'input_error' }) + } + filePaths = [undefined] + } + + for (const filePath of filePaths) { + const stream: Readable = filePath != null ? createReadStream(filePath, { encoding: 'utf-8' }) : process.stdin + + if (source_format === 'csv') { + const csvColumns = opts.csv_columns != null + ? opts.csv_columns.split(',').map(c => c.trim()).filter(Boolean) + : undefined + const parser = parseCsvStream({ + delimiter: opts.csv_delimiter ?? ',', + columns: csvColumns != null && csvColumns.length > 0 ? csvColumns : true, + from_line: opts.skip_header === true ? 2 : 1, + skip_empty_lines: true, + trim: true, + cast (value) { + if (value === 'true') return true + if (value === 'false') return false + if (value !== '' && !isNaN(Number(value))) return Number(value) + return value + } + }) + stream.pipe(parser) + for await (const record of parser) { + await addDoc(JSON.stringify(record)) + } + } else { + // ndjson / json: line-by-line for NDJSON, buffered fallback for JSON arrays + const rl = createInterface({ input: stream, crlfDelay: Infinity }) + let isJsonArray: boolean | null = null // null = not yet determined + let arrayBuf = '' + + for await (const line of rl) { + const trimmed = line.trim() + if (trimmed.length === 0) continue + + if (isJsonArray === null) { + isJsonArray = trimmed.startsWith('[') + } + + if (isJsonArray) { + arrayBuf += line + '\n' + continue + } + + try { + await addDoc(JSON.stringify(JSON.parse(trimmed))) + } catch { + throw new Error(`Failed to parse NDJSON line: ${trimmed.slice(0, 80)}`) + } + } + + if (isJsonArray === true && arrayBuf.trim().length > 0) { + const parsed: unknown = JSON.parse(arrayBuf) + if (!Array.isArray(parsed)) throw new Error('Expected a JSON array') + for (const doc of parsed) { + await addDoc(JSON.stringify(doc)) + } + } + } + } + + await flush() + + // Drain: acquire all slots to confirm every in-flight batch has finished. + for (let i = 0; i < concurrency; i++) { + await sem.acquire() + } + + if (errors.length > 0) throw errors[0] +} + function createBulkIngestHandler (deps: BulkIngestDeps = defaultDeps) { return async (parsed: { input?: BulkIngestInput; options: Record }): Promise => { const opts = parsed.input! @@ -172,56 +270,26 @@ function createBulkIngestHandler (deps: BulkIngestDeps = defaultDeps) { return missingConfigError(err) } - let docs: unknown[] - let filesProcessed: number - try { - const result = collectDocuments(opts) - docs = result.docs - filesProcessed = result.filesProcessed - } catch (err) { - return { - error: { - code: 'input_error', - message: err instanceof Error ? err.message : String(err) - } - } - } - - if (docs.length === 0) { - return { total: 0, succeeded: 0, failed: 0, retries: 0, elapsed_ms: 0 } - } - - const batches = splitIntoBatches(docs, opts.flush_bytes) - const reporter = new ProgressReporter() - reporter.filesProcessed = filesProcessed - - const { retries, retry_delay, index, pipeline, routing } = opts try { - await runWithConcurrency(batches, opts.concurrency, async (batch) => { - const ndjsonBody = buildBulkNdjsonBody(batch, { index, pipeline, routing }) - - const result = await retryWithBackoff( - async () => { - const res = await sendBatch(transport, ndjsonBody, index) - if (res.errors > 0) { - // Only retry if all items failed (likely a transient cluster issue) - // Partial failures are reported as-is - if (res.errors === res.total) { - throw new Error(`Bulk batch failed: ${res.errors}/${res.total} errors`) - } - } - return res - }, - { retries, delay: retry_delay } - ) - - reporter.report(result.total, result.errors) - return result - }) + await streamBulkIngest(opts, transport, reporter) } catch (err) { - // If retries exhausted, report what we have so far + const code = (err as { code?: string }).code + if (code === 'input_error' || (err instanceof Error && ( + err.message.startsWith('No files matched') || + err.message.startsWith('Provide only one') || + err.message.startsWith('No input provided') || + err.message.startsWith('Failed to parse') || + err.message.startsWith('Expected a JSON') + ))) { + return { + error: { + code: 'input_error', + message: err instanceof Error ? err.message : String(err) + } + } + } return transportError(err) } From 9b2992f0593de445c4133efc79363d8be68987ce Mon Sep 17 00:00:00 2001 From: margaretjgu Date: Fri, 31 Jul 2026 15:45:50 -0400 Subject: [PATCH 2/2] test: fix runCommand helper to handle Node test runner IPC on stdout --- test/es/helpers/bulk-ingest.test.ts | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/test/es/helpers/bulk-ingest.test.ts b/test/es/helpers/bulk-ingest.test.ts index 304f6271..ddc3c721 100644 --- a/test/es/helpers/bulk-ingest.test.ts +++ b/test/es/helpers/bulk-ingest.test.ts @@ -75,17 +75,25 @@ async function runCommand (args: string[], deps: BulkIngestDeps): Promise 0 ? errOutput : stdOutput - if (output.trim().length > 0) { - try { - return JSON.parse(output.trim()) - } catch { - return output.trim() + + if (errOutput.trim().length > 0) { + try { return JSON.parse(errOutput.trim()) } catch { return errOutput.trim() } + } + + for (let i = stdoutChunks.length - 1; i >= 0; i--) { + const t = stdoutChunks[i]!.trim() + if ((t.startsWith('{') || t.startsWith('[')) && t.length > 0) { + try { return JSON.parse(t) } catch {} } } + const stdOutput = stdoutChunks.join('') + if (stdOutput.trim().length > 0) return stdOutput.trim() return undefined }