@mailwoman/corpus 9.4.0 → 10.1.0

package/lib/parquet/readers.ts

@@ -0,0 +1,86 @@
+/**
+ * @copyright Sister Software
+ * @license AGPL-3.0
+ * @author Teffen Ellis, et al.
+ * @file Promise-based parquet reads. The iterator-based reads live in `./streams`.
+ *
+ *   {@linkcode readParquetRows} throws when the file is missing. {@linkcode tryReadParquetRows} returns `null`.
+ *   Both throw for an unreadable file, so a caller cannot mistake a read failure for an empty file.
+ */
+
+import { pathExists } from "@mailwoman/core/fs/readers"
+import type { PathBuilderLike } from "path-ts"
+
+import { escapeSQLString, openDuckDB } from "#parquet/duckdb"
+import { openParquetRowStream, type ParquetRowStreamOptions } from "#parquet/streams"
+
+/**
+ * Reads every row of a parquet file into memory.
+ *
+ * Use {@linkcode openParquetRowStream} for a file too large to hold in memory.
+ *
+ * @throws When the file is missing or `columns` includes a column the file lacks.
+ */
+export async function readParquetRows<T>(path: PathBuilderLike, options: ParquetRowStreamOptions = {}): Promise<T[]> {
+	if (!(await pathExists(path))) {
+		throw new Error(`No parquet file at ${path}`)
+	}
+
+	return await Array.fromAsync(openParquetRowStream<T>(path, options))
+}
+
+/**
+ * Reads every row of a parquet file, or returns `null` when the file is missing.
+ *
+ * @throws When the file cannot be parsed or `columns` includes a column the file lacks.
+ */
+export async function tryReadParquetRows<T>(
+	path: PathBuilderLike,
+	options: ParquetRowStreamOptions = {}
+): Promise<T[] | null> {
+	if (!(await pathExists(path))) return null
+
+	return await Array.fromAsync(openParquetRowStream<T>(path, options))
+}
+
+/**
+ * Counts the rows of a parquet file from its metadata, so the cost stays flat as the row count grows.
+ *
+ * The function opens its own DuckDB connection.
+ *
+ * @throws When the file is missing.
+ */
+export async function countParquetRows(path: PathBuilderLike): Promise<number> {
+	if (!(await pathExists(path))) {
+		throw new Error(`No parquet file at ${path}`)
+	}
+
+	using db = await openDuckDB()
+
+	const result = await db.runAndReadAll(`SELECT count(*) AS n FROM read_parquet('${escapeSQLString(path.toString())}')`)
+
+	const rows = result.getRowObjects() as Array<{ n: unknown }>
+
+	return Number(rows[0]?.n ?? 0)
+}
+
+/**
+ * Returns a parquet file's column names in file order.
+ *
+ * The query uses `DESCRIBE`, which lists logical columns.
+ * `parquet_schema` lists physical leaves instead, where a list column such as
+ * `tokens` appears only as its `element` child.
+ *
+ * @throws When the file is missing.
+ */
+export async function parquetColumnNames(path: PathBuilderLike): Promise<string[]> {
+	if (!(await pathExists(path))) {
+		throw new Error(`No parquet file at ${path}`)
+	}
+
+	using db = await openDuckDB()
+
+	const result = await db.runAndReadAll(`DESCRIBE SELECT * FROM read_parquet('${escapeSQLString(path.toString())}')`)
+
+	return (result.getRowObjects() as Array<{ column_name: unknown }>).map((row) => String(row.column_name))
+}

package/lib/parquet/schema.ts

