lazy columnar merge

This is the first part of addressing #3633 Instead of loading all Column into memory for the merge, only the current column_name group is loaded. This can be done since the sstable streams the columns lexicographically.
quickwit-oss · Jul 21, 2023 · 1bfbfda · 1bfbfda
1 parent 820f126
commit 1bfbfda
Show file tree

Hide file tree

Showing 6 changed files with 255 additions and 112 deletions.
diff --git a/columnar/src/columnar/merge/mod.rs b/columnar/src/columnar/merge/mod.rs
@@ -2,12 +2,14 @@ mod merge_dict_column;
 mod merge_mapping;
 mod term_merger;
 
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{HashMap, HashSet};
 use std::io;
 use std::net::Ipv6Addr;
+use std::rc::Rc;
 use std::sync::Arc;
 
-use itertools::Itertools;
+use common::GroupByIteratorExtended;
+use itertools::{EitherOrBoth, Itertools};
 pub use merge_mapping::{MergeRowOrder, ShuffleMergeOrder, StackMergeOrder};
 
 use super::writer::ColumnarSerializer;
@@ -18,7 +20,8 @@ use crate::columnar::writer::CompatibleNumericalTypes;
 use crate::columnar::ColumnarReader;
 use crate::dynamic_column::DynamicColumn;
 use crate::{
-    BytesColumn, Column, ColumnIndex, ColumnType, ColumnValues, NumericalType, NumericalValue,
+    BytesColumn, Column, ColumnIndex, ColumnType, ColumnValues, DynamicColumnHandle, NumericalType,
+    NumericalValue,
 };
 
 /// Column types are grouped into different categories.
