diff --git a/crates/s3select-api/src/object_store.rs b/crates/s3select-api/src/object_store.rs index 60d744ebc..76f5d881b 100644 --- a/crates/s3select-api/src/object_store.rs +++ b/crates/s3select-api/src/object_store.rs @@ -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 { @@ -328,7 +342,7 @@ fn map_storage_error(bucket: &str, object: &str, err: StorageError) -> o_Error { } } -fn scan_range_from_bounds(start: Option, end: Option, object_size: u64) -> Result> { +pub fn scan_range_from_bounds(start: Option, end: Option, object_size: u64) -> Result> { 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 ScanRangeState { fn push_record(&mut self, record: Vec, 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 { diff --git a/crates/s3select-api/src/query/session.rs b/crates/s3select-api/src/query/session.rs index 220be2c53..7410bacd1 100644 --- a/crates/s3select-api/src/query/session.rs +++ b/crates/s3select-api/src/query/session.rs @@ -131,26 +131,50 @@ fn test_parquet_bytes() -> QueryResult> { 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, + ids: &[i32], + names: &[&str], + ages: &[i32], + departments: &[&str], + salaries: &[i32], +) -> QueryResult { + 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() }) +} diff --git a/crates/s3select-query/src/data_source/table_source.rs b/crates/s3select-query/src/data_source/table_source.rs index 2760f9a44..cc07659b2 100644 --- a/crates/s3select-query/src/data_source/table_source.rs +++ b/crates/s3select-query/src/data_source/table_source.rs @@ -16,12 +16,10 @@ use std::any::Any; use std::borrow::Cow; use std::fmt::Display; use std::sync::Arc; -use std::write; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::Result as DFResult; -use datafusion::datasource::listing::ListingTable; use datafusion::datasource::{TableProvider, provider_as_source}; use datafusion::error::DataFusionError; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, TableProviderFilterPushDown, TableSource}; @@ -48,21 +46,11 @@ impl TableSourceAdapter { let table_name: String = table_name.into(); let table_handle = table_handle.into(); - let plan = match &table_handle { - // TableScan - TableHandle::External(t) => { - let table_source = provider_as_source(t.clone()); - 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()? - } - } + let table_source = provider_as_source(table_handle.provider()); + let plan = 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()); @@ -109,44 +97,32 @@ impl TableSource for TableSourceAdapter { } #[derive(Clone)] -pub enum TableHandle { - TableProvider(Arc), - External(Arc), -} +pub struct TableHandle(Arc); impl TableHandle { + fn provider(&self) -> Arc { + Arc::clone(&self.0) + } + pub fn schema(&self) -> SchemaRef { - match self { - Self::External(t) => t.schema(), - Self::TableProvider(t) => t.schema(), - } + self.0.schema() } pub fn supports_filters_pushdown(&self, filter: &[&Expr]) -> DFResult> { - match self { - Self::External(t) => t.supports_filters_pushdown(filter), - Self::TableProvider(t) => t.supports_filters_pushdown(filter), - } + self.0.supports_filters_pushdown(filter) } } impl From> for TableHandle { fn from(value: Arc) -> Self { - TableHandle::TableProvider(value) - } -} - -impl From> for TableHandle { - fn from(value: Arc) -> Self { - TableHandle::External(value) + Self(value) } } impl Display for TableHandle { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::External(e) => write!(f, "External({:?})", e.table_paths()), - Self::TableProvider(_) => write!(f, "TableProvider"), - } + let provider_name = std::any::type_name_of_val(self.0.as_ref()); + let short_name = provider_name.rsplit("::").next().unwrap_or(provider_name); + f.write_str(short_name) } } diff --git a/crates/s3select-query/src/dispatcher/manager.rs b/crates/s3select-query/src/dispatcher/manager.rs index 46d977c1d..b18daad49 100644 --- a/crates/s3select-query/src/dispatcher/manager.rs +++ b/crates/s3select-query/src/dispatcher/manager.rs @@ -26,7 +26,7 @@ use datafusion::{ record_batch::RecordBatch, }, datasource::{ - file_format::{csv::CsvFormat, json::JsonFormat, parquet::ParquetFormat}, + file_format::{csv::CsvFormat, json::JsonFormat}, listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl}, }, error::Result as DFResult, @@ -51,6 +51,7 @@ use s3s::dto::{FileHeaderInfo, SelectObjectContentInput}; use std::sync::LazyLock; use crate::{ + dispatcher::parquet_table::ParquetSelectTable, execution::factory::QueryExecutionFactoryRef, metadata::{ContextProviderExtension, MetadataProvider, TableHandleProviderRef, base_table::BaseTableProvider}, sql::logical::planner::DefaultLogicalPlanner, @@ -163,6 +164,15 @@ impl SimpleQueryDispatcher { } async fn build_scheme_provider(&self, session: &SessionCtx) -> QueryResult { + 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 table_path = ListingTableUrl::parse(path)?; let (listing_options, need_rename_volume_name, need_ignore_volume_name) = @@ -218,9 +228,6 @@ impl SimpleQueryDispatcher { need_rename_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() { let file_format = JsonFormat::default(); // Use the actual file extension from the object key so that files stored diff --git a/crates/s3select-query/src/dispatcher/mod.rs b/crates/s3select-query/src/dispatcher/mod.rs index 097451bcc..83ecdb005 100644 --- a/crates/s3select-query/src/dispatcher/mod.rs +++ b/crates/s3select-query/src/dispatcher/mod.rs @@ -13,3 +13,4 @@ // limitations under the License. pub mod manager; +mod parquet_table; diff --git a/crates/s3select-query/src/dispatcher/parquet_table.rs b/crates/s3select-query/src/dispatcher/parquet_table.rs new file mode 100644 index 000000000..ca65e69c2 --- /dev/null +++ b/crates/s3select-query/src/dispatcher/parquet_table.rs @@ -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>, +} + +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> { + 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 = 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>, + filters: &[Expr], + limit: Option, + ) -> DFResult> { + 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 = DataSourceExec::from_data_source(config); + Ok(plan) + } + + fn supports_filters_pushdown(&self, filters: &[&Expr]) -> DFResult> { + Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) + } +} + +fn parquet_access_plan( + input: &SelectObjectContentInput, + object_size: u64, + metadata: &ParquetMetaData, +) -> QueryResult>> { + 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 { + row_group.file_offset().and_then(non_negative_offset) +} + +fn non_negative_offset(offset: i64) -> Option { + 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 { + 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 { + 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::>(); + builder + .set_column_metadata(columns) + .build() + .expect("rewrite test row-group metadata") + } +} diff --git a/crates/s3select-query/src/metadata/base_table.rs b/crates/s3select-query/src/metadata/base_table.rs index 1a30b7ae1..fb814854d 100644 --- a/crates/s3select-query/src/metadata/base_table.rs +++ b/crates/s3select-query/src/metadata/base_table.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use datafusion::common::Result as DFResult; -use datafusion::datasource::listing::ListingTable; +use datafusion::datasource::TableProvider; use crate::data_source::table_source::TableHandle; @@ -25,7 +25,7 @@ use super::TableHandleProvider; pub struct BaseTableProvider {} impl TableHandleProvider for BaseTableProvider { - fn build_table_handle(&self, provider: Arc) -> DFResult { - Ok(TableHandle::External(provider)) + fn build_table_handle(&self, provider: Arc) -> DFResult { + Ok(provider.into()) } } diff --git a/crates/s3select-query/src/metadata/mod.rs b/crates/s3select-query/src/metadata/mod.rs index b5ec7c99e..5b518321e 100644 --- a/crates/s3select-query/src/metadata/mod.rs +++ b/crates/s3select-query/src/metadata/mod.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion::arrow::datatypes::DataType; 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::{AggregateUDF, ScalarUDF, TableSource, WindowUDF}; use datafusion::variable::VarType; @@ -39,11 +39,11 @@ pub trait ContextProviderExtension: ContextProvider { pub type TableHandleProviderRef = Arc; pub trait TableHandleProvider { - fn build_table_handle(&self, provider: Arc) -> DFResult; + fn build_table_handle(&self, provider: Arc) -> DFResult; } pub struct MetadataProvider { - provider: Arc, + provider: Arc, session: SessionCtx, config_options: ConfigOptions, func_manager: FuncMetaManagerRef, @@ -53,7 +53,7 @@ pub struct MetadataProvider { impl MetadataProvider { #[allow(clippy::too_many_arguments)] pub fn new( - provider: Arc, + provider: Arc, current_session_table_provider: TableHandleProviderRef, func_manager: FuncMetaManagerRef, session: SessionCtx, diff --git a/crates/s3select-query/src/test/integration_test.rs b/crates/s3select-query/src/test/integration_test.rs index 833b61de0..4ff8c8c9a 100644 --- a/crates/s3select-query/src/test/integration_test.rs +++ b/crates/s3select-query/src/test/integration_test.rs @@ -21,7 +21,7 @@ mod integration_tests { }; use s3s::dto::{ CSVInput, CSVOutput, ExpressionType, FileHeaderInfo, InputSerialization, JSONInput, JSONOutput, JSONType, - OutputSerialization, ParquetInput, SelectObjectContentInput, SelectObjectContentRequest, + OutputSerialization, ParquetInput, ScanRange, SelectObjectContentInput, SelectObjectContentRequest, }; use std::sync::Arc; @@ -54,8 +54,7 @@ mod integration_tests { } /// Build a `SelectObjectContentInput` targeting a JSON DOCUMENT file. - /// Uses `JSONType::DOCUMENT` so the NDJSON-flattening path in - /// `EcObjectStore` is exercised. + /// Uses `JSONType::DOCUMENT`, which keeps the document-style validation path. fn create_test_json_input(sql: &str) -> SelectObjectContentInput { SelectObjectContentInput { 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 { SelectObjectContentInput { bucket: "test-bucket".to_string(), @@ -132,7 +158,7 @@ mod integration_tests { async fn test_simple_select_query() { let sql = "SELECT * FROM S3Object"; 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 result = db.execute(&query).await; @@ -283,7 +309,6 @@ mod integration_tests { // ────────────────────────────────────────────── // JSON-input variants of all the above tests - // These exercise the JSONType::LINES (JSON lines) code path // ────────────────────────────────────────────── #[tokio::test] @@ -330,6 +355,110 @@ mod integration_tests { 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] async fn test_select_with_where_clause_json() { let sql = "SELECT name, age FROM S3Object WHERE age > 30"; diff --git a/rustfs/src/app/select_object.rs b/rustfs/src/app/select_object.rs index 3edac6f78..4f46f1be4 100644 --- a/rustfs/src/app/select_object.rs +++ b/rustfs/src/app/select_object.rs @@ -246,12 +246,27 @@ fn validate_scan_range(request: &SelectObjectContentRequest) -> S3Result<()> { if validate_scan_range_bounds(start, end, u64::MAX).is_err() { 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) { - return Err(invalid_scan_range_error()); - } + validate_scan_range_protocol(request).map_err(|_| invalid_scan_range_error())?; 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<()> { let Some(scan_range) = request.scan_range.as_ref() else { return Ok(()); @@ -721,6 +736,57 @@ mod tests { 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] fn validate_rejects_scan_range_start_after_object() { let mut input = base_input();