@@ -0,0 +1,214 @@
+/**
+ * @copyright Sister Software
+ * @license AGPL-3.0
+ * @author Teffen Ellis, et al.
+ * @file The corpus parquet schema defines its columns and logical types. It also defines the projection from a `LabeledRow`.
+ *
+ *   Every file in this family agrees on this one definition. A column added here and nowhere else fails to compile
+ *   against {@linkcode ParquetRow}, which keeps the writer, the reader and the manifest's `schema` key from drifting
+ *   apart.
+ *
+ *   Compression is `snappy` throughout. PyArrow uses it by default. It is also the standard ML-corpus codec. A reader
+ *   outside this repository therefore opens the files without configuration.
+ */
+
+import type { LabeledRow } from "#types"
+
+/**
+ * Row groups are written at this cadence within a parquet file.
+ */
+export const ROW_GROUP_SIZE = 50_000
+
+/**
+ * Rows per parquet file when a writer's caller names no limit.
+ *
+ * A file is closed here rather than grown, because `writeParquetFile` materializes
+ * one Arrow table and Arrow's list builder raises before a whole source fits:
+ * 4,228,212 rows raised where 1,603,143 wrote.
+ * The manifest records the value in force as `rows_per_slice`.
+ */
+export const ROWS_PER_FILE = 1_000_000
+
+/**
+ * Snappy is the codec selected for corpus parquet files.
+ */
+export const PARQUET_COMPRESSION = "SNAPPY"
+
+export interface ParquetFieldDefinition {
+	// oxlint-disable-next-line unicorn/text-encoding-identifier-case -- Parquet logical type name.
+	type: "UTF8" | "INT32"
+	compression: typeof PARQUET_COMPRESSION
+	repeated?: boolean
+	optional?: boolean
+}
+
+export type ParquetSchemaDefinition<T> = Record<Extract<keyof T, string>, ParquetFieldDefinition>
+
+/**
+ * A single Parquet row shape.
+ *
+ * The index signature allows callers to retain source fields before projection.
+ *
+ * Optional fields are represented as null in the Arrow table and read back as null.
+ */
+export interface ParquetRow {
+	raw: string
+	tokens: readonly string[]
+	labels: readonly string[]
+	span_starts: readonly number[]
+	span_ends: readonly number[]
+	span_tags: readonly string[]
+	country: string
+	locale?: string | null
+	source: string
+	source_id: string
+	corpus_version: string
+	license: string
+	register?: string | null
+	surface: string
+	recipe?: string | null
+	base_source_id?: string | null
+	[key: string]: unknown
+}
+
+/**
+ * Column names emitted into every parquet file.
+ *
+ * Matches `ParquetRow`.
+ */
+export const PARQUET_COLUMNS = [
+	"raw",
+	"tokens",
+	"labels",
+	"span_starts",
+	"span_ends",
+	"span_tags",
+	"country",
+	"locale",
+	"source",
+	"source_id",
+	"corpus_version",
+	"license",
+	"register",
+	"surface",
+	"recipe",
+	"base_source_id",
+] as const
+
+/**
+ * The DuckDB type each column is read and written as.
+ *
+ * Paired with {@linkcode PARQUET_COLUMNS} so a `read_json` column map and a `copy`
+ * select list are built from one list rather than two that can disagree.
+ */
+export const PARQUET_COLUMN_TYPES: Record<(typeof PARQUET_COLUMNS)[number], string> = {
+	raw: "VARCHAR",
+	tokens: "VARCHAR[]",
+	labels: "VARCHAR[]",
+	span_starts: "INTEGER[]",
+	span_ends: "INTEGER[]",
+	span_tags: "VARCHAR[]",
+	country: "VARCHAR",
+	locale: "VARCHAR",
+	source: "VARCHAR",
+	source_id: "VARCHAR",
+	corpus_version: "VARCHAR",
+	license: "VARCHAR",
+	register: "VARCHAR",
+	surface: "VARCHAR",
+	recipe: "VARCHAR",
+	base_source_id: "VARCHAR",
+}
+
+/* oxlint-disable unicorn/text-encoding-identifier-case -- `"UTF8"` below is a ParquetType enum member rather than a text-encoding identifier. A lowercase spelling does not type-check against ParquetSchemaDefinition.
+   The rule cannot distinguish the enum member from a text-encoding identifier. */
+
+/**
+ * Parquet schema for `LabeledRow`.
+ *
+ * Optional fields set `optional: true`.
+ * Repeated UTF8 columns capture the tokens and labels arrays.
+ */
+export const LABELED_ROW_SCHEMA: ParquetSchemaDefinition<ParquetRow> = {
+	raw: { type: "UTF8", compression: PARQUET_COMPRESSION },
+	tokens: { type: "UTF8", repeated: true, compression: PARQUET_COMPRESSION },
+	labels: { type: "UTF8", repeated: true, compression: PARQUET_COMPRESSION },
+	// Char-offset label spans, meaning parallel arrays over `raw` in UTF-16 code units. Each span is [start, end) with an exclusive end, sorted and non-overlapping. INT32 holds a short address string and round-trips as `number` where parquetjs INT64 would surface bigint.
+	span_starts: { type: "INT32", repeated: true, compression: PARQUET_COMPRESSION },
+	span_ends: { type: "INT32", repeated: true, compression: PARQUET_COMPRESSION },
+	span_tags: { type: "UTF8", repeated: true, compression: PARQUET_COMPRESSION },
+	country: { type: "UTF8", compression: PARQUET_COMPRESSION },
+	locale: { type: "UTF8", compression: PARQUET_COMPRESSION, optional: true },
+	source: { type: "UTF8", compression: PARQUET_COMPRESSION },
+	source_id: { type: "UTF8", compression: PARQUET_COMPRESSION },
+	corpus_version: { type: "UTF8", compression: PARQUET_COMPRESSION },
+	license: { type: "UTF8", compression: PARQUET_COMPRESSION },
+	register: { type: "UTF8", compression: PARQUET_COMPRESSION, optional: true },
+	surface: { type: "UTF8", compression: PARQUET_COMPRESSION },
+	recipe: { type: "UTF8", compression: PARQUET_COMPRESSION, optional: true },
+	base_source_id: { type: "UTF8", compression: PARQUET_COMPRESSION, optional: true },
+}
+
+/* oxlint-enable unicorn/text-encoding-identifier-case */
+
+/**
+ * Project a labeled row to the Parquet schema.
+ *
+ * The span triple is required because `alignRow` emits it on every labeled row.
+ * A row without it came from a producer that has not migrated.
+ *
+ * The writer would drop the labels from the file if it wrote that row.
+ * A thrown error identifies the row instead.
+ */
+export function rowToParquet(row: LabeledRow): ParquetRow {
+	const { span_starts, span_ends, span_tags } = row
+
+	if (span_starts === undefined || span_ends === undefined || span_tags === undefined) {
+		throw new Error(
+			`rowToParquet: row is missing the char-offset span triple (#519) — ` +
+				`span_starts=${span_starts !== undefined} span_ends=${span_ends !== undefined} span_tags=${span_tags !== undefined} ` +
+				`(source=${row.source}, source_id=${row.source_id}). ` +
+				`Every parquet-bound row must carry span_starts/span_ends/span_tags; ` +
+				`producers that emit tokens/labels only have not migrated to the v0.5.0 format.`
+		)
+	}
+
+	if (span_starts.length !== span_ends.length || span_starts.length !== span_tags.length) {
+		throw new Error(
+			`rowToParquet: span triple arrays are not parallel — ` +
+				`starts=${span_starts.length} ends=${span_ends.length} tags=${span_tags.length} ` +
+				`(source=${row.source}, source_id=${row.source_id})`
+		)
+	}
+
+	// The runner stamps the adapter's `surface` on every row that omits one, so an absent
+	// value here means the row reached parquet by a path that bypassed it.
+	// A default would record every such row as the publisher's own string.
+	if (!row.surface) {
+		throw new Error(
+			`rowToParquet: row carries no surface ` +
+				`(source=${row.source}, source_id=${row.source_id}). ` +
+				`An adapter declares CorpusAdapter.surface and the runner stamps it; ` +
+				`a recipe writing rows directly sets the field itself.`
+		)
+	}
+
+	return {
+		raw: row.raw,
+		tokens: row.tokens,
+		labels: row.labels,
+		span_starts,
+		span_ends,
+		span_tags,
+		country: row.country,
+		locale: row.locale ?? null,
+		source: row.source,
+		source_id: row.source_id,
+		corpus_version: row.corpus_version,
+		license: row.license,
+		register: row.register ?? null,
+		surface: row.surface,
+		recipe: row.recipe?.recipe ?? null,
+		base_source_id: row.recipe?.base_source_id ?? null,
+	}
+}

