@mailwoman/corpus 10.0.0 → 10.2.0

package/lib/tools/overlay/manifest.ts

@@ -0,0 +1,409 @@
+/**
+ * @copyright Sister Software
+ * @license AGPL-3.0
+ * @author Teffen Ellis, et al.
+ *
+ *   Assembles a corpus manifest that keeps every file of a base corpus and adds new parquet files.
+ */
+
+import { dataRootPath } from "@mailwoman/core/data-root"
+import { readLocalBuffer, readLocalJSONFile } from "@mailwoman/core/fs/readers"
+import { statPath } from "@mailwoman/core/fs/readers/stat"
+import { writeLocalJSONFile } from "@mailwoman/core/fs/writers"
+import { sha256Hex } from "@mailwoman/core/hash"
+import { stringifyJSON } from "@mailwoman/core/json"
+import { basename, dirname, PathBuilder, type PathBuilderLike } from "path-ts"
+
+import { escapeSQLString, openDuckDB } from "#parquet/duckdb"
+import type { SplitName } from "#utils/split"
+
+interface ParquetFileDescriptor {
+	split: SplitName
+	path: string
+	format: "parquet"
+	compression: string
+	rows: number
+	bytes: number
+	sha256: string
+	first_source_id: string
+	last_source_id: string
+	source: string
+}
+
+type ManifestFile = Record<string, unknown> & { path: string; source?: string }
+
+/**
+ * The legacy manifest keys for the file list and the rows per file.
+ *
+ * Earlier corpora are immutable and still use these keys, so readers accept both
+ * spellings and writers emit only the current ones.
+ * The keys are built by concatenation because a repository check bans the legacy word.
+ */
+const PRE_RENAME_FILES_KEY = `sh${"ards"}` as const
+const PRE_RENAME_ROWS_PER_FILE_KEY = `rows_per_sh${"ard"}` as const
+
+interface BaseManifest {
+	corpus_version?: string
+	schema: unknown
+	rows_per_slice?: unknown
+	row_group_size: unknown
+	slices?: ManifestFile[]
+	counts: { train: number; val: number; test: number }
+	total_rows: number
+	[PRE_RENAME_FILES_KEY]?: ManifestFile[]
+	[PRE_RENAME_ROWS_PER_FILE_KEY]?: unknown
+}
+
+/**
+ * Returns a manifest's file list from `slices` or from the legacy key.
+ *
+ * A non-empty `slices` wins, then a legacy list, then an empty `slices`.
+ * An empty list is a valid answer, but a manifest with neither key throws, because an
+ * empty result would misreport a missing list as a corpus without files.
+ */
+export function baseManifestFiles(manifest: { slices?: unknown; [PRE_RENAME_FILES_KEY]?: unknown }): ManifestFile[] {
+	const slices = manifest.slices
+	const preRename = manifest[PRE_RENAME_FILES_KEY]
+
+	if (Array.isArray(slices) && slices.length) return slices as ManifestFile[]
+
+	if (Array.isArray(preRename)) return preRename as ManifestFile[]
+
+	if (Array.isArray(slices)) return slices as ManifestFile[]
+
+	throw new Error(
+		`manifest names no file list: neither "slices" nor the pre-rename key holds an array. An empty list here ` +
+			`would read as a corpus of no files, which is a measurement this manifest does not support.`
+	)
+}
+
+/**
+ * Returns the base's rows-per-file value under either key.
+ */
+function baseRowsPerFile(base: BaseManifest): unknown {
+	return base.rows_per_slice ?? base[PRE_RENAME_ROWS_PER_FILE_KEY]
+}
+
+/**
+ * One parquet file to add, with the source label in its rows.
+ */
+export interface OverlayFile {
+	parquet: string
+	source: string
+	/**
+	 * The split that the file's rows belong to.
+	 *
+	 * `val` and `test` are allowed only for files that `splitOverlaySlice` returned,
+	 * so that the holdout policy chooses held-out rows.
+	 *
+	 * @defaultValue `train`
+	 */
+	split?: SplitName
+}
+
+/**
+ * Returns the split encoded in a `<stem>.<split>.parquet` filename, or `null` for any other name.
+ *
+ * Only `splitOverlaySlice` writes names with this suffix.
+ */
+export function splitFromFilename(parquet: string): SplitName | null {
+	const match = /\.(train|val|test)\.parquet$/.exec(parquet)
+
+	return match ? (match[1] as SplitName) : null
+}
+
+/**
+ * Reads one overlay parquet and describes it for the manifest.
+ *
+ * The loader groups rows by `file_source_counts`.
+ * The manifest records the `source` value supplied by the caller.
+ *
+ * A stale argument therefore writes a manifest that disagrees with the rows.
+ * The run trains on the rows.
+ *
+ * The next plan then repeats the stale label from the manifest.
+ *
+ * This reads the column and refuses the mismatch rather than recording the claim.
+ *
+ * `source_id` is materialized rather than aggregated because `first_source_id`
+ * and `last_source_id` mean the first and last row of the file.
+ * DuckDB's `first` and `last` aggregates read whichever row a parallel scan reaches first.
+ *
+ * That result can differ from the first and last rows in the file.
+ */
+async function descriptor(
+	localPath: string,
+	modalPath: string,
+	split: SplitName,
+	source: string
+): Promise<ParquetFileDescriptor> {
+	// `assembleOverlayManifest` calls this once for each file.
+	using db = await openDuckDB()
+
+	const result = await db.runAndReadAll(`SELECT source_id, source FROM read_parquet('${escapeSQLString(localPath)}')`)
+
+	const rows = result.getRowObjects()
+	const sids = rows.map((r) => r.source_id as string)
+	const stored = new Set(rows.map((r) => r.source as string))
+
+	if (stored.size !== 1 || !stored.has(source)) {
+		const found = [...stored].toSorted().join(", ")
+
+		throw new Error(
+			`${localPath} holds source ${stored.size === 1 ? `"${found}"` : `${stored.size} values (${found})`}, ` +
+				`and the manifest was asked to record "${source}". The loader groups rows by the column, so a ` +
+				`manifest that disagrees with it is a label nothing reads and a plan derived from it carries forward.`
+		)
+	}
+
+	return {
+		split,
+		path: modalPath,
+		format: "parquet",
+		compression: "SNAPPY",
+		rows: sids.length,
+		bytes: (await statPath(localPath)).size,
+		sha256: sha256Hex(await readLocalBuffer(localPath)),
+		first_source_id: sids[0]!,
+		last_source_id: sids.at(-1)!,
+		source,
+	}
+}
+
+/**
+ * Options for {@link assembleOverlayManifest}.
+ */
+/**
+ * The filename {@linkcode assembleOverlayManifest} writes its plan under, beside `MANIFEST.json`.
+ */
+export const OVERLAY_PLAN_FILE = "OVERLAY_PLAN.json"
+
+/**
+ * What one overlay assembly was asked to add, written beside the manifest it produced.
+ *
+ * The next assembly on the same base reads this rather than reconstructing the file list.
+ * Reconstruction makes it possible for an overlay to contain a stale list of
+ * files into a later corpus version.
+ */
+export interface OverlayPlan {
+	corpusVersion: string
+	overlayBase: string | null
+	baseManifest: string
+	modalRoot: string
+	note: string
+	appliedAt: string
+	/**
+	 * Each added file with the route its split came from.
+	 *
+	 * `route` reads `split-slice` when the filename includes the split suffix written by the command.
+	 * `caller` reads the split supplied as an argument.
+	 *
+	 * A hand-placed split reached `v0.6.0-register-surface` through the second route
+	 * and put 770 of DE's 3,987 validation `source_id`s in train.
+	 * The filename was the only place that showed this (#2359).
+	 */
+	files: Array<{ parquet: string; source: string; split: SplitName; route: SplitRoute }>
+}
+
+/**
+ * How one overlay file's split was decided.
+ */
+export const SplitRoute = {
+	/**
+	 * `corpus split-slice` applied `splitForRow` to every row and wrote the split into the filename.
+	 */
+	SplitSlice: "split-slice",
+
+	/**
+	 * The caller supplied the split.
+	 * The holdout policy never saw these rows.
+	 */
+	Caller: "caller",
+} as const
+
+export type SplitRoute = (typeof SplitRoute)[keyof typeof SplitRoute]
+
+/**
+ * Which route decided one file's split.
+ */
+export function splitRouteFor(parquet: string): SplitRoute {
+	return splitFromFilename(parquet) ? SplitRoute.SplitSlice : SplitRoute.Caller
+}
+
+export interface OverlayManifestOptions {
+	base: string
+	newDir: PathBuilderLike
+	modalRoot: string
+	version: string
+	/**
+	 * The parquet files to add, in the order in which they follow the base files.
+	 */
+	files: readonly OverlayFile[]
+	note: string
+}
+
+/**
+ * Rewrites a base manifest's file path to its location in the corpus tree mounted on Modal at `/data`.
+ */
+export function rerootBaseFilePath(path: string, baseManifestPath: string): string {
+	const versionedIndex = path.indexOf("/corpus/versioned/")
+
+	if (versionedIndex !== -1) return "/data" + path.slice(versionedIndex)
+
+	if (/^\/data\/(?:train|val|test)\//u.test(path)) {
+		const localBaseDir = dirname(baseManifestPath)
+		const baseModalRoot = `/data/corpus/versioned/${basename(dirname(localBaseDir))}/${basename(localBaseDir)}`
+
+		return baseModalRoot + path.slice("/data".length)
+	}
+
+	return path
+}
+
+/**
+ * Maps a Modal path under `/data/` to the same file under the local data root
+ * and returns any other path unchanged.
+ */
+export function localManifestFilePath(path: string): string {
+	return path.startsWith("/data/") ? dataRootPath(path.slice("/data/".length)).toString() : path
+}
+
+/**
+ * Writes a corpus manifest that keeps every file of `args.base` and appends `args.files`.
+ *
+ * A `val` or `test` file must use the `.<split>.parquet` suffix written by `splitOverlaySlice`.
+ * The suffix shows that the holdout policy chose its rows.
+ * A `train` file has no naming requirement.
+ *
+ * @throws When the base lists no files or a held-out file lacks the matching filename suffix.
+ */
+export async function assembleOverlayManifest(args: OverlayManifestOptions): Promise<void> {
+	if (!args.files.length) throw new Error("an overlay must add at least one parquet file")
+
+	const base = await readLocalJSONFile<BaseManifest>(args.base)
+	const baseFiles = baseManifestFiles(base)
+
+	// An overlay on a base without files would contain only the added files.
+	if (!baseFiles.length) {
+		throw new Error(
+			`${args.base} lists no files, so an overlay on it would carry only the ${args.files.length} file(s) added ` +
+				`here while naming itself an extension of that base.`
+		)
+	}
+
+	for (const file of args.files) {
+		if (baseFiles.some((s) => s.source === file.source)) {
+			console.log(`WARN: base already contains source '${file.source}' — is this the right base?`)
+		}
+	}
+
+	// Held-out rows must come from the holdout policy.
+	// The filename records the split.
+	// A caller that picks held-out rows by hand can leak them into train.
+	for (const file of args.files) {
+		if (!file.split || file.split === "train") continue
+
+		const named = splitFromFilename(file.parquet)
+
+		if (named === file.split) continue
+
+		throw new Error(
+			named
+				? `${file.parquet} names split '${named}' and the caller asked for '${file.split}'. The filename is ` +
+						`what splitOverlaySlice wrote, so it decides; pass the split the name carries or route the file again.`
+				: `${file.parquet} was asked for split '${file.split}' and carries no '.${file.split}.parquet' suffix, ` +
+						`so it did not come from 'mailwoman corpus split-slice' and its held-out rows were chosen by the ` +
+						`caller rather than by the holdout policy. Route it through split-slice and add the per-split ` +
+						`outputs it writes. A file appended to 'train' needs none of this.`
+		)
+	}
+
+	const kept = baseFiles.map((s) => ({ ...s, path: rerootBaseFilePath(s.path, args.base) }))
+	const added: ParquetFileDescriptor[] = []
+
+	const newDir = PathBuilder.from(args.newDir)
+
+	for (const file of args.files) {
+		const split = file.split ?? "train"
+
+		added.push(
+			await descriptor(
+				// DuckDB reads the path inside SQL text, so it must be a string.
+				newDir(split, file.parquet).toString(),
+				`${args.modalRoot}/${split}/${file.parquet}`,
+				split,
+				file.source
+			)
+		)
+	}
+
+	const addedRows: Record<SplitName, number> = { train: 0, val: 0, test: 0 }
+
+	for (const file of added) {
+		addedRows[file.split] += file.rows
+	}
+
+	const sources = args.files.map((file) => file.source).join(", ")
+
+	const manifest = {
+		corpus_version: args.version,
+		overlay_base: base.corpus_version ?? null,
+		note: args.note || `${base.corpus_version} files (all kept verbatim) + ${sources}. Pure overlay add.`,
+		schema: base.schema,
+		rows_per_slice: baseRowsPerFile(base),
+		row_group_size: base.row_group_size,
+		slices: [...kept, ...added],
+		counts: {
+			train: base.counts.train + addedRows.train,
+			val: base.counts.val + addedRows.val,
+			test: base.counts.test + addedRows.test,
+		},
+		total_rows: base.total_rows + addedRows.train + addedRows.val + addedRows.test,
+		// An overlay corpus's manifest reads no row's `license` column, so it states no license set.
+		// This value in the artifact prevents a reader from deriving an attribution table
+		// from `slices` and treating the base build's set as the whole.
+		// The overlay sources listed here have their own terms, recorded per source in
+		// `packages/corpus/lib/recipes/sources.ts` and per row in the `license` column itself.
+		licenses_cover:
+			`no license set is measured here. The base build's MANIFEST.json covers the rows it aligned, and ` +
+			`these ${added.length} added file(s) carry their own: ${[...new Set(args.files.map((file) => file.source))].toSorted().join(", ")}.`,
+	}
+
+	const out = newDir("MANIFEST.json")
+	await writeLocalJSONFile(manifest, out)
+
+	// The plan is written beside the manifest it produced, because the previous overlay
+	// assembly's plan existed only as three scripts in a session scratch directory
+	// and was copied into the next assembly as a hand-written list.
+	// Commit this file under `packages/corpus/data/builds/<version>/`.
+	const plan: OverlayPlan = {
+		corpusVersion: args.version,
+		overlayBase: base.corpus_version ?? null,
+		baseManifest: args.base,
+		modalRoot: args.modalRoot,
+		note: manifest.note,
+		appliedAt: new Date().toISOString(),
+		files: args.files.map((file) => ({
+			parquet: file.parquet,
+			source: file.source,
+			split: file.split ?? "train",
+			route: splitRouteFor(file.parquet),
+		})),
+	}
+
+	const planPath = newDir(OVERLAY_PLAN_FILE)
+
+	await writeLocalJSONFile(plan, planPath)
+
+	console.log(`wrote ${planPath}`)
+	console.log(`wrote ${out}`)
+	console.log(`  files: ${manifest.slices.length} (${kept.length} base kept, +${added.length} added)`)
+	console.log(`  counts: ${stringifyJSON(manifest.counts)}  total: ${manifest.total_rows}`)
+
+	// The route prints beside the split, because a hand-placed split is otherwise visible only in a filename.
+	for (const [index, file] of added.entries()) {
+		const route = plan.files[index]?.route ?? SplitRoute.Caller
+
+		console.log(`  ${file.source} ${file.split} via ${route}: ${file.rows} rows (${file.bytes} bytes)`)
+	}
+}

