Add job cancellation and dataset row limit

coordinator/internal/chunk/tsv.go

@@ -26,9 +26,19 @@ var ErrNoRows = fmt.Errorf("input has no data rows")
 // 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.
@@ -73,6 +83,9 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade
 				return err
 			}
 		}
+		if maxRows > 0 && index*rowsPerShard+rows == maxRows {
+			break
+		}
 	}
 	if err := sc.Err(); err != nil {
 		return fmt.Errorf("read rows: %w", err)