package/lib/parquet/shuffle.ts

@@ -0,0 +1,59 @@
+/**
+ * @copyright Sister Software
+ * @license AGPL-3.0
+ * @author Teffen Ellis, et al.
+ * @file A bounded streaming shuffle over the rows headed for parquet, so a row-group is a sample of its source.
+ */
+
+import { shuffleWith } from "@mailwoman/core/random"
+
+/**
+ * Rows held in memory while shuffling, when the caller names no window.
+ *
+ * Five 50,000-row groups span more than one output row-group and reach `wof-admin`'s
+ * whole 13-country set, while a window below `ROW_GROUP_SIZE` buys no wider mix.
+ */
+export const DEFAULT_SHUFFLE_WINDOW = 250_000
+
+/**
+ * Seed for the shuffle when the caller names none, fixed rather than drawn so two builds
+ * of one input write the same row order and a corpus is reproducible from its inputs.
+ */
+export const DEFAULT_SHUFFLE_SEED = 20_260_922
+
+/**
+ * Yield the input shuffled within a sliding window of `windowSize` rows, so memory is
+ * bounded by the window rather than the input; `windowSize <= 1` yields the input unchanged
+ * and consumes no random draw, keeping a no-shuffle build byte-identical.
+ */
+export async function* shuffleWithinWindow<T>(
+	rows: AsyncIterable<T>,
+	random: () => number,
+	windowSize: number = DEFAULT_SHUFFLE_WINDOW
+): AsyncGenerator<T> {
+	if (windowSize <= 1) {
+		yield* rows
+
+		return
+	}
+
+	const buffer: T[] = []
+
+	for await (const row of rows) {
+		if (buffer.length < windowSize) {
+			buffer.push(row)
+
+			continue
+		}
+
+		const index = Math.floor(random() * buffer.length)
+		const chosen = buffer[index]!
+		buffer[index] = row
+
+		yield chosen
+	}
+
+	shuffleWith(buffer, random)
+
+	yield* buffer
+}

