1use crate::analysis::sampling::ReadWatch;
12use crate::chart::chart_data::{COUNT_COLUMN, Counted, Tally, count_frame};
13use color_eyre::Result;
14use color_eyre::eyre::eyre;
15use polars::prelude::*;
16
17pub const TOP_N: usize = 1_000;
19
20pub const MAX_DISTINCT: usize = 2_000_000;
23
24pub const LARGE_SAMPLES: usize = 100;
27
28pub const HISTOGRAM_BINS: usize = 40;
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
34pub enum Order {
35 #[default]
37 Count,
38 Value,
40}
41
42impl Order {
43 pub fn toggled(self) -> Self {
44 match self {
45 Order::Count => Order::Value,
46 Order::Value => Order::Count,
47 }
48 }
49}
50
51#[derive(Debug, Clone, PartialEq)]
53pub enum Read {
54 Quick {
57 sample_rows: usize,
58 seed: u64,
59 remote: bool,
60 },
61 Exact,
63}
64
65pub struct Plan {
67 pub lf: LazyFrame,
69 pub column: String,
70 pub read: Read,
71 pub known_total: Option<usize>,
73 pub streaming: bool,
74}
75
76impl Plan {
77 pub fn run(&self, watch: &ReadWatch) -> Result<ValueCounts> {
79 let lf = self.lf.clone().select([col(self.column.as_str())]);
80 let dtype = lf
81 .clone()
82 .collect_schema()?
83 .get(self.column.as_str())
84 .cloned()
85 .ok_or_else(|| eyre!("no column {}", self.column))?;
86 if let Some((rows, seed)) = self.sample(&lf) {
87 let read = crate::analysis::sampling::sample_rows_counting(
88 &lf,
89 Some(rows),
90 self.known_total,
91 seed,
92 self.streaming,
93 Some(watch),
94 None,
95 )?;
96 watch.check()?;
97 let counted = count_frame(&read.rows.df, &self.column, MAX_DISTINCT)?;
98 let of = read.rows.sample_size.map(|_| read.rows.total_rows);
100 return ValueCounts::new(&self.column, dtype, counted, of);
101 }
102 let counted = stream_counts(&lf, &self.column, watch)?;
103 ValueCounts::new(&self.column, dtype, counted, None)
104 }
105
106 fn sample(&self, lf: &LazyFrame) -> Option<(usize, u64)> {
109 let Read::Quick {
110 sample_rows,
111 seed,
112 remote,
113 } = self.read
114 else {
115 return None;
116 };
117 if sample_rows == 0 {
118 return None;
119 }
120 let large = remote
121 || self
122 .known_total
123 .is_none_or(|rows| rows > sample_rows.saturating_mul(LARGE_SAMPLES));
124 (large && crate::analysis::sampling::slices_reach_into_the_scan(lf))
125 .then_some((sample_rows, seed))
126 }
127}
128
129fn stream_counts(lf: &LazyFrame, column: &str, watch: &ReadWatch) -> Result<Counted> {
132 let tally = crate::analysis::sampling::stream_fold(
133 lf.clone(),
134 Some(watch),
135 false,
136 Tally::new(column, MAX_DISTINCT),
137 |tally, batch| tally.observe(&batch),
138 )?;
139 watch.check()?;
141 Ok(tally.finish()?)
142}
143
144#[derive(Debug, Clone, Copy, PartialEq)]
147pub enum Number {
148 Int(i128),
149 Float(f64),
150}
151
152#[derive(Debug, Clone, PartialEq, Default)]
154pub struct Summary {
155 pub rows: usize,
157 pub distinct: usize,
159 pub nulls: usize,
160 pub sum: Option<Number>,
162 pub mean: Option<f64>,
163 pub min: Option<AnyValue<'static>>,
165 pub max: Option<AnyValue<'static>>,
166}
167
168impl Summary {
169 pub fn of(values: &Series, counts: &[u64]) -> PolarsResult<Self> {
172 let rows = counts.iter().sum::<u64>() as usize;
173 let nulls: u64 = values
174 .is_null()
175 .iter()
176 .zip(counts)
177 .filter(|(null, _)| null.unwrap_or(false))
178 .map(|(_, n)| n)
179 .sum();
180 let nulls = nulls as usize;
181 let distinct = values.len() - values.null_count();
182 let dtype = values.dtype();
183 let numeric = dtype.is_primitive_numeric() || matches!(dtype, DataType::Decimal(..));
184 let ordered = numeric || dtype.is_temporal();
185 let mut summary = Summary {
186 rows,
187 distinct,
188 nulls,
189 ..Summary::default()
190 };
191 if ordered && distinct > 0 {
192 summary.min = Some(values.min_reduce()?.value().clone().into_static());
193 summary.max = Some(values.max_reduce()?.value().clone().into_static());
194 }
195 if numeric {
196 let sum = weighted_sum(values, counts)?;
197 let present = rows - nulls;
198 summary.mean = (present > 0).then(|| {
199 let total = match sum {
200 Number::Int(n) => n as f64,
201 Number::Float(f) => f,
202 };
203 total / present as f64
204 });
205 summary.sum = Some(sum);
206 }
207 Ok(summary)
208 }
209}
210
211fn weighted_sum(values: &Series, counts: &[u64]) -> PolarsResult<Number> {
213 let dtype = values.dtype();
214 if dtype.is_integer() && !matches!(dtype, DataType::Int128) {
215 let total: i128 = if matches!(dtype, DataType::UInt64) {
218 values
219 .u64()?
220 .iter()
221 .zip(counts)
222 .filter_map(|(v, n)| v.map(|v| v as i128 * *n as i128))
223 .sum()
224 } else {
225 values
226 .cast(&DataType::Int64)?
227 .i64()?
228 .iter()
229 .zip(counts)
230 .filter_map(|(v, n)| v.map(|v| v as i128 * *n as i128))
231 .sum()
232 };
233 return Ok(Number::Int(total));
234 }
235 let total = values
236 .cast(&DataType::Float64)?
237 .f64()?
238 .iter()
239 .zip(counts)
240 .filter_map(|(v, n)| v.map(|v| v * *n as f64))
241 .sum();
242 Ok(Number::Float(total))
243}
244
245fn histogram_of(
250 column: &str,
251 values: &Series,
252 rows: &[u64],
253) -> Option<crate::chart::chart_data::HistogramData> {
254 use crate::chart::chart_data::{Clipped, HistogramBin, HistogramData, RowsRead, ValueRange};
255 let dtype = values.dtype();
256 if !dtype.is_primitive_numeric() {
257 return None;
258 }
259 let as_f64 = values.cast(&DataType::Float64).ok()?;
260 let mut pairs: Vec<(f64, u64)> = as_f64
261 .f64()
262 .ok()?
263 .iter()
264 .zip(rows)
265 .filter_map(|(v, n)| Some((v.filter(|v| v.is_finite())?, *n)))
266 .collect();
267 pairs.sort_by(|a, b| a.0.total_cmp(&b.0));
268 let (min, max) = (pairs.first()?.0, pairs.last()?.0);
269 let total: u64 = pairs.iter().map(|p| p.1).sum();
270 let at = |q: f64| {
272 let wanted = ((q * total as f64).ceil() as u64).max(1);
273 let mut seen = 0;
274 for (v, n) in &pairs {
275 seen += n;
276 if seen >= wanted {
277 return *v;
278 }
279 }
280 max
281 };
282 let (p1, p99) = (at(0.01), at(0.99));
283 let clip = p99 > p1 && (max - min) > 10.0 * (p99 - p1);
284 let (lo, hi) = if clip { (p1, p99) } else { (min, max) };
285 let (bins, width, x_min) = if dtype.is_integer() && hi - lo < HISTOGRAM_BINS as f64 {
286 ((hi - lo) as usize + 1, 1.0, lo - 0.5)
287 } else if hi > lo {
288 (HISTOGRAM_BINS, (hi - lo) / HISTOGRAM_BINS as f64, lo)
289 } else {
290 (1, 1.0, lo - 0.5)
291 };
292 let mut counts = vec![0.0_f64; bins];
293 let mut outside = 0;
294 for (v, n) in pairs {
295 if v < lo || v > hi {
296 outside += n as usize;
297 continue;
298 }
299 let bin = (((v - x_min) / width).floor().max(0.0) as usize).min(bins - 1);
300 counts[bin] += n as f64;
301 }
302 let max_count = counts.iter().copied().fold(0.0, f64::max);
303 Some(HistogramData {
304 column: column.to_string(),
305 bins: counts
306 .into_iter()
307 .enumerate()
308 .map(|(i, count)| HistogramBin {
309 center: x_min + (i as f64 + 0.5) * width,
310 count,
311 })
312 .collect(),
313 groups: Vec::new(),
314 other: false,
315 share: false,
316 x_min,
317 x_max: x_min + bins as f64 * width,
318 max_count,
319 rows: RowsRead {
320 total_rows: total as usize,
321 sample_size: None,
322 envelope_steps: None,
323 seed: None,
324 },
325 clipped: clip.then_some(Clipped {
326 range: ValueRange::Percentile1To99,
327 outside,
328 }),
329 })
330}
331
332#[derive(Debug, Clone, Copy, PartialEq, Eq)]
334pub enum LineKind {
335 Value(usize),
337 Null,
339 Other(usize),
341}
342
343#[derive(Debug, Clone, Copy, PartialEq, Eq)]
346pub struct Line {
347 pub kind: LineKind,
348 pub rows: u64,
349 pub cumulative: u64,
350}
351
352#[derive(Debug, Clone)]
355pub struct ValueCounts {
356 pub column: String,
357 pub dtype: DataType,
358 counts: DataFrame,
360 by_count: Vec<usize>,
362 by_value: Vec<usize>,
363 null_at: Option<usize>,
365 pub sampled_of: Option<usize>,
367 pub summary: Summary,
368 count_lines: Vec<Line>,
369 value_lines: Vec<Line>,
370 pub histogram: Option<crate::chart::chart_data::HistogramData>,
373}
374
375impl ValueCounts {
376 pub(crate) fn new(
377 column: &str,
378 dtype: DataType,
379 counted: Counted,
380 sampled_of: Option<usize>,
381 ) -> Result<Self> {
382 let counts = match counted {
383 Counted::All {
384 counts: Some(counts),
385 ..
386 } => counts,
387 Counted::All { counts: None, .. } => DataFrame::new_infer_height(vec![
388 Column::new_empty(column.into(), &dtype),
389 Column::new_empty(COUNT_COLUMN.into(), &DataType::UInt64),
390 ])?,
391 Counted::TooMany => {
392 return Err(eyre!(
393 "more than {} distinct values: counting stopped",
394 crate::numfmt::group_chrome(MAX_DISTINCT)
395 ));
396 }
397 };
398 let values = counts.column(column)?.as_materialized_series().clone();
399 let rows = row_counts(&counts)?;
400 let summary = Summary::of(&values, &rows)?;
401 let present = values.len() - values.null_count();
403 let by_value: Vec<usize> = values
404 .arg_sort(
405 SortOptions::default()
406 .with_nulls_last(true)
407 .with_maintain_order(true),
408 )
409 .iter()
410 .flatten()
411 .map(|i| i as usize)
412 .take(present)
413 .collect();
414 let null_at = values
415 .is_null()
416 .iter()
417 .position(|null| null.unwrap_or(false));
418 let mut by_count = by_value.clone();
420 by_count.sort_by(|&a, &b| rows[b].cmp(&rows[a]));
421 let null_rows = summary.nulls as u64;
422 let histogram = histogram_of(column, &values, &rows);
423 Ok(Self {
424 histogram,
425 column: column.to_string(),
426 dtype,
427 count_lines: listing(&by_count, &rows, null_rows, true),
428 value_lines: listing(&by_value, &rows, null_rows, false),
429 by_count,
430 by_value,
431 null_at,
432 counts,
433 sampled_of,
434 summary,
435 })
436 }
437
438 pub fn is_sample(&self) -> bool {
440 self.sampled_of.is_some()
441 }
442
443 pub fn lines(&self, order: Order) -> &[Line] {
445 match order {
446 Order::Count => &self.count_lines,
447 Order::Value => &self.value_lines,
448 }
449 }
450
451 pub fn value(&self, at: usize) -> PolarsResult<AnyValue<'static>> {
453 Ok(self.counts.column(&self.column)?.get(at)?.into_static())
454 }
455
456 pub fn table(&self, order: Order) -> PolarsResult<DataFrame> {
459 let rows = row_counts(&self.counts)?;
460 let picked: Vec<usize> = match order {
461 Order::Count => &self.by_count,
462 Order::Value => &self.by_value,
463 }
464 .iter()
465 .copied()
466 .chain(self.null_at)
467 .collect();
468 let values = self.counts.column(&self.column)?.take(&IdxCa::from_vec(
469 "order".into(),
470 picked.iter().map(|&i| i as IdxSize).collect(),
471 ))?;
472 let counts: Vec<u64> = picked.iter().map(|&i| rows[i]).collect();
473 let total = self.summary.rows.max(1) as f64;
474 let percent: Vec<f64> = counts.iter().map(|&n| n as f64 * 100.0 / total).collect();
475 let mut running = 0u64;
476 let cumulative: Vec<f64> = counts
477 .iter()
478 .map(|n| {
479 running += n;
480 running as f64 * 100.0 / total
481 })
482 .collect();
483 let mut names = vec![self.column.clone()];
485 let mut name = |wanted: &str| {
486 let mut name = wanted.to_string();
487 while names.contains(&name) {
488 name.push('_');
489 }
490 names.push(name.clone());
491 PlSmallStr::from(name)
492 };
493 DataFrame::new_infer_height(vec![
494 values,
495 Column::new(name("count"), counts),
496 Column::new(name("percent"), percent),
497 Column::new(name("cumulative_percent"), cumulative),
498 ])
499 }
500}
501
502fn row_counts(counts: &DataFrame) -> PolarsResult<Vec<u64>> {
503 Ok(counts
504 .column(COUNT_COLUMN)?
505 .u64()?
506 .into_no_null_iter()
507 .collect())
508}
509
510fn listing(order: &[usize], rows: &[u64], nulls: u64, by_count: bool) -> Vec<Line> {
515 let mut lines = Vec::with_capacity(order.len().min(TOP_N) + 2);
516 let mut cumulative = 0u64;
517 let mut push = |kind, n: u64, lines: &mut Vec<Line>| {
518 cumulative += n;
519 lines.push(Line {
520 kind,
521 rows: n,
522 cumulative,
523 });
524 };
525 let mut null_owed = nulls > 0;
526 for &at in order.iter().take(TOP_N) {
527 if null_owed && by_count && rows[at] < nulls {
528 push(LineKind::Null, nulls, &mut lines);
529 null_owed = false;
530 }
531 push(LineKind::Value(at), rows[at], &mut lines);
532 }
533 if null_owed {
534 push(LineKind::Null, nulls, &mut lines);
535 }
536 if order.len() > TOP_N {
537 let rest: u64 = order[TOP_N..].iter().map(|&at| rows[at]).sum();
538 push(LineKind::Other(order.len() - TOP_N), rest, &mut lines);
539 }
540 lines
541}
542
543#[cfg(test)]
544mod tests;