support spec char as delimiter

Signed-off-by: junxiang Mu <1948535941@qq.com>

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<SelectObjectContentInput>,
+    need_convert: bool,
+    delimiter: String,
 
     store: Arc<ECStore>,
 }
-
 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 {
             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<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
 where
     S: Stream<Item = Result<Bytes, std::io::Error>> + 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),
+        }
+    }
+}

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<SelectObjectContentInput>,
 }
 
 #[derive(Clone)]

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"