package/lib/parquet/span-sql.ts

@@ -0,0 +1,54 @@
+/**
+ * @copyright Sister Software
+ * @license AGPL-3.0
+ * @author Teffen Ellis, et al.
+ * @file The SQL that reads one component's text out of a corpus row, for a query that runs in DuckDB.
+ *
+ *   A parquet row stores component offsets in `span_starts` and `span_ends` arrays.
+ *   A third array, `span_tags`, identifies each component. The text comes from slicing `raw` at those offsets.
+ *   A query that groups or filters by a
+ *   component performs that slice in SQL instead of decoding every row in the process.
+ *
+ *   `holdoutComponents` in `#tools/overlay/split-slice` performs the same read in JavaScript.
+ *   The SQL and JavaScript operations use different string indexing. The difference is described below.
+ */
+
+/**
+ * The text of one component tag, or NULL when the row has no span for it.
+ *
+ * `list_position` is one-based and answers NULL for an absent tag.
+ * `substr` takes a one-based start with a length, so the start is `span_starts[i] + 1`
+ * and the length is the span's width.
+ *
+ * Known limitation and scope.
+ * `alignRow` records a span in UTF-16 code units.
+ *
+ * DuckDB's `substr` counts characters.
+ * The two agree for the Basic Multilingual Plane and diverge by one unit per astral character,
+ * so a row carrying one slices late by the number of astral characters before the span.
+ *
+ * In the first 40 train files of `v0.6.0-register-surface`, 6,982 of 40,000,000
+ * rows (0.0175%) contain an astral character.
+ * Among 11,521,648 US rows, 5,230 (0.0454%) contain one.
+ *
+ * `𐍀𐍂𐍉𐍆𐌹𐌳𐌰𐌹𐌽𐍃, RHODE ISLAND` read its region as `ND`, taken from `ISLAND`.
+ *
+ * So this is sound for an aggregate over millions of rows and unsound for a decision about one row.
+ * A caller that acts per row narrows with this and confirms with `holdoutComponents`, which
+ * slices with JavaScript string semantics and therefore reads the offsets as they were written.
+ */
+export function componentAtSpanSQL(tag: string, rawColumn = "raw"): string {
+	return componentAtIndexSQL(`list_position(span_tags, '${tag}')`, rawColumn)
+}
+
+/**
+ * The same slice against a span index the caller already computed.
+ *
+ * A query can hoist `list_position(span_tags, …)` into a subquery column and read it by name.
+ * A second evaluation in the projection would run the search twice per row.
+ *
+ * Every caveat on {@linkcode componentAtSpanSQL} applies here, because this is the expression it builds.
+ */
+export function componentAtIndexSQL(indexExpression: string, rawColumn = "raw"): string {
+	return `CASE WHEN ${indexExpression} IS NULL THEN NULL ELSE substr(${rawColumn}, span_starts[${indexExpression}] + 1, span_ends[${indexExpression}] - span_starts[${indexExpression}]) END`
+}

package/lib/parquet/streams.ts

