From 5187f91997eff78e8d2f597b832c19518c2a1f3c Mon Sep 17 00:00:00 2001 From: houseme Date: Sat, 8 Aug 2026 00:00:58 +0800 Subject: [PATCH] fix(s3select): replace deprecated parquet reader (#5811) Implement a local AsyncFileReader over DataFusion's object store re-export so Parquet metadata loading no longer uses the deprecated ParquetObjectReader adapter. Co-authored-by: heihutu Co-authored-by: zhi22915 --- Cargo.lock | 1 + crates/s3select-query/Cargo.toml | 1 + .../src/dispatcher/parquet_table.rs | 58 +++++++++++++++++-- 3 files changed, 55 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5577cff08..ab0fd1490 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10111,6 +10111,7 @@ version = "1.0.0-beta.12" dependencies = [ "async-recursion", "async-trait", + "bytes", "datafusion", "derive_builder", "futures", diff --git a/crates/s3select-query/Cargo.toml b/crates/s3select-query/Cargo.toml index 619997edb..61c687a50 100644 --- a/crates/s3select-query/Cargo.toml +++ b/crates/s3select-query/Cargo.toml @@ -45,6 +45,7 @@ hotpath.workspace = true rustfs-s3select-api = { workspace = true } async-recursion = { workspace = true } async-trait.workspace = true +bytes.workspace = true datafusion = { workspace = true, default-features = false, features = ["parquet", "recursive_protection", "sql"] } derive_builder = { workspace = true } futures = { workspace = true } diff --git a/crates/s3select-query/src/dispatcher/parquet_table.rs b/crates/s3select-query/src/dispatcher/parquet_table.rs index 909291a85..9ef1eac65 100644 --- a/crates/s3select-query/src/dispatcher/parquet_table.rs +++ b/crates/s3select-query/src/dispatcher/parquet_table.rs @@ -12,9 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::{fmt, sync::Arc}; +use std::{fmt, ops::Range, sync::Arc}; use async_trait::async_trait; +use bytes::Bytes; use datafusion::{ arrow::datatypes::SchemaRef, catalog::Session, @@ -27,13 +28,15 @@ use datafusion::{ }, execution::object_store::ObjectStoreUrl, logical_expr::{Expr, TableProviderFilterPushDown, TableType}, - object_store::{ObjectStoreExt, path::Path}, + object_store::{Error as ObjectStoreError, ObjectStore, ObjectStoreExt, path::Path}, parquet::{ - arrow::{ParquetRecordBatchStreamBuilder, async_reader::ParquetObjectReader}, - file::metadata::{ParquetMetaData, RowGroupMetaData}, + arrow::{ParquetRecordBatchStreamBuilder, arrow_reader::ArrowReaderOptions, async_reader::AsyncFileReader}, + errors::{ParquetError, Result as ParquetResult}, + file::metadata::{ParquetMetaData, ParquetMetaDataReader, RowGroupMetaData}, }, physical_plan::ExecutionPlan, }; +use futures::{FutureExt, TryFutureExt, future::BoxFuture}; use rustfs_s3select_api::{ QueryError, QueryResult, object_store::{SelectScanRange, scan_range_from_bounds}, @@ -60,6 +63,43 @@ impl fmt::Debug for ParquetSelectTable { } } +struct ObjectStoreParquetReader { + store: Arc, + path: Path, + file_size: u64, +} + +impl AsyncFileReader for ObjectStoreParquetReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, ParquetResult> { + self.store.get_range(&self.path, range).map_err(parquet_store_error).boxed() + } + + fn get_byte_ranges(&mut self, ranges: Vec>) -> BoxFuture<'_, ParquetResult>> { + async move { self.store.get_ranges(&self.path, &ranges).await.map_err(parquet_store_error) }.boxed() + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, ParquetResult>> { + async move { + let metadata_options = options.map(|options| options.metadata_options().clone()); + let mut metadata_reader = ParquetMetaDataReader::new().with_metadata_options(metadata_options); + + if let Some(options) = options { + metadata_reader = metadata_reader + .with_column_index_policy(options.column_index_policy()) + .with_offset_index_policy(options.offset_index_policy()); + } + + let file_size = self.file_size; + let metadata = metadata_reader.load_and_finish(self, file_size).await?; + Ok(Arc::new(metadata)) + } + .boxed() + } +} + impl ParquetSelectTable { pub async fn try_new(state: &dyn Session, input: &SelectObjectContentInput) -> QueryResult> { let table_path = ListingTableUrl::parse(format!("s3://{}/{}", input.bucket, input.key))?; @@ -68,7 +108,11 @@ impl ParquetSelectTable { 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 reader = ObjectStoreParquetReader { + store: Arc::clone(&store), + path: object_location, + file_size: object_meta.size, + }; let builder = ParquetRecordBatchStreamBuilder::new(reader) .await .map_err(query_store_error)?; @@ -166,6 +210,10 @@ fn non_negative_offset(offset: i64) -> Option { u64::try_from(offset).ok() } +fn parquet_store_error(err: ObjectStoreError) -> ParquetError { + ParquetError::External(Box::new(err)) +} + fn query_store_error(err: impl fmt::Display) -> QueryError { QueryError::StoreError { e: err.to_string() } }