From 0ac1095c70f62715bbfd81ac647400cb8c6dcdd3 Mon Sep 17 00:00:00 2001 From: junxiang Mu <1948535941@qq.com> Date: Tue, 6 May 2025 11:00:18 +0800 Subject: [PATCH] support spec char as delimiter Signed-off-by: junxiang Mu <1948535941@qq.com> --- .cargo/config.toml | 7 -- Cargo.lock | 2 + common/common/src/lib.rs | 3 + ecstore/Cargo.toml | 2 +- rustfs/src/storage/ecfs.rs | 5 +- s3select/api/Cargo.toml | 2 + s3select/api/src/object_store.rs | 127 ++++++++++++++++++++--- s3select/api/src/query/mod.rs | 4 +- s3select/api/src/query/session.rs | 2 +- s3select/query/src/dispatcher/manager.rs | 10 +- s3select/query/src/instance.rs | 14 +-- 11 files changed, 139 insertions(+), 39 deletions(-) delete mode 100644 .cargo/config.toml diff --git a/.cargo/config.toml b/.cargo/config.toml deleted file mode 100644 index a1c92ecf3..000000000 --- a/.cargo/config.toml +++ /dev/null @@ -1,7 +0,0 @@ -[target.x86_64-unknown-linux-gnu] -rustflags = [ - "-C", "link-arg=-fuse-ld=bfd" -] - -[target.x86_64-unknown-linux-musl] -linker = "x86_64-linux-musl-gcc" \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 50200628c..e4c51bc8c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -185,12 +185,14 @@ dependencies = [ "async-trait", "bytes", "chrono", + "common", "datafusion", "ecstore", "futures", "futures-core", "http", "object_store", + "pin-project-lite", "s3s", "snafu", "tokio", diff --git a/common/common/src/lib.rs b/common/common/src/lib.rs index a8998fee0..90250c5eb 100644 --- a/common/common/src/lib.rs +++ b/common/common/src/lib.rs @@ -2,6 +2,9 @@ pub mod error; pub mod globals; pub mod last_minute; +// is ',' +pub static DEFAULT_DELIMITER: u8 = 44; + /// Defers evaluation of a block of code until the end of the scope. #[macro_export] macro_rules! defer { diff --git a/ecstore/Cargo.toml b/ecstore/Cargo.toml index 8ba12ac1d..007150d40 100644 --- a/ecstore/Cargo.toml +++ b/ecstore/Cargo.toml @@ -82,4 +82,4 @@ winapi = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } [build-dependencies] -shadow-rs.workspace = true +shadow-rs = { workspace = true, features = ["build", "metadata"] } diff --git a/rustfs/src/storage/ecfs.rs b/rustfs/src/storage/ecfs.rs index e065fe7a1..fd7b4c0c9 100644 --- a/rustfs/src/storage/ecfs.rs +++ b/rustfs/src/storage/ecfs.rs @@ -62,6 +62,7 @@ use s3s::S3; use s3s::{S3Request, S3Response}; use std::fmt::Debug; use std::str::FromStr; +use std::sync::Arc; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use tokio_util::io::ReaderStream; @@ -1896,14 +1897,14 @@ impl S3 for FS { ) -> S3Result> { info!("handle select_object_content"); - let input = req.input; + let input = Arc::new(req.input); info!("{:?}", input); let db = make_rustfsms(input.clone(), false).await.map_err(|e| { error!("make db failed, {}", e.to_string()); s3_error!(InternalError, "{}", e.to_string()) })?; - let query = Query::new(Context { input: input.clone() }, input.request.expression); + let query = Query::new(Context { input: input.clone() }, input.request.expression.clone()); let result = db .execute(&query) .await diff --git a/s3select/api/Cargo.toml b/s3select/api/Cargo.toml index 27e260c41..125849cdc 100644 --- a/s3select/api/Cargo.toml +++ b/s3select/api/Cargo.toml @@ -7,12 +7,14 @@ edition.workspace = true async-trait.workspace = true bytes.workspace = true chrono.workspace = true +common.workspace = true datafusion = { workspace = true } ecstore.workspace = true futures = { workspace = true } futures-core = { workspace = true } http.workspace = true object_store = { workspace = true } +pin-project-lite.workspace = true s3s.workspace = true snafu = { workspace = true, features = ["backtrace"] } tokio.workspace = true diff --git a/s3select/api/src/object_store.rs b/s3select/api/src/object_store.rs index 7772936ce..52132bc1a 100644 --- a/s3select/api/src/object_store.rs +++ b/s3select/api/src/object_store.rs @@ -1,6 +1,7 @@ use async_trait::async_trait; use bytes::Bytes; use chrono::Utc; +use common::DEFAULT_DELIMITER; use ecstore::io::READ_BUFFER_SIZE; use ecstore::new_object_layer_fn; use ecstore::store::ECStore; @@ -24,29 +25,54 @@ use object_store::PutOptions; use object_store::PutPayload; use object_store::PutResult; use object_store::{Error as o_Error, Result}; +use pin_project_lite::pin_project; use s3s::dto::SelectObjectContentInput; use s3s::s3_error; use s3s::S3Result; use std::ops::Range; +use std::pin::Pin; use std::sync::Arc; +use std::task::ready; +use std::task::Poll; +use tokio::io::AsyncRead; use tokio_util::io::ReaderStream; use tracing::info; use transform_stream::AsyncTryStream; #[derive(Debug)] pub struct EcObjectStore { - input: SelectObjectContentInput, + input: Arc, + need_convert: bool, + delimiter: String, store: Arc, } - impl EcObjectStore { - pub fn new(input: SelectObjectContentInput) -> S3Result { + pub fn new(input: Arc) -> S3Result { let Some(store) = new_object_layer_fn() else { return Err(s3_error!(InternalError, "ec store not inited")); }; - Ok(Self { input, store }) + let (need_convert, delimiter) = if let Some(csv) = input.request.input_serialization.csv.as_ref() { + if let Some(delimiter) = csv.field_delimiter.as_ref() { + if delimiter.len() > 1 { + (true, delimiter.to_owned()) + } else { + (false, String::new()) + } + } else { + (false, String::new()) + } + } else { + (false, String::new()) + }; + + Ok(Self { + input, + need_convert, + delimiter, + store, + }) } } @@ -79,16 +105,6 @@ impl ObjectStore for EcObjectStore { source: "can not get object info".into(), })?; - // let stream = stream::unfold(reader.stream, |mut blob| async move { - // match blob.next().await { - // Some(Ok(chunk)) => { - // let bytes = chunk; - // Some((Ok(bytes), blob)) - // } - // _ => None, - // } - // }) - // .boxed(); let meta = ObjectMeta { location: location.clone(), last_modified: Utc::now(), @@ -98,10 +114,21 @@ impl ObjectStore for EcObjectStore { }; let attributes = Attributes::default(); - Ok(GetResult { - payload: object_store::GetResultPayload::Stream( + let payload = if self.need_convert { + object_store::GetResultPayload::Stream( + bytes_stream( + ReaderStream::with_capacity(ConvertStream::new(reader.stream, self.delimiter.clone()), READ_BUFFER_SIZE), + reader.object_info.size, + ) + .boxed(), + ) + } else { + object_store::GetResultPayload::Stream( bytes_stream(ReaderStream::with_capacity(reader.stream, READ_BUFFER_SIZE), reader.object_info.size).boxed(), - ), + ) + }; + Ok(GetResult { + payload, meta, range: 0..reader.object_info.size, attributes, @@ -154,6 +181,54 @@ impl ObjectStore for EcObjectStore { } } +pin_project! { + struct ConvertStream { + inner: R, + delimiter: Vec, + } +} + +impl ConvertStream { + fn new(inner: R, delimiter: String) -> Self { + ConvertStream { + inner, + delimiter: delimiter.as_bytes().to_vec(), + } + } +} + +impl AsyncRead for ConvertStream { + #[tracing::instrument(level = "debug", skip_all)] + fn poll_read( + self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &mut tokio::io::ReadBuf<'_>, + ) -> std::task::Poll> { + let me = self.project(); + ready!(Pin::new(&mut *me.inner).poll_read(cx, buf))?; + let bytes = buf.filled(); + let replaced = replace_symbol(me.delimiter, bytes); + buf.clear(); + buf.put_slice(&replaced); + Poll::Ready(Ok(())) + } +} + +fn replace_symbol(delimiter: &[u8], slice: &[u8]) -> Vec { + let mut result = Vec::with_capacity(slice.len()); + let mut i = 0; + while i < slice.len() { + if slice[i..].starts_with(delimiter) { + result.push(DEFAULT_DELIMITER); + i += delimiter.len(); + } else { + result.push(slice[i]); + i += 1; + } + } + result +} + pub fn bytes_stream(stream: S, content_length: usize) -> impl Stream> + Send + 'static where S: Stream> + Send + 'static, @@ -175,3 +250,21 @@ where Ok(()) }) } + +#[cfg(test)] +mod test { + use super::replace_symbol; + + #[test] + fn test_replace() { + let ss = String::from("dandan&&is&&best"); + let slice = ss.as_bytes(); + let delimiter = b"&&"; + println!("len: {}", "╦".len()); + let result = replace_symbol(delimiter, slice); + match String::from_utf8(result) { + Ok(s) => println!("slice: {}", s), + Err(e) => eprintln!("Error converting to string: {}", e), + } + } +} diff --git a/s3select/api/src/query/mod.rs b/s3select/api/src/query/mod.rs index 6ddd2dc86..93736be2b 100644 --- a/s3select/api/src/query/mod.rs +++ b/s3select/api/src/query/mod.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use s3s::dto::SelectObjectContentInput; pub mod analyzer; @@ -16,7 +18,7 @@ pub mod session; #[derive(Clone)] pub struct Context { // maybe we need transfer some info? - pub input: SelectObjectContentInput, + pub input: Arc, } #[derive(Clone)] diff --git a/s3select/api/src/query/session.rs b/s3select/api/src/query/session.rs index 286ee9f8e..581cdf39b 100644 --- a/s3select/api/src/query/session.rs +++ b/s3select/api/src/query/session.rs @@ -66,7 +66,7 @@ impl SessionCtxFactory { 9,Ivy,24,Marketing,4800 10,Jack,38,Finance,7500"; let data_bytes = data.as_bytes(); - // let data = r#""year"╦"gender"╦"ethnicity"╦"firstname"╦"count"╦"rank" + // let data = r#""year"╦"gender"╦"ethnicity"╦"firstname"╦"count"╦"rank" // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"SOPHIA"╦"119"╦"1" // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"CHLOE"╦"106"╦"2" // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"EMILY"╦"93"╦"3" diff --git a/s3select/query/src/dispatcher/manager.rs b/s3select/query/src/dispatcher/manager.rs index 80543ed69..ee5386e20 100644 --- a/s3select/query/src/dispatcher/manager.rs +++ b/s3select/query/src/dispatcher/manager.rs @@ -49,7 +49,7 @@ lazy_static! { #[derive(Clone)] pub struct SimpleQueryDispatcher { - input: SelectObjectContentInput, + input: Arc, // client for default tenant _default_table_provider: TableHandleProviderRef, session_factory: Arc, @@ -164,7 +164,9 @@ impl SimpleQueryDispatcher { .map(|e| e.as_bytes().first().copied().unwrap_or_default()), ); if let Some(delimiter) = csv.field_delimiter.as_ref() { - file_format = file_format.with_delimiter(delimiter.as_bytes().first().copied().unwrap_or_default()); + if delimiter.len() == 1 { + file_format = file_format.with_delimiter(delimiter.as_bytes()[0]); + } } // TODO waiting for processing @junxiang Mu // if csv.file_header_info.is_some() {} @@ -272,7 +274,7 @@ impl Stream for TrackedRecordBatchStream { #[derive(Default, Clone)] pub struct SimpleQueryDispatcherBuilder { - input: Option, + input: Option>, default_table_provider: Option, session_factory: Option>, parser: Option>, @@ -283,7 +285,7 @@ pub struct SimpleQueryDispatcherBuilder { } impl SimpleQueryDispatcherBuilder { - pub fn with_input(mut self, input: SelectObjectContentInput) -> Self { + pub fn with_input(mut self, input: Arc) -> Self { self.input = Some(input); self } diff --git a/s3select/query/src/instance.rs b/s3select/query/src/instance.rs index 44952a4ac..344920631 100644 --- a/s3select/query/src/instance.rs +++ b/s3select/query/src/instance.rs @@ -63,7 +63,7 @@ where } } -pub async fn make_rustfsms(input: SelectObjectContentInput, is_test: bool) -> QueryResult { +pub async fn make_rustfsms(input: Arc, is_test: bool) -> QueryResult { // init Function Manager, we can define some UDF if need let func_manager = SimpleFunctionMetadataManager::default(); // TODO session config need load global system config @@ -95,6 +95,8 @@ pub async fn make_rustfsms(input: SelectObjectContentInput, is_test: bool) -> Qu #[cfg(test)] mod tests { + use std::sync::Arc; + use api::{ query::{Context, Query}, server::dbms::DatabaseManagerSystem, @@ -111,7 +113,7 @@ mod tests { #[ignore] async fn test_simple_sql() { let sql = "select * from S3Object"; - let input = SelectObjectContentInput { + let input = Arc::new(SelectObjectContentInput { bucket: "dandan".to_string(), expected_bucket_owner: None, key: "test.csv".to_string(), @@ -135,7 +137,7 @@ mod tests { request_progress: None, scan_range: None, }, - }; + }); let db = make_rustfsms(input.clone(), true).await.unwrap(); let query = Query::new(Context { input }, sql.to_string()); @@ -167,8 +169,8 @@ mod tests { #[tokio::test] #[ignore] async fn test_func_sql() { - let sql = "SELECT s._1 FROM S3Object s"; - let input = SelectObjectContentInput { + let sql = "SELECT * FROM S3Object s"; + let input = Arc::new(SelectObjectContentInput { bucket: "dandan".to_string(), expected_bucket_owner: None, key: "test.csv".to_string(), @@ -194,7 +196,7 @@ mod tests { request_progress: None, scan_range: None, }, - }; + }); let db = make_rustfsms(input.clone(), true).await.unwrap(); let query = Query::new(Context { input }, sql.to_string());