@@ -0,0 +1,113 @@
+/**
+ * @copyright Sister Software
+ * @license AGPL-3.0
+ * @author Teffen Ellis, et al.
+ * @file Pull-based parquet reading, the same split `@mailwoman/core/fs` draws between `/readers` and `/streams`.
+ *
+ *   A stream returns immediately and the work happens as the consumer pulls, so a caller can walk a file larger than
+ *   memory. A caller that wants the rows, or only their count, uses `./readers` instead and awaits a value.
+ */
+
+import type { PathBuilderLike } from "path-ts"
+
+import { escapeSQLIdentifier, escapeSQLString, openDuckDB } from "#parquet/duckdb"
+
+/**
+ * DuckDB hands a list column back as `{ items: [...] }`.
+ *
+ * Unwrap it so a row reads the way the schema declares it, recursively,
+ * because a nested list arrives nested the same way.
+ */
+export function normalizeDuckDBValue(value: unknown): unknown {
+	if (value && typeof value === "object" && "items" in value && Array.isArray(value.items)) {
+		return value.items.map(normalizeDuckDBValue)
+	}
+
+	// A DuckDB STRUCT arrives as `{ entries: { field: value } }` rather than as
+	// the plain object its fields describe.
+	// The raw value makes `row.address_levels[0].value` read `undefined` on every row.
+	// A consumer counting values then reports that the file does not populate the column.
+	// Overture's `address_levels` and `sources` are both lists of structs.
+	// Unwrapped values made the audit report all 25,914,431 Italian rows as lacking locality.
+	if (value && typeof value === "object" && "entries" in value && value.entries && typeof value.entries === "object") {
+		return Object.fromEntries(
+			Object.entries(value.entries as Record<string, unknown>).map(([key, entry]) => [key, normalizeDuckDBValue(entry)])
+		)
+	}
+
+	if (Array.isArray(value)) return value.map(normalizeDuckDBValue)
+
+	return value
+}
+
+/**
+ * A row limit has to be a non-negative safe integer before it reaches a `limit` clause,
+ * because it is interpolated rather than bound.
+ */
+export function validateRowLimit(value: number): number {
+	if (!Number.isSafeInteger(value) || value < 0)
+		throw new Error(`Parquet row limit must be a non-negative integer: ${value}`)
+
+	return value
+}
+
+export interface ParquetRowStreamOptions {
+	/**
+	 * Columns to project.
+	 *
+	 * Every column when omitted.
+	 *
+	 * A column absent from the file raises rather than coming back absent: a projection that
+	 * silently drops a field hands the consumer a well-formed row with the field missing.
+	 * The consumer reads that as "this row has no value there" instead of "this file has no such column".
+	 */
+	columns?: readonly string[]
+	/**
+	 * Stop after this many rows.
+	 */
+	limit?: number
+}
+
+/**
+ * Open a row stream over a local parquet file, yielding rows in DuckDB-managed chunks.
+ *
+ * DuckDB opens the path itself and exposes its DataChunks through `fetchChunk()`.
+ * Rows are converted and yielded one chunk at a time, so neither the complete Parquet file
+ * nor the complete result set is copied into JavaScript memory.
+ */
+export async function* openParquetRowStream<T>(
+	path: PathBuilderLike,
+	options: ParquetRowStreamOptions = {}
+): AsyncGenerator<T> {
+	const { columns, limit } = options
+	// Disposed when this generator finishes or a consumer abandons it,
+	// since `for await` calls `.return()` on break.
+	using db = await openDuckDB()
+	const projection = columns?.length ? columns.map(escapeSQLIdentifier).join(", ") : "*"
+	const limitClause = limit === undefined ? "" : ` LIMIT ${validateRowLimit(limit)}`
+	const sql = `SELECT ${projection} FROM read_parquet('${escapeSQLString(path.toString())}')${limitClause}`
+
+	try {
+		const stream = await db.stream(sql)
+		const columnNames = stream.columnNames()
+
+		for (let chunk = await stream.fetchChunk(); chunk && chunk.rowCount > 0; chunk = await stream.fetchChunk()) {
+			const rows = chunk.getRowObjects(columnNames) as Record<string, unknown>[]
+
+			for (const row of rows) {
+				yield Object.fromEntries(Object.entries(row).map(([key, value]) => [key, normalizeDuckDBValue(value)])) as T
+			}
+		}
+	} catch (error) {
+		if (columns && /Binder Error|Referenced column|does not exist/.test(String(error))) {
+			const missing =
+				String(error).match(/(?:Referenced column|field named) ["']([^"']+)["']/)?.[1] ?? columns.join(", ")
+
+			throw new Error(`Parquet projection requested column absent from the file schema: ${missing}`, {
+				cause: error,
+			})
+		}
+
+		throw error
+	}
+}