106 lines
3.0 KiB
Go
106 lines
3.0 KiB
Go
// Package chunk splits a tabular input into deterministic shards. It is generic
|
|
// row splitting only — no workload semantics (SMILES, top-k) live here.
|
|
package chunk
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
)
|
|
|
|
// ErrNoRows is returned when the input has a header but no data rows: a job with
|
|
// zero tasks could never complete, so it is rejected at the source.
|
|
var ErrNoRows = fmt.Errorf("input has no data rows")
|
|
|
|
// SplitTSV reads a header-plus-rows text stream and cuts it into shards of at
|
|
// most rowsPerShard data rows. Every shard repeats the header, so a worker can
|
|
// parse its shard in isolation. emit is called once per shard, in order, with a
|
|
// reader over that shard's bytes; the reader is valid only for the duration of
|
|
// the call.
|
|
//
|
|
// Splitting is deterministic: the same input and rowsPerShard always produce the
|
|
// same shards, byte for byte — which is what lets chunk_index refer to a stable
|
|
// piece and makes a re-run reproducible.
|
|
//
|
|
// Only one shard is buffered at a time, so memory is bounded by shard size (a
|
|
// worker-sized slice of the data), not by the size of the whole dataset.
|
|
func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reader) error) error {
|
|
return SplitTSVLimit(r, rowsPerShard, 0, emit)
|
|
}
|
|
|
|
// SplitTSVLimit behaves like SplitTSV but emits no more than maxRows data rows.
|
|
// A maxRows value of zero means unlimited. This lets an operator make a small,
|
|
// representative pipeline check without materialising a second dataset file.
|
|
func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
|
|
if rowsPerShard <= 0 {
|
|
return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard)
|
|
}
|
|
if maxRows < 0 {
|
|
return fmt.Errorf("maxRows must be non-negative, got %d", maxRows)
|
|
}
|
|
|
|
sc := bufio.NewScanner(r)
|
|
// Allow long lines: a SMILES row can be far wider than bufio's 64 KB default.
|
|
sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
|
|
|
|
if !sc.Scan() {
|
|
if err := sc.Err(); err != nil {
|
|
return fmt.Errorf("read header: %w", err)
|
|
}
|
|
return ErrNoRows // completely empty input
|
|
}
|
|
header := append([]byte(nil), sc.Bytes()...)
|
|
|
|
var (
|
|
buf bytes.Buffer
|
|
rows int
|
|
index int
|
|
)
|
|
|
|
// flush emits the buffered shard and resets for the next one.
|
|
flush := func() error {
|
|
if err := emit(index, bytes.NewReader(buf.Bytes())); err != nil {
|
|
return err
|
|
}
|
|
index++
|
|
buf.Reset()
|
|
rows = 0
|
|
return nil
|
|
}
|
|
|
|
for sc.Scan() {
|
|
if rows == 0 {
|
|
buf.Write(header)
|
|
buf.WriteByte('\n')
|
|
}
|
|
buf.Write(sc.Bytes())
|
|
buf.WriteByte('\n')
|
|
rows++
|
|
|
|
if rows == rowsPerShard {
|
|
if err := flush(); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if maxRows > 0 && index*rowsPerShard+rows == maxRows {
|
|
break
|
|
}
|
|
}
|
|
if err := sc.Err(); err != nil {
|
|
return fmt.Errorf("read rows: %w", err)
|
|
}
|
|
|
|
// A partial final shard still has to go out.
|
|
if rows > 0 {
|
|
if err := flush(); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
if index == 0 {
|
|
return ErrNoRows // header only, no data
|
|
}
|
|
return nil
|
|
}
|