diff --git a/crates/s3select-api/src/csv_input.rs b/crates/s3select-api/src/csv_input.rs index b63d82d84..41206e675 100644 --- a/crates/s3select-api/src/csv_input.rs +++ b/crates/s3select-api/src/csv_input.rs @@ -17,7 +17,7 @@ use datafusion::object_store::{Error, Result}; use futures::{Stream, StreamExt, stream::BoxStream}; use transform_stream::AsyncTryStream; -use crate::SelectError; +use crate::{SelectError, input_stream::MAX_SELECT_RECORD_BYTES}; /// Arrow accepts byte-sized CSV controls. Unicode quotes need streaming normalization. pub fn csv_input_requires_normalization(quote: Option<&str>, escape: Option<&str>) -> bool { @@ -32,6 +32,7 @@ pub(crate) struct CsvSyntax<'a> { pub field: Option<&'a str>, pub record: Option<&'a str>, pub comment: Option, + pub allow_quoted_record_delimiter: bool, } #[derive(Clone, Copy, PartialEq, Eq)] @@ -45,8 +46,9 @@ enum State { } /// Emits ordinary CSV with every field quoted. This avoids reserving a sentinel -/// byte that might also appear in a UTF-8 field. Only a partial control token is -/// retained between chunks; neither records nor objects are buffered. +/// byte that might also appear in a UTF-8 field. Multiline input retains at most +/// one bounded record so schema inference always receives complete records. +/// Otherwise only a partial control token is retained between chunks. struct CsvInputNormalizer { quote: Vec, escape: Vec, @@ -54,6 +56,9 @@ struct CsvInputNormalizer { record: Vec, comment: Option, default_records: bool, + allow_quoted_record_delimiter: bool, + record_bytes: usize, + pending_record: Vec, state: State, record_start: bool, carry: Vec, @@ -89,6 +94,9 @@ impl CsvInputNormalizer { record, comment: csv.comment, default_records: csv.record.is_none(), + allow_quoted_record_delimiter: csv.allow_quoted_record_delimiter, + record_bytes: 0, + pending_record: Vec::new(), state: State::FieldStart, record_start: true, carry: Vec::new(), @@ -97,9 +105,11 @@ impl CsvInputNormalizer { } fn record_len(&self, bytes: &[u8]) -> usize { - if self.default_records && bytes.starts_with(b"\r\n") { + if self.default_records && self.field == b"\r\n" && bytes.starts_with(&self.field) { + 0 + } else if self.default_records && bytes.starts_with(b"\r\n") { 2 - } else if self.default_records && bytes.starts_with(b"\r") { + } else if self.default_records && self.field != b"\r" && bytes.starts_with(b"\r") { 1 } else if bytes.starts_with(&self.record) { self.record.len() @@ -108,10 +118,11 @@ impl CsvInputNormalizer { } } - fn push_value(output: &mut Vec, bytes: &[u8]) { + fn push_value(&self, output: &mut Vec, bytes: &[u8]) { for byte in bytes { - if *byte == b'"' { - output.push(b'"'); + // Schema inference treats backslashes as escapes when delimiting chunks. + if *byte == b'"' || (self.allow_quoted_record_delimiter && *byte == b'\\') { + output.push(*byte); } output.push(*byte); } @@ -125,9 +136,12 @@ impl CsvInputNormalizer { } else { bytes.len().saturating_sub(self.token_size - 1) }; - let mut output = Vec::with_capacity(bytes.len()); + let mut output = std::mem::take(&mut self.pending_record); + output.reserve(bytes.len()); + let mut complete_records_end = 0; let mut pos = 0; while pos < end { + let start = pos; let rest = &bytes[pos..]; let record_len = self.record_len(rest); let field = rest.starts_with(&self.field) && self.field.len() > record_len; @@ -141,10 +155,10 @@ impl CsvInputNormalizer { } } State::Escaped => { - if record_len > 0 { + if record_len > 0 && !self.allow_quoted_record_delimiter { return Err(SelectError::CsvParsingError); } - Self::push_value(&mut output, &rest[..1]); + self.push_value(&mut output, &rest[..1]); self.state = State::Quoted; pos += 1; } @@ -157,14 +171,14 @@ impl CsvInputNormalizer { pos += self.escape.len(); } State::Quoted => { - if record_len > 0 { + if record_len > 0 && !self.allow_quoted_record_delimiter { return Err(SelectError::CsvParsingError); } - Self::push_value(&mut output, &rest[..1]); + self.push_value(&mut output, &rest[..1]); pos += 1; } State::AfterQuote if rest.starts_with(&self.quote) => { - Self::push_value(&mut output, &self.quote); + self.push_value(&mut output, &self.quote); self.state = State::Quoted; pos += self.quote.len(); } @@ -189,6 +203,9 @@ impl CsvInputNormalizer { output.push(if field { b',' } else { b'\n' }); self.state = State::FieldStart; self.record_start = !field; + if !field { + complete_records_end = output.len(); + } pos += if field { self.field.len() } else { record_len }; } _ => { @@ -197,10 +214,21 @@ impl CsvInputNormalizer { } self.state = State::Unquoted; self.record_start = false; - Self::push_value(&mut output, &rest[..1]); + self.push_value(&mut output, &rest[..1]); pos += 1; } } + if self.allow_quoted_record_delimiter { + // Quoted delimiters do not end a logical record or reset its size limit. + if self.record_start && self.state != State::Comment { + self.record_bytes = 0; + } else { + self.record_bytes += pos - start; + if self.record_bytes > MAX_SELECT_RECORD_BYTES { + return Err(SelectError::OverMaxRecordSize); + } + } + } } self.carry.extend_from_slice(&bytes[pos..]); if last { @@ -210,6 +238,12 @@ impl CsvInputNormalizer { State::FieldStart if !self.record_start => output.extend_from_slice(b"\"\""), State::FieldStart | State::Comment => {} } + } else if self.allow_quoted_record_delimiter { + if complete_records_end == 0 { + self.pending_record = output; + return Ok(Vec::new()); + } + self.pending_record = output.split_off(complete_records_end); } Ok(output) } @@ -275,6 +309,7 @@ mod tests { field: Some("界"), record: Some("^Y"), comment: Some(b'#'), + ..Default::default() }; let input = "#skipع界^Yعa界bع界\"literal\"^Yعline\nbreakع界end^Y"; let expected = "\"a界b\",\"\"\"literal\"\"\"\n\"line\nbreak\",\"end\"\n"; @@ -296,6 +331,167 @@ mod tests { } } + #[test] + fn quoted_record_delimiters_preserve_values_at_every_chunk_boundary() { + let cases = [ + ("\"", "\"", ",", None, "\"a\nb\",tail\n", "\"a\nb\",\"tail\"\n"), + ("\"", "\"", ",", None, "\"a\r\nb\",tail\r\n", "\"a\r\nb\",\"tail\"\n"), + ("\"", "\"", ",", Some("\r\n"), "\"a\r\nb\",tail\r\n", "\"a\r\nb\",\"tail\"\n"), + ("\"", "\"", ",", Some("|"), "\"a|b\",tail|", "\"a|b\",\"tail\"\n"), + ("\"", "\"", ",", Some("^Y"), "\"a^Yb\",tail^Y", "\"a^Yb\",\"tail\"\n"), + ("ع", "\\", "界", Some("^Y"), "عa^Ybع界tail^Y", "\"a^Yb\",\"tail\"\n"), + ("\"", "\\", ",", None, "\"a\\\nb\",tail\n", "\"a\nb\",\"tail\"\n"), + ("\"", "\"", ",", None, "\"a\"\"\nb\",tail", "\"a\"\"\nb\",\"tail\""), + ("\"", "\"", ",", None, "#skip\n\"a\n#b\",tail\n", "\"a\n#b\",\"tail\"\n"), + ]; + for (quote, escape, field, record, input, expected) in cases { + let csv = CsvSyntax { + quote: Some(quote), + escape: Some(escape), + field: Some(field), + record, + comment: Some(b'#'), + allow_quoted_record_delimiter: true, + }; + for chunk_size in 1..=input.len() { + assert_eq!( + normalize_chunks(&csv, input.as_bytes(), chunk_size), + expected.as_bytes(), + "input={input:?}, chunk_size={chunk_size}" + ); + } + } + } + + #[test] + fn quoted_record_delimiters_enforce_logical_record_size_across_chunks() { + for size in [ + MAX_SELECT_RECORD_BYTES - 1, + MAX_SELECT_RECORD_BYTES, + MAX_SELECT_RECORD_BYTES + 1, + ] { + let csv = CsvSyntax { + allow_quoted_record_delimiter: true, + ..Default::default() + }; + let mut input = vec![b'\n'; size]; + input[0] = b'"'; + input[size - 1] = b'"'; + let mut normalizer = CsvInputNormalizer::new(&csv); + let result = input + .chunks(64 * 1024) + .try_for_each(|chunk| normalizer.convert(chunk, false).map(|_| ())) + .and_then(|()| normalizer.convert(b"\nnext\n", true)); + if size > MAX_SELECT_RECORD_BYTES { + assert_eq!(result, Err(SelectError::OverMaxRecordSize)); + } else { + assert!(result.is_ok(), "a record of {size} bytes must be accepted: {result:?}"); + assert!(normalizer.carry.is_empty()); + } + } + } + + #[test] + fn quoted_record_delimiters_count_comment_bytes_across_chunks() { + for size in [ + MAX_SELECT_RECORD_BYTES - 1, + MAX_SELECT_RECORD_BYTES, + MAX_SELECT_RECORD_BYTES + 1, + ] { + for record in [None, Some("\r\n"), Some("^Y")] { + let delimiter = record.unwrap_or("\n"); + for terminated in [false, true] { + let csv = CsvSyntax { + record, + comment: Some(b'#'), + allow_quoted_record_delimiter: true, + ..Default::default() + }; + let mut input = vec![b'x'; size]; + input[0] = b'#'; + if terminated { + input.extend_from_slice(delimiter.as_bytes()); + input.extend_from_slice(b"#next"); + input.extend_from_slice(delimiter.as_bytes()); + input.extend_from_slice(b"value"); + } + let mut normalizer = CsvInputNormalizer::new(&csv); + let mut output = Vec::new(); + let result = input + .chunks(64 * 1024) + .try_for_each(|chunk| normalizer.convert(chunk, false).map(|bytes| output.extend(bytes))) + .and_then(|()| normalizer.convert(b"", true).map(|bytes| output.extend(bytes))); + if size > MAX_SELECT_RECORD_BYTES { + assert_eq!(result, Err(SelectError::OverMaxRecordSize)); + } else { + result.expect("comments within the logical-record limit must be accepted"); + assert_eq!(output, if terminated { b"\"value\"".as_slice() } else { b"" }); + } + } + } + } + } + + #[test] + fn quoted_record_delimiters_preserve_carriage_return_fields() { + let csv = CsvSyntax { + field: Some("\r"), + allow_quoted_record_delimiter: true, + ..Default::default() + }; + for (input, expected) in [ + ("a\rb\n", "\"a\",\"b\"\n"), + ("a\rb\r\n", "\"a\",\"b\"\n"), + ("\"a\rb\"\rtail\n", "\"a\rb\",\"tail\"\n"), + ("a\r", "\"a\",\"\""), + ] { + for size in 1..=input.len() { + assert_eq!( + normalize_chunks(&csv, input.as_bytes(), size), + expected.as_bytes(), + "input={input:?}, chunk={size}" + ); + } + } + } + + #[test] + fn quoted_record_delimiters_preserve_crlf_fields() { + let csv = CsvSyntax { + field: Some("\r\n"), + comment: Some(b'#'), + allow_quoted_record_delimiter: true, + ..Default::default() + }; + for (input, expected) in [ + ("a\r\nb\n", "\"a\",\"b\"\n"), + ("\"a\r\nb\"\r\ntail\n", "\"a\r\nb\",\"tail\"\n"), + ("a\r\n", "\"a\",\"\""), + ("#skip\r\na\r\nb\n", "\"a\",\"b\"\n"), + ] { + for size in 1..=input.len() { + assert_eq!( + normalize_chunks(&csv, input.as_bytes(), size), + expected.as_bytes(), + "input={input:?}, chunk={size}" + ); + } + } + } + + #[test] + fn quoted_record_delimiters_do_not_allow_unterminated_fields() { + for input in ["\"a\nb", "\"a\nb\\"] { + let csv = CsvSyntax { + escape: Some("\\"), + allow_quoted_record_delimiter: true, + ..Default::default() + }; + let mut normalizer = CsvInputNormalizer::new(&csv); + assert_eq!(normalizer.convert(input.as_bytes(), true), Err(SelectError::CsvParsingError)); + } + } + #[test] fn unicode_csv_quotes_preserve_omitted_syntax_defaults() { assert!(!csv_input_requires_normalization(None, None)); @@ -372,6 +568,7 @@ where #[cfg(test)] mod stream_tests { use super::*; + use futures::TryStreamExt; use std::sync::{ Arc, atomic::{AtomicBool, AtomicUsize, Ordering}, @@ -406,6 +603,47 @@ mod stream_tests { assert_eq!(polls.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn quoted_record_delimiters_preserve_backslashes_through_schema_chunking() { + let data = b"path\\,\"line\nbreak\"\n"; + let csv = CsvSyntax { + allow_quoted_record_delimiter: true, + ..Default::default() + }; + for size in 1..=data.len() { + let chunks = data + .chunks(size) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect::>(); + let stream = normalize_csv_stream(futures::stream::iter(chunks), &csv); + let output = datafusion::object_store::delimited::newline_delimited_stream(stream) + .try_collect::>() + .await + .unwrap_or_else(|error| panic!("schema inference failed at chunk size {size}: {error}")); + assert_eq!(output.concat(), b"\"path\\\\\",\"line\nbreak\"\n", "chunk size={size}"); + } + } + + #[tokio::test] + async fn quoted_record_delimiters_report_unterminated_fields_after_partial_output() { + let source = futures::stream::iter([Ok(Bytes::from_static(b"first,row\n\"line\nbreak"))]); + let csv = CsvSyntax { + allow_quoted_record_delimiter: true, + ..Default::default() + }; + let mut stream = normalize_csv_stream(source, &csv); + let prefix = stream.next().await.expect("partial output").expect("valid prefix"); + assert!(prefix.starts_with(b"\"first\",\"row\"\n")); + let error = stream + .next() + .await + .expect("terminal failure") + .expect_err("unclosed field must fail"); + let Error::Generic { source, .. } = error else { panic!("expected a typed CSV parse error") }; + assert_eq!(source.downcast_ref::(), Some(&SelectError::CsvParsingError)); + assert!(stream.next().await.is_none()); + } + #[tokio::test] async fn unicode_csv_quotes_preserve_source_errors_after_partial_output() { let source = futures::stream::iter([ diff --git a/crates/s3select-api/src/object_store.rs b/crates/s3select-api/src/object_store.rs index 43e87dc41..ab1947760 100644 --- a/crates/s3select-api/src/object_store.rs +++ b/crates/s3select-api/src/object_store.rs @@ -16,7 +16,10 @@ use crate::{ PrepareSelectObjectSnapshotError, QueryError, SELECT_DEFAULT_READ_BUFFER_SIZE, SelectError, SelectGetObjectReader, SelectInputMetrics, SelectObjectOptions, SelectObjectSnapshot, SelectObjectSnapshotReadError, SelectStorageError, SelectStore, SnapshotConsistencyError, - input_stream::{CompressionFormat, compressed_input_reader, compressed_input_stream, processed_bytes_limit}, + input_stream::{ + CompressionFormat, SELECT_DECODE_CHUNK_BYTES, compressed_input_reader, compressed_input_stream, input_io_error, + processed_bytes_limit, + }, query::{ ast::{JsonPathSegment, JsonSource}, parser::RustFsDialect, @@ -352,9 +355,11 @@ impl EcObjectStore { S: Stream> + Send + 'static, { if let Some(csv) = self.input.request.input_serialization.csv.as_ref() - && csv_input_requires_normalization(csv.quote_character.as_deref(), csv.quote_escape_character.as_deref()) + && (csv.allow_quoted_record_delimiter.unwrap_or(false) + || csv_input_requires_normalization(csv.quote_character.as_deref(), csv.quote_escape_character.as_deref())) { let syntax = CsvSyntax { + allow_quoted_record_delimiter: csv.allow_quoted_record_delimiter.unwrap_or(false), quote: csv.quote_character.as_deref(), escape: csv.quote_escape_character.as_deref(), field: csv.field_delimiter.as_deref(), @@ -925,6 +930,26 @@ impl ObjectStore for EcObjectStore { self.query_tracker.clone(), ); GetResultPayload::Stream(stream) + } else if self + .input + .request + .input_serialization + .csv + .as_ref() + .is_some_and(|csv| csv.allow_quoted_record_delimiter.unwrap_or(false)) + { + // The CSV normalizer owns logical-record limits; a physical-line + // tracker would count default CRLF terminators as record data. + let reader = compressed_input_reader( + reader.stream, + original_size, + compression, + Arc::clone(&self.input_metrics), + max_processed_bytes, + query_guard, + ); + let stream = ReaderStream::with_capacity(reader, SELECT_DECODE_CHUNK_BYTES).map_err(input_io_error); + GetResultPayload::Stream(self.convert_csv_stream(stream)) } else { let input_record_delimiter = if self.input.request.input_serialization.csv.is_some() { self.record_delimiter() @@ -3010,6 +3035,141 @@ mod test { assert_eq!(input_metrics.snapshot().bytes_processed, 2); } + #[tokio::test] + async fn quoted_record_delimiters_preserve_compressed_values_and_metrics() { + const BUCKET: &str = "s3select-quoted-record-stream"; + let data = b"name,kind\n\"line\nbreak\",tail\n"; + let env = crate::storage_api::select_test_ecstore_env().await; + env.make_bucket(BUCKET, false).await; + for (object, compression) in [ + ("plain.csv", None), + ("gzip.csv", Some(CompressionFormat::Gzip)), + ("bzip.csv", Some(CompressionFormat::Bzip2)), + ] { + let bytes = match compression { + Some(format) => encode_compressed_fixture(format, data).await, + None => data.to_vec(), + }; + let raw_size = bytes.len(); + let mut reader = SelectPutObjReader::from_vec(bytes); + env.ecstore + .put_object(BUCKET, object, &mut reader, &Default::default()) + .await + .expect("write multiline CSV"); + let mut input = (*csv_input(BUCKET, object)).clone(); + let csv = input.request.input_serialization.csv.as_mut().expect("CSV input"); + csv.file_header_info = Some(FileHeaderInfo::from_static(FileHeaderInfo::USE)); + csv.allow_quoted_record_delimiter = Some(true); + input.request.input_serialization.compression_type = compression.map(|format| { + CompressionType::from_static(match format { + CompressionFormat::Gzip => CompressionType::GZIP, + CompressionFormat::Bzip2 => CompressionType::BZIP2, + }) + }); + let metrics = Arc::new(SelectInputMetrics::default()); + let store = EcObjectStore::build_with_snapshot( + Arc::new(input), + Arc::new(GreedyMemoryPool::new(1024 * 1024)), + None, + Arc::clone(&metrics), + prepare_test_snapshot(BUCKET, object).await, + JsonSource::default(), + ) + .expect("snapshot store"); + let result = store + .get_opts(&Path::from(object), GetOptions::default()) + .await + .expect("open multiline CSV stream"); + let GetResultPayload::Stream(stream) = result.payload else { panic!("CSV must remain streaming") }; + let output = stream + .try_collect::>() + .await + .expect("normalize multiline CSV") + .concat(); + assert_eq!(output, b"\"name\",\"kind\"\n\"line\nbreak\",\"tail\"\n", "object={object}"); + let measured = metrics.snapshot(); + assert_eq!(measured.bytes_scanned, u64::try_from(raw_size).expect("raw length")); + assert_eq!(measured.bytes_processed, u64::try_from(data.len()).expect("decoded length")); + } + } + + #[tokio::test] + async fn quoted_record_delimiters_preserve_compressed_crlf_size_boundaries() { + use crate::input_stream::MAX_SELECT_RECORD_BYTES; + + const BUCKET: &str = "select-quoted-crlf-limit"; + let env = crate::storage_api::select_test_ecstore_env().await; + env.make_bucket(BUCKET, false).await; + for compression in [None, Some(CompressionFormat::Gzip), Some(CompressionFormat::Bzip2)] { + for size in [ + MAX_SELECT_RECORD_BYTES - 1, + MAX_SELECT_RECORD_BYTES, + MAX_SELECT_RECORD_BYTES + 1, + ] { + for multiline in [false, true] { + let mut data = vec![b'a'; size]; + if multiline { + data[0] = b'"'; + data[size / 2] = b'\n'; + data[size - 1] = b'"'; + } + data.extend_from_slice(b"\r\n"); + let bytes = match compression { + Some(format) => encode_compressed_fixture(format, &data).await, + None => data.clone(), + }; + let raw_size = bytes.len(); + let object = format!("{compression:?}-{size}-{multiline}"); + let mut reader = SelectPutObjReader::from_vec(bytes); + env.ecstore + .put_object(BUCKET, &object, &mut reader, &Default::default()) + .await + .expect("write boundary fixture"); + let mut input = (*csv_input(BUCKET, &object)).clone(); + let csv = input.request.input_serialization.csv.as_mut().expect("CSV input"); + csv.file_header_info = Some(FileHeaderInfo::from_static(FileHeaderInfo::NONE)); + csv.allow_quoted_record_delimiter = Some(true); + input.request.input_serialization.compression_type = compression.map(|format| { + CompressionType::from_static(match format { + CompressionFormat::Gzip => CompressionType::GZIP, + CompressionFormat::Bzip2 => CompressionType::BZIP2, + }) + }); + let metrics = Arc::new(SelectInputMetrics::default()); + let store = EcObjectStore::build_with_snapshot( + Arc::new(input), + Arc::new(GreedyMemoryPool::new(4 * MAX_SELECT_RECORD_BYTES)), + None, + Arc::clone(&metrics), + prepare_test_snapshot(BUCKET, &object).await, + JsonSource::default(), + ) + .expect("boundary store"); + let result = store + .get_opts(&Path::from(object.as_str()), GetOptions::default()) + .await + .expect("open boundary stream"); + let GetResultPayload::Stream(stream) = result.payload else { panic!("CSV stream") }; + let output = stream.try_collect::>().await; + if size > MAX_SELECT_RECORD_BYTES { + let error = output.expect_err("logical records over 1 MiB must fail, even with embedded newlines"); + let object_store::Error::Generic { source, .. } = error else { panic!("typed record size error") }; + assert_eq!(source.downcast_ref::(), Some(&SelectError::OverMaxRecordSize)); + } else { + let value = if multiline { &data[1..size - 1] } else { &data[..size] }; + let mut expected = vec![b'"']; + expected.extend_from_slice(value); + expected.extend_from_slice(b"\"\n"); + assert_eq!(output.expect("records within the limit must pass").concat(), expected, "object={object}"); + let measured = metrics.snapshot(); + assert_eq!(measured.bytes_scanned, u64::try_from(raw_size).expect("raw length")); + assert_eq!(measured.bytes_processed, u64::try_from(data.len()).expect("decoded length")); + } + } + } + } + } + #[tokio::test] async fn unicode_csv_quotes_preserve_raw_offsets_and_metrics() { const BUCKET: &str = "s3select-unicode-csv-stream"; diff --git a/crates/s3select-api/src/query/session.rs b/crates/s3select-api/src/query/session.rs index 124c0c85d..1553d5731 100644 --- a/crates/s3select-api/src/query/session.rs +++ b/crates/s3select-api/src/query/session.rs @@ -458,7 +458,11 @@ impl SessionCtxFactory { input_metrics.is_some() && context.input.request.input_serialization.parquet.is_none(); let normalized_csv_requires_single_file_scan = context.input.request.input_serialization.csv.as_ref().is_some_and(|csv| { - crate::csv_input_requires_normalization(csv.quote_character.as_deref(), csv.quote_escape_character.as_deref()) + csv.allow_quoted_record_delimiter.unwrap_or(false) + || crate::csv_input_requires_normalization( + csv.quote_character.as_deref(), + csv.quote_escape_character.as_deref(), + ) }); let config = if normalized_csv_requires_single_file_scan || custom_two_byte_record_delimiter @@ -911,6 +915,24 @@ mod tests { assert!(!session.inner().config().options().optimizer.repartition_file_scans); } + #[tokio::test] + async fn quoted_record_delimiters_disable_file_scan_repartition() { + let mut context = test_context(); + Arc::make_mut(&mut context.input) + .request + .input_serialization + .csv + .as_mut() + .expect("CSV input") + .allow_quoted_record_delimiter = Some(true); + let session = SessionCtxFactory::new(true) + .with_target_partitions(4) + .create_session_ctx(&context) + .await + .expect("multiline CSV session"); + assert!(!session.inner().config().options().optimizer.repartition_file_scans); + } + #[tokio::test] async fn unicode_csv_quotes_disable_file_scan_repartition() { let mut context = test_context(); diff --git a/crates/s3select-query/src/dispatcher/manager.rs b/crates/s3select-query/src/dispatcher/manager.rs index 22a733fe7..06e21393c 100644 --- a/crates/s3select-query/src/dispatcher/manager.rs +++ b/crates/s3select-query/src/dispatcher/manager.rs @@ -478,13 +478,15 @@ impl SimpleQueryDispatcher { if let Some(quote) = csv.quote_character.as_ref() { file_format = file_format.with_quote(quote.as_bytes().first().copied().unwrap_or_default()); } - if rustfs_s3select_api::csv_input_requires_normalization( - csv.quote_character.as_deref(), - csv.quote_escape_character.as_deref(), - ) { + if csv.allow_quoted_record_delimiter.unwrap_or(false) + || rustfs_s3select_api::csv_input_requires_normalization( + csv.quote_character.as_deref(), + csv.quote_escape_character.as_deref(), + ) + { file_format = file_format .with_quote(b'"') - .with_escape(None) + .with_escape(csv.allow_quoted_record_delimiter.unwrap_or(false).then_some(b'\\')) .with_delimiter(b',') .with_terminator(Some(b'\n')) .with_comment(None) @@ -1582,6 +1584,84 @@ mod tests { } } + #[tokio::test] + async fn quoted_record_delimiters_reach_arrow_without_changing_field_values() { + let cases = [ + ( + "\"", + "\"", + ",", + "\n", + "\"a\nb\",tail\nnext,row\n", + vec![vec!["a\nb", "tail"], vec!["next", "row"]], + ), + ("\"", "\"", ",", "\r\n", "\"a\r\nb\",tail\r\n", vec![vec!["a\r\nb", "tail"]]), + ("\"", "\"", ",", "|", "\"a|b\",tail|", vec![vec!["a|b", "tail"]]), + ("\"", "\"", ",", "^Y", "\"a^Yb\",tail^Y", vec![vec!["a^Yb", "tail"]]), + ("ع", "\\", "界", "^Y", "عa^Ybع界عline\nbreakع^Y", vec![vec!["a^Yb", "line\nbreak"]]), + ("'", "\\", ";", "\n", "'a\\'\nb';tail\n", vec![vec!["a'\nb", "tail"]]), + ("\"", "\"", ",", "\n", "\"a\"\"\nb\",tail", vec![vec!["a\"\nb", "tail"]]), + ("\"", "\"", ",", "\n", "path\\,\"line\nbreak\"\n", vec![vec!["path\\", "line\nbreak"]]), + ("\"", "\"", "\r", "\n", "a\rb\n", vec![vec!["a", "b"]]), + ("\"", "\"", "\r\n", "\n", "a\r\nb\n", vec![vec!["a", "b"]]), + ]; + let env = snapshot_test_env().await; + for (index, (quote, escape, field, record, data, expected)) in cases.into_iter().enumerate() { + for header in [FileHeaderInfo::NONE, FileHeaderInfo::USE, FileHeaderInfo::IGNORE] { + let mut input = test_input(); + input.bucket = format!("select-quoted-records-{index}-{}", header.to_ascii_lowercase()); + input.key = "records".to_owned(); + let csv = input.request.input_serialization.csv.as_mut().expect("CSV input"); + csv.allow_quoted_record_delimiter = Some(true); + csv.file_header_info = Some(FileHeaderInfo::from_static(header)); + if index != 0 { + csv.quote_character = Some(quote.to_owned()); + csv.quote_escape_character = Some(escape.to_owned()); + csv.field_delimiter = Some(field.to_owned()); + csv.record_delimiter = (!matches!(field, "\r" | "\r\n")).then(|| record.to_owned()); + } + let data = if header == FileHeaderInfo::NONE { + data.to_owned() + } else { + format!("{quote}first{record}name{quote}{field}second{record}{data}") + }; + env.make_bucket(&input.bucket, false).await; + env.put_object_bytes(&input.bucket, &input.key, data.as_bytes().to_vec()) + .await; + let snapshot = env.prepare_select_object_snapshot(&input.bucket, &input.key).await; + let input = Arc::new(input); + let dispatcher = production_dispatcher(Arc::clone(&input)); + let query = Query::new_with_snapshot(QueryContext { input }, "SELECT * FROM S3Object".to_owned(), snapshot); + let output = dispatcher.execute_query(&query).await.expect("execute multiline CSV query"); + let mut stream = output.into_record_batch_stream().expect("record stream"); + let mut rows = Vec::new(); + while let Some(batch) = stream.next().await { + let batch = batch.expect("Arrow must receive valid UTF-8 fields"); + if header == FileHeaderInfo::USE { + assert_eq!(batch.schema().field(0).name(), &format!("first{record}name")); + } + for row in 0..batch.num_rows() { + rows.push( + batch + .columns() + .iter() + .map(|column| { + column + .as_any() + .downcast_ref::() + .expect("CSV string column") + .value(row) + .to_owned() + }) + .collect::>(), + ); + } + } + assert_eq!(rows, expected, "fixture={index}, header={header}"); + } + } + } + #[tokio::test] async fn select_uses_input_serialization_independently_of_object_extension() { for (key, json) in [ diff --git a/rustfs/src/app/select_object.rs b/rustfs/src/app/select_object.rs index 9e67b90ec..9a51272f5 100644 --- a/rustfs/src/app/select_object.rs +++ b/rustfs/src/app/select_object.rs @@ -38,7 +38,7 @@ use s3s::dto::{ SelectObjectContentEventStream, SelectObjectContentInput, SelectObjectContentOutput, SelectObjectContentRequest, Stats, StatsEvent, }; -use s3s::{S3Error, S3ErrorCode, S3Request, S3Response, S3Result, s3_error}; +use s3s::{S3Error, S3ErrorCode, S3Request, S3Response, S3Result}; use std::{ fmt, future::poll_fn, @@ -606,7 +606,15 @@ fn validate_select_request(headers: &http::HeaderMap, input: &mut SelectObjectCo normalize_input_serialization(&mut input.request.input_serialization)?; let compressed_input = is_compressed_input(&input.request.input_serialization); - if compressed_input && input.request.scan_range.as_ref().is_some_and(is_noop_scan_range) { + if compressed_input + && input.request.scan_range.as_ref().is_some_and(is_noop_scan_range) + && !input + .request + .input_serialization + .csv + .as_ref() + .is_some_and(|csv| csv.allow_quoted_record_delimiter.unwrap_or(false)) + { input.request.scan_range = None; } validate_scan_range(&input.request)?; @@ -670,12 +678,6 @@ fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()> .get_or_insert_with(|| CompressionType::from_static(CompressionType::NONE)); if let Some(csv) = input.csv.as_mut() { - if csv.allow_quoted_record_delimiter.unwrap_or(false) { - return Err(s3_error!( - NotImplemented, - "CSV AllowQuotedRecordDelimiter is not supported by SelectObjectContent" - )); - } let file_header_info = csv .file_header_info .get_or_insert_with(|| FileHeaderInfo::from_static(FileHeaderInfo::NONE)); @@ -755,7 +757,7 @@ fn validate_scan_range_protocol(request: &SelectObjectContentRequest) -> Result< input_serialization.json.as_ref(), input_serialization.parquet.as_ref(), ) { - (Some(_), None, None) => true, + (Some(csv), None, None) => !csv.allow_quoted_record_delimiter.unwrap_or(false), (None, Some(json), None) if !is_json_document(json) => true, (None, None, Some(_)) => true, _ => false, @@ -3533,6 +3535,66 @@ mod tests { assert_eq!(error.message(), Some(INVALID_SCAN_RANGE_MESSAGE)); } + #[test] + fn validate_accepts_quoted_record_delimiter_options() { + for allow in [None, Some(false), Some(true)] { + let mut input = base_input(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("CSV input") + .allow_quoted_record_delimiter = allow; + validate_select_request(&HeaderMap::new(), &mut input).expect("CSV option should validate without ScanRange"); + assert_eq!( + input + .request + .input_serialization + .csv + .as_ref() + .expect("CSV input") + .allow_quoted_record_delimiter, + allow + ); + } + } + + #[test] + fn validate_rejects_quoted_record_delimiters_with_scan_range() { + for compression in [CompressionType::NONE, CompressionType::GZIP, CompressionType::BZIP2] { + for range in [ + ScanRange { + start: Some(0), + end: None, + }, + ScanRange { + start: Some(0), + end: Some(10), + }, + ScanRange { + start: None, + end: Some(10), + }, + ] { + let mut input = base_input(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("CSV input") + .allow_quoted_record_delimiter = Some(true); + input.request.input_serialization.compression_type = Some(CompressionType::from_static(compression)); + input.request.scan_range = Some(range); + let error = validate_select_request(&HeaderMap::new(), &mut input) + .expect_err("quoted record delimiters cannot be combined with ScanRange"); + assert_eq!(error.code(), &S3ErrorCode::InvalidRequestParameter); + assert_eq!(error.message(), Some(INVALID_SCAN_RANGE_MESSAGE)); + } + } + } + #[test] fn validate_accepts_single_unicode_csv_input_quotes() { for quote in ["ع", "界", "🦀"] { @@ -4006,14 +4068,23 @@ mod tests { #[test] fn validate_allows_scan_range_for_csv_as_request_parameter() { - let mut input = base_input(); - input.request.scan_range = Some(ScanRange { - start: Some(0), - end: Some(10), - }); + for allow in [None, Some(false)] { + let mut input = base_input(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("CSV input") + .allow_quoted_record_delimiter = allow; + input.request.scan_range = Some(ScanRange { + start: Some(0), + end: Some(10), + }); - validate_select_request(&HeaderMap::new(), &mut input).expect("csv scan range should validate"); - validate_scan_range_for_object_size(&input.request, 16).expect("csv scan range should validate against object size"); + validate_select_request(&HeaderMap::new(), &mut input).expect("csv scan range should validate"); + validate_scan_range_for_object_size(&input.request, 16).expect("csv scan range should validate against object size"); + } } #[test]