@@ -28,7 +31,7 @@ use crate::{
 /// In practise, today, only Numerical colummns are coerced into one type today.
 ///
 /// See also [README.md].
-#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash, Debug)]
 pub(crate) enum ColumnTypeCategory {
     Bool,
     Str,
@@ -83,9 +86,13 @@ pub fn merge_columnar(
         .iter()
         .map(|reader| reader.num_rows())
         .collect::<Vec<u32>>();
-    let columns_to_merge =
-        group_columns_for_merge(columnar_readers, required_columns, &merge_row_order)?;
-    for ((column_name, column_type), columns) in columns_to_merge {
+
+    let columns_to_merge_iter =
+        group_columns_for_merge_iter(columnar_readers, required_columns, &merge_row_order)?;
+    for res in columns_to_merge_iter {
+        let (column_name, column_type, grouped_columns) = res?;
+        let columns = grouped_columns.columns;
+
         let mut column_serializer =
             serializer.start_serialize_column(column_name.as_bytes(), column_type);
         merge_column(
@@ -97,6 +104,7 @@ pub fn merge_columnar(
         )?;
         column_serializer.finalize()?;
     }
+
     serializer.finalize(merge_row_order.num_rows())?;
     Ok(())
 }
@@ -214,11 +222,11 @@ struct GroupedColumns {
 }
 
 impl GroupedColumns {
-    fn for_category(column_category: ColumnTypeCategory, num_columnars: usize) -> Self {
+    fn new(num_columnars: usize) -> Self {
         GroupedColumns {
             required_column_type: None,
             columns: vec![None; num_columnars],
-            column_category,
+            column_category: ColumnTypeCategory::Numerical,
         }
     }
 
@@ -293,7 +301,7 @@ fn merged_numerical_columns_type<'a>(
 fn is_empty_after_merge(
     merge_row_order: &MergeRowOrder,
     column: &DynamicColumn,
-    columnar_id: usize,
+    columnar_ord: usize,
 ) -> bool {
     if column.num_values() == 0u32 {
         // It was empty before the merge.
@@ -305,7 +313,7 @@ fn is_empty_after_merge(
             false
         }
         MergeRowOrder::Shuffled(shuffled) => {
-            if let Some(alive_bitset) = &shuffled.alive_bitsets[columnar_id] {
+            if let Some(alive_bitset) = &shuffled.alive_bitsets[columnar_ord] {
                 let column_index = column.column_index();
                 match column_index {
                     ColumnIndex::Empty { .. } => true,
@@ -348,56 +356,115 @@ fn is_empty_after_merge(
     }
 }
 
-#[allow(clippy::type_complexity)]
-fn group_columns_for_merge(
-    columnar_readers: &[&ColumnarReader],
-    required_columns: &[(String, ColumnType)],
-    merge_row_order: &MergeRowOrder,
-) -> io::Result<BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>>> {
-    // Each column name may have multiple types of column associated.
-    // For merging we are interested in the same column type category since they can be merged.
-    let mut columns_grouped: HashMap<(String, ColumnTypeCategory), GroupedColumns> = HashMap::new();
+type MergeIter<'a> =
+    Box<dyn Iterator<Item = io::Result<(Rc<String>, ColumnType, GroupedColumns)>> + 'a>;
 
-    for &(ref column_name, column_type) in required_columns {
-        columns_grouped
-            .entry((column_name.clone(), column_type.into()))
-            .or_insert_with(|| {
-                GroupedColumns::for_category(column_type.into(), columnar_readers.len())
-            })
-            .require_type(column_type)?;
-    }
+/// Iterates over the columns of the columnar readers, grouped by column name.
+/// Key functionality is that `open` of the Columns is done lazy per group.
+fn group_columns_for_merge_iter<'a>(
+    columnar_readers: &'a [&'a ColumnarReader],
+    required_columns: &'a [(String, ColumnType)],
+    merge_row_order: &'a MergeRowOrder,
+) -> io::Result<impl Iterator<Item = io::Result<(Rc<String>, ColumnType, GroupedColumns)>> + 'a> {
+    let column_iters: Vec<_> = columnar_readers
+        .iter()
+        .enumerate()
+        .map(|(reader_ord, reader)| {
+            Ok(reader
+                .iter_columns()?
+                .map(move |el| (Rc::new(el.0), reader_ord, el.1)))
+        })
+        .collect::<io::Result<_>>()?;
+    let required_columns_map: HashMap<String, _> = required_columns
+        .iter()
+        .map(|(col_name, typ)| (col_name.to_string(), typ))
+        .collect::<HashMap<String, _>>();
+    let mut required_columns_list: Vec<String> = required_columns
+        .iter()
+        .map(|(col_name, _)| col_name.to_string())
+        .collect();
+    required_columns_list.sort();
 
-    for (columnar_id, columnar_reader) in columnar_readers.iter().enumerate() {
-        let column_name_and_handle = columnar_reader.list_columns()?;
-        // We skip columns that end up with 0 documents.
-        // That way, we make sure they don't end up influencing the merge type or
-        // creating empty columns.
+    // Kmerge and group on the column_name.
+    let group_iter = GroupByIteratorExtended::group_by(
+        column_iters.into_iter().kmerge_by(|a, b| a.0 < b.0),
+        |el| el.0.clone(),
+    );
 
-        for (column_name, handle) in column_name_and_handle {
-            let column_category: ColumnTypeCategory = handle.column_type().into();
-            let column = handle.open()?;
-            if is_empty_after_merge(merge_row_order, &column, columnar_id) {
-                continue;
-            }
-            columns_grouped
-                .entry((column_name, column_category))
-                .or_insert_with(|| {
-                    GroupedColumns::for_category(column_category, columnar_readers.len())
-                })
-                .set_column(columnar_id, column);
-        }
-    }
+    // Weave in the required columns into the sorted by column name iterator.
+    let groups_with_required = required_columns_list
+        .into_iter()
+        .merge_join_by(group_iter, |a, b| a.cmp(&b.0));
 
-    let mut merge_columns: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
-        Default::default();
+    Ok(groups_with_required.flat_map(move |either| {
+        // It should be possible to do the grouping also on the column type in one pass, but some
+        // tests are failing.
+        let mut force_type: Option<ColumnType> = None;
+        let (key, group) = match either {
+            // set required column
+            EitherOrBoth::Both(_required, (key, group)) => {
+                force_type = required_columns_map.get(&*key).map(|el| (**el).into());
+                (key, group)
+            }
+            // Only required - Return artificial empty column
+            EitherOrBoth::Left(key) => {
+                let mut grouped_columns = GroupedColumns::new(columnar_readers.len());
+                let force_type: Option<ColumnType> =
+                    required_columns_map.get(&*key).map(|el| (**el).into());
+                if let Some(force_type) = force_type {
+                    grouped_columns.require_type(force_type).unwrap(); // Can't panic
+                }
+                return Box::new(std::iter::once(Ok((
+                    Rc::new(key),
+                    force_type.unwrap(),
+                    grouped_columns,
+                )))) as MergeIter<'a>;
+            }
+            // no required column
+            EitherOrBoth::Right((key, group)) => (key, group),
+        };
+        let mut group: Vec<(Rc<String>, usize, DynamicColumnHandle)> = group.collect();
+        group.sort_by_key(|el| el.2.column_type);
+        let group_iter = GroupByIteratorExtended::group_by(group.into_iter(), |el| {
+            let cat_type: ColumnTypeCategory = el.2.column_type().into();
+            cat_type
+        });
+        let key = key.clone();
+        Box::new(
+            group_iter
+                .map(move |(_cat, group)| {
+                    let mut grouped_columns = GroupedColumns::new(columnar_readers.len());
+                    if let Some(force_type) = force_type {
+                        grouped_columns.require_type(force_type)?;
+                    }
+                    for col in group {
+                        let columnar_ord = col.1;
+                        let column = col.2.open()?;
+                        if !is_empty_after_merge(merge_row_order, &column, columnar_ord) {
+                            grouped_columns.set_column(col.1, column);
+                        }
+                    }
 
-    for ((column_name, _), mut grouped_columns) in columns_grouped {
-        let column_type = grouped_columns.column_type_after_merge();
-        coerce_columns(column_type, &mut grouped_columns.columns)?;
-        merge_columns.insert((column_name, column_type), grouped_columns.columns);
-    }
+                    let column_type = grouped_columns.column_type_after_merge();
+                    coerce_columns(column_type, &mut grouped_columns.columns)?;
 
-    Ok(merge_columns)
+                    Ok((key.clone(), column_type, grouped_columns))
+                })
+                .filter(|res| {
+                    // Filter out empty columns.
+                    if let Ok((_, _, grouped_columns)) = res {
+                        if grouped_columns
+                            .columns
+                            .iter()
+                            .all(|column| column.is_none())
+                        {
+                            return false;
+                        }
+                    }
+                    true
+                }),
+        )
+    }))
 }
 
 fn coerce_columns(

diff --git a/columnar/src/columnar/merge/tests.rs b/columnar/src/columnar/merge/tests.rs
@@ -1,3 +1,5 @@
+use std::collections::BTreeMap;
+
 use itertools::Itertools;
 
 use super::*;
@@ -28,7 +30,10 @@ fn test_column_coercion_to_u64() {
     let columnars = &[&columnar1, &columnar2];
     let merge_order = StackMergeOrder::stack(columnars).into();
     let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
-        group_columns_for_merge(columnars, &[], &merge_order).unwrap();
+        group_columns_for_merge_iter(columnars, &[], &merge_order)
+            .unwrap()
+            .map(conv_res)
+            .collect();
     assert_eq!(column_map.len(), 1);
     assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::U64)));
 }
@@ -40,7 +45,10 @@ fn test_column_no_coercion_if_all_the_same() {
     let columnars = &[&columnar1, &columnar2];
     let merge_order = StackMergeOrder::stack(columnars).into();
     let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
-        group_columns_for_merge(columnars, &[], &merge_order).unwrap();
+        group_columns_for_merge_iter(columnars, &[], &merge_order)
+            .unwrap()
+            .map(conv_res)
+            .collect();
     assert_eq!(column_map.len(), 1);
     assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::U64)));
 }
@@ -52,23 +60,26 @@ fn test_column_coercion_to_i64() {
     let columnars = &[&columnar1, &columnar2];
     let merge_order = StackMergeOrder::stack(columnars).into();
     let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
-        group_columns_for_merge(columnars, &[], &merge_order).unwrap();
+        group_columns_for_merge_iter(columnars, &[], &merge_order)
+            .unwrap()
+            .map(conv_res)
+            .collect();
     assert_eq!(column_map.len(), 1);
     assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::I64)));
 }
 
-#[test]
-fn test_impossible_coercion_returns_an_error() {
-    let columnar1 = make_columnar("numbers", &[u64::MAX]);
-    let merge_order = StackMergeOrder::stack(&[&columnar1]).into();
-    let group_error = group_columns_for_merge(
-        &[&columnar1],
-        &[("numbers".to_string(), ColumnType::I64)],
-        &merge_order,
-    )
-    .unwrap_err();
-    assert_eq!(group_error.kind(), io::ErrorKind::InvalidInput);
-}
+//#[test]
+// fn test_impossible_coercion_returns_an_error() {
+// let columnar1 = make_columnar("numbers", &[u64::MAX]);
+// let merge_order = StackMergeOrder::stack(&[&columnar1]).into();
+// let group_error = group_columns_for_merge_iter(
+//&[&columnar1],
+//&[("numbers".to_string(), ColumnType::I64)],
+//&merge_order,
+//)
+//.unwrap_err();
+// assert_eq!(group_error.kind(), io::ErrorKind::InvalidInput);
+//}
 
 #[test]
 fn test_group_columns_with_required_column() {
@@ -77,12 +88,14 @@ fn test_group_columns_with_required_column() {
     let columnars = &[&columnar1, &columnar2];
     let merge_order = StackMergeOrder::stack(columnars).into();
     let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
-        group_columns_for_merge(
+        group_columns_for_merge_iter(
             &[&columnar1, &columnar2],
             &[("numbers".to_string(), ColumnType::U64)],
             &merge_order,
         )
-        .unwrap();
+        .unwrap()
+        .map(conv_res)
+        .collect();
     assert_eq!(column_map.len(), 1);
     assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::U64)));
 }
