Skip to content

feat(helpers): streaming bulk-ingest with bounded memory #495

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

Draft
margaretjgu wants to merge 2 commits into
base: main
Choose a base branch
from
+211 −135
Draft

feat(helpers): streaming bulk-ingest with bounded memory #495

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cd0023b
feat: streaming bulk-ingest bounded by flush_bytes * concurrency
margaretjgu Jul 31, 2026
9b2992f
test: fix runCommand helper to handle Node test runner IPC on stdout
margaretjgu Jul 31, 2026
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
322 changes: 195 additions & 127 deletions src/es/helpers/bulk-ingest.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'

Expand Down Expand Up @@ -51,44 +49,23 @@ const inputSchema = z.object({

type BulkIngestInput = z.infer<typeof inputSchema>

/**
* 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<void> {
if (this.count > 0) { this.count--; return }
await new Promise<void>(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. */
Expand All @@ -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,
Expand All @@ -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<void> {
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<void> => {
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<void> => {
if (bufBytes === 0) return
const body = buf
buf = ''
bufBytes = 0
await submitBatch(body)
}

const addDoc = async (docJson: string): Promise<void> => {
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<string | undefined>

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<string, string | number | boolean> }): Promise<JsonValue> => {
const opts = parsed.input!
Expand All @@ -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)
}

Expand Down
Loading
Loading