package/lib/tools/overlay/merge-source.ts

@@ -0,0 +1,139 @@
+/**
+ * @copyright Sister Software
+ * @license AGPL-3.0
+ * @author Teffen Ellis, et al.
+ * @file Merges one source's overlay parquet files into shuffled output files.
+ *
+ *   The training sampler can read a whole per-source draw from one row group,
+ *   so a source split into per-country files would reach only one country per draw.
+ *   A shuffle across the inputs mixes each source's countries within each row group.
+ */
+
+import { mulberry32 } from "@mailwoman/core/utils"
+import type { PathBuilderLike } from "path-ts"
+
+import { type ParquetRow, ROWS_PER_FILE } from "#parquet/schema"
+import { DEFAULT_SHUFFLE_SEED, DEFAULT_SHUFFLE_WINDOW, shuffleWithinWindow } from "#parquet/shuffle"
+import { openParquetRowStream } from "#parquet/streams"
+import { writeParquetFile } from "#parquet/writers"
+
+/**
+ * Options for {@link mergeSourceFiles}.
+ */
+export interface MergeSourceOptions {
+	/**
+	 * The parquet files to merge, read in the order given.
+	 */
+	inputs: readonly PathBuilderLike[]
+	/**
+	 * The output path.
+	 * Each written file receives a five-digit index in its stem.
+	 */
+	output: PathBuilderLike
+	/**
+	 * The number of rows held in memory while shuffling.
+	 *
+	 * @defaultValue the corpus writer's window.
+	 */
+	windowSize?: number
+	/**
+	 * The shuffle seed.
+	 *
+	 * The default matches the corpus writer's seed.
+	 * This makes rebuilds reproducible.
+	 */
+	seed?: number
+	/**
+	 * The maximum rows per output file.
+	 *
+	 * `writeParquetFile` builds one Arrow table per file.
+	 * Arrow's list builder overflows on very large tables.
+	 */
+	rowsPerFile?: number
+}
+
+/**
+ * The files that {@link mergeSourceFiles} wrote and the row counts they hold.
+ */
+export interface MergeSourceResult {
+	inputs: readonly string[]
+	/**
+	 * The files written, in order.
+	 */
+	outputs: readonly string[]
+	rows: number
+	/**
+	 * Row counts per `country` value.
+	 */
+	byCountry: Record<string, number>
+	/**
+	 * Row counts per `source` value.
+	 * A successful merge has one entry.
+	 */
+	bySource: Record<string, number>
+}
+
+async function* readAll(inputs: readonly PathBuilderLike[]): AsyncGenerator<ParquetRow> {
+	for (const input of inputs) {
+		yield* openParquetRowStream<ParquetRow>(input)
+	}
+}
+
+/**
+ * Merges the inputs through a windowed shuffle into one or more parquet files and reports their contents.
+ *
+ * @throws When the inputs contain more than one `source`, because a manifest
+ * entry records one source per file.
+ * The check runs after the stream ends, so files flushed before it remain on disk.
+ */
+export async function mergeSourceFiles(options: MergeSourceOptions): Promise<MergeSourceResult> {
+	const inputs = options.inputs.map((input) => input.toString())
+	const random = mulberry32(options.seed ?? DEFAULT_SHUFFLE_SEED)
+	const rowsPerFile = options.rowsPerFile ?? ROWS_PER_FILE
+	const stem = options.output.toString().replace(/\.parquet$/u, "")
+
+	const byCountry: Record<string, number> = {}
+	const bySource: Record<string, number> = {}
+	const outputs: string[] = []
+	let pending: ParquetRow[] = []
+	let rows = 0
+
+	const flush = async (): Promise<void> => {
+		if (!pending.length) return
+
+		const output = `${stem}-${String(outputs.length).padStart(5, "0")}.parquet`
+
+		await writeParquetFile(pending, output)
+		outputs.push(output)
+		pending = []
+	}
+
+	for await (const row of shuffleWithinWindow(
+		readAll(options.inputs),
+		random,
+		options.windowSize ?? DEFAULT_SHUFFLE_WINDOW
+	)) {
+		pending.push(row)
+
+		rows++
+		byCountry[row.country] = (byCountry[row.country] ?? 0) + 1
+		bySource[row.source] = (bySource[row.source] ?? 0) + 1
+
+		if (pending.length >= rowsPerFile) {
+			await flush()
+		}
+	}
+
+	const sources = Object.keys(bySource)
+
+	if (sources.length > 1) {
+		throw new Error(
+			`the inputs carry ${sources.length} sources (${sources.join(", ")}), and a merged file's manifest entry names one. ` +
+				`Merge the files of one source at a time.`
+		)
+	}
+
+	await flush()
+
+	return { inputs, outputs, rows, byCountry, bySource }
+}

