feat: add fast paths for substr and string column concat (#19628)

* perf: add fast paths for substr and string column concat

* test(functions): cover substr fast paths

* fix(expression): compact sparse string buffers in concat

Author: KKould

src/query/expression/src/kernels/concat.rs

@@ -49,6 +49,8 @@ use crate::types::map::KvColumnBuilder;
 use crate::types::nullable::NullableColumn;
 use crate::types::number::NumberColumn;
 use crate::types::opaque::OpaqueType;
+use crate::types::string::StringColumn;
+use crate::types::string::StringColumnBuilder;
 use crate::types::timestamp_tz::TimestampTzType;
 use crate::types::vector::VectorColumnBuilder;
 use crate::with_decimal_mapped_type;
@@ -242,11 +244,14 @@ impl Column {
                     Column::Vector(builder.build())
                 }
             }),
+            Column::String(_) => Column::String(Self::concat_string_types(
+                columns.map(|col| col.into_string().unwrap()),
+                capacity,
+            )),
             Column::Variant(_)
             | Column::Geometry(_)
             | Column::Geography(_)
             | Column::Binary(_)
-            | Column::String(_)
             | Column::Bitmap(_) => {
                 Self::concat_use_arrow(columns, first_column.data_type(), capacity)
             }
@@ -268,6 +273,17 @@ impl Column {
         builder.into()
     }
 