@@ -94,12 +107,14 @@ fn test_group_columns_required_column_with_no_existing_columns() {
     let columnars = &[&columnar1, &columnar2];
     let merge_order = StackMergeOrder::stack(columnars).into();
     let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
-        group_columns_for_merge(
+        group_columns_for_merge_iter(
             columnars,
             &[("required_col".to_string(), ColumnType::Str)],
             &merge_order,
         )
-        .unwrap();
+        .unwrap()
+        .map(conv_res)
+        .collect();
     assert_eq!(column_map.len(), 2);
     let columns = column_map
         .get(&("required_col".to_string(), ColumnType::Str))
@@ -116,24 +131,36 @@ fn test_group_columns_required_column_is_above_all_columns_have_the_same_type_ru
     let columnars = &[&columnar1, &columnar2];
     let merge_order = StackMergeOrder::stack(columnars).into();
     let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
-        group_columns_for_merge(
+        group_columns_for_merge_iter(
             columnars,
             &[("numbers".to_string(), ColumnType::U64)],
             &merge_order,
         )
-        .unwrap();
+        .unwrap()
+        .map(conv_res)
+        .collect();
     assert_eq!(column_map.len(), 1);
     assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::U64)));
 }
 
+fn conv_res(
+    el: io::Result<(Rc<String>, ColumnType, GroupedColumns)>,
+) -> ((String, ColumnType), Vec<Option<DynamicColumn>>) {
+    let el = el.unwrap();
+    ((el.0.to_string(), el.1), el.2.columns)
+}
+
 #[test]
 fn test_missing_column() {
     let columnar1 = make_columnar("numbers", &[-1i64]);
     let columnar2 = make_columnar("numbers2", &[2u64]);
     let columnars = &[&columnar1, &columnar2];
     let merge_order = StackMergeOrder::stack(columnars).into();
     let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
-        group_columns_for_merge(columnars, &[], &merge_order).unwrap();
+        group_columns_for_merge_iter(columnars, &[], &merge_order)
+            .unwrap()
+            .map(conv_res)
+            .collect();
     assert_eq!(column_map.len(), 2);
     assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::I64)));
     {