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:
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
@@ -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<dyn TableProvider>),
External(Arc<ListingTable>),
}
pub struct TableHandle(Arc<dyn TableProvider>);
impl TableHandle {
fn provider(&self) -> Arc<dyn TableProvider> {
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<Vec<TableProviderFilterPushDown>> {
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<Arc<dyn TableProvider>> for TableHandle {
fn from(value: Arc<dyn TableProvider>) -> Self {
TableHandle::TableProvider(value)
}
}
impl From<Arc<ListingTable>> for TableHandle {
fn from(value: Arc<ListingTable>) -> 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)
}
}
@@ -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<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 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
@@ -13,3 +13,4 @@
// limitations under the License.
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 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<ListingTable>) -> DFResult<TableHandle> {
Ok(TableHandle::External(provider))
fn build_table_handle(&self, provider: Arc<dyn TableProvider>) -> DFResult<TableHandle> {
Ok(provider.into())
}
}
+4 -4
View File
@@ -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<dyn TableHandleProvider + Send + Sync>;
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 {
provider: Arc<ListingTable>,
provider: Arc<dyn TableProvider>,
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<ListingTable>,
provider: Arc<dyn TableProvider>,
current_session_table_provider: TableHandleProviderRef,
func_manager: FuncMetaManagerRef,
session: SessionCtx,
@@ -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";