+    pub fn concat_string_types(
+        cols: impl Iterator<Item = StringColumn>,
+        num_rows: usize,
+    ) -> StringColumn {
+        let mut builder = StringColumnBuilder::with_capacity(num_rows);
+        for col in cols {
+            builder.append_column_for_concat(&col);
+        }
+        builder.build()
+    }
+
     fn concat_opaque_column<I, const N: usize>(columns: I, capacity: usize) -> Column
     where I: Iterator<Item = Column> + TrustedLen + Clone {
         let buffer = Self::concat_primitive_types(
@@ -307,3 +323,53 @@ impl Column {
         T::upcast_column_with_type(T::build_column(builder), data_type)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::Column;
+    use crate::types::string::StringColumn;
+
+    #[test]
+    fn test_concat_string_columns_fast_path() {
+        let left = Column::String(StringColumn::from_iter([
+            "existing short",
+            "existing long string that definitely exceeds inline storage!",
+        ]));
+        let right = Column::String(StringColumn::from_iter([
+            "append short",
+            "append another extremely long string to force buffer usage!!!",
+        ]));
+
+        let result = Column::concat_columns(vec![left, right].into_iter()).unwrap();
+        let result = result.into_string().unwrap();
+        let values = result.iter().collect::<Vec<_>>();
+
+        assert_eq!(values, vec![
+            "existing short",
+            "existing long string that definitely exceeds inline storage!",
+            "append short",
+            "append another extremely long string to force buffer usage!!!",
+        ]);
+    }
+
+    #[test]
+    fn test_concat_string_columns_compacts_sparse_buffers() {
+        let long = "x".repeat(20_000);
+        let source = StringColumn::from_iter((0..8).map(|idx| format!("{idx}-{long}")));
+        let source_total_buffer_len = source.total_buffer_len();
+        let sparse_left = Column::String(source.clone().sliced(0, 1));
+        let sparse_right = Column::String(source.clone().sliced(7, 1));
+
+        let result = Column::concat_columns(vec![sparse_left, sparse_right].into_iter()).unwrap();
+        let result = result.into_string().unwrap();
+
+        assert_eq!(result.iter().collect::<Vec<_>>(), vec![
+            format!("0-{long}"),
+            format!("7-{long}"),
+        ]);
+        assert!(
+            result.total_buffer_len() < source_total_buffer_len,
+            "concat should compact sparse string buffers instead of retaining them"
+        );
+    }
+}

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

@@ -342,6 +342,21 @@ impl StringColumnBuilder {
         self.data.extend_values(other.iter());
     }
 
+    pub fn append_column_for_concat(&mut self, other: &StringColumn) {
+        debug_assert!(
+            self.row_buffer.is_empty(),
+            "append_column_for_concat expects no pending row data"
+        );
+
+        let source = other.clone().maybe_gc();
+        unsafe {
+            self.data.append_views_unchecked(
+                source.views().as_slice().iter(),
+                source.data_buffers().as_ref(),
+            );
+        }
+    }
+
     pub fn build(self) -> StringColumn {
         self.data.into()
     }
@@ -397,6 +412,29 @@ impl StringColumnBuilder {
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::StringColumn;
+    use super::StringColumnBuilder;
+
+    #[test]
+    fn test_append_column_copies_sparse_string_buffers() {
+        let long = "x".repeat(20_000);
+        let source = StringColumn::from_iter((0..8).map(|idx| format!("{idx}-{long}")));
+        let sparse = source.sliced(7, 1);
+
+        let mut builder = StringColumnBuilder::with_capacity(1);
+        builder.append_column(&sparse);
+        let result = builder.build();
+
+        assert_eq!(result.iter().collect::<Vec<_>>(), vec![format!("7-{long}")]);
+        assert!(
+            result.total_buffer_len() < sparse.total_buffer_len(),
+            "append_column should not retain the full source buffers"
+        );
+    }
+}
+
 impl NullableColumnBuilder<StringType> {
     pub fn take_from_views(
         views: &[ColumnView<NullableType<StringType>>],

src/query/functions/src/scalars/string.rs

@@ -1056,6 +1056,16 @@ fn substr(builder: &mut StringColumnBuilder, str: &str, pos: i64, len: u64) {
         return;
     }
 
+    if str.is_ascii() {
+        substr_ascii(builder, str, pos, len);
+        return;
+    }
+
+    if len == 1 && pos > 0 {
+        substr_one_char(builder, str, pos);
+        return;
+    }
+
     let char_len = str.chars().count();
     let start = if pos > 0 {
         (pos - 1).min(char_len as i64) as usize
@@ -1068,3 +1078,81 @@ fn substr(builder: &mut StringColumnBuilder, str: &str, pos: i64, len: u64) {
     builder.put_char_iter(str.chars().skip(start).take(len as usize));
     builder.commit_row();
 }
+
+#[inline]
+fn substr_ascii(builder: &mut StringColumnBuilder, str: &str, pos: i64, len: u64) {
+    let byte_len = str.len();
+    let start = if pos > 0 {
+        (pos - 1).min(byte_len as i64) as usize
+    } else {
+        byte_len
+            .checked_sub(pos.unsigned_abs() as usize)
+            .unwrap_or(byte_len)
+    };
+
+    let end = start.saturating_add(len as usize).min(byte_len);
+    builder.put_slice(&str.as_bytes()[start..end]);
+    builder.commit_row();
+}
+
+#[inline]
+fn substr_one_char(builder: &mut StringColumnBuilder, str: &str, pos: i64) {
+    let target = (pos - 1) as usize;
+    let bytes = str.as_bytes();
+    let mut byte_idx = 0;
+    let mut char_idx = 0;
+
+    while byte_idx < bytes.len() {
+        if char_idx == target {
+            let start = byte_idx;
+            byte_idx += 1;
+            while byte_idx < bytes.len() && is_utf8_continuation_byte(bytes[byte_idx]) {
+                byte_idx += 1;
+            }
+            builder.put_slice(&bytes[start..byte_idx]);
+            builder.commit_row();
+            return;
+        }
+
+        byte_idx += 1;
+        while byte_idx < bytes.len() && is_utf8_continuation_byte(bytes[byte_idx]) {
+            byte_idx += 1;
+        }
+        char_idx += 1;
+    }
+
+    builder.commit_row();
+}
+
+#[inline]
+fn is_utf8_continuation_byte(byte: u8) -> bool {
+    (byte & 0b1100_0000) == 0b1000_0000
+}
+
+#[cfg(test)]
+mod tests {
+    use databend_common_expression::types::string::StringColumnBuilder;
+
+    use super::substr;
+
+    fn eval_substr(input: &str, pos: i64, len: u64) -> String {
+        let mut builder = StringColumnBuilder::with_capacity(1);
+        substr(&mut builder, input, pos, len);
+        builder.build_scalar()
+    }
+
+    #[test]
+    fn test_substr_ascii_fast_path() {
+        assert_eq!(eval_substr("abcdef", 2, 3), "bcd");
+        assert_eq!(eval_substr("abcdef", -2, 2), "ef");
+        assert_eq!(eval_substr("abcdef", 20, 1), "");
+    }
+
+    #[test]
+    fn test_substr_single_char_fast_path() {
+        assert_eq!(eval_substr("abcdef", 3, 1), "c");
+        assert_eq!(eval_substr("你好世界", 3, 1), "世");
+        assert_eq!(eval_substr("a你b", 2, 1), "你");
+        assert_eq!(eval_substr("こんにちは", 2, 1), "ん");
+    }
+}

src/query/functions/tests/it/scalars/string.rs

@@ -736,8 +736,14 @@ fn test_substr(file: &mut impl Write) {
     run_ast(file, "substr('Sakila' from -4 for 2)", &[]);
     run_ast(file, "substr('sakila' FROM -4)", &[]);
     run_ast(file, "substr('abc',2)", &[]);
+    run_ast(file, "substr('abcdef', 2, 3)", &[]);
+    run_ast(file, "substr('abcdef', -2, 2)", &[]);
+    run_ast(file, "substr('abcdef', 20, 1)", &[]);
     run_ast(file, "substr('你好世界', 3)", &[]);
+    run_ast(file, "substr('你好世界', 3, 1)", &[]);
+    run_ast(file, "substr('a你b', 2, 1)", &[]);
     run_ast(file, "substr('こんにちは', 2)", &[]);
+    run_ast(file, "substr('こんにちは', 2, 1)", &[]);
     run_ast(file, "substr('abc', pos, len)", &[
         (
             "pos",

src/query/functions/tests/it/scalars/testdata/string.txt

@@ -3594,6 +3594,33 @@ output domain  : {"bc"..="bc"}
 output         : 'bc'
 
 
+ast            : substr('abcdef', 2, 3)
+raw expr       : substr('abcdef', 2, 3)
+checked expr   : substr<String, Int64, UInt64>("abcdef", CAST<UInt8>(2_u8 AS Int64), CAST<UInt8>(3_u8 AS UInt64))
+optimized expr : "bcd"
+output type    : String
+output domain  : {"bcd"..="bcd"}
+output         : 'bcd'
+
+
+ast            : substr('abcdef', -2, 2)
+raw expr       : substr('abcdef', -2, 2)
+checked expr   : substr<String, Int64, UInt64>("abcdef", CAST<Int8>(-2_i8 AS Int64), CAST<UInt8>(2_u8 AS UInt64))
+optimized expr : "ef"
+output type    : String
+output domain  : {"ef"..="ef"}
+output         : 'ef'
+
+
+ast            : substr('abcdef', 20, 1)
+raw expr       : substr('abcdef', 20, 1)
+checked expr   : substr<String, Int64, UInt64>("abcdef", CAST<UInt8>(20_u8 AS Int64), CAST<UInt8>(1_u8 AS UInt64))
+optimized expr : ""
+output type    : String
+output domain  : {""..=""}
+output         : ''
+
+
 ast            : substr('你好世界', 3)
 raw expr       : substr('你好世界', 3)
 checked expr   : substr<String, Int64>("你好世界", CAST<UInt8>(3_u8 AS Int64))
@@ -3603,6 +3630,24 @@ output domain  : {"世界"..="世界"}
 output         : '世界'
 
 
+ast            : substr('你好世界', 3, 1)
+raw expr       : substr('你好世界', 3, 1)
+checked expr   : substr<String, Int64, UInt64>("你好世界", CAST<UInt8>(3_u8 AS Int64), CAST<UInt8>(1_u8 AS UInt64))
+optimized expr : "世"
+output type    : String
+output domain  : {"世"..="世"}
+output         : '世'
+
+
+ast            : substr('a你b', 2, 1)
+raw expr       : substr('a你b', 2, 1)
+checked expr   : substr<String, Int64, UInt64>("a你b", CAST<UInt8>(2_u8 AS Int64), CAST<UInt8>(1_u8 AS UInt64))
+optimized expr : "你"
+output type    : String
+output domain  : {"你"..="你"}
+output         : '你'
+
+
 ast            : substr('こんにちは', 2)
 raw expr       : substr('こんにちは', 2)
 checked expr   : substr<String, Int64>("こんにちは", CAST<UInt8>(2_u8 AS Int64))
@@ -3612,6 +3657,15 @@ output domain  : {"んにちは"..="んにちは"}
 output         : 'んにちは'
 
 
+ast            : substr('こんにちは', 2, 1)
+raw expr       : substr('こんにちは', 2, 1)
+checked expr   : substr<String, Int64, UInt64>("こんにちは", CAST<UInt8>(2_u8 AS Int64), CAST<UInt8>(1_u8 AS UInt64))
+optimized expr : "ん"
+output type    : String
+output domain  : {"ん"..="ん"}
+output         : 'ん'
+
+
 ast            : substr('abc', pos, len)
 raw expr       : substr('abc', pos::Int8, len::UInt8)
 checked expr   : substr<String, Int64, UInt64>("abc", CAST<Int8>(pos AS Int64), CAST<UInt8>(len AS UInt64))