Skip to main content

datui_lib/table/
copy.rs

1//! Drilling into groups and values, the inspector's rows, and what a copy takes.
2
3use super::*;
4
5/// The grouped view and pipeline state a drill-down saved, so filters and sort inside
6/// the group apply to it and `drill_up` restores the grouped view.
7#[derive(Clone)]
8pub(super) struct GroupedView {
9    pub(super) lf: LazyFrame,
10    pub(super) base_lf: LazyFrame,
11    /// The schemas of `base_lf` and of the view, so coming back resolves neither.
12    base_schema: Arc<Schema>,
13    schema: Arc<Schema>,
14    pub(super) filters: Vec<FilterStatement>,
15    pub(super) sort_columns: Vec<String>,
16    pub(super) sort_descending: Vec<bool>,
17    pub(super) sort_ascending: bool,
18    /// Whether `lf` has the hidden drift column and its groups' meaning, so drilling up
19    /// restores those cells and their notes.
20    drift: bool,
21    drift_groups: Arc<Vec<crate::formats::schema_union::DriftGroup>>,
22    /// Whether `lf` numbers its rows itself (`#`).
23    view_numbered: bool,
24    notes: Vec<crate::notes::Note>,
25    pub(super) group_source: Option<GroupSource>,
26    /// Where the user was, so coming back puts the cursor on the group drilled into
27    /// with the key columns still frozen.
28    column_order: Vec<String>,
29    locked_columns_count: usize,
30    start_row: usize,
31    termcol_index: usize,
32    cursor_column: Option<String>,
33    selected: Option<usize>,
34    /// Drilled into from Value Counts rather than from a grouped row.
35    by_value: bool,
36    /// How `base_lf` was built, for Copy as Python.
37    base_steps: Vec<Step>,
38    lineage: Lineage,
39}
40
41/// The rows a grouped result came from and how its keys were computed, so a drill-down
42/// finds a group's rows even from aggregates. Recorded by the grouping query, not
43/// inferred from (possibly renamed) columns.
44#[derive(Clone)]
45pub(super) struct GroupSource {
46    /// The rows before grouping, after any filter the query applied first.
47    pub(super) rows: LazyFrame,
48    /// Each key's column in the result, with the expression that computes it from `rows`.
49    pub(super) keys: Vec<(PlSmallStr, Expr)>,
50    /// Columns `rows` carries only to compute keys, left out of a drill.
51    pub(super) scratch: Vec<PlSmallStr>,
52    /// Whether the result's list columns are each group's rows, as a `by` query's are.
53    /// A SQL result's lists are values it computed, such as `ARRAY_AGG`.
54    pub(super) rows_in_lists: bool,
55    /// The same as Copy as Python steps: how `rows` was built, and each key as
56    /// Python code, aliases undone. None where the script cannot say.
57    pub(super) python_rows: Option<Vec<Step>>,
58    pub(super) python_keys: Vec<Option<String>>,
59    /// Which loaded column each column of `rows` is.
60    pub(super) lineage: Lineage,
61}
62
63/// One field the row inspector lists: a column of the frame on screen.
64#[derive(Debug, Clone, PartialEq)]
65pub struct InspectField {
66    pub name: String,
67    pub dtype: DataType,
68    /// Hidden from the table, so not among the rows read for it.
69    pub hidden: bool,
70}
71
72impl InspectField {
73    /// Whether the table's rows hold this field's value: a hidden column is not
74    /// read for them, and a binary one is read as a stub.
75    pub fn buffered(&self) -> bool {
76        !self.hidden && !matches!(self.dtype, DataType::Binary)
77    }
78}
79
80/// The selected row as the buffer holds it; see [`DataTableState::inspect_row`].
81#[derive(Clone)]
82pub struct InspectRow {
83    /// The row's index in the view.
84    pub row: usize,
85    /// The frame it is a row of (`len_generation`), which a sort, a filter or a
86    /// query replaces.
87    pub frame: u64,
88    /// The row number the table shows.
89    pub display_row: usize,
90    /// The row, one column per table column, raw.
91    pub values: DataFrame,
92    /// Which drift group its file is in, when the files differ.
93    pub drift_group: Option<u32>,
94}
95
96/// What a null means where the files of a dataset differ.
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98pub enum NullKind {
99    Null,
100    Absent,
101    Conflict,
102}
103
104/// The row a drill into a group reads, from [`DataTableState::drill_row`].
105pub enum DrillRow {
106    /// Taken from the rows on screen.
107    Buffered(DataFrame),
108    /// Not on hand: the one-row frame to collect, off the UI thread.
109    Read(Box<LazyFrame>),
110}
111
112/// A group's rows, drilled from its row of a grouped view.
113pub(super) struct GroupRows {
114    lf: LazyFrame,
115    /// The key columns of the grouped view and the row's values in them, as text.
116    key_columns: Vec<String>,
117    key_values: Vec<String>,
118    /// Columns of `lf` that hold the keys as they stand, to lead the view.
119    lead: Vec<String>,
120    /// How `lf` is built, as Copy as Python steps.
121    steps: Vec<Step>,
122    /// Which loaded column each column of `lf` is.
123    lineage: Lineage,
124}
125
126/// The rows an export writes. See [`DataTableState::export_frame`].
127pub struct ExportFrame {
128    pub(super) lf: LazyFrame,
129    pub(super) files: Option<SourceFiles>,
130}
131
132/// The dataset's files in scan order, and the row each starts at.
133pub(super) struct SourceFiles {
134    pub(super) names: Arc<Vec<String>>,
135    pub(super) starts: Arc<Vec<usize>>,
136}
137
138impl ExportFrame {
139    /// Rows that are not the view's, such as a column's value counts.
140    pub fn of(lf: LazyFrame) -> Self {
141        Self { lf, files: None }
142    }
143}
144
145impl ExportFrame {
146    /// The name of the column an export adds when asked to say where each row is from.
147    pub const SOURCE_FILE_COLUMN: &'static str = "source_file";
148
149    /// The plan. Naming files reads the schema (possibly resolving the scan), so off the UI
150    /// thread; names map from the row index per batch, so streamed exports never hold
151    /// every row.
152    pub fn into_lazy(self) -> PolarsResult<LazyFrame> {
153        let Some(SourceFiles { names, starts }) = self.files else {
154            return Ok(self.lf);
155        };
156        let mut lf = self.lf;
157        let schema = lf.collect_schema()?;
158        let name = Self::free_name(schema.iter_names().map(|n| n.as_str()));
159        let index = crate::formats::schema_union::DRIFT_COLUMN;
160        let file_of = move |rows: Column| -> PolarsResult<Column> {
161            let rows = rows.strict_cast(&DataType::UInt64)?;
162            let named: StringChunked = rows
163                .u64()?
164                .iter()
165                .map(|row| {
166                    let row = row? as usize;
167                    let file = starts
168                        .partition_point(|&start| start <= row)
169                        .saturating_sub(1);
170                    names.get(file).map(String::as_str)
171                })
172                .collect();
173            Ok(named.with_name(rows.name().clone()).into_column())
174        };
175        // Added last, after the dataset's own columns, as the index comes off.
176        Ok(lf
177            .with_column(
178                col(index)
179                    .map(file_of, |_, field| {
180                        Ok(Field::new(field.name().clone(), DataType::String))
181                    })
182                    .alias(name),
183            )
184            .drop(by_name([index], true, false)))
185    }
186
187    /// A name for the source-file column no column has: `source_file` may exist already,
188    /// and adding a same-named column would silently replace it.
189    fn free_name<'a>(taken: impl Iterator<Item = &'a str>) -> String {
190        let taken: HashSet<&str> = taken.collect();
191        std::iter::once(Self::SOURCE_FILE_COLUMN.to_string())
192            .chain((1..).map(|n| format!("{}_{n}", Self::SOURCE_FILE_COLUMN)))
193            .find(|candidate| !taken.contains(candidate.as_str()))
194            .expect("some suffix is free")
195    }
196}
197
198impl DataTableState {
199    /// The group drilled into, as its key columns and their values.
200    pub fn drilled_group_key(&self) -> Option<(&[String], &[String])> {
201        let values = self.view.drilled_down_group_key.as_deref()?;
202        let columns = self
203            .view
204            .drilled_down_group_key_columns
205            .as_deref()
206            .unwrap_or_default();
207        Some((columns, values))
208    }
209
210    /// The columns and types of `df`, the table SQL runs against, from the known schema; a
211    /// drilled group or reshape only resolves its plan, reading nothing.
212    pub fn sql_table_columns(&self) -> Vec<(String, DataType)> {
213        let schema = if self.view.grouped.is_none() && self.view.reshaped_lf.is_none() {
214            Some(self.original_schema.clone())
215        } else {
216            self.query_root().collect_schema().ok()
217        };
218        schema
219            .map(|schema| {
220                schema
221                    .iter()
222                    .filter(|(name, _)| name.as_str() != crate::formats::schema_union::DRIFT_COLUMN)
223                    .map(|(name, dtype)| (name.to_string(), dtype.clone()))
224                    .collect()
225            })
226            .unwrap_or_default()
227    }
228
229    /// Rows `df` holds, when that is known without counting: the data as loaded,
230    /// once its count has come back.
231    pub fn sql_table_rows(&self) -> Option<usize> {
232        if self.view.grouped.is_some() || self.view.reshaped_lf.is_some() {
233            return None;
234        }
235        self.pristine_rows
236    }
237
238    pub fn get_active_fuzzy_query(&self) -> &str {
239        &self.view.active_fuzzy_query
240    }
241
242    pub fn last_pivot_spec(&self) -> Option<&PivotSpec> {
243        self.view.last_pivot_spec.as_ref()
244    }
245
246    pub fn last_melt_spec(&self) -> Option<&MeltSpec> {
247        self.view.last_melt_spec.as_ref()
248    }
249
250    /// What the pivot or melt in effect ran over. See the field.
251    pub fn reshape_source(&self) -> Option<&ReshapeSource> {
252        self.view.reshape_source.as_ref()
253    }
254
255    /// Whether the view is a grouping's result (a `by` query, SQL GROUP BY) whose rows drill
256    /// into groups; recorded by the query, never inferred from list columns.
257    pub fn is_grouped(&self) -> bool {
258        self.view.group_source.is_some()
259    }
260
261    /// Whether any column holds lists.
262    fn has_list_columns(&self) -> bool {
263        self.view
264            .schema
265            .iter()
266            .any(|(_, dtype)| matches!(dtype, DataType::List(_)))
267    }
268
269    fn group_key_columns(&self) -> Vec<String> {
270        self.view
271            .schema
272            .iter()
273            .filter(|(_, dtype)| !matches!(dtype, DataType::List(_)))
274            .map(|(name, _)| name.to_string())
275            .collect()
276    }
277
278    fn group_value_columns(&self) -> Vec<String> {
279        self.view
280            .schema
281            .iter()
282            .filter(|(_, dtype)| matches!(dtype, DataType::List(_)))
283            .map(|(name, _)| name.to_string())
284            .collect()
285    }
286
287    /// Names of binary columns in the source schema. Their values are not read into the display
288    /// buffer (the `‹binary›` stub stands in); the renderer uses this to style those cells.
289    pub fn binary_column_names(&self) -> std::collections::HashSet<String> {
290        self.view
291            .schema
292            .iter()
293            .filter(|(_, dtype)| matches!(dtype, DataType::Binary))
294            .map(|(name, _)| name.to_string())
295            .collect()
296    }
297
298    /// Estimated heap size in bytes of the currently buffered slice (locked + scrollable), if collected.
299    pub fn buffered_memory_bytes(&self) -> Option<usize> {
300        let locked = self
301            .view
302            .locked_df
303            .as_ref()
304            .map(|df| df.estimated_size())
305            .unwrap_or(0);
306        let scroll = self
307            .view
308            .df
309            .as_ref()
310            .map(|df| df.estimated_size())
311            .unwrap_or(0);
312        if locked == 0 && scroll == 0 {
313            None
314        } else {
315            Some(locked + scroll)
316        }
317    }
318
319    /// The rows on hand, first to past the last.
320    pub fn buffered_span(&self) -> (usize, usize) {
321        (self.view.buffered_start_row, self.view.buffered_end_row)
322    }
323
324    /// Rows in the buffer; 0 when none is loaded.
325    pub fn buffered_rows(&self) -> usize {
326        self.view
327            .buffered_end_row
328            .saturating_sub(self.view.buffered_start_row)
329    }
330
331    /// The first `limit` distinct non-null values of `column` in the display buffer, as
332    /// text, a long string by its start. Reads nothing; a column not buffered gives none.
333    pub(crate) fn buffered_values(&self, column: &str, limit: usize) -> Vec<String> {
334        let Some(series) = [self.view.df.as_ref(), self.view.locked_df.as_ref()]
335            .into_iter()
336            .flatten()
337            .find_map(|df| df.column(column).ok())
338        else {
339            return Vec::new();
340        };
341        let series = series.as_materialized_series();
342        let mut values = Vec::new();
343        for value in (0..series.len()).filter_map(|index| series.get(index).ok()) {
344            if values.len() == limit {
345                break;
346            }
347            let text = match value {
348                AnyValue::Null => continue,
349                // Drawn each frame beside a column's name: its start, not a whole
350                // article.
351                AnyValue::String(text) => {
352                    crate::exact::prefix(text, crate::exact::CELL_PREVIEW_BYTES).to_string()
353                }
354                AnyValue::List(items) => crate::exact::list_preview(&items),
355                // Drawn on the UI thread: Polars' display panics on a date past
356                // the calendar.
357                value => {
358                    crate::exact::past_calendar_text(&value).unwrap_or_else(|| value.to_string())
359                }
360            };
361            if !values.contains(&text) {
362                values.push(text);
363            }
364        }
365        values
366    }
367
368    /// Current scrollable display buffer. None until first collect().
369    pub fn display_df(&self) -> Option<&DataFrame> {
370        self.view.df.as_ref()
371    }
372
373    /// The display buffer's rows on screen, as the table draws them.
374    pub fn display_slice_df(&self) -> Option<DataFrame> {
375        let df = self.view.df.as_ref()?;
376        let offset = self
377            .view
378            .start_row
379            .saturating_sub(self.view.buffered_start_row);
380        let slice_len = self.visible_rows.min(df.height().saturating_sub(offset));
381        if offset < df.height() && slice_len > 0 {
382            Some(df.slice(offset as i64, slice_len))
383        } else {
384            None
385        }
386    }
387
388    /// The selected row with every display column, raw and in display order:
389    /// what a row copy carries.
390    pub fn copy_row_df(&self) -> Option<DataFrame> {
391        let df = self.view.buffered_df.as_ref()?;
392        let absolute = self.view.start_row + self.table_state.selected()?;
393        let offset = absolute.checked_sub(self.view.buffered_start_row)?;
394        if offset >= df.height() {
395            return None;
396        }
397        let names: Vec<&str> = self.view.column_order.iter().map(|s| s.as_str()).collect();
398        df.select(names).ok().map(|d| d.slice(offset as i64, 1))
399    }
400
401    /// The rows on screen with every display column, raw, in display order, ignoring
402    /// column scroll (scrolled-past identifiers make the rows readable).
403    pub fn copy_view_df(&self) -> Option<DataFrame> {
404        let df = self.view.buffered_df.as_ref()?;
405        let names: Vec<&str> = self.view.column_order.iter().map(|s| s.as_str()).collect();
406        let selected = df.select(names).ok()?;
407        let offset = self
408            .view
409            .start_row
410            .saturating_sub(self.view.buffered_start_row);
411        let len = self
412            .visible_rows
413            .min(selected.height().saturating_sub(offset));
414        (len > 0).then(|| selected.slice(offset as i64, len))
415    }
416
417    /// The selected row's value in `column`, exactly as stored ([`crate::exact`]): floats
418    /// as their round-tripping decimal, null as empty (as in exports), lists and structs as
419    /// JSON.
420    pub fn copy_cell_value(&self, column: &str) -> Option<String> {
421        let row = self.copy_row_df()?;
422        crate::exact::copy_text(row.column(column).ok()?).ok()
423    }
424
425    /// The selected row's number as the row-numbers column would print it.
426    pub fn selected_display_row(&self) -> Option<usize> {
427        Some(self.view.start_row + self.table_state.selected()? + self.row_start_index)
428    }
429
430    /// Rows times estimated row width (binary at base64 size), for the copy guard. `None`
431    /// until counted, or while a binary column's width is unknown (only a footer says).
432    pub fn estimated_copy_bytes(&self) -> Option<usize> {
433        let rows = self.num_rows_if_valid()?;
434        if rows == 0 {
435            return Some(0);
436        }
437        let base64 = |bytes: usize| bytes.div_ceil(3) * 4;
438        let footer_width = |name: &str| {
439            self.column_bytes
440                .iter()
441                .find(|(n, _)| n == name)
442                .map(|(_, w)| *w)
443        };
444        let mut row = self.bytes_per_row();
445        for name in &self.view.column_order {
446            match self.view.schema.get(name.as_str()) {
447                Some(DataType::Binary) => row += base64(footer_width(name)?),
448                // Buffered whole, so the buffer measured it with the row; base64 adds
449                // a third on top.
450                Some(dtype) if crate::export::nested_json::has_binary(dtype) => {
451                    let buffered = self.view.buffered_df.as_ref().and_then(|df| {
452                        let column = df.column(name).ok()?;
453                        (df.height() > 0)
454                            .then(|| column.as_materialized_series().estimated_size() / df.height())
455                    });
456                    row += buffered.or_else(|| footer_width(name)).unwrap_or(0) / 3;
457                }
458                _ => {}
459            }
460        }
461        Some(rows.saturating_mul(row))
462    }
463
464    /// Each row's drift group from the top of the view, for a frame `frame_rows` tall (the
465    /// whole frame, header included: extra entries are ignored, and `visible_rows` is 0
466    /// before the first frame). Empty once a query or reshape replaced the frame (its rows
467    /// stand for no file).
468    pub fn display_drift(&self, frame_rows: usize) -> Vec<u32> {
469        if !self.view.drift_column_present {
470            return Vec::new();
471        }
472        let Some(df) = self.view.buffered_df.as_ref() else {
473            return Vec::new();
474        };
475        let Ok(column) = df.column(crate::formats::schema_union::DRIFT_COLUMN) else {
476            return Vec::new();
477        };
478        let offset = self
479            .view
480            .start_row
481            .saturating_sub(self.view.buffered_start_row);
482        let len = frame_rows.min(column.len().saturating_sub(offset));
483        if len == 0 {
484            return Vec::new();
485        }
486        let slice = column.slice(offset as i64, len);
487        let Ok(rows) = slice.u32() else {
488            return Vec::new();
489        };
490        // The column holds each row's place in the dataset. The file it came from is
491        // the last one starting at or before it, and the file says what it is missing.
492        rows.iter()
493            .map(|row| self.file_group_of(row.unwrap_or(0) as usize))
494            .collect()
495    }
496
497    /// Maximum buffer size in rows (0 = no limit).
498    pub fn max_buffered_rows(&self) -> usize {
499        self.max_buffered_rows
500    }
501
502    /// Maximum buffer size in MiB (0 = no limit).
503    pub fn max_buffered_mb(&self) -> usize {
504        self.max_buffered_mb
505    }
506
507    /// Whether Enter on a row drills into its group: the view is a grouped result and
508    /// not already a group's rows.
509    pub fn can_drill_down(&self) -> bool {
510        !self.is_drilled_down() && self.is_grouped()
511    }
512
513    /// Whether a drill shows the lists of a row as the group's rows rather than
514    /// filtering the source: a `by` result that holds its groups as lists.
515    fn drills_lists(&self) -> bool {
516        self.has_list_columns()
517            && self
518                .view
519                .group_source
520                .as_ref()
521                .is_some_and(|s| s.rows_in_lists)
522    }
523
524    /// The columns of a row that a drill into its group reads: every column of a result
525    /// holding its groups as lists, the keys of one holding aggregates.
526    fn drill_columns(&self) -> Vec<String> {
527        if self.drills_lists() {
528            return self
529                .view
530                .schema
531                .iter_names()
532                .map(|n| n.to_string())
533                .collect();
534        }
535        self.view
536            .group_source
537            .iter()
538            .flat_map(|source| source.keys.iter().map(|(name, _)| name.to_string()))
539            .collect()
540    }
541
542    /// The columns the inspector lists for a row: the table's, in its order, then
543    /// the ones hidden from it, in schema order. Never the scan's own row index.
544    pub fn inspect_fields(&self) -> Vec<InspectField> {
545        let shown = self.view.column_order.iter().filter_map(|name| {
546            Some(InspectField {
547                name: name.clone(),
548                dtype: self.view.schema.get(name.as_str())?.clone(),
549                hidden: false,
550            })
551        });
552        let hidden = self
553            .view
554            .schema
555            .iter()
556            .filter(|(name, _)| {
557                name.as_str() != crate::formats::schema_union::DRIFT_COLUMN
558                    && !self.view.column_order.iter().any(|c| c == name.as_str())
559            })
560            .map(|(name, dtype)| InspectField {
561                name: name.to_string(),
562                dtype: dtype.clone(),
563                hidden: true,
564            });
565        shown.chain(hidden).collect()
566    }
567
568    /// The selected row as the buffer holds it, with the file group that says what
569    /// its nulls are. Reads nothing; `None` while the row is not on hand.
570    pub fn inspect_row(&self) -> Option<InspectRow> {
571        self.inspect_row_at(self.view.start_row + self.table_state.selected()?)
572    }
573
574    /// Row `row` of the view as the buffer holds it, as [`Self::inspect_row`] does
575    /// the selected one: Compare's next row. `None` while it is not on hand.
576    pub fn inspect_row_at(&self, row: usize) -> Option<InspectRow> {
577        let df = self.view.buffered_df.as_ref()?;
578        let offset = row.checked_sub(self.view.buffered_start_row)?;
579        if offset >= df.height() {
580            return None;
581        }
582        let names: Vec<&str> = self.view.column_order.iter().map(|s| s.as_str()).collect();
583        let values = df.select(names).ok()?.slice(offset as i64, 1);
584        let drift_group = self
585            .view
586            .drift_column_present
587            .then(|| df.column(crate::formats::schema_union::DRIFT_COLUMN).ok())
588            .flatten()
589            .and_then(|c| c.get(offset).ok())
590            .and_then(|v| v.extract::<usize>())
591            .map(|place| self.file_group_of(place));
592        Some(InspectRow {
593            row,
594            frame: self.view.len_generation,
595            display_row: row + self.row_start_index,
596            values,
597            drift_group,
598        })
599    }
600
601    /// The drift group of the file holding the dataset's row `place`.
602    fn file_group_of(&self, place: usize) -> u32 {
603        let file = self
604            .drift_file_starts
605            .partition_point(|&start| start <= place)
606            .saturating_sub(1);
607        self.drift_file_group.get(file).copied().unwrap_or(0)
608    }
609
610    /// What a null in `column` is, for a row of file group `group`: the data's own,
611    /// a file without the column, or a file holding it in another type.
612    pub fn null_kind(&self, column: &str, group: Option<u32>) -> NullKind {
613        let Some(group) = group.and_then(|g| self.view.drift_groups.get(g as usize)) else {
614            return NullKind::Null;
615        };
616        if group.absent.iter().any(|c| c == column) {
617            NullKind::Absent
618        } else if group.unread.iter().any(|c| c == column) {
619            NullKind::Conflict
620        } else {
621            NullKind::Null
622        }
623    }
624
625    /// The frame reading `columns` of view row `row` for the inspector, through the
626    /// buffer's window (a remote dataset reads only that file). Run off this thread.
627    pub fn inspect_read_lf(&self, row: usize, columns: &[String]) -> PolarsResult<LazyFrame> {
628        let exprs = columns.iter().map(|c| col(c.as_str())).collect();
629        self.window_lf(row, 1, exprs)
630    }
631
632    /// What drilling into the group on view row `group_index` reads, or `None` when not
633    /// grouped. The buffer holds the row; a missing column (hidden, binary stub) means
634    /// reading the row, which for an aggregate is the whole aggregate, so off the UI thread.
635    pub fn drill_row(&self, group_index: usize) -> Option<DrillRow> {
636        if !self.can_drill_down() {
637            return None;
638        }
639        let columns = self.drill_columns();
640        let buffered = self
641            .view
642            .buffered_df
643            .as_ref()
644            .filter(|_| {
645                (self.view.buffered_start_row..self.view.buffered_end_row).contains(&group_index)
646            })
647            .filter(|_| {
648                columns
649                    .iter()
650                    .all(|c| !matches!(self.view.schema.get(c.as_str()), Some(DataType::Binary)))
651            })
652            .and_then(|df| df.select(columns.iter().map(|c| c.as_str())).ok())
653            .map(|df| df.slice((group_index - self.view.buffered_start_row) as i64, 1))
654            .filter(|row| row.height() == 1);
655        Some(match buffered {
656            Some(row) => DrillRow::Buffered(row),
657            None => DrillRow::Read(Box::new(
658                self.visible_lf()
659                    .select(columns.iter().map(|c| col(c.as_str())).collect::<Vec<_>>())
660                    .slice(group_index as i64, 1),
661            )),
662        })
663    }
664
665    /// Show the group on view row `group_index`, reading the row on this thread if not
666    /// buffered; the app uses [`Self::drill_row`] so keys never wait.
667    pub fn drill_down_into_group(&mut self, group_index: usize) -> Result<()> {
668        let row = match self.drill_row(group_index) {
669            None => return Ok(()),
670            Some(DrillRow::Buffered(row)) => row,
671            Some(DrillRow::Read(lf)) => collect_lazy(*lf, self.polars_streaming)?,
672        };
673        self.drill_down_with_row(group_index, &row)
674    }
675
676    /// Show the group whose row `group_index` is `row` (from [`Self::drill_row`]): list
677    /// columns as rows, or for aggregates the source rows sharing the group's keys, key
678    /// columns first.
679    pub fn drill_down_with_row(&mut self, group_index: usize, row: &DataFrame) -> Result<()> {
680        if !self.can_drill_down() {
681            return Ok(());
682        }
683        if row.height() == 0 {
684            return Err(color_eyre::eyre::eyre!("Group index out of bounds"));
685        }
686        let mut group = if self.drills_lists() {
687            Self::group_from_lists(
688                row,
689                self.group_key_columns(),
690                self.group_value_columns(),
691                self.view.lineage.clone(),
692            )?
693        } else if let Some(source) = &self.view.group_source {
694            Self::group_from_source(source, row)?
695        } else {
696            return Ok(());
697        };
698        // A list form that also aggregates (`select a, n: count a by k`) holds its
699        // aggregates beside the keys; the query knows which columns are keys.
700        if let Some(source) = self
701            .view
702            .group_source
703            .as_ref()
704            .filter(|_| self.drills_lists())
705        {
706            let keys: Vec<&str> = source.keys.iter().map(|(n, _)| n.as_str()).collect();
707            (group.key_columns, group.key_values) = group
708                .key_columns
709                .into_iter()
710                .zip(group.key_values)
711                .filter(|(name, _)| keys.contains(&name.as_str()))
712                .unzip();
713        }
714        self.enter_group(group, group_index, false)
715    }
716
717    /// Show the view's rows with `value` in `column` (null matches null), like a group
718    /// drill: breadcrumb names the value, Esc returns. Inside a group, it narrows that group
719    /// and Esc returns to the view it was drilled from.
720    pub fn drill_into_value(&mut self, column: &str, value: AnyValue<'static>) -> Result<()> {
721        let dtype = self
722            .view
723            .schema
724            .get(column)
725            .cloned()
726            .ok_or_else(|| color_eyre::eyre::eyre!("no column {column}"))?;
727        let label = crate::exact::str_value(&value).to_string();
728        let mut steps = self.view_steps();
729        steps.push(match crate::export::python_script::py_value(&value) {
730            Some(literal) => Step::Matching(vec![(
731                format!("pl.col({})", crate::export::python_script::py_str(column)),
732                literal,
733            )]),
734            None => Step::Unreproducible(format!(
735                "drilled down to the rows where {column} is {label}, a value of a type not written as Python"
736            )),
737        });
738        let matches = col(column).eq_missing(lit(Scalar::new(dtype, value)));
739        let group = GroupRows {
740            lf: self.visible_lf().filter(matches),
741            key_columns: vec![column.to_string()],
742            key_values: vec![label],
743            lead: vec![column.to_string()],
744            steps,
745            lineage: self.view.lineage.clone(),
746        };
747        if !self.is_drilled_down() {
748            let index = self.view.start_row + self.table_state.selected().unwrap_or(0);
749            return self.enter_group(group, index, true);
750        }
751        let schema = group.lf.clone().collect_schema()?;
752        let order = std::mem::take(&mut self.view.column_order);
753        if let Some(keys) = self.view.drilled_down_group_key_columns.as_mut() {
754            keys.extend(group.key_columns);
755        }
756        if let Some(values) = self.view.drilled_down_group_key.as_mut() {
757            values.extend(group.key_values);
758        }
759        // The group's filters and sort are in the frame now.
760        self.view.filters.clear();
761        self.view.sort_columns.clear();
762        self.view.sort_descending.clear();
763        self.view.sort_ascending = true;
764        self.install_base(group.lf, schema);
765        self.view.base_steps = group.steps;
766        self.view.lineage = group.lineage;
767        self.view.column_order = order;
768        self.view.start_row = 0;
769        self.termcol_index = 0;
770        self.clear_column_moves();
771        self.settle_cursor();
772        self.table_state.select(Some(0));
773        self.collect();
774        Ok(())
775    }
776
777    /// Whether the drill on screen came from Value Counts.
778    pub fn drilled_into_value(&self) -> bool {
779        self.view.grouped.as_ref().is_some_and(|view| view.by_value)
780    }
781
782    /// Show `group`, the group on row `group_index` of the view, keeping the view to
783    /// come back to.
784    fn enter_group(&mut self, group: GroupRows, group_index: usize, by_value: bool) -> Result<()> {
785        let schema = group.lf.clone().collect_schema()?;
786        self.view.drilled_down_group_key = Some(group.key_values);
787        self.view.drilled_down_group_key_columns = Some(group.key_columns);
788
789        // The group becomes the pipeline root while drilled in, so a sidebar filter or
790        // sort applies within it instead of rebuilding the grouped view underneath.
791        self.view.grouped = Some(GroupedView {
792            lf: self.view.lf.clone(),
793            base_lf: self.view.base_lf.clone(),
794            base_schema: self.view.base_schema.clone(),
795            schema: self.view.schema.clone(),
796            filters: std::mem::take(&mut self.view.filters),
797            sort_columns: std::mem::take(&mut self.view.sort_columns),
798            sort_descending: std::mem::take(&mut self.view.sort_descending),
799            sort_ascending: self.view.sort_ascending,
800            drift: self.view.drift_column_present,
801            drift_groups: self.view.drift_groups.clone(),
802            view_numbered: self.view.view_numbered,
803            notes: self.view.notes.clone(),
804            group_source: self.view.group_source.take(),
805            column_order: self.view.column_order.clone(),
806            locked_columns_count: self.view.locked_columns_count,
807            start_row: self.view.start_row,
808            termcol_index: self.termcol_index,
809            cursor_column: self.view.cursor_column.clone(),
810            selected: self.table_state.selected(),
811            by_value,
812            base_steps: std::mem::take(&mut self.view.base_steps),
813            lineage: self.view.lineage.clone(),
814        });
815        self.view.sort_ascending = true;
816        self.install_base(group.lf, schema);
817        self.view.base_steps = group.steps;
818        self.view.lineage = group.lineage;
819        // Led by the keys, as a group drilled from lists is.
820        let rest: Vec<String> = std::mem::take(&mut self.view.column_order)
821            .into_iter()
822            .filter(|c| !group.lead.contains(c))
823            .collect();
824        self.view.column_order = group.lead.into_iter().chain(rest).collect();
825        self.view.drilled_down_group_index = Some(group_index);
826        self.view.start_row = 0;
827        self.termcol_index = 0;
828        self.clear_column_moves();
829        self.view.locked_columns_count = 0;
830        self.settle_cursor();
831        self.table_state.select(Some(0));
832        self.collect();
833
834        Ok(())
835    }
836
837    /// A group held as lists in `row`: the keys repeated beside the lists as columns.
838    fn group_from_lists(
839        row: &DataFrame,
840        key_columns: Vec<String>,
841        value_columns: Vec<String>,
842        lineage: Lineage,
843    ) -> Result<GroupRows> {
844        if value_columns.is_empty() {
845            return Err(color_eyre::eyre::eyre!("No value columns in grouped data"));
846        }
847        let row_count = match row.column(&value_columns[0])?.get(0)? {
848            AnyValue::List(list_series) => list_series.len(),
849            _ => 0,
850        };
851
852        let mut columns = Vec::new();
853        let mut key_values = Vec::new();
854        for col_name in &key_columns {
855            let key = row.column(col_name)?;
856            key_values.push(crate::exact::str_value(&key.get(0)?).to_string());
857            // Repeated in its own type, so a date stays a date and a null key null.
858            columns.push(key.new_from_index(0, row_count));
859        }
860        for col_name in &value_columns {
861            if let AnyValue::List(list_series) = row.column(col_name)?.get(0)? {
862                columns.push(list_series.with_name(col_name.as_str().into()).into());
863            }
864        }
865        let group = key_columns
866            .iter()
867            .zip(&key_values)
868            .map(|(c, v)| format!("{c} = {v}"))
869            .collect::<Vec<_>>()
870            .join(", ");
871        Ok(GroupRows {
872            lf: DataFrame::new_infer_height(columns)?.lazy(),
873            key_columns,
874            key_values,
875            // Already first.
876            lead: Vec::new(),
877            steps: vec![Step::Unreproducible(format!(
878                "drilled down into the group {group}, read from the grouped result's lists: \
879                 not written as Python"
880            ))],
881            // The lists keep the result's names.
882            lineage,
883        })
884    }
885
886    /// A group of an aggregated result in `row`: the source rows whose keys equal the
887    /// row's, a null key matching nulls.
888    fn group_from_source(source: &GroupSource, row: &DataFrame) -> Result<GroupRows> {
889        let mut predicate: Option<Expr> = None;
890        let mut key_columns = Vec::new();
891        let mut key_values = Vec::new();
892        let mut lead = Vec::new();
893        let mut matching = Vec::new();
894        for (i, (name, expr)) in source.keys.iter().enumerate() {
895            let column = row.column(name)?;
896            let value = column.get(0)?.into_static();
897            matching.push(
898                source
899                    .python_keys
900                    .get(i)
901                    .cloned()
902                    .flatten()
903                    .zip(crate::export::python_script::py_value(&value)),
904            );
905            key_columns.push(name.to_string());
906            key_values.push(crate::exact::str_value(&value).to_string());
907            // The key as the query computed it, against the value it produced; the alias
908            // only named the result's column.
909            let key = expr.clone().meta().undo_aliases();
910            if let Expr::Column(source_column) = &key
911                && !source.scratch.contains(source_column)
912            {
913                lead.push(source_column.to_string());
914            }
915            let matches = key.eq_missing(lit(Scalar::new(column.dtype().clone(), value)));
916            predicate = Some(match predicate {
917                Some(all) => all.and(matches),
918                None => matches,
919            });
920        }
921        let rows = source.rows.clone();
922        let mut lf = match predicate {
923            Some(predicate) => rows.filter(predicate),
924            None => rows,
925        };
926        if !source.scratch.is_empty() {
927            lf = lf.drop(by_name(source.scratch.iter().cloned(), true, false));
928        }
929        let matching: Option<Vec<(String, String)>> = matching.into_iter().collect();
930        let steps = match (&source.python_rows, matching) {
931            (Some(rows), Some(matching)) => {
932                let mut steps = rows.clone();
933                steps.push(Step::Matching(matching));
934                if !source.scratch.is_empty() {
935                    steps.push(Step::Drop(
936                        source.scratch.iter().map(|c| c.to_string()).collect(),
937                    ));
938                }
939                steps
940            }
941            _ => vec![Step::Unreproducible(format!(
942                "drilled down into the group {}: not written as Python",
943                key_columns
944                    .iter()
945                    .zip(&key_values)
946                    .map(|(c, v)| format!("{c} = {v}"))
947                    .collect::<Vec<_>>()
948                    .join(", ")
949            ))],
950        };
951        Ok(GroupRows {
952            lf,
953            key_columns,
954            key_values,
955            lead,
956            steps,
957            lineage: source.lineage.clone(),
958        })
959    }
960
961    pub fn drill_up(&mut self) -> Result<()> {
962        let Some(view) = self.view.grouped.take() else {
963            return Err(color_eyre::eyre::eyre!("Not in drill-down mode"));
964        };
965        self.invalidate_num_rows();
966        // The buffer holds the group's rows; kept, it would stand in for the grouped
967        // view wherever the view fits inside it.
968        self.drop_buffer();
969        self.view.observed_bytes_per_row = None;
970        self.widths.relearn();
971        self.view.lf = view.lf;
972        self.view.unsorted_lf = None;
973        self.view.base_lf = view.base_lf;
974        self.view.base_schema = view.base_schema;
975        self.view.base_steps = view.base_steps;
976        self.view.lineage = view.lineage;
977        self.view.filters = view.filters;
978        self.view.sort_columns = view.sort_columns;
979        self.view.sort_descending = view.sort_descending;
980        self.view.sort_ascending = view.sort_ascending;
981        self.view.drift_column_present = view.drift;
982        self.view.drift_groups = view.drift_groups;
983        self.view.view_numbered = view.view_numbered;
984        self.view.notes = view.notes;
985        self.view.group_source = view.group_source;
986        // The restored frame already excludes what its filter and sort excluded; rederive
987        // those notes so they cannot be stale.
988        self.view.view_notes = self.view_notes_only();
989        self.view.schema = view.schema;
990        self.view.column_order = view.column_order;
991        self.view.locked_columns_count = view.locked_columns_count;
992        self.view.drilled_down_group_index = None;
993        self.view.drilled_down_group_key = None;
994        self.view.drilled_down_group_key_columns = None;
995        self.view.start_row = view.start_row;
996        self.termcol_index = view.termcol_index;
997        self.clear_column_moves();
998        self.view.cursor_column = view.cursor_column;
999        self.settle_cursor();
1000        self.table_state.select(view.selected);
1001        self.collect();
1002        Ok(())
1003    }
1004}