use super::*;
#[derive(Clone)]
pub(super) struct GroupedView {
pub(super) lf: LazyFrame,
pub(super) base_lf: LazyFrame,
base_schema: Arc<Schema>,
schema: Arc<Schema>,
pub(super) filters: Vec<FilterStatement>,
pub(super) sort_columns: Vec<String>,
pub(super) sort_descending: Vec<bool>,
pub(super) sort_ascending: bool,
drift: bool,
drift_groups: Arc<Vec<crate::formats::schema_union::DriftGroup>>,
view_numbered: bool,
notes: Vec<crate::notes::Note>,
pub(super) group_source: Option<GroupSource>,
column_order: Vec<String>,
locked_columns_count: usize,
start_row: usize,
termcol_index: usize,
cursor_column: Option<String>,
selected: Option<usize>,
by_value: bool,
steps: Vec<DataFrame>,
base_steps: Vec<Step>,
lineage: Lineage,
}
#[derive(Clone)]
pub(super) struct GroupSource {
pub(super) rows: LazyFrame,
pub(super) keys: Vec<(PlSmallStr, Expr)>,
pub(super) scratch: Vec<PlSmallStr>,
pub(super) rows_in_lists: bool,
pub(super) python_rows: Option<Vec<Step>>,
pub(super) python_keys: Vec<Option<String>>,
pub(super) lineage: Lineage,
}
#[derive(Debug, Clone, PartialEq)]
pub struct InspectField {
pub name: String,
pub dtype: DataType,
pub hidden: bool,
}
impl InspectField {
pub fn buffered(&self) -> bool {
!self.hidden && !matches!(self.dtype, DataType::Binary)
}
}
#[derive(Clone)]
pub struct InspectRow {
pub row: usize,
pub frame: u64,
pub display_row: usize,
pub values: DataFrame,
pub drift_group: Option<u32>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NullKind {
Null,
Absent,
Conflict,
}
#[derive(Debug, Clone)]
pub struct DrillPlace {
by_group: bool,
steps: Vec<DataFrame>,
filters: Vec<FilterStatement>,
sort_columns: Vec<String>,
sort_descending: Vec<bool>,
column_order: Vec<String>,
locked_columns_count: usize,
}
impl DrillPlace {
pub fn by_group(&self) -> bool {
self.by_group
}
pub fn describe(&self) -> String {
self.steps.first().map(describe).unwrap_or_default()
}
}
fn describe(keys: &DataFrame) -> String {
keys.columns()
.iter()
.map(|c| {
let value = c.get(0).map(|v| crate::exact::str_value(&v).to_string());
format!("{} = {}", c.name(), value.unwrap_or_default())
})
.collect::<Vec<_>>()
.join(", ")
}
#[derive(Debug)]
pub struct DrillGone(pub String);
impl std::fmt::Display for DrillGone {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "No rows with {} in the current files", self.0)
}
}
impl std::error::Error for DrillGone {}
const GROUP_ROW: &str = "__datui_group_row";
pub enum DrillRow {
Buffered(DataFrame),
Read(Box<LazyFrame>),
}
pub(super) struct GroupRows {
lf: LazyFrame,
key_columns: Vec<String>,
key_values: Vec<String>,
lead: Vec<String>,
steps: Vec<Step>,
lineage: Lineage,
}
pub struct ExportFrame {
pub(super) lf: LazyFrame,
pub(super) files: Option<SourceFiles>,
}
pub(super) struct SourceFiles {
pub(super) names: Arc<Vec<String>>,
pub(super) starts: Arc<Vec<usize>>,
}
impl ExportFrame {
pub fn of(lf: LazyFrame) -> Self {
Self { lf, files: None }
}
}
impl ExportFrame {
pub const SOURCE_FILE_COLUMN: &'static str = "source_file";
pub fn into_lazy(self) -> PolarsResult<LazyFrame> {
let Some(SourceFiles { names, starts }) = self.files else {
return Ok(self.lf);
};
let mut lf = self.lf;
let schema = lf.collect_schema()?;
let name = Self::free_name(schema.iter_names().map(|n| n.as_str()));
let index = crate::formats::schema_union::DRIFT_COLUMN;
let file_of = move |rows: Column| -> PolarsResult<Column> {
let rows = rows.strict_cast(&DataType::UInt64)?;
let named: StringChunked = rows
.u64()?
.iter()
.map(|row| {
let row = row? as usize;
let file = starts
.partition_point(|&start| start <= row)
.saturating_sub(1);
names.get(file).map(String::as_str)
})
.collect();
Ok(named.with_name(rows.name().clone()).into_column())
};
Ok(lf
.with_column(
col(index)
.map(file_of, |_, field| {
Ok(Field::new(field.name().clone(), DataType::String))
})
.alias(name),
)
.drop(by_name([index], true, false)))
}
fn free_name<'a>(taken: impl Iterator<Item = &'a str>) -> String {
let taken: HashSet<&str> = taken.collect();
std::iter::once(Self::SOURCE_FILE_COLUMN.to_string())
.chain((1..).map(|n| format!("{}_{n}", Self::SOURCE_FILE_COLUMN)))
.find(|candidate| !taken.contains(candidate.as_str()))
.expect("some suffix is free")
}
}
impl DataTableState {
pub fn drilled_group_key(&self) -> Option<(&[String], &[String])> {
let values = self.view.drilled_down_group_key.as_deref()?;
let columns = self
.view
.drilled_down_group_key_columns
.as_deref()
.unwrap_or_default();
Some((columns, values))
}
pub fn drill_place(&self) -> Option<DrillPlace> {
let grouped = self.view.grouped.as_ref()?;
Some(DrillPlace {
by_group: !grouped.by_value,
steps: grouped.steps.clone(),
filters: self.view.filters.clone(),
sort_columns: self.view.sort_columns.clone(),
sort_descending: self.view.sort_descending.clone(),
column_order: self.view.column_order.clone(),
locked_columns_count: self.view.locked_columns_count,
})
}
pub fn grouped_column_order(&self) -> Option<(&[String], usize)> {
let grouped = self.view.grouped.as_ref()?;
Some((&grouped.column_order, grouped.locked_columns_count))
}
pub fn find_group(&self, place: &DrillPlace) -> Option<Result<LazyFrame>> {
let keys = place.steps.first().filter(|_| place.by_group)?;
if !self.can_drill_down() {
return Some(Err(color_eyre::eyre::eyre!("the view is not grouped")));
}
let mut matches = lit(true);
for key in keys.columns() {
let value = key.get(0).ok().map(|v| v.into_static());
let Some((dtype, value)) =
value.and_then(|v| self.value_in_column(key.name().as_str(), v))
else {
return Some(Err(DrillGone(place.describe()).into()));
};
let scalar = Scalar::new(dtype, value);
matches = matches.and(col(key.name().clone()).eq_missing(lit(scalar)));
}
let columns = std::iter::once(GROUP_ROW.to_string()).chain(self.drill_columns());
Some(Ok(self
.visible_lf()
.with_row_index(GROUP_ROW, None)
.filter(matches)
.select(columns.map(|c| col(c.as_str())).collect::<Vec<_>>())
.slice(0, 1)))
}
fn value_in_column(
&self,
column: &str,
value: AnyValue<'static>,
) -> Option<(DataType, AnyValue<'static>)> {
let dtype = self.view.schema.get(column)?.clone();
if value.dtype() == dtype || value.is_null() {
return Some((dtype, value));
}
let cast = Series::from_any_values(PlSmallStr::EMPTY, &[value], true)
.and_then(|s| s.strict_cast(&dtype))
.ok()?;
let value = cast.get(0).ok()?.into_static();
Some((dtype, value))
}
pub fn found_group(read: DataFrame) -> Result<Option<(usize, DataFrame)>> {
if read.height() == 0 {
return Ok(None);
}
let index = read
.column(GROUP_ROW)?
.get(0)?
.extract::<usize>()
.ok_or_else(|| color_eyre::eyre::eyre!("no row position"))?;
Ok(Some((index, read.drop(GROUP_ROW)?)))
}
pub fn redrill(&mut self, place: &DrillPlace, found: Option<(usize, DataFrame)>) -> Result<()> {
let mut steps = place.steps.iter();
if place.by_group {
steps.next();
let Some((index, row)) = found else {
return Ok(());
};
self.drill_down_with_row(index, &row)?;
}
for keys in steps {
let Some(key) = keys.columns().first() else {
continue;
};
let column = key.name().as_str();
let Some((_, value)) = self.value_in_column(column, key.get(0)?.into_static()) else {
return Err(DrillGone(describe(keys)).into());
};
self.drill_into_value(column, value)?;
}
if !place.filters.is_empty() {
self.filter(place.filters.clone());
}
if !place.sort_columns.is_empty() {
self.sort_by(place.sort_columns.clone(), place.sort_descending.clone());
}
if let Some(error) = self.error().cloned() {
return Err(color_eyre::eyre::eyre!("{error}"));
}
let shown: HashSet<&str> = self.view.schema.iter_names().map(|n| n.as_str()).collect();
if place
.column_order
.iter()
.all(|c| shown.contains(c.as_str()))
{
self.set_column_order(place.column_order.clone());
self.set_locked_columns(place.locked_columns_count);
}
Ok(())
}
pub fn sql_table_columns(&self) -> Vec<(String, DataType)> {
let schema = if self.view.grouped.is_none() && self.view.reshaped_lf.is_none() {
Some(self.original_schema.clone())
} else {
self.query_root().collect_schema().ok()
};
schema
.map(|schema| {
schema
.iter()
.filter(|(name, _)| name.as_str() != crate::formats::schema_union::DRIFT_COLUMN)
.map(|(name, dtype)| (name.to_string(), dtype.clone()))
.collect()
})
.unwrap_or_default()
}
pub fn sql_table_rows(&self) -> Option<usize> {
if self.view.grouped.is_some() || self.view.reshaped_lf.is_some() {
return None;
}
self.pristine_rows
}
pub fn get_active_fuzzy_query(&self) -> &str {
&self.view.active_fuzzy_query
}
pub fn last_pivot_spec(&self) -> Option<&PivotSpec> {
self.view.last_pivot_spec.as_ref()
}
pub fn last_melt_spec(&self) -> Option<&MeltSpec> {
self.view.last_melt_spec.as_ref()
}
pub fn reshape_source(&self) -> Option<&ReshapeSource> {
self.view.reshape_source.as_ref()
}
pub fn is_grouped(&self) -> bool {
self.view.group_source.is_some()
}
fn has_list_columns(&self) -> bool {
self.view
.schema
.iter()
.any(|(_, dtype)| matches!(dtype, DataType::List(_)))
}
fn group_key_columns(&self) -> Vec<String> {
self.view
.schema
.iter()
.filter(|(_, dtype)| !matches!(dtype, DataType::List(_)))
.map(|(name, _)| name.to_string())
.collect()
}
fn group_value_columns(&self) -> Vec<String> {
self.view
.schema
.iter()
.filter(|(_, dtype)| matches!(dtype, DataType::List(_)))
.map(|(name, _)| name.to_string())
.collect()
}
pub fn binary_column_names(&self) -> std::collections::HashSet<String> {
self.view
.schema
.iter()
.filter(|(_, dtype)| matches!(dtype, DataType::Binary))
.map(|(name, _)| name.to_string())
.collect()
}
pub fn buffered_memory_bytes(&self) -> Option<usize> {
let locked = self
.view
.locked_df
.as_ref()
.map(|df| df.estimated_size())
.unwrap_or(0);
let scroll = self
.view
.df
.as_ref()
.map(|df| df.estimated_size())
.unwrap_or(0);
if locked == 0 && scroll == 0 {
None
} else {
Some(locked + scroll)
}
}
pub fn buffered_span(&self) -> (usize, usize) {
(self.view.buffered_start_row, self.view.buffered_end_row)
}
pub fn buffered_rows(&self) -> usize {
self.view
.buffered_end_row
.saturating_sub(self.view.buffered_start_row)
}
pub(crate) fn buffered_values(&self, column: &str, limit: usize) -> Vec<String> {
let Some(series) = [self.view.df.as_ref(), self.view.locked_df.as_ref()]
.into_iter()
.flatten()
.find_map(|df| df.column(column).ok())
else {
return Vec::new();
};
let series = series.as_materialized_series();
let mut values = Vec::new();
for value in (0..series.len()).filter_map(|index| series.get(index).ok()) {
if values.len() == limit {
break;
}
let text = match value {
AnyValue::Null => continue,
AnyValue::String(text) => {
crate::exact::prefix(text, crate::exact::CELL_PREVIEW_BYTES).to_string()
}
AnyValue::List(items) => crate::exact::list_preview(&items),
value => {
crate::exact::past_calendar_text(&value).unwrap_or_else(|| value.to_string())
}
};
if !values.contains(&text) {
values.push(text);
}
}
values
}
pub fn display_df(&self) -> Option<&DataFrame> {
self.view.df.as_ref()
}
pub fn display_slice_df(&self) -> Option<DataFrame> {
let df = self.view.df.as_ref()?;
let offset = self
.view
.start_row
.saturating_sub(self.view.buffered_start_row);
let slice_len = self.visible_rows.min(df.height().saturating_sub(offset));
if offset < df.height() && slice_len > 0 {
Some(df.slice(offset as i64, slice_len))
} else {
None
}
}
pub fn copy_row_df(&self) -> Option<DataFrame> {
let df = self.view.buffered_df.as_ref()?;
let absolute = self.view.start_row + self.table_state.selected()?;
let offset = absolute.checked_sub(self.view.buffered_start_row)?;
if offset >= df.height() {
return None;
}
let names: Vec<&str> = self.view.column_order.iter().map(|s| s.as_str()).collect();
df.select(names).ok().map(|d| d.slice(offset as i64, 1))
}
pub fn copy_view_df(&self) -> Option<DataFrame> {
let df = self.view.buffered_df.as_ref()?;
let names: Vec<&str> = self.view.column_order.iter().map(|s| s.as_str()).collect();
let selected = df.select(names).ok()?;
let offset = self
.view
.start_row
.saturating_sub(self.view.buffered_start_row);
let len = self
.visible_rows
.min(selected.height().saturating_sub(offset));
(len > 0).then(|| selected.slice(offset as i64, len))
}
pub fn copy_cell_value(&self, column: &str) -> Option<String> {
let row = self.copy_row_df()?;
crate::exact::copy_text(row.column(column).ok()?).ok()
}
pub fn selected_display_row(&self) -> Option<usize> {
Some(self.view.start_row + self.table_state.selected()? + self.row_start_index)
}
pub fn estimated_copy_bytes(&self) -> Option<usize> {
let rows = self.num_rows_if_valid()?;
if rows == 0 {
return Some(0);
}
let base64 = |bytes: usize| bytes.div_ceil(3) * 4;
let footer_width = |name: &str| {
self.column_bytes
.iter()
.find(|(n, _)| n == name)
.map(|(_, w)| *w)
};
let mut row = self.bytes_per_row();
for name in &self.view.column_order {
match self.view.schema.get(name.as_str()) {
Some(DataType::Binary) => row += base64(footer_width(name)?),
Some(dtype) if crate::export::nested_json::has_binary(dtype) => {
let buffered = self.view.buffered_df.as_ref().and_then(|df| {
let column = df.column(name).ok()?;
(df.height() > 0)
.then(|| column.as_materialized_series().estimated_size() / df.height())
});
row += buffered.or_else(|| footer_width(name)).unwrap_or(0) / 3;
}
_ => {}
}
}
Some(rows.saturating_mul(row))
}
pub fn display_drift(&self, frame_rows: usize) -> Vec<u32> {
if !self.view.drift_column_present {
return Vec::new();
}
let Some(df) = self.view.buffered_df.as_ref() else {
return Vec::new();
};
let Ok(column) = df.column(crate::formats::schema_union::DRIFT_COLUMN) else {
return Vec::new();
};
let offset = self
.view
.start_row
.saturating_sub(self.view.buffered_start_row);
let len = frame_rows.min(column.len().saturating_sub(offset));
if len == 0 {
return Vec::new();
}
let slice = column.slice(offset as i64, len);
let Ok(rows) = slice.u32() else {
return Vec::new();
};
rows.iter()
.map(|row| self.file_group_of(row.unwrap_or(0) as usize))
.collect()
}
pub fn max_buffered_rows(&self) -> usize {
self.max_buffered_rows
}
pub fn max_buffered_mb(&self) -> usize {
self.max_buffered_mb
}
pub fn can_drill_down(&self) -> bool {
!self.is_drilled_down() && self.is_grouped()
}
fn drills_lists(&self) -> bool {
self.has_list_columns()
&& self
.view
.group_source
.as_ref()
.is_some_and(|s| s.rows_in_lists)
}
fn drill_columns(&self) -> Vec<String> {
if self.drills_lists() {
return self
.view
.schema
.iter_names()
.map(|n| n.to_string())
.collect();
}
self.view
.group_source
.iter()
.flat_map(|source| source.keys.iter().map(|(name, _)| name.to_string()))
.collect()
}
pub fn inspect_fields(&self) -> Vec<InspectField> {
let shown = self.view.column_order.iter().filter_map(|name| {
Some(InspectField {
name: name.clone(),
dtype: self.view.schema.get(name.as_str())?.clone(),
hidden: false,
})
});
let hidden = self
.view
.schema
.iter()
.filter(|(name, _)| {
name.as_str() != crate::formats::schema_union::DRIFT_COLUMN
&& !self.view.column_order.iter().any(|c| c == name.as_str())
})
.map(|(name, dtype)| InspectField {
name: name.to_string(),
dtype: dtype.clone(),
hidden: true,
});
shown.chain(hidden).collect()
}
pub fn inspect_row(&self) -> Option<InspectRow> {
self.inspect_row_at(self.view.start_row + self.table_state.selected()?)
}
pub fn inspect_row_at(&self, row: usize) -> Option<InspectRow> {
let df = self.view.buffered_df.as_ref()?;
let offset = row.checked_sub(self.view.buffered_start_row)?;
if offset >= df.height() {
return None;
}
let names: Vec<&str> = self.view.column_order.iter().map(|s| s.as_str()).collect();
let values = df.select(names).ok()?.slice(offset as i64, 1);
let drift_group = self
.view
.drift_column_present
.then(|| df.column(crate::formats::schema_union::DRIFT_COLUMN).ok())
.flatten()
.and_then(|c| c.get(offset).ok())
.and_then(|v| v.extract::<usize>())
.map(|place| self.file_group_of(place));
Some(InspectRow {
row,
frame: self.view.len_generation,
display_row: row + self.row_start_index,
values,
drift_group,
})
}
fn file_group_of(&self, place: usize) -> u32 {
let file = self
.drift_file_starts
.partition_point(|&start| start <= place)
.saturating_sub(1);
self.drift_file_group.get(file).copied().unwrap_or(0)
}
pub fn null_kind(&self, column: &str, group: Option<u32>) -> NullKind {
let Some(group) = group.and_then(|g| self.view.drift_groups.get(g as usize)) else {
return NullKind::Null;
};
if group.absent.iter().any(|c| c == column) {
NullKind::Absent
} else if group.unread.iter().any(|c| c == column) {
NullKind::Conflict
} else {
NullKind::Null
}
}
pub fn inspect_read_lf(&self, row: usize, columns: &[String]) -> PolarsResult<LazyFrame> {
let exprs = columns.iter().map(|c| col(c.as_str())).collect();
self.window_lf(row, 1, exprs)
}
pub fn drill_row(&self, group_index: usize) -> Option<DrillRow> {
if !self.can_drill_down() {
return None;
}
let columns = self.drill_columns();
let buffered = self
.view
.buffered_df
.as_ref()
.filter(|_| {
(self.view.buffered_start_row..self.view.buffered_end_row).contains(&group_index)
})
.filter(|_| {
columns
.iter()
.all(|c| !matches!(self.view.schema.get(c.as_str()), Some(DataType::Binary)))
})
.and_then(|df| df.select(columns.iter().map(|c| c.as_str())).ok())
.map(|df| df.slice((group_index - self.view.buffered_start_row) as i64, 1))
.filter(|row| row.height() == 1);
Some(match buffered {
Some(row) => DrillRow::Buffered(row),
None => DrillRow::Read(Box::new(
self.visible_lf()
.select(columns.iter().map(|c| col(c.as_str())).collect::<Vec<_>>())
.slice(group_index as i64, 1),
)),
})
}
pub fn drill_down_into_group(&mut self, group_index: usize) -> Result<()> {
let row = match self.drill_row(group_index) {
None => return Ok(()),
Some(DrillRow::Buffered(row)) => row,
Some(DrillRow::Read(lf)) => collect_lazy(*lf, self.polars_streaming)?,
};
self.drill_down_with_row(group_index, &row)
}
pub fn drill_down_with_row(&mut self, group_index: usize, row: &DataFrame) -> Result<()> {
if !self.can_drill_down() {
return Ok(());
}
if row.height() == 0 {
return Err(color_eyre::eyre::eyre!("Group index out of bounds"));
}
let mut group = if self.drills_lists() {
Self::group_from_lists(
row,
self.group_key_columns(),
self.group_value_columns(),
self.view.lineage.clone(),
)?
} else if let Some(source) = &self.view.group_source {
Self::group_from_source(source, row)?
} else {
return Ok(());
};
if let Some(source) = self
.view
.group_source
.as_ref()
.filter(|_| self.drills_lists())
{
let keys: Vec<&str> = source.keys.iter().map(|(n, _)| n.as_str()).collect();
(group.key_columns, group.key_values) = group
.key_columns
.into_iter()
.zip(group.key_values)
.filter(|(name, _)| keys.contains(&name.as_str()))
.unzip();
}
let keys = row.select(group.key_columns.iter().map(String::as_str))?;
self.enter_group(group, group_index, false, keys)
}
pub fn drill_into_value(&mut self, column: &str, value: AnyValue<'static>) -> Result<()> {
let (dtype, value) = self
.value_in_column(column, value)
.ok_or_else(|| color_eyre::eyre::eyre!("{column} has no values of that type"))?;
let label = crate::exact::str_value(&value).to_string();
let mut steps = self.view_steps();
steps.push(match crate::export::python_script::py_value(&value) {
Some(literal) => Step::Matching(vec![(
format!("pl.col({})", crate::export::python_script::py_str(column)),
literal,
)]),
None => Step::Unreproducible(format!(
"drilled down to the rows where {column} is {label}, a value of a type not written as Python"
)),
});
let keys = DataFrame::new(
1,
vec![Column::new_scalar(
column.into(),
Scalar::new(dtype.clone(), value.clone()),
1,
)],
)?;
let matches = col(column).eq_missing(lit(Scalar::new(dtype, value)));
let group = GroupRows {
lf: self.visible_lf().filter(matches),
key_columns: vec![column.to_string()],
key_values: vec![label],
lead: vec![column.to_string()],
steps,
lineage: self.view.lineage.clone(),
};
if !self.is_drilled_down() {
let index = self.view.start_row + self.table_state.selected().unwrap_or(0);
return self.enter_group(group, index, true, keys);
}
let schema = group.lf.clone().collect_schema()?;
let order = std::mem::take(&mut self.view.column_order);
if let Some(keys) = self.view.drilled_down_group_key_columns.as_mut() {
keys.extend(group.key_columns);
}
if let Some(values) = self.view.drilled_down_group_key.as_mut() {
values.extend(group.key_values);
}
if let Some(grouped) = self.view.grouped.as_mut() {
grouped.steps.push(keys);
}
self.view.filters.clear();
self.view.sort_columns.clear();
self.view.sort_descending.clear();
self.view.sort_ascending = true;
self.install_base(group.lf, schema);
self.view.base_steps = group.steps;
self.view.lineage = group.lineage;
self.view.column_order = order;
self.view.start_row = 0;
self.termcol_index = 0;
self.clear_column_moves();
self.settle_cursor();
self.table_state.select(Some(0));
self.collect();
Ok(())
}
pub fn drilled_into_value(&self) -> bool {
self.view.grouped.as_ref().is_some_and(|view| view.by_value)
}
fn enter_group(
&mut self,
group: GroupRows,
group_index: usize,
by_value: bool,
keys: DataFrame,
) -> Result<()> {
let schema = group.lf.clone().collect_schema()?;
self.view.drilled_down_group_key = Some(group.key_values);
self.view.drilled_down_group_key_columns = Some(group.key_columns);
self.view.grouped = Some(GroupedView {
lf: self.view.lf.clone(),
base_lf: self.view.base_lf.clone(),
base_schema: self.view.base_schema.clone(),
schema: self.view.schema.clone(),
filters: std::mem::take(&mut self.view.filters),
sort_columns: std::mem::take(&mut self.view.sort_columns),
sort_descending: std::mem::take(&mut self.view.sort_descending),
sort_ascending: self.view.sort_ascending,
drift: self.view.drift_column_present,
drift_groups: self.view.drift_groups.clone(),
view_numbered: self.view.view_numbered,
notes: self.view.notes.clone(),
group_source: self.view.group_source.take(),
column_order: self.view.column_order.clone(),
locked_columns_count: self.view.locked_columns_count,
start_row: self.view.start_row,
termcol_index: self.termcol_index,
cursor_column: self.view.cursor_column.clone(),
selected: self.table_state.selected(),
by_value,
steps: vec![keys],
base_steps: std::mem::take(&mut self.view.base_steps),
lineage: self.view.lineage.clone(),
});
self.view.sort_ascending = true;
self.install_base(group.lf, schema);
self.view.base_steps = group.steps;
self.view.lineage = group.lineage;
let rest: Vec<String> = std::mem::take(&mut self.view.column_order)
.into_iter()
.filter(|c| !group.lead.contains(c))
.collect();
self.view.column_order = group.lead.into_iter().chain(rest).collect();
self.view.drilled_down_group_index = Some(group_index);
self.view.start_row = 0;
self.termcol_index = 0;
self.clear_column_moves();
self.view.locked_columns_count = 0;
self.settle_cursor();
self.table_state.select(Some(0));
self.collect();
Ok(())
}
fn group_from_lists(
row: &DataFrame,
key_columns: Vec<String>,
value_columns: Vec<String>,
lineage: Lineage,
) -> Result<GroupRows> {
if value_columns.is_empty() {
return Err(color_eyre::eyre::eyre!("No value columns in grouped data"));
}
let row_count = match row.column(&value_columns[0])?.get(0)? {
AnyValue::List(list_series) => list_series.len(),
_ => 0,
};
let mut columns = Vec::new();
let mut key_values = Vec::new();
for col_name in &key_columns {
let key = row.column(col_name)?;
key_values.push(crate::exact::str_value(&key.get(0)?).to_string());
columns.push(key.new_from_index(0, row_count));
}
for col_name in &value_columns {
if let AnyValue::List(list_series) = row.column(col_name)?.get(0)? {
columns.push(list_series.with_name(col_name.as_str().into()).into());
}
}
let group = key_columns
.iter()
.zip(&key_values)
.map(|(c, v)| format!("{c} = {v}"))
.collect::<Vec<_>>()
.join(", ");
Ok(GroupRows {
lf: DataFrame::new_infer_height(columns)?.lazy(),
key_columns,
key_values,
lead: Vec::new(),
steps: vec![Step::Unreproducible(format!(
"drilled down into the group {group}, read from the grouped result's lists: \
not written as Python"
))],
lineage,
})
}
fn group_from_source(source: &GroupSource, row: &DataFrame) -> Result<GroupRows> {
let mut predicate: Option<Expr> = None;
let mut key_columns = Vec::new();
let mut key_values = Vec::new();
let mut lead = Vec::new();
let mut matching = Vec::new();
for (i, (name, expr)) in source.keys.iter().enumerate() {
let column = row.column(name)?;
let value = column.get(0)?.into_static();
matching.push(
source
.python_keys
.get(i)
.cloned()
.flatten()
.zip(crate::export::python_script::py_value(&value)),
);
key_columns.push(name.to_string());
key_values.push(crate::exact::str_value(&value).to_string());
let key = expr.clone().meta().undo_aliases();
if let Expr::Column(source_column) = &key
&& !source.scratch.contains(source_column)
{
lead.push(source_column.to_string());
}
let matches = key.eq_missing(lit(Scalar::new(column.dtype().clone(), value)));
predicate = Some(match predicate {
Some(all) => all.and(matches),
None => matches,
});
}
let rows = source.rows.clone();
let mut lf = match predicate {
Some(predicate) => rows.filter(predicate),
None => rows,
};
if !source.scratch.is_empty() {
lf = lf.drop(by_name(source.scratch.iter().cloned(), true, false));
}
let matching: Option<Vec<(String, String)>> = matching.into_iter().collect();
let steps = match (&source.python_rows, matching) {
(Some(rows), Some(matching)) => {
let mut steps = rows.clone();
steps.push(Step::Matching(matching));
if !source.scratch.is_empty() {
steps.push(Step::Drop(
source.scratch.iter().map(|c| c.to_string()).collect(),
));
}
steps
}
_ => vec![Step::Unreproducible(format!(
"drilled down into the group {}: not written as Python",
key_columns
.iter()
.zip(&key_values)
.map(|(c, v)| format!("{c} = {v}"))
.collect::<Vec<_>>()
.join(", ")
))],
};
Ok(GroupRows {
lf,
key_columns,
key_values,
lead,
steps,
lineage: source.lineage.clone(),
})
}
pub fn drill_up(&mut self) -> Result<()> {
let Some(view) = self.view.grouped.take() else {
return Err(color_eyre::eyre::eyre!("Not in drill-down mode"));
};
self.invalidate_num_rows();
self.drop_buffer();
self.view.observed_bytes_per_row = None;
self.widths.relearn();
self.view.lf = view.lf;
self.view.unsorted_lf = None;
self.view.base_lf = view.base_lf;
self.view.base_schema = view.base_schema;
self.view.base_steps = view.base_steps;
self.view.lineage = view.lineage;
self.view.filters = view.filters;
self.view.sort_columns = view.sort_columns;
self.view.sort_descending = view.sort_descending;
self.view.sort_ascending = view.sort_ascending;
self.view.drift_column_present = view.drift;
self.view.drift_groups = view.drift_groups;
self.view.view_numbered = view.view_numbered;
self.view.notes = view.notes;
self.view.group_source = view.group_source;
self.view.view_notes = self.view_notes_only();
self.view.schema = view.schema;
self.view.column_order = view.column_order;
self.view.locked_columns_count = view.locked_columns_count;
self.view.drilled_down_group_index = None;
self.view.drilled_down_group_key = None;
self.view.drilled_down_group_key_columns = None;
self.view.start_row = view.start_row;
self.termcol_index = view.termcol_index;
self.clear_column_moves();
self.view.cursor_column = view.cursor_column;
self.settle_cursor();
self.table_state.select(view.selected);
self.collect();
Ok(())
}
}