support spec char as delimiter

Signed-off-by: junxiang Mu <1948535941@qq.com>
This commit is contained in:
junxiang Mu
2025-05-06 11:00:18 +08:00
parent fc47ca9dd2
commit 0ac1095c70
11 changed files with 139 additions and 39 deletions
-7
View File
@@ -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"
Generated
+2
View File
@@ -185,12 +185,14 @@ dependencies = [
"async-trait", "async-trait",
"bytes", "bytes",
"chrono", "chrono",
"common",
"datafusion", "datafusion",
"ecstore", "ecstore",
"futures", "futures",
"futures-core", "futures-core",
"http", "http",
"object_store", "object_store",
"pin-project-lite",
"s3s", "s3s",
"snafu", "snafu",
"tokio", "tokio",
+3
View File
@@ -2,6 +2,9 @@ pub mod error;
pub mod globals; pub mod globals;
pub mod last_minute; pub mod last_minute;
// is ','
pub static DEFAULT_DELIMITER: u8 = 44;
/// Defers evaluation of a block of code until the end of the scope. /// Defers evaluation of a block of code until the end of the scope.
#[macro_export] #[macro_export]
macro_rules! defer { macro_rules! defer {
+1 -1
View File
@@ -82,4 +82,4 @@ winapi = { workspace = true }
tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }
[build-dependencies] [build-dependencies]
shadow-rs.workspace = true shadow-rs = { workspace = true, features = ["build", "metadata"] }
+3 -2
View File
@@ -62,6 +62,7 @@ use s3s::S3;
use s3s::{S3Request, S3Response}; use s3s::{S3Request, S3Response};
use std::fmt::Debug; use std::fmt::Debug;
use std::str::FromStr; use std::str::FromStr;
use std::sync::Arc;
use tokio::sync::mpsc; use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream; use tokio_stream::wrappers::ReceiverStream;
use tokio_util::io::ReaderStream; use tokio_util::io::ReaderStream;
@@ -1896,14 +1897,14 @@ impl S3 for FS {
) -> S3Result<S3Response<SelectObjectContentOutput>> { ) -> S3Result<S3Response<SelectObjectContentOutput>> {
info!("handle select_object_content"); info!("handle select_object_content");
let input = req.input; let input = Arc::new(req.input);
info!("{:?}", input); info!("{:?}", input);
let db = make_rustfsms(input.clone(), false).await.map_err(|e| { let db = make_rustfsms(input.clone(), false).await.map_err(|e| {
error!("make db failed, {}", e.to_string()); error!("make db failed, {}", e.to_string());
s3_error!(InternalError, "{}", 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 let result = db
.execute(&query) .execute(&query)
.await .await
+2
View File
@@ -7,12 +7,14 @@ edition.workspace = true
async-trait.workspace = true async-trait.workspace = true
bytes.workspace = true bytes.workspace = true
chrono.workspace = true chrono.workspace = true
common.workspace = true
datafusion = { workspace = true } datafusion = { workspace = true }
ecstore.workspace = true ecstore.workspace = true
futures = { workspace = true } futures = { workspace = true }
futures-core = { workspace = true } futures-core = { workspace = true }
http.workspace = true http.workspace = true
object_store = { workspace = true } object_store = { workspace = true }
pin-project-lite.workspace = true
s3s.workspace = true s3s.workspace = true
snafu = { workspace = true, features = ["backtrace"] } snafu = { workspace = true, features = ["backtrace"] }
tokio.workspace = true tokio.workspace = true
+110 -17
View File
@@ -1,6 +1,7 @@
use async_trait::async_trait; use async_trait::async_trait;
use bytes::Bytes; use bytes::Bytes;
use chrono::Utc; use chrono::Utc;
use common::DEFAULT_DELIMITER;
use ecstore::io::READ_BUFFER_SIZE; use ecstore::io::READ_BUFFER_SIZE;
use ecstore::new_object_layer_fn; use ecstore::new_object_layer_fn;
use ecstore::store::ECStore; use ecstore::store::ECStore;
@@ -24,29 +25,54 @@ use object_store::PutOptions;
use object_store::PutPayload; use object_store::PutPayload;
use object_store::PutResult; use object_store::PutResult;
use object_store::{Error as o_Error, Result}; use object_store::{Error as o_Error, Result};
use pin_project_lite::pin_project;
use s3s::dto::SelectObjectContentInput; use s3s::dto::SelectObjectContentInput;
use s3s::s3_error; use s3s::s3_error;
use s3s::S3Result; use s3s::S3Result;
use std::ops::Range; use std::ops::Range;
use std::pin::Pin;
use std::sync::Arc; use std::sync::Arc;
use std::task::ready;
use std::task::Poll;
use tokio::io::AsyncRead;
use tokio_util::io::ReaderStream; use tokio_util::io::ReaderStream;
use tracing::info; use tracing::info;
use transform_stream::AsyncTryStream; use transform_stream::AsyncTryStream;
#[derive(Debug)] #[derive(Debug)]
pub struct EcObjectStore { pub struct EcObjectStore {
input: SelectObjectContentInput, input: Arc<SelectObjectContentInput>,
need_convert: bool,
delimiter: String,
store: Arc<ECStore>, store: Arc<ECStore>,
} }
impl EcObjectStore { impl EcObjectStore {
pub fn new(input: SelectObjectContentInput) -> S3Result<Self> { pub fn new(input: Arc<SelectObjectContentInput>) -> S3Result<Self> {
let Some(store) = new_object_layer_fn() else { let Some(store) = new_object_layer_fn() else {
return Err(s3_error!(InternalError, "ec store not inited")); 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(), 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 { let meta = ObjectMeta {
location: location.clone(), location: location.clone(),
last_modified: Utc::now(), last_modified: Utc::now(),
@@ -98,10 +114,21 @@ impl ObjectStore for EcObjectStore {
}; };
let attributes = Attributes::default(); let attributes = Attributes::default();
Ok(GetResult { let payload = if self.need_convert {
payload: object_store::GetResultPayload::Stream( 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(), bytes_stream(ReaderStream::with_capacity(reader.stream, READ_BUFFER_SIZE), reader.object_info.size).boxed(),
), )
};
Ok(GetResult {
payload,
meta, meta,
range: 0..reader.object_info.size, range: 0..reader.object_info.size,
attributes, attributes,
@@ -154,6 +181,54 @@ impl ObjectStore for EcObjectStore {
} }
} }
pin_project! {
struct ConvertStream<R> {
inner: R,
delimiter: Vec<u8>,
}
}
impl<R> ConvertStream<R> {
fn new(inner: R, delimiter: String) -> Self {
ConvertStream {
inner,
delimiter: delimiter.as_bytes().to_vec(),
}
}
}
impl<R: AsyncRead + Unpin> AsyncRead for ConvertStream<R> {
#[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<std::io::Result<()>> {
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<u8> {
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<S>(stream: S, content_length: usize) -> impl Stream<Item = Result<Bytes>> + Send + 'static pub fn bytes_stream<S>(stream: S, content_length: usize) -> impl Stream<Item = Result<Bytes>> + Send + 'static
where where
S: Stream<Item = Result<Bytes, std::io::Error>> + Send + 'static, S: Stream<Item = Result<Bytes, std::io::Error>> + Send + 'static,
@@ -175,3 +250,21 @@ where
Ok(()) 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),
}
}
}
+3 -1
View File
@@ -1,3 +1,5 @@
use std::sync::Arc;
use s3s::dto::SelectObjectContentInput; use s3s::dto::SelectObjectContentInput;
pub mod analyzer; pub mod analyzer;
@@ -16,7 +18,7 @@ pub mod session;
#[derive(Clone)] #[derive(Clone)]
pub struct Context { pub struct Context {
// maybe we need transfer some info? // maybe we need transfer some info?
pub input: SelectObjectContentInput, pub input: Arc<SelectObjectContentInput>,
} }
#[derive(Clone)] #[derive(Clone)]
+1 -1
View File
@@ -66,7 +66,7 @@ impl SessionCtxFactory {
9,Ivy,24,Marketing,4800 9,Ivy,24,Marketing,4800
10,Jack,38,Finance,7500"; 10,Jack,38,Finance,7500";
let data_bytes = data.as_bytes(); 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"╦"SOPHIA"╦"119"╦"1"
// "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"CHLOE"╦"106"╦"2" // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"CHLOE"╦"106"╦"2"
// "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"EMILY"╦"93"╦"3" // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"EMILY"╦"93"╦"3"
+6 -4
View File
@@ -49,7 +49,7 @@ lazy_static! {
#[derive(Clone)] #[derive(Clone)]
pub struct SimpleQueryDispatcher { pub struct SimpleQueryDispatcher {
input: SelectObjectContentInput, input: Arc<SelectObjectContentInput>,
// client for default tenant // client for default tenant
_default_table_provider: TableHandleProviderRef, _default_table_provider: TableHandleProviderRef,
session_factory: Arc<SessionCtxFactory>, session_factory: Arc<SessionCtxFactory>,
@@ -164,7 +164,9 @@ impl SimpleQueryDispatcher {
.map(|e| e.as_bytes().first().copied().unwrap_or_default()), .map(|e| e.as_bytes().first().copied().unwrap_or_default()),
); );
if let Some(delimiter) = csv.field_delimiter.as_ref() { 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 // TODO waiting for processing @junxiang Mu
// if csv.file_header_info.is_some() {} // if csv.file_header_info.is_some() {}
@@ -272,7 +274,7 @@ impl Stream for TrackedRecordBatchStream {
#[derive(Default, Clone)] #[derive(Default, Clone)]
pub struct SimpleQueryDispatcherBuilder { pub struct SimpleQueryDispatcherBuilder {
input: Option<SelectObjectContentInput>, input: Option<Arc<SelectObjectContentInput>>,
default_table_provider: Option<TableHandleProviderRef>, default_table_provider: Option<TableHandleProviderRef>,
session_factory: Option<Arc<SessionCtxFactory>>, session_factory: Option<Arc<SessionCtxFactory>>,
parser: Option<Arc<dyn Parser + Send + Sync>>, parser: Option<Arc<dyn Parser + Send + Sync>>,
@@ -283,7 +285,7 @@ pub struct SimpleQueryDispatcherBuilder {
} }
impl SimpleQueryDispatcherBuilder { impl SimpleQueryDispatcherBuilder {
pub fn with_input(mut self, input: SelectObjectContentInput) -> Self { pub fn with_input(mut self, input: Arc<SelectObjectContentInput>) -> Self {
self.input = Some(input); self.input = Some(input);
self self
} }
+8 -6
View File
@@ -63,7 +63,7 @@ where
} }
} }
pub async fn make_rustfsms(input: SelectObjectContentInput, is_test: bool) -> QueryResult<impl DatabaseManagerSystem> { pub async fn make_rustfsms(input: Arc<SelectObjectContentInput>, is_test: bool) -> QueryResult<impl DatabaseManagerSystem> {
// init Function Manager, we can define some UDF if need // init Function Manager, we can define some UDF if need
let func_manager = SimpleFunctionMetadataManager::default(); let func_manager = SimpleFunctionMetadataManager::default();
// TODO session config need load global system config // 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)] #[cfg(test)]
mod tests { mod tests {
use std::sync::Arc;
use api::{ use api::{
query::{Context, Query}, query::{Context, Query},
server::dbms::DatabaseManagerSystem, server::dbms::DatabaseManagerSystem,
@@ -111,7 +113,7 @@ mod tests {
#[ignore] #[ignore]
async fn test_simple_sql() { async fn test_simple_sql() {
let sql = "select * from S3Object"; let sql = "select * from S3Object";
let input = SelectObjectContentInput { let input = Arc::new(SelectObjectContentInput {
bucket: "dandan".to_string(), bucket: "dandan".to_string(),
expected_bucket_owner: None, expected_bucket_owner: None,
key: "test.csv".to_string(), key: "test.csv".to_string(),
@@ -135,7 +137,7 @@ mod tests {
request_progress: None, request_progress: None,
scan_range: None, scan_range: None,
}, },
}; });
let db = make_rustfsms(input.clone(), true).await.unwrap(); let db = make_rustfsms(input.clone(), true).await.unwrap();
let query = Query::new(Context { input }, sql.to_string()); let query = Query::new(Context { input }, sql.to_string());
@@ -167,8 +169,8 @@ mod tests {
#[tokio::test] #[tokio::test]
#[ignore] #[ignore]
async fn test_func_sql() { async fn test_func_sql() {
let sql = "SELECT s._1 FROM S3Object s"; let sql = "SELECT * FROM S3Object s";
let input = SelectObjectContentInput { let input = Arc::new(SelectObjectContentInput {
bucket: "dandan".to_string(), bucket: "dandan".to_string(),
expected_bucket_owner: None, expected_bucket_owner: None,
key: "test.csv".to_string(), key: "test.csv".to_string(),
@@ -194,7 +196,7 @@ mod tests {
request_progress: None, request_progress: None,
scan_range: None, scan_range: None,
}, },
}; });
let db = make_rustfsms(input.clone(), true).await.unwrap(); let db = make_rustfsms(input.clone(), true).await.unwrap();
let query = Query::new(Context { input }, sql.to_string()); let query = Query::new(Context { input }, sql.to_string());