Skip to main content

datui_lib/
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_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::chart_data::{COUNT_COLUMN, Counted, Tally, count_frame};
12use crate::sampling::ReadWatch;
13use color_eyre::Result;
14use color_eyre::eyre::eyre;
15use polars::prelude::*;
16use std::sync::{Arc, Mutex};
17
18/// Values listed one per line; the rest are summed into one `other` line.
19pub const TOP_N: usize = 1_000;
20
21/// Distinct values a count keeps before it stops: past this the column is an
22/// identifier, and the counts would grow with the table.
23pub const MAX_DISTINCT: usize = 2_000_000;
24
25/// A local view this many times the sample size is sampled first, where a sample
26/// reads less of it: 10,000,000 rows at the default sample size.
27pub const LARGE_SAMPLES: usize = 100;
28
29/// Bins of the histogram view. An integer column spanning fewer values than this
30/// takes a bin per value instead.
31pub const HISTOGRAM_BINS: usize = 40;
32
33/// Which way the values are listed.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
35pub enum Order {
36    /// Most rows first; equal counts in value order.
37    #[default]
38    Count,
39    /// In the column's own order: text A to Z, numbers ascending.
40    Value,
41}
42
43impl Order {
44    pub fn toggled(self) -> Self {
45        match self {
46            Order::Count => Order::Value,
47            Order::Value => Order::Count,
48        }
49    }
50}
51
52/// How a count reads the view.
53#[derive(Debug, Clone, PartialEq)]
54pub enum Read {
55    /// Exact, unless the view is large or remote and a sample reads less of it:
56    /// then `sample_rows` of it, picked with `seed`.
57    Quick {
58        sample_rows: usize,
59        seed: u64,
60        remote: bool,
61    },
62    /// Every row.
63    Exact,
64}
65
66/// A count to run off the UI thread: the view, the column, and how to read it.
67pub struct Plan {
68    /// The view as filtered and queried. Its order does not change a count.
69    pub lf: LazyFrame,
70    pub column: String,
71    pub read: Read,
72    /// The view's row count, when the table knows it.
73    pub known_total: Option<usize>,
74    pub streaming: bool,
75}
76
77impl Plan {
78    /// Count the column, stopping when `watch` says to and counting the rows read.
79    pub fn run(&self, watch: &ReadWatch) -> Result<ValueCounts> {
80        let lf = self.lf.clone().select([col(self.column.as_str())]);
81        let dtype = lf
82            .clone()
83            .collect_schema()?
84            .get(self.column.as_str())
85            .cloned()
86            .ok_or_else(|| eyre!("no column {}", self.column))?;
87        if let Some((rows, seed)) = self.sample(&lf) {
88            let read = crate::statistics::sample_rows_counting(
89                &lf,
90                Some(rows),
91                self.known_total,
92                seed,
93                self.streaming,
94                Some(watch),
95                None,
96            )?;
97            watch.check()?;
98            let counted = count_frame(&read.rows.df, &self.column, MAX_DISTINCT)?;
99            // A view no larger than the sample was read whole: that is exact.
100            let of = read.rows.sample_size.map(|_| read.rows.total_rows);
101            return ValueCounts::new(&self.column, dtype, counted, of);
102        }
103        let counted = stream_counts(&lf, &self.column, watch)?;
104        ValueCounts::new(&self.column, dtype, counted, None)
105    }
106
107    /// The sample to read first, if one is worth it: the view is large or remote,
108    /// and a sample of it reads only part of it.
109    fn sample(&self, lf: &LazyFrame) -> Option<(usize, u64)> {
110        let Read::Quick {
111            sample_rows,
112            seed,
113            remote,
114        } = self.read
115        else {
116            return None;
117        };
118        if sample_rows == 0 {
119            return None;
120        }
121        let large = remote
122            || self
123                .known_total
124                .is_none_or(|rows| rows > sample_rows.saturating_mul(LARGE_SAMPLES));
125        (large && crate::statistics::slices_reach_into_the_scan(lf)).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 state = Arc::new(Mutex::new(Tally::new(column, MAX_DISTINCT)));
133    let tally = Arc::clone(&state);
134    let seen = watch.clone();
135    let sink = lf.clone().sink_batches(
136        PlanCallback::new(move |batch: DataFrame| {
137            if seen.stopped() {
138                return Ok(true);
139            }
140            seen.saw(batch.height());
141            tally
142                .lock()
143                .map_err(|_| PolarsError::ComputeError("count lock failed".into()))?
144                .observe(&batch)
145        }),
146        false,
147        None,
148    )?;
149    // Streaming whatever the setting: a count per value is all this holds.
150    crate::statistics::collect_lazy(sink, true)?;
151    // Part of the view counted is not a count of it.
152    watch.check()?;
153    let tally = std::mem::replace(
154        &mut *state.lock().unwrap_or_else(|e| e.into_inner()),
155        Tally::new(column, MAX_DISTINCT),
156    );
157    Ok(tally.finish()?)
158}
159
160/// A number the summary adds up: whole for integer columns, so a large sum keeps
161/// every digit.
162#[derive(Debug, Clone, Copy, PartialEq)]
163pub enum Number {
164    Int(i128),
165    Float(f64),
166}
167
168/// The header strip: what can be said of the column from its counts alone.
169#[derive(Debug, Clone, PartialEq, Default)]
170pub struct Summary {
171    /// Rows counted, nulls included.
172    pub rows: usize,
173    /// Distinct values, null not among them.
174    pub distinct: usize,
175    pub nulls: usize,
176    /// Numbers only.
177    pub sum: Option<Number>,
178    pub mean: Option<f64>,
179    /// Numbers and dates, times and durations.
180    pub min: Option<AnyValue<'static>>,
181    pub max: Option<AnyValue<'static>>,
182}
183
184impl Summary {
185    /// The summary of a column whose distinct values are `values` (null among
186    /// them at most once), each standing for `counts` rows.
187    pub fn of(values: &Series, counts: &[u64]) -> PolarsResult<Self> {
188        let rows = counts.iter().sum::<u64>() as usize;
189        let nulls: u64 = values
190            .is_null()
191            .iter()
192            .zip(counts)
193            .filter(|(null, _)| null.unwrap_or(false))
194            .map(|(_, n)| n)
195            .sum();
196        let nulls = nulls as usize;
197        let distinct = values.len() - values.null_count();
198        let dtype = values.dtype();
199        let numeric = dtype.is_primitive_numeric() || matches!(dtype, DataType::Decimal(..));
200        let ordered = numeric || dtype.is_temporal();
201        let mut summary = Summary {
202            rows,
203            distinct,
204            nulls,
205            ..Summary::default()
206        };
207        if ordered && distinct > 0 {
208            summary.min = Some(values.min_reduce()?.value().clone().into_static());
209            summary.max = Some(values.max_reduce()?.value().clone().into_static());
210        }
211        if numeric {
212            let sum = weighted_sum(values, counts)?;
213            let present = rows - nulls;
214            summary.mean = (present > 0).then(|| {
215                let total = match sum {
216                    Number::Int(n) => n as f64,
217                    Number::Float(f) => f,
218                };
219                total / present as f64
220            });
221            summary.sum = Some(sum);
222        }
223        Ok(summary)
224    }
225}
226
227/// Each value times the rows holding it, added up: whole for integers.
228fn weighted_sum(values: &Series, counts: &[u64]) -> PolarsResult<Number> {
229    let dtype = values.dtype();
230    if dtype.is_integer() && !matches!(dtype, DataType::Int128) {
231        // Unsigned 64-bit values past i64 have their own path; every other integer
232        // fits in i64.
233        let total: i128 = if matches!(dtype, DataType::UInt64) {
234            values
235                .u64()?
236                .iter()
237                .zip(counts)
238                .filter_map(|(v, n)| v.map(|v| v as i128 * *n as i128))
239                .sum()
240        } else {
241            values
242                .cast(&DataType::Int64)?
243                .i64()?
244                .iter()
245                .zip(counts)
246                .filter_map(|(v, n)| v.map(|v| v as i128 * *n as i128))
247                .sum()
248        };
249        return Ok(Number::Int(total));
250    }
251    let total = values
252        .cast(&DataType::Float64)?
253        .f64()?
254        .iter()
255        .zip(counts)
256        .filter_map(|(v, n)| v.map(|v| v * *n as f64))
257        .sum();
258    Ok(Number::Float(total))
259}
260
261/// A number column's counts in bins: every value with its rows. The bins span the
262/// values, or the 1st to the 99th percentile when the tails reach ten times past
263/// it, and the values outside are counted. An integer column of few values has a
264/// bin per value.
265fn histogram_of(
266    column: &str,
267    values: &Series,
268    rows: &[u64],
269) -> Option<crate::chart_data::HistogramData> {
270    use crate::chart_data::{Clipped, HistogramBin, HistogramData, RowsRead, ValueRange};
271    let dtype = values.dtype();
272    if !dtype.is_primitive_numeric() {
273        return None;
274    }
275    let as_f64 = values.cast(&DataType::Float64).ok()?;
276    let mut pairs: Vec<(f64, u64)> = as_f64
277        .f64()
278        .ok()?
279        .iter()
280        .zip(rows)
281        .filter_map(|(v, n)| Some((v.filter(|v| v.is_finite())?, *n)))
282        .collect();
283    pairs.sort_by(|a, b| a.0.total_cmp(&b.0));
284    let (min, max) = (pairs.first()?.0, pairs.last()?.0);
285    let total: u64 = pairs.iter().map(|p| p.1).sum();
286    // The value at quantile `q`, weighted by rows.
287    let at = |q: f64| {
288        let wanted = ((q * total as f64).ceil() as u64).max(1);
289        let mut seen = 0;
290        for (v, n) in &pairs {
291            seen += n;
292            if seen >= wanted {
293                return *v;
294            }
295        }
296        max
297    };
298    let (p1, p99) = (at(0.01), at(0.99));
299    let clip = p99 > p1 && (max - min) > 10.0 * (p99 - p1);
300    let (lo, hi) = if clip { (p1, p99) } else { (min, max) };
301    let (bins, width, x_min) = if dtype.is_integer() && hi - lo < HISTOGRAM_BINS as f64 {
302        ((hi - lo) as usize + 1, 1.0, lo - 0.5)
303    } else if hi > lo {
304        (HISTOGRAM_BINS, (hi - lo) / HISTOGRAM_BINS as f64, lo)
305    } else {
306        (1, 1.0, lo - 0.5)
307    };
308    let mut counts = vec![0.0_f64; bins];
309    let mut outside = 0;
310    for (v, n) in pairs {
311        if v < lo || v > hi {
312            outside += n as usize;
313            continue;
314        }
315        let bin = (((v - x_min) / width).floor().max(0.0) as usize).min(bins - 1);
316        counts[bin] += n as f64;
317    }
318    let max_count = counts.iter().copied().fold(0.0, f64::max);
319    Some(HistogramData {
320        column: column.to_string(),
321        bins: counts
322            .into_iter()
323            .enumerate()
324            .map(|(i, count)| HistogramBin {
325                center: x_min + (i as f64 + 0.5) * width,
326                count,
327            })
328            .collect(),
329        groups: Vec::new(),
330        other: false,
331        share: false,
332        x_min,
333        x_max: x_min + bins as f64 * width,
334        max_count,
335        rows: RowsRead {
336            total_rows: total as usize,
337            sample_size: None,
338            envelope_steps: None,
339            seed: None,
340        },
341        clipped: clip.then_some(Clipped {
342            range: ValueRange::Percentile1To99,
343            outside,
344        }),
345    })
346}
347
348/// What one line of the listing stands for.
349#[derive(Debug, Clone, Copy, PartialEq, Eq)]
350pub enum LineKind {
351    /// The value at this row of the counts.
352    Value(usize),
353    /// The rows with no value.
354    Null,
355    /// The values past the top ones, this many of them.
356    Other(usize),
357}
358
359/// One line of the listing: what it stands for, its rows, and the rows of it and
360/// every line above it.
361#[derive(Debug, Clone, Copy, PartialEq, Eq)]
362pub struct Line {
363    pub kind: LineKind,
364    pub rows: u64,
365    pub cumulative: u64,
366}
367
368/// A column's values counted: every distinct value with its rows, the summary,
369/// and the listing in either order.
370#[derive(Debug, Clone)]
371pub struct ValueCounts {
372    pub column: String,
373    pub dtype: DataType,
374    /// Every distinct value and its rows: `column`, then [`COUNT_COLUMN`].
375    counts: DataFrame,
376    /// The rows of `counts` holding a value, by count and by value.
377    by_count: Vec<usize>,
378    by_value: Vec<usize>,
379    /// The row of `counts` holding null, if any row is null.
380    null_at: Option<usize>,
381    /// The rows of the view when the counts are of a sample of them.
382    pub sampled_of: Option<usize>,
383    pub summary: Summary,
384    count_lines: Vec<Line>,
385    value_lines: Vec<Line>,
386    /// A number column's counts in bins, for the histogram view; made with the
387    /// counts, off the UI thread.
388    pub histogram: Option<crate::chart_data::HistogramData>,
389}
390
391impl ValueCounts {
392    pub(crate) fn new(
393        column: &str,
394        dtype: DataType,
395        counted: Counted,
396        sampled_of: Option<usize>,
397    ) -> Result<Self> {
398        let counts = match counted {
399            Counted::All {
400                counts: Some(counts),
401                ..
402            } => counts,
403            Counted::All { counts: None, .. } => DataFrame::new_infer_height(vec![
404                Column::new_empty(column.into(), &dtype),
405                Column::new_empty(COUNT_COLUMN.into(), &DataType::UInt64),
406            ])?,
407            Counted::TooMany => {
408                return Err(eyre!(
409                    "more than {} distinct values: counting stopped",
410                    crate::numfmt::group_chrome(MAX_DISTINCT)
411                ));
412            }
413        };
414        let values = counts.column(column)?.as_materialized_series().clone();
415        let rows = row_counts(&counts)?;
416        let summary = Summary::of(&values, &rows)?;
417        // Nulls sort last, and are left out: they have a line of their own.
418        let present = values.len() - values.null_count();
419        let by_value: Vec<usize> = values
420            .arg_sort(
421                SortOptions::default()
422                    .with_nulls_last(true)
423                    .with_maintain_order(true),
424            )
425            .iter()
426            .flatten()
427            .map(|i| i as usize)
428            .take(present)
429            .collect();
430        let null_at = values
431            .is_null()
432            .iter()
433            .position(|null| null.unwrap_or(false));
434        // Stable from value order, so values with equal counts list in value order.
435        let mut by_count = by_value.clone();
436        by_count.sort_by(|&a, &b| rows[b].cmp(&rows[a]));
437        let null_rows = summary.nulls as u64;
438        let histogram = histogram_of(column, &values, &rows);
439        Ok(Self {
440            histogram,
441            column: column.to_string(),
442            dtype,
443            count_lines: listing(&by_count, &rows, null_rows, true),
444            value_lines: listing(&by_value, &rows, null_rows, false),
445            by_count,
446            by_value,
447            null_at,
448            counts,
449            sampled_of,
450            summary,
451        })
452    }
453
454    /// Whether the counts are of a sample of the view.
455    pub fn is_sample(&self) -> bool {
456        self.sampled_of.is_some()
457    }
458
459    /// The listing in `order`: the top values, then the nulls, then the rest.
460    pub fn lines(&self, order: Order) -> &[Line] {
461        match order {
462            Order::Count => &self.count_lines,
463            Order::Value => &self.value_lines,
464        }
465    }
466
467    /// The value at row `at` of the counts.
468    pub fn value(&self, at: usize) -> PolarsResult<AnyValue<'static>> {
469        Ok(self.counts.column(&self.column)?.get(at)?.into_static())
470    }
471
472    /// Every value with its rows, in `order`, the nulls after them: what a copy or an
473    /// export of the counts writes. Nothing is summed into an `other` line.
474    pub fn table(&self, order: Order) -> PolarsResult<DataFrame> {
475        let rows = row_counts(&self.counts)?;
476        let picked: Vec<usize> = match order {
477            Order::Count => &self.by_count,
478            Order::Value => &self.by_value,
479        }
480        .iter()
481        .copied()
482        .chain(self.null_at)
483        .collect();
484        let values = self.counts.column(&self.column)?.take(&IdxCa::from_vec(
485            "order".into(),
486            picked.iter().map(|&i| i as IdxSize).collect(),
487        ))?;
488        let counts: Vec<u64> = picked.iter().map(|&i| rows[i]).collect();
489        let total = self.summary.rows.max(1) as f64;
490        let percent: Vec<f64> = counts.iter().map(|&n| n as f64 * 100.0 / total).collect();
491        let mut running = 0u64;
492        let cumulative: Vec<f64> = counts
493            .iter()
494            .map(|n| {
495                running += n;
496                running as f64 * 100.0 / total
497            })
498            .collect();
499        // Named so none takes the counted column's own name.
500        let mut names = vec![self.column.clone()];
501        let mut name = |wanted: &str| {
502            let mut name = wanted.to_string();
503            while names.contains(&name) {
504                name.push('_');
505            }
506            names.push(name.clone());
507            PlSmallStr::from(name)
508        };
509        DataFrame::new_infer_height(vec![
510            values,
511            Column::new(name("count"), counts),
512            Column::new(name("percent"), percent),
513            Column::new(name("cumulative_percent"), cumulative),
514        ])
515    }
516}
517
518fn row_counts(counts: &DataFrame) -> PolarsResult<Vec<u64>> {
519    Ok(counts
520        .column(COUNT_COLUMN)?
521        .u64()?
522        .into_no_null_iter()
523        .collect())
524}
525
526/// The listing of values in `order` (rows of the counts, nulls not among them):
527/// the first [`TOP_N`] with the nulls among them, then the rest summed into one
528/// line. By count, the nulls rank by their rows, after values with as many; by
529/// value, they come last, as a sort puts them.
530fn listing(order: &[usize], rows: &[u64], nulls: u64, by_count: bool) -> Vec<Line> {
531    let mut lines = Vec::with_capacity(order.len().min(TOP_N) + 2);
532    let mut cumulative = 0u64;
533    let mut push = |kind, n: u64, lines: &mut Vec<Line>| {
534        cumulative += n;
535        lines.push(Line {
536            kind,
537            rows: n,
538            cumulative,
539        });
540    };
541    let mut null_owed = nulls > 0;
542    for &at in order.iter().take(TOP_N) {
543        if null_owed && by_count && rows[at] < nulls {
544            push(LineKind::Null, nulls, &mut lines);
545            null_owed = false;
546        }
547        push(LineKind::Value(at), rows[at], &mut lines);
548    }
549    if null_owed {
550        push(LineKind::Null, nulls, &mut lines);
551    }
552    if order.len() > TOP_N {
553        let rest: u64 = order[TOP_N..].iter().map(|&at| rows[at]).sum();
554        push(LineKind::Other(order.len() - TOP_N), rest, &mut lines);
555    }
556    lines
557}
558
559#[cfg(test)]
560mod tests {
561    use super::*;
562
563    fn count(df: DataFrame, column: &str) -> ValueCounts {
564        let plan = Plan {
565            lf: df.lazy(),
566            column: column.to_string(),
567            read: Read::Exact,
568            known_total: None,
569            streaming: false,
570        };
571        plan.run(&ReadWatch::default()).unwrap()
572    }
573
574    fn values(counts: &ValueCounts, order: Order) -> Vec<(String, u64)> {
575        counts
576            .lines(order)
577            .iter()
578            .map(|line| {
579                let label = match line.kind {
580                    LineKind::Value(at) => {
581                        crate::exact::str_value(&counts.value(at).unwrap()).into_owned()
582                    }
583                    LineKind::Null => "null".to_string(),
584                    LineKind::Other(n) => format!("other {n}"),
585                };
586                (label, line.rows)
587            })
588            .collect()
589    }
590
591    #[test]
592    fn counts_list_by_rows_then_by_value_with_nulls_on_their_own_line() {
593        let df = df!("k" => [Some("b"), Some("a"), None, Some("b"), Some("c"), Some("a"), Some("b"), None])
594            .unwrap();
595        let counts = count(df, "k");
596        let pair = |s: &str, n| (s.to_string(), n);
597        // By count the nulls rank by their rows, after values with as many.
598        assert_eq!(
599            values(&counts, Order::Count),
600            [pair("b", 3), pair("a", 2), pair("null", 2), pair("c", 1)]
601        );
602        assert_eq!(
603            values(&counts, Order::Value),
604            [pair("a", 2), pair("b", 3), pair("c", 1), pair("null", 2)]
605        );
606        let lines = counts.lines(Order::Count);
607        assert_eq!(
608            lines.iter().map(|l| l.cumulative).collect::<Vec<_>>(),
609            [3, 5, 7, 8]
610        );
611        assert_eq!(counts.summary.rows, 8);
612        assert_eq!(counts.summary.distinct, 3);
613        assert_eq!(counts.summary.nulls, 2);
614        // Text has no sum, mean or range.
615        assert_eq!(counts.summary.sum, None);
616        assert_eq!(counts.summary.min, None);
617        assert!(!counts.is_sample());
618    }
619
620    #[test]
621    fn past_the_top_values_the_rest_are_one_line() {
622        let ids: Vec<i64> = (0..TOP_N as i64 + 5).chain([0, 0, 1]).collect();
623        let counts = count(df!("id" => ids).unwrap(), "id");
624        let lines = counts.lines(Order::Count);
625        assert_eq!(lines.len(), TOP_N + 1);
626        assert_eq!(lines[0].rows, 3, "0 is the most common");
627        assert_eq!(lines[1].rows, 2);
628        let other = lines.last().unwrap();
629        assert_eq!(other.kind, LineKind::Other(5));
630        assert_eq!(other.rows, 5);
631        assert_eq!(other.cumulative, TOP_N as u64 + 8);
632        assert_eq!(counts.summary.distinct, TOP_N + 5);
633        // The table holds every value, nothing summed.
634        let table = counts.table(Order::Count).unwrap();
635        assert_eq!(table.height(), TOP_N + 5);
636        assert_eq!(
637            table.get_column_names(),
638            ["id", "count", "percent", "cumulative_percent"]
639        );
640        let cumulative = table.column("cumulative_percent").unwrap().f64().unwrap();
641        assert!((cumulative.get(TOP_N + 4).unwrap() - 100.0).abs() < 1e-9);
642    }
643
644    #[test]
645    fn integer_summary_is_exact_and_whole() {
646        let df =
647            df!("n" => [Some(3i64), Some(1), None, Some(3), Some(i64::MAX), Some(-2)]).unwrap();
648        let summary = count(df, "n").summary;
649        assert_eq!(summary.rows, 6);
650        assert_eq!(summary.distinct, 4);
651        assert_eq!(summary.nulls, 1);
652        let sum = 3 + 1 + 3 + i64::MAX as i128 - 2;
653        assert_eq!(summary.sum, Some(Number::Int(sum)));
654        assert_eq!(summary.mean, Some(sum as f64 / 5.0));
655        assert_eq!(summary.min, Some(AnyValue::Int64(-2)));
656        assert_eq!(summary.max, Some(AnyValue::Int64(i64::MAX)));
657    }
658
659    #[test]
660    fn summary_math_over_counts() {
661        // 2.5 twice, 1.0 once, 4.0 three times: the counts weight the sum.
662        let values = Series::new("x".into(), [2.5f64, 1.0, 4.0]);
663        let summary = Summary::of(&values, &[2, 1, 3]).unwrap();
664        assert_eq!(summary.rows, 6);
665        assert_eq!(summary.distinct, 3);
666        assert_eq!(summary.nulls, 0);
667        assert_eq!(summary.sum, Some(Number::Float(18.0)));
668        assert_eq!(summary.mean, Some(3.0));
669        assert_eq!(summary.min, Some(AnyValue::Float64(1.0)));
670        assert_eq!(summary.max, Some(AnyValue::Float64(4.0)));
671
672        // Unsigned values past i64 add up whole.
673        let values = Series::new("u".into(), [u64::MAX, 1]);
674        let summary = Summary::of(&values, &[2, 1]).unwrap();
675        assert_eq!(summary.sum, Some(Number::Int(u64::MAX as i128 * 2 + 1)));
676
677        // Only nulls: nothing to add, no mean, no range.
678        let values = Series::new_null("z".into(), 1)
679            .cast(&DataType::Int32)
680            .unwrap();
681        let summary = Summary::of(&values, &[4]).unwrap();
682        assert_eq!((summary.rows, summary.distinct, summary.nulls), (4, 0, 4));
683        assert_eq!(summary.sum, Some(Number::Int(0)));
684        assert_eq!(summary.mean, None);
685        assert_eq!(summary.min, None);
686    }
687
688    #[test]
689    fn dates_have_a_range_and_no_sum() {
690        let dates = Series::new("d".into(), [19000i32, 19005, 18999])
691            .cast(&DataType::Date)
692            .unwrap();
693        let summary = Summary::of(&dates, &[1, 1, 1]).unwrap();
694        assert_eq!(summary.sum, None);
695        assert_eq!(summary.min, Some(AnyValue::Date(18999)));
696        assert_eq!(summary.max, Some(AnyValue::Date(19005)));
697    }
698
699    #[test]
700    fn an_empty_view_counts_nothing() {
701        let df = df!("k" => Vec::<i32>::new()).unwrap();
702        let counts = count(df, "k");
703        assert!(counts.lines(Order::Count).is_empty());
704        assert_eq!(counts.summary.rows, 0);
705        assert_eq!(counts.summary.mean, None);
706    }
707
708    #[test]
709    fn a_stopped_count_is_no_count() {
710        let watch = ReadWatch::default();
711        watch.stop();
712        let plan = Plan {
713            lf: df!("k" => [1, 2, 3]).unwrap().lazy(),
714            column: "k".to_string(),
715            read: Read::Exact,
716            known_total: None,
717            streaming: false,
718        };
719        assert!(plan.run(&watch).is_err());
720    }
721
722    #[test]
723    fn a_small_or_in_memory_view_is_counted_exactly() {
724        // In memory, a sample would read every row: the count is exact.
725        let plan = Plan {
726            lf: df!("k" => (0..50i32).collect::<Vec<_>>()).unwrap().lazy(),
727            column: "k".to_string(),
728            read: Read::Quick {
729                sample_rows: 10,
730                seed: 1,
731                remote: true,
732            },
733            known_total: Some(50),
734            streaming: false,
735        };
736        let counts = plan.run(&ReadWatch::default()).unwrap();
737        assert!(!counts.is_sample());
738        assert_eq!(counts.summary.rows, 50);
739    }
740
741    #[test]
742    fn a_large_parquet_file_is_sampled_first_and_says_of_how_many() {
743        let dir = tempfile::tempdir().unwrap();
744        let path = dir.path().join("big.parquet");
745        let mut df = df!("k" => (0..20_000i64).map(|i| i % 7).collect::<Vec<_>>()).unwrap();
746        let file = std::fs::File::create(&path).unwrap();
747        ParquetWriter::new(file)
748            .with_row_group_size(Some(1_000))
749            .finish(&mut df)
750            .unwrap();
751        let lf =
752            LazyFrame::scan_parquet(PlRefPath::try_from_path(&path).unwrap(), Default::default())
753                .unwrap();
754        let plan = |read| Plan {
755            lf: lf.clone(),
756            column: "k".to_string(),
757            read,
758            known_total: Some(20_000),
759            streaming: false,
760        };
761        let quick = Read::Quick {
762            sample_rows: 2_000,
763            seed: 7,
764            remote: true,
765        };
766        let sampled = plan(quick).run(&ReadWatch::default()).unwrap();
767        assert_eq!(sampled.sampled_of, Some(20_000));
768        assert_eq!(sampled.summary.rows, 2_000);
769        let exact = plan(Read::Exact).run(&ReadWatch::default()).unwrap();
770        assert!(!exact.is_sample());
771        assert_eq!(exact.summary.rows, 20_000);
772        assert_eq!(exact.summary.distinct, 7);
773    }
774
775    /// Every line's rows against a group-by of the same frame, and the rows a
776    /// drill into each line's value finds, for the types that group oddly: NaN and
777    /// -0.0, empty and escaped text, categoricals, dates, lists, structs, decimals.
778    #[test]
779    fn counts_and_drills_agree_with_a_group_by_for_every_kind_of_value() {
780        let mixed = df!(
781            "f" => [Some(1.5f64), Some(f64::NAN), None, Some(f64::NAN), Some(-0.0), Some(0.1 + 0.2)],
782            "s" => [Some(""), Some("a\tb"), Some("x\ny"), None, Some(""), Some("'\"\\")],
783            "b" => [Some(true), Some(false), None, Some(true), Some(true), None],
784            "l" => [Some(Series::new("".into(), [1i32, 2])), None, Some(Series::new("".into(), [1i32, 2])), Some(Series::new("".into(), Vec::<i32>::new())), None, None],
785            "n" => [1.25f64, 1.25, 3.5, 3.5, 3.5, 0.0],
786            "d" => [Some(19000i32), None, Some(19000), Some(1), None, Some(1)],
787        )
788        .unwrap()
789        .lazy()
790        .with_columns([
791            col("s")
792                .cast(DataType::from_categories(Categories::global()))
793                .alias("c"),
794            col("n").cast(DataType::Decimal(10, 2)).alias("x"),
795            col("d").cast(DataType::Date),
796            as_struct(vec![col("b"), col("d")]).alias("st"),
797        ])
798        .collect()
799        .unwrap();
800        for column in ["f", "s", "b", "l", "c", "x", "d", "st"] {
801            let counts = count(mixed.clone(), column);
802            let groups = mixed
803                .clone()
804                .lazy()
805                .group_by([col(column)])
806                .agg([len()])
807                .collect()
808                .unwrap()
809                .height();
810            let summary = &counts.summary;
811            assert_eq!(
812                summary.distinct + usize::from(summary.nulls > 0),
813                groups,
814                "{column}"
815            );
816            let lines = counts.lines(Order::Count);
817            assert_eq!(lines.last().unwrap().cumulative, 6, "{column}");
818            let dtype = mixed.schema().get(column).unwrap().clone();
819            for line in lines {
820                let value = match line.kind {
821                    LineKind::Value(at) => counts.value(at).unwrap(),
822                    LineKind::Null => AnyValue::Null,
823                    LineKind::Other(_) => unreachable!(),
824                };
825                let found = mixed
826                    .clone()
827                    .lazy()
828                    .filter(col(column).eq_missing(lit(Scalar::new(dtype.clone(), value))))
829                    .collect()
830                    .unwrap()
831                    .height();
832                assert_eq!(found as u64, line.rows, "{column} {:?}", line.kind);
833            }
834        }
835    }
836}