Skip to main content

datui_lib/chart/
chart_data.rs

1//! Prepare chart data from a LazyFrame: read the chart's columns, then make points,
2//! bins or statistics. Every chart reads through `read_columns`: up to a row limit,
3//! spread across the table by the analysis sampler (not its first rows), returning
4//! [`RowsRead`] so the chart can say when it shows a sample.
5
6use chrono::{DateTime, Datelike, NaiveDate, NaiveDateTime, NaiveTime};
7use color_eyre::Result;
8use polars::chunked_array::cast::CastOptions;
9use polars::datatypes::{DataType, TimeUnit};
10use polars::prelude::*;
11use std::f64::consts::PI;
12use std::sync::atomic::{AtomicBool, Ordering};
13use std::sync::{Arc, Mutex};
14
15use crate::widgets::axis_numbers::{AxisFormat, format_bar_value, table_number_format};
16
17/// Describes how x-axis numeric values map to temporal types for label formatting.
18#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
19pub enum XAxisTemporalKind {
20    #[default]
21    Numeric,
22    Date,       // x = days since Unix epoch (f64)
23    DatetimeUs, // x = microseconds since epoch
24    DatetimeMs,
25    DatetimeNs,
26    Time, // x = nanoseconds since midnight
27}
28
29fn x_axis_temporal_kind(dtype: &DataType) -> XAxisTemporalKind {
30    match dtype {
31        DataType::Date => XAxisTemporalKind::Date,
32        DataType::Datetime(unit, _) => match unit {
33            TimeUnit::Nanoseconds => XAxisTemporalKind::DatetimeNs,
34            TimeUnit::Microseconds => XAxisTemporalKind::DatetimeUs,
35            TimeUnit::Milliseconds => XAxisTemporalKind::DatetimeMs,
36        },
37        DataType::Time => XAxisTemporalKind::Time,
38        _ => XAxisTemporalKind::Numeric,
39    }
40}
41
42/// Returns the x-axis temporal kind for a column from the schema (for axis label formatting when no data is loaded yet).
43pub fn x_axis_temporal_kind_for_column(schema: &Schema, x_column: &str) -> XAxisTemporalKind {
44    schema
45        .get(x_column)
46        .map(x_axis_temporal_kind)
47        .unwrap_or(XAxisTemporalKind::Numeric)
48}
49
50/// An x value as the date and time it stands for, when `kind` is a date or datetime.
51pub(crate) fn x_datetime(v: f64, kind: XAxisTemporalKind) -> Option<NaiveDateTime> {
52    const UNIX_EPOCH_CE_DAYS: i32 = 719_163;
53    match kind {
54        XAxisTemporalKind::Date => NaiveDate::from_num_days_from_ce_opt(
55            UNIX_EPOCH_CE_DAYS.saturating_add(v.trunc() as i32),
56        )
57        .map(|d| d.and_time(NaiveTime::MIN)),
58        XAxisTemporalKind::DatetimeUs => {
59            DateTime::from_timestamp_micros(v.trunc() as i64).map(|dt| dt.naive_utc())
60        }
61        XAxisTemporalKind::DatetimeMs => {
62            DateTime::from_timestamp_millis(v.trunc() as i64).map(|dt| dt.naive_utc())
63        }
64        XAxisTemporalKind::DatetimeNs => {
65            DateTime::from_timestamp_millis((v.trunc() as i64) / 1_000_000).map(|dt| dt.naive_utc())
66        }
67        XAxisTemporalKind::Numeric | XAxisTemporalKind::Time => None,
68    }
69}
70
71/// An x value as a time of day, when `kind` is a time.
72pub(crate) fn x_time(v: f64) -> Option<NaiveTime> {
73    let nsecs = v.trunc() as u64;
74    NaiveTime::from_num_seconds_from_midnight_opt(
75        (nsecs / 1_000_000_000) as u32,
76        (nsecs % 1_000_000_000) as u32,
77    )
78}
79
80/// An x tick at detail `level` (0 fullest), `None` past the shortest form. Narrow axes
81/// step down until labels fit: a date to year-month then year (or month-day when
82/// `bounds` fall in one year); a datetime to its date (or the minute within one day);
83/// a time to the minute. Numbers, and times out of range, use `numbers`.
84pub fn x_axis_label_at(
85    v: f64,
86    kind: XAxisTemporalKind,
87    bounds: (f64, f64),
88    level: usize,
89    numbers: &AxisFormat,
90) -> Option<String> {
91    if kind == XAxisTemporalKind::Numeric {
92        return numbers.label(v, level);
93    }
94    if kind == XAxisTemporalKind::Time {
95        let pattern = ["%H:%M:%S", "%H:%M"].get(level)?;
96        return Some(match x_time(v) {
97            Some(t) => t.format(pattern).to_string(),
98            None => numbers.label(v, level)?,
99        });
100    }
101    let Some(at) = x_datetime(v, kind) else {
102        return numbers.label(v, level);
103    };
104    let ends = x_datetime(bounds.0, kind).zip(x_datetime(bounds.1, kind));
105    let one_day = ends.is_some_and(|(a, b)| a.date() == b.date());
106    let one_year = ends.is_some_and(|(a, b)| a.year() == b.year());
107    let dates: &[&str] = if one_year {
108        &["%Y-%m-%d", "%m-%d"]
109    } else {
110        &["%Y-%m-%d", "%Y-%m", "%Y"]
111    };
112    let patterns: Vec<&str> = if kind == XAxisTemporalKind::Date {
113        dates.to_vec()
114    } else if one_day {
115        vec!["%Y-%m-%d %H:%M", "%H:%M"]
116    } else {
117        std::iter::once("%Y-%m-%d %H:%M")
118            .chain(dates.iter().copied())
119            .collect()
120    };
121    patterns.get(level).map(|p| at.format(p).to_string())
122}
123
124/// How a chart reads its rows.
125#[derive(Clone, Debug)]
126pub struct ChartSampling {
127    /// Rows to read; `None` reads every row.
128    pub limit: Option<usize>,
129    /// The view's row count when the table already knows it, which saves a count.
130    pub known_total: Option<usize>,
131    /// The shared analysis seed, so a chart and Describe draw alike.
132    pub seed: u64,
133    pub streaming: bool,
134    /// Whether the view may be read whole twice for a line's envelope: not an in-place
135    /// object-store scan, where the envelope would download everything twice. See
136    /// [`prepare_chart_data`].
137    pub full_passes: bool,
138    /// The rows already read from this view.
139    pub held: HeldRows,
140    /// Set once nobody wants the result: a streamed count stops at its next batch.
141    pub cancel: Arc<AtomicBool>,
142}
143
144impl ChartSampling {
145    /// Up to `limit` rows, with the analysis tools' default seed.
146    #[cfg(test)]
147    pub fn rows(limit: Option<usize>) -> Self {
148        Self {
149            limit,
150            known_total: None,
151            seed: crate::analysis::sampling::Sample::default().seed,
152            streaming: false,
153            full_passes: true,
154            held: HeldRows::default(),
155            cancel: Arc::default(),
156        }
157    }
158}
159
160/// The rows a chart last read from one view: other bins, ranges, bandwidths or charts
161/// over the same columns draw from them, describing one sample. Shared with the reading
162/// worker; the view's owner starts a new one when the view changes.
163#[derive(Clone, Default)]
164pub struct HeldRows(Arc<Mutex<Holding>>);
165
166#[derive(Default)]
167struct Holding {
168    rows: Option<Held>,
169    /// Rows per category, from a count of the whole view: exact whatever the sample
170    /// size, so another order or size draws from them rather than counting again.
171    counts: Vec<HeldCounts>,
172}
173
174struct Held {
175    limit: Option<usize>,
176    seed: u64,
177    df: DataFrame,
178    rows: RowsRead,
179}
180
181struct HeldCounts {
182    category: String,
183    counted: Counted,
184}
185
186impl std::fmt::Debug for HeldRows {
187    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
188        f.write_str("HeldRows")
189    }
190}
191
192/// What a chart read: the rows the table has, and how many of them were sampled when
193/// that was fewer.
194#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
195pub struct RowsRead {
196    pub total_rows: usize,
197    pub sample_size: Option<usize>,
198    /// For a line over more rows than its sample size, the steps along X each drawn as
199    /// their low and high value instead of sampling (see [`prepare_chart_data`]).
200    pub envelope_steps: Option<usize>,
201    /// The seed the sample was drawn with, when it is a sample.
202    pub seed: Option<u64>,
203}
204
205/// Which values a histogram, box plot or KDE draws. Outliers far from the body squash
206/// it into a bin or two; a percentile range leaves them out and says how many.
207#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
208#[serde(rename_all = "snake_case")]
209pub enum ValueRange {
210    #[default]
211    All,
212    /// The 1st to the 99th percentile.
213    Percentile1To99,
214}
215
216impl ValueRange {
217    pub const ALL: [Self; 2] = [Self::All, Self::Percentile1To99];
218
219    pub fn label(self) -> &'static str {
220        match self {
221            Self::All => "All",
222            Self::Percentile1To99 => "p1-p99",
223        }
224    }
225
226    fn quantiles(self) -> Option<(f64, f64)> {
227        match self {
228            Self::All => None,
229            Self::Percentile1To99 => Some((0.01, 0.99)),
230        }
231    }
232}
233
234/// A range that left values out: which, and how many.
235#[derive(Clone, Copy, Debug, PartialEq, Eq)]
236pub struct Clipped {
237    pub range: ValueRange,
238    pub outside: usize,
239}
240
241/// The chart's notes under the plot, one line each: that it is a sample (with its
242/// seed) and how many values a range left out. Empty when it shows everything.
243pub fn chart_notes(rows: &RowsRead, clipped: Option<&Clipped>, middot: &str) -> Vec<String> {
244    let mut notes = Vec::new();
245    if let Some(steps) = rows.envelope_steps {
246        notes.push(format!(
247            "min and max of {} rows in {} steps",
248            crate::home::discover::format_rows(rows.total_rows),
249            crate::numfmt::group_chrome(steps)
250        ));
251    }
252    if let Some(n) = rows.sample_size {
253        let mut note = format!(
254            "sample of {} of {} rows",
255            crate::numfmt::group_chrome(n),
256            crate::home::discover::format_rows(rows.total_rows)
257        );
258        if let Some(seed) = rows.seed {
259            note.push_str(&format!(" {middot} seed {seed}"));
260        }
261        notes.push(note);
262    }
263    if let Some(clipped) = clipped {
264        let noun = if clipped.outside == 1 {
265            "value"
266        } else {
267            "values"
268        };
269        notes.push(format!(
270            "{} {noun} outside {}",
271            crate::numfmt::group_chrome(clipped.outside),
272            clipped.range.label()
273        ));
274    }
275    notes
276}
277
278/// Read `columns` (each once) through the analysis sampler: every row up to the limit,
279/// past it a seeded spread (runs of one Parquet/IPC file, or one streamed pass), never
280/// the head. Held rows of the same size and seed are reused if they have the columns;
281/// otherwise held columns are read along so returning to one costs nothing.
282fn read_columns(
283    lf: &LazyFrame,
284    columns: &[&str],
285    sampling: &ChartSampling,
286) -> Result<(DataFrame, RowsRead)> {
287    let mut unique: Vec<PlSmallStr> = Vec::with_capacity(columns.len());
288    for c in columns {
289        if !unique.iter().any(|u| u == c) {
290            unique.push((*c).into());
291        }
292    }
293    let mut holding = sampling.held.0.lock().unwrap_or_else(|e| e.into_inner());
294    if let Some(h) = holding
295        .rows
296        .as_ref()
297        .filter(|h| h.limit == sampling.limit && h.seed == sampling.seed)
298    {
299        if unique.iter().all(|c| h.df.column(c).is_ok()) {
300            return Ok((h.df.select(unique.iter().cloned())?, h.rows));
301        }
302        for c in h.df.get_column_names() {
303            if !unique.contains(c) {
304                unique.push(c.clone());
305            }
306        }
307    }
308    let lf = lf
309        .clone()
310        .select(unique.iter().map(|c| col(c.clone())).collect::<Vec<_>>());
311    let read = crate::analysis::sampling::analysis_rows(
312        &lf,
313        sampling.limit,
314        sampling.known_total,
315        sampling.seed,
316        sampling.streaming,
317    )?;
318    let rows = RowsRead {
319        total_rows: read.total_rows,
320        sample_size: read.sample_size,
321        envelope_steps: None,
322        seed: read.sample_size.map(|_| sampling.seed),
323    };
324    holding.rows = Some(Held {
325        limit: sampling.limit,
326        seed: sampling.seed,
327        df: read.df.clone(),
328        rows,
329    });
330    Ok((read.df, rows))
331}
332
333/// A column's values as `f64`, one per row; null, NaN and infinities are `None`.
334fn f64_values(df: &DataFrame, column: &str) -> Result<Vec<Option<f64>>> {
335    let cast = df.column(column)?.cast(&DataType::Float64)?;
336    Ok(cast
337        .f64()?
338        .iter()
339        .map(|v| v.filter(|v| v.is_finite()))
340        .collect())
341}
342
343/// X as `f64`, one per row: numbers as they are, temporal types as their ordinal (see
344/// [`XAxisTemporalKind`]).
345fn x_values(df: &DataFrame, column: &str, dtype: &DataType) -> Result<Vec<Option<f64>>> {
346    match dtype {
347        DataType::Datetime(_, _) | DataType::Date | DataType::Time => {
348            let ordinal = df.column(column)?.cast(&DataType::Int64)?;
349            Ok(ordinal.i64()?.iter().map(|v| v.map(|v| v as f64)).collect())
350        }
351        _ => f64_values(df, column),
352    }
353}
354
355/// X's type in `schema`, or an error naming it.
356fn x_dtype<'a>(schema: &'a Schema, x: &str) -> Result<&'a DataType> {
357    schema
358        .get(x)
359        .ok_or_else(|| color_eyre::eyre::eyre!("x column '{}' not in schema", x))
360}
361
362/// The rows with an X, as `(x, row)`, in X order; rows sharing an X keep their table
363/// order.
364fn x_order(df: &DataFrame, x: &str, dtype: &DataType) -> Result<Vec<(f64, usize)>> {
365    let mut order: Vec<(f64, usize)> = x_values(df, x, dtype)?
366        .into_iter()
367        .enumerate()
368        .filter_map(|(i, x)| x.map(|x| (x, i)))
369        .collect();
370    order.sort_by(|a, b| a.0.total_cmp(&b.0));
371    Ok(order)
372}
373
374/// Series built a point at a time in X order, each breaking its line after a gap:
375/// a point with no Y.
376struct SeriesBuilder {
377    series: Vec<Vec<(f64, f64)>>,
378    breaks: Vec<Vec<usize>>,
379    gap: Vec<bool>,
380}
381
382impl SeriesBuilder {
383    fn new(n: usize) -> Self {
384        Self {
385            series: vec![Vec::new(); n],
386            breaks: vec![Vec::new(); n],
387            gap: vec![false; n],
388        }
389    }
390
391    fn push(&mut self, s: usize, x: f64, y: Option<f64>) {
392        match y {
393            Some(y) => {
394                if self.gap[s] && !self.series[s].is_empty() {
395                    self.breaks[s].push(self.series[s].len());
396                }
397                self.gap[s] = false;
398                self.series[s].push((x, y));
399            }
400            None => self.gap[s] = true,
401        }
402    }
403}
404
405/// A column's values as text, as a category is labeled.
406fn text_labels(df: &DataFrame, column: &str) -> Result<Series> {
407    Ok(crate::past_calendar::cast_text(
408        df.column(column)?.as_materialized_series(),
409        CastOptions::NonStrict,
410    )?)
411}
412
413/// A count column as `u64`s, null as none.
414fn counts_of(df: &DataFrame, column: &str) -> Result<Vec<u64>> {
415    Ok(df
416        .column(column)?
417        .cast(&DataType::UInt64)?
418        .u64()?
419        .iter()
420        .map(|n| n.unwrap_or(0))
421        .collect())
422}
423
424impl RowsRead {
425    /// Every one of `total_rows` rows, read whole.
426    pub fn every(total_rows: usize) -> Self {
427        Self {
428            total_rows,
429            ..Self::default()
430        }
431    }
432}
433
434/// The X column's range alone, for a line or scatter chart's axes before any Y.
435#[derive(Debug, Clone)]
436pub struct ChartXRangeResult {
437    pub x_min: f64,
438    pub x_max: f64,
439    pub x_axis_kind: XAxisTemporalKind,
440    pub rows: RowsRead,
441}
442
443/// Loads only the x column and returns its min/max (for axis display when no y is selected).
444pub fn prepare_chart_x_range(
445    lf: &LazyFrame,
446    schema: &Schema,
447    x_column: &str,
448    sampling: &ChartSampling,
449) -> Result<ChartXRangeResult> {
450    let x_dtype = x_dtype(schema, x_column)?;
451    let x_axis_kind = x_axis_temporal_kind(x_dtype);
452    let (df, rows) = read_columns(lf, &[x_column], sampling)?;
453    let (x_min, x_max) = x_values(&df, x_column, x_dtype)?
454        .into_iter()
455        .flatten()
456        .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), x| {
457            (lo.min(x), hi.max(x))
458        });
459    let (x_min, x_max) = if x_max >= x_min {
460        (x_min, x_max)
461    } else {
462        (0.0, 1.0)
463    };
464    Ok(ChartXRangeResult {
465        x_min,
466        x_max,
467        x_axis_kind,
468        rows,
469    })
470}
471
472/// Result of preparing chart data: series points and x-axis kind for label formatting.
473pub struct ChartDataResult {
474    /// One per y column, in X order.
475    pub series: Vec<Vec<(f64, f64)>>,
476    /// Per series, the indices of `series` where a line starts again after a gap: a row
477    /// whose y is null is left out of that series alone, and a line does not bridge it.
478    pub breaks: Vec<Vec<usize>>,
479    pub x_axis_kind: XAxisTemporalKind,
480    pub rows: RowsRead,
481}
482
483/// A series split at its breaks: the runs a line joins.
484pub fn segments<'a>(points: &'a [(f64, f64)], breaks: &[usize]) -> Vec<&'a [(f64, f64)]> {
485    let mut out = Vec::with_capacity(breaks.len() + 1);
486    let mut start = 0;
487    for &b in breaks {
488        if b > start && b <= points.len() {
489            out.push(&points[start..b]);
490            start = b;
491        }
492    }
493    if start < points.len() {
494        out.push(&points[start..]);
495    }
496    out
497}
498
499/// Histogram bin (center and count).
500#[derive(Clone, Debug)]
501pub struct HistogramBin {
502    pub center: f64,
503    pub count: f64,
504}
505
506/// One group's bins of a histogram split by a color: its count (or share) per bin,
507/// on the bins of the whole.
508#[derive(Clone, Debug)]
509pub struct HistogramGroup {
510    pub name: String,
511    pub counts: Vec<f64>,
512}
513
514/// Histogram data for a single column.
515#[derive(Clone, Debug)]
516pub struct HistogramData {
517    pub column: String,
518    /// Every row's bins; with `share`, each bin's share of the rows.
519    pub bins: Vec<HistogramBin>,
520    /// Per color group, on the same bins: drawn as outlines over one another.
521    pub groups: Vec<HistogramGroup>,
522    /// The last group is Other: every value of the color without a group of its own.
523    pub other: bool,
524    /// Each bin is a share of its group's rows (or of all rows), not a count.
525    pub share: bool,
526    pub x_min: f64,
527    pub x_max: f64,
528    pub max_count: f64,
529    pub rows: RowsRead,
530    pub clipped: Option<Clipped>,
531}
532
533/// KDE series and bounds.
534#[derive(Clone, Debug)]
535pub struct KdeSeries {
536    pub name: String,
537    pub points: Vec<(f64, f64)>,
538}
539
540#[derive(Clone, Debug)]
541pub struct KdeData {
542    pub series: Vec<KdeSeries>,
543    /// The last series is Other: every value of the color without a series of its own.
544    pub other: bool,
545    pub x_min: f64,
546    pub x_max: f64,
547    pub y_max: f64,
548    pub rows: RowsRead,
549    pub clipped: Option<Clipped>,
550}
551
552/// Box plot stats for a column.
553#[derive(Clone, Debug)]
554pub struct BoxPlotStats {
555    pub name: String,
556    pub min: f64,
557    pub q1: f64,
558    pub median: f64,
559    pub q3: f64,
560    pub max: f64,
561}
562
563/// A box and its whiskers as line segments, the box `half` either side of `center`
564/// and the caps `cap` either side; Y in data values, X in whatever unit `center` is.
565pub struct BoxMarks {
566    /// The box, corner to corner and back to the first.
567    pub outline: [(f64, f64); 5],
568    pub median: [(f64, f64); 2],
569    /// Minimum to the first quartile, and the third quartile to the maximum.
570    pub low: [(f64, f64); 2],
571    pub high: [(f64, f64); 2],
572    pub low_cap: [(f64, f64); 2],
573    pub high_cap: [(f64, f64); 2],
574}
575
576impl BoxPlotStats {
577    pub fn marks(&self, center: f64, half: f64, cap: f64) -> BoxMarks {
578        let (left, right) = (center - half, center + half);
579        BoxMarks {
580            outline: [
581                (left, self.q1),
582                (right, self.q1),
583                (right, self.q3),
584                (left, self.q3),
585                (left, self.q1),
586            ],
587            median: [(left, self.median), (right, self.median)],
588            low: [(center, self.min), (center, self.q1)],
589            high: [(center, self.q3), (center, self.max)],
590            low_cap: [(center - cap, self.min), (center + cap, self.min)],
591            high_cap: [(center - cap, self.max), (center + cap, self.max)],
592        }
593    }
594}
595
596#[derive(Clone, Debug)]
597pub struct BoxPlotData {
598    pub stats: Vec<BoxPlotStats>,
599    pub y_min: f64,
600    pub y_max: f64,
601    pub rows: RowsRead,
602    pub clipped: Option<Clipped>,
603    /// One box per category: how many categories there are, of which the largest
604    /// have a box. 0 for a box per column.
605    pub of: usize,
606}
607
608impl HistogramData {
609    /// Each group's bins as the outline of its bars: up the left edge of each bin,
610    /// across its top, and down at the end.
611    pub fn step_outlines(&self) -> Vec<Vec<(f64, f64)>> {
612        let n = self.bins.len().max(1);
613        let width = (self.x_max - self.x_min) / n as f64;
614        self.groups
615            .iter()
616            .map(|group| {
617                let mut points = vec![(self.x_min, 0.0)];
618                for (i, &count) in group.counts.iter().enumerate() {
619                    let x0 = self.x_min + i as f64 * width;
620                    points.push((x0, count));
621                    points.push((x0 + width, count));
622                }
623                points.push((self.x_max, 0.0));
624                points
625            })
626            .collect()
627    }
628}
629
630/// Where Other is among `n` series: the last, when there is one.
631pub fn other_at(other: bool, n: usize) -> Option<usize> {
632    (other && n > 0).then(|| n - 1)
633}
634
635/// The order `n` series are drawn in: Other first, under the series drawn over it.
636pub fn drawing_order(n: usize, other: Option<usize>) -> impl Iterator<Item = usize> {
637    other
638        .filter(|&o| o < n)
639        .into_iter()
640        .chain((0..n).filter(move |&i| Some(i) != other))
641}
642
643/// Heatmap data for two numeric columns.
644#[derive(Clone, Debug)]
645pub struct HeatmapData {
646    pub x_column: String,
647    pub y_column: String,
648    pub x_min: f64,
649    pub x_max: f64,
650    pub y_min: f64,
651    pub y_max: f64,
652    pub x_bins: usize,
653    pub y_bins: usize,
654    pub counts: Vec<Vec<f64>>,
655    pub max_count: f64,
656    pub rows: RowsRead,
657}
658
659/// XY series from the LazyFrame, X cast to f64 (temporals as ordinals). Nulls drop per
660/// series (a null X drops the row, a null Y breaks that series' line). Points in X
661/// order, ties in table order. With `envelope`, a view larger than the sample size is
662/// cut into steps, each drawing its low and high Y, keeping peaks a sample would miss;
663/// it reads twice, so only where [`ChartSampling::full_passes`] allows, stopping on
664/// [`ChartSampling::cancel`].
665pub fn prepare_chart_data(
666    lf: &LazyFrame,
667    schema: &Schema,
668    x_column: &str,
669    y_columns: &[String],
670    sampling: &ChartSampling,
671    envelope: bool,
672) -> Result<ChartDataResult> {
673    if y_columns.is_empty() {
674        return Ok(ChartDataResult {
675            series: Vec::new(),
676            breaks: Vec::new(),
677            x_axis_kind: XAxisTemporalKind::Numeric,
678            rows: RowsRead::default(),
679        });
680    }
681
682    let x_dtype = x_dtype(schema, x_column)?;
683    let x_axis_kind = x_axis_temporal_kind(x_dtype);
684
685    let mut counted = None;
686    if envelope
687        && sampling.full_passes
688        && let Some(limit) = sampling.limit.filter(|&n| n > 0)
689        && sampling.known_total.is_none_or(|n| n > limit)
690    {
691        match envelope_series(lf, x_column, x_dtype, y_columns, limit, sampling)? {
692            Envelope::Drawn {
693                series,
694                breaks,
695                rows,
696                steps,
697            } => {
698                return Ok(ChartDataResult {
699                    series,
700                    breaks,
701                    x_axis_kind,
702                    rows: RowsRead {
703                        total_rows: rows,
704                        sample_size: None,
705                        envelope_steps: Some(steps),
706                        seed: None,
707                    },
708                });
709            }
710            // The first pass counted the view: the sample takes it whole.
711            Envelope::Fits(rows) => counted = Some(rows),
712        }
713    }
714    let counted_sampling;
715    let sampling = match counted {
716        Some(rows) => {
717            counted_sampling = ChartSampling {
718                known_total: Some(rows),
719                ..sampling.clone()
720            };
721            &counted_sampling
722        }
723        None => sampling,
724    };
725
726    let mut columns = vec![x_column];
727    columns.extend(y_columns.iter().map(String::as_str));
728    let (df, rows) = read_columns(lf, &columns, sampling)?;
729
730    let order = x_order(&df, x_column, x_dtype)?;
731    let mut built = SeriesBuilder::new(y_columns.len());
732    for (s, y_column) in y_columns.iter().enumerate() {
733        let ys = f64_values(&df, y_column)?;
734        for &(x, i) in &order {
735            built.push(s, x, ys[i]);
736        }
737    }
738
739    Ok(ChartDataResult {
740        series: built.series,
741        breaks: built.breaks,
742        x_axis_kind,
743        rows,
744    })
745}
746
747/// What [`envelope_series`] found.
748enum Envelope {
749    /// Per Y column, its points and where its line breaks, from `rows` rows in
750    /// `steps` steps.
751    Drawn {
752        series: Vec<Vec<(f64, f64)>>,
753        breaks: Vec<Vec<usize>>,
754        rows: usize,
755        steps: usize,
756    },
757    /// The view is no more rows than the sample size, this many: no envelope.
758    Fits(usize),
759}
760
761/// What a pass stopped by [`until_cancelled`] fails with.
762const ENVELOPE_CANCELLED: &str = "chart cancelled";
763
764/// `e`, failing the query once `cancel` is set: a streamed pass stops at its next
765/// morsel rather than reading on for a chart nobody waits for.
766fn until_cancelled(e: Expr, cancel: &Arc<AtomicBool>) -> Expr {
767    let cancel = Arc::clone(cancel);
768    e.map(
769        move |c: Column| {
770            polars_ensure!(!cancel.load(Ordering::Relaxed), ComputeError: ENVELOPE_CANCELLED);
771            Ok(c)
772        },
773        |_, field| Ok(field.clone()),
774    )
775}
776
777/// Collect a pass of the envelope, streamed whatever the setting: it holds a few
778/// numbers per step. A pass stopped by `cancel` is an error that says so.
779fn envelope_pass(lf: LazyFrame, cancel: &Arc<AtomicBool>) -> Result<DataFrame> {
780    crate::analysis::statistics::collect_lazy(lf, true).map_err(|e| {
781        if cancel.load(Ordering::Relaxed) {
782            color_eyre::eyre::eyre!(ENVELOPE_CANCELLED)
783        } else {
784            e.into()
785        }
786    })
787}
788
789/// Per Y column, half `limit` steps along X, each its lowest and highest finite Y at the
790/// step's lowest X, in X order: two streamed passes (rows and X bounds, then a group per
791/// step). Rows without X are dropped; a step without a series' value breaks its line.
792fn envelope_series(
793    lf: &LazyFrame,
794    x_column: &str,
795    x_dtype: &DataType,
796    y_columns: &[String],
797    limit: usize,
798    sampling: &ChartSampling,
799) -> Result<Envelope> {
800    let cancel = &sampling.cancel;
801    // Temporal X as its ordinal, as `x_values` reads it.
802    let x = match x_dtype {
803        DataType::Datetime(_, _) | DataType::Date | DataType::Time | DataType::Duration(_) => {
804            col(x_column).cast(DataType::Int64).cast(DataType::Float64)
805        }
806        _ => col(x_column).cast(DataType::Float64),
807    };
808    // Not finite is null, so the aggregations below are a plain min and max, which
809    // stream; a filter inside them does not.
810    let finite = |e: Expr| {
811        when(e.clone().is_finite())
812            .then(e)
813            .otherwise(lit(NULL).cast(DataType::Float64))
814    };
815    let x = finite(until_cancelled(x, cancel)).alias("__x");
816    let bounds = envelope_pass(
817        lf.clone().select([
818            len().alias("rows"),
819            x.clone().min().alias("lo"),
820            x.clone().max().alias("hi"),
821        ]),
822        cancel,
823    )?;
824    let rows = bounds
825        .column("rows")?
826        .cast(&DataType::UInt64)?
827        .u64()?
828        .get(0)
829        .unwrap_or(0) as usize;
830    if rows <= limit {
831        return Ok(Envelope::Fits(rows));
832    }
833    let steps = (limit / 2).max(1);
834    let n = y_columns.len();
835    let drawn = |series, breaks| Envelope::Drawn {
836        series,
837        breaks,
838        rows,
839        steps,
840    };
841    let bound = |name: &str| -> Result<Option<f64>> { Ok(bounds.column(name)?.f64()?.get(0)) };
842    let (Some(lo), Some(hi)) = (bound("lo")?, bound("hi")?) else {
843        return Ok(drawn(vec![Vec::new(); n], vec![Vec::new(); n]));
844    };
845    let per_x = if hi > lo {
846        steps as f64 / (hi - lo)
847    } else {
848        0.0
849    };
850    let lf = lf
851        .clone()
852        .select(
853            std::iter::once(x)
854                .chain(y_columns.iter().enumerate().map(|(i, y)| {
855                    finite(col(y.as_str()).cast(DataType::Float64)).alias(format!("__y{i}"))
856                }))
857                .collect::<Vec<_>>(),
858        )
859        // A row filter, not a filter inside the select: that would leave X shorter than
860        // the Y columns beside it.
861        .filter(col("__x").is_not_null());
862    let step = ((col("__x") - lit(lo)) * lit(per_x))
863        .floor()
864        .cast(DataType::Int64)
865        .clip(lit(0i64), lit(steps as i64 - 1))
866        .alias("__step");
867    let mut aggs = vec![col("__x").min()];
868    for i in 0..n {
869        let y = col(format!("__y{i}"));
870        aggs.push(y.clone().min().alias(format!("__lo{i}")));
871        aggs.push(y.max().alias(format!("__hi{i}")));
872    }
873    let df = envelope_pass(
874        lf.group_by([step])
875            .agg(aggs)
876            .sort(["__step"], Default::default()),
877        cancel,
878    )?;
879    let xs = df.column("__x")?.f64()?.clone();
880    let mut series = Vec::with_capacity(n);
881    let mut breaks = Vec::with_capacity(n);
882    for i in 0..n {
883        let lows = df.column(&format!("__lo{i}"))?.f64()?.clone();
884        let highs = df.column(&format!("__hi{i}"))?.f64()?.clone();
885        let mut points = Vec::with_capacity(xs.len() * 2);
886        let mut starts = Vec::new();
887        let mut gap = false;
888        for ((x, low), high) in xs.iter().zip(lows.iter()).zip(highs.iter()) {
889            let (Some(x), Some(low), Some(high)) = (x, low, high) else {
890                gap = true;
891                continue;
892            };
893            if gap && !points.is_empty() {
894                starts.push(points.len());
895            }
896            gap = false;
897            points.push((x, low));
898            if high != low {
899                points.push((x, high));
900            }
901        }
902        series.push(points);
903        breaks.push(starts);
904    }
905    Ok(drawn(series, breaks))
906}
907
908/// Each column's finite values, read in one pass; nulls are dropped per column.
909fn read_values(
910    lf: &LazyFrame,
911    column: &str,
912    sampling: &ChartSampling,
913) -> Result<(Vec<f64>, RowsRead)> {
914    let (df, rows) = read_columns(lf, &[column], sampling)?;
915    Ok((
916        f64_values(&df, column)?.into_iter().flatten().collect(),
917        rows,
918    ))
919}
920
921/// Sort `values` and keep those inside `range`; returns how many were left out.
922fn sort_and_clip(values: &mut Vec<f64>, range: ValueRange) -> usize {
923    values.sort_by(f64::total_cmp);
924    let Some((low, high)) = range.quantiles() else {
925        return 0;
926    };
927    if values.is_empty() {
928        return 0;
929    }
930    let (low, high) = (quantile(values, low), quantile(values, high));
931    let before = values.len();
932    values.retain(|v| (low..=high).contains(v));
933    before - values.len()
934}
935
936fn clipped(range: ValueRange, outside: usize) -> Option<Clipped> {
937    (range != ValueRange::All).then_some(Clipped { range, outside })
938}
939
940/// A histogram of `column`, split by `color` into groups on the same bins. With `share`,
941/// a bin is its share of its group's rows so groups of different sizes compare. The
942/// range is the whole column's, so every group clips alike.
943pub fn prepare_histogram_by(
944    lf: &LazyFrame,
945    column: &str,
946    bins: usize,
947    range: ValueRange,
948    share: bool,
949    color: Option<ColorSplit<'_>>,
950    sampling: &ChartSampling,
951) -> Result<HistogramData> {
952    let (values, rows) = read_split(lf, column, color, sampling)?;
953    let mut all: Vec<f64> = values.iter().map(|(v, _)| *v).collect();
954    let outside = sort_and_clip(&mut all, range);
955    let clipped = clipped(range, outside);
956    let mut data = HistogramData {
957        column: column.to_string(),
958        bins: Vec::new(),
959        groups: Vec::new(),
960        other: false,
961        share,
962        x_min: 0.0,
963        x_max: 1.0,
964        max_count: 0.0,
965        rows,
966        clipped,
967    };
968    let (Some(&lo), Some(&hi)) = (all.first(), all.last()) else {
969        return Ok(data);
970    };
971    let span = hi - lo;
972    let bin_count = if span <= f64::EPSILON { 1 } else { bins.max(1) };
973    let bin_width = if span <= f64::EPSILON {
974        1.0
975    } else {
976        span / bin_count as f64
977    };
978    (data.x_min, data.x_max) = if span <= f64::EPSILON {
979        (lo - 0.5, hi + 0.5)
980    } else {
981        (lo, hi)
982    };
983    let bin_of = |v: f64| {
984        if span <= f64::EPSILON {
985            0
986        } else {
987            (((v - lo) / bin_width).floor().max(0.0) as usize).min(bin_count - 1)
988        }
989    };
990    let groups = color.map_or(0, |c| c.series());
991    let mut total = vec![0.0_f64; bin_count];
992    let mut by_group = vec![vec![0.0_f64; bin_count]; groups];
993    for (v, group) in values {
994        // The range is the whole view's; the bins count the groups drawn.
995        if !(lo..=hi).contains(&v) || (color.is_some() && group.is_none()) {
996            continue;
997        }
998        let bin = bin_of(v);
999        total[bin] += 1.0;
1000        if let Some(g) = group {
1001            by_group[g][bin] += 1.0;
1002        }
1003    }
1004    let as_share = |counts: &mut Vec<f64>| {
1005        let n: f64 = counts.iter().sum();
1006        if share && n > 0.0 {
1007            counts.iter_mut().for_each(|c| *c /= n);
1008        }
1009    };
1010    as_share(&mut total);
1011    by_group.iter_mut().for_each(as_share);
1012    let center = |i: usize| {
1013        if span <= f64::EPSILON {
1014            lo
1015        } else {
1016            lo + (i as f64 + 0.5) * bin_width
1017        }
1018    };
1019    data.bins = total
1020        .iter()
1021        .enumerate()
1022        .map(|(i, &count)| HistogramBin {
1023            center: center(i),
1024            count,
1025        })
1026        .collect();
1027    let max = |counts: &[f64]| counts.iter().copied().fold(0.0_f64, f64::max);
1028    if let Some(color) = color {
1029        data.groups = color
1030            .names()
1031            .into_iter()
1032            .zip(by_group)
1033            .map(|(name, counts)| HistogramGroup { name, counts })
1034            .collect();
1035        data.other = color.other;
1036        data.max_count = data
1037            .groups
1038            .iter()
1039            .map(|g| max(&g.counts))
1040            .fold(0.0, f64::max);
1041    } else {
1042        data.max_count = max(&total);
1043    }
1044    Ok(data)
1045}
1046
1047fn quantile(sorted: &[f64], q: f64) -> f64 {
1048    if sorted.is_empty() {
1049        return 0.0;
1050    }
1051    let n = sorted.len();
1052    if n == 1 {
1053        return sorted[0];
1054    }
1055    let pos = q.clamp(0.0, 1.0) * (n as f64 - 1.0);
1056    let idx = pos.floor() as usize;
1057    let next = pos.ceil() as usize;
1058    if idx == next {
1059        sorted[idx]
1060    } else {
1061        let lower = sorted[idx];
1062        let upper = sorted[next];
1063        let weight = pos - idx as f64;
1064        lower + (upper - lower) * weight
1065    }
1066}
1067
1068/// The five numbers of sorted `values`, or `None` when there are none.
1069fn box_stats(name: String, values: &[f64]) -> Option<BoxPlotStats> {
1070    let (min, max) = (*values.first()?, *values.last()?);
1071    Some(BoxPlotStats {
1072        name,
1073        min,
1074        q1: quantile(values, 0.25),
1075        median: quantile(values, 0.5),
1076        q3: quantile(values, 0.75),
1077        max,
1078    })
1079}
1080
1081/// Box plot data from stats, its bounds taken from them.
1082fn box_data(stats: Vec<BoxPlotStats>, rows: RowsRead, clipped: Option<Clipped>) -> BoxPlotData {
1083    let mut y_min = stats.iter().map(|s| s.min).fold(f64::INFINITY, f64::min);
1084    let mut y_max = stats
1085        .iter()
1086        .map(|s| s.max)
1087        .fold(f64::NEG_INFINITY, f64::max);
1088    if stats.is_empty() {
1089        (y_min, y_max) = (0.0, 1.0);
1090    } else if y_max <= y_min {
1091        y_max = y_min + 1.0;
1092    }
1093    BoxPlotData {
1094        stats,
1095        y_min,
1096        y_max,
1097        rows,
1098        clipped,
1099        of: 0,
1100    }
1101}
1102
1103/// The box plot of one numeric column.
1104pub fn prepare_box_plot_data(
1105    lf: &LazyFrame,
1106    column: &str,
1107    range: ValueRange,
1108    sampling: &ChartSampling,
1109) -> Result<BoxPlotData> {
1110    let (mut values, rows) = read_values(lf, column, sampling)?;
1111    let outside = sort_and_clip(&mut values, range);
1112    let stats = box_stats(column.to_string(), &values).into_iter().collect();
1113    Ok(box_data(stats, rows, clipped(range, outside)))
1114}
1115
1116/// Prepare one box of `column` per group of `by`: a box per category. The range is
1117/// each group's own, as each box describes its group.
1118pub fn prepare_box_by(
1119    lf: &LazyFrame,
1120    column: &str,
1121    by: ColorSplit<'_>,
1122    range: ValueRange,
1123    sampling: &ChartSampling,
1124) -> Result<BoxPlotData> {
1125    let (values, rows) = read_split(lf, column, Some(by), sampling)?;
1126    let mut groups = vec![Vec::new(); by.groups.len()];
1127    for (v, group) in values {
1128        if let Some(g) = group {
1129            groups[g].push(v);
1130        }
1131    }
1132    let mut outside = 0;
1133    let mut stats = Vec::new();
1134    for (name, mut values) in by.groups.iter().zip(groups) {
1135        outside += sort_and_clip(&mut values, range);
1136        stats.extend(box_stats(group_label(name), &values));
1137    }
1138    Ok(box_data(stats, rows, clipped(range, outside)))
1139}
1140
1141fn kde_bandwidth(values: &[f64]) -> f64 {
1142    if values.len() <= 1 {
1143        return 1.0;
1144    }
1145    let n = values.len() as f64;
1146    let mean = values.iter().sum::<f64>() / n;
1147    let var = values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n;
1148    let std = var.sqrt();
1149    if std <= f64::EPSILON {
1150        return 1.0;
1151    }
1152    1.06 * std * n.powf(-0.2)
1153}
1154
1155/// The density of sorted `values` at 200 points, from three bandwidths below the
1156/// least to three above the greatest; `None` when there are none.
1157fn kde_series(name: String, values: &[f64], bandwidth_factor: f64) -> Option<KdeSeries> {
1158    let (min, max) = (*values.first()?, *values.last()?);
1159    let bandwidth = (kde_bandwidth(values) * bandwidth_factor).max(f64::EPSILON);
1160    let x_start = min - 3.0 * bandwidth;
1161    let x_end = max + 3.0 * bandwidth;
1162    let samples = 200_usize;
1163    let step = (x_end - x_start) / (samples.saturating_sub(1).max(1) as f64);
1164    let inv = 1.0 / ((values.len() as f64) * bandwidth * (2.0 * PI).sqrt());
1165    let points = (0..samples)
1166        .map(|i| {
1167            let x = x_start + i as f64 * step;
1168            let sum: f64 = values
1169                .iter()
1170                .map(|&v| {
1171                    let u = (x - v) / bandwidth;
1172                    (-0.5 * u * u).exp()
1173                })
1174                .sum();
1175            (x, inv * sum)
1176        })
1177        .collect();
1178    Some(KdeSeries { name, points })
1179}
1180
1181/// KDE data from its series, its bounds taken from them.
1182fn kde_data(series: Vec<KdeSeries>, rows: RowsRead, clipped: Option<Clipped>) -> KdeData {
1183    let points = || series.iter().flat_map(|s| s.points.iter());
1184    let mut x_min = points().map(|p| p.0).fold(f64::INFINITY, f64::min);
1185    let mut x_max = points().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max);
1186    let mut y_max = points().map(|p| p.1).fold(f64::NEG_INFINITY, f64::max);
1187    if series.is_empty() {
1188        (x_min, x_max, y_max) = (0.0, 1.0, 1.0);
1189    }
1190    if x_max <= x_min {
1191        x_max = x_min + 1.0;
1192    }
1193    if y_max <= 0.0 {
1194        y_max = 1.0;
1195    }
1196    KdeData {
1197        series,
1198        other: false,
1199        x_min,
1200        x_max,
1201        y_max,
1202        rows,
1203        clipped,
1204    }
1205}
1206
1207/// The density curve of one numeric column.
1208pub fn prepare_kde_data(
1209    lf: &LazyFrame,
1210    column: &str,
1211    bandwidth_factor: f64,
1212    range: ValueRange,
1213    sampling: &ChartSampling,
1214) -> Result<KdeData> {
1215    let (mut values, rows) = read_values(lf, column, sampling)?;
1216    let outside = sort_and_clip(&mut values, range);
1217    let series = kde_series(column.to_string(), &values, bandwidth_factor)
1218        .into_iter()
1219        .collect();
1220    Ok(kde_data(series, rows, clipped(range, outside)))
1221}
1222
1223/// Prepare one density curve of `column` per group of `color`. The range is the
1224/// whole column's, so every curve is clipped alike.
1225pub fn prepare_kde_by(
1226    lf: &LazyFrame,
1227    column: &str,
1228    bandwidth_factor: f64,
1229    range: ValueRange,
1230    color: ColorSplit<'_>,
1231    sampling: &ChartSampling,
1232) -> Result<KdeData> {
1233    let (values, rows) = read_split(lf, column, Some(color), sampling)?;
1234    let mut all: Vec<f64> = values.iter().map(|(v, _)| *v).collect();
1235    let outside = sort_and_clip(&mut all, range);
1236    let (lo, hi) = match (all.first(), all.last()) {
1237        (Some(&lo), Some(&hi)) => (lo, hi),
1238        _ => (f64::INFINITY, f64::NEG_INFINITY),
1239    };
1240    let mut groups = vec![Vec::new(); color.series()];
1241    for (v, group) in values {
1242        if let Some(g) = group
1243            && (lo..=hi).contains(&v)
1244        {
1245            groups[g].push(v);
1246        }
1247    }
1248    let last = color.series().saturating_sub(1);
1249    let mut other = false;
1250    let series = color
1251        .names()
1252        .into_iter()
1253        .zip(groups)
1254        .enumerate()
1255        .filter_map(|(i, (name, mut values))| {
1256            values.sort_by(f64::total_cmp);
1257            let series = kde_series(name, &values, bandwidth_factor)?;
1258            other = color.other && i == last;
1259            Some(series)
1260        })
1261        .collect();
1262    Ok(KdeData {
1263        other,
1264        ..kde_data(series, rows, clipped(range, outside))
1265    })
1266}
1267
1268/// Prepare heatmap data for two numeric columns. A row counts when both are present.
1269pub fn prepare_heatmap_data(
1270    lf: &LazyFrame,
1271    x_column: &str,
1272    y_column: &str,
1273    bins: usize,
1274    sampling: &ChartSampling,
1275) -> Result<HeatmapData> {
1276    let (df, rows) = read_columns(lf, &[x_column, y_column], sampling)?;
1277    let pairs: Vec<(f64, f64)> = f64_values(&df, x_column)?
1278        .into_iter()
1279        .zip(f64_values(&df, y_column)?)
1280        .filter_map(|(x, y)| Some((x?, y?)))
1281        .collect();
1282    let x_bins = bins.max(1);
1283    let y_bins = bins.max(1);
1284    if pairs.is_empty() {
1285        return Ok(HeatmapData {
1286            x_column: x_column.to_string(),
1287            y_column: y_column.to_string(),
1288            x_min: 0.0,
1289            x_max: 1.0,
1290            y_min: 0.0,
1291            y_max: 1.0,
1292            x_bins,
1293            y_bins,
1294            counts: vec![vec![0.0; x_bins]; y_bins],
1295            max_count: 0.0,
1296            rows,
1297        });
1298    }
1299    let mut x_min = f64::INFINITY;
1300    let mut x_max = f64::NEG_INFINITY;
1301    let mut y_min = f64::INFINITY;
1302    let mut y_max = f64::NEG_INFINITY;
1303    for (x, y) in &pairs {
1304        x_min = x_min.min(*x);
1305        x_max = x_max.max(*x);
1306        y_min = y_min.min(*y);
1307        y_max = y_max.max(*y);
1308    }
1309    if x_max <= x_min {
1310        x_max = x_min + 1.0;
1311    }
1312    if y_max <= y_min {
1313        y_max = y_min + 1.0;
1314    }
1315    let mut counts = vec![vec![0.0_f64; x_bins]; y_bins];
1316    let x_range = x_max - x_min;
1317    let y_range = y_max - y_min;
1318    for (x, y) in pairs {
1319        let xi =
1320            (((x - x_min) / x_range * x_bins as f64).floor().max(0.0) as usize).min(x_bins - 1);
1321        let yi =
1322            (((y - y_min) / y_range * y_bins as f64).floor().max(0.0) as usize).min(y_bins - 1);
1323        counts[yi][xi] += 1.0;
1324    }
1325    let max_count = counts
1326        .iter()
1327        .flat_map(|row| row.iter())
1328        .cloned()
1329        .fold(0.0_f64, f64::max);
1330    Ok(HeatmapData {
1331        x_column: x_column.to_string(),
1332        y_column: y_column.to_string(),
1333        x_min,
1334        x_max,
1335        y_min,
1336        y_max,
1337        x_bins,
1338        y_bins,
1339        counts,
1340        max_count,
1341        rows,
1342    })
1343}
1344
1345/// The order of a bar chart's bars, top to bottom.
1346#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
1347#[serde(rename_all = "snake_case")]
1348pub enum BarOrder {
1349    /// Largest value first.
1350    #[default]
1351    Value,
1352    /// The category's own order: text A to Z, numbers ascending, false before true.
1353    Label,
1354}
1355
1356impl BarOrder {
1357    pub const ALL: [Self; 2] = [Self::Value, Self::Label];
1358
1359    pub fn label(self) -> &'static str {
1360        match self {
1361            Self::Value => "Value",
1362            Self::Label => "Label",
1363        }
1364    }
1365}
1366
1367/// Most bars a bar chart prepares; the rest are counted, not drawn. More than a tall
1368/// terminal shows, few enough that an export stays legible.
1369pub const BAR_CAP: usize = 100;
1370
1371/// Column types a bar chart takes as its category: text, categories, booleans and
1372/// integers. A float or a timestamp is a measure, not a category.
1373pub fn is_category_dtype(dtype: &DataType) -> bool {
1374    matches!(
1375        dtype,
1376        DataType::String | DataType::Categorical(_, _) | DataType::Enum(_, _) | DataType::Boolean
1377    ) || dtype.is_integer()
1378}
1379
1380/// Categories a count keeps before it stops: past this the column is an identifier,
1381/// not a category, and a count per value would hold as much as the table.
1382pub const COUNT_CATEGORY_CAP: usize = 100_000;
1383
1384/// One bar: its category (`None` for null) and value. Split by color, a row of bars per
1385/// group (`None` for an empty group), valued at their total.
1386#[derive(Clone, Debug, PartialEq)]
1387pub struct Bar {
1388    pub label: Option<String>,
1389    pub value: f64,
1390    pub by_group: Vec<Option<f64>>,
1391}
1392
1393/// A bar chart: one bar per category, in order, up to [`BAR_CAP`].
1394#[derive(Clone, Debug)]
1395pub struct BarData {
1396    pub category: String,
1397    pub value_column: String,
1398    pub bars: Vec<Bar>,
1399    /// Categories past the cap, not in `bars`.
1400    pub more: usize,
1401    /// Categories left out because their value is null.
1402    pub no_value: usize,
1403    pub rows: RowsRead,
1404    /// The value column's type: an integer column prints whole numbers.
1405    pub value_dtype: DataType,
1406    /// Rows counted, when a count read past the sample size: the counts are of every
1407    /// row of the view, and the note says so.
1408    pub counted: Option<usize>,
1409    /// The color groups each bar is split into, in color order; empty when not split.
1410    pub groups: Vec<String>,
1411    /// The last group is Other: every value of the color without a group of its own.
1412    pub other: bool,
1413    /// What an aggregate read, said under the plot: `all 336,776 rows`.
1414    pub rows_note: Option<String>,
1415}
1416
1417impl BarData {
1418    /// The format the table prints the value column in, or plain where it shows the
1419    /// column unformatted.
1420    pub fn value_format(
1421        &self,
1422        settings: &crate::numfmt::NumberFormatSettings,
1423    ) -> crate::numfmt::NumberFormat {
1424        table_number_format(settings, &self.value_column, &self.value_dtype)
1425    }
1426
1427    /// Each bar's value as the table prints the value column.
1428    pub fn value_labels(&self, settings: &crate::numfmt::NumberFormatSettings) -> Vec<String> {
1429        self.labels_in(&self.value_format(settings))
1430    }
1431
1432    /// Each bar's value in `format`: whole for an integer column or a count.
1433    pub fn labels_in(&self, format: &crate::numfmt::NumberFormat) -> Vec<String> {
1434        let integer = self.value_dtype.is_integer();
1435        self.bars
1436            .iter()
1437            .map(|b| format_bar_value(b.value, integer, format))
1438            .collect()
1439    }
1440}
1441
1442/// A column name as SQL reads it: bare when it is a plain lowercase identifier, quoted
1443/// otherwise.
1444fn sql_ident(name: &str) -> String {
1445    let plain = name
1446        .chars()
1447        .next()
1448        .is_some_and(|c| c.is_ascii_lowercase() || c == '_')
1449        && name
1450            .chars()
1451            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_');
1452    if plain {
1453        name.to_string()
1454    } else {
1455        format!("\"{}\"", name.replace('"', "\"\""))
1456    }
1457}
1458
1459/// A bar chart, one bar per row: category from `category`, length from `value`. Expects
1460/// a grouped result and refuses a repeated category rather than guess. A null category
1461/// is its own bar; a null value is left out, counted in `no_value`.
1462pub fn prepare_bar_data(
1463    lf: &LazyFrame,
1464    category: &str,
1465    value: &str,
1466    order: BarOrder,
1467    cap: usize,
1468    sampling: &ChartSampling,
1469) -> Result<BarData> {
1470    let (df, rows) = read_columns(lf, &[category, value], sampling)?;
1471    let categories = df.column(category)?.as_materialized_series().clone();
1472    let labels_series = text_labels(&df, category)?;
1473    let labels: Vec<Option<&str>> = labels_series.str()?.iter().collect();
1474
1475    let mut seen: std::collections::HashMap<Option<&str>, usize> =
1476        std::collections::HashMap::with_capacity(labels.len());
1477    for label in &labels {
1478        *seen.entry(*label).or_default() += 1;
1479    }
1480    if seen.len() < labels.len() {
1481        let read = match rows.sample_size {
1482            Some(n) => format!("a sample of {} rows", crate::numfmt::group_chrome(n)),
1483            None => format!("{} rows", crate::numfmt::group_chrome(labels.len())),
1484        };
1485        let (c, v) = (sql_ident(category), sql_ident(value));
1486        // The q form only where it reads the names as they are.
1487        let q = if c == category && v == value {
1488            format!(" (or select avg {value} by {category})")
1489        } else {
1490            String::new()
1491        };
1492        return Err(color_eyre::eyre::eyre!(
1493            "{category} repeats: {} categories in {read}. A bar takes one row per category, \
1494             so group first: SELECT {c}, AVG({v}) FROM df GROUP BY {c}{q}, or choose Count \
1495             for the rows per category",
1496            crate::numfmt::group_chrome(seen.len()),
1497        ));
1498    }
1499
1500    let values = f64_values(&df, value)?;
1501    let (bars, more, no_value) = order_bars(&categories, &labels, &values, order, cap);
1502    Ok(BarData {
1503        category: category.to_string(),
1504        value_column: value.to_string(),
1505        bars,
1506        more,
1507        no_value,
1508        rows,
1509        value_dtype: df.column(value)?.dtype().clone(),
1510        counted: None,
1511        groups: Vec::new(),
1512        other: false,
1513        rows_note: None,
1514    })
1515}
1516
1517/// A bar per row in `order`, up to `cap`: the bars, how many are past the cap, and how
1518/// many rows had no value. Equal values keep row order.
1519fn order_bars(
1520    categories: &Series,
1521    labels: &[Option<&str>],
1522    values: &[Option<f64>],
1523    order: BarOrder,
1524    cap: usize,
1525) -> (Vec<Bar>, usize, usize) {
1526    let row_order: Vec<usize> = match order {
1527        BarOrder::Value => (0..labels.len()).collect(),
1528        BarOrder::Label => label_order(categories),
1529    };
1530    let mut no_value = 0;
1531    let mut bars: Vec<Bar> = row_order
1532        .into_iter()
1533        .filter_map(|i| match values[i] {
1534            Some(value) => Some(Bar {
1535                label: labels[i].map(str::to_string),
1536                value,
1537                by_group: Vec::new(),
1538            }),
1539            None => {
1540                no_value += 1;
1541                None
1542            }
1543        })
1544        .collect();
1545    if order == BarOrder::Value {
1546        // Stable: equal values keep row order.
1547        bars.sort_by(|a, b| b.value.total_cmp(&a.value));
1548    }
1549    let more = bars.len().saturating_sub(cap);
1550    bars.truncate(cap);
1551    (bars, more, no_value)
1552}
1553
1554/// Rows in the category's own order: text A to Z, numbers ascending, an enum in its
1555/// order, nulls last.
1556fn label_order(categories: &Series) -> Vec<usize> {
1557    categories
1558        .arg_sort(
1559            SortOptions::default()
1560                .with_nulls_last(true)
1561                .with_maintain_order(true),
1562        )
1563        .iter()
1564        .flatten()
1565        .map(|i| i as usize)
1566        .collect()
1567}
1568
1569/// A bar chart of rows per category, exact: the whole view is counted in one streamed
1570/// pass keeping only counts (or the held rows, if they are the whole view). Past
1571/// [`COUNT_CATEGORY_CAP`] categories the count stops and says so.
1572pub fn prepare_bar_counts(
1573    lf: &LazyFrame,
1574    category: &str,
1575    order: BarOrder,
1576    cap: usize,
1577    sampling: &ChartSampling,
1578) -> Result<BarData> {
1579    count_bars(lf, category, order, cap, COUNT_CATEGORY_CAP, sampling)
1580}
1581
1582fn count_bars(
1583    lf: &LazyFrame,
1584    category: &str,
1585    order: BarOrder,
1586    cap: usize,
1587    max_categories: usize,
1588    sampling: &ChartSampling,
1589) -> Result<BarData> {
1590    let counted = match held_counts(sampling, category, max_categories)? {
1591        Some(counted) => counted,
1592        None => {
1593            // A view the sample size takes whole is read as the other charts read it,
1594            // so they draw from the same rows after; a larger one is streamed.
1595            let fits = sampling
1596                .limit
1597                .zip(sampling.known_total)
1598                .is_some_and(|(n, total)| total <= n);
1599            let whole = if fits {
1600                let (df, rows) = read_columns(lf, &[category], sampling)?;
1601                rows.sample_size.is_none().then_some(df)
1602            } else {
1603                None
1604            };
1605            let counted = match whole {
1606                Some(df) => count_frame(&df, category, max_categories)?,
1607                None => stream_counts(lf, category, max_categories, &sampling.cancel)?,
1608            };
1609            hold_counts(sampling, category, &counted);
1610            counted
1611        }
1612    };
1613    let (counts, total) = match counted {
1614        Counted::All { counts, rows } => (counts, rows),
1615        Counted::TooMany => {
1616            return Err(color_eyre::eyre::eyre!(
1617                "more than {} categories of {category}: counting stopped. Count by a \
1618                 column with fewer values",
1619                crate::numfmt::group_chrome(max_categories)
1620            ));
1621        }
1622    };
1623    let data = |bars, more| BarData {
1624        category: category.to_string(),
1625        value_column: "count".to_string(),
1626        bars,
1627        more,
1628        no_value: 0,
1629        rows: RowsRead::every(total),
1630        value_dtype: DataType::UInt64,
1631        counted: sampling.limit.is_some_and(|n| total > n).then_some(total),
1632        groups: Vec::new(),
1633        other: false,
1634        rows_note: None,
1635    };
1636    let Some(counts) = counts else {
1637        return Ok(data(Vec::new(), 0));
1638    };
1639    // Label order first, so categories with equal counts come A to Z.
1640    let by_label: Vec<IdxSize> = label_order(counts.column(category)?.as_materialized_series())
1641        .into_iter()
1642        .map(|i| i as IdxSize)
1643        .collect();
1644    let counts = counts.take(&IdxCa::from_vec("order".into(), by_label))?;
1645    let categories = counts.column(category)?.as_materialized_series().clone();
1646    let labels_series = text_labels(&counts, category)?;
1647    let labels: Vec<Option<&str>> = labels_series.str()?.iter().collect();
1648    let values: Vec<Option<f64>> = counts
1649        .column(COUNT_COLUMN)?
1650        .u64()?
1651        .iter()
1652        .map(|n| n.map(|n| n as f64))
1653        .collect();
1654    let (bars, more, _) = order_bars(&categories, &labels, &values, order, cap);
1655    Ok(data(bars, more))
1656}
1657
1658/// The view's counts of `category` without reading it: counted before, or counted now
1659/// from rows held that are the whole view.
1660fn held_counts(
1661    sampling: &ChartSampling,
1662    category: &str,
1663    max_categories: usize,
1664) -> Result<Option<Counted>> {
1665    let holding = sampling.held.0.lock().unwrap_or_else(|e| e.into_inner());
1666    if let Some(held) = holding.counts.iter().find(|h| h.category == category) {
1667        return Ok(Some(held.counted.clone()));
1668    }
1669    let Some(whole) = holding
1670        .rows
1671        .as_ref()
1672        .filter(|h| h.rows.sample_size.is_none() && h.df.column(category).is_ok())
1673    else {
1674        return Ok(None);
1675    };
1676    let counted = count_frame(&whole.df, category, max_categories)?;
1677    drop(holding);
1678    hold_counts(sampling, category, &counted);
1679    Ok(Some(counted))
1680}
1681
1682fn hold_counts(sampling: &ChartSampling, category: &str, counted: &Counted) {
1683    let mut holding = sampling.held.0.lock().unwrap_or_else(|e| e.into_inner());
1684    holding.counts.retain(|h| h.category != category);
1685    if holding.counts.len() >= HELD_COUNTS {
1686        holding.counts.remove(0);
1687    }
1688    holding.counts.push(HeldCounts {
1689        category: category.to_string(),
1690        counted: counted.clone(),
1691    });
1692}
1693
1694/// Count the categories of rows already in memory.
1695pub(crate) fn count_frame(
1696    df: &DataFrame,
1697    category: &str,
1698    max_categories: usize,
1699) -> Result<Counted> {
1700    let mut tally = Tally::new(category, max_categories);
1701    tally.observe(&df.select([category])?)?;
1702    Ok(tally.finish()?)
1703}
1704
1705/// What a count found: one row per category with its rows in [`COUNT_COLUMN`] (none
1706/// when the view has no rows), and the rows counted; or more categories than it keeps.
1707#[derive(Clone)]
1708pub(crate) enum Counted {
1709    All {
1710        counts: Option<DataFrame>,
1711        rows: usize,
1712    },
1713    TooMany,
1714}
1715
1716pub(crate) const COUNT_COLUMN: &str = "__datui_bar_count";
1717
1718/// Counts held per view: enough to go back and forth between a few categories, each
1719/// up to [`COUNT_CATEGORY_CAP`] rows.
1720const HELD_COUNTS: usize = 4;
1721
1722/// Rows piled up unmerged before a merge, at least: a run of new categories costs a
1723/// merge now and then rather than one per batch.
1724const MERGE_AFTER: usize = 1 << 16;
1725
1726/// Rows per category, added up batch by batch. Holds one row per category and the
1727/// batches since the last merge, never the rows themselves.
1728pub(crate) struct Tally {
1729    category: PlSmallStr,
1730    max: usize,
1731    counts: Option<DataFrame>,
1732    /// Rows of `counts` as last merged: one per category.
1733    merged: usize,
1734    rows: usize,
1735    too_many: bool,
1736    /// Stopped before the end of the view because nobody wants the count.
1737    cancelled: bool,
1738}
1739
1740impl Tally {
1741    pub(crate) fn new(category: &str, max: usize) -> Self {
1742        Self {
1743            category: category.into(),
1744            max,
1745            counts: None,
1746            merged: 0,
1747            rows: 0,
1748            too_many: false,
1749            cancelled: false,
1750        }
1751    }
1752
1753    /// Count a batch of the category column. True once there are more categories than
1754    /// the count keeps, which stops the read.
1755    pub(crate) fn observe(&mut self, batch: &DataFrame) -> PolarsResult<bool> {
1756        if self.too_many {
1757            return Ok(true);
1758        }
1759        self.rows += batch.height();
1760        let part = group_counts(batch, &self.category, false)?;
1761        let mut counts = match self.counts.take() {
1762            Some(mut counts) => {
1763                counts.vstack_mut(&part)?;
1764                counts
1765            }
1766            None => part,
1767        };
1768        if counts.height() - self.merged >= self.merged.max(MERGE_AFTER) {
1769            counts = group_counts(&counts, &self.category, true)?;
1770            self.merged = counts.height();
1771            self.too_many = self.merged > self.max;
1772        }
1773        self.counts = Some(counts);
1774        Ok(self.too_many)
1775    }
1776
1777    pub(crate) fn finish(self) -> PolarsResult<Counted> {
1778        let counts = match self.counts {
1779            Some(counts) => Some(group_counts(&counts, &self.category, true)?),
1780            None => None,
1781        };
1782        if self.too_many || counts.as_ref().is_some_and(|c| c.height() > self.max) {
1783            return Ok(Counted::TooMany);
1784        }
1785        Ok(Counted::All {
1786            counts,
1787            rows: self.rows,
1788        })
1789    }
1790}
1791
1792/// One row per category of `df` with how many rows it stands for: its rows, or when
1793/// `summed` the counts they already carry, added up. A null category is one of them.
1794fn group_counts(df: &DataFrame, category: &str, summed: bool) -> PolarsResult<DataFrame> {
1795    let by = df.group_by([category])?;
1796    let groups = by.get_groups();
1797    let counts: Vec<u64> = if summed {
1798        let carried: Vec<u64> = df
1799            .column(COUNT_COLUMN)?
1800            .u64()?
1801            .into_no_null_iter()
1802            .collect();
1803        groups
1804            .iter()
1805            .map(|group| match group {
1806                GroupsIndicator::Idx((_, rows)) => rows.iter().map(|&i| carried[i as usize]).sum(),
1807                GroupsIndicator::Slice([first, len]) => {
1808                    carried[first as usize..(first + len) as usize].iter().sum()
1809                }
1810            })
1811            .collect()
1812    } else {
1813        groups.iter().map(|group| group.len() as u64).collect()
1814    };
1815    let mut columns = by.keys();
1816    columns.push(Column::new(COUNT_COLUMN.into(), counts));
1817    DataFrame::new_infer_height(columns)
1818}
1819
1820/// Count categories in one streamed pass, stopping past `max` or on `cancel` (a large
1821/// table can take minutes, and the next chart waits).
1822fn stream_counts(
1823    lf: &LazyFrame,
1824    category: &str,
1825    max: usize,
1826    cancel: &Arc<AtomicBool>,
1827) -> Result<Counted> {
1828    let state = Arc::new(Mutex::new(Tally::new(category, max)));
1829    let callback_state = Arc::clone(&state);
1830    let callback_cancel = Arc::clone(cancel);
1831    let sink = lf.clone().select([col(category)]).sink_batches(
1832        PlanCallback::new(move |batch: DataFrame| {
1833            let mut tally = callback_state
1834                .lock()
1835                .map_err(|_| PolarsError::ComputeError("count lock failed".into()))?;
1836            if callback_cancel.load(Ordering::Relaxed) {
1837                tally.cancelled = true;
1838                return Ok(true);
1839            }
1840            tally.observe(&batch)
1841        }),
1842        false,
1843        None,
1844    )?;
1845    // Streaming whatever the setting: a count per category is all this holds.
1846    crate::analysis::statistics::collect_lazy(sink, true)?;
1847    let tally = std::mem::replace(
1848        &mut *state.lock().unwrap_or_else(|e| e.into_inner()),
1849        Tally::new(category, max),
1850    );
1851    // Part of the view counted is not a count of it. A pass that finished before it
1852    // was told to stop is whole, and kept.
1853    if tally.cancelled {
1854        return Err(color_eyre::eyre::eyre!("count cancelled"));
1855    }
1856    Ok(tally.finish()?)
1857}
1858
1859// ----- Color: one series per value -----
1860
1861/// A color split: the column and the values given a group each, in color order (`None`
1862/// for nulls). With `other`, all other values form one more group, [`OTHER`].
1863#[derive(Clone, Copy, Debug)]
1864pub struct ColorSplit<'a> {
1865    pub column: &'a str,
1866    pub groups: &'a [Option<String>],
1867    pub other: bool,
1868}
1869
1870/// The name of the group of every value not given one of its own.
1871pub const OTHER: &str = "Other";
1872
1873impl ColorSplit<'_> {
1874    /// How many groups the rows fall in: the values', and Other.
1875    pub fn series(&self) -> usize {
1876        self.groups.len() + usize::from(self.other)
1877    }
1878
1879    /// Each group's name, as a legend writes it: Other last.
1880    pub fn names(&self) -> Vec<String> {
1881        let mut names: Vec<String> = self.groups.iter().map(group_label).collect();
1882        if self.other {
1883            names.push(OTHER.to_string());
1884        }
1885        names
1886    }
1887}
1888
1889/// A group's name as a legend writes it.
1890pub fn group_label(value: &Option<String>) -> String {
1891    value.clone().unwrap_or_else(|| "null".to_string())
1892}
1893
1894/// Each row's group: its index in `split.groups`, Other's after, or `None`. Compared as
1895/// text, as the value picker lists them.
1896fn row_groups(df: &DataFrame, split: ColorSplit<'_>) -> Result<Vec<Option<usize>>> {
1897    let text = text_labels(df, split.column)?;
1898    let index: std::collections::HashMap<Option<&str>, usize> = split
1899        .groups
1900        .iter()
1901        .enumerate()
1902        .map(|(i, g)| (g.as_deref(), i))
1903        .collect();
1904    let other = split.other.then_some(split.groups.len());
1905    Ok(text
1906        .str()?
1907        .iter()
1908        .map(|v| index.get(&v).copied().or(other))
1909        .collect())
1910}
1911
1912/// Values with the group each is in, and what was read for them.
1913type SplitValues = (Vec<(f64, Option<usize>)>, RowsRead);
1914
1915/// `column`'s finite values as read for a chart, each with its group when split
1916/// (`None` for a row of a value no group has, which still counts toward a range).
1917fn read_split(
1918    lf: &LazyFrame,
1919    column: &str,
1920    split: Option<ColorSplit<'_>>,
1921    sampling: &ChartSampling,
1922) -> Result<SplitValues> {
1923    let mut columns = vec![column];
1924    if let Some(split) = split {
1925        columns.push(split.column);
1926    }
1927    let (df, rows) = read_columns(lf, &columns, sampling)?;
1928    let values = f64_values(&df, column)?;
1929    let groups = split.map(|s| row_groups(&df, s)).transpose()?;
1930    let out = values
1931        .into_iter()
1932        .enumerate()
1933        .filter_map(|(i, v)| {
1934            let v = v?;
1935            Some((v, groups.as_ref().and_then(|groups| groups[i])))
1936        })
1937        .collect();
1938    Ok((out, rows))
1939}
1940
1941/// A column's values with the rows holding each, most rows first (equal counts in
1942/// the column's order), counted over the whole view.
1943#[derive(Clone, Debug, Default, PartialEq)]
1944pub struct ValueRows {
1945    pub values: Vec<(Option<String>, u64)>,
1946    /// Rows counted.
1947    pub rows: usize,
1948}
1949
1950/// Count `column`'s values over the whole view, or take the count held for it: one
1951/// streamed pass that keeps a count per value, the one a bar chart of counts makes.
1952pub fn value_rows(lf: &LazyFrame, column: &str, sampling: &ChartSampling) -> Result<ValueRows> {
1953    let counted = match held_counts(sampling, column, COUNT_CATEGORY_CAP)? {
1954        Some(counted) => counted,
1955        None => {
1956            let counted = stream_counts(lf, column, COUNT_CATEGORY_CAP, &sampling.cancel)?;
1957            hold_counts(sampling, column, &counted);
1958            counted
1959        }
1960    };
1961    let (counts, rows) = match counted {
1962        Counted::All { counts, rows } => (counts, rows),
1963        Counted::TooMany => {
1964            return Err(color_eyre::eyre::eyre!(
1965                "more than {} values of {column}: choose a column with fewer",
1966                crate::numfmt::group_chrome(COUNT_CATEGORY_CAP)
1967            ));
1968        }
1969    };
1970    let Some(counts) = counts else {
1971        return Ok(ValueRows {
1972            values: Vec::new(),
1973            rows,
1974        });
1975    };
1976    let by_label: Vec<IdxSize> = label_order(counts.column(column)?.as_materialized_series())
1977        .into_iter()
1978        .map(|i| i as IdxSize)
1979        .collect();
1980    let counts = counts.take(&IdxCa::from_vec("order".into(), by_label))?;
1981    let labels = text_labels(&counts, column)?;
1982    let mut values: Vec<(Option<String>, u64)> = labels
1983        .str()?
1984        .iter()
1985        .zip(counts.column(COUNT_COLUMN)?.u64()?.iter())
1986        .map(|(label, n)| (label.map(str::to_string), n.unwrap_or(0)))
1987        .collect();
1988    // Stable, so equal counts keep the column's order.
1989    values.sort_by_key(|v| std::cmp::Reverse(v.1));
1990    Ok(ValueRows { values, rows })
1991}
1992
1993/// The groups a color makes: the values picked, in the order picked, or else the
1994/// largest by rows; at most `most`, one per series color.
1995pub fn color_groups(
1996    rows: &ValueRows,
1997    picked: &[Option<String>],
1998    most: usize,
1999) -> Vec<Option<String>> {
2000    if !picked.is_empty() {
2001        return picked.iter().take(most).cloned().collect();
2002    }
2003    rows.values
2004        .iter()
2005        .take(most)
2006        .map(|(value, _)| value.clone())
2007        .collect()
2008}
2009
2010/// In a plan: each row's group as its place among `split.groups` (UInt32), Other's
2011/// after them, null for a value that has none.
2012fn group_expr(split: ColorSplit<'_>) -> Expr {
2013    let text = crate::past_calendar::text_expr(col(split.column), CastOptions::NonStrict);
2014    let mut out = match split.other {
2015        true => lit(split.groups.len() as u32).cast(DataType::UInt32),
2016        false => lit(NULL).cast(DataType::UInt32),
2017    };
2018    for (i, group) in split.groups.iter().enumerate().rev() {
2019        let matches = match group {
2020            Some(value) => text.clone().eq(lit(value.clone())),
2021            None => col(split.column).is_null(),
2022        };
2023        out = when(matches).then(lit(i as u32)).otherwise(out);
2024    }
2025    out
2026}
2027
2028/// Series of a line or scatter chart, each named: a Y column's or a color group's.
2029#[derive(Clone, Debug, Default)]
2030pub struct GroupedSeries {
2031    pub names: Vec<String>,
2032    pub series: Vec<Vec<(f64, f64)>>,
2033    /// Per series, where its line starts again after a gap.
2034    pub breaks: Vec<Vec<usize>>,
2035    pub x_axis_kind: XAxisTemporalKind,
2036    pub rows: RowsRead,
2037    /// The last series is Other: every value of the color without a series of its own.
2038    pub other: bool,
2039}
2040
2041/// A line or scatter chart of `y` split by `color`, from the rows a chart samples:
2042/// one series per group, in X order, each breaking where its Y has no value.
2043pub fn prepare_xy_by(
2044    lf: &LazyFrame,
2045    schema: &Schema,
2046    x: &str,
2047    y: &str,
2048    color: ColorSplit<'_>,
2049    sampling: &ChartSampling,
2050) -> Result<GroupedSeries> {
2051    let x_dtype = x_dtype(schema, x)?;
2052    let (df, rows) = read_columns(lf, &[x, y, color.column], sampling)?;
2053    let ys = f64_values(&df, y)?;
2054    let groups = row_groups(&df, color)?;
2055    let mut built = SeriesBuilder::new(color.series());
2056    for (x, i) in x_order(&df, x, x_dtype)? {
2057        if let Some(g) = groups[i] {
2058            built.push(g, x, ys[i]);
2059        }
2060    }
2061    Ok(GroupedSeries {
2062        names: color.names(),
2063        series: built.series,
2064        breaks: built.breaks,
2065        x_axis_kind: x_axis_temporal_kind(x_dtype),
2066        rows,
2067        other: color.other,
2068    })
2069}
2070
2071/// Most points an aggregated chart keeps. Past this X is close to a value per row,
2072/// and a time bucket is what it needs.
2073pub const AGGREGATE_POINTS_MAX: usize = 200_000;
2074
2075/// What an aggregated line or scatter chart groups by and makes of the rows.
2076#[derive(Clone, Copy, Debug)]
2077pub struct AggregateSpec<'a> {
2078    pub x: &'a str,
2079    pub time_unit: crate::chart::chart_modal::TimeUnit,
2080    pub ys: &'a [String],
2081    pub aggregate: crate::chart::chart_modal::Aggregate,
2082    /// The percentile a quantile takes.
2083    pub quantile: u8,
2084    pub cumulative: crate::chart::chart_modal::Cumulative,
2085    pub color: Option<ColorSplit<'a>>,
2086}
2087
2088/// Y as an aggregate reads it: as numbers, or as it is for a distinct count, which
2089/// counts strings and dates too.
2090fn y_values(y: Expr, aggregate: crate::chart::chart_modal::Aggregate) -> Expr {
2091    if aggregate.takes_any_y() {
2092        y
2093    } else {
2094        y.cast(DataType::Float64)
2095    }
2096}
2097
2098/// The row index first and last read the rows' order by.
2099const ROW_ORDER: &str = "__i";
2100
2101/// `values`' aggregate in a plan: quantiles at `quantile` percent; first and last by
2102/// [`ROW_ORDER`] (which the plan must carry), so the view's order holds.
2103fn aggregate_expr(
2104    values: Expr,
2105    aggregate: crate::chart::chart_modal::Aggregate,
2106    quantile: u8,
2107) -> Expr {
2108    use crate::chart::chart_modal::Aggregate;
2109    let in_order = || {
2110        values
2111            .clone()
2112            .sort_by([col(ROW_ORDER)], SortMultipleOptions::default())
2113            .drop_nulls()
2114    };
2115    match aggregate {
2116        // Nulls are no value: a group of only nulls has none, a gap.
2117        Aggregate::Distinct => values.drop_nulls().n_unique().cast(DataType::Float64),
2118        Aggregate::Sum => values.sum(),
2119        Aggregate::Mean => values.mean(),
2120        Aggregate::Median => values.median(),
2121        // The sample deviation: null for a group of one, which draws no point.
2122        Aggregate::Stdev => values.std(1),
2123        Aggregate::Quantile => {
2124            values.quantile(lit(f64::from(quantile) / 100.0), QuantileMethod::Linear)
2125        }
2126        Aggregate::Min => values.min(),
2127        Aggregate::Max => values.max(),
2128        Aggregate::First => in_order().first(),
2129        Aggregate::Last => in_order().last(),
2130        Aggregate::None | Aggregate::Count => len().cast(DataType::Float64),
2131    }
2132}
2133
2134/// `lf` with [`ROW_ORDER`] when `aggregate` reads the rows' order.
2135fn with_row_order(lf: &LazyFrame, aggregate: crate::chart::chart_modal::Aggregate) -> LazyFrame {
2136    if aggregate.follows_row_order() {
2137        lf.clone().with_row_index(ROW_ORDER, None)
2138    } else {
2139        lf.clone()
2140    }
2141}
2142
2143/// The rows a group-by aggregates: `select` (keys and values), with the row order first
2144/// and last need and each row's color group as a key; ungrouped rows dropped. Returns
2145/// the plan and its keys.
2146fn group_plan(
2147    lf: &LazyFrame,
2148    (mut select, mut keys): (Vec<Expr>, Vec<Expr>),
2149    aggregate: crate::chart::chart_modal::Aggregate,
2150    color: Option<ColorSplit<'_>>,
2151) -> (LazyFrame, Vec<Expr>) {
2152    if aggregate.follows_row_order() {
2153        select.push(col(ROW_ORDER));
2154    }
2155    if let Some(color) = color {
2156        select.push(group_expr(color).alias("__g"));
2157        keys.push(col("__g"));
2158    }
2159    let mut plan = with_row_order(lf, aggregate).select(select);
2160    if color.is_some() {
2161        plan = plan.filter(col("__g").is_not_null());
2162    }
2163    (plan, keys)
2164}
2165
2166/// Collect an aggregate's plan, streamed regardless of setting (a group-by holds a row
2167/// per group, and streaming checks `cancel` between morsels); a plan streaming cannot
2168/// take runs in memory to the end. A cancelled pass errors saying so.
2169fn aggregate_pass(lf: LazyFrame, sampling: &ChartSampling) -> Result<DataFrame> {
2170    crate::analysis::statistics::collect_lazy(lf, true).map_err(|e| {
2171        if sampling.cancel.load(Ordering::Relaxed) {
2172            color_eyre::eyre::eyre!(ENVELOPE_CANCELLED)
2173        } else {
2174            e.into()
2175        }
2176    })
2177}
2178
2179/// Rows of X a sample reads to judge how many values it has.
2180const GROUPS_SAMPLE: usize = 20_000;
2181
2182/// Refuse, before grouping, an X with more values than a chart can draw (a group per
2183/// near-unique value would hold the table), judged from a sample's distinct share.
2184fn refuse_too_many_groups(
2185    lf: &LazyFrame,
2186    x: &str,
2187    most: usize,
2188    sampling: &ChartSampling,
2189) -> Result<()> {
2190    let read = crate::analysis::sampling::analysis_rows(
2191        &lf.clone().select([col(x)]),
2192        Some(GROUPS_SAMPLE),
2193        sampling.known_total,
2194        sampling.seed,
2195        sampling.streaming,
2196    )?;
2197    let distinct = read.df.column(x)?.n_unique()?;
2198    let read_rows = read.df.height().max(1);
2199    let estimate = match read.sample_size {
2200        Some(_) => distinct as f64 / read_rows as f64 * read.total_rows as f64,
2201        None => distinct as f64,
2202    };
2203    if estimate > most as f64 {
2204        return Err(color_eyre::eyre::eyre!(
2205            "about {} values of {x}: more than a chart can draw. Bucket X by a time \
2206             unit, or choose a column with fewer values",
2207            crate::numfmt::group_chrome(estimate as usize)
2208        ));
2209    }
2210    Ok(())
2211}
2212
2213/// A line or scatter of Y aggregated per X (per time bucket for temporal X) and per
2214/// color group, in one lazy group-by over the view. A count needs no Y; other aggregates
2215/// draw a series per Y column (or color group of the first). Cumulative: each point is
2216/// the running total through its X, per series: a running sum, or compounded rates
2217/// `(1 + y1)(1 + y2)... - 1` (via summed `ln(1 + y)`); the aggregate is ignored, except
2218/// a count runs as a row count.
2219pub fn prepare_aggregate_xy(
2220    lf: &LazyFrame,
2221    schema: &Schema,
2222    spec: &AggregateSpec<'_>,
2223    sampling: &ChartSampling,
2224) -> Result<GroupedSeries> {
2225    use crate::chart::chart_modal::{Aggregate, Cumulative};
2226    let x_dtype = x_dtype(schema, spec.x)?;
2227    let mut x = col(spec.x);
2228    let bucketed = spec.time_unit.every().is_some()
2229        && matches!(x_dtype, DataType::Date | DataType::Datetime(_, _));
2230    if let Some(every) = spec.time_unit.every().filter(|_| bucketed) {
2231        x = x.dt().truncate(lit(every));
2232    }
2233    if !bucketed {
2234        refuse_too_many_groups(lf, spec.x, AGGREGATE_POINTS_MAX, sampling)?;
2235    }
2236    let x = until_cancelled(x, &sampling.cancel).alias("__x");
2237    let count = spec.aggregate == Aggregate::Count;
2238    let ys: &[String] = match (count, spec.color) {
2239        (true, _) => &[],
2240        (false, Some(_)) => &spec.ys[..spec.ys.len().min(1)],
2241        (false, None) => spec.ys,
2242    };
2243    let mut select = vec![x];
2244    let keys = vec![col("__x")];
2245    for (i, y) in ys.iter().enumerate() {
2246        select.push(y_values(col(y.as_str()), spec.aggregate).alias(format!("__y{i}")));
2247    }
2248    let (plan, keys) = group_plan(lf, (select, keys), spec.aggregate, spec.color);
2249    let plan = plan.filter(col("__x").is_not_null());
2250    let mut aggs = vec![len().alias("__n")];
2251    for i in 0..ys.len() {
2252        let y = col(format!("__y{i}"));
2253        let made = match spec.cumulative {
2254            Cumulative::Off => aggregate_expr(y.clone(), spec.aggregate, spec.quantile),
2255            Cumulative::Sum => y.clone().sum(),
2256            Cumulative::Compound => (lit(1.0) + y.clone()).log(lit(std::f64::consts::E)).sum(),
2257        };
2258        aggs.push(made.alias(format!("__a{i}")));
2259        // The values behind it: none is a gap, not the zero a sum of nothing is.
2260        aggs.push(y.count().alias(format!("__c{i}")));
2261    }
2262    let df = aggregate_pass(
2263        plan.group_by_stable(keys)
2264            .agg(aggs)
2265            .sort(["__x"], Default::default()),
2266        sampling,
2267    )?;
2268    if df.height() > AGGREGATE_POINTS_MAX {
2269        return Err(color_eyre::eyre::eyre!(
2270            "{} points: more than a chart can draw. Bucket X by a time unit, or \
2271             choose an X with fewer values",
2272            crate::numfmt::group_chrome(df.height())
2273        ));
2274    }
2275    let xs: Vec<Option<f64>> = x_values(&df, "__x", x_dtype)?;
2276    let counts = counts_of(&df, "__n")?;
2277    let groups: Option<Vec<Option<u32>>> = match spec.color {
2278        Some(_) => Some(df.column("__g")?.u32()?.iter().collect()),
2279        None => None,
2280    };
2281    let values: Vec<Vec<Option<f64>>> = if count {
2282        vec![counts.iter().map(|&n| Some(n as f64)).collect()]
2283    } else {
2284        (0..ys.len())
2285            .map(|i| {
2286                let made = df.column(&format!("__a{i}"))?.f64()?.clone();
2287                let behind = counts_of(&df, &format!("__c{i}"))?;
2288                Ok(made
2289                    .iter()
2290                    .zip(behind)
2291                    .map(|(v, n)| {
2292                        let v = v.filter(|_| n > 0)?;
2293                        // A bucket's compound return, from its log sum.
2294                        Some(if spec.cumulative == Cumulative::Compound {
2295                            v.exp_m1()
2296                        } else {
2297                            v
2298                        })
2299                    })
2300                    .collect())
2301            })
2302            .collect::<Result<_>>()?
2303    };
2304    let names: Vec<String> = match spec.color {
2305        Some(color) => color.names(),
2306        None if count => vec!["count".to_string()],
2307        None => ys.to_vec(),
2308    };
2309    let mut built = SeriesBuilder::new(names.len());
2310    let mut push = |s: usize, x: f64, y: Option<f64>| built.push(s, x, y.filter(|y| y.is_finite()));
2311    for (row, x) in xs.iter().enumerate() {
2312        let Some(x) = *x else { continue };
2313        match &groups {
2314            Some(groups) => {
2315                if let Some(g) = groups[row] {
2316                    push(g as usize, x, values[0][row]);
2317                }
2318            }
2319            None => {
2320                for (s, column) in values.iter().enumerate() {
2321                    push(s, x, column[row]);
2322                }
2323            }
2324        }
2325    }
2326    // A count of rows runs as a count, whichever way the totals were asked to run.
2327    let how = match spec.cumulative {
2328        Cumulative::Compound if count => Cumulative::Sum,
2329        how => how,
2330    };
2331    for points in &mut built.series {
2332        accumulate(points, how);
2333    }
2334    Ok(GroupedSeries {
2335        names,
2336        series: built.series,
2337        breaks: built.breaks,
2338        x_axis_kind: x_axis_temporal_kind(x_dtype),
2339        rows: RowsRead::every(counts.iter().sum::<u64>() as usize),
2340        other: spec.color.is_some_and(|c| c.other),
2341    })
2342}
2343
2344/// Make `points` cumulative along X: a running sum, or returns compounded (each
2345/// value a rate; the point is what 1 grew to, less 1).
2346pub fn accumulate(points: &mut [(f64, f64)], how: crate::chart::chart_modal::Cumulative) {
2347    use crate::chart::chart_modal::Cumulative;
2348    let mut total = 0.0;
2349    for (_, y) in points.iter_mut() {
2350        total = match how {
2351            Cumulative::Off => return,
2352            Cumulative::Sum => total + *y,
2353            Cumulative::Compound => (1.0 + total) * (1.0 + *y) - 1.0,
2354        };
2355        *y = total;
2356    }
2357}
2358
2359/// What an aggregated bar chart groups by and makes of the rows.
2360#[derive(Clone, Copy, Debug)]
2361pub struct BarAggregate<'a> {
2362    pub category: &'a str,
2363    /// The Y column; none for a count.
2364    pub value: Option<&'a str>,
2365    pub aggregate: crate::chart::chart_modal::Aggregate,
2366    /// The percentile a quantile takes.
2367    pub quantile: u8,
2368    pub color: Option<ColorSplit<'a>>,
2369    pub order: BarOrder,
2370    pub cap: usize,
2371}
2372
2373/// A bar chart of `value` aggregated per category (and color group) in one lazy
2374/// group-by; a count needs no value column, and uncolored is the exact count chart.
2375pub fn prepare_bar_aggregate(
2376    lf: &LazyFrame,
2377    schema: &Schema,
2378    spec: &BarAggregate<'_>,
2379    sampling: &ChartSampling,
2380) -> Result<BarData> {
2381    use crate::chart::chart_modal::Aggregate;
2382    let BarAggregate {
2383        category,
2384        value,
2385        aggregate,
2386        quantile,
2387        color,
2388        order,
2389        cap,
2390    } = *spec;
2391    let count = aggregate == Aggregate::Count;
2392    if count && color.is_none() {
2393        return prepare_bar_counts(lf, category, order, cap, sampling);
2394    }
2395    let value = match value {
2396        Some(value) if !count => Some(value),
2397        None if !count => return Err(color_eyre::eyre::eyre!("Pick a Y column")),
2398        _ => None,
2399    };
2400    let value_dtype = match value {
2401        Some(v) => schema
2402            .get(v)
2403            .cloned()
2404            .ok_or_else(|| color_eyre::eyre::eyre!("no column {v}"))?,
2405        None => DataType::UInt64,
2406    };
2407    let mut select = vec![until_cancelled(col(category), &sampling.cancel)];
2408    if let Some(value) = value {
2409        select.push(y_values(col(value), aggregate).alias("__v"));
2410    }
2411    let (plan, keys) = group_plan(lf, (select, vec![col(category)]), aggregate, color);
2412    let measure = match value {
2413        Some(_) => aggregate_expr(col("__v"), aggregate, quantile),
2414        None => len().cast(DataType::Float64),
2415    };
2416    // The values behind each bar: none is no bar, not the zero a sum of nothing is.
2417    let behind = match value {
2418        Some(_) => col("__v").count(),
2419        None => len(),
2420    };
2421    refuse_too_many_groups(lf, category, COUNT_CATEGORY_CAP, sampling)?;
2422    let df = aggregate_pass(
2423        // Stable, so bars of equal value keep one order from run to run.
2424        plan.group_by_stable(keys).agg([
2425            len().alias("__n"),
2426            measure.alias("__a"),
2427            behind.alias("__c"),
2428        ]),
2429        sampling,
2430    )?;
2431    let rows = counts_of(&df, "__n")?.iter().sum::<u64>() as usize;
2432    // A sum, least or greatest of whole numbers is whole; a count always is.
2433    let whole = aggregate.is_count()
2434        || (value_dtype.is_integer()
2435            && matches!(
2436                aggregate,
2437                Aggregate::Sum
2438                    | Aggregate::Min
2439                    | Aggregate::Max
2440                    | Aggregate::First
2441                    | Aggregate::Last
2442            ));
2443    let categories = df.column(category)?.as_materialized_series().clone();
2444    let labels_series = text_labels(&df, category)?;
2445    let labels: Vec<Option<&str>> = labels_series.str()?.iter().collect();
2446    let behind = counts_of(&df, "__c")?;
2447    // NaN and infinities draw nothing true: no bar.
2448    let measures: Vec<Option<f64>> = df
2449        .column("__a")?
2450        .f64()?
2451        .iter()
2452        .zip(&behind)
2453        .map(|(v, &n)| v.filter(|v| v.is_finite() && n > 0))
2454        .collect();
2455    let value_column = match value {
2456        Some(value) => format!("{} {value}", aggregate.named(quantile)),
2457        None => "count".to_string(),
2458    };
2459    let mut data = BarData {
2460        category: category.to_string(),
2461        value_column,
2462        bars: Vec::new(),
2463        more: 0,
2464        no_value: 0,
2465        rows: RowsRead::every(rows),
2466        value_dtype: if whole {
2467            DataType::Int64
2468        } else {
2469            DataType::Float64
2470        },
2471        counted: None,
2472        groups: Vec::new(),
2473        other: false,
2474        rows_note: None,
2475    };
2476    let too_many = || {
2477        color_eyre::eyre::eyre!(
2478            "more than {} categories of {category}: choose a column with fewer",
2479            crate::numfmt::group_chrome(COUNT_CATEGORY_CAP)
2480        )
2481    };
2482    let Some(color) = color else {
2483        if df.height() > COUNT_CATEGORY_CAP {
2484            return Err(too_many());
2485        }
2486        let (bars, more, no_value) = order_bars(&categories, &labels, &measures, order, cap);
2487        (data.bars, data.more, data.no_value) = (bars, more, no_value);
2488        return Ok(data);
2489    };
2490    // One bar row per category, a value per group: the category's first row stands
2491    // for it when ordering by label.
2492    let groups: Vec<Option<u32>> = df.column("__g")?.u32()?.iter().collect();
2493    let mut at: std::collections::HashMap<Option<&str>, usize> = Default::default();
2494    let mut firsts: Vec<IdxSize> = Vec::new();
2495    let mut rows_of: Vec<Vec<Option<f64>>> = Vec::new();
2496    for (row, label) in labels.iter().enumerate() {
2497        let i = *at.entry(*label).or_insert_with(|| {
2498            firsts.push(row as IdxSize);
2499            rows_of.push(vec![None; color.series()]);
2500            rows_of.len() - 1
2501        });
2502        if let Some(g) = groups[row] {
2503            rows_of[i][g as usize] = measures[row];
2504        }
2505    }
2506    if rows_of.len() > COUNT_CATEGORY_CAP {
2507        return Err(too_many());
2508    }
2509    let unique = categories.take(&IdxCa::from_vec("firsts".into(), firsts.clone()))?;
2510    let unique_labels: Vec<Option<&str>> = firsts.iter().map(|&r| labels[r as usize]).collect();
2511    // A count or a sum adds up across groups; any other measure orders by its
2512    // largest group.
2513    let totals: Vec<Option<f64>> = rows_of
2514        .iter()
2515        .map(|values| {
2516            let present = values.iter().flatten();
2517            if count || aggregate == Aggregate::Sum {
2518                Some(present.sum())
2519            } else {
2520                present.copied().reduce(f64::max)
2521            }
2522        })
2523        .collect();
2524    let (mut bars, more, no_value) = order_bars(&unique, &unique_labels, &totals, order, cap);
2525    // `order_bars` names each bar by its label; give each its groups back.
2526    let by_label: std::collections::HashMap<Option<&str>, usize> = unique_labels
2527        .iter()
2528        .enumerate()
2529        .map(|(i, l)| (*l, i))
2530        .collect();
2531    for bar in &mut bars {
2532        if let Some(&i) = by_label.get(&bar.label.as_deref()) {
2533            bar.by_group = rows_of[i].clone();
2534        }
2535    }
2536    data.bars = bars;
2537    data.more = more;
2538    data.no_value = no_value;
2539    data.groups = color.names();
2540    data.other = color.other;
2541    Ok(data)
2542}
2543
2544#[cfg(test)]
2545mod tests;