gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-08-13 16:46:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(s3select): validate scan_range protocol and parquet overlap (#3176)

Browse Source 
* fix(s3select): enforce scan range protocol and parquet overlap

* fix(s3select): select row groups by start offset

---------

Co-authored-by: houseme <housemecn@gmail.com>
This commit is contained in:
GatewayJ
2026-06-03 23:57:17 +08:00
committed by GitHub
parent 7b57d5b217
commit 785d53fce8
10 changed files with 613 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/s3select-api/src/object_store.rs
+32 -18
View File
@@ -83,11 +83,25 @@ pub struct EcObjectStore {
}
#[derive(Clone, Copy, Debug)]
struct SelectScanRange {
pub struct SelectScanRange {
start: u64,
end: u64,
}
impl SelectScanRange {
pub const fn new(start: u64, end: u64) -> Self {
Self { start, end }
}
pub const fn start(&self) -> u64 {
self.start
}
pub const fn end(&self) -> u64 {
self.end
}
}
#[derive(Clone, Copy, Debug)]
pub struct InvalidScanRange;
@@ -305,7 +319,7 @@ fn http_range_spec_from_start(start: u64) -> HTTPRangeSpec {
}
fn scan_range_read_start(scan_range: SelectScanRange, delimiter: &[u8]) -> u64 {
scan_range.start.saturating_sub(delimiter.len() as u64)
scan_range.start().saturating_sub(delimiter.len() as u64)
}
fn find_delimiter(bytes: &[u8], delimiter: &[u8]) -> Option<usize> {
@@ -328,7 +342,7 @@ fn map_storage_error(bucket: &str, object: &str, err: StorageError) -> o_Error {
}
}
fn scan_range_from_bounds(start: Option<i64>, end: Option<i64>, object_size: u64) -> Result<Option<SelectScanRange>> {
pub fn scan_range_from_bounds(start: Option<i64>, end: Option<i64>, object_size: u64) -> Result<Option<SelectScanRange>> {
parse_scan_range_from_bounds(start, end, object_size).map_err(|_| invalid_scan_range_store_error())
}
@@ -362,14 +376,14 @@ fn parse_scan_range_from_bounds(
if start > 0 {
return Err(InvalidScanRange);
}
return Ok(Some(SelectScanRange { start: 0, end: 0 }));
return Ok(Some(SelectScanRange::new(0, 0)));
}
if start >= object_size {
return Err(InvalidScanRange);
}
}
if object_size == 0 {
return Ok(Some(SelectScanRange { start: 0, end: 0 }));
return Ok(Some(SelectScanRange::new(0, 0)));
}
let last_byte = object_size - 1;
@@ -382,7 +396,7 @@ fn parse_scan_range_from_bounds(
}
(None, None) => return Ok(None),
};
Ok(Some(SelectScanRange { start, end }))
Ok(Some(SelectScanRange::new(start, end)))
}
fn invalid_scan_range_store_error() -> o_Error {
@@ -477,7 +491,7 @@ impl ObjectStore for EcObjectStore {
let delimiter = self.record_delimiter();
let include_header = self.csv_has_header();
let read_start = scan_range_read_start(scan_range, &delimiter);
let header = if include_header && scan_range.start > 0 {
let header = if include_header && scan_range.start() > 0 {
Some(self.read_header_record(original_size, &delimiter, &opts).await?)
} else {
None
@@ -827,11 +841,11 @@ impl<S> ScanRangeState<S> {
fn push_record(&mut self, record: Vec<u8>, record_start: u64) {
let include_header = self.include_header && record_start == 0;
let include_record = record_start >= self.range.start && record_start <= self.range.end;
let include_record = record_start >= self.range.start() && record_start <= self.range.end();
if include_header || include_record {
self.pending.push_back(Bytes::from(record));
} else {
if record_start > self.range.end {
if record_start > self.range.end() {
self.done = true;
}
}
@@ -1103,7 +1117,7 @@ mod test {
#[tokio::test]
async fn test_scan_range_stream_keeps_header_and_selected_record() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2\n1,a\n2,b\n3,c\n"))]);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange { start: 10, end: 11 }, true, 0);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(10, 11), true, 0);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1114,7 +1128,7 @@ mod test {
#[tokio::test]
async fn test_scan_range_stream_skips_record_when_start_is_in_middle() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange { start: 2, end: 7 }, false, 0);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(2, 7), false, 0);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1125,7 +1139,7 @@ mod test {
#[tokio::test]
async fn test_scan_range_stream_keeps_record_when_end_is_in_middle() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange { start: 0, end: 5 }, false, 0);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 5), false, 0);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1136,7 +1150,7 @@ mod test {
#[tokio::test]
async fn test_scan_range_stream_uses_base_offset_for_range_reader() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"\n2,b\n3,c\n"))]);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange { start: 4, end: 7 }, false, 3);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(4, 7), false, 3);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1151,7 +1165,7 @@ mod test {
Ok::<_, std::io::Error>(Bytes::from_static(b"\n1,a\r\n2,b\r")),
Ok::<_, std::io::Error>(Bytes::from_static(b"\n3,c\r\n")),
]);
let mut stream = scan_range_stream(chunks, b"\r\n".to_vec(), SelectScanRange { start: 12, end: 14 }, true, 0);
let mut stream = scan_range_stream(chunks, b"\r\n".to_vec(), SelectScanRange::new(12, 14), true, 0);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1161,7 +1175,7 @@ mod test {
#[test]
fn test_scan_range_read_start_keeps_full_delimiter_boundary() {
let range = SelectScanRange { start: 10, end: 20 };
let range = SelectScanRange::new(10, 20);
assert_eq!(scan_range_read_start(range, b"\n"), 9);
assert_eq!(scan_range_read_start(range, b"\r\n"), 8);
assert_eq!(scan_range_read_start(range, b"abcdef"), 4);
@@ -1176,8 +1190,8 @@ mod test {
#[test]
fn test_scan_range_end_only_uses_aws_suffix_semantics() {
let range = scan_range_from_bounds(None, Some(35), 100).unwrap().unwrap();
assert_eq!(range.start, 65);
assert_eq!(range.end, 99);
assert_eq!(range.start(), 65);
assert_eq!(range.end(), 99);
}
#[test]
@@ -1234,7 +1248,7 @@ mod test {
#[tokio::test]
async fn test_scan_range_output_can_convert_field_delimiter() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"a&&1\nb&&2\n"))]);
let stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange { start: 0, end: 10 }, false, 0);
let stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 10), false, 0);
let mut stream = convert_field_delimiter_stream(stream, Some("&&".to_string()));
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
crates/s3select-api/src/query/session.rs
+36 -12
View File
@@ -131,26 +131,50 @@ fn test_parquet_bytes() -> QueryResult<Vec<u8>> {
Field::new("department", DataType::Utf8, false),
Field::new("salary", DataType::Int32, false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])),
Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie", "Diana", "Eve"])),
Arc::new(Int32Array::from(vec![25, 30, 35, 22, 28])),
Arc::new(StringArray::from(vec!["HR", "IT", "Finance", "Marketing", "IT"])),
Arc::new(Int32Array::from(vec![5000, 6000, 7000, 4500, 5500])),
],
)
.map_err(|e| QueryError::StoreError { e: e.to_string() })?;
let first_batch =
test_parquet_batch(Arc::clone(&schema), &[1, 2], &["Alice", "Bob"], &[25, 30], &["HR", "IT"], &[5000, 6000])?;
let second_batch = test_parquet_batch(
Arc::clone(&schema),
&[3, 4, 5],
&["Charlie", "Diana", "Eve"],
&[35, 22, 28],
&["Finance", "Marketing", "IT"],
&[7000, 4500, 5500],
)?;
let mut bytes = Vec::new();
{
let mut writer =
ArrowWriter::try_new(&mut bytes, schema, None).map_err(|e| QueryError::StoreError { e: e.to_string() })?;
writer
.write(&batch)
.write(&first_batch)
.map_err(|e| QueryError::StoreError { e: e.to_string() })?;
writer.flush().map_err(|e| QueryError::StoreError { e: e.to_string() })?;
writer
.write(&second_batch)
.map_err(|e| QueryError::StoreError { e: e.to_string() })?;
writer.close().map_err(|e| QueryError::StoreError { e: e.to_string() })?;
}
Ok(bytes)
}
fn test_parquet_batch(
schema: Arc<Schema>,
ids: &[i32],
names: &[&str],
ages: &[i32],
departments: &[&str],
salaries: &[i32],
) -> QueryResult<RecordBatch> {
RecordBatch::try_new(
schema,
vec![
Arc::new(Int32Array::from(ids.to_vec())),
Arc::new(StringArray::from(names.to_vec())),
Arc::new(Int32Array::from(ages.to_vec())),
Arc::new(StringArray::from(departments.to_vec())),
Arc::new(Int32Array::from(salaries.to_vec())),
],
)
.map_err(|e| QueryError::StoreError { e: e.to_string() })
}