gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-10-03 20:20:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(s3select): support quoted record delimiters (#8203)

Browse Source 
* feat(s3select): support quoted record delimiters

* fix(s3select): preserve CSV comment limits and CR fields

* fix(s3select): preserve CRLF fields and compressed record limits
This commit is contained in:
GatewayJ
2026-10-03 11:10:11 +08:00
committed by GitHub
parent 5efa6254eb
commit f07142910c
5 changed files with 610 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/s3select-api/src/csv_input.rs
+253 -15
View File
@@ -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<u8>,
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<u8>,
escape: Vec<u8>,
@@ -54,6 +56,9 @@ struct CsvInputNormalizer {
record: Vec<u8>,
comment: Option<u8>,
default_records: bool,
allow_quoted_record_delimiter: bool,
record_bytes: usize,
pending_record: Vec<u8>,
state: State,
record_start: bool,
carry: Vec<u8>,
@@ -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<u8>, bytes: &[u8]) {
fn push_value(&self, output: &mut Vec<u8>, 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::<Vec<_>>();
let stream = normalize_csv_stream(futures::stream::iter(chunks), &csv);
let output = datafusion::object_store::delimited::newline_delimited_stream(stream)
.try_collect::<Vec<_>>()
.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::<SelectError>(), 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([
crates/s3select-api/src/object_store.rs
+162 -2
View File
@@ -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<Item = Result<Bytes>> + 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::<Vec<_>>()
.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::<Vec<_>>().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::<SelectError>(), 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";
crates/s3select-api/src/query/session.rs
+23 -1
View File
@@ -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();
crates/s3select-query/src/dispatcher/manager.rs
+85 -5
View File
@@ -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::<StringArray>()
.expect("CSV string column")
.value(row)
.to_owned()
})
.collect::<Vec<_>>(),
);
}
}
assert_eq!(rows, expected, "fixture={index}, header={header}");
}
}
}
#[tokio::test]
async fn select_uses_input_serialization_independently_of_object_extension() {
for (key, json) in [
rustfs/src/app/select_object.rs
+87 -16
View File
@@ -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]