fix

Author: zhyass

src/query/expression/src/utils/arrow.rs

@@ -94,7 +94,7 @@ pub fn deserialize_column(bytes: &[u8]) -> Result<Column> {
     read_column(&mut cursor)
 }
 
-pub fn read_column<R: Read + Seek>(r: &mut R) -> Result<Column> {
+fn read_column<R: Read + Seek>(r: &mut R) -> Result<Column> {
     let mut reader = FileReaderBuilder::new().build(r)?;
     let schema = reader.schema();
     let f = DataField::try_from(schema.field(0))?;

src/query/service/src/interpreters/interpreter_table_recluster.rs

@@ -319,7 +319,6 @@ impl ReclusterTableInterpreter {
         let rows_per_block =
             block_thresholds.calc_rows_per_block(total_bytes, total_rows, Some(block_per_seg));
         let total_partitions = std::cmp::max(total_rows / rows_per_block, 1);
-        settings.set_max_block_size(rows_per_block as u64)?;
         warn!("Do hilbert recluster, total_bytes: {}, total_rows: {}, total_partitions: {}, rows_per_block: {}",
             total_bytes, total_rows, total_partitions, rows_per_block);
 

src/query/service/src/spillers/serialize.rs

@@ -15,19 +15,24 @@
 use std::io::Write;
 use std::sync::Arc;
 
+use arrow_array::ArrayRef;
+use arrow_ipc::reader::FileReaderBuilder;
 use arrow_schema::Schema;
 use buf_list::BufList;
 use buf_list::Cursor;
 use bytes::Buf;
 use databend_common_base::base::Alignment;
 use databend_common_base::base::DmaWriteBuf;
+use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
-use databend_common_expression::arrow::read_column;
 use databend_common_expression::arrow::write_column;
 use databend_common_expression::infer_table_schema;
+use databend_common_expression::types::DataType;
+use databend_common_expression::BlockEntry;
 use databend_common_expression::DataBlock;
 use databend_common_expression::DataField;
 use databend_common_expression::DataSchema;
+use databend_common_expression::Value;
 use opendal::Buffer;
 use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
 use parquet::arrow::ArrowWriter;
@@ -101,21 +106,9 @@ impl BlocksEncoder {
     }
 }
 
-pub(super) fn deserialize_block(columns_layout: &Layout, mut data: Buffer) -> Result<DataBlock> {
+pub(super) fn deserialize_block(columns_layout: &Layout, data: Buffer) -> Result<DataBlock> {
     match columns_layout {
-        Layout::ArrowIpc(layout) => {
-            let columns = layout
-                .iter()
-                .map(|&layout| {
-                    let ls = BufList::from_iter(data.slice(0..layout));
-                    data.advance(layout);
-                    let mut cursor = Cursor::new(ls);
-                    read_column(&mut cursor)
-                })
-                .collect::<Result<Vec<_>, _>>()?;
-
-            Ok(DataBlock::new_from_columns(columns))
-        }
+        Layout::ArrowIpc(layout) => bare_blocks_from_arrow_ipc(layout, data),
         Layout::Parquet => bare_blocks_from_parquet(Reader(data)),
         Layout::Aggregate => unreachable!(),
     }
@@ -131,14 +124,51 @@ fn fake_data_schema(block: &DataBlock) -> DataSchema {
     DataSchema::new(fields)
 }
 
+fn bare_blocks_from_arrow_ipc(layout: &[usize], mut data: Buffer) -> Result<DataBlock> {
+    assert!(!layout.is_empty());
+    let mut columns = Vec::with_capacity(layout.len());
+    let mut read_array = |layout: usize| -> Result<(ArrayRef, DataType)> {
+        let ls = BufList::from_iter(data.slice(0..layout));
+        data.advance(layout);
+        let mut reader = FileReaderBuilder::new().build(Cursor::new(ls))?;
+        let schema = reader.schema();
+        let f = DataField::try_from(schema.field(0))?;
+        let data_type = f.data_type().clone();
+        let col = reader
+            .next()
+            .ok_or_else(|| ErrorCode::Internal("expected one arrow array"))??
+            .remove_column(0);
+        Ok((col, data_type))
+    };
+    let (array, data_type) = read_array(layout[0])?;
+    let num_rows = array.len();
+    let val = Value::from_arrow_rs(array, &data_type)?;
+    columns.push(BlockEntry::new(data_type, val));
+    for &layout in layout.iter().skip(1) {
+        let (array, data_type) = read_array(layout)?;
+        let val = Value::from_arrow_rs(array, &data_type)?;
+        columns.push(BlockEntry::new(data_type, val));
+    }
+    Ok(DataBlock::new(columns, num_rows))
+}
+
 /// Deserialize bare data block from parquet format.
 fn bare_blocks_from_parquet<R: ChunkReader + 'static>(data: R) -> Result<DataBlock> {
     let reader = ParquetRecordBatchReader::try_new(data, usize::MAX)?;
     let mut blocks = Vec::new();
     for record_batch in reader {
         let record_batch = record_batch?;
         let schema = DataSchema::try_from(record_batch.schema().as_ref())?;
-        let (block, _) = DataBlock::from_record_batch(&schema, &record_batch)?;
+        let num_rows = record_batch.num_rows();
+        let mut columns = Vec::with_capacity(record_batch.num_columns());
+        for (array, field) in record_batch.columns().iter().zip(schema.fields()) {
+            let data_type = field.data_type();
+            columns.push(BlockEntry::new(
+                data_type.clone(),
+                Value::from_arrow_rs(array.clone(), data_type)?,
+            ))
+        }
+        let block = DataBlock::new(columns, num_rows);
         blocks.push(block);
     }