fix: copy from CSV OOM when file is large.

Author: youngsofun

src/query/storages/stage/Cargo.toml

@@ -34,7 +34,6 @@ databend-common-storages-parquet = { workspace = true }
 databend-storages-common-stage = { workspace = true }
 databend-storages-common-table-meta = { workspace = true }
 enum-as-inner = { workspace = true }
-futures = { workspace = true }
 jsonb = { workspace = true }
 lexical-core = { workspace = true }
 log = { workspace = true }

src/query/storages/stage/src/read/row_based/processors/reader.rs

@@ -23,24 +23,21 @@ use databend_common_exception::Result;
 use databend_common_expression::DataBlock;
 use databend_common_pipeline_sources::PrefetchAsyncSource;
 use databend_storages_common_stage::SingleFilePartition;
-use futures::AsyncRead;
-use futures::AsyncReadExt;
 use log::debug;
 use opendal::Operator;
 
 use crate::read::row_based::batch::BytesBatch;
 
 struct FileState {
     file: SingleFilePartition,
-    reader: opendal::FuturesAsyncReader,
+    reader: opendal::Reader,
     offset: usize,
 }
 
 pub struct BytesReader {
     table_ctx: Arc<dyn TableContext>,
     op: Operator,
     read_batch_size: usize,
-    io_size: usize,
     file_state: Option<FileState>,
     prefetch_num: usize,
 }
@@ -52,36 +49,26 @@ impl BytesReader {
         read_batch_size: usize,
         prefetch_num: usize,
     ) -> Result<Self> {
-        // TODO: Use 8MiB as default IO size for now, we can extract as a new config.
-        let default_io_size: usize = 8 * 1024 * 1024;
-        // Calculate the IO size, which:
-        //
-        // - is the multiple of read_batch_size.
-        // - is larger or equal to default_io_size.
-        let io_size = default_io_size.div_ceil(read_batch_size) * read_batch_size;
-
         Ok(Self {
             table_ctx,
             op,
             read_batch_size,
-            io_size,
             file_state: None,
             prefetch_num,
         })
     }
 
     pub async fn read_batch(&mut self) -> Result<DataBlock> {
         if let Some(state) = &mut self.file_state {
-            let end = state.file.size.min(self.read_batch_size + state.offset);
-            let mut buffer = vec![0u8; end - state.offset];
-            let n = read_full(&mut state.reader, &mut buffer[..]).await?;
+            let end = state.file.size.min(self.read_batch_size + state.offset) as u64;
+            let buffer = state.reader.read(state.offset as u64..end).await?.to_vec();
+            let n = buffer.len();
             if n == 0 {
                 return Err(ErrorCode::BadBytes(format!(
                     "Unexpected EOF {} expect {} bytes, read only {} bytes.",
                     state.file.path, state.file.size, state.offset
                 )));
             };
-            buffer.truncate(n);
 
             Profile::record_usize_profile(ProfileStatisticsName::ScanBytes, n);
             self.table_ctx
@@ -129,15 +116,8 @@ impl PrefetchAsyncSource for BytesReader {
             };
             let file = SingleFilePartition::from_part(&part)?.clone();
 
-            let reader = self
-                .op
-                .reader_with(&file.path)
-                .chunk(self.io_size)
-                // TODO: Use 4 concurrent for test, let's extract as a new setting.
-                .concurrent(4)
-                .await?
-                .into_futures_async_read(0..file.size as u64)
-                .await?;
+            let reader = self.op.reader(&file.path).await?;
+
             self.file_state = Some(FileState {
                 file,
                 reader,
@@ -150,18 +130,3 @@ impl PrefetchAsyncSource for BytesReader {
         }
     }
 }
-
-#[async_backtrace::framed]
-pub async fn read_full<R: AsyncRead + Unpin>(reader: &mut R, buf: &mut [u8]) -> Result<usize> {
-    let mut buf = &mut buf[0..];
-    let mut n = 0;
-    while !buf.is_empty() {
-        let read = reader.read(buf).await?;
-        if read == 0 {
-            break;
-        }
-        n += read;
-        buf = &mut buf[read..]
-    }
-    Ok(n)
-}

src/query/storages/stage/src/read/row_based/read_pipeline.rs

@@ -55,7 +55,7 @@ impl RowBasedReadPipelineBuilder<'_> {
         let batch_size = settings.get_input_read_buffer_size()? as usize;
         pipeline.add_source(
             |output| {
-                let reader = BytesReader::try_create(ctx.clone(), operator.clone(), batch_size, 1)?;
+                let reader = BytesReader::try_create(ctx.clone(), operator.clone(), batch_size, 2)?;
                 PrefetchAsyncSourcer::create(ctx.clone(), output, reader)
             },
             num_threads,