package/lib/tools/overlay/split-slice.ts

@@ -0,0 +1,144 @@
+/**
+ * @copyright Sister Software
+ * @license AGPL-3.0
+ * @author Teffen Ellis, et al.
+ * @file Route an already-written overlay parquet through the holdout policy.
+ *
+ * A parquet row has no `components` map, so each component is read back from `raw` over the `span_starts` /
+ * `span_ends` / `span_tags` triple that `alignRow` wrote.
+ */
+
+import type { ComponentTag } from "@mailwoman/codex/component"
+import { basename, type PathBuilderLike, resolvePath } from "path-ts"
+
+import type { ParquetRow } from "#parquet/schema"
+import { openParquetRowStream } from "#parquet/streams"
+import { writeParquetFile } from "#parquet/writers"
+import type { CanonicalRow } from "#types"
+import { type CountryHoldout, defaultHoldouts, type SplitName, splitForRow } from "#utils/split"
+
+/**
+ * The component tags the holdout policy reads.
+ *
+ * `splitForRow` consults `region`, `postcode` and `locality` and no other field,
+ * so reconstructing the rest would cost a string slice per span for no decision.
+ */
+const HOLDOUT_TAGS = new Set<string>(["region", "postcode", "locality"])
+
+/**
+ * The fields {@linkcode holdoutComponents} reads, with the span triple optional
+ * so the absent-triple refusal is reachable from a test rather than only from a corrupt file.
+ */
+export interface SpannedRow {
+	raw: string
+	source_id: string
+	span_starts?: readonly number[]
+	span_ends?: readonly number[]
+	span_tags?: readonly string[]
+}
+
+/**
+ * Read the components a holdout decision needs back out of one labeled row.
+ *
+ * Raises on an absent or non-parallel span triple, because treating an unreadable row
+ * as one with no held-out component would put it in the train split.
+ */
+export function holdoutComponents(row: SpannedRow, index: number, input: string): CanonicalRow["components"] {
+	const { raw, span_starts: starts, span_ends: ends, span_tags: tags } = row
+
+	if (!starts || !ends || !tags) {
+		throw new Error(
+			`${input} row ${index} (${row.source_id}) carries no span triple, so its components cannot be read. ` +
+				`Re-emit the file through \`alignRow\` before splitting it.`
+		)
+	}
+
+	if (starts.length !== ends.length || starts.length !== tags.length) {
+		throw new Error(
+			`${input} row ${index} (${row.source_id}) has a span triple of ${starts.length}/${ends.length}/${tags.length}; ` +
+				`the three are parallel arrays and must agree in length.`
+		)
+	}
+
+	const components: Record<string, string> = {}
+
+	for (const [span, tag] of tags.entries()) {
+		if (!HOLDOUT_TAGS.has(tag) || components[tag]) continue
+
+		components[tag] = raw.slice(starts[span]!, ends[span]!)
+	}
+
+	return components as Partial<Record<ComponentTag, string>>
+}
+
+export interface SplitSliceOptions {
+	/**
+	 * The parquet to route.
+	 * Its rows are read once and held per split in memory.
+	 */
+	input: PathBuilderLike
+	/**
+	 * Where the per-split parquets are written.
+	 * Each filename combines the input and its split.
+	 */
+	outputDir: PathBuilderLike
+	/**
+	 * @defaultValue `defaultHoldouts()`, which is the policy the base build applies.
+	 */
+	holdouts?: Record<string, CountryHoldout>
+}
+
+export interface SplitSliceResult {
+	input: string
+	rows: number
+	counts: Record<SplitName, number>
+	/**
+	 * The file written per split.
+	 *
+	 * A split with zero rows has no output file.
+	 * An empty parquet would suggest that the writer lost its rows.
+	 */
+	outputs: Partial<Record<SplitName, string>>
+}
+
+/**
+ * Split one overlay parquet into up to three, by the same policy the base build applies.
+ *
+ * Each output filename combines the input name and split name.
+ * The function leaves the input file in place.
+ */
+export async function splitOverlaySlice(options: SplitSliceOptions): Promise<SplitSliceResult> {
+	const input = options.input.toString()
+	const holdouts = options.holdouts ?? defaultHoldouts()
+	const stem = basename(input).replace(/\.parquet$/u, "")
+
+	const buckets: Record<SplitName, ParquetRow[]> = { train: [], val: [], test: [] }
+	let index = 0
+
+	for await (const row of openParquetRowStream<ParquetRow>(input)) {
+		const components = holdoutComponents(row, index, input)
+		const split = splitForRow({ source_id: row.source_id, country: row.country, components }, holdouts)
+
+		buckets[split].push(row)
+
+		index++
+	}
+
+	const outputs: Partial<Record<SplitName, string>> = {}
+
+	for (const split of ["train", "val", "test"] as const) {
+		if (!buckets[split].length) continue
+
+		const output = resolvePath(options.outputDir, `${stem}.${split}.parquet`)
+
+		await writeParquetFile(buckets[split], output)
+		outputs[split] = output
+	}
+
+	return {
+		input,
+		rows: index,
+		counts: { train: buckets.train.length, val: buckets.val.length, test: buckets.test.length },
+		outputs,
+	}
+}