mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-28 16:07:05 +00:00
fix(s3select): validate scan_range protocol and parquet overlap (#3176)
* 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:
@@ -83,11 +83,25 @@ pub struct EcObjectStore {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone, Copy, Debug)]
|
#[derive(Clone, Copy, Debug)]
|
||||||
struct SelectScanRange {
|
pub struct SelectScanRange {
|
||||||
start: u64,
|
start: u64,
|
||||||
end: 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)]
|
#[derive(Clone, Copy, Debug)]
|
||||||
pub struct InvalidScanRange;
|
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 {
|
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> {
|
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())
|
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 {
|
if start > 0 {
|
||||||
return Err(InvalidScanRange);
|
return Err(InvalidScanRange);
|
||||||
}
|
}
|
||||||
return Ok(Some(SelectScanRange { start: 0, end: 0 }));
|
return Ok(Some(SelectScanRange::new(0, 0)));
|
||||||
}
|
}
|
||||||
if start >= object_size {
|
if start >= object_size {
|
||||||
return Err(InvalidScanRange);
|
return Err(InvalidScanRange);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if object_size == 0 {
|
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;
|
let last_byte = object_size - 1;
|
||||||
@@ -382,7 +396,7 @@ fn parse_scan_range_from_bounds(
|
|||||||
}
|
}
|
||||||
(None, None) => return Ok(None),
|
(None, None) => return Ok(None),
|
||||||
};
|
};
|
||||||
Ok(Some(SelectScanRange { start, end }))
|
Ok(Some(SelectScanRange::new(start, end)))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn invalid_scan_range_store_error() -> o_Error {
|
fn invalid_scan_range_store_error() -> o_Error {
|
||||||
@@ -477,7 +491,7 @@ impl ObjectStore for EcObjectStore {
|
|||||||
let delimiter = self.record_delimiter();
|
let delimiter = self.record_delimiter();
|
||||||
let include_header = self.csv_has_header();
|
let include_header = self.csv_has_header();
|
||||||
let read_start = scan_range_read_start(scan_range, &delimiter);
|
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?)
|
Some(self.read_header_record(original_size, &delimiter, &opts).await?)
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
@@ -827,11 +841,11 @@ impl<S> ScanRangeState<S> {
|
|||||||
|
|
||||||
fn push_record(&mut self, record: Vec<u8>, record_start: u64) {
|
fn push_record(&mut self, record: Vec<u8>, record_start: u64) {
|
||||||
let include_header = self.include_header && record_start == 0;
|
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 {
|
if include_header || include_record {
|
||||||
self.pending.push_back(Bytes::from(record));
|
self.pending.push_back(Bytes::from(record));
|
||||||
} else {
|
} else {
|
||||||
if record_start > self.range.end {
|
if record_start > self.range.end() {
|
||||||
self.done = true;
|
self.done = true;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1103,7 +1117,7 @@ mod test {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_scan_range_stream_keeps_header_and_selected_record() {
|
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 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();
|
let mut output = Vec::new();
|
||||||
while let Some(bytes) = stream.next().await {
|
while let Some(bytes) = stream.next().await {
|
||||||
output.extend_from_slice(&bytes.unwrap());
|
output.extend_from_slice(&bytes.unwrap());
|
||||||
@@ -1114,7 +1128,7 @@ mod test {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_scan_range_stream_skips_record_when_start_is_in_middle() {
|
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 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();
|
let mut output = Vec::new();
|
||||||
while let Some(bytes) = stream.next().await {
|
while let Some(bytes) = stream.next().await {
|
||||||
output.extend_from_slice(&bytes.unwrap());
|
output.extend_from_slice(&bytes.unwrap());
|
||||||
@@ -1125,7 +1139,7 @@ mod test {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_scan_range_stream_keeps_record_when_end_is_in_middle() {
|
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 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();
|
let mut output = Vec::new();
|
||||||
while let Some(bytes) = stream.next().await {
|
while let Some(bytes) = stream.next().await {
|
||||||
output.extend_from_slice(&bytes.unwrap());
|
output.extend_from_slice(&bytes.unwrap());
|
||||||
@@ -1136,7 +1150,7 @@ mod test {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_scan_range_stream_uses_base_offset_for_range_reader() {
|
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 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();
|
let mut output = Vec::new();
|
||||||
while let Some(bytes) = stream.next().await {
|
while let Some(bytes) = stream.next().await {
|
||||||
output.extend_from_slice(&bytes.unwrap());
|
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"\n1,a\r\n2,b\r")),
|
||||||
Ok::<_, std::io::Error>(Bytes::from_static(b"\n3,c\r\n")),
|
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();
|
let mut output = Vec::new();
|
||||||
while let Some(bytes) = stream.next().await {
|
while let Some(bytes) = stream.next().await {
|
||||||
output.extend_from_slice(&bytes.unwrap());
|
output.extend_from_slice(&bytes.unwrap());
|
||||||
@@ -1161,7 +1175,7 @@ mod test {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_scan_range_read_start_keeps_full_delimiter_boundary() {
|
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"\n"), 9);
|
||||||
assert_eq!(scan_range_read_start(range, b"\r\n"), 8);
|
assert_eq!(scan_range_read_start(range, b"\r\n"), 8);
|
||||||
assert_eq!(scan_range_read_start(range, b"abcdef"), 4);
|
assert_eq!(scan_range_read_start(range, b"abcdef"), 4);
|
||||||
@@ -1176,8 +1190,8 @@ mod test {
|
|||||||
#[test]
|
#[test]
|
||||||
fn test_scan_range_end_only_uses_aws_suffix_semantics() {
|
fn test_scan_range_end_only_uses_aws_suffix_semantics() {
|
||||||
let range = scan_range_from_bounds(None, Some(35), 100).unwrap().unwrap();
|
let range = scan_range_from_bounds(None, Some(35), 100).unwrap().unwrap();
|
||||||
assert_eq!(range.start, 65);
|
assert_eq!(range.start(), 65);
|
||||||
assert_eq!(range.end, 99);
|
assert_eq!(range.end(), 99);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -1234,7 +1248,7 @@ mod test {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_scan_range_output_can_convert_field_delimiter() {
|
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 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 stream = convert_field_delimiter_stream(stream, Some("&&".to_string()));
|
||||||
let mut output = Vec::new();
|
let mut output = Vec::new();
|
||||||
while let Some(bytes) = stream.next().await {
|
while let Some(bytes) = stream.next().await {
|
||||||
|
|||||||
@@ -131,26 +131,50 @@ fn test_parquet_bytes() -> QueryResult<Vec<u8>> {
|
|||||||
Field::new("department", DataType::Utf8, false),
|
Field::new("department", DataType::Utf8, false),
|
||||||
Field::new("salary", DataType::Int32, false),
|
Field::new("salary", DataType::Int32, false),
|
||||||
]));
|
]));
|
||||||
let batch = RecordBatch::try_new(
|
let first_batch =
|
||||||
schema.clone(),
|
test_parquet_batch(Arc::clone(&schema), &[1, 2], &["Alice", "Bob"], &[25, 30], &["HR", "IT"], &[5000, 6000])?;
|
||||||
vec![
|
let second_batch = test_parquet_batch(
|
||||||
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])),
|
Arc::clone(&schema),
|
||||||
Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie", "Diana", "Eve"])),
|
&[3, 4, 5],
|
||||||
Arc::new(Int32Array::from(vec![25, 30, 35, 22, 28])),
|
&["Charlie", "Diana", "Eve"],
|
||||||
Arc::new(StringArray::from(vec!["HR", "IT", "Finance", "Marketing", "IT"])),
|
&[35, 22, 28],
|
||||||
Arc::new(Int32Array::from(vec![5000, 6000, 7000, 4500, 5500])),
|
&["Finance", "Marketing", "IT"],
|
||||||
],
|
&[7000, 4500, 5500],
|
||||||
)
|
)?;
|
||||||
.map_err(|e| QueryError::StoreError { e: e.to_string() })?;
|
|
||||||
|
|
||||||
let mut bytes = Vec::new();
|
let mut bytes = Vec::new();
|
||||||
{
|
{
|
||||||
let mut writer =
|
let mut writer =
|
||||||
ArrowWriter::try_new(&mut bytes, schema, None).map_err(|e| QueryError::StoreError { e: e.to_string() })?;
|
ArrowWriter::try_new(&mut bytes, schema, None).map_err(|e| QueryError::StoreError { e: e.to_string() })?;
|
||||||
writer
|
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() })?;
|
.map_err(|e| QueryError::StoreError { e: e.to_string() })?;
|
||||||
writer.close().map_err(|e| QueryError::StoreError { e: e.to_string() })?;
|
writer.close().map_err(|e| QueryError::StoreError { e: e.to_string() })?;
|
||||||
}
|
}
|
||||||
Ok(bytes)
|
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() })
|
||||||
|
}
|
||||||
|
|||||||
@@ -16,12 +16,10 @@ use std::any::Any;
|
|||||||
use std::borrow::Cow;
|
use std::borrow::Cow;
|
||||||
use std::fmt::Display;
|
use std::fmt::Display;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::write;
|
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use datafusion::arrow::datatypes::SchemaRef;
|
use datafusion::arrow::datatypes::SchemaRef;
|
||||||
use datafusion::common::Result as DFResult;
|
use datafusion::common::Result as DFResult;
|
||||||
use datafusion::datasource::listing::ListingTable;
|
|
||||||
use datafusion::datasource::{TableProvider, provider_as_source};
|
use datafusion::datasource::{TableProvider, provider_as_source};
|
||||||
use datafusion::error::DataFusionError;
|
use datafusion::error::DataFusionError;
|
||||||
use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, TableProviderFilterPushDown, TableSource};
|
use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, TableProviderFilterPushDown, TableSource};
|
||||||
@@ -48,21 +46,11 @@ impl TableSourceAdapter {
|
|||||||
let table_name: String = table_name.into();
|
let table_name: String = table_name.into();
|
||||||
|
|
||||||
let table_handle = table_handle.into();
|
let table_handle = table_handle.into();
|
||||||
let plan = match &table_handle {
|
let table_source = provider_as_source(table_handle.provider());
|
||||||
// TableScan
|
let plan = if let Some(plan) = table_source.get_logical_plan() {
|
||||||
TableHandle::External(t) => {
|
LogicalPlanBuilder::from(plan.into_owned()).build()?
|
||||||
let table_source = provider_as_source(t.clone());
|
} else {
|
||||||
LogicalPlanBuilder::scan(table_ref, table_source, None)?.build()?
|
LogicalPlanBuilder::scan(table_ref, table_source, None)?.build()?
|
||||||
}
|
|
||||||
// TableScan
|
|
||||||
TableHandle::TableProvider(t) => {
|
|
||||||
let table_source = provider_as_source(t.clone());
|
|
||||||
if let Some(plan) = table_source.get_logical_plan() {
|
|
||||||
LogicalPlanBuilder::from(plan.into_owned()).build()?
|
|
||||||
} else {
|
|
||||||
LogicalPlanBuilder::scan(table_ref, table_source, None)?.build()?
|
|
||||||
}
|
|
||||||
}
|
|
||||||
};
|
};
|
||||||
|
|
||||||
debug!("Table source logical plan node of {}:\n{}", table_name, plan.display_indent_schema());
|
debug!("Table source logical plan node of {}:\n{}", table_name, plan.display_indent_schema());
|
||||||
@@ -109,44 +97,32 @@ impl TableSource for TableSourceAdapter {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
pub enum TableHandle {
|
pub struct TableHandle(Arc<dyn TableProvider>);
|
||||||
TableProvider(Arc<dyn TableProvider>),
|
|
||||||
External(Arc<ListingTable>),
|
|
||||||
}
|
|
||||||
|
|
||||||
impl TableHandle {
|
impl TableHandle {
|
||||||
|
fn provider(&self) -> Arc<dyn TableProvider> {
|
||||||
|
Arc::clone(&self.0)
|
||||||
|
}
|
||||||
|
|
||||||
pub fn schema(&self) -> SchemaRef {
|
pub fn schema(&self) -> SchemaRef {
|
||||||
match self {
|
self.0.schema()
|
||||||
Self::External(t) => t.schema(),
|
|
||||||
Self::TableProvider(t) => t.schema(),
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn supports_filters_pushdown(&self, filter: &[&Expr]) -> DFResult<Vec<TableProviderFilterPushDown>> {
|
pub fn supports_filters_pushdown(&self, filter: &[&Expr]) -> DFResult<Vec<TableProviderFilterPushDown>> {
|
||||||
match self {
|
self.0.supports_filters_pushdown(filter)
|
||||||
Self::External(t) => t.supports_filters_pushdown(filter),
|
|
||||||
Self::TableProvider(t) => t.supports_filters_pushdown(filter),
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl From<Arc<dyn TableProvider>> for TableHandle {
|
impl From<Arc<dyn TableProvider>> for TableHandle {
|
||||||
fn from(value: Arc<dyn TableProvider>) -> Self {
|
fn from(value: Arc<dyn TableProvider>) -> Self {
|
||||||
TableHandle::TableProvider(value)
|
Self(value)
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl From<Arc<ListingTable>> for TableHandle {
|
|
||||||
fn from(value: Arc<ListingTable>) -> Self {
|
|
||||||
TableHandle::External(value)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Display for TableHandle {
|
impl Display for TableHandle {
|
||||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
match self {
|
let provider_name = std::any::type_name_of_val(self.0.as_ref());
|
||||||
Self::External(e) => write!(f, "External({:?})", e.table_paths()),
|
let short_name = provider_name.rsplit("::").next().unwrap_or(provider_name);
|
||||||
Self::TableProvider(_) => write!(f, "TableProvider"),
|
f.write_str(short_name)
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -26,7 +26,7 @@ use datafusion::{
|
|||||||
record_batch::RecordBatch,
|
record_batch::RecordBatch,
|
||||||
},
|
},
|
||||||
datasource::{
|
datasource::{
|
||||||
file_format::{csv::CsvFormat, json::JsonFormat, parquet::ParquetFormat},
|
file_format::{csv::CsvFormat, json::JsonFormat},
|
||||||
listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl},
|
listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl},
|
||||||
},
|
},
|
||||||
error::Result as DFResult,
|
error::Result as DFResult,
|
||||||
@@ -51,6 +51,7 @@ use s3s::dto::{FileHeaderInfo, SelectObjectContentInput};
|
|||||||
use std::sync::LazyLock;
|
use std::sync::LazyLock;
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
|
dispatcher::parquet_table::ParquetSelectTable,
|
||||||
execution::factory::QueryExecutionFactoryRef,
|
execution::factory::QueryExecutionFactoryRef,
|
||||||
metadata::{ContextProviderExtension, MetadataProvider, TableHandleProviderRef, base_table::BaseTableProvider},
|
metadata::{ContextProviderExtension, MetadataProvider, TableHandleProviderRef, base_table::BaseTableProvider},
|
||||||
sql::logical::planner::DefaultLogicalPlanner,
|
sql::logical::planner::DefaultLogicalPlanner,
|
||||||
@@ -163,6 +164,15 @@ impl SimpleQueryDispatcher {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn build_scheme_provider(&self, session: &SessionCtx) -> QueryResult<MetadataProvider> {
|
async fn build_scheme_provider(&self, session: &SessionCtx) -> QueryResult<MetadataProvider> {
|
||||||
|
if self.input.request.input_serialization.parquet.is_some() {
|
||||||
|
let provider = ParquetSelectTable::try_new(session.inner(), self.input.as_ref()).await?;
|
||||||
|
let current_session_table_provider = self.build_table_handle_provider()?;
|
||||||
|
let metadata_provider =
|
||||||
|
MetadataProvider::new(provider, current_session_table_provider, self.func_manager.clone(), session.clone());
|
||||||
|
|
||||||
|
return Ok(metadata_provider);
|
||||||
|
}
|
||||||
|
|
||||||
let path = format!("s3://{}/{}", self.input.bucket, self.input.key);
|
let path = format!("s3://{}/{}", self.input.bucket, self.input.key);
|
||||||
let table_path = ListingTableUrl::parse(path)?;
|
let table_path = ListingTableUrl::parse(path)?;
|
||||||
let (listing_options, need_rename_volume_name, need_ignore_volume_name) =
|
let (listing_options, need_rename_volume_name, need_ignore_volume_name) =
|
||||||
@@ -218,9 +228,6 @@ impl SimpleQueryDispatcher {
|
|||||||
need_rename_volume_name,
|
need_rename_volume_name,
|
||||||
need_ignore_volume_name,
|
need_ignore_volume_name,
|
||||||
)
|
)
|
||||||
} else if self.input.request.input_serialization.parquet.is_some() {
|
|
||||||
let file_format = ParquetFormat::new();
|
|
||||||
(ListingOptions::new(Arc::new(file_format)).with_file_extension(".parquet"), false, false)
|
|
||||||
} else if self.input.request.input_serialization.json.is_some() {
|
} else if self.input.request.input_serialization.json.is_some() {
|
||||||
let file_format = JsonFormat::default();
|
let file_format = JsonFormat::default();
|
||||||
// Use the actual file extension from the object key so that files stored
|
// Use the actual file extension from the object key so that files stored
|
||||||
|
|||||||
@@ -13,3 +13,4 @@
|
|||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
pub mod manager;
|
pub mod manager;
|
||||||
|
mod parquet_table;
|
||||||
|
|||||||
@@ -0,0 +1,307 @@
|
|||||||
|
// Copyright 2024 RustFS Team
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
use std::{any::Any, fmt, sync::Arc};
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use datafusion::{
|
||||||
|
arrow::datatypes::SchemaRef,
|
||||||
|
catalog::Session,
|
||||||
|
common::Result as DFResult,
|
||||||
|
datasource::{
|
||||||
|
TableProvider,
|
||||||
|
listing::{ListingTableUrl, PartitionedFile},
|
||||||
|
physical_plan::{FileScanConfigBuilder, ParquetSource, parquet::ParquetAccessPlan},
|
||||||
|
source::DataSourceExec,
|
||||||
|
},
|
||||||
|
execution::object_store::ObjectStoreUrl,
|
||||||
|
logical_expr::{Expr, TableProviderFilterPushDown, TableType},
|
||||||
|
object_store::{ObjectStoreExt, path::Path},
|
||||||
|
parquet::{
|
||||||
|
arrow::{ParquetRecordBatchStreamBuilder, async_reader::ParquetObjectReader},
|
||||||
|
file::metadata::{ParquetMetaData, RowGroupMetaData},
|
||||||
|
},
|
||||||
|
physical_plan::ExecutionPlan,
|
||||||
|
};
|
||||||
|
use rustfs_s3select_api::{
|
||||||
|
QueryError, QueryResult,
|
||||||
|
object_store::{SelectScanRange, scan_range_from_bounds},
|
||||||
|
};
|
||||||
|
use s3s::dto::SelectObjectContentInput;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ParquetSelectTable {
|
||||||
|
schema: SchemaRef,
|
||||||
|
object_store_url: ObjectStoreUrl,
|
||||||
|
object_path: String,
|
||||||
|
object_size: u64,
|
||||||
|
access_plan: Option<Arc<ParquetAccessPlan>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl fmt::Debug for ParquetSelectTable {
|
||||||
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
|
f.debug_struct("ParquetSelectTable")
|
||||||
|
.field("object_store_url", &self.object_store_url)
|
||||||
|
.field("object_path", &self.object_path)
|
||||||
|
.field("object_size", &self.object_size)
|
||||||
|
.field("has_access_plan", &self.access_plan.is_some())
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ParquetSelectTable {
|
||||||
|
pub async fn try_new(state: &dyn Session, input: &SelectObjectContentInput) -> QueryResult<Arc<dyn TableProvider>> {
|
||||||
|
let table_path = ListingTableUrl::parse(format!("s3://{}/{}", input.bucket, input.key))?;
|
||||||
|
let object_store_url = table_path.object_store();
|
||||||
|
let object_location = Path::from(input.key.clone());
|
||||||
|
let store = state.runtime_env().object_store(&object_store_url)?;
|
||||||
|
let object_meta = store.head(&object_location).await.map_err(query_store_error)?;
|
||||||
|
|
||||||
|
let reader = ParquetObjectReader::new(Arc::clone(&store), object_location).with_file_size(object_meta.size);
|
||||||
|
let builder = ParquetRecordBatchStreamBuilder::new(reader)
|
||||||
|
.await
|
||||||
|
.map_err(query_store_error)?;
|
||||||
|
let schema = Arc::clone(builder.schema());
|
||||||
|
let metadata = Arc::clone(builder.metadata());
|
||||||
|
let access_plan = parquet_access_plan(input, object_meta.size, metadata.as_ref())?;
|
||||||
|
|
||||||
|
Ok(Arc::new(Self {
|
||||||
|
schema,
|
||||||
|
object_store_url,
|
||||||
|
object_path: input.key.clone(),
|
||||||
|
object_size: object_meta.size,
|
||||||
|
access_plan,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn partitioned_file(&self) -> PartitionedFile {
|
||||||
|
let file = PartitionedFile::new(self.object_path.clone(), self.object_size);
|
||||||
|
if let Some(access_plan) = self.access_plan.as_ref() {
|
||||||
|
let extensions: Arc<dyn Any + Send + Sync> = access_plan.clone();
|
||||||
|
file.with_extensions(extensions)
|
||||||
|
} else {
|
||||||
|
file
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl TableProvider for ParquetSelectTable {
|
||||||
|
fn as_any(&self) -> &dyn Any {
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> SchemaRef {
|
||||||
|
Arc::clone(&self.schema)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn table_type(&self) -> TableType {
|
||||||
|
TableType::Base
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn scan(
|
||||||
|
&self,
|
||||||
|
_state: &dyn Session,
|
||||||
|
projection: Option<&Vec<usize>>,
|
||||||
|
filters: &[Expr],
|
||||||
|
limit: Option<usize>,
|
||||||
|
) -> DFResult<Arc<dyn ExecutionPlan>> {
|
||||||
|
let scan_limit = if filters.is_empty() { limit } else { None };
|
||||||
|
let file_source = Arc::new(ParquetSource::new(Arc::clone(&self.schema)));
|
||||||
|
let config = FileScanConfigBuilder::new(self.object_store_url.clone(), file_source)
|
||||||
|
.with_file(self.partitioned_file())
|
||||||
|
.with_projection_indices(projection.cloned())?
|
||||||
|
.with_limit(scan_limit)
|
||||||
|
.build();
|
||||||
|
let plan: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(config);
|
||||||
|
Ok(plan)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn supports_filters_pushdown(&self, filters: &[&Expr]) -> DFResult<Vec<TableProviderFilterPushDown>> {
|
||||||
|
Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parquet_access_plan(
|
||||||
|
input: &SelectObjectContentInput,
|
||||||
|
object_size: u64,
|
||||||
|
metadata: &ParquetMetaData,
|
||||||
|
) -> QueryResult<Option<Arc<ParquetAccessPlan>>> {
|
||||||
|
let Some(scan_range) = input.request.scan_range.as_ref() else {
|
||||||
|
return Ok(None);
|
||||||
|
};
|
||||||
|
let scan_range = scan_range_from_bounds(scan_range.start, scan_range.end, object_size).map_err(query_store_error)?;
|
||||||
|
Ok(scan_range.map(|range| Arc::new(access_plan_for_scan_range(range, metadata))))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn access_plan_for_scan_range(scan_range: SelectScanRange, metadata: &ParquetMetaData) -> ParquetAccessPlan {
|
||||||
|
let mut access_plan = ParquetAccessPlan::new_none(metadata.num_row_groups());
|
||||||
|
for (idx, row_group) in metadata.row_groups().iter().enumerate() {
|
||||||
|
// S3 Select processes a parquet row group when its on-disk start offset
|
||||||
|
// falls inside the requested scan range.
|
||||||
|
if let Some(start) = row_group_start_offset(row_group) {
|
||||||
|
if start >= scan_range.start() && start <= scan_range.end() {
|
||||||
|
access_plan.scan(idx);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// If row-group start offset is unavailable, keep existing behavior and
|
||||||
|
// scan conservatively.
|
||||||
|
access_plan.scan(idx);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
access_plan
|
||||||
|
}
|
||||||
|
|
||||||
|
fn row_group_start_offset(row_group: &RowGroupMetaData) -> Option<u64> {
|
||||||
|
row_group.file_offset().and_then(non_negative_offset)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn non_negative_offset(offset: i64) -> Option<u64> {
|
||||||
|
u64::try_from(offset).ok()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn query_store_error(err: impl fmt::Display) -> QueryError {
|
||||||
|
QueryError::StoreError { e: err.to_string() }
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use datafusion::{
|
||||||
|
arrow::{
|
||||||
|
array::Int32Array,
|
||||||
|
datatypes::{DataType, Field, Schema, SchemaRef},
|
||||||
|
record_batch::RecordBatch,
|
||||||
|
},
|
||||||
|
parquet::arrow::{ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder},
|
||||||
|
};
|
||||||
|
use std::{
|
||||||
|
fs::File,
|
||||||
|
sync::Arc,
|
||||||
|
time::{SystemTime, UNIX_EPOCH},
|
||||||
|
};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn access_plan_selects_row_group_by_row_group_start() {
|
||||||
|
let metadata = two_row_group_metadata();
|
||||||
|
let first_start = row_group_start_offset(&metadata.row_groups()[0]).expect("first row group should have start offset");
|
||||||
|
let second_start = row_group_start_offset(&metadata.row_groups()[1]).expect("second row group should have start offset");
|
||||||
|
|
||||||
|
let first_start_plan = access_plan_for_scan_range(SelectScanRange::new(first_start, first_start), metadata.as_ref());
|
||||||
|
assert!(first_start_plan.should_scan(0));
|
||||||
|
let second_start_plan = access_plan_for_scan_range(SelectScanRange::new(second_start, second_start), metadata.as_ref());
|
||||||
|
assert!(second_start_plan.should_scan(1));
|
||||||
|
if first_start != second_start {
|
||||||
|
assert!(!first_start_plan.should_scan(1));
|
||||||
|
assert!(!second_start_plan.should_scan(0));
|
||||||
|
}
|
||||||
|
|
||||||
|
let ((lower_start, lower_idx), (higher_start, higher_idx)) = if first_start <= second_start {
|
||||||
|
((first_start, 0usize), (second_start, 1usize))
|
||||||
|
} else {
|
||||||
|
((second_start, 1usize), (first_start, 0usize))
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(before_higher) = higher_start.checked_sub(1)
|
||||||
|
&& lower_start <= before_higher
|
||||||
|
{
|
||||||
|
let boundary_plan = access_plan_for_scan_range(SelectScanRange::new(lower_start, before_higher), metadata.as_ref());
|
||||||
|
assert!(boundary_plan.should_scan(lower_idx));
|
||||||
|
assert!(!boundary_plan.should_scan(higher_idx));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn access_plan_uses_row_group_start_not_column_span_overlap() {
|
||||||
|
let metadata = synthetic_overlap_metadata();
|
||||||
|
// Range ends inside the first row group byte span, but should only include
|
||||||
|
// the row group whose start offset is within the requested range.
|
||||||
|
let plan = access_plan_for_scan_range(SelectScanRange::new(100, 190), metadata.as_ref());
|
||||||
|
assert!(plan.should_scan(0));
|
||||||
|
assert!(!plan.should_scan(1));
|
||||||
|
}
|
||||||
|
|
||||||
|
fn two_row_group_metadata() -> Arc<ParquetMetaData> {
|
||||||
|
let now = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.expect("system time should be after unix epoch")
|
||||||
|
.as_nanos();
|
||||||
|
let path = std::env::temp_dir().join(format!("rustfs_s3select_parquet_scan_range_{now}.parquet"));
|
||||||
|
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
|
||||||
|
|
||||||
|
{
|
||||||
|
let file = File::create(&path).expect("create parquet test file");
|
||||||
|
let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).expect("create parquet writer");
|
||||||
|
writer
|
||||||
|
.write(&single_i32_batch(Arc::clone(&schema), 1))
|
||||||
|
.expect("write first row group");
|
||||||
|
writer.flush().expect("flush first row group");
|
||||||
|
writer
|
||||||
|
.write(&single_i32_batch(Arc::clone(&schema), 2))
|
||||||
|
.expect("write second row group");
|
||||||
|
writer.close().expect("close parquet writer");
|
||||||
|
}
|
||||||
|
|
||||||
|
let file = File::open(&path).expect("open parquet test file");
|
||||||
|
let metadata = ParquetRecordBatchReaderBuilder::try_new(file)
|
||||||
|
.expect("read parquet metadata")
|
||||||
|
.metadata()
|
||||||
|
.clone();
|
||||||
|
std::fs::remove_file(&path).expect("remove parquet test file");
|
||||||
|
metadata
|
||||||
|
}
|
||||||
|
|
||||||
|
fn single_i32_batch(schema: SchemaRef, value: i32) -> RecordBatch {
|
||||||
|
RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![value]))]).expect("create test record batch")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn synthetic_overlap_metadata() -> Arc<ParquetMetaData> {
|
||||||
|
let metadata = two_row_group_metadata();
|
||||||
|
let mut metadata_builder = metadata.as_ref().clone().into_builder();
|
||||||
|
let mut row_groups = metadata_builder.take_row_groups();
|
||||||
|
assert_eq!(row_groups.len(), 2, "test metadata should contain two row groups");
|
||||||
|
|
||||||
|
let first = row_group_with_offsets(row_groups.remove(0), 100, 100, 250);
|
||||||
|
let second = row_group_with_offsets(row_groups.remove(0), 500, 160, 90);
|
||||||
|
let row_groups = vec![first, second];
|
||||||
|
|
||||||
|
metadata_builder = metadata_builder.set_row_groups(row_groups);
|
||||||
|
Arc::new(metadata_builder.build())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn row_group_with_offsets(
|
||||||
|
row_group: RowGroupMetaData,
|
||||||
|
file_offset: i64,
|
||||||
|
column_offset: i64,
|
||||||
|
column_len: i64,
|
||||||
|
) -> RowGroupMetaData {
|
||||||
|
let mut builder = row_group.into_builder().set_file_offset(file_offset);
|
||||||
|
let columns = builder
|
||||||
|
.take_columns()
|
||||||
|
.into_iter()
|
||||||
|
.map(|column| {
|
||||||
|
column
|
||||||
|
.into_builder()
|
||||||
|
.set_data_page_offset(column_offset)
|
||||||
|
.set_total_compressed_size(column_len)
|
||||||
|
.build()
|
||||||
|
.expect("rewrite test column metadata")
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
builder
|
||||||
|
.set_column_metadata(columns)
|
||||||
|
.build()
|
||||||
|
.expect("rewrite test row-group metadata")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -15,7 +15,7 @@
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use datafusion::common::Result as DFResult;
|
use datafusion::common::Result as DFResult;
|
||||||
use datafusion::datasource::listing::ListingTable;
|
use datafusion::datasource::TableProvider;
|
||||||
|
|
||||||
use crate::data_source::table_source::TableHandle;
|
use crate::data_source::table_source::TableHandle;
|
||||||
|
|
||||||
@@ -25,7 +25,7 @@ use super::TableHandleProvider;
|
|||||||
pub struct BaseTableProvider {}
|
pub struct BaseTableProvider {}
|
||||||
|
|
||||||
impl TableHandleProvider for BaseTableProvider {
|
impl TableHandleProvider for BaseTableProvider {
|
||||||
fn build_table_handle(&self, provider: Arc<ListingTable>) -> DFResult<TableHandle> {
|
fn build_table_handle(&self, provider: Arc<dyn TableProvider>) -> DFResult<TableHandle> {
|
||||||
Ok(TableHandle::External(provider))
|
Ok(provider.into())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ use std::sync::Arc;
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use datafusion::arrow::datatypes::DataType;
|
use datafusion::arrow::datatypes::DataType;
|
||||||
use datafusion::common::Result as DFResult;
|
use datafusion::common::Result as DFResult;
|
||||||
use datafusion::datasource::listing::ListingTable;
|
use datafusion::datasource::TableProvider;
|
||||||
use datafusion::logical_expr::var_provider::is_system_variables;
|
use datafusion::logical_expr::var_provider::is_system_variables;
|
||||||
use datafusion::logical_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF};
|
use datafusion::logical_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF};
|
||||||
use datafusion::variable::VarType;
|
use datafusion::variable::VarType;
|
||||||
@@ -39,11 +39,11 @@ pub trait ContextProviderExtension: ContextProvider {
|
|||||||
pub type TableHandleProviderRef = Arc<dyn TableHandleProvider + Send + Sync>;
|
pub type TableHandleProviderRef = Arc<dyn TableHandleProvider + Send + Sync>;
|
||||||
|
|
||||||
pub trait TableHandleProvider {
|
pub trait TableHandleProvider {
|
||||||
fn build_table_handle(&self, provider: Arc<ListingTable>) -> DFResult<TableHandle>;
|
fn build_table_handle(&self, provider: Arc<dyn TableProvider>) -> DFResult<TableHandle>;
|
||||||
}
|
}
|
||||||
|
|
||||||
pub struct MetadataProvider {
|
pub struct MetadataProvider {
|
||||||
provider: Arc<ListingTable>,
|
provider: Arc<dyn TableProvider>,
|
||||||
session: SessionCtx,
|
session: SessionCtx,
|
||||||
config_options: ConfigOptions,
|
config_options: ConfigOptions,
|
||||||
func_manager: FuncMetaManagerRef,
|
func_manager: FuncMetaManagerRef,
|
||||||
@@ -53,7 +53,7 @@ pub struct MetadataProvider {
|
|||||||
impl MetadataProvider {
|
impl MetadataProvider {
|
||||||
#[allow(clippy::too_many_arguments)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
pub fn new(
|
pub fn new(
|
||||||
provider: Arc<ListingTable>,
|
provider: Arc<dyn TableProvider>,
|
||||||
current_session_table_provider: TableHandleProviderRef,
|
current_session_table_provider: TableHandleProviderRef,
|
||||||
func_manager: FuncMetaManagerRef,
|
func_manager: FuncMetaManagerRef,
|
||||||
session: SessionCtx,
|
session: SessionCtx,
|
||||||
|
|||||||
@@ -21,7 +21,7 @@ mod integration_tests {
|
|||||||
};
|
};
|
||||||
use s3s::dto::{
|
use s3s::dto::{
|
||||||
CSVInput, CSVOutput, ExpressionType, FileHeaderInfo, InputSerialization, JSONInput, JSONOutput, JSONType,
|
CSVInput, CSVOutput, ExpressionType, FileHeaderInfo, InputSerialization, JSONInput, JSONOutput, JSONType,
|
||||||
OutputSerialization, ParquetInput, SelectObjectContentInput, SelectObjectContentRequest,
|
OutputSerialization, ParquetInput, ScanRange, SelectObjectContentInput, SelectObjectContentRequest,
|
||||||
};
|
};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
@@ -54,8 +54,7 @@ mod integration_tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Build a `SelectObjectContentInput` targeting a JSON DOCUMENT file.
|
/// Build a `SelectObjectContentInput` targeting a JSON DOCUMENT file.
|
||||||
/// Uses `JSONType::DOCUMENT` so the NDJSON-flattening path in
|
/// Uses `JSONType::DOCUMENT`, which keeps the document-style validation path.
|
||||||
/// `EcObjectStore` is exercised.
|
|
||||||
fn create_test_json_input(sql: &str) -> SelectObjectContentInput {
|
fn create_test_json_input(sql: &str) -> SelectObjectContentInput {
|
||||||
SelectObjectContentInput {
|
SelectObjectContentInput {
|
||||||
bucket: "test-bucket".to_string(),
|
bucket: "test-bucket".to_string(),
|
||||||
@@ -83,6 +82,33 @@ mod integration_tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn create_test_json_lines_input(sql: &str) -> SelectObjectContentInput {
|
||||||
|
SelectObjectContentInput {
|
||||||
|
bucket: "test-bucket".to_string(),
|
||||||
|
expected_bucket_owner: None,
|
||||||
|
key: "test.json".to_string(),
|
||||||
|
sse_customer_algorithm: None,
|
||||||
|
sse_customer_key: None,
|
||||||
|
sse_customer_key_md5: None,
|
||||||
|
request: SelectObjectContentRequest {
|
||||||
|
expression: sql.to_string(),
|
||||||
|
expression_type: ExpressionType::from_static("SQL"),
|
||||||
|
input_serialization: InputSerialization {
|
||||||
|
json: Some(JSONInput {
|
||||||
|
type_: Some(JSONType::from_static(JSONType::LINES)),
|
||||||
|
}),
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
output_serialization: OutputSerialization {
|
||||||
|
json: Some(JSONOutput::default()),
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
request_progress: None,
|
||||||
|
scan_range: None,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn create_test_parquet_input(sql: &str) -> SelectObjectContentInput {
|
fn create_test_parquet_input(sql: &str) -> SelectObjectContentInput {
|
||||||
SelectObjectContentInput {
|
SelectObjectContentInput {
|
||||||
bucket: "test-bucket".to_string(),
|
bucket: "test-bucket".to_string(),
|
||||||
@@ -132,7 +158,7 @@ mod integration_tests {
|
|||||||
async fn test_simple_select_query() {
|
async fn test_simple_select_query() {
|
||||||
let sql = "SELECT * FROM S3Object";
|
let sql = "SELECT * FROM S3Object";
|
||||||
let input = create_test_input(sql);
|
let input = create_test_input(sql);
|
||||||
let db = get_global_db(input.clone(), true).await.unwrap();
|
let db = get_global_db(input.clone(), true).await.expect("create csv test database");
|
||||||
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
||||||
|
|
||||||
let result = db.execute(&query).await;
|
let result = db.execute(&query).await;
|
||||||
@@ -283,7 +309,6 @@ mod integration_tests {
|
|||||||
|
|
||||||
// ──────────────────────────────────────────────
|
// ──────────────────────────────────────────────
|
||||||
// JSON-input variants of all the above tests
|
// JSON-input variants of all the above tests
|
||||||
// These exercise the JSONType::LINES (JSON lines) code path
|
|
||||||
// ──────────────────────────────────────────────
|
// ──────────────────────────────────────────────
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -330,6 +355,110 @@ mod integration_tests {
|
|||||||
assert_eq!(total_rows, 3);
|
assert_eq!(total_rows, 3);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_simple_select_query_parquet_with_scan_range_filters_row_groups() {
|
||||||
|
let sql = "SELECT name, age FROM S3Object WHERE age > 25";
|
||||||
|
let mut input = create_test_parquet_input(sql);
|
||||||
|
input.request.scan_range = Some(ScanRange {
|
||||||
|
start: Some(0),
|
||||||
|
end: Some(1),
|
||||||
|
});
|
||||||
|
let db = get_global_db(input.clone(), true)
|
||||||
|
.await
|
||||||
|
.expect("create parquet scan range test database");
|
||||||
|
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
||||||
|
|
||||||
|
let result = db.execute(&query).await;
|
||||||
|
assert!(result.is_ok());
|
||||||
|
|
||||||
|
let output = result
|
||||||
|
.expect("execute parquet scan range query")
|
||||||
|
.result()
|
||||||
|
.chunk_result()
|
||||||
|
.await
|
||||||
|
.expect("collect parquet scan range query output");
|
||||||
|
let total_rows: usize = output.iter().map(|batch| batch.num_rows()).sum();
|
||||||
|
assert_eq!(total_rows, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_simple_select_query_parquet_with_full_scan_range() {
|
||||||
|
let sql = "SELECT * FROM S3Object";
|
||||||
|
let mut input = create_test_parquet_input(sql);
|
||||||
|
input.request.scan_range = Some(ScanRange {
|
||||||
|
start: Some(0),
|
||||||
|
end: Some(1024),
|
||||||
|
});
|
||||||
|
let db = get_global_db(input.clone(), true)
|
||||||
|
.await
|
||||||
|
.expect("create parquet full scan range database");
|
||||||
|
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
||||||
|
|
||||||
|
let result = db.execute(&query).await;
|
||||||
|
assert!(result.is_ok());
|
||||||
|
|
||||||
|
let output = result
|
||||||
|
.expect("execute parquet full scan range query")
|
||||||
|
.result()
|
||||||
|
.chunk_result()
|
||||||
|
.await
|
||||||
|
.expect("collect parquet full scan range output");
|
||||||
|
let total_rows: usize = output.iter().map(|batch| batch.num_rows()).sum();
|
||||||
|
assert_eq!(total_rows, 5);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_simple_select_query_csv_with_scan_range() {
|
||||||
|
let sql = "SELECT name, age FROM S3Object LIMIT 20";
|
||||||
|
let mut input = create_test_input(sql);
|
||||||
|
input.request.scan_range = Some(ScanRange {
|
||||||
|
start: Some(0),
|
||||||
|
end: Some(1024),
|
||||||
|
});
|
||||||
|
let db = get_global_db(input.clone(), true)
|
||||||
|
.await
|
||||||
|
.expect("create csv scan range test database");
|
||||||
|
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
||||||
|
|
||||||
|
let result = db.execute(&query).await;
|
||||||
|
assert!(result.is_ok());
|
||||||
|
|
||||||
|
let output = result
|
||||||
|
.expect("execute csv scan range query")
|
||||||
|
.result()
|
||||||
|
.chunk_result()
|
||||||
|
.await
|
||||||
|
.expect("collect csv scan range output");
|
||||||
|
let total_rows: usize = output.iter().map(|batch| batch.num_rows()).sum();
|
||||||
|
assert!(total_rows > 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_simple_select_query_json_with_scan_range() {
|
||||||
|
let sql = "SELECT name, age FROM S3Object LIMIT 20";
|
||||||
|
let mut input = create_test_json_lines_input(sql);
|
||||||
|
input.request.scan_range = Some(ScanRange {
|
||||||
|
start: Some(0),
|
||||||
|
end: Some(1024),
|
||||||
|
});
|
||||||
|
let db = get_global_db(input.clone(), true)
|
||||||
|
.await
|
||||||
|
.expect("create json scan range test database");
|
||||||
|
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
|
||||||
|
|
||||||
|
let result = db.execute(&query).await;
|
||||||
|
assert!(result.is_ok());
|
||||||
|
|
||||||
|
let output = result
|
||||||
|
.expect("execute json scan range query")
|
||||||
|
.result()
|
||||||
|
.chunk_result()
|
||||||
|
.await
|
||||||
|
.expect("collect json scan range output");
|
||||||
|
let total_rows: usize = output.iter().map(|batch| batch.num_rows()).sum();
|
||||||
|
assert!(total_rows > 0);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_select_with_where_clause_json() {
|
async fn test_select_with_where_clause_json() {
|
||||||
let sql = "SELECT name, age FROM S3Object WHERE age > 30";
|
let sql = "SELECT name, age FROM S3Object WHERE age > 30";
|
||||||
|
|||||||
@@ -246,12 +246,27 @@ fn validate_scan_range(request: &SelectObjectContentRequest) -> S3Result<()> {
|
|||||||
if validate_scan_range_bounds(start, end, u64::MAX).is_err() {
|
if validate_scan_range_bounds(start, end, u64::MAX).is_err() {
|
||||||
return Err(invalid_scan_range_error());
|
return Err(invalid_scan_range_error());
|
||||||
}
|
}
|
||||||
if request.input_serialization.parquet.is_some() || request.input_serialization.json.as_ref().is_some_and(is_json_document) {
|
validate_scan_range_protocol(request).map_err(|_| invalid_scan_range_error())?;
|
||||||
return Err(invalid_scan_range_error());
|
|
||||||
}
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn validate_scan_range_protocol(request: &SelectObjectContentRequest) -> Result<(), ()> {
|
||||||
|
let input_serialization = &request.input_serialization;
|
||||||
|
|
||||||
|
let supports_scan_range = match (
|
||||||
|
input_serialization.csv.as_ref(),
|
||||||
|
input_serialization.json.as_ref(),
|
||||||
|
input_serialization.parquet.as_ref(),
|
||||||
|
) {
|
||||||
|
(Some(_), None, None) => true,
|
||||||
|
(None, Some(json), None) if !is_json_document(json) => true,
|
||||||
|
(None, None, Some(_)) => true,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
if supports_scan_range { Ok(()) } else { Err(()) }
|
||||||
|
}
|
||||||
|
|
||||||
fn validate_scan_range_for_object_size(request: &SelectObjectContentRequest, object_size: u64) -> S3Result<()> {
|
fn validate_scan_range_for_object_size(request: &SelectObjectContentRequest, object_size: u64) -> S3Result<()> {
|
||||||
let Some(scan_range) = request.scan_range.as_ref() else {
|
let Some(scan_range) = request.scan_range.as_ref() else {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
@@ -721,6 +736,57 @@ mod tests {
|
|||||||
assert_eq!(err.message(), Some(INVALID_SCAN_RANGE_MESSAGE));
|
assert_eq!(err.message(), Some(INVALID_SCAN_RANGE_MESSAGE));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn validate_allows_scan_range_for_json_lines_as_request_parameter() {
|
||||||
|
let mut input = base_input();
|
||||||
|
input.request.input_serialization = InputSerialization {
|
||||||
|
csv: None,
|
||||||
|
json: Some(JSONInput {
|
||||||
|
type_: Some(JSONType::from_static(JSONType::LINES)),
|
||||||
|
}),
|
||||||
|
parquet: None,
|
||||||
|
compression_type: None,
|
||||||
|
};
|
||||||
|
input.request.scan_range = Some(ScanRange {
|
||||||
|
start: Some(0),
|
||||||
|
end: Some(10),
|
||||||
|
});
|
||||||
|
|
||||||
|
validate_select_request(&HeaderMap::new(), &mut input).expect("json lines scan range should validate");
|
||||||
|
validate_scan_range_for_object_size(&input.request, 16)
|
||||||
|
.expect("json lines scan range should validate against object size");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn validate_allows_scan_range_for_parquet_as_request_parameter() {
|
||||||
|
let mut input = base_input();
|
||||||
|
input.request.input_serialization = InputSerialization {
|
||||||
|
csv: None,
|
||||||
|
json: None,
|
||||||
|
parquet: Some(ParquetInput {}),
|
||||||
|
compression_type: None,
|
||||||
|
};
|
||||||
|
input.request.scan_range = Some(ScanRange {
|
||||||
|
start: Some(0),
|
||||||
|
end: Some(10),
|
||||||
|
});
|
||||||
|
|
||||||
|
validate_select_request(&HeaderMap::new(), &mut input).expect("parquet scan range should validate");
|
||||||
|
validate_scan_range_for_object_size(&input.request, 16).expect("parquet scan range should validate against object size");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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),
|
||||||
|
});
|
||||||
|
|
||||||
|
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]
|
#[test]
|
||||||
fn validate_rejects_scan_range_start_after_object() {
|
fn validate_rejects_scan_range_start_after_object() {
|
||||||
let mut input = base_input();
|
let mut input = base_input();
|
||||||
|
|||||||
Reference in New Issue
Block a user