Skip to main content

datui_lib/analysis/
value_counts.rs

1//! Value counts: how many rows of the view hold each value of one column, and a
2//! summary of the column, from one read of that column.
3//!
4//! The count is exact: one streamed pass over the column that keeps a count per
5//! value and no rows (`crate::chart::chart_data::Tally`). A view too large to count at
6//! once, where the sampler can read part of it (one Parquet or IPC file, read a few
7//! row groups at a time), is sampled first instead and says so; counting every row
8//! is then the user's call. Where the sampler would stream every row anyway, the
9//! exact count is the same read and is what runs.
10
11use crate::analysis::sampling::ReadWatch;
12use crate::chart::chart_data::{COUNT_COLUMN, Counted, Tally, count_frame};
13use color_eyre::Result;
14use color_eyre::eyre::eyre;
15use polars::prelude::*;
16
17/// Values listed one per line; the rest are summed into one `other` line.
18pub const TOP_N: usize = 1_000;
19
20/// Distinct values a count keeps before it stops: past this the column is an
21/// identifier, and the counts would grow with the table.
22pub const MAX_DISTINCT: usize = 2_000_000;
23
24/// A local view this many times the sample size is sampled first, where a sample
25/// reads less of it: 10,000,000 rows at the default sample size.
26pub const LARGE_SAMPLES: usize = 100;
27
28/// Bins of the histogram view. An integer column spanning fewer values than this
29/// takes a bin per value instead.
30pub const HISTOGRAM_BINS: usize = 40;
31
32/// Which way the values are listed.
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
34pub enum Order {
35    /// Most rows first; equal counts in value order.
36    #[default]
37    Count,
38    /// In the column's own order: text A to Z, numbers ascending.
39    Value,
40}
41
42impl Order {
43    pub fn toggled(self) -> Self {
44        match self {
45            Order::Count => Order::Value,
46            Order::Value => Order::Count,
47        }
48    }
49}
50
51/// How a count reads the view.
52#[derive(Debug, Clone, PartialEq)]
53pub enum Read {
54    /// Exact, unless the view is large or remote and a sample reads less of it:
55    /// then `sample_rows` of it, picked with `seed`.
56    Quick {
57        sample_rows: usize,
58        seed: u64,
59        remote: bool,
60    },
61    /// Every row.
62    Exact,
63}
64
65/// A count to run off the UI thread: the view, the column, and how to read it.
66pub struct Plan {
67    /// The view as filtered and queried. Its order does not change a count.
68    pub lf: LazyFrame,
69    pub column: String,
70    pub read: Read,
71    /// The view's row count, when the table knows it.
72    pub known_total: Option<usize>,
73    pub streaming: bool,
74}
75
76impl Plan {
77    /// Count the column, stopping when `watch` says to and counting the rows read.
78    pub fn run(&self, watch: &ReadWatch) -> Result<ValueCounts> {
79        let lf = self.lf.clone().select([col(self.column.as_str())]);
80        let dtype = lf
81            .clone()
82            .collect_schema()?
83            .get(self.column.as_str())
84            .cloned()
85            .ok_or_else(|| eyre!("no column {}", self.column))?;
86        if let Some((rows, seed)) = self.sample(&lf) {
87            let read = crate::analysis::sampling::sample_rows_counting(
88                &lf,
89                Some(rows),
90                self.known_total,
91                seed,
92                self.streaming,
93                Some(watch),
94                None,
95            )?;
96            watch.check()?;
97            let counted = count_frame(&read.rows.df, &self.column, MAX_DISTINCT)?;
98            // A view no larger than the sample was read whole: that is exact.
99            let of = read.rows.sample_size.map(|_| read.rows.total_rows);
100            return ValueCounts::new(&self.column, dtype, counted, of);
101        }
102        let counted = stream_counts(&lf, &self.column, watch)?;
103        ValueCounts::new(&self.column, dtype, counted, None)
104    }
105
106    /// The sample to read first, if one is worth it: the view is large or remote,
107    /// and a sample of it reads only part of it.
108    fn sample(&self, lf: &LazyFrame) -> Option<(usize, u64)> {
109        let Read::Quick {
110            sample_rows,
111            seed,
112            remote,
113        } = self.read
114        else {
115            return None;
116        };
117        if sample_rows == 0 {
118            return None;
119        }
120        let large = remote
121            || self
122                .known_total
123                .is_none_or(|rows| rows > sample_rows.saturating_mul(LARGE_SAMPLES));
124        (large && crate::analysis::sampling::slices_reach_into_the_scan(lf))
125            .then_some((sample_rows, seed))
126    }
127}
128
129/// Count `column` of `lf` in one streamed pass, stopping between batches when
130/// `watch` says to.
131fn stream_counts(lf: &LazyFrame, column: &str, watch: &ReadWatch) -> Result<Counted> {
132    let tally = crate::analysis::sampling::stream_fold(
133        lf.clone(),
134        Some(watch),
135        false,
136        Tally::new(column, MAX_DISTINCT),
137        |tally, batch| tally.observe(&batch),
138    )?;
139    // Part of the view counted is not a count of it.
140    watch.check()?;
141    Ok(tally.finish()?)
142}
143
144/// A number the summary adds up: whole for integer columns, so a large sum keeps
145/// every digit.
146#[derive(Debug, Clone, Copy, PartialEq)]
147pub enum Number {
148    Int(i128),
149    Float(f64),
150}
151
152/// The header strip: what can be said of the column from its counts alone.
153#[derive(Debug, Clone, PartialEq, Default)]
154pub struct Summary {
155    /// Rows counted, nulls included.
156    pub rows: usize,
157    /// Distinct values, null not among them.
158    pub distinct: usize,
159    pub nulls: usize,
160    /// Numbers only.
161    pub sum: Option<Number>,
162    pub mean: Option<f64>,
163    /// Numbers and dates, times and durations.
164    pub min: Option<AnyValue<'static>>,
165    pub max: Option<AnyValue<'static>>,
166}
167
168impl Summary {
169    /// The summary of a column whose distinct values are `values` (null among
170    /// them at most once), each standing for `counts` rows.
171    pub fn of(values: &Series, counts: &[u64]) -> PolarsResult<Self> {
172        let rows = counts.iter().sum::<u64>() as usize;
173        let nulls: u64 = values
174            .is_null()
175            .iter()
176            .zip(counts)
177            .filter(|(null, _)| null.unwrap_or(false))
178            .map(|(_, n)| n)
179            .sum();
180        let nulls = nulls as usize;
181        let distinct = values.len() - values.null_count();
182        let dtype = values.dtype();
183        let numeric = dtype.is_primitive_numeric() || matches!(dtype, DataType::Decimal(..));
184        let ordered = numeric || dtype.is_temporal();
185        let mut summary = Summary {
186            rows,
187            distinct,
188            nulls,
189            ..Summary::default()
190        };
191        if ordered && distinct > 0 {
192            summary.min = Some(values.min_reduce()?.value().clone().into_static());
193            summary.max = Some(values.max_reduce()?.value().clone().into_static());
194        }
195        if numeric {
196            let sum = weighted_sum(values, counts)?;
197            let present = rows - nulls;
198            summary.mean = (present > 0).then(|| {
199                let total = match sum {
200                    Number::Int(n) => n as f64,
201                    Number::Float(f) => f,
202                };
203                total / present as f64
204            });
205            summary.sum = Some(sum);
206        }
207        Ok(summary)
208    }
209}
210
211/// Each value times the rows holding it, added up: whole for integers.
212fn weighted_sum(values: &Series, counts: &[u64]) -> PolarsResult<Number> {
213    let dtype = values.dtype();
214    if dtype.is_integer() && !matches!(dtype, DataType::Int128) {
215        // Unsigned 64-bit values past i64 have their own path; every other integer
216        // fits in i64.
217        let total: i128 = if matches!(dtype, DataType::UInt64) {
218            values
219                .u64()?
220                .iter()
221                .zip(counts)
222                .filter_map(|(v, n)| v.map(|v| v as i128 * *n as i128))
223                .sum()
224        } else {
225            values
226                .cast(&DataType::Int64)?
227                .i64()?
228                .iter()
229                .zip(counts)
230                .filter_map(|(v, n)| v.map(|v| v as i128 * *n as i128))
231                .sum()
232        };
233        return Ok(Number::Int(total));
234    }
235    let total = values
236        .cast(&DataType::Float64)?
237        .f64()?
238        .iter()
239        .zip(counts)
240        .filter_map(|(v, n)| v.map(|v| v * *n as f64))
241        .sum();
242    Ok(Number::Float(total))
243}
244
245/// A number column's counts in bins: every value with its rows. The bins span the
246/// values, or the 1st to the 99th percentile when the tails reach ten times past
247/// it, and the values outside are counted. An integer column of few values has a
248/// bin per value.
249fn histogram_of(
250    column: &str,
251    values: &Series,
252    rows: &[u64],
253) -> Option<crate::chart::chart_data::HistogramData> {
254    use crate::chart::chart_data::{Clipped, HistogramBin, HistogramData, RowsRead, ValueRange};
255    let dtype = values.dtype();
256    if !dtype.is_primitive_numeric() {
257        return None;
258    }
259    let as_f64 = values.cast(&DataType::Float64).ok()?;
260    let mut pairs: Vec<(f64, u64)> = as_f64
261        .f64()
262        .ok()?
263        .iter()
264        .zip(rows)
265        .filter_map(|(v, n)| Some((v.filter(|v| v.is_finite())?, *n)))
266        .collect();
267    pairs.sort_by(|a, b| a.0.total_cmp(&b.0));
268    let (min, max) = (pairs.first()?.0, pairs.last()?.0);
269    let total: u64 = pairs.iter().map(|p| p.1).sum();
270    // The value at quantile `q`, weighted by rows.
271    let at = |q: f64| {
272        let wanted = ((q * total as f64).ceil() as u64).max(1);
273        let mut seen = 0;
274        for (v, n) in &pairs {
275            seen += n;
276            if seen >= wanted {
277                return *v;
278            }
279        }
280        max
281    };
282    let (p1, p99) = (at(0.01), at(0.99));
283    let clip = p99 > p1 && (max - min) > 10.0 * (p99 - p1);
284    let (lo, hi) = if clip { (p1, p99) } else { (min, max) };
285    let (bins, width, x_min) = if dtype.is_integer() && hi - lo < HISTOGRAM_BINS as f64 {
286        ((hi - lo) as usize + 1, 1.0, lo - 0.5)
287    } else if hi > lo {
288        (HISTOGRAM_BINS, (hi - lo) / HISTOGRAM_BINS as f64, lo)
289    } else {
290        (1, 1.0, lo - 0.5)
291    };
292    let mut counts = vec![0.0_f64; bins];
293    let mut outside = 0;
294    for (v, n) in pairs {
295        if v < lo || v > hi {
296            outside += n as usize;
297            continue;
298        }
299        let bin = (((v - x_min) / width).floor().max(0.0) as usize).min(bins - 1);
300        counts[bin] += n as f64;
301    }
302    let max_count = counts.iter().copied().fold(0.0, f64::max);
303    Some(HistogramData {
304        column: column.to_string(),
305        bins: counts
306            .into_iter()
307            .enumerate()
308            .map(|(i, count)| HistogramBin {
309                center: x_min + (i as f64 + 0.5) * width,
310                count,
311            })
312            .collect(),
313        groups: Vec::new(),
314        other: false,
315        share: false,
316        x_min,
317        x_max: x_min + bins as f64 * width,
318        max_count,
319        rows: RowsRead {
320            total_rows: total as usize,
321            sample_size: None,
322            envelope_steps: None,
323            seed: None,
324        },
325        clipped: clip.then_some(Clipped {
326            range: ValueRange::Percentile1To99,
327            outside,
328        }),
329    })
330}
331
332/// What one line of the listing stands for.
333#[derive(Debug, Clone, Copy, PartialEq, Eq)]
334pub enum LineKind {
335    /// The value at this row of the counts.
336    Value(usize),
337    /// The rows with no value.
338    Null,
339    /// The values past the top ones, this many of them.
340    Other(usize),
341}
342
343/// One line of the listing: what it stands for, its rows, and the rows of it and
344/// every line above it.
345#[derive(Debug, Clone, Copy, PartialEq, Eq)]
346pub struct Line {
347    pub kind: LineKind,
348    pub rows: u64,
349    pub cumulative: u64,
350}
351
352/// A column's values counted: every distinct value with its rows, the summary,
353/// and the listing in either order.
354#[derive(Debug, Clone)]
355pub struct ValueCounts {
356    pub column: String,
357    pub dtype: DataType,
358    /// Every distinct value and its rows: `column`, then [`COUNT_COLUMN`].
359    counts: DataFrame,
360    /// The rows of `counts` holding a value, by count and by value.
361    by_count: Vec<usize>,
362    by_value: Vec<usize>,
363    /// The row of `counts` holding null, if any row is null.
364    null_at: Option<usize>,
365    /// The rows of the view when the counts are of a sample of them.
366    pub sampled_of: Option<usize>,
367    pub summary: Summary,
368    count_lines: Vec<Line>,
369    value_lines: Vec<Line>,
370    /// A number column's counts in bins, for the histogram view; made with the
371    /// counts, off the UI thread.
372    pub histogram: Option<crate::chart::chart_data::HistogramData>,
373}
374
375impl ValueCounts {
376    pub(crate) fn new(
377        column: &str,
378        dtype: DataType,
379        counted: Counted,
380        sampled_of: Option<usize>,
381    ) -> Result<Self> {
382        let counts = match counted {
383            Counted::All {
384                counts: Some(counts),
385                ..
386            } => counts,
387            Counted::All { counts: None, .. } => DataFrame::new_infer_height(vec![
388                Column::new_empty(column.into(), &dtype),
389                Column::new_empty(COUNT_COLUMN.into(), &DataType::UInt64),
390            ])?,
391            Counted::TooMany => {
392                return Err(eyre!(
393                    "more than {} distinct values: counting stopped",
394                    crate::numfmt::group_chrome(MAX_DISTINCT)
395                ));
396            }
397        };
398        let values = counts.column(column)?.as_materialized_series().clone();
399        let rows = row_counts(&counts)?;
400        let summary = Summary::of(&values, &rows)?;
401        // Nulls sort last, and are left out: they have a line of their own.
402        let present = values.len() - values.null_count();
403        let by_value: Vec<usize> = values
404            .arg_sort(
405                SortOptions::default()
406                    .with_nulls_last(true)
407                    .with_maintain_order(true),
408            )
409            .iter()
410            .flatten()
411            .map(|i| i as usize)
412            .take(present)
413            .collect();
414        let null_at = values
415            .is_null()
416            .iter()
417            .position(|null| null.unwrap_or(false));
418        // Stable from value order, so values with equal counts list in value order.
419        let mut by_count = by_value.clone();
420        by_count.sort_by(|&a, &b| rows[b].cmp(&rows[a]));
421        let null_rows = summary.nulls as u64;
422        let histogram = histogram_of(column, &values, &rows);
423        Ok(Self {
424            histogram,
425            column: column.to_string(),
426            dtype,
427            count_lines: listing(&by_count, &rows, null_rows, true),
428            value_lines: listing(&by_value, &rows, null_rows, false),
429            by_count,
430            by_value,
431            null_at,
432            counts,
433            sampled_of,
434            summary,
435        })
436    }
437
438    /// Whether the counts are of a sample of the view.
439    pub fn is_sample(&self) -> bool {
440        self.sampled_of.is_some()
441    }
442
443    /// The listing in `order`: the top values, then the nulls, then the rest.
444    pub fn lines(&self, order: Order) -> &[Line] {
445        match order {
446            Order::Count => &self.count_lines,
447            Order::Value => &self.value_lines,
448        }
449    }
450
451    /// The value at row `at` of the counts.
452    pub fn value(&self, at: usize) -> PolarsResult<AnyValue<'static>> {
453        Ok(self.counts.column(&self.column)?.get(at)?.into_static())
454    }
455
456    /// Every value with its rows, in `order`, the nulls after them: what a copy or an
457    /// export of the counts writes. Nothing is summed into an `other` line.
458    pub fn table(&self, order: Order) -> PolarsResult<DataFrame> {
459        let rows = row_counts(&self.counts)?;
460        let picked: Vec<usize> = match order {
461            Order::Count => &self.by_count,
462            Order::Value => &self.by_value,
463        }
464        .iter()
465        .copied()
466        .chain(self.null_at)
467        .collect();
468        let values = self.counts.column(&self.column)?.take(&IdxCa::from_vec(
469            "order".into(),
470            picked.iter().map(|&i| i as IdxSize).collect(),
471        ))?;
472        let counts: Vec<u64> = picked.iter().map(|&i| rows[i]).collect();
473        let total = self.summary.rows.max(1) as f64;
474        let percent: Vec<f64> = counts.iter().map(|&n| n as f64 * 100.0 / total).collect();
475        let mut running = 0u64;
476        let cumulative: Vec<f64> = counts
477            .iter()
478            .map(|n| {
479                running += n;
480                running as f64 * 100.0 / total
481            })
482            .collect();
483        // Named so none takes the counted column's own name.
484        let mut names = vec![self.column.clone()];
485        let mut name = |wanted: &str| {
486            let mut name = wanted.to_string();
487            while names.contains(&name) {
488                name.push('_');
489            }
490            names.push(name.clone());
491            PlSmallStr::from(name)
492        };
493        DataFrame::new_infer_height(vec![
494            values,
495            Column::new(name("count"), counts),
496            Column::new(name("percent"), percent),
497            Column::new(name("cumulative_percent"), cumulative),
498        ])
499    }
500}
501
502fn row_counts(counts: &DataFrame) -> PolarsResult<Vec<u64>> {
503    Ok(counts
504        .column(COUNT_COLUMN)?
505        .u64()?
506        .into_no_null_iter()
507        .collect())
508}
509
510/// The listing of values in `order` (rows of the counts, nulls not among them):
511/// the first [`TOP_N`] with the nulls among them, then the rest summed into one
512/// line. By count, the nulls rank by their rows, after values with as many; by
513/// value, they come last, as a sort puts them.
514fn listing(order: &[usize], rows: &[u64], nulls: u64, by_count: bool) -> Vec<Line> {
515    let mut lines = Vec::with_capacity(order.len().min(TOP_N) + 2);
516    let mut cumulative = 0u64;
517    let mut push = |kind, n: u64, lines: &mut Vec<Line>| {
518        cumulative += n;
519        lines.push(Line {
520            kind,
521            rows: n,
522            cumulative,
523        });
524    };
525    let mut null_owed = nulls > 0;
526    for &at in order.iter().take(TOP_N) {
527        if null_owed && by_count && rows[at] < nulls {
528            push(LineKind::Null, nulls, &mut lines);
529            null_owed = false;
530        }
531        push(LineKind::Value(at), rows[at], &mut lines);
532    }
533    if null_owed {
534        push(LineKind::Null, nulls, &mut lines);
535    }
536    if order.len() > TOP_N {
537        let rest: u64 = order[TOP_N..].iter().map(|&at| rows[at]).sum();
538        push(LineKind::Other(order.len() - TOP_N), rest, &mut lines);
539    }
540    lines
541}
542
543#[cfg(test)]
544mod tests;