Skip to content

Commit f37ccf8

Browse files
committed
fix: fix incorrect splitting with line delimited streaming
In some cases, valid CSV in datafusion would return: `Generic { store: "LineDelimiter", source: UnterminatedString }` due to incorrect logic. records_ends is a double ended iterator, so when calling next_back() the quoting/escaping logic would run in reverse, corrupting the internal state.
1 parent 0c380a1 commit f37ccf8

1 file changed

Lines changed: 55 additions & 16 deletions

File tree

src/delimited.rs

Lines changed: 55 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -77,22 +77,27 @@ impl LineDelimiter {
7777

7878
let is_escape = &mut self.is_escape;
7979
let is_quote = &mut self.is_quote;
80-
let mut record_ends = val.iter().enumerate().filter_map(|(idx, v)| {
81-
if *is_escape {
82-
*is_escape = false;
83-
None
84-
} else if *v == ESCAPE {
85-
*is_escape = true;
86-
None
87-
} else if *v == QUOTE {
88-
*is_quote = !*is_quote;
89-
None
90-
} else if *is_quote {
91-
None
92-
} else {
93-
(*v == NEWLINE).then_some(idx + 1)
94-
}
95-
});
80+
let mut record_ends = val
81+
.iter()
82+
.enumerate()
83+
.filter_map(|(idx, v)| {
84+
if *is_escape {
85+
*is_escape = false;
86+
None
87+
} else if *v == ESCAPE {
88+
*is_escape = true;
89+
None
90+
} else if *v == QUOTE {
91+
*is_quote = !*is_quote;
92+
None
93+
} else if *is_quote {
94+
None
95+
} else {
96+
(*v == NEWLINE).then_some(idx + 1)
97+
}
98+
})
99+
.collect::<Vec<_>>()
100+
.into_iter();
96101

97102
let start_offset = match self.remainder.is_empty() {
98103
true => 0,
@@ -270,4 +275,38 @@ mod tests {
270275
]
271276
)
272277
}
278+
279+
#[tokio::test]
280+
async fn test_delimiter_quotes_stream() {
281+
let input = vec!["x,y,z\n,\"new\nline\",\"with ", "space\""];
282+
let input_stream =
283+
futures_util::stream::iter(input.into_iter().map(|s| Ok(Bytes::from(s))));
284+
let stream = newline_delimited_stream(input_stream);
285+
286+
let results: Vec<_> = stream.try_collect().await.unwrap();
287+
assert_eq!(
288+
results,
289+
vec![
290+
Bytes::from("x,y,z\n"),
291+
Bytes::from(",\"new\nline\",\"with space\"")
292+
]
293+
)
294+
}
295+
296+
#[tokio::test]
297+
async fn test_delimiter_escape_stream() {
298+
let input = vec!["hello\n\n\"\\ttabulated\"", "world"];
299+
let input_stream =
300+
futures_util::stream::iter(input.into_iter().map(|s| Ok(Bytes::from(s))));
301+
let stream = newline_delimited_stream(input_stream);
302+
303+
let results: Vec<_> = stream.try_collect().await.unwrap();
304+
assert_eq!(
305+
results,
306+
vec![
307+
Bytes::from("hello\n\n"),
308+
Bytes::from("\"\\ttabulated\"world")
309+
]
310+
)
311+
}
273312
}

0 commit comments

Comments
 (0)