Merge pull request #10731 from sundy-li/with_capacity_hint

Author: sundy-li

src/query/expression/src/types/string.rs

@@ -246,7 +246,7 @@ impl StringColumnBuilder {
         let mut offsets = Vec::with_capacity(len + 1);
         offsets.push(0);
         StringColumnBuilder {
-            need_estimated: data_capacity == 0,
+            need_estimated: data_capacity == 0 && len > 0,
             data: Vec::with_capacity(data_capacity),
             offsets,
         }
@@ -330,7 +330,7 @@ impl StringColumnBuilder {
             let bytes_per_row = self.data.len() / 64 + 1;
             let bytes_estimate = bytes_per_row * self.offsets.capacity();
 
-            const MAX_HINT_SIZE: usize = 1000000000;
+            const MAX_HINT_SIZE: usize = 1_000_000_000;
             // if we are more than 10% over the capacity, we reserve more
             if bytes_estimate < MAX_HINT_SIZE
                 && bytes_estimate as f64 > self.data.capacity() as f64 * 1.10f64

src/query/expression/src/values.rs

@@ -1927,6 +1927,16 @@ impl ColumnBuilder {
     }
 
     pub fn with_capacity(ty: &DataType, capacity: usize) -> ColumnBuilder {
+        ColumnBuilder::with_capacity_hint(ty, capacity, true)
+    }
+
+    /// Create a new column builder with capacity and enable_datasize_hint
+    /// enable_datasize_hint is used in StringColumnBuilder to decide whether to pre-allocate values
+    pub fn with_capacity_hint(
+        ty: &DataType,
+        capacity: usize,
+        enable_datasize_hint: bool,
+    ) -> ColumnBuilder {
         match ty {
             DataType::Null => ColumnBuilder::Null { len: 0 },
             DataType::EmptyArray => ColumnBuilder::EmptyArray { len: 0 },
@@ -1939,27 +1949,28 @@ impl ColumnBuilder {
             }
             DataType::Boolean => ColumnBuilder::Boolean(MutableBitmap::with_capacity(capacity)),
             DataType::String => {
-                ColumnBuilder::String(StringColumnBuilder::with_capacity(capacity, 0))
+                let data_capacity = if enable_datasize_hint { 0 } else { capacity };
+                ColumnBuilder::String(StringColumnBuilder::with_capacity(capacity, data_capacity))
             }
             DataType::Timestamp => ColumnBuilder::Timestamp(Vec::with_capacity(capacity)),
             DataType::Date => ColumnBuilder::Date(Vec::with_capacity(capacity)),
             DataType::Nullable(ty) => ColumnBuilder::Nullable(Box::new(NullableColumnBuilder {
-                builder: Self::with_capacity(ty, capacity),
+                builder: Self::with_capacity_hint(ty, capacity, enable_datasize_hint),
                 validity: MutableBitmap::with_capacity(capacity),
             })),
             DataType::Array(ty) => {
                 let mut offsets = Vec::with_capacity(capacity + 1);
                 offsets.push(0);
                 ColumnBuilder::Array(Box::new(ArrayColumnBuilder {
-                    builder: Self::with_capacity(ty, 0),
+                    builder: Self::with_capacity_hint(ty, 0, enable_datasize_hint),
                     offsets,
                 }))
             }
             DataType::Map(ty) => {
                 let mut offsets = Vec::with_capacity(capacity + 1);
                 offsets.push(0);
                 ColumnBuilder::Map(Box::new(ArrayColumnBuilder {
-                    builder: Self::with_capacity(ty, 0),
+                    builder: Self::with_capacity_hint(ty, 0, enable_datasize_hint),
                     offsets,
                 }))
             }
@@ -1968,12 +1979,15 @@ impl ColumnBuilder {
                 ColumnBuilder::Tuple(
                     fields
                         .iter()
-                        .map(|field| Self::with_capacity(field, capacity))
+                        .map(|field| {
+                            Self::with_capacity_hint(field, capacity, enable_datasize_hint)
+                        })
                         .collect(),
                 )
             }
             DataType::Variant => {
-                ColumnBuilder::Variant(StringColumnBuilder::with_capacity(capacity, 0))
+                let data_capacity = if enable_datasize_hint { 0 } else { capacity };
+                ColumnBuilder::Variant(StringColumnBuilder::with_capacity(capacity, data_capacity))
             }
             DataType::Generic(_) => {
                 unreachable!("unable to initialize column builder for generic type")

src/query/pipeline/sources/src/input_formats/input_format_text.rs

@@ -463,7 +463,11 @@ impl<T: InputFormatTextBase> BlockBuilder<T> {
             .mutable_columns
             .iter_mut()
             .map(|col| {
-                let empty_builder = ColumnBuilder::with_capacity(&col.data_type(), 0);
+                let empty_builder = ColumnBuilder::with_capacity_hint(
+                    &col.data_type(),
+                    self.ctx.block_compact_thresholds.min_rows_per_block,
+                    false,
+                );
                 std::mem::replace(col, empty_builder).build()
             })
             .collect();
@@ -504,9 +508,10 @@ impl<T: InputFormatTextBase> BlockBuilderTrait for BlockBuilder<T> {
             .fields()
             .iter()
             .map(|f| {
-                ColumnBuilder::with_capacity(
+                ColumnBuilder::with_capacity_hint(
                     &f.data_type().into(),
                     ctx.block_compact_thresholds.min_rows_per_block,
+                    false,
                 )
             })
             .collect();