zeppelin_embed/lifecycle/
group_count.rs

1//! Live-document counts grouped by one typed attribute (ZE-229).
2
3use std::collections::BTreeMap;
4use std::collections::btree_map::Entry;
5
6use crate::meta::{Column, ColumnId, ColumnStore, ColumnType, DocBitmap};
7
8use super::{AdmittedDocumentRead, QueryError, Store, StoreError};
9
10/// Largest group `limit` accepted by [`Store::count_documents_grouped`].
11pub const MAX_DOCUMENT_GROUP_LIMIT: usize = 1 << 16;
12
13/// The value shared by every document of one group.
14#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
15pub enum DocumentGroupValue {
16    /// A `U64` attribute value.
17    U64(u64),
18    /// An `I64` attribute value.
19    I64(i64),
20    /// A `DictionaryString` or `RawString` attribute value.
21    String(String),
22}
23
24/// One group and its exact live-document count.
25#[derive(Clone, Debug, Eq, PartialEq)]
26pub struct DocumentGroup {
27    /// Attribute value shared by the group's documents.
28    pub value: DocumentGroupValue,
29    /// Number of matching live documents with this value; never zero.
30    pub count: u64,
31}
32
33/// Grouped live-document counts from one pinned store generation.
34#[derive(Clone, Debug, Eq, PartialEq)]
35pub struct DocumentGroupCounts {
36    /// Groups in ascending value order: numeric for integers, byte order for
37    /// strings.
38    pub groups: Vec<DocumentGroup>,
39    /// Matching live documents whose group attribute is null.
40    pub missing: u64,
41    /// All matching live documents; the group counts plus `missing`.
42    pub count: u64,
43    /// Store generation pinned for every group.
44    pub generation: u64,
45}
46
47enum Tally {
48    U64(BTreeMap<u64, u64>),
49    I64(BTreeMap<i64, u64>),
50    String(BTreeMap<String, u64>),
51}
52
53impl Store {
54    /// Counts live documents matching an optional predicate and timestamp
55    /// range, grouped by the value of `column`.
56    ///
57    /// `column` must be a `U64`, `I64`, `DictionaryString` or `RawString`
58    /// schema attribute. Documents where it is null are counted in
59    /// [`DocumentGroupCounts::missing`], not as a group. More than `limit`
60    /// distinct values fail the whole call with
61    /// [`StoreError::GroupLimitExceeded`]; no group is ever dropped.
62    pub fn count_documents_grouped(
63        &self,
64        predicate: Option<&crate::meta::Predicate>,
65        timestamp_range: Option<(i64, i64)>,
66        column: ColumnId,
67        limit: usize,
68    ) -> Result<DocumentGroupCounts, QueryError> {
69        let invalid = |detail: String| QueryError::Store(StoreError::InvalidScan { detail });
70        if limit == 0 || limit > MAX_DOCUMENT_GROUP_LIMIT {
71            return Err(invalid(format!(
72                "group limit must be in 1..={MAX_DOCUMENT_GROUP_LIMIT}, received {limit}"
73            )));
74        }
75        let definition = self.schema.column(column).ok_or_else(|| {
76            invalid(format!(
77                "group-by attribute {} is not in the schema",
78                column.get()
79            ))
80        })?;
81        let nullable = definition.is_nullable();
82        let mut tally = match definition.column_type() {
83            ColumnType::U64 => Tally::U64(BTreeMap::new()),
84            ColumnType::I64 => Tally::I64(BTreeMap::new()),
85            ColumnType::DictionaryString | ColumnType::RawString => Tally::String(BTreeMap::new()),
86            other @ (ColumnType::Id128 | ColumnType::F64 | ColumnType::Bool) => {
87                return Err(invalid(format!(
88                    "group-by attribute {} has type {other:?}; grouping supports \
89                     U64, I64, DictionaryString and RawString",
90                    column.get()
91                )));
92            }
93        };
94        let predicate = super::document_scan_predicate(predicate, timestamp_range)?;
95        if let Some(predicate) = &predicate {
96            crate::planner::validate_predicate(predicate, &self.schema)
97                .map_err(|error| invalid(error.to_string()))?;
98        }
99        let AdmittedDocumentRead {
100            generation,
101            active,
102            snapshot,
103            active_query,
104        } = self.admit_document_read().map_err(QueryError::Store)?;
105        let mut missing = 0_u64;
106        for segment in snapshot.segments() {
107            let rows = super::sealed_scan_rows(segment, predicate.as_ref())?;
108            if rows.is_empty() {
109                continue;
110            }
111            let columns = segment.query_columns().map_err(QueryError::Store)?;
112            missing = add(
113                missing,
114                tally.add_rows(&columns, &rows, column, nullable, limit)?,
115            )?;
116        }
117        let columns = super::active_scan_columns(&active, &self.schema)?;
118        let alive = active.alive().map_err(QueryError::Store)?;
119        let rows = match &predicate {
120            Some(predicate) => {
121                crate::meta::evaluate(predicate, &columns, &alive).map_err(|error| {
122                    QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Columns(
123                        error.to_string(),
124                    )))
125                })?
126            }
127            None => alive.alive_bitmap().clone(),
128        };
129        missing = add(
130            missing,
131            tally.add_rows(&columns, &rows, column, nullable, limit)?,
132        )?;
133        drop(active_query);
134        let groups = tally.into_groups()?;
135        let count = groups
136            .iter()
137            .try_fold(missing, |total, group| add(total, group.count))?;
138        Ok(DocumentGroupCounts {
139            groups,
140            missing,
141            count,
142            generation,
143        })
144    }
145}
146
147impl Tally {
148    /// Adds every row of `rows` and returns how many had a null value.
149    fn add_rows(
150        &mut self,
151        columns: &ColumnStore,
152        rows: &DocBitmap,
153        column: ColumnId,
154        nullable: bool,
155        limit: usize,
156    ) -> Result<u64, QueryError> {
157        let values = columns.column(column).ok_or_else(|| {
158            corrupt(format!(
159                "group-by attribute {} is missing from a segment",
160                column.get()
161            ))
162        })?;
163        let mut missing = 0_u64;
164        for row in rows.iter() {
165            let present = match (&mut *self, values) {
166                (Self::U64(map), Column::U64(values)) => values
167                    .get(row)
168                    .map(|value| bump(map, value, limit))
169                    .transpose()?,
170                (Self::I64(map), Column::I64(values)) => values
171                    .get(row)
172                    .map(|value| bump(map, value, limit))
173                    .transpose()?,
174                (Self::String(map), Column::DictionaryString(values)) => values
175                    .get(row)
176                    .map(|value| bump_string(map, value, limit))
177                    .transpose()?,
178                (Self::String(map), Column::RawString(values)) => values
179                    .get(row)
180                    .map(|value| bump_string(map, value, limit))
181                    .transpose()?,
182                (_, values) => {
183                    return Err(corrupt(format!(
184                        "group-by attribute {} is stored as {:?}",
185                        column.get(),
186                        values.column_type()
187                    )));
188                }
189            };
190            if present.is_none() {
191                if !nullable {
192                    return Err(corrupt(format!(
193                        "required group-by attribute {} is null at row {row}",
194                        column.get()
195                    )));
196                }
197                missing = add(missing, 1)?;
198            }
199        }
200        Ok(missing)
201    }
202
203    fn into_groups(self) -> Result<Vec<DocumentGroup>, QueryError> {
204        fn collect<K>(
205            map: BTreeMap<K, u64>,
206            value: impl Fn(K) -> DocumentGroupValue,
207        ) -> Result<Vec<DocumentGroup>, QueryError> {
208            let mut groups = Vec::new();
209            groups.try_reserve_exact(map.len()).map_err(|_| {
210                QueryError::Store(StoreError::AllocationFailed {
211                    needed: super::allocation_bytes::<DocumentGroup>(map.len()),
212                    component: "grouped count groups",
213                })
214            })?;
215            groups.extend(map.into_iter().map(|(key, count)| DocumentGroup {
216                value: value(key),
217                count,
218            }));
219            Ok(groups)
220        }
221        match self {
222            Self::U64(map) => collect(map, DocumentGroupValue::U64),
223            Self::I64(map) => collect(map, DocumentGroupValue::I64),
224            Self::String(map) => collect(map, DocumentGroupValue::String),
225        }
226    }
227}
228
229fn bump<K: Ord>(map: &mut BTreeMap<K, u64>, key: K, limit: usize) -> Result<(), QueryError> {
230    let len = map.len();
231    match map.entry(key) {
232        Entry::Occupied(mut entry) => {
233            *entry.get_mut() = add(*entry.get(), 1)?;
234        }
235        Entry::Vacant(_) if len >= limit => {
236            return Err(QueryError::Store(StoreError::GroupLimitExceeded { limit }));
237        }
238        Entry::Vacant(entry) => {
239            entry.insert(1);
240        }
241    }
242    Ok(())
243}
244
245fn bump_string(map: &mut BTreeMap<String, u64>, key: &str, limit: usize) -> Result<(), QueryError> {
246    if let Some(count) = map.get_mut(key) {
247        *count = add(*count, 1)?;
248        return Ok(());
249    }
250    if map.len() >= limit {
251        return Err(QueryError::Store(StoreError::GroupLimitExceeded { limit }));
252    }
253    let key = super::copy_string(key, "grouped count value").map_err(QueryError::Store)?;
254    map.insert(key, 1);
255    Ok(())
256}
257
258fn add(total: u64, count: u64) -> Result<u64, QueryError> {
259    total
260        .checked_add(count)
261        .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))
262}
263
264fn corrupt(detail: String) -> QueryError {
265    QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
266        detail,
267    )))
268}
zeppelin_embed/lifecycle/group_count.rs

zeppelin_embed/lifecycle/
group_count.rs
