1use super::*;
4
5#[derive(Clone)]
8pub(super) struct GroupedView {
9 pub(super) lf: LazyFrame,
10 pub(super) base_lf: LazyFrame,
11 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 drift: bool,
21 drift_groups: Arc<Vec<crate::formats::schema_union::DriftGroup>>,
22 view_numbered: bool,
24 notes: Vec<crate::notes::Note>,
25 pub(super) group_source: Option<GroupSource>,
26 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 by_value: bool,
36 base_steps: Vec<Step>,
38 lineage: Lineage,
39}
40
41#[derive(Clone)]
45pub(super) struct GroupSource {
46 pub(super) rows: LazyFrame,
48 pub(super) keys: Vec<(PlSmallStr, Expr)>,
50 pub(super) scratch: Vec<PlSmallStr>,
52 pub(super) rows_in_lists: bool,
55 pub(super) python_rows: Option<Vec<Step>>,
58 pub(super) python_keys: Vec<Option<String>>,
59 pub(super) lineage: Lineage,
61}
62
63#[derive(Debug, Clone, PartialEq)]
65pub struct InspectField {
66 pub name: String,
67 pub dtype: DataType,
68 pub hidden: bool,
70}
71
72impl InspectField {
73 pub fn buffered(&self) -> bool {
76 !self.hidden && !matches!(self.dtype, DataType::Binary)
77 }
78}
79
80#[derive(Clone)]
82pub struct InspectRow {
83 pub row: usize,
85 pub frame: u64,
88 pub display_row: usize,
90 pub values: DataFrame,
92 pub drift_group: Option<u32>,
94}
95
96#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98pub enum NullKind {
99 Null,
100 Absent,
101 Conflict,
102}
103
104pub enum DrillRow {
106 Buffered(DataFrame),
108 Read(Box<LazyFrame>),
110}
111
112pub(super) struct GroupRows {
114 lf: LazyFrame,
115 key_columns: Vec<String>,
117 key_values: Vec<String>,
118 lead: Vec<String>,
120 steps: Vec<Step>,
122 lineage: Lineage,
124}
125
126pub struct ExportFrame {
128 pub(super) lf: LazyFrame,
129 pub(super) files: Option<SourceFiles>,
130}
131
132pub(super) struct SourceFiles {
134 pub(super) names: Arc<Vec<String>>,
135 pub(super) starts: Arc<Vec<usize>>,
136}
137
138impl ExportFrame {
139 pub fn of(lf: LazyFrame) -> Self {
141 Self { lf, files: None }
142 }
143}
144
145impl ExportFrame {
146 pub const SOURCE_FILE_COLUMN: &'static str = "source_file";
148
149 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 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 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 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 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 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 pub fn reshape_source(&self) -> Option<&ReshapeSource> {
252 self.view.reshape_source.as_ref()
253 }
254
255 pub fn is_grouped(&self) -> bool {
258 self.view.group_source.is_some()
259 }
260
261 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 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 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 pub fn buffered_span(&self) -> (usize, usize) {
321 (self.view.buffered_start_row, self.view.buffered_end_row)
322 }
323
324 pub fn buffered_rows(&self) -> usize {
326 self.view
327 .buffered_end_row
328 .saturating_sub(self.view.buffered_start_row)
329 }
330
331 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 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 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 pub fn display_df(&self) -> Option<&DataFrame> {
370 self.view.df.as_ref()
371 }
372
373 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 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 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 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 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 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 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 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 rows.iter()
493 .map(|row| self.file_group_of(row.unwrap_or(0) as usize))
494 .collect()
495 }
496
497 pub fn max_buffered_rows(&self) -> usize {
499 self.max_buffered_rows
500 }
501
502 pub fn max_buffered_mb(&self) -> usize {
504 self.max_buffered_mb
505 }
506
507 pub fn can_drill_down(&self) -> bool {
510 !self.is_drilled_down() && self.is_grouped()
511 }
512
513 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 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 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 pub fn inspect_row(&self) -> Option<InspectRow> {
571 self.inspect_row_at(self.view.start_row + self.table_state.selected()?)
572 }
573
574 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 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 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 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 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 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 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 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 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 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 pub fn drilled_into_value(&self) -> bool {
779 self.view.grouped.as_ref().is_some_and(|view| view.by_value)
780 }
781
782 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 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 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 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 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 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 lineage,
883 })
884 }
885
886 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 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 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 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}