Skip to content

Commit 0b6784b

Browse files
dfa1claude
andcommitted
fix(csv): single-pass streaming — buffer first chunk for schema inference
Previously two passes: first full-file inference pass gave no feedback (appeared frozen on large files), then write pass. Now single pass: buffer first chunkSize rows for schema inference, write them, stream the rest. No separate inference pass — first progress update appears after the first 65K rows. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent d5280ae commit 0b6784b

1 file changed

Lines changed: 89 additions & 87 deletions

File tree

csv/src/main/java/io/github/dfa1/vortex/csv/CsvImporter.java

Lines changed: 89 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,10 @@
2222
/// Provide a schema via [ImportOptions#withSchema] to skip inference.
2323
/// Empty cell values are treated as 0 / false / "" for typed columns.
2424
///
25-
/// Import is streaming: rows are read and written in chunks of [ImportOptions#chunkSize]
26-
/// without loading the entire file into memory. Schema inference requires a first
27-
/// sequential pass over the file; writing is a second pass. When a schema is provided
28-
/// via [ImportOptions#withSchema] only one pass is needed.
25+
/// Import is single-pass streaming. The first [ImportOptions#chunkSize] data rows
26+
/// are buffered to infer the schema (or skipped when a schema is provided), then
27+
/// all rows — including those first rows — are written in chunks. Memory usage is
28+
/// O(chunkSize) regardless of file size.
2929
public final class CsvImporter {
3030

3131
private CsvImporter() {
@@ -42,45 +42,85 @@ public static void importCsv(Path csvPath, Path vortexPath) throws IOException {
4242

4343
/// Imports a CSV file to a Vortex file.
4444
///
45-
/// Rows are streamed in chunks — the file is never fully loaded into memory.
46-
/// If no schema is provided, a first streaming pass infers column types; a
47-
/// second pass writes the data. Progress is reported via
48-
/// {@link ProgressListener#onProgress(long, long)} with `rowsTotal = -1`
49-
/// (total unknown) on each completed chunk.
45+
/// The file is read exactly once. The first chunk of rows is buffered for schema
46+
/// inference (O(chunkSize) memory); remaining rows stream directly to the writer.
47+
/// Progress is reported via {@link ProgressListener#onProgress(long, long)} with
48+
/// `rowsTotal = -1` (total unknown) after each chunk completes.
5049
///
5150
/// @param csvPath path to the source CSV file
5251
/// @param vortexPath path to write the output Vortex file
5352
/// @param options import configuration
5453
/// @throws IOException if reading or writing fails
55-
/// @throws IllegalArgumentException if the CSV file is empty
54+
/// @throws IllegalArgumentException if the CSV file has no data rows
5655
public static void importCsv(Path csvPath, Path vortexPath, ImportOptions options) throws IOException {
57-
String[] headers = readHeader(csvPath, options);
56+
int chunkSize = options.chunkSize();
5857

59-
DType.Struct schema = options.schema() != null
60-
? options.schema()
61-
: inferSchemaStreaming(csvPath, options, headers);
58+
try (CsvReader<CsvRecord> reader = csvReader(csvPath, options)) {
59+
// Read header row (if present) and buffer the first chunk of data rows.
60+
// Both happen in a single pass so the reader position advances correctly.
61+
String[] headers = null;
62+
List<String[]> firstChunk = new ArrayList<>(chunkSize);
63+
boolean expectHeader = options.hasHeader();
6264

63-
try (FileChannel channel = FileChannel.open(
64-
vortexPath, StandardOpenOption.CREATE, StandardOpenOption.WRITE,
65-
StandardOpenOption.TRUNCATE_EXISTING);
66-
VortexWriter writer = VortexWriter.create(channel, schema, options.writeOptions())) {
67-
writeChunksStreaming(csvPath, options, schema, writer);
68-
}
69-
}
65+
for (CsvRecord record : reader) {
66+
if (expectHeader) {
67+
headers = record.getFields().toArray(String[]::new);
68+
expectHeader = false;
69+
continue;
70+
}
71+
firstChunk.add(record.getFields().toArray(String[]::new));
72+
if (firstChunk.size() == chunkSize) {
73+
break;
74+
}
75+
}
76+
77+
if (firstChunk.isEmpty()) {
78+
throw new IllegalArgumentException("CSV file has no data rows");
79+
}
7080

71-
private static String[] readHeader(Path path, ImportOptions options) throws IOException {
72-
try (CsvReader<CsvRecord> reader = csvReader(path, options)) {
73-
CsvRecord first = reader.stream().findFirst().orElse(null);
74-
if (first == null) {
75-
throw new IllegalArgumentException("CSV file is empty");
81+
// Generate synthetic column names when the file has no header row.
82+
if (headers == null) {
83+
headers = generateHeaders(firstChunk.get(0).length);
84+
}
85+
86+
DType.Struct schema = options.schema() != null
87+
? options.schema()
88+
: inferSchemaFromRows(headers, firstChunk);
89+
90+
try (FileChannel channel = FileChannel.open(
91+
vortexPath, StandardOpenOption.CREATE, StandardOpenOption.WRITE,
92+
StandardOpenOption.TRUNCATE_EXISTING);
93+
VortexWriter writer = VortexWriter.create(channel, schema, options.writeOptions())) {
94+
95+
long totalRows = 0;
96+
97+
// Write the buffered first chunk.
98+
writer.writeChunk(buildChunk(schema, firstChunk));
99+
totalRows += firstChunk.size();
100+
reportProgress(options, totalRows);
101+
firstChunk.clear();
102+
103+
// Stream the rest of the file through the still-open reader.
104+
List<String[]> chunk = new ArrayList<>(chunkSize);
105+
for (CsvRecord record : reader) {
106+
chunk.add(record.getFields().toArray(String[]::new));
107+
if (chunk.size() == chunkSize) {
108+
writer.writeChunk(buildChunk(schema, chunk));
109+
totalRows += chunk.size();
110+
reportProgress(options, totalRows);
111+
chunk.clear();
112+
}
113+
}
114+
if (!chunk.isEmpty()) {
115+
writer.writeChunk(buildChunk(schema, chunk));
116+
totalRows += chunk.size();
117+
reportProgress(options, totalRows);
118+
}
76119
}
77-
String[] fields = first.getFields().toArray(String[]::new);
78-
return options.hasHeader() ? fields : generateHeaders(fields.length);
79120
}
80121
}
81122

82-
private static DType.Struct inferSchemaStreaming(Path path, ImportOptions options, String[] headers)
83-
throws IOException {
123+
private static DType.Struct inferSchemaFromRows(String[] headers, List<String[]> rows) {
84124
int colCount = headers.length;
85125
boolean[] canBeLong = new boolean[colCount];
86126
boolean[] canBeDouble = new boolean[colCount];
@@ -89,37 +129,29 @@ private static DType.Struct inferSchemaStreaming(Path path, ImportOptions option
89129
Arrays.fill(canBeDouble, true);
90130
Arrays.fill(canBeBool, true);
91131

92-
try (CsvReader<CsvRecord> reader = csvReader(path, options)) {
93-
boolean skipFirst = options.hasHeader();
94-
for (CsvRecord record : reader) {
95-
if (skipFirst) {
96-
skipFirst = false;
132+
for (String[] row : rows) {
133+
for (int c = 0; c < colCount; c++) {
134+
String val = safeGet(row, c);
135+
if (val.isEmpty()) {
97136
continue;
98137
}
99-
String[] fields = record.getFields().toArray(String[]::new);
100-
for (int c = 0; c < colCount; c++) {
101-
String val = safeGet(fields, c);
102-
if (val.isEmpty()) {
103-
continue;
138+
if (canBeLong[c]) {
139+
try {
140+
Long.parseLong(val);
141+
} catch (NumberFormatException e) {
142+
canBeLong[c] = false;
104143
}
105-
if (canBeLong[c]) {
106-
try {
107-
Long.parseLong(val);
108-
} catch (NumberFormatException e) {
109-
canBeLong[c] = false;
110-
}
111-
}
112-
if (canBeDouble[c]) {
113-
try {
114-
Double.parseDouble(val);
115-
} catch (NumberFormatException e) {
116-
canBeDouble[c] = false;
117-
}
144+
}
145+
if (canBeDouble[c]) {
146+
try {
147+
Double.parseDouble(val);
148+
} catch (NumberFormatException e) {
149+
canBeDouble[c] = false;
118150
}
119-
if (canBeBool[c]) {
120-
if (!val.equalsIgnoreCase("true") && !val.equalsIgnoreCase("false")) {
121-
canBeBool[c] = false;
122-
}
151+
}
152+
if (canBeBool[c]) {
153+
if (!val.equalsIgnoreCase("true") && !val.equalsIgnoreCase("false")) {
154+
canBeBool[c] = false;
123155
}
124156
}
125157
}
@@ -133,36 +165,6 @@ private static DType.Struct inferSchemaStreaming(Path path, ImportOptions option
133165
return new DType.Struct(names, types, false);
134166
}
135167

136-
private static void writeChunksStreaming(Path path, ImportOptions options, DType.Struct schema,
137-
VortexWriter writer) throws IOException {
138-
int chunkSize = options.chunkSize();
139-
List<String[]> chunk = new ArrayList<>(chunkSize);
140-
long totalRows = 0;
141-
142-
try (CsvReader<CsvRecord> reader = csvReader(path, options)) {
143-
boolean skipFirst = options.hasHeader();
144-
for (CsvRecord record : reader) {
145-
if (skipFirst) {
146-
skipFirst = false;
147-
continue;
148-
}
149-
chunk.add(record.getFields().toArray(String[]::new));
150-
if (chunk.size() == chunkSize) {
151-
writer.writeChunk(buildChunk(schema, chunk));
152-
totalRows += chunk.size();
153-
reportProgress(options, totalRows);
154-
chunk.clear();
155-
}
156-
}
157-
}
158-
159-
if (!chunk.isEmpty()) {
160-
writer.writeChunk(buildChunk(schema, chunk));
161-
totalRows += chunk.size();
162-
reportProgress(options, totalRows);
163-
}
164-
}
165-
166168
private static void reportProgress(ImportOptions options, long totalRows) {
167169
if (options.progressListener() != null) {
168170
options.progressListener().onProgress(totalRows, -1);

0 commit comments

Comments
 (0)