1use chrono::{DateTime, Datelike, NaiveDate, NaiveDateTime, NaiveTime};
10use color_eyre::Result;
11use polars::chunked_array::cast::CastOptions;
12use polars::datatypes::{DataType, TimeUnit};
13use polars::prelude::*;
14use std::f64::consts::PI;
15use std::sync::atomic::{AtomicBool, Ordering};
16use std::sync::{Arc, Mutex};
17
18#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
20pub enum XAxisTemporalKind {
21 #[default]
22 Numeric,
23 Date, DatetimeUs, DatetimeMs,
26 DatetimeNs,
27 Time, }
29
30fn x_axis_temporal_kind(dtype: &DataType) -> XAxisTemporalKind {
31 match dtype {
32 DataType::Date => XAxisTemporalKind::Date,
33 DataType::Datetime(unit, _) => match unit {
34 TimeUnit::Nanoseconds => XAxisTemporalKind::DatetimeNs,
35 TimeUnit::Microseconds => XAxisTemporalKind::DatetimeUs,
36 TimeUnit::Milliseconds => XAxisTemporalKind::DatetimeMs,
37 },
38 DataType::Time => XAxisTemporalKind::Time,
39 _ => XAxisTemporalKind::Numeric,
40 }
41}
42
43pub fn x_axis_temporal_kind_for_column(schema: &Schema, x_column: &str) -> XAxisTemporalKind {
45 schema
46 .get(x_column)
47 .map(x_axis_temporal_kind)
48 .unwrap_or(XAxisTemporalKind::Numeric)
49}
50
51const MAX_AXIS_PLACES: i32 = 6;
53const MAX_MANTISSA_PLACES: i32 = 12;
55const SCIENTIFIC_FROM: f64 = 1e15;
57
58#[derive(Clone, Debug, Default, PartialEq)]
61pub struct AxisNumbers {
62 pub format: crate::numfmt::NumberFormat,
63 pub whole: bool,
64}
65
66impl AxisNumbers {
67 pub fn column(
69 settings: &crate::numfmt::NumberFormatSettings,
70 schema: Option<&Schema>,
71 column: &str,
72 ) -> Self {
73 match schema.and_then(|s| s.get(column)) {
74 Some(dtype) => Self {
75 format: table_number_format(settings, column, dtype),
76 whole: dtype.is_integer(),
77 },
78 None => Self::default(),
79 }
80 }
81
82 pub fn columns(
84 settings: &crate::numfmt::NumberFormatSettings,
85 schema: Option<&Schema>,
86 columns: &[String],
87 ) -> Self {
88 let mut each = columns.iter().map(|c| Self::column(settings, schema, c));
89 let Some(first) = each.next() else {
90 return Self::default();
91 };
92 let whole = first.whole && each.all(|n| n.whole);
93 Self { whole, ..first }
94 }
95
96 pub fn count(settings: &crate::numfmt::NumberFormatSettings) -> Self {
98 Self {
99 format: table_number_format(settings, "Count", &DataType::UInt64),
100 whole: true,
101 }
102 }
103
104 pub fn measure(settings: &crate::numfmt::NumberFormatSettings, name: &str) -> Self {
106 Self {
107 format: table_number_format(settings, name, &DataType::Float64),
108 whole: false,
109 }
110 }
111
112 pub fn fractional(self) -> Self {
114 Self {
115 whole: false,
116 ..self
117 }
118 }
119}
120
121#[derive(Clone, Debug)]
125pub struct AxisFormat {
126 format: crate::numfmt::NumberFormat,
127 full: Notation,
128 short: Option<Notation>,
130 zero_below: f64,
133}
134
135#[derive(Clone, Copy, Debug, PartialEq)]
136enum Notation {
137 Fixed {
139 places: usize,
140 unit: f64,
141 suffix: &'static str,
142 },
143 Scientific { places: usize },
145 Prefixed,
149}
150
151impl AxisFormat {
152 pub fn new(ticks: &[f64], numbers: &AxisNumbers) -> Self {
154 let ticks: Vec<f64> = ticks.iter().copied().filter(|v| v.is_finite()).collect();
155 let top = ticks.iter().fold(0.0_f64, |top, v| top.max(v.abs()));
156 let gap = ticks
158 .windows(2)
159 .map(|w| (w[1] - w[0]).abs())
160 .filter(|gap| *gap > 0.0)
161 .fold(f64::INFINITY, f64::min);
162 let places = if top >= SCIENTIFIC_FROM {
163 None
164 } else if numbers.whole {
165 Some(0)
166 } else {
167 fixed_places(&ticks, top, gap)
168 };
169 let (full, short) = match places {
170 Some(places) => (
171 Notation::Fixed {
172 places,
173 unit: 1.0,
174 suffix: "",
175 },
176 short_notation(&ticks, top, gap),
177 ),
178 None => {
179 let apart = if gap.is_finite() && top > 0.0 {
181 (magnitude(top) - magnitude(gap)).clamp(0, MAX_MANTISSA_PLACES) as usize
182 } else {
183 0
184 };
185 (
186 Notation::Scientific {
187 places: apart.max(2),
188 },
189 Some(Notation::Scientific { places: apart }),
190 )
191 }
192 };
193 Self {
194 format: numbers.format.clone(),
195 full,
196 short,
197 zero_below: if gap.is_finite() { gap * 1e-9 } else { 0.0 },
198 }
199 }
200
201 pub fn log(ticks: &[f64], numbers: &AxisNumbers) -> Self {
207 let ticks: Vec<f64> = ticks.iter().copied().filter(|v| v.is_finite()).collect();
208 let top = ticks.iter().fold(0.0_f64, |top, v| top.max(v.abs()));
209 let full = if top >= SCIENTIFIC_FROM {
210 Notation::Scientific { places: 0 }
211 } else {
212 Notation::Fixed {
213 places: fewest_places(&ticks, 1.0, 0, MAX_AXIS_PLACES),
214 unit: 1.0,
215 suffix: "",
216 }
217 };
218 Self {
219 format: numbers.format.clone(),
220 full,
221 short: (1e3..SCIENTIFIC_FROM)
222 .contains(&top)
223 .then_some(Notation::Prefixed),
224 zero_below: 0.0,
225 }
226 }
227
228 pub fn ends_and_middle([lo, hi]: [f64; 2], numbers: &AxisNumbers) -> Self {
231 Self::new(&[lo, (lo + hi) / 2.0, hi], numbers)
232 }
233
234 pub fn label(&self, v: f64, level: usize) -> Option<String> {
237 let notation = match level {
238 0 => self.full,
239 1 => self.short?,
240 _ => return None,
241 };
242 Some(self.write(v, notation))
243 }
244
245 fn write(&self, v: f64, notation: Notation) -> String {
246 let (places, unit, suffix) = match notation {
247 Notation::Scientific { places } => {
248 let v = if v.abs() < self.zero_below { 0.0 } else { v };
249 return scientific(v, places, self.format.decimal_sep);
250 }
251 Notation::Fixed { .. } | Notation::Prefixed if !v.is_finite() => {
252 return v.to_string();
253 }
254 Notation::Prefixed => {
255 let (unit, suffix) = [(1e12, "T"), (1e9, "G"), (1e6, "M"), (1e3, "k")]
256 .into_iter()
257 .find(|(unit, _)| v.abs() >= *unit)
258 .unwrap_or((1.0, ""));
259 let places = fewest_places(&[v], unit, 0, 2);
260 (places, unit, suffix)
261 }
262 Notation::Fixed {
263 places,
264 unit,
265 suffix,
266 } => (places, unit, suffix),
267 };
268 let fixed = crate::numfmt::NumberFormat {
269 float_precision: Some(places as u8),
270 ..self.format.clone()
271 };
272 let mut out = String::new();
273 fixed.write_f64(v / unit, &mut String::new(), &mut out);
274 if !out.chars().any(|c| matches!(c, '1'..='9')) {
277 if unit > 1.0 {
278 return "0".to_string();
279 }
280 out.retain(|c| c != '-');
281 }
282 out.push_str(suffix);
283 out
284 }
285}
286
287fn fixed_places(ticks: &[f64], top: f64, gap: f64) -> Option<usize> {
292 let figures = if top > 0.0 { 2 - magnitude(top) } else { 0 };
293 let apart = if gap.is_finite() { -magnitude(gap) } else { 0 };
294 let places = figures.max(apart).max(0);
295 (places <= MAX_AXIS_PLACES).then(|| fewest_places(ticks, 1.0, apart.max(0), places))
296}
297
298fn fewest_places(ticks: &[f64], unit: f64, least: i32, most: i32) -> usize {
301 let exact = |places: i32| {
302 ticks.iter().all(|v| {
303 let scaled = v / unit * 10f64.powi(places);
304 (scaled - scaled.round()).abs() <= 1e-9 * scaled.abs().max(1.0)
306 })
307 };
308 (least..most).find(|&p| exact(p)).unwrap_or(most) as usize
309}
310
311fn short_notation(ticks: &[f64], top: f64, gap: f64) -> Option<Notation> {
316 let (unit, suffix) = [(1e12, "T"), (1e9, "G"), (1e6, "M"), (1e3, "k")]
317 .into_iter()
318 .find(|(unit, _)| top >= *unit)?;
319 let figures = 1 - magnitude(top / unit);
320 let apart = if gap.is_finite() {
321 -magnitude(gap / unit)
322 } else {
323 0
324 };
325 let most = figures.max(apart).clamp(0, 2);
326 Some(Notation::Fixed {
327 places: fewest_places(ticks, unit, apart.clamp(0, most), most),
328 unit,
329 suffix,
330 })
331}
332
333fn magnitude(v: f64) -> i32 {
337 (v.log10() + 1e-9).floor() as i32
338}
339
340fn scientific(v: f64, places: usize, decimal_sep: char) -> String {
342 let v = if v == 0.0 { 0.0 } else { v };
344 let text = format!("{v:.places$e}");
345 if decimal_sep == '.' {
346 text
347 } else {
348 text.replacen('.', decimal_sep.encode_utf8(&mut [0; 4]), 1)
349 }
350}
351
352pub fn table_number_format(
354 settings: &crate::numfmt::NumberFormatSettings,
355 column: &str,
356 dtype: &DataType,
357) -> crate::numfmt::NumberFormat {
358 match settings.formatter_for(column, dtype) {
359 crate::numfmt::CellFormatter::Number(format) => format,
360 crate::numfmt::CellFormatter::Passthrough => crate::numfmt::NumberFormat::PLAIN,
361 }
362}
363
364pub(crate) fn x_datetime(v: f64, kind: XAxisTemporalKind) -> Option<NaiveDateTime> {
366 const UNIX_EPOCH_CE_DAYS: i32 = 719_163;
367 match kind {
368 XAxisTemporalKind::Date => NaiveDate::from_num_days_from_ce_opt(
369 UNIX_EPOCH_CE_DAYS.saturating_add(v.trunc() as i32),
370 )
371 .map(|d| d.and_time(NaiveTime::MIN)),
372 XAxisTemporalKind::DatetimeUs => {
373 DateTime::from_timestamp_micros(v.trunc() as i64).map(|dt| dt.naive_utc())
374 }
375 XAxisTemporalKind::DatetimeMs => {
376 DateTime::from_timestamp_millis(v.trunc() as i64).map(|dt| dt.naive_utc())
377 }
378 XAxisTemporalKind::DatetimeNs => {
379 DateTime::from_timestamp_millis((v.trunc() as i64) / 1_000_000).map(|dt| dt.naive_utc())
380 }
381 XAxisTemporalKind::Numeric | XAxisTemporalKind::Time => None,
382 }
383}
384
385pub(crate) fn x_time(v: f64) -> Option<NaiveTime> {
387 let nsecs = v.trunc() as u64;
388 NaiveTime::from_num_seconds_from_midnight_opt(
389 (nsecs / 1_000_000_000) as u32,
390 (nsecs % 1_000_000_000) as u32,
391 )
392}
393
394pub fn x_axis_label_at(
401 v: f64,
402 kind: XAxisTemporalKind,
403 bounds: (f64, f64),
404 level: usize,
405 numbers: &AxisFormat,
406) -> Option<String> {
407 if kind == XAxisTemporalKind::Numeric {
408 return numbers.label(v, level);
409 }
410 if kind == XAxisTemporalKind::Time {
411 let pattern = ["%H:%M:%S", "%H:%M"].get(level)?;
412 return Some(match x_time(v) {
413 Some(t) => t.format(pattern).to_string(),
414 None => numbers.label(v, level)?,
415 });
416 }
417 let Some(at) = x_datetime(v, kind) else {
418 return numbers.label(v, level);
419 };
420 let ends = x_datetime(bounds.0, kind).zip(x_datetime(bounds.1, kind));
421 let one_day = ends.is_some_and(|(a, b)| a.date() == b.date());
422 let one_year = ends.is_some_and(|(a, b)| a.year() == b.year());
423 let dates: &[&str] = if one_year {
424 &["%Y-%m-%d", "%m-%d"]
425 } else {
426 &["%Y-%m-%d", "%Y-%m", "%Y"]
427 };
428 let patterns: Vec<&str> = if kind == XAxisTemporalKind::Date {
429 dates.to_vec()
430 } else if one_day {
431 vec!["%Y-%m-%d %H:%M", "%H:%M"]
432 } else {
433 std::iter::once("%Y-%m-%d %H:%M")
434 .chain(dates.iter().copied())
435 .collect()
436 };
437 patterns.get(level).map(|p| at.format(p).to_string())
438}
439
440#[derive(Clone, Debug)]
442pub struct ChartSampling {
443 pub limit: Option<usize>,
445 pub known_total: Option<usize>,
447 pub seed: u64,
449 pub streaming: bool,
450 pub full_passes: bool,
454 pub held: HeldRows,
456 pub cancel: Arc<AtomicBool>,
458}
459
460impl ChartSampling {
461 pub fn rows(limit: Option<usize>) -> Self {
463 Self {
464 limit,
465 known_total: None,
466 seed: crate::sampling::Sample::default().seed,
467 streaming: false,
468 full_passes: true,
469 held: HeldRows::default(),
470 cancel: Arc::default(),
471 }
472 }
473}
474
475#[derive(Clone, Default)]
480pub struct HeldRows(Arc<Mutex<Holding>>);
481
482#[derive(Default)]
483struct Holding {
484 rows: Option<Held>,
485 counts: Vec<HeldCounts>,
488}
489
490struct Held {
491 limit: Option<usize>,
492 seed: u64,
493 df: DataFrame,
494 rows: RowsRead,
495}
496
497struct HeldCounts {
498 category: String,
499 counted: Counted,
500}
501
502impl std::fmt::Debug for HeldRows {
503 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
504 f.write_str("HeldRows")
505 }
506}
507
508#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
511pub struct RowsRead {
512 pub total_rows: usize,
513 pub sample_size: Option<usize>,
514 pub envelope_steps: Option<usize>,
518 pub seed: Option<u64>,
520}
521
522#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
525#[serde(rename_all = "snake_case")]
526pub enum ValueRange {
527 #[default]
528 All,
529 Percentile1To99,
531}
532
533impl ValueRange {
534 pub const ALL: [Self; 2] = [Self::All, Self::Percentile1To99];
535
536 pub fn label(self) -> &'static str {
537 match self {
538 Self::All => "All",
539 Self::Percentile1To99 => "p1-p99",
540 }
541 }
542
543 fn quantiles(self) -> Option<(f64, f64)> {
544 match self {
545 Self::All => None,
546 Self::Percentile1To99 => Some((0.01, 0.99)),
547 }
548 }
549}
550
551#[derive(Clone, Copy, Debug, PartialEq, Eq)]
553pub struct Clipped {
554 pub range: ValueRange,
555 pub outside: usize,
556}
557
558pub fn chart_notes(rows: &RowsRead, clipped: Option<&Clipped>, middot: &str) -> Vec<String> {
562 let mut notes = Vec::new();
563 if let Some(steps) = rows.envelope_steps {
564 notes.push(format!(
565 "min and max of {} rows in {} steps",
566 crate::discover::format_rows(rows.total_rows),
567 crate::numfmt::group_chrome(steps)
568 ));
569 }
570 if let Some(n) = rows.sample_size {
571 let mut note = format!(
572 "sample of {} of {} rows",
573 crate::numfmt::group_chrome(n),
574 crate::discover::format_rows(rows.total_rows)
575 );
576 if let Some(seed) = rows.seed {
577 note.push_str(&format!(" {middot} seed {seed}"));
578 }
579 notes.push(note);
580 }
581 if let Some(clipped) = clipped {
582 let noun = if clipped.outside == 1 {
583 "value"
584 } else {
585 "values"
586 };
587 notes.push(format!(
588 "{} {noun} outside {}",
589 crate::numfmt::group_chrome(clipped.outside),
590 clipped.range.label()
591 ));
592 }
593 notes
594}
595
596fn read_columns(
604 lf: &LazyFrame,
605 columns: &[&str],
606 sampling: &ChartSampling,
607) -> Result<(DataFrame, RowsRead)> {
608 let mut unique: Vec<PlSmallStr> = Vec::with_capacity(columns.len());
609 for c in columns {
610 if !unique.iter().any(|u| u == c) {
611 unique.push((*c).into());
612 }
613 }
614 let mut holding = sampling.held.0.lock().unwrap_or_else(|e| e.into_inner());
615 if let Some(h) = holding
616 .rows
617 .as_ref()
618 .filter(|h| h.limit == sampling.limit && h.seed == sampling.seed)
619 {
620 if unique.iter().all(|c| h.df.column(c).is_ok()) {
621 return Ok((h.df.select(unique.iter().cloned())?, h.rows));
622 }
623 for c in h.df.get_column_names() {
624 if !unique.contains(c) {
625 unique.push(c.clone());
626 }
627 }
628 }
629 let lf = lf
630 .clone()
631 .select(unique.iter().map(|c| col(c.clone())).collect::<Vec<_>>());
632 let read = crate::statistics::analysis_rows(
633 &lf,
634 sampling.limit,
635 sampling.known_total,
636 sampling.seed,
637 sampling.streaming,
638 )?;
639 let rows = RowsRead {
640 total_rows: read.total_rows,
641 sample_size: read.sample_size,
642 envelope_steps: None,
643 seed: read.sample_size.map(|_| sampling.seed),
644 };
645 holding.rows = Some(Held {
646 limit: sampling.limit,
647 seed: sampling.seed,
648 df: read.df.clone(),
649 rows,
650 });
651 Ok((read.df, rows))
652}
653
654fn f64_values(df: &DataFrame, column: &str) -> Result<Vec<Option<f64>>> {
656 let cast = df.column(column)?.cast(&DataType::Float64)?;
657 Ok(cast
658 .f64()?
659 .iter()
660 .map(|v| v.filter(|v| v.is_finite()))
661 .collect())
662}
663
664fn x_values(df: &DataFrame, column: &str, dtype: &DataType) -> Result<Vec<Option<f64>>> {
667 match dtype {
668 DataType::Datetime(_, _) | DataType::Date | DataType::Time => {
669 let ordinal = df.column(column)?.cast(&DataType::Int64)?;
670 Ok(ordinal.i64()?.iter().map(|v| v.map(|v| v as f64)).collect())
671 }
672 _ => f64_values(df, column),
673 }
674}
675
676pub struct ChartXRangeResult {
678 pub x_min: f64,
679 pub x_max: f64,
680 pub x_axis_kind: XAxisTemporalKind,
681 pub rows: RowsRead,
682}
683
684pub fn prepare_chart_x_range(
686 lf: &LazyFrame,
687 schema: &Schema,
688 x_column: &str,
689 sampling: &ChartSampling,
690) -> Result<ChartXRangeResult> {
691 let x_dtype = schema
692 .get(x_column)
693 .ok_or_else(|| color_eyre::eyre::eyre!("x column '{}' not in schema", x_column))?;
694 let x_axis_kind = x_axis_temporal_kind(x_dtype);
695 let (df, rows) = read_columns(lf, &[x_column], sampling)?;
696 let (x_min, x_max) = x_values(&df, x_column, x_dtype)?
697 .into_iter()
698 .flatten()
699 .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), x| {
700 (lo.min(x), hi.max(x))
701 });
702 let (x_min, x_max) = if x_max >= x_min {
703 (x_min, x_max)
704 } else {
705 (0.0, 1.0)
706 };
707 Ok(ChartXRangeResult {
708 x_min,
709 x_max,
710 x_axis_kind,
711 rows,
712 })
713}
714
715pub struct ChartDataResult {
717 pub series: Vec<Vec<(f64, f64)>>,
719 pub breaks: Vec<Vec<usize>>,
722 pub x_axis_kind: XAxisTemporalKind,
723 pub rows: RowsRead,
724}
725
726pub fn segments<'a>(points: &'a [(f64, f64)], breaks: &[usize]) -> Vec<&'a [(f64, f64)]> {
728 let mut out = Vec::with_capacity(breaks.len() + 1);
729 let mut start = 0;
730 for &b in breaks {
731 if b > start && b <= points.len() {
732 out.push(&points[start..b]);
733 start = b;
734 }
735 }
736 if start < points.len() {
737 out.push(&points[start..]);
738 }
739 out
740}
741
742#[derive(Clone, Debug)]
744pub struct HistogramBin {
745 pub center: f64,
746 pub count: f64,
747}
748
749#[derive(Clone, Debug)]
752pub struct HistogramGroup {
753 pub name: String,
754 pub counts: Vec<f64>,
755}
756
757#[derive(Clone, Debug)]
759pub struct HistogramData {
760 pub column: String,
761 pub bins: Vec<HistogramBin>,
763 pub groups: Vec<HistogramGroup>,
765 pub other: bool,
767 pub share: bool,
769 pub x_min: f64,
770 pub x_max: f64,
771 pub max_count: f64,
772 pub rows: RowsRead,
773 pub clipped: Option<Clipped>,
774}
775
776#[derive(Clone, Debug)]
778pub struct KdeSeries {
779 pub name: String,
780 pub points: Vec<(f64, f64)>,
781}
782
783#[derive(Clone, Debug)]
784pub struct KdeData {
785 pub series: Vec<KdeSeries>,
786 pub other: bool,
788 pub x_min: f64,
789 pub x_max: f64,
790 pub y_max: f64,
791 pub rows: RowsRead,
792 pub clipped: Option<Clipped>,
793}
794
795#[derive(Clone, Debug)]
797pub struct BoxPlotStats {
798 pub name: String,
799 pub min: f64,
800 pub q1: f64,
801 pub median: f64,
802 pub q3: f64,
803 pub max: f64,
804}
805
806#[derive(Clone, Debug)]
807pub struct BoxPlotData {
808 pub stats: Vec<BoxPlotStats>,
809 pub y_min: f64,
810 pub y_max: f64,
811 pub rows: RowsRead,
812 pub clipped: Option<Clipped>,
813 pub of: usize,
816}
817
818#[derive(Clone, Debug)]
820pub struct HeatmapData {
821 pub x_column: String,
822 pub y_column: String,
823 pub x_min: f64,
824 pub x_max: f64,
825 pub y_min: f64,
826 pub y_max: f64,
827 pub x_bins: usize,
828 pub y_bins: usize,
829 pub counts: Vec<Vec<f64>>,
830 pub max_count: f64,
831 pub rows: RowsRead,
832}
833
834pub fn prepare_chart_data(
846 lf: &LazyFrame,
847 schema: &Schema,
848 x_column: &str,
849 y_columns: &[String],
850 sampling: &ChartSampling,
851 envelope: bool,
852) -> Result<ChartDataResult> {
853 if y_columns.is_empty() {
854 return Ok(ChartDataResult {
855 series: Vec::new(),
856 breaks: Vec::new(),
857 x_axis_kind: XAxisTemporalKind::Numeric,
858 rows: RowsRead::default(),
859 });
860 }
861
862 let x_dtype = schema
863 .get(x_column)
864 .ok_or_else(|| color_eyre::eyre::eyre!("x column '{}' not in schema", x_column))?;
865 let x_axis_kind = x_axis_temporal_kind(x_dtype);
866
867 let mut counted = None;
868 if envelope
869 && sampling.full_passes
870 && let Some(limit) = sampling.limit.filter(|&n| n > 0)
871 && sampling.known_total.is_none_or(|n| n > limit)
872 {
873 match envelope_series(lf, x_column, x_dtype, y_columns, limit, sampling)? {
874 Envelope::Drawn {
875 series,
876 breaks,
877 rows,
878 steps,
879 } => {
880 return Ok(ChartDataResult {
881 series,
882 breaks,
883 x_axis_kind,
884 rows: RowsRead {
885 total_rows: rows,
886 sample_size: None,
887 envelope_steps: Some(steps),
888 seed: None,
889 },
890 });
891 }
892 Envelope::Fits(rows) => counted = Some(rows),
894 }
895 }
896 let counted_sampling;
897 let sampling = match counted {
898 Some(rows) => {
899 counted_sampling = ChartSampling {
900 known_total: Some(rows),
901 ..sampling.clone()
902 };
903 &counted_sampling
904 }
905 None => sampling,
906 };
907
908 let mut columns = vec![x_column];
909 columns.extend(y_columns.iter().map(String::as_str));
910 let (df, rows) = read_columns(lf, &columns, sampling)?;
911
912 let mut order: Vec<(f64, usize)> = x_values(&df, x_column, x_dtype)?
913 .into_iter()
914 .enumerate()
915 .filter_map(|(i, x)| x.map(|x| (x, i)))
916 .collect();
917 order.sort_by(|a, b| a.0.total_cmp(&b.0));
919
920 let mut series = Vec::with_capacity(y_columns.len());
921 let mut breaks = Vec::with_capacity(y_columns.len());
922 for y_column in y_columns {
923 let ys = f64_values(&df, y_column)?;
924 let mut points = Vec::with_capacity(order.len());
925 let mut starts = Vec::new();
926 let mut gap = false;
927 for &(x, i) in &order {
928 match ys[i] {
929 Some(y) => {
930 if gap && !points.is_empty() {
931 starts.push(points.len());
932 }
933 gap = false;
934 points.push((x, y));
935 }
936 None => gap = true,
937 }
938 }
939 series.push(points);
940 breaks.push(starts);
941 }
942
943 Ok(ChartDataResult {
944 series,
945 breaks,
946 x_axis_kind,
947 rows,
948 })
949}
950
951enum Envelope {
953 Drawn {
956 series: Vec<Vec<(f64, f64)>>,
957 breaks: Vec<Vec<usize>>,
958 rows: usize,
959 steps: usize,
960 },
961 Fits(usize),
963}
964
965const ENVELOPE_CANCELLED: &str = "chart cancelled";
967
968fn until_cancelled(e: Expr, cancel: &Arc<AtomicBool>) -> Expr {
971 let cancel = Arc::clone(cancel);
972 e.map(
973 move |c: Column| {
974 polars_ensure!(!cancel.load(Ordering::Relaxed), ComputeError: ENVELOPE_CANCELLED);
975 Ok(c)
976 },
977 |_, field| Ok(field.clone()),
978 )
979}
980
981fn envelope_pass(lf: LazyFrame, cancel: &Arc<AtomicBool>) -> Result<DataFrame> {
984 crate::statistics::collect_lazy(lf, true).map_err(|e| {
985 if cancel.load(Ordering::Relaxed) {
986 color_eyre::eyre::eyre!(ENVELOPE_CANCELLED)
987 } else {
988 e.into()
989 }
990 })
991}
992
993fn envelope_series(
998 lf: &LazyFrame,
999 x_column: &str,
1000 x_dtype: &DataType,
1001 y_columns: &[String],
1002 limit: usize,
1003 sampling: &ChartSampling,
1004) -> Result<Envelope> {
1005 let cancel = &sampling.cancel;
1006 let x = match x_dtype {
1008 DataType::Datetime(_, _) | DataType::Date | DataType::Time | DataType::Duration(_) => {
1009 col(x_column).cast(DataType::Int64).cast(DataType::Float64)
1010 }
1011 _ => col(x_column).cast(DataType::Float64),
1012 };
1013 let finite = |e: Expr| {
1016 when(e.clone().is_finite())
1017 .then(e)
1018 .otherwise(lit(NULL).cast(DataType::Float64))
1019 };
1020 let x = finite(until_cancelled(x, cancel)).alias("__x");
1021 let bounds = envelope_pass(
1022 lf.clone().select([
1023 len().alias("rows"),
1024 x.clone().min().alias("lo"),
1025 x.clone().max().alias("hi"),
1026 ]),
1027 cancel,
1028 )?;
1029 let rows = bounds
1030 .column("rows")?
1031 .cast(&DataType::UInt64)?
1032 .u64()?
1033 .get(0)
1034 .unwrap_or(0) as usize;
1035 if rows <= limit {
1036 return Ok(Envelope::Fits(rows));
1037 }
1038 let steps = (limit / 2).max(1);
1039 let n = y_columns.len();
1040 let drawn = |series, breaks| Envelope::Drawn {
1041 series,
1042 breaks,
1043 rows,
1044 steps,
1045 };
1046 let bound = |name: &str| -> Result<Option<f64>> { Ok(bounds.column(name)?.f64()?.get(0)) };
1047 let (Some(lo), Some(hi)) = (bound("lo")?, bound("hi")?) else {
1048 return Ok(drawn(vec![Vec::new(); n], vec![Vec::new(); n]));
1049 };
1050 let per_x = if hi > lo {
1051 steps as f64 / (hi - lo)
1052 } else {
1053 0.0
1054 };
1055 let lf = lf
1056 .clone()
1057 .select(
1058 std::iter::once(x)
1059 .chain(y_columns.iter().enumerate().map(|(i, y)| {
1060 finite(col(y.as_str()).cast(DataType::Float64)).alias(format!("__y{i}"))
1061 }))
1062 .collect::<Vec<_>>(),
1063 )
1064 .filter(col("__x").is_not_null());
1067 let step = ((col("__x") - lit(lo)) * lit(per_x))
1068 .floor()
1069 .cast(DataType::Int64)
1070 .clip(lit(0i64), lit(steps as i64 - 1))
1071 .alias("__step");
1072 let mut aggs = vec![col("__x").min()];
1073 for i in 0..n {
1074 let y = col(format!("__y{i}"));
1075 aggs.push(y.clone().min().alias(format!("__lo{i}")));
1076 aggs.push(y.max().alias(format!("__hi{i}")));
1077 }
1078 let df = envelope_pass(
1079 lf.group_by([step])
1080 .agg(aggs)
1081 .sort(["__step"], Default::default()),
1082 cancel,
1083 )?;
1084 let xs = df.column("__x")?.f64()?.clone();
1085 let mut series = Vec::with_capacity(n);
1086 let mut breaks = Vec::with_capacity(n);
1087 for i in 0..n {
1088 let lows = df.column(&format!("__lo{i}"))?.f64()?.clone();
1089 let highs = df.column(&format!("__hi{i}"))?.f64()?.clone();
1090 let mut points = Vec::with_capacity(xs.len() * 2);
1091 let mut starts = Vec::new();
1092 let mut gap = false;
1093 for ((x, low), high) in xs.iter().zip(lows.iter()).zip(highs.iter()) {
1094 let (Some(x), Some(low), Some(high)) = (x, low, high) else {
1095 gap = true;
1096 continue;
1097 };
1098 if gap && !points.is_empty() {
1099 starts.push(points.len());
1100 }
1101 gap = false;
1102 points.push((x, low));
1103 if high != low {
1104 points.push((x, high));
1105 }
1106 }
1107 series.push(points);
1108 breaks.push(starts);
1109 }
1110 Ok(drawn(series, breaks))
1111}
1112
1113fn read_values(
1115 lf: &LazyFrame,
1116 columns: &[&str],
1117 sampling: &ChartSampling,
1118) -> Result<(Vec<Vec<f64>>, RowsRead)> {
1119 let (df, rows) = read_columns(lf, columns, sampling)?;
1120 let values = columns
1121 .iter()
1122 .map(|c| Ok(f64_values(&df, c)?.into_iter().flatten().collect()))
1123 .collect::<Result<Vec<Vec<f64>>>>()?;
1124 Ok((values, rows))
1125}
1126
1127fn sort_and_clip(values: &mut Vec<f64>, range: ValueRange) -> usize {
1129 values.sort_by(f64::total_cmp);
1130 let Some((low, high)) = range.quantiles() else {
1131 return 0;
1132 };
1133 if values.is_empty() {
1134 return 0;
1135 }
1136 let (low, high) = (quantile(values, low), quantile(values, high));
1137 let before = values.len();
1138 values.retain(|v| (low..=high).contains(v));
1139 before - values.len()
1140}
1141
1142fn clipped(range: ValueRange, outside: usize) -> Option<Clipped> {
1143 (range != ValueRange::All).then_some(Clipped { range, outside })
1144}
1145
1146pub fn prepare_histogram_data(
1148 lf: &LazyFrame,
1149 column: &str,
1150 bins: usize,
1151 range: ValueRange,
1152 sampling: &ChartSampling,
1153) -> Result<HistogramData> {
1154 prepare_histogram_by(lf, column, bins, range, false, None, sampling)
1155}
1156
1157pub fn prepare_histogram_by(
1162 lf: &LazyFrame,
1163 column: &str,
1164 bins: usize,
1165 range: ValueRange,
1166 share: bool,
1167 color: Option<ColorSplit<'_>>,
1168 sampling: &ChartSampling,
1169) -> Result<HistogramData> {
1170 let (values, rows) = read_split(lf, column, color, sampling)?;
1171 let mut all: Vec<f64> = values.iter().map(|(v, _)| *v).collect();
1172 let outside = sort_and_clip(&mut all, range);
1173 let clipped = clipped(range, outside);
1174 let mut data = HistogramData {
1175 column: column.to_string(),
1176 bins: Vec::new(),
1177 groups: Vec::new(),
1178 other: false,
1179 share,
1180 x_min: 0.0,
1181 x_max: 1.0,
1182 max_count: 0.0,
1183 rows,
1184 clipped,
1185 };
1186 let (Some(&lo), Some(&hi)) = (all.first(), all.last()) else {
1187 return Ok(data);
1188 };
1189 let span = hi - lo;
1190 let bin_count = if span <= f64::EPSILON { 1 } else { bins.max(1) };
1191 let bin_width = if span <= f64::EPSILON {
1192 1.0
1193 } else {
1194 span / bin_count as f64
1195 };
1196 (data.x_min, data.x_max) = if span <= f64::EPSILON {
1197 (lo - 0.5, hi + 0.5)
1198 } else {
1199 (lo, hi)
1200 };
1201 let bin_of = |v: f64| {
1202 if span <= f64::EPSILON {
1203 0
1204 } else {
1205 (((v - lo) / bin_width).floor().max(0.0) as usize).min(bin_count - 1)
1206 }
1207 };
1208 let groups = color.map_or(0, |c| c.series());
1209 let mut total = vec![0.0_f64; bin_count];
1210 let mut by_group = vec![vec![0.0_f64; bin_count]; groups];
1211 for (v, group) in values {
1212 if !(lo..=hi).contains(&v) || (color.is_some() && group.is_none()) {
1214 continue;
1215 }
1216 let bin = bin_of(v);
1217 total[bin] += 1.0;
1218 if let Some(g) = group {
1219 by_group[g][bin] += 1.0;
1220 }
1221 }
1222 let as_share = |counts: &mut Vec<f64>| {
1223 let n: f64 = counts.iter().sum();
1224 if share && n > 0.0 {
1225 counts.iter_mut().for_each(|c| *c /= n);
1226 }
1227 };
1228 as_share(&mut total);
1229 by_group.iter_mut().for_each(as_share);
1230 let center = |i: usize| {
1231 if span <= f64::EPSILON {
1232 lo
1233 } else {
1234 lo + (i as f64 + 0.5) * bin_width
1235 }
1236 };
1237 data.bins = total
1238 .iter()
1239 .enumerate()
1240 .map(|(i, &count)| HistogramBin {
1241 center: center(i),
1242 count,
1243 })
1244 .collect();
1245 let max = |counts: &[f64]| counts.iter().copied().fold(0.0_f64, f64::max);
1246 if let Some(color) = color {
1247 data.groups = color
1248 .names()
1249 .into_iter()
1250 .zip(by_group)
1251 .map(|(name, counts)| HistogramGroup { name, counts })
1252 .collect();
1253 data.other = color.other;
1254 data.max_count = data
1255 .groups
1256 .iter()
1257 .map(|g| max(&g.counts))
1258 .fold(0.0, f64::max);
1259 } else {
1260 data.max_count = max(&total);
1261 }
1262 Ok(data)
1263}
1264
1265fn quantile(sorted: &[f64], q: f64) -> f64 {
1266 if sorted.is_empty() {
1267 return 0.0;
1268 }
1269 let n = sorted.len();
1270 if n == 1 {
1271 return sorted[0];
1272 }
1273 let pos = q.clamp(0.0, 1.0) * (n as f64 - 1.0);
1274 let idx = pos.floor() as usize;
1275 let next = pos.ceil() as usize;
1276 if idx == next {
1277 sorted[idx]
1278 } else {
1279 let lower = sorted[idx];
1280 let upper = sorted[next];
1281 let weight = pos - idx as f64;
1282 lower + (upper - lower) * weight
1283 }
1284}
1285
1286fn box_stats(name: String, values: &[f64]) -> Option<BoxPlotStats> {
1288 let (min, max) = (*values.first()?, *values.last()?);
1289 Some(BoxPlotStats {
1290 name,
1291 min,
1292 q1: quantile(values, 0.25),
1293 median: quantile(values, 0.5),
1294 q3: quantile(values, 0.75),
1295 max,
1296 })
1297}
1298
1299fn box_data(stats: Vec<BoxPlotStats>, rows: RowsRead, clipped: Option<Clipped>) -> BoxPlotData {
1301 let mut y_min = stats.iter().map(|s| s.min).fold(f64::INFINITY, f64::min);
1302 let mut y_max = stats
1303 .iter()
1304 .map(|s| s.max)
1305 .fold(f64::NEG_INFINITY, f64::max);
1306 if stats.is_empty() {
1307 (y_min, y_max) = (0.0, 1.0);
1308 } else if y_max <= y_min {
1309 y_max = y_min + 1.0;
1310 }
1311 BoxPlotData {
1312 stats,
1313 y_min,
1314 y_max,
1315 rows,
1316 clipped,
1317 of: 0,
1318 }
1319}
1320
1321pub fn prepare_box_plot_data<T: AsRef<str>>(
1323 lf: &LazyFrame,
1324 columns: &[T],
1325 range: ValueRange,
1326 sampling: &ChartSampling,
1327) -> Result<BoxPlotData> {
1328 let col_refs: Vec<&str> = columns.iter().map(|c| c.as_ref()).collect();
1329 let (columns_values, rows) = if col_refs.is_empty() {
1330 (Vec::new(), RowsRead::default())
1331 } else {
1332 read_values(lf, &col_refs, sampling)?
1333 };
1334 let mut stats = Vec::new();
1335 let mut outside = 0;
1336 for (column, mut values) in col_refs.iter().zip(columns_values) {
1337 outside += sort_and_clip(&mut values, range);
1338 stats.extend(box_stats((*column).to_string(), &values));
1339 }
1340 Ok(box_data(stats, rows, clipped(range, outside)))
1341}
1342
1343pub fn prepare_box_by(
1346 lf: &LazyFrame,
1347 column: &str,
1348 by: ColorSplit<'_>,
1349 range: ValueRange,
1350 sampling: &ChartSampling,
1351) -> Result<BoxPlotData> {
1352 let (values, rows) = read_split(lf, column, Some(by), sampling)?;
1353 let mut groups = vec![Vec::new(); by.groups.len()];
1354 for (v, group) in values {
1355 if let Some(g) = group {
1356 groups[g].push(v);
1357 }
1358 }
1359 let mut outside = 0;
1360 let mut stats = Vec::new();
1361 for (name, mut values) in by.groups.iter().zip(groups) {
1362 outside += sort_and_clip(&mut values, range);
1363 stats.extend(box_stats(group_label(name), &values));
1364 }
1365 Ok(box_data(stats, rows, clipped(range, outside)))
1366}
1367
1368fn kde_bandwidth(values: &[f64]) -> f64 {
1369 if values.len() <= 1 {
1370 return 1.0;
1371 }
1372 let n = values.len() as f64;
1373 let mean = values.iter().sum::<f64>() / n;
1374 let var = values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n;
1375 let std = var.sqrt();
1376 if std <= f64::EPSILON {
1377 return 1.0;
1378 }
1379 1.06 * std * n.powf(-0.2)
1380}
1381
1382fn kde_series(name: String, values: &[f64], bandwidth_factor: f64) -> Option<KdeSeries> {
1385 let (min, max) = (*values.first()?, *values.last()?);
1386 let bandwidth = (kde_bandwidth(values) * bandwidth_factor).max(f64::EPSILON);
1387 let x_start = min - 3.0 * bandwidth;
1388 let x_end = max + 3.0 * bandwidth;
1389 let samples = 200_usize;
1390 let step = (x_end - x_start) / (samples.saturating_sub(1).max(1) as f64);
1391 let inv = 1.0 / ((values.len() as f64) * bandwidth * (2.0 * PI).sqrt());
1392 let points = (0..samples)
1393 .map(|i| {
1394 let x = x_start + i as f64 * step;
1395 let sum: f64 = values
1396 .iter()
1397 .map(|&v| {
1398 let u = (x - v) / bandwidth;
1399 (-0.5 * u * u).exp()
1400 })
1401 .sum();
1402 (x, inv * sum)
1403 })
1404 .collect();
1405 Some(KdeSeries { name, points })
1406}
1407
1408fn kde_data(series: Vec<KdeSeries>, rows: RowsRead, clipped: Option<Clipped>) -> KdeData {
1410 let points = || series.iter().flat_map(|s| s.points.iter());
1411 let mut x_min = points().map(|p| p.0).fold(f64::INFINITY, f64::min);
1412 let mut x_max = points().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max);
1413 let mut y_max = points().map(|p| p.1).fold(f64::NEG_INFINITY, f64::max);
1414 if series.is_empty() {
1415 (x_min, x_max, y_max) = (0.0, 1.0, 1.0);
1416 }
1417 if x_max <= x_min {
1418 x_max = x_min + 1.0;
1419 }
1420 if y_max <= 0.0 {
1421 y_max = 1.0;
1422 }
1423 KdeData {
1424 series,
1425 other: false,
1426 x_min,
1427 x_max,
1428 y_max,
1429 rows,
1430 clipped,
1431 }
1432}
1433
1434pub fn prepare_kde_data<T: AsRef<str>>(
1436 lf: &LazyFrame,
1437 columns: &[T],
1438 bandwidth_factor: f64,
1439 range: ValueRange,
1440 sampling: &ChartSampling,
1441) -> Result<KdeData> {
1442 let col_refs: Vec<&str> = columns.iter().map(|c| c.as_ref()).collect();
1443 let (columns_values, rows) = if col_refs.is_empty() {
1444 (Vec::new(), RowsRead::default())
1445 } else {
1446 read_values(lf, &col_refs, sampling)?
1447 };
1448 let mut series = Vec::new();
1449 let mut outside = 0;
1450 for (column, mut values) in col_refs.iter().zip(columns_values) {
1451 outside += sort_and_clip(&mut values, range);
1452 series.extend(kde_series((*column).to_string(), &values, bandwidth_factor));
1453 }
1454 Ok(kde_data(series, rows, clipped(range, outside)))
1455}
1456
1457pub fn prepare_kde_by(
1460 lf: &LazyFrame,
1461 column: &str,
1462 bandwidth_factor: f64,
1463 range: ValueRange,
1464 color: ColorSplit<'_>,
1465 sampling: &ChartSampling,
1466) -> Result<KdeData> {
1467 let (values, rows) = read_split(lf, column, Some(color), sampling)?;
1468 let mut all: Vec<f64> = values.iter().map(|(v, _)| *v).collect();
1469 let outside = sort_and_clip(&mut all, range);
1470 let (lo, hi) = match (all.first(), all.last()) {
1471 (Some(&lo), Some(&hi)) => (lo, hi),
1472 _ => (f64::INFINITY, f64::NEG_INFINITY),
1473 };
1474 let mut groups = vec![Vec::new(); color.series()];
1475 for (v, group) in values {
1476 if let Some(g) = group
1477 && (lo..=hi).contains(&v)
1478 {
1479 groups[g].push(v);
1480 }
1481 }
1482 let last = color.series().saturating_sub(1);
1483 let mut other = false;
1484 let series = color
1485 .names()
1486 .into_iter()
1487 .zip(groups)
1488 .enumerate()
1489 .filter_map(|(i, (name, mut values))| {
1490 values.sort_by(f64::total_cmp);
1491 let series = kde_series(name, &values, bandwidth_factor)?;
1492 other = color.other && i == last;
1493 Some(series)
1494 })
1495 .collect();
1496 Ok(KdeData {
1497 other,
1498 ..kde_data(series, rows, clipped(range, outside))
1499 })
1500}
1501
1502pub fn prepare_heatmap_data(
1504 lf: &LazyFrame,
1505 x_column: &str,
1506 y_column: &str,
1507 bins: usize,
1508 sampling: &ChartSampling,
1509) -> Result<HeatmapData> {
1510 let (df, rows) = read_columns(lf, &[x_column, y_column], sampling)?;
1511 let pairs: Vec<(f64, f64)> = f64_values(&df, x_column)?
1512 .into_iter()
1513 .zip(f64_values(&df, y_column)?)
1514 .filter_map(|(x, y)| Some((x?, y?)))
1515 .collect();
1516 let x_bins = bins.max(1);
1517 let y_bins = bins.max(1);
1518 if pairs.is_empty() {
1519 return Ok(HeatmapData {
1520 x_column: x_column.to_string(),
1521 y_column: y_column.to_string(),
1522 x_min: 0.0,
1523 x_max: 1.0,
1524 y_min: 0.0,
1525 y_max: 1.0,
1526 x_bins,
1527 y_bins,
1528 counts: vec![vec![0.0; x_bins]; y_bins],
1529 max_count: 0.0,
1530 rows,
1531 });
1532 }
1533 let mut x_min = f64::INFINITY;
1534 let mut x_max = f64::NEG_INFINITY;
1535 let mut y_min = f64::INFINITY;
1536 let mut y_max = f64::NEG_INFINITY;
1537 for (x, y) in &pairs {
1538 x_min = x_min.min(*x);
1539 x_max = x_max.max(*x);
1540 y_min = y_min.min(*y);
1541 y_max = y_max.max(*y);
1542 }
1543 if x_max <= x_min {
1544 x_max = x_min + 1.0;
1545 }
1546 if y_max <= y_min {
1547 y_max = y_min + 1.0;
1548 }
1549 let mut counts = vec![vec![0.0_f64; x_bins]; y_bins];
1550 let x_range = x_max - x_min;
1551 let y_range = y_max - y_min;
1552 for (x, y) in pairs {
1553 let xi =
1554 (((x - x_min) / x_range * x_bins as f64).floor().max(0.0) as usize).min(x_bins - 1);
1555 let yi =
1556 (((y - y_min) / y_range * y_bins as f64).floor().max(0.0) as usize).min(y_bins - 1);
1557 counts[yi][xi] += 1.0;
1558 }
1559 let max_count = counts
1560 .iter()
1561 .flat_map(|row| row.iter())
1562 .cloned()
1563 .fold(0.0_f64, f64::max);
1564 Ok(HeatmapData {
1565 x_column: x_column.to_string(),
1566 y_column: y_column.to_string(),
1567 x_min,
1568 x_max,
1569 y_min,
1570 y_max,
1571 x_bins,
1572 y_bins,
1573 counts,
1574 max_count,
1575 rows,
1576 })
1577}
1578
1579#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
1581#[serde(rename_all = "snake_case")]
1582pub enum BarOrder {
1583 #[default]
1585 Value,
1586 Label,
1588}
1589
1590impl BarOrder {
1591 pub const ALL: [Self; 2] = [Self::Value, Self::Label];
1592
1593 pub fn label(self) -> &'static str {
1594 match self {
1595 Self::Value => "Value",
1596 Self::Label => "Label",
1597 }
1598 }
1599}
1600
1601pub const BAR_CAP: usize = 100;
1604
1605pub fn is_category_dtype(dtype: &DataType) -> bool {
1608 matches!(
1609 dtype,
1610 DataType::String | DataType::Categorical(_, _) | DataType::Enum(_, _) | DataType::Boolean
1611 ) || dtype.is_integer()
1612}
1613
1614#[derive(Clone, Debug, PartialEq, Eq)]
1617pub enum BarValue {
1618 Count,
1619 Column(String),
1620}
1621
1622impl BarValue {
1623 pub fn label(&self) -> &str {
1624 match self {
1625 Self::Count => "Count",
1626 Self::Column(column) => column,
1627 }
1628 }
1629}
1630
1631pub const COUNT_CATEGORY_CAP: usize = 100_000;
1634
1635#[derive(Clone, Debug, PartialEq)]
1639pub struct Bar {
1640 pub label: Option<String>,
1641 pub value: f64,
1642 pub by_group: Vec<Option<f64>>,
1643}
1644
1645#[derive(Clone, Debug)]
1647pub struct BarData {
1648 pub category: String,
1649 pub value_column: String,
1650 pub bars: Vec<Bar>,
1651 pub more: usize,
1653 pub no_value: usize,
1655 pub rows: RowsRead,
1656 pub value_dtype: DataType,
1658 pub counted: Option<usize>,
1661 pub groups: Vec<String>,
1663 pub other: bool,
1665 pub rows_note: Option<String>,
1667}
1668
1669impl BarData {
1670 pub fn value_format(
1673 &self,
1674 settings: &crate::numfmt::NumberFormatSettings,
1675 ) -> crate::numfmt::NumberFormat {
1676 table_number_format(settings, &self.value_column, &self.value_dtype)
1677 }
1678
1679 pub fn value_labels(&self, settings: &crate::numfmt::NumberFormatSettings) -> Vec<String> {
1681 self.labels_in(&self.value_format(settings))
1682 }
1683
1684 pub fn labels_in(&self, format: &crate::numfmt::NumberFormat) -> Vec<String> {
1686 let integer = self.value_dtype.is_integer();
1687 self.bars
1688 .iter()
1689 .map(|b| format_bar_value(b.value, integer, format))
1690 .collect()
1691 }
1692}
1693
1694pub fn format_bar_value(v: f64, integer: bool, format: &crate::numfmt::NumberFormat) -> String {
1699 let places = format.float_precision.unwrap_or(2);
1700 let smallest = 0.5 * 10f64.powi(-i32::from(places));
1701 if !v.is_finite() || v.abs() >= 1e15 || (!integer && v != 0.0 && v.abs() < smallest) {
1702 return scientific(v, 2, format.decimal_sep);
1703 }
1704 let mut out = String::new();
1705 if integer {
1706 format.write_i64(v as i64, &mut out);
1707 } else {
1708 let fixed = crate::numfmt::NumberFormat {
1709 float_precision: Some(places),
1710 ..format.clone()
1711 };
1712 fixed.write_f64(v, &mut String::new(), &mut out);
1713 }
1714 out
1715}
1716
1717fn sql_ident(name: &str) -> String {
1720 let plain = name
1721 .chars()
1722 .next()
1723 .is_some_and(|c| c.is_ascii_lowercase() || c == '_')
1724 && name
1725 .chars()
1726 .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_');
1727 if plain {
1728 name.to_string()
1729 } else {
1730 format!("\"{}\"", name.replace('"', "\"\""))
1731 }
1732}
1733
1734pub fn prepare_bar_data(
1739 lf: &LazyFrame,
1740 category: &str,
1741 value: &str,
1742 order: BarOrder,
1743 cap: usize,
1744 sampling: &ChartSampling,
1745) -> Result<BarData> {
1746 let (df, rows) = read_columns(lf, &[category, value], sampling)?;
1747 let categories = df.column(category)?.as_materialized_series().clone();
1748 let labels_series = crate::past_calendar::cast_text(&categories, CastOptions::NonStrict)?;
1749 let labels: Vec<Option<&str>> = labels_series.str()?.iter().collect();
1750
1751 let mut seen: std::collections::HashMap<Option<&str>, usize> =
1752 std::collections::HashMap::with_capacity(labels.len());
1753 for label in &labels {
1754 *seen.entry(*label).or_default() += 1;
1755 }
1756 if seen.len() < labels.len() {
1757 let read = match rows.sample_size {
1758 Some(n) => format!("a sample of {} rows", crate::numfmt::group_chrome(n)),
1759 None => format!("{} rows", crate::numfmt::group_chrome(labels.len())),
1760 };
1761 let (c, v) = (sql_ident(category), sql_ident(value));
1762 let q = if c == category && v == value {
1764 format!(" (or select avg {value} by {category})")
1765 } else {
1766 String::new()
1767 };
1768 return Err(color_eyre::eyre::eyre!(
1769 "{category} repeats: {} categories in {read}. A bar takes one row per category, \
1770 so group first: SELECT {c}, AVG({v}) FROM df GROUP BY {c}{q}, or choose Count \
1771 for the rows per category",
1772 crate::numfmt::group_chrome(seen.len()),
1773 ));
1774 }
1775
1776 let values = f64_values(&df, value)?;
1777 let (bars, more, no_value) = order_bars(&categories, &labels, &values, order, cap);
1778 Ok(BarData {
1779 category: category.to_string(),
1780 value_column: value.to_string(),
1781 bars,
1782 more,
1783 no_value,
1784 rows,
1785 value_dtype: df.column(value)?.dtype().clone(),
1786 counted: None,
1787 groups: Vec::new(),
1788 other: false,
1789 rows_note: None,
1790 })
1791}
1792
1793fn order_bars(
1796 categories: &Series,
1797 labels: &[Option<&str>],
1798 values: &[Option<f64>],
1799 order: BarOrder,
1800 cap: usize,
1801) -> (Vec<Bar>, usize, usize) {
1802 let row_order: Vec<usize> = match order {
1803 BarOrder::Value => (0..labels.len()).collect(),
1804 BarOrder::Label => label_order(categories),
1805 };
1806 let mut no_value = 0;
1807 let mut bars: Vec<Bar> = row_order
1808 .into_iter()
1809 .filter_map(|i| match values[i] {
1810 Some(value) => Some(Bar {
1811 label: labels[i].map(str::to_string),
1812 value,
1813 by_group: Vec::new(),
1814 }),
1815 None => {
1816 no_value += 1;
1817 None
1818 }
1819 })
1820 .collect();
1821 if order == BarOrder::Value {
1822 bars.sort_by(|a, b| b.value.total_cmp(&a.value));
1824 }
1825 let more = bars.len().saturating_sub(cap);
1826 bars.truncate(cap);
1827 (bars, more, no_value)
1828}
1829
1830fn label_order(categories: &Series) -> Vec<usize> {
1833 categories
1834 .arg_sort(
1835 SortOptions::default()
1836 .with_nulls_last(true)
1837 .with_maintain_order(true),
1838 )
1839 .iter()
1840 .flatten()
1841 .map(|i| i as usize)
1842 .collect()
1843}
1844
1845pub fn prepare_bar_counts(
1851 lf: &LazyFrame,
1852 category: &str,
1853 order: BarOrder,
1854 cap: usize,
1855 sampling: &ChartSampling,
1856) -> Result<BarData> {
1857 count_bars(lf, category, order, cap, COUNT_CATEGORY_CAP, sampling)
1858}
1859
1860fn count_bars(
1861 lf: &LazyFrame,
1862 category: &str,
1863 order: BarOrder,
1864 cap: usize,
1865 max_categories: usize,
1866 sampling: &ChartSampling,
1867) -> Result<BarData> {
1868 let counted = match held_counts(sampling, category, max_categories)? {
1869 Some(counted) => counted,
1870 None => {
1871 let fits = sampling
1874 .limit
1875 .zip(sampling.known_total)
1876 .is_some_and(|(n, total)| total <= n);
1877 let whole = if fits {
1878 let (df, rows) = read_columns(lf, &[category], sampling)?;
1879 rows.sample_size.is_none().then_some(df)
1880 } else {
1881 None
1882 };
1883 let counted = match whole {
1884 Some(df) => count_frame(&df, category, max_categories)?,
1885 None => stream_counts(lf, category, max_categories, &sampling.cancel)?,
1886 };
1887 hold_counts(sampling, category, &counted);
1888 counted
1889 }
1890 };
1891 let (counts, total) = match counted {
1892 Counted::All { counts, rows } => (counts, rows),
1893 Counted::TooMany => {
1894 return Err(color_eyre::eyre::eyre!(
1895 "more than {} categories of {category}: counting stopped. Count by a \
1896 column with fewer values",
1897 crate::numfmt::group_chrome(max_categories)
1898 ));
1899 }
1900 };
1901 let data = |bars, more| BarData {
1902 category: category.to_string(),
1903 value_column: "count".to_string(),
1904 bars,
1905 more,
1906 no_value: 0,
1907 rows: RowsRead {
1908 total_rows: total,
1909 sample_size: None,
1910 envelope_steps: None,
1911 seed: None,
1912 },
1913 value_dtype: DataType::UInt64,
1914 counted: sampling.limit.is_some_and(|n| total > n).then_some(total),
1915 groups: Vec::new(),
1916 other: false,
1917 rows_note: None,
1918 };
1919 let Some(counts) = counts else {
1920 return Ok(data(Vec::new(), 0));
1921 };
1922 let by_label: Vec<IdxSize> = label_order(counts.column(category)?.as_materialized_series())
1924 .into_iter()
1925 .map(|i| i as IdxSize)
1926 .collect();
1927 let counts = counts.take(&IdxCa::from_vec("order".into(), by_label))?;
1928 let categories = counts.column(category)?.as_materialized_series().clone();
1929 let labels_series = crate::past_calendar::cast_text(&categories, CastOptions::NonStrict)?;
1930 let labels: Vec<Option<&str>> = labels_series.str()?.iter().collect();
1931 let values: Vec<Option<f64>> = counts
1932 .column(COUNT_COLUMN)?
1933 .u64()?
1934 .iter()
1935 .map(|n| n.map(|n| n as f64))
1936 .collect();
1937 let (bars, more, _) = order_bars(&categories, &labels, &values, order, cap);
1938 Ok(data(bars, more))
1939}
1940
1941fn held_counts(
1944 sampling: &ChartSampling,
1945 category: &str,
1946 max_categories: usize,
1947) -> Result<Option<Counted>> {
1948 let holding = sampling.held.0.lock().unwrap_or_else(|e| e.into_inner());
1949 if let Some(held) = holding.counts.iter().find(|h| h.category == category) {
1950 return Ok(Some(held.counted.clone()));
1951 }
1952 let Some(whole) = holding
1953 .rows
1954 .as_ref()
1955 .filter(|h| h.rows.sample_size.is_none() && h.df.column(category).is_ok())
1956 else {
1957 return Ok(None);
1958 };
1959 let counted = count_frame(&whole.df, category, max_categories)?;
1960 drop(holding);
1961 hold_counts(sampling, category, &counted);
1962 Ok(Some(counted))
1963}
1964
1965fn hold_counts(sampling: &ChartSampling, category: &str, counted: &Counted) {
1966 let mut holding = sampling.held.0.lock().unwrap_or_else(|e| e.into_inner());
1967 holding.counts.retain(|h| h.category != category);
1968 if holding.counts.len() >= HELD_COUNTS {
1969 holding.counts.remove(0);
1970 }
1971 holding.counts.push(HeldCounts {
1972 category: category.to_string(),
1973 counted: counted.clone(),
1974 });
1975}
1976
1977pub(crate) fn count_frame(
1979 df: &DataFrame,
1980 category: &str,
1981 max_categories: usize,
1982) -> Result<Counted> {
1983 let mut tally = Tally::new(category, max_categories);
1984 tally.observe(&df.select([category])?)?;
1985 Ok(tally.finish()?)
1986}
1987
1988#[derive(Clone)]
1991pub(crate) enum Counted {
1992 All {
1993 counts: Option<DataFrame>,
1994 rows: usize,
1995 },
1996 TooMany,
1997}
1998
1999pub(crate) const COUNT_COLUMN: &str = "__datui_bar_count";
2000
2001const HELD_COUNTS: usize = 4;
2004
2005const MERGE_AFTER: usize = 1 << 16;
2008
2009pub(crate) struct Tally {
2012 category: PlSmallStr,
2013 max: usize,
2014 counts: Option<DataFrame>,
2015 merged: usize,
2017 rows: usize,
2018 too_many: bool,
2019 cancelled: bool,
2021}
2022
2023impl Tally {
2024 pub(crate) fn new(category: &str, max: usize) -> Self {
2025 Self {
2026 category: category.into(),
2027 max,
2028 counts: None,
2029 merged: 0,
2030 rows: 0,
2031 too_many: false,
2032 cancelled: false,
2033 }
2034 }
2035
2036 pub(crate) fn observe(&mut self, batch: &DataFrame) -> PolarsResult<bool> {
2039 if self.too_many {
2040 return Ok(true);
2041 }
2042 self.rows += batch.height();
2043 let part = group_counts(batch, &self.category, false)?;
2044 let mut counts = match self.counts.take() {
2045 Some(mut counts) => {
2046 counts.vstack_mut(&part)?;
2047 counts
2048 }
2049 None => part,
2050 };
2051 if counts.height() - self.merged >= self.merged.max(MERGE_AFTER) {
2052 counts = group_counts(&counts, &self.category, true)?;
2053 self.merged = counts.height();
2054 self.too_many = self.merged > self.max;
2055 }
2056 self.counts = Some(counts);
2057 Ok(self.too_many)
2058 }
2059
2060 pub(crate) fn finish(self) -> PolarsResult<Counted> {
2061 let counts = match self.counts {
2062 Some(counts) => Some(group_counts(&counts, &self.category, true)?),
2063 None => None,
2064 };
2065 if self.too_many || counts.as_ref().is_some_and(|c| c.height() > self.max) {
2066 return Ok(Counted::TooMany);
2067 }
2068 Ok(Counted::All {
2069 counts,
2070 rows: self.rows,
2071 })
2072 }
2073}
2074
2075fn group_counts(df: &DataFrame, category: &str, summed: bool) -> PolarsResult<DataFrame> {
2078 let by = df.group_by([category])?;
2079 let groups = by.get_groups();
2080 let counts: Vec<u64> = if summed {
2081 let carried: Vec<u64> = df
2082 .column(COUNT_COLUMN)?
2083 .u64()?
2084 .into_no_null_iter()
2085 .collect();
2086 groups
2087 .iter()
2088 .map(|group| match group {
2089 GroupsIndicator::Idx((_, rows)) => rows.iter().map(|&i| carried[i as usize]).sum(),
2090 GroupsIndicator::Slice([first, len]) => {
2091 carried[first as usize..(first + len) as usize].iter().sum()
2092 }
2093 })
2094 .collect()
2095 } else {
2096 groups.iter().map(|group| group.len() as u64).collect()
2097 };
2098 let mut columns = by.keys();
2099 columns.push(Column::new(COUNT_COLUMN.into(), counts));
2100 DataFrame::new_infer_height(columns)
2101}
2102
2103fn stream_counts(
2107 lf: &LazyFrame,
2108 category: &str,
2109 max: usize,
2110 cancel: &Arc<AtomicBool>,
2111) -> Result<Counted> {
2112 let state = Arc::new(Mutex::new(Tally::new(category, max)));
2113 let callback_state = Arc::clone(&state);
2114 let callback_cancel = Arc::clone(cancel);
2115 let sink = lf.clone().select([col(category)]).sink_batches(
2116 PlanCallback::new(move |batch: DataFrame| {
2117 let mut tally = callback_state
2118 .lock()
2119 .map_err(|_| PolarsError::ComputeError("count lock failed".into()))?;
2120 if callback_cancel.load(Ordering::Relaxed) {
2121 tally.cancelled = true;
2122 return Ok(true);
2123 }
2124 tally.observe(&batch)
2125 }),
2126 false,
2127 None,
2128 )?;
2129 crate::statistics::collect_lazy(sink, true)?;
2131 let tally = std::mem::replace(
2132 &mut *state.lock().unwrap_or_else(|e| e.into_inner()),
2133 Tally::new(category, max),
2134 );
2135 if tally.cancelled {
2138 return Err(color_eyre::eyre::eyre!("count cancelled"));
2139 }
2140 Ok(tally.finish()?)
2141}
2142
2143#[derive(Clone, Copy, Debug)]
2149pub struct ColorSplit<'a> {
2150 pub column: &'a str,
2151 pub groups: &'a [Option<String>],
2152 pub other: bool,
2153}
2154
2155pub const OTHER: &str = "Other";
2157
2158impl ColorSplit<'_> {
2159 pub fn series(&self) -> usize {
2161 self.groups.len() + usize::from(self.other)
2162 }
2163
2164 pub fn names(&self) -> Vec<String> {
2166 let mut names: Vec<String> = self.groups.iter().map(group_label).collect();
2167 if self.other {
2168 names.push(OTHER.to_string());
2169 }
2170 names
2171 }
2172}
2173
2174pub fn group_label(value: &Option<String>) -> String {
2176 value.clone().unwrap_or_else(|| "null".to_string())
2177}
2178
2179fn row_groups(df: &DataFrame, split: ColorSplit<'_>) -> Result<Vec<Option<usize>>> {
2183 let series = df.column(split.column)?.as_materialized_series();
2184 let text = crate::past_calendar::cast_text(series, CastOptions::NonStrict)?;
2185 let index: std::collections::HashMap<Option<&str>, usize> = split
2186 .groups
2187 .iter()
2188 .enumerate()
2189 .map(|(i, g)| (g.as_deref(), i))
2190 .collect();
2191 let other = split.other.then_some(split.groups.len());
2192 Ok(text
2193 .str()?
2194 .iter()
2195 .map(|v| index.get(&v).copied().or(other))
2196 .collect())
2197}
2198
2199type SplitValues = (Vec<(f64, Option<usize>)>, RowsRead);
2201
2202fn read_split(
2205 lf: &LazyFrame,
2206 column: &str,
2207 split: Option<ColorSplit<'_>>,
2208 sampling: &ChartSampling,
2209) -> Result<SplitValues> {
2210 let mut columns = vec![column];
2211 if let Some(split) = split {
2212 columns.push(split.column);
2213 }
2214 let (df, rows) = read_columns(lf, &columns, sampling)?;
2215 let values = f64_values(&df, column)?;
2216 let groups = split.map(|s| row_groups(&df, s)).transpose()?;
2217 let out = values
2218 .into_iter()
2219 .enumerate()
2220 .filter_map(|(i, v)| {
2221 let v = v?;
2222 Some((v, groups.as_ref().and_then(|groups| groups[i])))
2223 })
2224 .collect();
2225 Ok((out, rows))
2226}
2227
2228#[derive(Clone, Debug, Default, PartialEq)]
2231pub struct ValueRows {
2232 pub values: Vec<(Option<String>, u64)>,
2233 pub rows: usize,
2235}
2236
2237pub fn value_rows(lf: &LazyFrame, column: &str, sampling: &ChartSampling) -> Result<ValueRows> {
2240 let counted = match held_counts(sampling, column, COUNT_CATEGORY_CAP)? {
2241 Some(counted) => counted,
2242 None => {
2243 let counted = stream_counts(lf, column, COUNT_CATEGORY_CAP, &sampling.cancel)?;
2244 hold_counts(sampling, column, &counted);
2245 counted
2246 }
2247 };
2248 let (counts, rows) = match counted {
2249 Counted::All { counts, rows } => (counts, rows),
2250 Counted::TooMany => {
2251 return Err(color_eyre::eyre::eyre!(
2252 "more than {} values of {column}: choose a column with fewer",
2253 crate::numfmt::group_chrome(COUNT_CATEGORY_CAP)
2254 ));
2255 }
2256 };
2257 let Some(counts) = counts else {
2258 return Ok(ValueRows {
2259 values: Vec::new(),
2260 rows,
2261 });
2262 };
2263 let by_label: Vec<IdxSize> = label_order(counts.column(column)?.as_materialized_series())
2264 .into_iter()
2265 .map(|i| i as IdxSize)
2266 .collect();
2267 let counts = counts.take(&IdxCa::from_vec("order".into(), by_label))?;
2268 let labels = crate::past_calendar::cast_text(
2269 counts.column(column)?.as_materialized_series(),
2270 CastOptions::NonStrict,
2271 )?;
2272 let mut values: Vec<(Option<String>, u64)> = labels
2273 .str()?
2274 .iter()
2275 .zip(counts.column(COUNT_COLUMN)?.u64()?.iter())
2276 .map(|(label, n)| (label.map(str::to_string), n.unwrap_or(0)))
2277 .collect();
2278 values.sort_by_key(|v| std::cmp::Reverse(v.1));
2280 Ok(ValueRows { values, rows })
2281}
2282
2283pub fn color_groups(
2286 rows: &ValueRows,
2287 picked: &[Option<String>],
2288 most: usize,
2289) -> Vec<Option<String>> {
2290 if !picked.is_empty() {
2291 return picked.iter().take(most).cloned().collect();
2292 }
2293 rows.values
2294 .iter()
2295 .take(most)
2296 .map(|(value, _)| value.clone())
2297 .collect()
2298}
2299
2300fn group_expr(split: ColorSplit<'_>) -> Expr {
2303 let text = crate::past_calendar::text_expr(col(split.column), CastOptions::NonStrict);
2304 let mut out = match split.other {
2305 true => lit(split.groups.len() as u32).cast(DataType::UInt32),
2306 false => lit(NULL).cast(DataType::UInt32),
2307 };
2308 for (i, group) in split.groups.iter().enumerate().rev() {
2309 let matches = match group {
2310 Some(value) => text.clone().eq(lit(value.clone())),
2311 None => col(split.column).is_null(),
2312 };
2313 out = when(matches).then(lit(i as u32)).otherwise(out);
2314 }
2315 out
2316}
2317
2318#[derive(Clone, Debug, Default)]
2320pub struct GroupedSeries {
2321 pub names: Vec<String>,
2322 pub series: Vec<Vec<(f64, f64)>>,
2323 pub breaks: Vec<Vec<usize>>,
2325 pub x_axis_kind: XAxisTemporalKind,
2326 pub rows: RowsRead,
2327 pub other: bool,
2329}
2330
2331pub fn prepare_xy_by(
2334 lf: &LazyFrame,
2335 schema: &Schema,
2336 x: &str,
2337 y: &str,
2338 color: ColorSplit<'_>,
2339 sampling: &ChartSampling,
2340) -> Result<GroupedSeries> {
2341 let x_dtype = schema
2342 .get(x)
2343 .ok_or_else(|| color_eyre::eyre::eyre!("x column '{}' not in schema", x))?;
2344 let (df, rows) = read_columns(lf, &[x, y, color.column], sampling)?;
2345 let xs = x_values(&df, x, x_dtype)?;
2346 let ys = f64_values(&df, y)?;
2347 let groups = row_groups(&df, color)?;
2348 let mut order: Vec<(f64, usize)> = xs
2349 .into_iter()
2350 .enumerate()
2351 .filter_map(|(i, x)| x.map(|x| (x, i)))
2352 .collect();
2353 order.sort_by(|a, b| a.0.total_cmp(&b.0));
2354 let n = color.series();
2355 let mut series = vec![Vec::new(); n];
2356 let mut breaks = vec![Vec::new(); n];
2357 let mut gap = vec![false; n];
2358 for (x, i) in order {
2359 let Some(g) = groups[i] else { continue };
2360 match ys[i] {
2361 Some(y) => {
2362 if gap[g] && !series[g].is_empty() {
2363 breaks[g].push(series[g].len());
2364 }
2365 gap[g] = false;
2366 series[g].push((x, y));
2367 }
2368 None => gap[g] = true,
2369 }
2370 }
2371 Ok(GroupedSeries {
2372 names: color.names(),
2373 series,
2374 breaks,
2375 x_axis_kind: x_axis_temporal_kind(x_dtype),
2376 rows,
2377 other: color.other,
2378 })
2379}
2380
2381pub const AGGREGATE_POINTS_MAX: usize = 200_000;
2384
2385#[derive(Clone, Copy, Debug)]
2387pub struct AggregateSpec<'a> {
2388 pub x: &'a str,
2389 pub time_unit: crate::chart_modal::TimeUnit,
2390 pub ys: &'a [String],
2391 pub aggregate: crate::chart_modal::Aggregate,
2392 pub quantile: u8,
2394 pub cumulative: crate::chart_modal::Cumulative,
2395 pub color: Option<ColorSplit<'a>>,
2396}
2397
2398fn y_values(y: Expr, aggregate: crate::chart_modal::Aggregate) -> Expr {
2401 if aggregate.takes_any_y() {
2402 y
2403 } else {
2404 y.cast(DataType::Float64)
2405 }
2406}
2407
2408const ROW_ORDER: &str = "__i";
2410
2411fn aggregate_expr(values: Expr, aggregate: crate::chart_modal::Aggregate, quantile: u8) -> Expr {
2415 use crate::chart_modal::Aggregate;
2416 let in_order = || {
2417 values
2418 .clone()
2419 .sort_by([col(ROW_ORDER)], SortMultipleOptions::default())
2420 .drop_nulls()
2421 };
2422 match aggregate {
2423 Aggregate::Distinct => values.drop_nulls().n_unique().cast(DataType::Float64),
2425 Aggregate::Sum => values.sum(),
2426 Aggregate::Mean => values.mean(),
2427 Aggregate::Median => values.median(),
2428 Aggregate::Stdev => values.std(1),
2430 Aggregate::Quantile => {
2431 values.quantile(lit(f64::from(quantile) / 100.0), QuantileMethod::Linear)
2432 }
2433 Aggregate::Min => values.min(),
2434 Aggregate::Max => values.max(),
2435 Aggregate::First => in_order().first(),
2436 Aggregate::Last => in_order().last(),
2437 Aggregate::None | Aggregate::Count => len().cast(DataType::Float64),
2438 }
2439}
2440
2441fn with_row_order(lf: &LazyFrame, aggregate: crate::chart_modal::Aggregate) -> LazyFrame {
2443 if aggregate.follows_row_order() {
2444 lf.clone().with_row_index(ROW_ORDER, None)
2445 } else {
2446 lf.clone()
2447 }
2448}
2449
2450fn aggregate_pass(lf: LazyFrame, sampling: &ChartSampling) -> Result<DataFrame> {
2455 crate::statistics::collect_lazy(lf, true).map_err(|e| {
2456 if sampling.cancel.load(Ordering::Relaxed) {
2457 color_eyre::eyre::eyre!(ENVELOPE_CANCELLED)
2458 } else {
2459 e.into()
2460 }
2461 })
2462}
2463
2464const GROUPS_SAMPLE: usize = 20_000;
2466
2467fn refuse_too_many_groups(
2471 lf: &LazyFrame,
2472 x: &str,
2473 most: usize,
2474 sampling: &ChartSampling,
2475) -> Result<()> {
2476 let read = crate::statistics::analysis_rows(
2477 &lf.clone().select([col(x)]),
2478 Some(GROUPS_SAMPLE),
2479 sampling.known_total,
2480 sampling.seed,
2481 sampling.streaming,
2482 )?;
2483 let distinct = read.df.column(x)?.n_unique()?;
2484 let read_rows = read.df.height().max(1);
2485 let estimate = match read.sample_size {
2486 Some(_) => distinct as f64 / read_rows as f64 * read.total_rows as f64,
2487 None => distinct as f64,
2488 };
2489 if estimate > most as f64 {
2490 return Err(color_eyre::eyre::eyre!(
2491 "about {} values of {x}: more than a chart can draw. Bucket X by a time \
2492 unit, or choose a column with fewer values",
2493 crate::numfmt::group_chrome(estimate as usize)
2494 ));
2495 }
2496 Ok(())
2497}
2498
2499pub fn prepare_aggregate_xy(
2511 lf: &LazyFrame,
2512 schema: &Schema,
2513 spec: &AggregateSpec<'_>,
2514 sampling: &ChartSampling,
2515) -> Result<GroupedSeries> {
2516 use crate::chart_modal::{Aggregate, Cumulative};
2517 let x_dtype = schema
2518 .get(spec.x)
2519 .ok_or_else(|| color_eyre::eyre::eyre!("x column '{}' not in schema", spec.x))?;
2520 let mut x = col(spec.x);
2521 let bucketed = spec.time_unit.every().is_some()
2522 && matches!(x_dtype, DataType::Date | DataType::Datetime(_, _));
2523 if let Some(every) = spec.time_unit.every().filter(|_| bucketed) {
2524 x = x.dt().truncate(lit(every));
2525 }
2526 if !bucketed {
2527 refuse_too_many_groups(lf, spec.x, AGGREGATE_POINTS_MAX, sampling)?;
2528 }
2529 let x = until_cancelled(x, &sampling.cancel).alias("__x");
2530 let count = spec.aggregate == Aggregate::Count;
2531 let ys: &[String] = match (count, spec.color) {
2532 (true, _) => &[],
2533 (false, Some(_)) => &spec.ys[..spec.ys.len().min(1)],
2534 (false, None) => spec.ys,
2535 };
2536 let mut select = vec![x];
2537 let mut keys = vec![col("__x")];
2538 for (i, y) in ys.iter().enumerate() {
2539 select.push(y_values(col(y.as_str()), spec.aggregate).alias(format!("__y{i}")));
2540 }
2541 let plan = with_row_order(lf, spec.aggregate);
2542 if spec.aggregate.follows_row_order() {
2543 select.push(col(ROW_ORDER));
2544 }
2545 if let Some(color) = spec.color {
2546 select.push(group_expr(color).alias("__g"));
2547 keys.push(col("__g"));
2548 }
2549 let mut plan = plan.select(select).filter(col("__x").is_not_null());
2550 if spec.color.is_some() {
2551 plan = plan.filter(col("__g").is_not_null());
2552 }
2553 let mut aggs = vec![len().alias("__n")];
2554 for i in 0..ys.len() {
2555 let y = col(format!("__y{i}"));
2556 let made = match spec.cumulative {
2557 Cumulative::Off => aggregate_expr(y.clone(), spec.aggregate, spec.quantile),
2558 Cumulative::Sum => y.clone().sum(),
2559 Cumulative::Compound => (lit(1.0) + y.clone()).log(lit(std::f64::consts::E)).sum(),
2560 };
2561 aggs.push(made.alias(format!("__a{i}")));
2562 aggs.push(y.count().alias(format!("__c{i}")));
2564 }
2565 let df = aggregate_pass(
2566 plan.group_by_stable(keys)
2567 .agg(aggs)
2568 .sort(["__x"], Default::default()),
2569 sampling,
2570 )?;
2571 if df.height() > AGGREGATE_POINTS_MAX {
2572 return Err(color_eyre::eyre::eyre!(
2573 "{} points: more than a chart can draw. Bucket X by a time unit, or \
2574 choose an X with fewer values",
2575 crate::numfmt::group_chrome(df.height())
2576 ));
2577 }
2578 let xs: Vec<Option<f64>> = x_values(&df, "__x", x_dtype)?;
2579 let counts: Vec<u64> = df
2580 .column("__n")?
2581 .cast(&DataType::UInt64)?
2582 .u64()?
2583 .iter()
2584 .map(|n| n.unwrap_or(0))
2585 .collect();
2586 let groups: Option<Vec<Option<u32>>> = match spec.color {
2587 Some(_) => Some(df.column("__g")?.u32()?.iter().collect()),
2588 None => None,
2589 };
2590 let values: Vec<Vec<Option<f64>>> = if count {
2591 vec![counts.iter().map(|&n| Some(n as f64)).collect()]
2592 } else {
2593 (0..ys.len())
2594 .map(|i| {
2595 let made = df.column(&format!("__a{i}"))?.f64()?.clone();
2596 let behind = df.column(&format!("__c{i}"))?.cast(&DataType::UInt64)?;
2597 let behind = behind.u64()?;
2598 Ok(made
2599 .iter()
2600 .zip(behind.iter())
2601 .map(|(v, n)| {
2602 let v = v.filter(|_| n.unwrap_or(0) > 0)?;
2603 Some(if spec.cumulative == Cumulative::Compound {
2605 v.exp_m1()
2606 } else {
2607 v
2608 })
2609 })
2610 .collect())
2611 })
2612 .collect::<Result<_>>()?
2613 };
2614 let names: Vec<String> = match spec.color {
2615 Some(color) => color.names(),
2616 None if count => vec!["count".to_string()],
2617 None => ys.to_vec(),
2618 };
2619 let n = names.len();
2620 let mut series = vec![Vec::new(); n];
2621 let mut breaks = vec![Vec::new(); n];
2622 let mut gap = vec![false; n];
2623 let mut push = |s: usize, x: f64, y: Option<f64>| match y.filter(|y| y.is_finite()) {
2624 Some(y) => {
2625 if gap[s] && !series[s].is_empty() {
2626 breaks[s].push(series[s].len());
2627 }
2628 gap[s] = false;
2629 series[s].push((x, y));
2630 }
2631 None => gap[s] = true,
2632 };
2633 for (row, x) in xs.iter().enumerate() {
2634 let Some(x) = *x else { continue };
2635 match &groups {
2636 Some(groups) => {
2637 if let Some(g) = groups[row] {
2638 push(g as usize, x, values[0][row]);
2639 }
2640 }
2641 None => {
2642 for (s, column) in values.iter().enumerate() {
2643 push(s, x, column[row]);
2644 }
2645 }
2646 }
2647 }
2648 let how = match spec.cumulative {
2650 Cumulative::Compound if count => Cumulative::Sum,
2651 how => how,
2652 };
2653 for points in &mut series {
2654 accumulate(points, how);
2655 }
2656 Ok(GroupedSeries {
2657 names,
2658 series,
2659 breaks,
2660 x_axis_kind: x_axis_temporal_kind(x_dtype),
2661 rows: RowsRead {
2662 total_rows: counts.iter().sum::<u64>() as usize,
2663 sample_size: None,
2664 envelope_steps: None,
2665 seed: None,
2666 },
2667 other: spec.color.is_some_and(|c| c.other),
2668 })
2669}
2670
2671pub fn accumulate(points: &mut [(f64, f64)], how: crate::chart_modal::Cumulative) {
2674 use crate::chart_modal::Cumulative;
2675 let mut total = 0.0;
2676 for (_, y) in points.iter_mut() {
2677 total = match how {
2678 Cumulative::Off => return,
2679 Cumulative::Sum => total + *y,
2680 Cumulative::Compound => (1.0 + total) * (1.0 + *y) - 1.0,
2681 };
2682 *y = total;
2683 }
2684}
2685
2686#[derive(Clone, Copy, Debug)]
2688pub struct BarAggregate<'a> {
2689 pub category: &'a str,
2690 pub value: Option<&'a str>,
2692 pub aggregate: crate::chart_modal::Aggregate,
2693 pub quantile: u8,
2695 pub color: Option<ColorSplit<'a>>,
2696 pub order: BarOrder,
2697 pub cap: usize,
2698}
2699
2700pub fn prepare_bar_aggregate(
2704 lf: &LazyFrame,
2705 spec: &BarAggregate<'_>,
2706 sampling: &ChartSampling,
2707) -> Result<BarData> {
2708 use crate::chart_modal::Aggregate;
2709 let BarAggregate {
2710 category,
2711 value,
2712 aggregate,
2713 quantile,
2714 color,
2715 order,
2716 cap,
2717 } = *spec;
2718 let count = aggregate == Aggregate::Count;
2719 if count && color.is_none() {
2720 return prepare_bar_counts(lf, category, order, cap, sampling);
2721 }
2722 let value = match value {
2723 Some(value) if !count => Some(value),
2724 None if !count => return Err(color_eyre::eyre::eyre!("Pick a Y column")),
2725 _ => None,
2726 };
2727 let schema = lf.clone().collect_schema()?;
2728 let value_dtype = match value {
2729 Some(v) => schema
2730 .get(v)
2731 .cloned()
2732 .ok_or_else(|| color_eyre::eyre::eyre!("no column {v}"))?,
2733 None => DataType::UInt64,
2734 };
2735 let mut select = vec![until_cancelled(col(category), &sampling.cancel)];
2736 let mut keys = vec![col(category)];
2737 if let Some(value) = value {
2738 select.push(y_values(col(value), aggregate).alias("__v"));
2739 }
2740 if aggregate.follows_row_order() {
2741 select.push(col(ROW_ORDER));
2742 }
2743 if let Some(color) = color {
2744 select.push(group_expr(color).alias("__g"));
2745 keys.push(col("__g"));
2746 }
2747 let mut plan = with_row_order(lf, aggregate).select(select);
2748 if color.is_some() {
2749 plan = plan.filter(col("__g").is_not_null());
2750 }
2751 let measure = match value {
2752 Some(_) => aggregate_expr(col("__v"), aggregate, quantile),
2753 None => len().cast(DataType::Float64),
2754 };
2755 let behind = match value {
2757 Some(_) => col("__v").count(),
2758 None => len(),
2759 };
2760 refuse_too_many_groups(lf, category, COUNT_CATEGORY_CAP, sampling)?;
2761 let df = aggregate_pass(
2762 plan.group_by_stable(keys).agg([
2764 len().alias("__n"),
2765 measure.alias("__a"),
2766 behind.alias("__c"),
2767 ]),
2768 sampling,
2769 )?;
2770 let rows: usize = df
2771 .column("__n")?
2772 .cast(&DataType::UInt64)?
2773 .u64()?
2774 .iter()
2775 .map(|n| n.unwrap_or(0) as usize)
2776 .sum();
2777 let whole = aggregate.is_count()
2779 || (value_dtype.is_integer()
2780 && matches!(
2781 aggregate,
2782 Aggregate::Sum
2783 | Aggregate::Min
2784 | Aggregate::Max
2785 | Aggregate::First
2786 | Aggregate::Last
2787 ));
2788 let categories = df.column(category)?.as_materialized_series().clone();
2789 let labels_series = crate::past_calendar::cast_text(&categories, CastOptions::NonStrict)?;
2790 let labels: Vec<Option<&str>> = labels_series.str()?.iter().collect();
2791 let behind: Vec<u64> = df
2792 .column("__c")?
2793 .cast(&DataType::UInt64)?
2794 .u64()?
2795 .iter()
2796 .map(|n| n.unwrap_or(0))
2797 .collect();
2798 let measures: Vec<Option<f64>> = df
2800 .column("__a")?
2801 .f64()?
2802 .iter()
2803 .zip(&behind)
2804 .map(|(v, &n)| v.filter(|v| v.is_finite() && n > 0))
2805 .collect();
2806 let value_column = match value {
2807 Some(value) => format!("{} {value}", aggregate.named(quantile)),
2808 None => "count".to_string(),
2809 };
2810 let mut data = BarData {
2811 category: category.to_string(),
2812 value_column,
2813 bars: Vec::new(),
2814 more: 0,
2815 no_value: 0,
2816 rows: RowsRead {
2817 total_rows: rows,
2818 sample_size: None,
2819 envelope_steps: None,
2820 seed: None,
2821 },
2822 value_dtype: if whole {
2823 DataType::Int64
2824 } else {
2825 DataType::Float64
2826 },
2827 counted: None,
2828 groups: Vec::new(),
2829 other: false,
2830 rows_note: None,
2831 };
2832 let too_many = || {
2833 color_eyre::eyre::eyre!(
2834 "more than {} categories of {category}: choose a column with fewer",
2835 crate::numfmt::group_chrome(COUNT_CATEGORY_CAP)
2836 )
2837 };
2838 let Some(color) = color else {
2839 if df.height() > COUNT_CATEGORY_CAP {
2840 return Err(too_many());
2841 }
2842 let (bars, more, no_value) = order_bars(&categories, &labels, &measures, order, cap);
2843 (data.bars, data.more, data.no_value) = (bars, more, no_value);
2844 return Ok(data);
2845 };
2846 let groups: Vec<Option<u32>> = df.column("__g")?.u32()?.iter().collect();
2849 let mut at: std::collections::HashMap<Option<&str>, usize> = Default::default();
2850 let mut firsts: Vec<IdxSize> = Vec::new();
2851 let mut rows_of: Vec<Vec<Option<f64>>> = Vec::new();
2852 for (row, label) in labels.iter().enumerate() {
2853 let i = *at.entry(*label).or_insert_with(|| {
2854 firsts.push(row as IdxSize);
2855 rows_of.push(vec![None; color.series()]);
2856 rows_of.len() - 1
2857 });
2858 if let Some(g) = groups[row] {
2859 rows_of[i][g as usize] = measures[row];
2860 }
2861 }
2862 if rows_of.len() > COUNT_CATEGORY_CAP {
2863 return Err(too_many());
2864 }
2865 let unique = categories.take(&IdxCa::from_vec("firsts".into(), firsts.clone()))?;
2866 let unique_labels: Vec<Option<&str>> = firsts.iter().map(|&r| labels[r as usize]).collect();
2867 let totals: Vec<Option<f64>> = rows_of
2870 .iter()
2871 .map(|values| {
2872 let present = values.iter().flatten();
2873 if count || aggregate == Aggregate::Sum {
2874 Some(present.sum())
2875 } else {
2876 present.copied().reduce(f64::max)
2877 }
2878 })
2879 .collect();
2880 let (mut bars, more, no_value) = order_bars(&unique, &unique_labels, &totals, order, cap);
2881 let by_label: std::collections::HashMap<Option<&str>, usize> = unique_labels
2883 .iter()
2884 .enumerate()
2885 .map(|(i, l)| (*l, i))
2886 .collect();
2887 for bar in &mut bars {
2888 if let Some(&i) = by_label.get(&bar.label.as_deref()) {
2889 bar.by_group = rows_of[i].clone();
2890 }
2891 }
2892 data.bars = bars;
2893 data.more = more;
2894 data.no_value = no_value;
2895 data.groups = color.names();
2896 data.other = color.other;
2897 Ok(data)
2898}
2899
2900#[cfg(test)]
2901mod tests {
2902 use super::*;
2903
2904 fn tick_labels(ticks: &[f64], numbers: &AxisNumbers, level: usize) -> Vec<String> {
2906 let format = AxisFormat::new(ticks, numbers);
2907 ticks
2908 .iter()
2909 .map(|&v| format.label(v, level).unwrap())
2910 .collect()
2911 }
2912
2913 fn preset(name: &str) -> AxisNumbers {
2914 AxisNumbers {
2915 format: crate::numfmt::NumberFormat::preset(name).unwrap(),
2916 whole: false,
2917 }
2918 }
2919
2920 #[test]
2923 fn log_axis_labels_write_each_tick_exactly() {
2924 let ticks = [0.0, 1.0, 10.0, 100.0, 1e3, 1e4, 2e5, 1e6, 5e9];
2925 let labels = |numbers: &AxisNumbers, level| {
2926 let format = AxisFormat::log(&ticks, numbers);
2927 ticks
2928 .iter()
2929 .map(|&v| format.label(v, level).unwrap())
2930 .collect::<Vec<_>>()
2931 };
2932 assert_eq!(
2933 labels(&preset("thousands"), 0),
2934 [
2935 "0",
2936 "1",
2937 "10",
2938 "100",
2939 "1,000",
2940 "10,000",
2941 "200,000",
2942 "1,000,000",
2943 "5,000,000,000"
2944 ]
2945 );
2946 assert_eq!(
2947 labels(&AxisNumbers::default(), 1),
2948 ["0", "1", "10", "100", "1k", "10k", "200k", "1M", "5G"]
2949 );
2950 let short = [0.0, 0.25, 0.5, 0.75, 1.0];
2951 let format = AxisFormat::log(&short, &AxisNumbers::default());
2952 let written: Vec<_> = short.iter().map(|&v| format.label(v, 0).unwrap()).collect();
2953 assert_eq!(written, ["0.00", "0.25", "0.50", "0.75", "1.00"]);
2954 assert_eq!(
2955 format.label(1.0, 1),
2956 None,
2957 "no shorter form under a thousand"
2958 );
2959 let huge = AxisFormat::log(&[1e15, 2e16, 1e18], &AxisNumbers::default());
2960 assert_eq!(huge.label(2e16, 0).as_deref(), Some("2e16"));
2961 }
2962
2963 #[test]
2965 fn axis_labels_step_down_to_shorter_forms() {
2966 let plain = AxisNumbers::default();
2967 let ticks = [0.0, 12_345.0, 24_690.0];
2968 assert_eq!(tick_labels(&ticks, &plain, 0), ["0", "12345", "24690"]);
2969 assert_eq!(tick_labels(&ticks, &plain, 1), ["0", "12k", "25k"]);
2970 let ticks = [-1500.0, 0.0, 1500.0];
2971 assert_eq!(tick_labels(&ticks, &plain, 1), ["-1.5k", "0", "1.5k"]);
2972 let ticks = [0.0, 1.5e9, 3e9];
2973 assert_eq!(tick_labels(&ticks, &plain, 1), ["0", "1.5G", "3.0G"]);
2974 let format = AxisFormat::new(&[0.0, 5.0], &plain);
2976 assert_eq!(format.label(5.0, 1), None);
2977 assert_eq!(format.label(5.0, 2), None);
2978
2979 let (lo, hi) = (18262.0, 20088.0);
2981 let numbers = AxisFormat::new(&[], &plain);
2982 let date =
2983 |v, bounds, level| x_axis_label_at(v, XAxisTemporalKind::Date, bounds, level, &numbers);
2984 let forms: Vec<_> = (0..).map_while(|level| date(hi, (lo, hi), level)).collect();
2985 assert_eq!(forms, ["2024-12-31", "2024-12", "2024"]);
2986 assert_eq!(date(lo, (lo, lo + 30.0), 1).as_deref(), Some("01-01"));
2988
2989 let us = 86_400.0 * 1e6;
2990 let kind = XAxisTemporalKind::DatetimeUs;
2991 let at = |v, bounds, level| x_axis_label_at(v, kind, bounds, level, &numbers);
2992 let forms: Vec<_> = (0..)
2993 .map_while(|level| at(lo * us, (lo * us, hi * us), level))
2994 .collect();
2995 assert_eq!(forms, ["2020-01-01 00:00", "2020-01-01", "2020-01", "2020"]);
2996 let day = (lo * us, lo * us + 3600e6);
2998 assert_eq!(at(lo * us + 3600e6, day, 1).as_deref(), Some("01:00"));
2999 }
3000
3001 #[test]
3004 fn an_axis_keeps_one_notation_and_precision() {
3005 let plain = AxisNumbers::default();
3006 let density = tick_labels(&[0.0, 0.00651, 0.01302], &plain, 0);
3007 assert_eq!(density, ["0.0000", "0.0065", "0.0130"]);
3008 let density = tick_labels(&[0.0, 0.00451, 0.00902], &plain, 0);
3009 assert_eq!(density, ["0.00000", "0.00451", "0.00902"]);
3010 let round = tick_labels(&[0.0, 0.006, 0.012], &plain, 0);
3013 assert_eq!(round, ["0.000", "0.006", "0.012"]);
3014 let round = tick_labels(&[0.0, 0.1 * 3.0, 0.6], &plain, 0);
3015 assert_eq!(round, ["0.0", "0.3", "0.6"]);
3016 let round = tick_labels(&[0.0, 20.0, 40.0, 60.0, 80.0], &plain, 0);
3017 assert_eq!(round, ["0", "20", "40", "60", "80"]);
3018 let round = tick_labels(&[0.0, 2_000.0, 4_000.0], &plain, 1);
3019 assert_eq!(round, ["0", "2k", "4k"]);
3020 let tiny = tick_labels(&[0.0, 2.5e-8, 5e-8], &plain, 0);
3022 assert_eq!(tiny, ["0.00e0", "2.50e-8", "5.00e-8"]);
3023 assert_eq!(
3025 tick_labels(&[3.21, 50.17, 97.2], &plain, 0),
3026 ["3.2", "50.2", "97.2"]
3027 );
3028 let close = tick_labels(&[1000.1, 1000.2, 1000.3], &plain, 0);
3029 assert_eq!(close, ["1000.1", "1000.2", "1000.3"]);
3030 assert_eq!(tick_labels(&[-0.0001, 1.0], &plain, 0), ["0.00", "1.00"]);
3032 let padded = [-0.5, 249.75, 500.0];
3033 assert_eq!(tick_labels(&padded, &plain, 0), ["0", "250", "500"]);
3034 let padded = [-500.0, 24_750.0, 50_000.0];
3035 assert_eq!(tick_labels(&padded, &plain, 1), ["0", "25k", "50k"]);
3036 assert_eq!(tick_labels(&[-0.0, 5e-8], &plain, 0), ["0.00e0", "5.00e-8"]);
3037 let stepped = [-2e-8, 1.3e-24, 2e-8];
3039 assert_eq!(
3040 tick_labels(&stepped, &plain, 0),
3041 ["-2.00e-8", "0.00e0", "2.00e-8"]
3042 );
3043 let close = tick_labels(&[1.0, 1.000_000_1], &plain, 0);
3046 assert_eq!(close, ["1.0000000e0", "1.0000001e0"]);
3047 let nanoseconds = [1.727e18, 1.727_05e18, 1.7271e18];
3048 let labels = tick_labels(&nanoseconds, &plain, 0);
3049 assert_eq!(labels, ["1.72700e18", "1.72705e18", "1.72710e18"]);
3050 assert_eq!(tick_labels(&nanoseconds, &plain, 1), labels);
3051 let huge = [0.0, 5e15];
3053 assert_eq!(tick_labels(&huge, &plain, 0), ["0.00e0", "5.00e15"]);
3054 assert_eq!(tick_labels(&huge, &plain, 1), ["0e0", "5e15"]);
3055 let huge = [1e15, 1.5e15, 2e15];
3056 assert_eq!(
3057 tick_labels(&huge, &plain, 1),
3058 ["1.0e15", "1.5e15", "2.0e15"]
3059 );
3060 let whole = AxisNumbers {
3062 whole: true,
3063 ..preset("thousands")
3064 };
3065 assert_eq!(
3066 tick_labels(&[0.0, 2161.0, 4322.0], &whole, 0),
3067 ["0", "2,161", "4,322"]
3068 );
3069 }
3070
3071 #[test]
3074 fn axis_labels_take_the_table_number_style() {
3075 let european = preset("european");
3076 let ticks = [12_000.0, 12_300.0, 12_600.0];
3077 assert_eq!(
3078 tick_labels(&ticks, &european, 0),
3079 ["12.000", "12.300", "12.600"]
3080 );
3081 assert_eq!(
3082 tick_labels(&ticks, &european, 1),
3083 ["12,0k", "12,3k", "12,6k"]
3084 );
3085 let ticks = [0.0, 0.25, 0.5];
3086 assert_eq!(tick_labels(&ticks, &european, 0), ["0,00", "0,25", "0,50"]);
3087 assert_eq!(
3088 tick_labels(&[0.0, 5e-8], &european, 0),
3089 ["0,00e0", "5,00e-8"]
3090 );
3091
3092 let thousands = preset("thousands");
3093 let ticks = [0.0, 6172.4, 12345.0];
3094 assert_eq!(tick_labels(&ticks, &thousands, 0), ["0", "6,172", "12,345"]);
3095 let ticks = [0.0, 12_300.0];
3096 assert_eq!(tick_labels(&ticks, &thousands, 1), ["0", "12k"]);
3097 }
3098
3099 fn all_rows() -> ChartSampling {
3100 ChartSampling::rows(Some(10_000))
3101 }
3102
3103 fn xy(lf: &LazyFrame, x: &str, ys: &[&str], sampling: &ChartSampling) -> ChartDataResult {
3104 let schema = lf.clone().collect_schema().unwrap();
3105 let ys: Vec<String> = ys.iter().map(|s| s.to_string()).collect();
3106 prepare_chart_data(lf, schema.as_ref(), x, &ys, sampling, false).unwrap()
3107 }
3108
3109 #[test]
3112 fn a_long_line_is_drawn_as_its_envelope() {
3113 let n = 100_000usize;
3114 let x: Vec<i64> = (0..n as i64).collect();
3115 let y: Vec<f64> = (0..n)
3116 .map(|i| match i {
3117 54_321 => 9.0,
3119 _ => ((i as f64) / 50.0).sin(),
3120 })
3121 .collect();
3122 let lf = df!("x" => &x, "y" => &y).unwrap().lazy();
3123 let schema = lf.clone().collect_schema().unwrap();
3124 let sampling = ChartSampling::rows(Some(1_000));
3125 let result =
3126 prepare_chart_data(&lf, schema.as_ref(), "x", &["y".into()], &sampling, true).unwrap();
3127 assert_eq!(
3128 result.rows,
3129 RowsRead {
3130 total_rows: n,
3131 sample_size: None,
3132 envelope_steps: Some(500),
3133 seed: None,
3134 }
3135 );
3136 let points = &result.series[0];
3137 assert!(points.len() <= 1_000, "{} points", points.len());
3138 let top = points.iter().map(|p| p.1).fold(f64::MIN, f64::max);
3139 assert_eq!(top, 9.0, "the spike is kept");
3140 let bottom = points.iter().map(|p| p.1).fold(f64::MAX, f64::min);
3141 assert!(bottom < -0.99, "so is every trough: {bottom}");
3142 assert!(points.windows(2).all(|w| w[0].0 <= w[1].0), "in X order");
3143 assert_eq!(
3144 chart_notes(&result.rows, None, "·"),
3145 ["min and max of 100k rows in 500 steps"]
3146 );
3147
3148 let sampling = ChartSampling::rows(Some(200_000));
3150 let result =
3151 prepare_chart_data(&lf, schema.as_ref(), "x", &["y".into()], &sampling, true).unwrap();
3152 assert_eq!(result.rows.envelope_steps, None);
3153 assert_eq!(result.series[0].len(), n);
3154 }
3155
3156 #[test]
3160 fn an_envelope_is_sampled_instead_where_full_reads_cost_and_stops_when_cancelled() {
3161 let n = 10_000i64;
3162 let lf = df!("x" => (0..n).collect::<Vec<_>>(), "y" => (0..n).collect::<Vec<_>>())
3163 .unwrap()
3164 .lazy();
3165 let schema = lf.clone().collect_schema().unwrap();
3166 let remote = ChartSampling {
3167 full_passes: false,
3168 ..ChartSampling::rows(Some(100))
3169 };
3170 let result =
3171 prepare_chart_data(&lf, schema.as_ref(), "x", &["y".into()], &remote, true).unwrap();
3172 assert_eq!(result.rows.envelope_steps, None);
3173 assert_eq!(result.rows.sample_size, Some(100));
3174
3175 let cancelled = ChartSampling::rows(Some(100));
3176 cancelled.cancel.store(true, Ordering::Relaxed);
3177 let err = prepare_chart_data(&lf, schema.as_ref(), "x", &["y".into()], &cancelled, true)
3178 .err()
3179 .expect("a cancelled envelope is not drawn");
3180 assert_eq!(err.to_string(), "chart cancelled");
3181 }
3182
3183 #[test]
3185 fn an_envelope_places_temporal_x_by_its_ordinal() {
3186 let days: Vec<i32> = (0..1_000).collect();
3187 let lf = df!("d" => &days, "y" => (0..1_000).map(f64::from).collect::<Vec<_>>())
3188 .unwrap()
3189 .lazy()
3190 .with_column(col("d").cast(DataType::Date));
3191 let schema = lf.clone().collect_schema().unwrap();
3192 let result = prepare_chart_data(
3193 &lf,
3194 schema.as_ref(),
3195 "d",
3196 &["y".into()],
3197 &ChartSampling::rows(Some(100)),
3198 true,
3199 )
3200 .unwrap();
3201 assert_eq!(result.rows.envelope_steps, Some(50));
3202 let points = &result.series[0];
3203 assert_eq!(points.first(), Some(&(0.0, 0.0)));
3204 assert_eq!(points.last().map(|p| p.1), Some(999.0));
3205 assert!(
3206 points.iter().all(|&(x, y)| y >= x && y < x + 20.0),
3207 "{points:?}"
3208 );
3209 }
3210
3211 #[test]
3213 fn an_envelope_breaks_where_a_series_has_no_values() {
3214 let x: Vec<i64> = (0..100).collect();
3215 let y: Vec<Option<f64>> = (0..100)
3216 .map(|i| (!(40..60).contains(&i)).then_some(i as f64))
3217 .collect();
3218 let lf = df!("x" => &x, "y" => &y).unwrap().lazy();
3219 let schema = lf.clone().collect_schema().unwrap();
3220 let sampling = ChartSampling::rows(Some(20));
3221 let result =
3222 prepare_chart_data(&lf, schema.as_ref(), "x", &["y".into()], &sampling, true).unwrap();
3223 assert_eq!(result.rows.envelope_steps, Some(10));
3224 assert_eq!(result.breaks[0].len(), 1, "one gap: {:?}", result.series[0]);
3225 }
3226
3227 #[test]
3229 fn an_envelope_leaves_out_rows_with_no_x() {
3230 let x: Vec<f64> = (0..100)
3231 .map(|i| if i % 10 == 0 { f64::NAN } else { i as f64 })
3232 .collect();
3233 let y: Vec<f64> = (0..100).map(|i| i as f64).collect();
3234 let lf = df!("x" => &x, "y" => &y).unwrap().lazy();
3235 let schema = lf.clone().collect_schema().unwrap();
3236 let sampling = ChartSampling::rows(Some(20));
3237 let result =
3238 prepare_chart_data(&lf, schema.as_ref(), "x", &["y".into()], &sampling, true).unwrap();
3239 let ys: Vec<f64> = result.series[0].iter().map(|p| p.1).collect();
3240 assert!(!ys.contains(&0.0) && !ys.contains(&50.0), "{ys:?}");
3241 assert_eq!(ys.iter().cloned().fold(f64::MIN, f64::max), 99.0);
3242 }
3243
3244 #[test]
3245 fn prepare_empty_y_columns() {
3246 let lf = df!("x" => &[1.0_f64, 2.0], "y" => &[10.0, 20.0])
3247 .unwrap()
3248 .lazy();
3249 let result = xy(&lf, "x", &[], &all_rows());
3250 assert!(result.series.is_empty());
3251 assert_eq!(result.x_axis_kind, XAxisTemporalKind::Numeric);
3252 }
3253
3254 #[test]
3255 fn prepare_small_data() {
3256 let lf = df!(
3257 "x" => &[1.0_f64, 2.0, 3.0],
3258 "a" => &[10.0_f64, 20.0, 30.0],
3259 "b" => &[100.0_f64, 200.0, 300.0]
3260 )
3261 .unwrap()
3262 .lazy();
3263 let result = xy(&lf, "x", &["a", "b"], &all_rows());
3264 assert_eq!(result.series.len(), 2);
3265 assert_eq!(
3266 result.series[0],
3267 vec![(1.0, 10.0), (2.0, 20.0), (3.0, 30.0)]
3268 );
3269 assert_eq!(
3270 result.series[1],
3271 vec![(1.0, 100.0), (2.0, 200.0), (3.0, 300.0)]
3272 );
3273 assert_eq!(result.x_axis_kind, XAxisTemporalKind::Numeric);
3274 assert_eq!(
3275 result.rows,
3276 RowsRead {
3277 total_rows: 3,
3278 sample_size: None,
3279 envelope_steps: None,
3280 seed: None,
3281 },
3282 "every row read: nothing to say"
3283 );
3284 assert!(chart_notes(&result.rows, None, "·").is_empty());
3285 }
3286
3287 #[test]
3288 fn prepare_skips_nan() {
3289 let lf = df!(
3290 "x" => &[1.0_f64, 2.0, 3.0],
3291 "y" => &[10.0_f64, f64::NAN, 30.0]
3292 )
3293 .unwrap()
3294 .lazy();
3295 let result = xy(&lf, "x", &["y"], &all_rows());
3296 assert_eq!(result.series[0], vec![(1.0, 10.0), (3.0, 30.0)]);
3297 }
3298
3299 #[test]
3300 fn prepare_missing_x_column_errors() {
3301 let lf = df!("x" => &[1.0_f64], "y" => &[2.0_f64]).unwrap().lazy();
3302 let schema = lf.clone().collect_schema().unwrap();
3303 let result = prepare_chart_data(
3304 &lf,
3305 schema.as_ref(),
3306 "missing",
3307 &["y".into()],
3308 &all_rows(),
3309 false,
3310 );
3311 assert!(result.is_err());
3312 }
3313
3314 #[test]
3317 fn over_the_limit_a_chart_reads_a_spread_sample_and_says_so() {
3318 let n = 50_000_i64;
3319 let lf = df!(
3320 "x" => (0..n).collect::<Vec<_>>(),
3321 "y" => (0..n).map(|v| v * 2).collect::<Vec<_>>()
3322 )
3323 .unwrap()
3324 .lazy();
3325 let result = xy(&lf, "x", &["y"], &ChartSampling::rows(Some(1_000)));
3326 let points = &result.series[0];
3327 assert_eq!(points.len(), 1_000);
3328 let last_x = points.last().unwrap().0;
3329 assert!(
3330 last_x > (n as f64) * 0.9,
3331 "the sample reaches the end of the table, got {last_x}"
3332 );
3333 assert_eq!(
3334 result.rows,
3335 RowsRead {
3336 total_rows: n as usize,
3337 sample_size: Some(1_000),
3338 envelope_steps: None,
3339 seed: Some(crate::sampling::Sample::default().seed),
3340 }
3341 );
3342 let seed = crate::sampling::Sample::default().seed;
3344 assert_eq!(
3345 chart_notes(&result.rows, None, "·"),
3346 [format!("sample of 1,000 of 50k rows · seed {seed}")]
3347 );
3348 assert_eq!(
3349 chart_notes(&result.rows, None, "-"),
3350 [format!("sample of 1,000 of 50k rows - seed {seed}")]
3351 );
3352
3353 let every = xy(&lf, "x", &["y"], &ChartSampling::rows(None));
3355 assert_eq!(every.series[0].len(), n as usize);
3356 assert_eq!(every.rows.sample_size, None);
3357 }
3358
3359 #[test]
3362 fn a_parquet_file_is_sampled_in_runs() {
3363 let dir = tempfile::tempdir().unwrap();
3364 let n = 100_000_i64;
3365 let mut df = df!(
3366 "id" => (0..n).collect::<Vec<_>>(),
3367 "fare" => (0..n).map(|v| v as f64).collect::<Vec<_>>(),
3368 "other" => vec!["x"; n as usize]
3369 )
3370 .unwrap();
3371 let path = dir.path().join("trips.parquet");
3372 ParquetWriter::new(std::fs::File::create(&path).unwrap())
3373 .with_row_group_size(Some(1_000))
3374 .finish(&mut df)
3375 .unwrap();
3376 let lf =
3377 LazyFrame::scan_parquet(PlRefPath::try_from_path(&path).unwrap(), Default::default())
3378 .unwrap();
3379 assert!(crate::statistics::slices_reach_into_the_scan(
3380 &lf.clone().select([col("id"), col("fare")])
3381 ));
3382 let data = prepare_histogram_data(
3383 &lf,
3384 "fare",
3385 10,
3386 ValueRange::All,
3387 &ChartSampling::rows(Some(2_000)),
3388 )
3389 .unwrap();
3390 assert_eq!(
3391 data.rows,
3392 RowsRead {
3393 total_rows: n as usize,
3394 sample_size: Some(2_000),
3395 envelope_steps: None,
3396 seed: Some(crate::sampling::Sample::default().seed),
3397 }
3398 );
3399 assert!(data.x_max > 90_000.0, "reaches the end: {}", data.x_max);
3400 }
3401
3402 #[test]
3405 fn the_sample_is_seeded() {
3406 let lf = df!("x" => (0..20_000_i64).collect::<Vec<_>>(), "y" => vec![1.0_f64; 20_000])
3407 .unwrap()
3408 .lazy();
3409 let a = xy(&lf, "x", &["y"], &ChartSampling::rows(Some(500)));
3410 let b = xy(&lf, "x", &["y"], &ChartSampling::rows(Some(500)));
3411 assert_eq!(a.series, b.series);
3412 let other = ChartSampling {
3413 seed: 7,
3414 ..ChartSampling::rows(Some(500))
3415 };
3416 let c = xy(&lf, "x", &["y"], &other);
3417 assert_ne!(a.series, c.series);
3418 }
3419
3420 #[test]
3424 fn rows_already_read_are_not_read_again() {
3425 let dir = tempfile::tempdir().unwrap();
3426 let path = dir.path().join("fares.csv");
3427 let write = |a: i64, b: i64| {
3428 let rows: String = (0..100).map(|i| format!("{},{}\n", i + a, i + b)).collect();
3429 std::fs::write(&path, format!("a,b\n{rows}")).unwrap();
3430 };
3431 write(0, 0);
3432 let lf = LazyCsvReader::new(PlRefPath::try_from_path(&path).unwrap())
3433 .finish()
3434 .unwrap();
3435 let sampling = ChartSampling::rows(Some(10_000));
3436 let first = prepare_histogram_data(&lf, "a", 10, ValueRange::All, &sampling).unwrap();
3437 assert_eq!(first.x_min, 0.0);
3438
3439 write(1_000, 1_000);
3440 let held =
3441 prepare_histogram_data(&lf, "a", 5, ValueRange::Percentile1To99, &sampling).unwrap();
3442 assert!(
3443 held.x_max < 100.0,
3444 "drawn from the rows held: {}",
3445 held.x_max
3446 );
3447 let boxed = prepare_box_plot_data(&lf, &["a"], ValueRange::All, &sampling).unwrap();
3448 assert_eq!(boxed.stats[0].max, 99.0);
3449
3450 let with_b = prepare_histogram_data(&lf, "b", 10, ValueRange::All, &sampling).unwrap();
3451 assert_eq!(with_b.x_min, 1_000.0, "b was not held: read");
3452 let a_again = prepare_histogram_data(&lf, "a", 10, ValueRange::All, &sampling).unwrap();
3453 assert_eq!(a_again.x_min, 1_000.0, "read along with b");
3454
3455 write(5_000, 5_000);
3456 let other_size = ChartSampling {
3457 limit: Some(50),
3458 ..sampling.clone()
3459 };
3460 let resampled = prepare_histogram_data(&lf, "a", 10, ValueRange::All, &other_size).unwrap();
3461 assert!(resampled.x_min >= 5_000.0, "another size reads afresh");
3462 }
3463
3464 #[test]
3467 fn points_come_in_x_order() {
3468 let lf = df!(
3469 "year" => &[2001_i64, 1999, 2003, 2000, 2002],
3470 "count" => &[1.0_f64, 2.0, 3.0, 4.0, 5.0]
3471 )
3472 .unwrap()
3473 .lazy();
3474 let result = xy(&lf, "year", &["count"], &all_rows());
3475 let xs: Vec<f64> = result.series[0].iter().map(|p| p.0).collect();
3476 assert_eq!(xs, [1999.0, 2000.0, 2001.0, 2002.0, 2003.0]);
3477 assert_eq!(result.series[0][0], (1999.0, 2.0));
3478 assert!(result.breaks[0].is_empty());
3479 }
3480
3481 #[test]
3484 fn x_as_a_y_series_charts_rather_than_failing() {
3485 let lf = df!("x" => &[1.0_f64, 2.0], "y" => &[3.0_f64, 4.0])
3486 .unwrap()
3487 .lazy();
3488 let result = xy(&lf, "x", &["x", "y"], &all_rows());
3489 assert_eq!(result.series[0], vec![(1.0, 1.0), (2.0, 2.0)]);
3490 assert_eq!(result.series[1], vec![(1.0, 3.0), (2.0, 4.0)]);
3491 }
3492
3493 #[test]
3495 fn a_missing_y_column_is_an_error() {
3496 let lf = df!("x" => &[1.0_f64], "y" => &[2.0_f64]).unwrap().lazy();
3497 let schema = lf.clone().collect_schema().unwrap();
3498 let result = prepare_chart_data(
3499 &lf,
3500 schema.as_ref(),
3501 "x",
3502 &["gone".into()],
3503 &all_rows(),
3504 false,
3505 );
3506 assert!(result.is_err());
3507 }
3508
3509 #[test]
3512 fn nulls_drop_per_series_and_break_the_line() {
3513 let lf = df!(
3514 "year" => &[Some(1880_i64), Some(1881), Some(1882), None, Some(1883), Some(1884)],
3515 "emma" => &[Some(10.0_f64), Some(11.0), Some(12.0), Some(99.0), Some(13.0), Some(14.0)],
3516 "jennifer" => &[None, None, Some(5.0_f64), Some(99.0), None, Some(7.0)]
3517 )
3518 .unwrap()
3519 .lazy();
3520 let result = xy(&lf, "year", &["emma", "jennifer"], &all_rows());
3521 assert_eq!(
3522 result.series[0],
3523 vec![
3524 (1880.0, 10.0),
3525 (1881.0, 11.0),
3526 (1882.0, 12.0),
3527 (1883.0, 13.0),
3528 (1884.0, 14.0)
3529 ],
3530 "Emma keeps the years Jennifer is missing; the null year is gone"
3531 );
3532 assert!(result.breaks[0].is_empty());
3533 assert_eq!(result.series[1], vec![(1882.0, 5.0), (1884.0, 7.0)]);
3534 assert_eq!(result.breaks[1], [1], "1883 is missing: the line breaks");
3535 assert_eq!(
3536 segments(&result.series[1], &result.breaks[1]),
3537 vec![&[(1882.0, 5.0)][..], &[(1884.0, 7.0)][..]]
3538 );
3539 }
3540
3541 #[test]
3542 fn segments_split_at_breaks() {
3543 let points = [(0.0, 0.0), (1.0, 1.0), (2.0, 2.0), (3.0, 3.0)];
3544 assert_eq!(segments(&points, &[]), vec![&points[..]]);
3545 assert_eq!(
3546 segments(&points, &[1, 3]),
3547 vec![&points[..1], &points[1..3], &points[3..]]
3548 );
3549 assert!(segments(&[], &[]).is_empty());
3550 }
3551
3552 #[test]
3554 fn a_date_x_is_ordinal() {
3555 let lf = df!("d" => &[Some(1_i32), None, Some(0)], "y" => &[1.0_f64, 2.0, 3.0])
3556 .unwrap()
3557 .lazy()
3558 .with_column(col("d").cast(DataType::Date));
3559 let result = xy(&lf, "d", &["y"], &all_rows());
3560 assert_eq!(result.x_axis_kind, XAxisTemporalKind::Date);
3561 assert_eq!(result.series[0], vec![(0.0, 3.0), (1.0, 1.0)]);
3562 }
3563
3564 fn with_outliers() -> LazyFrame {
3565 let mut v: Vec<f64> = (1..=100).map(f64::from).collect();
3567 v.push(-10_000.0);
3568 v.push(50_000.0);
3569 df!("fare" => v).unwrap().lazy()
3570 }
3571
3572 #[test]
3574 fn a_histogram_range_clips_the_tails_and_counts_them() {
3575 let lf = with_outliers();
3576 let all = prepare_histogram_data(&lf, "fare", 10, ValueRange::All, &all_rows()).unwrap();
3577 assert_eq!(all.x_min, -10_000.0);
3578 assert!(all.clipped.is_none());
3579
3580 let clipped =
3581 prepare_histogram_data(&lf, "fare", 10, ValueRange::Percentile1To99, &all_rows())
3582 .unwrap();
3583 assert_eq!((clipped.x_min, clipped.x_max), (2.0, 99.0));
3586 let outside = clipped.clipped.unwrap().outside;
3587 assert_eq!(outside, 4);
3588 let counted: f64 = clipped.bins.iter().map(|b| b.count).sum();
3589 assert_eq!(counted as usize + outside, 102);
3590 assert_eq!(
3591 chart_notes(&clipped.rows, clipped.clipped.as_ref(), "·"),
3592 ["4 values outside p1-p99"]
3593 );
3594 }
3595
3596 #[test]
3597 fn box_plot_and_kde_take_the_range_too() {
3598 let lf = with_outliers();
3599 let boxed = prepare_box_plot_data(&lf, &["fare"], ValueRange::Percentile1To99, &all_rows())
3600 .unwrap();
3601 assert!(boxed.stats[0].min > 0.0 && boxed.stats[0].max <= 100.0);
3602 assert!(boxed.clipped.unwrap().outside >= 2);
3603
3604 let kde = prepare_kde_data(
3605 &lf,
3606 &["fare"],
3607 1.0,
3608 ValueRange::Percentile1To99,
3609 &all_rows(),
3610 )
3611 .unwrap();
3612 assert!(kde.x_min > -1_000.0 && kde.x_max < 1_000.0);
3613 assert!(kde.clipped.unwrap().outside >= 2);
3614
3615 let whole = prepare_box_plot_data(&lf, &["fare"], ValueRange::All, &all_rows()).unwrap();
3616 assert_eq!(whole.stats[0].min, -10_000.0);
3617 }
3618
3619 #[test]
3621 fn a_histogram_keeps_every_value_of_its_column() {
3622 let lf = df!("a" => &[Some(1.0_f64), Some(2.0), None, Some(4.0)])
3623 .unwrap()
3624 .lazy();
3625 let data = prepare_histogram_data(&lf, "a", 5, ValueRange::All, &all_rows()).unwrap();
3626 let counted: f64 = data.bins.iter().map(|b| b.count).sum();
3627 assert_eq!(counted, 3.0);
3628 }
3629
3630 #[test]
3631 fn prepare_x_range_numeric() {
3632 let lf = df!("x" => &[10.0_f64, 20.0, 5.0, 30.0]).unwrap().lazy();
3633 let schema = lf.clone().collect_schema().unwrap();
3634 let r = prepare_chart_x_range(&lf, schema.as_ref(), "x", &all_rows()).unwrap();
3635 assert_eq!(r.x_min, 5.0);
3636 assert_eq!(r.x_max, 30.0);
3637 assert_eq!(r.x_axis_kind, XAxisTemporalKind::Numeric);
3638 }
3639
3640 #[test]
3641 fn prepare_x_range_empty_returns_placeholder() {
3642 let lf = df!("x" => &[1.0_f64]).unwrap().lazy().slice(0, 0);
3643 let schema = lf.clone().collect_schema().unwrap();
3644 let r = prepare_chart_x_range(&lf, schema.as_ref(), "x", &all_rows()).unwrap();
3645 assert_eq!(r.x_min, 0.0);
3646 assert_eq!(r.x_max, 1.0);
3647 }
3648
3649 fn bars(lf: &LazyFrame, order: BarOrder, cap: usize) -> BarData {
3650 prepare_bar_data(lf, "carrier", "delay", order, cap, &all_rows()).unwrap()
3651 }
3652
3653 fn labels(data: &BarData) -> Vec<Option<&str>> {
3654 data.bars.iter().map(|b| b.label.as_deref()).collect()
3655 }
3656
3657 #[test]
3660 fn bars_order_by_value_or_label_and_cap_the_rest() {
3661 let lf = df!(
3662 "carrier" => &["UA", "AA", "DL", "B6", "AS"],
3663 "delay" => &[3.5_f64, 0.4, 1.6, 9.5, -9.9]
3664 )
3665 .unwrap()
3666 .lazy();
3667 let by_value = bars(&lf, BarOrder::Value, BAR_CAP);
3668 assert_eq!(
3669 labels(&by_value),
3670 [Some("B6"), Some("UA"), Some("DL"), Some("AA"), Some("AS")]
3671 );
3672 assert_eq!(by_value.bars[4].value, -9.9);
3673 assert_eq!(by_value.more, 0);
3674
3675 let by_label = bars(&lf, BarOrder::Label, BAR_CAP);
3676 assert_eq!(
3677 labels(&by_label),
3678 [Some("AA"), Some("AS"), Some("B6"), Some("DL"), Some("UA")]
3679 );
3680
3681 let capped = bars(&lf, BarOrder::Value, 2);
3682 assert_eq!(labels(&capped), [Some("B6"), Some("UA")]);
3683 assert_eq!(capped.more, 3, "the three smallest are counted, not drawn");
3684 let capped = bars(&lf, BarOrder::Label, 2);
3685 assert_eq!(labels(&capped), [Some("AA"), Some("AS")]);
3686 assert_eq!(capped.more, 3);
3687 }
3688
3689 #[test]
3692 fn bar_categories_keep_their_type_and_nulls_are_counted() {
3693 let lf = df!(
3694 "carrier" => &[Some(10_i64), Some(9), None, Some(100), Some(2)],
3695 "delay" => &[Some(1.0_f64), Some(2.0), Some(3.0), None, Some(1.0)]
3696 )
3697 .unwrap()
3698 .lazy();
3699 let data = bars(&lf, BarOrder::Label, BAR_CAP);
3700 assert_eq!(labels(&data), [Some("2"), Some("9"), Some("10"), None]);
3701 assert_eq!(data.no_value, 1, "100 has no value");
3702
3703 let data = bars(&lf, BarOrder::Value, BAR_CAP);
3704 assert_eq!(
3705 labels(&data),
3706 [None, Some("9"), Some("10"), Some("2")],
3707 "ties keep table order"
3708 );
3709 }
3710
3711 #[test]
3715 fn a_date_category_past_the_calendar_is_labeled_by_its_stored_number() {
3716 let paris = TimeZone::opt_try_new(Some("Europe/Paris")).unwrap();
3717 let datetime = |unit, zone: Option<TimeZone>| {
3718 Series::new("at".into(), [i64::MIN + 1, 0])
3719 .cast(&DataType::Datetime(unit, zone))
3720 .unwrap()
3721 };
3722 for (at, labels_in_order) in [
3723 (
3724 Series::new("at".into(), [i32::MAX, 0])
3725 .cast(&DataType::Date)
3726 .unwrap(),
3727 ["1970-01-01", "2147483647 days since 1970-01-01"],
3728 ),
3729 (
3730 datetime(TimeUnit::Milliseconds, None),
3731 [
3732 "-9223372036854775807 ms since 1970-01-01 UTC",
3733 "1970-01-01 00:00:00.000",
3734 ],
3735 ),
3736 (
3737 datetime(TimeUnit::Microseconds, paris),
3738 [
3739 "-9223372036854775807 us since 1970-01-01 UTC",
3740 "1970-01-01 01:00:00.000000+01:00",
3741 ],
3742 ),
3743 ] {
3744 let lf = DataFrame::new_infer_height(vec![
3745 at.into_column(),
3746 Column::new("n".into(), [1i64, 2]),
3747 ])
3748 .unwrap()
3749 .lazy();
3750 let by_value =
3751 prepare_bar_data(&lf, "at", "n", BarOrder::Label, BAR_CAP, &all_rows()).unwrap();
3752 let counted =
3753 prepare_bar_counts(&lf, "at", BarOrder::Label, BAR_CAP, &all_rows()).unwrap();
3754 for data in [by_value, counted] {
3755 assert_eq!(labels(&data), labels_in_order.map(Some));
3756 }
3757 }
3758 }
3759
3760 #[test]
3762 fn a_repeated_category_is_refused() {
3763 let lf = df!(
3764 "species" => &["Adelie", "Adelie", "Gentoo"],
3765 "body_mass_g" => &[3750_i64, 3800, 5000]
3766 )
3767 .unwrap()
3768 .lazy();
3769 let err = prepare_bar_data(
3770 &lf,
3771 "species",
3772 "body_mass_g",
3773 BarOrder::Value,
3774 BAR_CAP,
3775 &all_rows(),
3776 )
3777 .unwrap_err()
3778 .to_string();
3779 assert!(
3780 err.contains("species repeats: 2 categories in 3 rows"),
3781 "{err}"
3782 );
3783 assert!(
3784 err.contains("SELECT species, AVG(body_mass_g) FROM df GROUP BY species"),
3785 "SQL first: {err}"
3786 );
3787 assert!(
3788 err.contains("(or select avg body_mass_g by species)"),
3789 "{err}"
3790 );
3791
3792 let lf = df!("Species" => &["a", "a"], "mass g" => &[1_i64, 2])
3794 .unwrap()
3795 .lazy();
3796 let err = prepare_bar_data(
3797 &lf,
3798 "Species",
3799 "mass g",
3800 BarOrder::Value,
3801 BAR_CAP,
3802 &all_rows(),
3803 )
3804 .unwrap_err()
3805 .to_string();
3806 assert!(
3807 err.ends_with(
3808 r#"SELECT "Species", AVG("mass g") FROM df GROUP BY "Species", or choose Count for the rows per category"#
3809 ),
3810 "{err}"
3811 );
3812 }
3813
3814 #[test]
3816 fn a_new_bar_order_does_not_read_again() {
3817 let dir = tempfile::tempdir().unwrap();
3818 let path = dir.path().join("delays.csv");
3819 std::fs::write(&path, "carrier,delay\nUA,3.5\nAA,0.4\n").unwrap();
3820 let lf = LazyCsvReader::new(PlRefPath::try_from_path(&path).unwrap())
3821 .finish()
3822 .unwrap();
3823 let sampling = all_rows();
3824 let first =
3825 prepare_bar_data(&lf, "carrier", "delay", BarOrder::Value, BAR_CAP, &sampling).unwrap();
3826 assert_eq!(labels(&first), [Some("UA"), Some("AA")]);
3827 std::fs::write(&path, "carrier,delay\nZZ,1.0\n").unwrap();
3828 let again =
3829 prepare_bar_data(&lf, "carrier", "delay", BarOrder::Label, BAR_CAP, &sampling).unwrap();
3830 assert_eq!(
3831 labels(&again),
3832 [Some("AA"), Some("UA")],
3833 "from the rows held"
3834 );
3835 }
3836
3837 #[test]
3839 fn booleans_and_categoricals_chart_as_categories() {
3840 let lf = df!("flag" => &[true, false], "n" => &[5_i64, 7])
3841 .unwrap()
3842 .lazy();
3843 let data =
3844 prepare_bar_data(&lf, "flag", "n", BarOrder::Label, BAR_CAP, &all_rows()).unwrap();
3845 assert_eq!(labels(&data), [Some("false"), Some("true")]);
3846
3847 let lf = df!("kind" => &["b", "a"], "n" => &[5_i64, 7])
3848 .unwrap()
3849 .lazy()
3850 .with_column(col("kind").cast(DataType::from_categories(Categories::global())));
3851 let schema = lf.clone().collect_schema().unwrap();
3852 assert!(is_category_dtype(schema.get("kind").unwrap()));
3853 let data =
3854 prepare_bar_data(&lf, "kind", "n", BarOrder::Value, BAR_CAP, &all_rows()).unwrap();
3855 assert_eq!(labels(&data), [Some("a"), Some("b")]);
3856 assert!(!is_category_dtype(&DataType::Float64));
3857 assert!(is_category_dtype(&DataType::UInt8));
3858 }
3859
3860 #[test]
3863 fn bar_values_follow_the_table_number_format() {
3864 use crate::numfmt::NumberFormat;
3865 let plain = NumberFormat::PLAIN;
3866 let thousands = NumberFormat::preset("thousands").unwrap();
3867 let european = NumberFormat::preset("european").unwrap();
3868 assert_eq!(format_bar_value(1_234_567.0, true, &plain), "1234567");
3869 assert_eq!(format_bar_value(1_234_567.0, true, &thousands), "1,234,567");
3870 assert_eq!(format_bar_value(22.0, false, &plain), "22.00");
3871 assert_eq!(format_bar_value(-9.9296, false, &plain), "-9.93");
3872 assert_eq!(format_bar_value(4213.7, false, &thousands), "4,213.70");
3873 assert_eq!(format_bar_value(4213.7, false, &european), "4.213,70");
3874 let one_place = NumberFormat {
3875 float_precision: Some(1),
3876 ..thousands
3877 };
3878 assert_eq!(format_bar_value(4213.74, false, &one_place), "4,213.7");
3879 assert_eq!(format_bar_value(0.0, false, &plain), "0.00");
3880 assert_eq!(format_bar_value(0.001, false, &plain), "1.00e-3");
3881
3882 let lf = df!("carrier" => &["UA", "AA"], "delay" => &[1234.5_f64, 7.0])
3883 .unwrap()
3884 .lazy();
3885 let data = bars(&lf, BarOrder::Value, BAR_CAP);
3886 let mut settings = crate::numfmt::NumberFormatSettings {
3887 format: NumberFormat::preset("thousands").unwrap(),
3888 ..Default::default()
3889 };
3890 assert_eq!(data.value_labels(&settings), ["1,234.50", "7.00"]);
3891 settings.enabled = false;
3892 assert_eq!(
3893 data.value_labels(&settings),
3894 ["1234.50", "7.00"],
3895 "F turns it off"
3896 );
3897 }
3898
3899 fn species(n_adelie: usize, n_gentoo: usize, n_chinstrap: usize, n_null: usize) -> LazyFrame {
3900 let mut species: Vec<Option<&str>> = Vec::new();
3901 let mut left = [
3903 (Some("Adelie"), n_adelie),
3904 (Some("Gentoo"), n_gentoo),
3905 (Some("Chinstrap"), n_chinstrap),
3906 (None, n_null),
3907 ];
3908 while left.iter().any(|(_, n)| *n > 0) {
3909 for (name, n) in &mut left {
3910 if *n > 0 {
3911 species.push(*name);
3912 *n -= 1;
3913 }
3914 }
3915 }
3916 df!("species" => species).unwrap().lazy()
3917 }
3918
3919 fn counts(data: &BarData) -> Vec<(Option<&str>, f64)> {
3920 data.bars
3921 .iter()
3922 .map(|b| (b.label.as_deref(), b.value))
3923 .collect()
3924 }
3925
3926 #[test]
3930 fn counts_are_exact_past_the_sample_size() {
3931 let lf = species(30_000, 15_000, 4_999, 1);
3932 let sampling = ChartSampling::rows(Some(1_000));
3933 let data = prepare_bar_counts(&lf, "species", BarOrder::Value, BAR_CAP, &sampling).unwrap();
3934 assert_eq!(
3935 counts(&data),
3936 [
3937 (Some("Adelie"), 30_000.0),
3938 (Some("Gentoo"), 15_000.0),
3939 (Some("Chinstrap"), 4_999.0),
3940 (None, 1.0)
3941 ]
3942 );
3943 assert_eq!(data.counted, Some(50_000), "counted past the sample size");
3944 assert_eq!(data.rows.sample_size, None);
3945 assert_eq!(data.value_column, "count");
3946 assert_eq!(
3947 data.value_labels(&crate::numfmt::NumberFormatSettings {
3948 format: crate::numfmt::NumberFormat::preset("thousands").unwrap(),
3949 ..Default::default()
3950 }),
3951 ["30,000", "15,000", "4,999", "1"],
3952 "whole numbers"
3953 );
3954
3955 let by_label =
3956 prepare_bar_counts(&lf, "species", BarOrder::Label, BAR_CAP, &sampling).unwrap();
3957 assert_eq!(
3958 counts(&by_label),
3959 [
3960 (Some("Adelie"), 30_000.0),
3961 (Some("Chinstrap"), 4_999.0),
3962 (Some("Gentoo"), 15_000.0),
3963 (None, 1.0)
3964 ],
3965 "the null category last"
3966 );
3967
3968 let every = prepare_bar_counts(
3970 &lf,
3971 "species",
3972 BarOrder::Value,
3973 BAR_CAP,
3974 &ChartSampling::rows(None),
3975 )
3976 .unwrap();
3977 assert_eq!(every.counted, None);
3978 let small = species(152, 124, 68, 0);
3979 let data =
3980 prepare_bar_counts(&small, "species", BarOrder::Value, BAR_CAP, &all_rows()).unwrap();
3981 assert_eq!(
3982 counts(&data),
3983 [
3984 (Some("Adelie"), 152.0),
3985 (Some("Gentoo"), 124.0),
3986 (Some("Chinstrap"), 68.0)
3987 ]
3988 );
3989 assert_eq!(data.counted, None);
3990 }
3991
3992 #[test]
3995 fn counts_cap_their_bars_and_stop_past_the_category_cap() {
3996 let lf = df!("carrier" => &["UA", "B6", "AA", "AA", "DL", "B6", "AA"])
3997 .unwrap()
3998 .lazy();
3999 let data = prepare_bar_counts(&lf, "carrier", BarOrder::Value, 2, &all_rows()).unwrap();
4000 assert_eq!(counts(&data), [(Some("AA"), 3.0), (Some("B6"), 2.0)]);
4001 assert_eq!(data.more, 2, "DL and UA are counted, not drawn");
4002 let data =
4003 prepare_bar_counts(&lf, "carrier", BarOrder::Value, BAR_CAP, &all_rows()).unwrap();
4004 assert_eq!(
4005 counts(&data)[2..],
4006 [(Some("DL"), 1.0), (Some("UA"), 1.0)],
4007 "ties A to Z"
4008 );
4009
4010 let err = count_bars(&lf, "carrier", BarOrder::Value, BAR_CAP, 3, &all_rows())
4011 .unwrap_err()
4012 .to_string();
4013 assert_eq!(
4014 err,
4015 "more than 3 categories of carrier: counting stopped. Count by a column with \
4016 fewer values"
4017 );
4018 let data = count_bars(&lf, "carrier", BarOrder::Value, BAR_CAP, 4, &all_rows()).unwrap();
4019 assert_eq!(data.bars.len(), 4, "four is not more than four");
4020 }
4021
4022 #[test]
4025 fn a_tally_merges_batches_and_stops_past_its_cap() {
4026 let batch = |ids: std::ops::Range<i64>| df!("id" => ids.collect::<Vec<_>>()).unwrap();
4027 let mut tally = Tally::new("id", 200_000);
4028 assert!(!tally.observe(&batch(0..70_000)).unwrap());
4029 assert_eq!(tally.merged, 70_000, "merged once the batches pile up");
4030 assert!(!tally.observe(&batch(0..10)).unwrap());
4031 let Counted::All { counts, rows } = tally.finish().unwrap() else {
4032 panic!("under the cap");
4033 };
4034 assert_eq!(rows, 70_010);
4035 let counts = counts.unwrap();
4036 assert_eq!(counts.height(), 70_000);
4037 let total: u64 = counts
4038 .column(COUNT_COLUMN)
4039 .unwrap()
4040 .u64()
4041 .unwrap()
4042 .sum()
4043 .unwrap();
4044 assert_eq!(total, 70_010);
4045
4046 let mut tally = Tally::new("id", 1_000);
4047 assert!(
4048 tally.observe(&batch(0..70_000)).unwrap(),
4049 "past the cap: stop reading"
4050 );
4051 assert!(matches!(tally.finish().unwrap(), Counted::TooMany));
4052 }
4053
4054 #[test]
4057 fn a_streamed_count_stops_past_its_cap_or_when_cancelled() {
4058 let ids = df!("id" => (0..200_000i64).collect::<Vec<_>>())
4059 .unwrap()
4060 .lazy();
4061 let cancel = Arc::default();
4062 assert!(matches!(
4063 stream_counts(&ids, "id", 1_000, &cancel).unwrap(),
4064 Counted::TooMany
4065 ));
4066
4067 let lf = species(30_000, 15_000, 4_999, 1);
4068 let sampling = ChartSampling::rows(Some(1_000));
4069 sampling.cancel.store(true, Ordering::Relaxed);
4070 let err = prepare_bar_counts(&lf, "species", BarOrder::Value, BAR_CAP, &sampling)
4071 .unwrap_err()
4072 .to_string();
4073 assert_eq!(err, "count cancelled");
4074 assert!(sampling.held.0.lock().unwrap().counts.is_empty());
4075 sampling.cancel.store(false, Ordering::Relaxed);
4076 let data = prepare_bar_counts(&lf, "species", BarOrder::Value, BAR_CAP, &sampling).unwrap();
4077 assert_eq!(data.counted, Some(50_000));
4078 }
4079
4080 #[test]
4083 fn counts_come_from_the_rows_held_and_are_held() {
4084 let dir = tempfile::tempdir().unwrap();
4085 let path = dir.path().join("flights.csv");
4086 std::fs::write(&path, "carrier,delay\nUA,1\nUA,2\nAA,3\n").unwrap();
4087 let lf = LazyCsvReader::new(PlRefPath::try_from_path(&path).unwrap())
4088 .finish()
4089 .unwrap();
4090 let sampling = all_rows();
4091 prepare_histogram_data(&lf, "delay", 10, ValueRange::All, &sampling).unwrap();
4092 std::fs::write(&path, "carrier,delay\nZZ,1\n").unwrap();
4093 let data = prepare_bar_counts(&lf, "carrier", BarOrder::Value, BAR_CAP, &sampling).unwrap();
4095 assert_eq!(counts(&data), [(Some("ZZ"), 1.0)]);
4096
4097 let dir = tempfile::tempdir().unwrap();
4098 let path = dir.path().join("flights.csv");
4099 std::fs::write(&path, "carrier,delay\nUA,1\nUA,2\nAA,3\n").unwrap();
4100 let lf = LazyCsvReader::new(PlRefPath::try_from_path(&path).unwrap())
4101 .finish()
4102 .unwrap();
4103 let sampling = all_rows();
4104 assert!(
4106 prepare_bar_data(&lf, "carrier", "delay", BarOrder::Value, BAR_CAP, &sampling).is_err()
4107 );
4108 std::fs::write(&path, "carrier,delay\nZZ,1\n").unwrap();
4109 let data = prepare_bar_counts(&lf, "carrier", BarOrder::Value, BAR_CAP, &sampling).unwrap();
4110 assert_eq!(
4111 counts(&data),
4112 [(Some("UA"), 2.0), (Some("AA"), 1.0)],
4113 "counted from the rows held"
4114 );
4115 sampling.held.0.lock().unwrap().rows = None;
4117 let data = prepare_bar_counts(&lf, "carrier", BarOrder::Label, BAR_CAP, &sampling).unwrap();
4118 assert_eq!(
4119 counts(&data),
4120 [(Some("AA"), 1.0), (Some("UA"), 2.0)],
4121 "another order from the count held"
4122 );
4123
4124 let sampling = ChartSampling {
4127 known_total: Some(3),
4128 ..all_rows()
4129 };
4130 std::fs::write(&path, "carrier,delay\nUA,1\nUA,2\nAA,3\n").unwrap();
4131 let data = prepare_bar_counts(&lf, "carrier", BarOrder::Value, BAR_CAP, &sampling).unwrap();
4132 assert_eq!(counts(&data), [(Some("UA"), 2.0), (Some("AA"), 1.0)]);
4133 let holding = sampling.held.0.lock().unwrap();
4134 let held = holding.rows.as_ref().expect("the rows read are held");
4135 assert_eq!(held.df.column("carrier").unwrap().len(), 3);
4136 }
4137
4138 fn returns() -> LazyFrame {
4143 let days: Vec<i32> = (0..90).collect();
4144 let n = days.len();
4145 let mut df = df!(
4146 "date" => days.iter().chain(&days).map(|d| 19723 + d).collect::<Vec<i32>>(),
4147 "symbol" => std::iter::repeat_n("A", n).chain(std::iter::repeat_n("B", n)).collect::<Vec<_>>(),
4148 "ret" => std::iter::repeat_n(1.0, n).chain(std::iter::repeat_n(2.0, n)).collect::<Vec<f64>>()
4149 )
4150 .unwrap();
4151 df.apply("date", |c| c.cast(&DataType::Date).unwrap())
4152 .unwrap();
4153 df.lazy()
4154 }
4155
4156 fn aggregate(
4157 lf: &LazyFrame,
4158 unit: crate::chart_modal::TimeUnit,
4159 aggregate: crate::chart_modal::Aggregate,
4160 cumulative: crate::chart_modal::Cumulative,
4161 color: Option<ColorSplit<'_>>,
4162 ) -> GroupedSeries {
4163 let schema = lf.clone().collect_schema().unwrap();
4164 let ys = ["ret".to_string()];
4165 prepare_aggregate_xy(
4166 lf,
4167 schema.as_ref(),
4168 &AggregateSpec {
4169 x: "date",
4170 time_unit: unit,
4171 ys: &ys,
4172 aggregate,
4173 quantile: 90,
4174 cumulative,
4175 color,
4176 },
4177 &all_rows(),
4178 )
4179 .unwrap()
4180 }
4181
4182 #[test]
4185 fn a_time_bucket_aggregates_every_row_per_month_and_color() {
4186 use crate::chart_modal::{Aggregate, Cumulative, TimeUnit};
4187 let lf = returns();
4188 let groups = [Some("A".to_string()), Some("B".to_string())];
4189 let split = ColorSplit {
4190 column: "symbol",
4191 groups: &groups,
4192 other: false,
4193 };
4194 let sum = aggregate(
4195 &lf,
4196 TimeUnit::Month,
4197 Aggregate::Sum,
4198 Cumulative::Off,
4199 Some(split),
4200 );
4201 assert_eq!(sum.names, ["A", "B"]);
4202 assert_eq!(sum.x_axis_kind, XAxisTemporalKind::Date);
4203 assert_eq!(sum.rows.total_rows, 180, "every row, no sample");
4204 let xs: Vec<f64> = sum.series[0].iter().map(|p| p.0).collect();
4206 assert_eq!(xs, [19723.0, 19754.0, 19783.0]);
4207 let a: Vec<f64> = sum.series[0].iter().map(|p| p.1).collect();
4208 let b: Vec<f64> = sum.series[1].iter().map(|p| p.1).collect();
4209 assert_eq!(a, [31.0, 29.0, 30.0]);
4210 assert_eq!(b, [62.0, 58.0, 60.0]);
4211
4212 let mean = aggregate(
4213 &lf,
4214 TimeUnit::Month,
4215 Aggregate::Mean,
4216 Cumulative::Off,
4217 Some(split),
4218 );
4219 assert!(mean.series[1].iter().all(|p| p.1 == 2.0));
4220 let count = aggregate(
4221 &lf,
4222 TimeUnit::Quarter,
4223 Aggregate::Count,
4224 Cumulative::Off,
4225 None,
4226 );
4227 assert_eq!(count.names, ["count"]);
4228 assert_eq!(
4229 count.series[0],
4230 [(19723.0, 180.0)],
4231 "one quarter, both symbols"
4232 );
4233 let weeks = aggregate(&lf, TimeUnit::Week, Aggregate::Max, Cumulative::Off, None);
4234 assert_eq!(weeks.series[0].len(), 13);
4236 assert!(weeks.series[0].iter().all(|p| p.1 == 2.0));
4237 }
4238
4239 #[test]
4242 fn cumulative_sums_or_compounds_along_x() {
4243 use crate::chart_modal::{Aggregate, Cumulative, TimeUnit};
4244 let lf = returns();
4245 let groups = [Some("A".to_string())];
4246 let split = ColorSplit {
4247 column: "symbol",
4248 groups: &groups,
4249 other: false,
4250 };
4251 let running = aggregate(
4252 &lf,
4253 TimeUnit::Month,
4254 Aggregate::Sum,
4255 Cumulative::Sum,
4256 Some(split),
4257 );
4258 let ys: Vec<f64> = running.series[0].iter().map(|p| p.1).collect();
4259 assert_eq!(ys, [31.0, 60.0, 90.0]);
4260 assert_eq!(running.names, ["A"], "only the groups picked");
4261
4262 let mut points = vec![(0.0, 0.1), (1.0, 0.1), (2.0, -0.5)];
4263 accumulate(&mut points, Cumulative::Compound);
4264 let ys: Vec<f64> = points.iter().map(|p| (p.1 * 1e6).round() / 1e6).collect();
4265 assert_eq!(ys, [0.1, 0.21, -0.395]);
4266 let mut points = vec![(0.0, 3.0), (1.0, 4.0)];
4267 accumulate(&mut points, Cumulative::Off);
4268 assert_eq!(points, [(0.0, 3.0), (1.0, 4.0)]);
4269 }
4270
4271 #[test]
4275 fn compound_runs_over_the_rows_of_each_bucket() {
4276 use crate::chart_modal::{Aggregate, Cumulative, TimeUnit};
4277 let mut df = df!(
4278 "date" => (0..90).map(|d| 19723 + d).collect::<Vec<i32>>(),
4279 "ret" => vec![0.01; 90]
4280 )
4281 .unwrap();
4282 df.apply("date", |c| c.cast(&DataType::Date).unwrap())
4283 .unwrap();
4284 let lf = df.lazy();
4285 for how in [Aggregate::Mean, Aggregate::Sum, Aggregate::Max] {
4286 let out = aggregate(&lf, TimeUnit::Month, how, Cumulative::Compound, None);
4287 let ys: Vec<f64> = out.series[0].iter().map(|p| p.1).collect();
4288 let want = [
4289 1.01f64.powi(31) - 1.0,
4290 1.01f64.powi(60) - 1.0,
4291 1.01f64.powi(90) - 1.0,
4292 ];
4293 for (y, w) in ys.iter().zip(want) {
4294 assert!((y - w).abs() < 1e-9, "{how:?}: {ys:?}");
4295 }
4296 }
4297 let rows = aggregate(
4298 &lf,
4299 TimeUnit::Month,
4300 Aggregate::Count,
4301 Cumulative::Compound,
4302 None,
4303 );
4304 let ys: Vec<f64> = rows.series[0].iter().map(|p| p.1).collect();
4305 assert_eq!(ys, [31.0, 60.0, 90.0], "a count runs as a count");
4306 }
4307
4308 #[test]
4310 fn a_bucket_with_no_values_is_a_gap() {
4311 use crate::chart_modal::{Aggregate, Cumulative, TimeUnit};
4312 let lf = df!(
4313 "date" => [0i32, 0, 1, 2],
4314 "ret" => [Some(1.0), Some(2.0), None, Some(4.0)]
4315 )
4316 .unwrap()
4317 .lazy()
4318 .with_column(col("date").cast(DataType::Date));
4319 let out = aggregate(&lf, TimeUnit::Day, Aggregate::Sum, Cumulative::Off, None);
4320 assert_eq!(out.series[0], [(0.0, 3.0), (2.0, 4.0)]);
4321 assert_eq!(out.breaks[0], [1], "the line breaks over day 1");
4322 }
4323
4324 #[test]
4327 fn an_x_of_too_many_values_is_refused_first() {
4328 use crate::chart_modal::{Aggregate, Cumulative};
4329 let n = AGGREGATE_POINTS_MAX as i64 + 10_000;
4330 let lf = df!("x" => (0..n).collect::<Vec<i64>>(), "ret" => vec![1.0; n as usize])
4331 .unwrap()
4332 .lazy();
4333 let schema = lf.clone().collect_schema().unwrap();
4334 let ys = ["ret".to_string()];
4335 let err = prepare_aggregate_xy(
4336 &lf,
4337 schema.as_ref(),
4338 &AggregateSpec {
4339 x: "x",
4340 time_unit: crate::chart_modal::TimeUnit::None,
4341 ys: &ys,
4342 aggregate: Aggregate::Mean,
4343 quantile: 90,
4344 cumulative: Cumulative::Off,
4345 color: None,
4346 },
4347 &all_rows(),
4348 )
4349 .unwrap_err();
4350 assert!(err.to_string().contains("values of x"), "{err}");
4351 }
4352
4353 #[test]
4356 fn color_takes_the_largest_groups_or_the_ones_picked() {
4357 let values: Vec<String> = (0..9)
4358 .flat_map(|i| std::iter::repeat_n(format!("v{i}"), 10 + i))
4359 .chain(std::iter::once("v0".to_string()))
4360 .collect();
4361 let lf = df!("c" => values).unwrap().lazy();
4362 let rows = value_rows(&lf, "c", &all_rows()).unwrap();
4363 assert_eq!(rows.values.len(), 9);
4364 assert_eq!(rows.rows, 10 + 11 + 12 + 13 + 14 + 15 + 16 + 17 + 18 + 1);
4365 assert_eq!(rows.values[0], (Some("v8".to_string()), 18));
4366 let top = color_groups(&rows, &[], 7);
4367 assert_eq!(
4368 top,
4369 ["v8", "v7", "v6", "v5", "v4", "v3", "v2"]
4370 .map(|v| Some(v.to_string()))
4371 .to_vec()
4372 );
4373 assert_eq!(rows.values[7], (Some("v0".to_string()), 11));
4375 let picked = [Some("v1".to_string()), None];
4376 assert_eq!(color_groups(&rows, &picked, 7), picked);
4377 assert_eq!(color_groups(&rows, &[], 3).len(), 3);
4379 assert_eq!(color_groups(&rows, &picked, 1), [Some("v1".to_string())]);
4380 }
4381
4382 #[test]
4385 fn a_color_splits_the_sampled_points() {
4386 let lf = df!(
4387 "x" => [3i64, 1, 2, 1, 2],
4388 "y" => [30.0, 10.0, 20.0, 1.0, 2.0],
4389 "c" => ["a", "a", "a", "b", "z"]
4390 )
4391 .unwrap()
4392 .lazy();
4393 let schema = lf.clone().collect_schema().unwrap();
4394 let groups = [Some("a".to_string()), Some("b".to_string())];
4395 let split = ColorSplit {
4396 column: "c",
4397 groups: &groups,
4398 other: false,
4399 };
4400 let out = prepare_xy_by(&lf, schema.as_ref(), "x", "y", split, &all_rows()).unwrap();
4401 assert_eq!(out.names, ["a", "b"]);
4402 assert_eq!(out.series[0], [(1.0, 10.0), (2.0, 20.0), (3.0, 30.0)]);
4403 assert_eq!(out.series[1], [(1.0, 1.0)]);
4404 }
4405
4406 #[test]
4410 fn stdev_quantile_first_and_last_per_x() {
4411 use crate::chart_modal::{Aggregate, Cumulative, TimeUnit};
4412 let lf = df!(
4414 "x" => [1i64, 1, 1, 1, 2, 3, 3],
4415 "y" => [Some(4.0), Some(1.0), Some(3.0), Some(2.0), Some(9.0), Some(5.0), None],
4416 "t" => [3i64, 1, 4, 2, 1, 2, 1],
4417 "c" => ["a", "b", "a", "b", "a", "a", "a"]
4418 )
4419 .unwrap()
4420 .lazy();
4421 let schema = lf.clone().collect_schema().unwrap();
4422 let ys = ["y".to_string()];
4423 let run = |lf: &LazyFrame, aggregate, quantile| {
4424 let out = prepare_aggregate_xy(
4425 lf,
4426 schema.as_ref(),
4427 &AggregateSpec {
4428 x: "x",
4429 time_unit: TimeUnit::None,
4430 ys: &ys,
4431 aggregate,
4432 quantile,
4433 cumulative: Cumulative::Off,
4434 color: None,
4435 },
4436 &all_rows(),
4437 )
4438 .unwrap();
4439 out.series[0].clone()
4440 };
4441 let stdev = run(&lf, Aggregate::Stdev, 90);
4443 assert_eq!(stdev.len(), 1, "x=2 and x=3 have one value each: {stdev:?}");
4444 assert!((stdev[0].1 - (5.0f64 / 3.0).sqrt()).abs() < 1e-12);
4445 let p90 = run(&lf, Aggregate::Quantile, 90);
4447 assert!((p90[0].1 - 3.7).abs() < 1e-12, "{p90:?}");
4448 assert_eq!(p90[1], (2.0, 9.0));
4449 let p25 = run(&lf, Aggregate::Quantile, 25);
4450 assert!((p25[0].1 - 1.75).abs() < 1e-12, "{p25:?}");
4451 assert_eq!(
4454 run(&lf, Aggregate::First, 90),
4455 [(1.0, 4.0), (2.0, 9.0), (3.0, 5.0)]
4456 );
4457 assert_eq!(
4458 run(&lf, Aggregate::Last, 90),
4459 [(1.0, 2.0), (2.0, 9.0), (3.0, 5.0)]
4460 );
4461 let sorted = lf.clone().sort(["t"], Default::default());
4463 assert_eq!(run(&sorted, Aggregate::First, 90)[0], (1.0, 1.0));
4464 assert_eq!(run(&sorted, Aggregate::Last, 90)[0], (1.0, 3.0));
4465 let groups = [Some("a".to_string()), Some("b".to_string())];
4467 let bars = prepare_bar_aggregate(
4468 &lf,
4469 &BarAggregate {
4470 category: "x",
4471 value: Some("y"),
4472 aggregate: Aggregate::Last,
4473 quantile: 90,
4474 color: Some(ColorSplit {
4475 column: "c",
4476 groups: &groups,
4477 other: false,
4478 }),
4479 order: BarOrder::Label,
4480 cap: BAR_CAP,
4481 },
4482 &all_rows(),
4483 )
4484 .unwrap();
4485 assert_eq!(bars.bars[0].by_group, [Some(3.0), Some(2.0)]);
4486 assert_eq!(bars.value_column, "last y");
4487 let p = prepare_bar_aggregate(
4488 &lf,
4489 &BarAggregate {
4490 category: "x",
4491 value: Some("y"),
4492 aggregate: Aggregate::Quantile,
4493 quantile: 90,
4494 color: None,
4495 order: BarOrder::Label,
4496 cap: BAR_CAP,
4497 },
4498 &all_rows(),
4499 )
4500 .unwrap();
4501 assert_eq!(p.value_column, "p90 y");
4502 }
4503
4504 #[test]
4507 fn distinct_counts_any_y_per_x() {
4508 use crate::chart_modal::{Aggregate, Cumulative, TimeUnit};
4509 let mut df = df!(
4510 "date" => [19723i32, 19723, 19723, 19724, 19724, 19754, 19755],
4511 "name" => [Some("Ann"), Some("Bo"), Some("Ann"), None, None, Some("Cy"), Some("Di")],
4512 "sex" => ["F", "M", "F", "F", "M", "M", "M"]
4513 )
4514 .unwrap();
4515 df.apply("date", |c| c.cast(&DataType::Date).unwrap())
4516 .unwrap();
4517 let lf = df.lazy();
4518 let schema = lf.clone().collect_schema().unwrap();
4519 let ys = ["name".to_string()];
4520 let distinct = |unit, color| {
4521 prepare_aggregate_xy(
4522 &lf,
4523 schema.as_ref(),
4524 &AggregateSpec {
4525 x: "date",
4526 time_unit: unit,
4527 ys: &ys,
4528 aggregate: Aggregate::Distinct,
4529 quantile: 90,
4530 cumulative: Cumulative::Off,
4531 color,
4532 },
4533 &all_rows(),
4534 )
4535 .unwrap()
4536 };
4537 let by_day = distinct(TimeUnit::Day, None);
4538 let ys_of = |s: &[(f64, f64)]| s.iter().map(|p| p.1).collect::<Vec<_>>();
4539 assert_eq!(ys_of(&by_day.series[0]), [2.0, 1.0, 1.0]);
4541 assert_eq!(by_day.breaks[0], [1], "the day of nulls breaks the line");
4542 let by_month = distinct(TimeUnit::Month, None);
4543 assert_eq!(ys_of(&by_month.series[0]), [2.0, 2.0], "Ann, Bo; Cy, Di");
4544 let groups = [Some("F".to_string()), Some("M".to_string())];
4545 let split = ColorSplit {
4546 column: "sex",
4547 groups: &groups,
4548 other: false,
4549 };
4550 let colored = distinct(TimeUnit::Month, Some(split));
4551 assert_eq!(colored.names, ["F", "M"]);
4552 assert_eq!(ys_of(&colored.series[0]), [1.0], "Ann");
4553 assert_eq!(ys_of(&colored.series[1]), [1.0, 2.0], "Bo; Cy, Di");
4554 let bars = prepare_bar_aggregate(
4556 &lf,
4557 &BarAggregate {
4558 category: "sex",
4559 value: Some("name"),
4560 aggregate: Aggregate::Distinct,
4561 quantile: 90,
4562 color: None,
4563 order: BarOrder::Label,
4564 cap: BAR_CAP,
4565 },
4566 &all_rows(),
4567 )
4568 .unwrap();
4569 let values: Vec<f64> = bars.bars.iter().map(|b| b.value).collect();
4570 assert_eq!(values, [1.0, 3.0]);
4571 assert!(bars.value_dtype.is_integer());
4572 assert_eq!(bars.value_column, "distinct name");
4573 }
4574
4575 #[test]
4579 fn other_gathers_every_value_without_a_series() {
4580 use crate::chart_modal::{Aggregate, Cumulative, TimeUnit};
4581 let lf = df!(
4582 "x" => [1i64, 1, 2, 2, 3, 3],
4583 "y" => [10.0, 1.0, 20.0, 2.0, 30.0, 4.0],
4584 "c" => [Some("a"), Some("b"), Some("a"), Some("z"), Some("a"), None]
4585 )
4586 .unwrap()
4587 .lazy();
4588 let schema = lf.clone().collect_schema().unwrap();
4589 let groups = [Some("a".to_string())];
4590 let split = |other| ColorSplit {
4591 column: "c",
4592 groups: &groups,
4593 other,
4594 };
4595 let out = prepare_xy_by(&lf, schema.as_ref(), "x", "y", split(true), &all_rows()).unwrap();
4597 assert_eq!(out.names, ["a", OTHER]);
4598 assert!(out.other);
4599 assert_eq!(out.series[1], [(1.0, 1.0), (2.0, 2.0), (3.0, 4.0)]);
4600 let out = prepare_xy_by(&lf, schema.as_ref(), "x", "y", split(false), &all_rows()).unwrap();
4601 assert_eq!(out.names, ["a"]);
4602 assert!(!out.other);
4603 let ys = ["y".to_string()];
4605 let line = |other| {
4606 prepare_aggregate_xy(
4607 &lf,
4608 schema.as_ref(),
4609 &AggregateSpec {
4610 x: "x",
4611 time_unit: TimeUnit::None,
4612 ys: &ys,
4613 aggregate: Aggregate::Sum,
4614 quantile: 90,
4615 cumulative: Cumulative::Off,
4616 color: Some(split(other)),
4617 },
4618 &all_rows(),
4619 )
4620 .unwrap()
4621 };
4622 let on = line(true);
4623 assert_eq!(on.names, ["a", OTHER]);
4624 assert_eq!(on.series[1], [(1.0, 1.0), (2.0, 2.0), (3.0, 4.0)]);
4625 assert_eq!(on.rows.total_rows, 6, "every row is in a series");
4626 let off = line(false);
4627 assert_eq!(off.names, ["a"]);
4628 assert_eq!(off.rows.total_rows, 3);
4629 let bars = |other| {
4631 let spec = BarAggregate {
4632 category: "x",
4633 value: Some("y"),
4634 aggregate: Aggregate::Mean,
4635 quantile: 90,
4636 color: Some(split(other)),
4637 order: BarOrder::Label,
4638 cap: BAR_CAP,
4639 };
4640 prepare_bar_aggregate(&lf, &spec, &all_rows()).unwrap()
4641 };
4642 let on = bars(true);
4643 assert_eq!(on.groups, ["a", OTHER]);
4644 assert!(on.other);
4645 let by: Vec<Vec<Option<f64>>> = on.bars.iter().map(|b| b.by_group.clone()).collect();
4646 assert_eq!(
4647 by,
4648 [
4649 vec![Some(10.0), Some(1.0)],
4650 vec![Some(20.0), Some(2.0)],
4651 vec![Some(30.0), Some(4.0)]
4652 ]
4653 );
4654 let off = bars(false);
4655 assert_eq!(off.groups, ["a"]);
4656 assert!(off.bars.iter().all(|b| b.by_group.len() == 1));
4657 }
4658
4659 #[test]
4662 fn bars_aggregate_per_category_and_color() {
4663 use crate::chart_modal::Aggregate;
4664 let lf = df!(
4665 "carrier" => ["UA", "UA", "UA", "AA", "AA"],
4666 "origin" => ["EWR", "EWR", "JFK", "EWR", "JFK"],
4667 "delay" => [10.0, 20.0, 5.0, 1.0, 50.0]
4668 )
4669 .unwrap()
4670 .lazy();
4671 let groups = [Some("EWR".to_string()), Some("JFK".to_string())];
4672 let split = ColorSplit {
4673 column: "origin",
4674 groups: &groups,
4675 other: false,
4676 };
4677 let spec = BarAggregate {
4678 category: "carrier",
4679 value: Some("delay"),
4680 aggregate: Aggregate::Mean,
4681 quantile: 90,
4682 color: Some(split),
4683 order: BarOrder::Value,
4684 cap: BAR_CAP,
4685 };
4686 let data = prepare_bar_aggregate(&lf, &spec, &all_rows()).unwrap();
4687 assert_eq!(data.groups, ["EWR", "JFK"]);
4688 assert_eq!(data.value_column, "mean delay");
4689 assert_eq!(data.rows.total_rows, 5);
4690 let bars: Vec<(Option<&str>, Vec<Option<f64>>)> = data
4691 .bars
4692 .iter()
4693 .map(|b| (b.label.as_deref(), b.by_group.clone()))
4694 .collect();
4695 assert_eq!(
4696 bars,
4697 [
4698 (Some("AA"), vec![Some(1.0), Some(50.0)]),
4699 (Some("UA"), vec![Some(15.0), Some(5.0)])
4700 ],
4701 "AA's largest group is larger"
4702 );
4703 let count = BarAggregate {
4704 value: None,
4705 aggregate: Aggregate::Count,
4706 quantile: 90,
4707 ..spec
4708 };
4709 let data = prepare_bar_aggregate(&lf, &count, &all_rows()).unwrap();
4710 assert_eq!(data.bars[0].label.as_deref(), Some("UA"));
4711 assert_eq!(data.bars[0].value, 3.0, "a count adds up across groups");
4712 assert!(data.value_dtype.is_integer(), "counts print whole");
4713 let odd = df!(
4715 "carrier" => ["UA", "AA", "DL"],
4716 "delay" => [Some(f64::NAN), None, Some(1.0)]
4717 )
4718 .unwrap()
4719 .lazy();
4720 let mean = BarAggregate {
4721 color: None,
4722 ..spec
4723 };
4724 let data = prepare_bar_aggregate(&odd, &mean, &all_rows()).unwrap();
4725 let labels: Vec<Option<&str>> = data.bars.iter().map(|b| b.label.as_deref()).collect();
4726 assert_eq!(labels, [Some("DL")]);
4727 assert_eq!(data.no_value, 2);
4728 let sum = BarAggregate {
4729 aggregate: Aggregate::Sum,
4730 quantile: 90,
4731 color: None,
4732 ..spec
4733 };
4734 let data = prepare_bar_aggregate(&lf, &sum, &all_rows()).unwrap();
4735 assert_eq!(
4736 data.bars
4737 .iter()
4738 .map(|b| (b.label.as_deref(), b.value))
4739 .collect::<Vec<_>>(),
4740 [(Some("AA"), 51.0), (Some("UA"), 35.0)]
4741 );
4742 }
4743
4744 #[test]
4747 fn a_histogram_splits_into_groups_on_shared_bins() {
4748 let lf = df!(
4749 "v" => [0.0, 1.0, 2.0, 3.0, 0.0, 0.0],
4750 "g" => ["a", "a", "a", "a", "b", "b"]
4751 )
4752 .unwrap()
4753 .lazy();
4754 let groups = [Some("a".to_string()), Some("b".to_string())];
4755 let split = ColorSplit {
4756 column: "g",
4757 groups: &groups,
4758 other: false,
4759 };
4760 let data =
4761 prepare_histogram_by(&lf, "v", 3, ValueRange::All, true, Some(split), &all_rows())
4762 .unwrap();
4763 assert_eq!(data.bins.len(), 3);
4764 assert_eq!(data.groups.len(), 2);
4765 assert_eq!(data.groups[0].counts, [0.25, 0.25, 0.5]);
4766 assert_eq!(data.groups[1].counts, [1.0, 0.0, 0.0]);
4767 assert_eq!(data.max_count, 1.0);
4768 let total: f64 = data.bins.iter().map(|b| b.count).sum();
4769 assert!((total - 1.0).abs() < 1e-9, "the whole is a share too");
4770 }
4771
4772 #[test]
4774 fn a_box_per_category() {
4775 let lf = df!(
4776 "v" => [1.0, 2.0, 3.0, 10.0, 20.0],
4777 "k" => ["x", "x", "x", "y", "y"]
4778 )
4779 .unwrap()
4780 .lazy();
4781 let groups = [Some("y".to_string()), Some("x".to_string())];
4782 let split = ColorSplit {
4783 column: "k",
4784 groups: &groups,
4785 other: false,
4786 };
4787 let data = prepare_box_by(&lf, "v", split, ValueRange::All, &all_rows()).unwrap();
4788 let names: Vec<&str> = data.stats.iter().map(|s| s.name.as_str()).collect();
4789 assert_eq!(names, ["y", "x"]);
4790 assert_eq!(data.stats[1].median, 2.0);
4791 assert_eq!((data.y_min, data.y_max), (1.0, 20.0));
4792 }
4793}