mirror of
https://github.com/rustfs/rustfs.git
synced 2026-09-04 19:25:40 +00:00
fix(s3select): support two-byte CSV record delimiters (#5565)
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -349,9 +349,17 @@ impl SessionCtxFactory {
|
|||||||
};
|
};
|
||||||
let rt = RuntimeEnvBuilder::new().with_memory_limit(memory_limit_bytes, 1.0).build()?;
|
let rt = RuntimeEnvBuilder::new().with_memory_limit(memory_limit_bytes, 1.0).build()?;
|
||||||
let config = SessionConfig::new().with_target_partitions(self.target_partitions);
|
let config = SessionConfig::new().with_target_partitions(self.target_partitions);
|
||||||
|
let custom_two_byte_record_delimiter = context
|
||||||
|
.input
|
||||||
|
.request
|
||||||
|
.input_serialization
|
||||||
|
.csv
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|csv| csv.record_delimiter.as_deref())
|
||||||
|
.is_some_and(|delimiter| delimiter.len() == 2 && delimiter.as_bytes() != b"\r\n");
|
||||||
let scan_range_requires_single_file_scan =
|
let scan_range_requires_single_file_scan =
|
||||||
context.input.request.scan_range.is_some() && context.input.request.input_serialization.parquet.is_none();
|
context.input.request.scan_range.is_some() && context.input.request.input_serialization.parquet.is_none();
|
||||||
let config = if scan_range_requires_single_file_scan {
|
let config = if custom_two_byte_record_delimiter || scan_range_requires_single_file_scan {
|
||||||
config.with_repartition_file_scans(false)
|
config.with_repartition_file_scans(false)
|
||||||
} else {
|
} else {
|
||||||
config
|
config
|
||||||
@@ -586,6 +594,66 @@ mod tests {
|
|||||||
assert!(!session.inner().config().options().optimizer.repartition_file_scans);
|
assert!(!session.inner().config().options().optimizer.repartition_file_scans);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn csv_scan_range_disables_file_repartitioning() {
|
||||||
|
let mut context = test_context();
|
||||||
|
Arc::make_mut(&mut context.input).request.scan_range = Some(ScanRange {
|
||||||
|
start: Some(0),
|
||||||
|
end: Some(0),
|
||||||
|
});
|
||||||
|
|
||||||
|
let session = SessionCtxFactory::new(true)
|
||||||
|
.with_target_partitions(2)
|
||||||
|
.create_session_ctx(&context)
|
||||||
|
.await
|
||||||
|
.expect("CSV ScanRange session should be created");
|
||||||
|
|
||||||
|
assert!(!session.inner().config().options().optimizer.repartition_file_scans);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn two_byte_csv_record_delimiter_disables_file_scan_repartition() {
|
||||||
|
let mut context = test_context();
|
||||||
|
Arc::get_mut(&mut context.input)
|
||||||
|
.expect("test context input should be uniquely owned")
|
||||||
|
.request
|
||||||
|
.input_serialization
|
||||||
|
.csv
|
||||||
|
.as_mut()
|
||||||
|
.expect("test context should use CSV")
|
||||||
|
.record_delimiter = Some("^Y".to_string());
|
||||||
|
let session = SessionCtxFactory::new(true)
|
||||||
|
.with_target_partitions(4)
|
||||||
|
.create_session_ctx(&context)
|
||||||
|
.await
|
||||||
|
.expect("session should be created");
|
||||||
|
|
||||||
|
assert!(!session.inner().config().options().optimizer.repartition_file_scans);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn crlf_record_delimiter_retains_file_scan_repartition() {
|
||||||
|
let mut context = test_context();
|
||||||
|
Arc::get_mut(&mut context.input)
|
||||||
|
.expect("test context input should be uniquely owned")
|
||||||
|
.request
|
||||||
|
.input_serialization
|
||||||
|
.csv
|
||||||
|
.as_mut()
|
||||||
|
.expect("test context should use CSV")
|
||||||
|
.record_delimiter = Some("\r\n".to_string());
|
||||||
|
let session = SessionCtxFactory::new(true)
|
||||||
|
.with_target_partitions(4)
|
||||||
|
.create_session_ctx(&context)
|
||||||
|
.await
|
||||||
|
.expect("session should be created");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
session.inner().config().options().optimizer.repartition_file_scans,
|
||||||
|
SessionConfig::new().options().optimizer.repartition_file_scans
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn session_factory_zero_target_partitions_uses_datafusion_default() {
|
async fn session_factory_zero_target_partitions_uses_datafusion_default() {
|
||||||
let factory = SessionCtxFactory::new(true);
|
let factory = SessionCtxFactory::new(true);
|
||||||
|
|||||||
@@ -250,7 +250,8 @@ fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()>
|
|||||||
validate_single_byte(csv.comments.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
validate_single_byte(csv.comments.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
||||||
validate_single_byte(csv.quote_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
validate_single_byte(csv.quote_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
||||||
validate_single_byte(csv.quote_escape_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
validate_single_byte(csv.quote_escape_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
||||||
validate_record_delimiter(csv.record_delimiter.as_deref())?;
|
validate_input_record_delimiter(csv.record_delimiter.as_deref())?;
|
||||||
|
validate_input_delimiter_pair(csv.field_delimiter.as_deref(), csv.record_delimiter.as_deref())?;
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(json) = input.json.as_mut() {
|
if let Some(json) = input.json.as_mut() {
|
||||||
@@ -276,7 +277,7 @@ fn normalize_output_serialization(output: &mut OutputSerialization) -> S3Result<
|
|||||||
validate_single_byte(csv.field_delimiter.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
validate_single_byte(csv.field_delimiter.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
||||||
validate_single_byte(csv.quote_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
validate_single_byte(csv.quote_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
||||||
validate_single_byte(csv.quote_escape_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
validate_single_byte(csv.quote_escape_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?;
|
||||||
validate_record_delimiter(csv.record_delimiter.as_deref())?;
|
validate_output_record_delimiter(csv.record_delimiter.as_deref())?;
|
||||||
if let Some(quote_fields) = csv.quote_fields.as_ref()
|
if let Some(quote_fields) = csv.quote_fields.as_ref()
|
||||||
&& !matches!(quote_fields.as_str(), QuoteFields::ALWAYS | QuoteFields::ASNEEDED)
|
&& !matches!(quote_fields.as_str(), QuoteFields::ALWAYS | QuoteFields::ASNEEDED)
|
||||||
{
|
{
|
||||||
@@ -351,7 +352,7 @@ fn validate_single_byte(value: Option<&str>, code: S3ErrorCode) -> S3Result<()>
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn validate_record_delimiter(value: Option<&str>) -> S3Result<()> {
|
fn validate_output_record_delimiter(value: Option<&str>) -> S3Result<()> {
|
||||||
if let Some(value) = value
|
if let Some(value) = value
|
||||||
&& value.len() != 1
|
&& value.len() != 1
|
||||||
&& value != "\r\n"
|
&& value != "\r\n"
|
||||||
@@ -361,6 +362,30 @@ fn validate_record_delimiter(value: Option<&str>) -> S3Result<()> {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn validate_input_record_delimiter(value: Option<&str>) -> S3Result<()> {
|
||||||
|
if let Some(value) = value
|
||||||
|
&& !(1..=2).contains(&value.len())
|
||||||
|
{
|
||||||
|
return Err(S3Error::new(S3ErrorCode::InvalidRequestParameter));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_input_delimiter_pair(field_delimiter: Option<&str>, record_delimiter: Option<&str>) -> S3Result<()> {
|
||||||
|
let field_delimiter = field_delimiter.unwrap_or(",");
|
||||||
|
let record_delimiter = record_delimiter.unwrap_or("\n");
|
||||||
|
let normalized_field_delimiter = if field_delimiter.len() > 1 { "," } else { field_delimiter };
|
||||||
|
let normalized_record_delimiter = if record_delimiter.len() == 2 {
|
||||||
|
"\r\n"
|
||||||
|
} else {
|
||||||
|
record_delimiter
|
||||||
|
};
|
||||||
|
if record_delimiter.starts_with(field_delimiter) || normalized_record_delimiter.contains(normalized_field_delimiter) {
|
||||||
|
return Err(S3Error::new(S3ErrorCode::InvalidRequestParameter));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
async fn preflight_select_object(headers: &http::HeaderMap, input: &SelectObjectContentInput) -> S3Result<SelectObjectMetadata> {
|
async fn preflight_select_object(headers: &http::HeaderMap, input: &SelectObjectContentInput) -> S3Result<SelectObjectMetadata> {
|
||||||
let opts = get_opts(&input.bucket, &input.key, None, None, headers)
|
let opts = get_opts(&input.bucket, &input.key, None, None, headers)
|
||||||
.await
|
.await
|
||||||
@@ -1052,6 +1077,120 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn validate_accepts_two_byte_csv_input_record_delimiter() {
|
||||||
|
let mut input = base_input();
|
||||||
|
input
|
||||||
|
.request
|
||||||
|
.input_serialization
|
||||||
|
.csv
|
||||||
|
.as_mut()
|
||||||
|
.expect("base input should use CSV")
|
||||||
|
.record_delimiter = Some("^Y".to_string());
|
||||||
|
|
||||||
|
assert!(validate_select_request(&HeaderMap::new(), &mut input).is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn validate_accepts_supported_multibyte_csv_delimiter_pairs() {
|
||||||
|
let mut input = base_input();
|
||||||
|
let csv = input
|
||||||
|
.request
|
||||||
|
.input_serialization
|
||||||
|
.csv
|
||||||
|
.as_mut()
|
||||||
|
.expect("base input should use CSV");
|
||||||
|
csv.field_delimiter = Some("\r\n".to_string());
|
||||||
|
csv.record_delimiter = Some("^Y".to_string());
|
||||||
|
|
||||||
|
assert!(validate_select_request(&HeaderMap::new(), &mut input).is_ok());
|
||||||
|
|
||||||
|
let csv = input
|
||||||
|
.request
|
||||||
|
.input_serialization
|
||||||
|
.csv
|
||||||
|
.as_mut()
|
||||||
|
.expect("input should still use CSV");
|
||||||
|
csv.field_delimiter = Some("\nX".to_string());
|
||||||
|
csv.record_delimiter = None;
|
||||||
|
assert!(validate_select_request(&HeaderMap::new(), &mut input).is_ok());
|
||||||
|
|
||||||
|
let csv = input
|
||||||
|
.request
|
||||||
|
.input_serialization
|
||||||
|
.csv
|
||||||
|
.as_mut()
|
||||||
|
.expect("input should still use CSV");
|
||||||
|
csv.field_delimiter = Some("aa".to_string());
|
||||||
|
csv.record_delimiter = Some("a".to_string());
|
||||||
|
assert!(validate_select_request(&HeaderMap::new(), &mut input).is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn validate_rejects_csv_input_record_delimiter_outside_one_to_two_bytes() {
|
||||||
|
for delimiter in ["", "^YZ"] {
|
||||||
|
let mut input = base_input();
|
||||||
|
input
|
||||||
|
.request
|
||||||
|
.input_serialization
|
||||||
|
.csv
|
||||||
|
.as_mut()
|
||||||
|
.expect("base input should use CSV")
|
||||||
|
.record_delimiter = Some(delimiter.to_string());
|
||||||
|
|
||||||
|
let err = validate_select_request(&HeaderMap::new(), &mut input)
|
||||||
|
.expect_err("record delimiter outside the supported length must be rejected");
|
||||||
|
|
||||||
|
assert_eq!(err.code(), &S3ErrorCode::InvalidRequestParameter);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn validate_rejects_overlapping_csv_input_delimiters() {
|
||||||
|
for (field_delimiter, record_delimiter) in [
|
||||||
|
(Some("^"), Some("^")),
|
||||||
|
(Some("\r"), Some("\r\n")),
|
||||||
|
(Some("a"), Some("aa")),
|
||||||
|
(None, Some(",")),
|
||||||
|
(Some("\n"), None),
|
||||||
|
(Some("||"), Some(",")),
|
||||||
|
(Some("\n"), Some("^Y")),
|
||||||
|
(Some("\r"), Some("^Y")),
|
||||||
|
] {
|
||||||
|
let mut input = base_input();
|
||||||
|
let csv = input
|
||||||
|
.request
|
||||||
|
.input_serialization
|
||||||
|
.csv
|
||||||
|
.as_mut()
|
||||||
|
.expect("base input should use CSV");
|
||||||
|
csv.field_delimiter = field_delimiter.map(str::to_string);
|
||||||
|
csv.record_delimiter = record_delimiter.map(str::to_string);
|
||||||
|
|
||||||
|
let err = validate_select_request(&HeaderMap::new(), &mut input)
|
||||||
|
.expect_err("overlapping field and record delimiters must be rejected");
|
||||||
|
|
||||||
|
assert_eq!(err.code(), &S3ErrorCode::InvalidRequestParameter);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn validate_keeps_csv_output_record_delimiter_restriction() {
|
||||||
|
let mut input = base_input();
|
||||||
|
input
|
||||||
|
.request
|
||||||
|
.output_serialization
|
||||||
|
.csv
|
||||||
|
.as_mut()
|
||||||
|
.expect("base output should use CSV")
|
||||||
|
.record_delimiter = Some("^Y".to_string());
|
||||||
|
|
||||||
|
let err =
|
||||||
|
validate_select_request(&HeaderMap::new(), &mut input).expect_err("multi-byte output delimiter must remain rejected");
|
||||||
|
|
||||||
|
assert_eq!(err.code(), &S3ErrorCode::InvalidRequestParameter);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn json_encoder_outputs_line_delimited_records() {
|
fn json_encoder_outputs_line_delimited_records() {
|
||||||
let schema =
|
let schema =
|
||||||
|
|||||||
Reference in New Issue
Block a user