1use 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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
19pub enum XAxisTemporalKind {
20 #[default]
21 Numeric,
22 Date, DatetimeUs, DatetimeMs,
25 DatetimeNs,
26 Time, }
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
42pub 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
50pub(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
71pub(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
80pub 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#[derive(Clone, Debug)]
126pub struct ChartSampling {
127 pub limit: Option<usize>,
129 pub known_total: Option<usize>,
131 pub seed: u64,
133 pub streaming: bool,
134 pub full_passes: bool,
138 pub held: HeldRows,
140 pub cancel: Arc<AtomicBool>,
142}
143
144impl ChartSampling {
145 #[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#[derive(Clone, Default)]
164pub struct HeldRows(Arc<Mutex<Holding>>);
165
166#[derive(Default)]
167struct Holding {
168 rows: Option<Held>,
169 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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
195pub struct RowsRead {
196 pub total_rows: usize,
197 pub sample_size: Option<usize>,
198 pub envelope_steps: Option<usize>,
201 pub seed: Option<u64>,
203}
204
205#[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 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
236pub struct Clipped {
237 pub range: ValueRange,
238 pub outside: usize,
239}
240
241pub 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
278fn 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
333fn 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
343fn 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
355fn 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
362fn 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
374struct 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
405fn 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
413fn 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 pub fn every(total_rows: usize) -> Self {
427 Self {
428 total_rows,
429 ..Self::default()
430 }
431 }
432}
433
434#[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
443pub 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
472pub struct ChartDataResult {
474 pub series: Vec<Vec<(f64, f64)>>,
476 pub breaks: Vec<Vec<usize>>,
479 pub x_axis_kind: XAxisTemporalKind,
480 pub rows: RowsRead,
481}
482
483pub 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#[derive(Clone, Debug)]
501pub struct HistogramBin {
502 pub center: f64,
503 pub count: f64,
504}
505
506#[derive(Clone, Debug)]
509pub struct HistogramGroup {
510 pub name: String,
511 pub counts: Vec<f64>,
512}
513
514#[derive(Clone, Debug)]
516pub struct HistogramData {
517 pub column: String,
518 pub bins: Vec<HistogramBin>,
520 pub groups: Vec<HistogramGroup>,
522 pub other: bool,
524 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#[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 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#[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
563pub struct BoxMarks {
566 pub outline: [(f64, f64); 5],
568 pub median: [(f64, f64); 2],
569 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 pub of: usize,
606}
607
608impl HistogramData {
609 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
630pub fn other_at(other: bool, n: usize) -> Option<usize> {
632 (other && n > 0).then(|| n - 1)
633}
634
635pub 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#[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
659pub 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 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
747enum Envelope {
749 Drawn {
752 series: Vec<Vec<(f64, f64)>>,
753 breaks: Vec<Vec<usize>>,
754 rows: usize,
755 steps: usize,
756 },
757 Fits(usize),
759}
760
761const ENVELOPE_CANCELLED: &str = "chart cancelled";
763
764fn 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
777fn 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
789fn 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 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 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 .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
908fn 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
921fn 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
940pub 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 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
1068fn 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
1081fn 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
1103pub 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
1116pub 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
1155fn 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
1181fn 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
1207pub 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
1223pub 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
1268pub 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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
1347#[serde(rename_all = "snake_case")]
1348pub enum BarOrder {
1349 #[default]
1351 Value,
1352 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
1367pub const BAR_CAP: usize = 100;
1370
1371pub 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
1380pub const COUNT_CATEGORY_CAP: usize = 100_000;
1383
1384#[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#[derive(Clone, Debug)]
1395pub struct BarData {
1396 pub category: String,
1397 pub value_column: String,
1398 pub bars: Vec<Bar>,
1399 pub more: usize,
1401 pub no_value: usize,
1403 pub rows: RowsRead,
1404 pub value_dtype: DataType,
1406 pub counted: Option<usize>,
1409 pub groups: Vec<String>,
1411 pub other: bool,
1413 pub rows_note: Option<String>,
1415}
1416
1417impl BarData {
1418 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 pub fn value_labels(&self, settings: &crate::numfmt::NumberFormatSettings) -> Vec<String> {
1429 self.labels_in(&self.value_format(settings))
1430 }
1431
1432 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
1442fn 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
1459pub 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 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
1517fn 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 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
1554fn 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
1569pub 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 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 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
1658fn 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
1694pub(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#[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
1718const HELD_COUNTS: usize = 4;
1721
1722const MERGE_AFTER: usize = 1 << 16;
1725
1726pub(crate) struct Tally {
1729 category: PlSmallStr,
1730 max: usize,
1731 counts: Option<DataFrame>,
1732 merged: usize,
1734 rows: usize,
1735 too_many: bool,
1736 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 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
1792fn 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
1820fn 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 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 if tally.cancelled {
1854 return Err(color_eyre::eyre::eyre!("count cancelled"));
1855 }
1856 Ok(tally.finish()?)
1857}
1858
1859#[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
1870pub const OTHER: &str = "Other";
1872
1873impl ColorSplit<'_> {
1874 pub fn series(&self) -> usize {
1876 self.groups.len() + usize::from(self.other)
1877 }
1878
1879 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
1889pub fn group_label(value: &Option<String>) -> String {
1891 value.clone().unwrap_or_else(|| "null".to_string())
1892}
1893
1894fn 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
1912type SplitValues = (Vec<(f64, Option<usize>)>, RowsRead);
1914
1915fn 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#[derive(Clone, Debug, Default, PartialEq)]
1944pub struct ValueRows {
1945 pub values: Vec<(Option<String>, u64)>,
1946 pub rows: usize,
1948}
1949
1950pub 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 values.sort_by_key(|v| std::cmp::Reverse(v.1));
1990 Ok(ValueRows { values, rows })
1991}
1992
1993pub 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
2010fn 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#[derive(Clone, Debug, Default)]
2030pub struct GroupedSeries {
2031 pub names: Vec<String>,
2032 pub series: Vec<Vec<(f64, f64)>>,
2033 pub breaks: Vec<Vec<usize>>,
2035 pub x_axis_kind: XAxisTemporalKind,
2036 pub rows: RowsRead,
2037 pub other: bool,
2039}
2040
2041pub 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
2071pub const AGGREGATE_POINTS_MAX: usize = 200_000;
2074
2075#[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 pub quantile: u8,
2084 pub cumulative: crate::chart::chart_modal::Cumulative,
2085 pub color: Option<ColorSplit<'a>>,
2086}
2087
2088fn 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
2098const ROW_ORDER: &str = "__i";
2100
2101fn 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 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 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
2134fn 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
2143fn 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
2166fn 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
2179const GROUPS_SAMPLE: usize = 20_000;
2181
2182fn 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
2213pub 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 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 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 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
2344pub 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#[derive(Clone, Copy, Debug)]
2361pub struct BarAggregate<'a> {
2362 pub category: &'a str,
2363 pub value: Option<&'a str>,
2365 pub aggregate: crate::chart::chart_modal::Aggregate,
2366 pub quantile: u8,
2368 pub color: Option<ColorSplit<'a>>,
2369 pub order: BarOrder,
2370 pub cap: usize,
2371}
2372
2373pub 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 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 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 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 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 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 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 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;