1use super::*;
5
6pub(super) fn pivot_agg_expr(agg: PivotAggregation, values: Expr) -> Expr {
8 match agg {
9 PivotAggregation::Last => values.last(),
10 PivotAggregation::First => values.first(),
11 PivotAggregation::Min => values.min(),
12 PivotAggregation::Max => values.max(),
13 PivotAggregation::Avg => values.mean(),
14 PivotAggregation::Med => values.median(),
15 PivotAggregation::Std => values.std(1),
16 PivotAggregation::Count => values.len(),
17 }
18}
19
20pub const PIVOT_COLUMN_LIMIT: usize = 10_000;
23
24pub struct PivotJob {
26 pub(super) view: LazyFrame,
27 pub(super) spec: PivotSpec,
28 pub(super) streaming: bool,
29}
30
31impl PivotJob {
32 pub(crate) fn new(view: LazyFrame, spec: PivotSpec, streaming: bool) -> Self {
34 Self {
35 view,
36 spec,
37 streaming,
38 }
39 }
40
41 pub fn run(self) -> Result<DataFrame> {
47 let on = self.spec.pivot_column.as_str();
48 let value = self.spec.value_column.as_str();
49 let index: Vec<PlSmallStr> = if self.spec.index.is_empty() {
50 self.view
51 .clone()
52 .collect_schema()?
53 .iter_names()
54 .filter(|name| name.as_str() != on && name.as_str() != value)
55 .cloned()
56 .collect()
57 } else {
58 self.spec.index.iter().map(PlSmallStr::from).collect()
59 };
60 let keys: Vec<Expr> = index
63 .iter()
64 .cloned()
65 .chain([PlSmallStr::from(on)])
66 .map(Expr::Column)
67 .collect();
68 let cells = collect_lazy(
69 self.view.group_by_stable(keys).agg([pivot_agg_expr(
70 self.spec.aggregation,
71 Expr::Column(PlSmallStr::from(value)),
72 )
73 .alias(value)]),
74 self.streaming,
75 )?;
76 let on_columns = cells
77 .clone()
78 .lazy()
79 .select([Expr::Column(PlSmallStr::from(on))])
80 .unique(None, UniqueKeepStrategy::Any)
81 .sort([on], SortMultipleOptions::default().with_nulls_last(true))
82 .collect()?;
83 if on_columns.height() > PIVOT_COLUMN_LIMIT {
86 return Err(color_eyre::eyre::eyre!(
87 "Pivot would make {} columns from {on}; the limit is {}. Filter first, or pivot a column with fewer values",
88 numfmt::group_chrome(on_columns.height()),
89 numfmt::group_chrome(PIVOT_COLUMN_LIMIT),
90 ));
91 }
92 let (cells, on_columns) = Self::pivot_dates_as_text(cells, on_columns, on)?;
93 let cell = match self.spec.aggregation {
96 PivotAggregation::Count => element().sum(),
97 _ => element().first(),
98 };
99 let pivoted = cells
100 .lazy()
101 .pivot(
102 by_name([on], true, false),
103 Arc::new(on_columns),
104 by_name(index, true, false),
105 by_name([value], true, false),
106 cell,
107 true,
108 PlSmallStr::from_static("_"),
109 PivotColumnNaming::Auto,
110 )
111 .collect()?;
112 Ok(pivoted)
113 }
114
115 fn pivot_dates_as_text(
119 cells: DataFrame,
120 on_columns: DataFrame,
121 on: &str,
122 ) -> Result<(DataFrame, DataFrame)> {
123 let values = on_columns.column(on)?.as_materialized_series();
124 if crate::exact::calendar_without_out_of_range(values)?.is_none() {
125 return Ok((cells, on_columns));
126 }
127 let text = |mut df: DataFrame| -> Result<DataFrame> {
128 let values = df.column(on)?.as_materialized_series();
129 let values = crate::past_calendar::cast_text(
130 values,
131 polars::chunked_array::cast::CastOptions::NonStrict,
132 )?;
133 df.with_column(values.into_column())?;
134 Ok(df)
135 };
136 Ok((text(cells)?, text(on_columns)?))
137 }
138}
139
140pub(super) type Lineage = Option<Arc<Vec<(String, String)>>>;
144
145pub(super) fn traced(root: &Lineage, pairs: Vec<(String, String)>) -> Lineage {
148 let pairs = match root {
149 None => pairs,
150 Some(root) => pairs
151 .into_iter()
152 .filter_map(|(shown, from)| {
153 root.iter()
154 .find(|(name, _)| *name == from)
155 .map(|(_, loaded)| (shown, loaded.clone()))
156 })
157 .collect(),
158 };
159 Some(Arc::new(pairs))
160}
161
162pub(super) fn passed_through(exprs: &[Expr]) -> Vec<(String, String)> {
165 exprs
166 .iter()
167 .filter_map(|e| {
168 let Expr::Column(from) = e.clone().meta().undo_aliases() else {
169 return None;
170 };
171 let shown = e.clone().meta().output_name().ok()?;
172 Some((shown.to_string(), from.to_string()))
173 })
174 .collect()
175}
176
177pub(super) enum ActiveQuery {
179 Dsl(String),
180 #[cfg(feature = "sql")]
181 Sql(String),
182 Fuzzy(String),
183}
184
185pub(crate) fn fuzzy_token_regex(token: &str) -> String {
187 let inner: String =
188 token
189 .chars()
190 .map(|c| regex::escape(&c.to_string()))
191 .fold(String::new(), |mut s, e| {
192 if !s.is_empty() {
193 s.push_str(".*");
194 }
195 s.push_str(&e);
196 s
197 });
198 format!("(?i).*{}.*", inner)
199}
200
201pub(super) fn sort_options(descending: Vec<bool>) -> SortMultipleOptions {
206 let n = descending.len();
207 SortMultipleOptions::default()
208 .with_order_descending_multi(descending)
209 .with_nulls_last_multi(vec![true; n])
210 .with_maintain_order(true)
211}
212
213#[derive(Debug, Clone, Copy, PartialEq, Eq)]
215pub struct ViewChange {
216 pub filters: bool,
217 pub sort: bool,
218}
219
220impl DataTableState {
221 pub fn plan_pivot(&self, spec: &PivotSpec) -> PivotJob {
224 PivotJob {
225 view: self.visible_lf(),
226 spec: spec.clone(),
227 streaming: self.polars_streaming,
228 }
229 }
230
231 pub fn install_pivot(&mut self, spec: &PivotSpec, pivoted: DataFrame) -> Result<()> {
233 let index = if spec.index.is_empty() {
234 self.view
236 .schema
237 .iter_names()
238 .map(|n| n.to_string())
239 .filter(|n| {
240 n != &spec.pivot_column
241 && n != &spec.value_column
242 && n != crate::formats::schema_union::DRIFT_COLUMN
243 })
244 .collect()
245 } else {
246 spec.index.clone()
247 };
248 let kept = index.clone();
249 let step = Step::Pivot {
250 index,
251 on: spec.pivot_column.clone(),
252 values: spec.value_column.clone(),
253 aggregation: spec.aggregation,
254 };
255 self.view.last_pivot_spec = Some(spec.clone());
256 self.view.last_melt_spec = None;
257 self.replace_lf_after_reshape(pivoted.lazy(), step, &kept)
258 }
259
260 pub fn pivot(&mut self, spec: &PivotSpec) -> Result<()> {
263 let pivoted = self.plan_pivot(spec).run()?;
264 self.install_pivot(spec, pivoted)
265 }
266
267 pub(crate) fn melt_lf(view: LazyFrame, spec: &MeltSpec) -> Result<LazyFrame> {
270 let on = cols(spec.value_columns.iter().map(|s| s.as_str()));
271 let index = cols(spec.index.iter().map(|s| s.as_str()));
272 let args = UnpivotArgsDSL {
273 on: Some(on),
274 index,
275 variable_name: Some(PlSmallStr::from(spec.variable_name.as_str())),
276 value_name: Some(PlSmallStr::from(spec.value_name.as_str())),
277 };
278 Ok(Self::melt_dates_as_text(view, spec, &args)?.unpivot(args))
279 }
280
281 pub fn melt(&mut self, spec: &MeltSpec) -> Result<()> {
283 let lf = Self::melt_lf(self.visible_lf(), spec)?;
284 let step = Step::Melt {
285 index: spec.index.clone(),
286 on: spec.value_columns.clone(),
287 variable_name: spec.variable_name.clone(),
288 value_name: spec.value_name.clone(),
289 };
290 self.view.last_melt_spec = Some(spec.clone());
291 self.view.last_pivot_spec = None;
292 self.replace_lf_after_reshape(lf, step, &spec.index)?;
293 Ok(())
294 }
295
296 fn melt_dates_as_text(
299 view: LazyFrame,
300 spec: &MeltSpec,
301 args: &UnpivotArgsDSL,
302 ) -> Result<LazyFrame> {
303 let schema = view.clone().collect_schema()?;
304 let melted = view.clone().unpivot(args.clone()).collect_schema()?;
305 if melted.get(spec.value_name.as_str()) != Some(&DataType::String) {
306 return Ok(view);
307 }
308 let texts: Vec<Expr> = spec
309 .value_columns
310 .iter()
311 .filter(|name| {
312 schema
313 .get(name.as_str())
314 .is_some_and(crate::past_calendar::can_leave_calendar)
315 })
316 .map(|name| {
317 crate::past_calendar::text_expr(
318 Expr::Column(PlSmallStr::from(name.as_str())),
319 polars::chunked_array::cast::CastOptions::NonStrict,
320 )
321 })
322 .collect();
323 Ok(if texts.is_empty() {
324 view
325 } else {
326 view.with_columns(texts)
327 })
328 }
329
330 fn replace_lf_after_reshape(
333 &mut self,
334 lf: LazyFrame,
335 step: Step,
336 kept: &[String],
337 ) -> Result<()> {
338 let schema = lf.clone().collect_schema()?;
339 let lineage = traced(
340 &self.view.lineage,
341 kept.iter().map(|c| (c.clone(), c.clone())).collect(),
342 );
343 let mut steps = self.view_steps();
344 steps.push(step);
345 let text = |q: &str| Some(q.trim().to_string()).filter(|q| !q.is_empty());
348 let source = ReshapeSource {
349 query: text(&self.view.active_query),
350 sql_query: text(&self.view.active_sql_query),
351 fuzzy_query: text(&self.view.active_fuzzy_query),
352 filters: self.view.filters.clone(),
353 sort_columns: self.view.sort_columns.clone(),
354 sort_descending: self.view.sort_descending.clone(),
355 };
356 self.view.reshape_source =
357 (self.view.reshaped_lf.is_none() && !source.is_empty()).then_some(source);
358 self.view.reshaped_lf = Some(lf.clone());
359 self.install_base(lf, schema);
360 self.view.base_steps = steps.clone();
361 self.view.reshape_steps = Some(steps);
362 self.view.lineage = lineage.clone();
363 self.view.reshape_lineage = lineage;
364 self.reset_view_state(0);
365 self.error = None;
366 self.view.df = None;
367 self.view.locked_df = None;
368 self.collect();
369 Ok(())
370 }
371
372 fn typed_filters(&self) -> Vec<SidebarFilter> {
374 self.view
375 .filters
376 .iter()
377 .map(|f| SidebarFilter::typed_in(f, &self.view.schema, &self.view.column_order))
378 .collect()
379 }
380
381 pub fn read_python(&self) -> &[String] {
383 &self.read_python
384 }
385
386 pub(crate) fn unfit_to_count(
389 &self,
390 ) -> Option<(LazyFrame, Vec<crate::formats::column_types::Typed>)> {
391 if self.unfit_notes.is_some() || self.typing.typed.is_empty() {
392 return None;
393 }
394 Some((self.typing.source.clone()?, self.typing.typed.clone()))
395 }
396
397 pub fn column_changes(&self) -> &[crate::formats::column_types::ColumnChange] {
399 &self.view.column_changes
400 }
401
402 pub fn retyped_columns(&self) -> Vec<String> {
404 self.view
405 .column_changes
406 .iter()
407 .filter(|c| matches!(c.change, crate::formats::column_types::Change::Typed(_)))
408 .map(|c| c.name.clone())
409 .collect()
410 }
411
412 pub fn type_as_read(&self, column: &str) -> Option<DataType> {
414 self.view
415 .base_schema
416 .get(column)
417 .or_else(|| self.view.schema.get(column))
418 .cloned()
419 }
420
421 pub fn values_on_screen(&self, column: &str, n: usize) -> Vec<String> {
424 let from_buffer = self.column_type_of(column).is_none();
425 let df = if from_buffer {
426 self.view.buffered_df.clone()
427 } else {
428 self.view
429 .base_lf
430 .clone()
431 .select([col(column)])
432 .limit(n as IdxSize * 10)
433 .collect()
434 .ok()
435 };
436 let Some(values) = df.and_then(|df| df.column(column).ok().cloned()) else {
437 return Vec::new();
438 };
439 let Ok(text) = values.cast(&DataType::String) else {
440 return Vec::new();
441 };
442 let Ok(text) = text.str().cloned() else {
443 return Vec::new();
444 };
445 text.iter()
446 .flatten()
447 .map(str::trim)
448 .filter(|v| !v.is_empty())
449 .take(n)
450 .map(str::to_string)
451 .collect()
452 }
453
454 pub fn column_type_of(
456 &self,
457 column: &str,
458 ) -> Option<&crate::formats::column_types::ColumnType> {
459 self.view
460 .column_changes
461 .iter()
462 .find_map(|c| match &c.change {
463 crate::formats::column_types::Change::Typed(ty) if c.name == column => Some(ty),
464 _ => None,
465 })
466 }
467
468 pub fn set_column_type(
471 &mut self,
472 column: &str,
473 ty: Option<crate::formats::column_types::ColumnType>,
474 ) {
475 use crate::formats::column_types::{Change, ColumnChange};
476 self.view
477 .column_changes
478 .retain(|c| !(c.name == column && matches!(c.change, Change::Typed(_))));
479 if let Some(ty) = ty {
480 self.view.column_changes.push(ColumnChange {
481 name: column.to_string(),
482 change: Change::Typed(ty),
483 });
484 }
485 self.column_changes_changed();
486 }
487
488 pub fn add_made_column(
491 &mut self,
492 derived: crate::formats::column_types::Derived,
493 ) -> std::result::Result<(), String> {
494 use crate::formats::column_types::{Change, ColumnChange};
495 if self.view.schema.contains(&derived.name) {
496 return Err(format!("a column is named {} already", derived.name));
497 }
498 for from in &derived.from {
499 if !self.view.schema.contains(from) {
500 return Err(format!("no column {from}"));
501 }
502 }
503 let first = derived.from[0].clone();
504 let at = self
505 .view
506 .column_order
507 .iter()
508 .position(|c| *c == first)
509 .unwrap_or(self.view.column_order.len());
510 self.view.column_order.insert(at, derived.name.clone());
511 self.view.column_changes.push(ColumnChange {
512 name: derived.name,
513 change: Change::Made {
514 from: derived.from,
515 kind: derived.kind.name().to_string(),
516 format: derived.format,
517 },
518 });
519 self.column_changes_changed();
520 Ok(())
521 }
522
523 pub fn set_column_changes(
526 &mut self,
527 changes: &[crate::formats::column_types::ColumnChange],
528 ) -> Vec<String> {
529 self.view.column_changes = Vec::new();
530 let base = self.view.base_schema.clone();
531 let mut known: Vec<String> = base.iter_names().map(|n| n.to_string()).collect();
532 let mut dropped = Vec::new();
533 for change in changes {
534 let fits = match &change.change {
535 crate::formats::column_types::Change::Typed(_) => known.contains(&change.name),
536 crate::formats::column_types::Change::Made { from, .. } => {
537 from.iter().all(|f| known.contains(f)) && change.derived().is_some()
538 }
539 };
540 if fits {
541 if !known.contains(&change.name) {
542 known.push(change.name.clone());
543 }
544 self.view.column_changes.push(change.clone());
545 } else {
546 dropped.push(change.name.clone());
547 }
548 }
549 self.view.changes_dropped = if dropped.is_empty() {
550 Vec::new()
551 } else {
552 vec![crate::notes::Note {
553 summary: format!(
554 "view steps left out, no such column: {}",
555 crate::notes::some_names(&dropped)
556 ),
557 scope: "the view's column types".to_string(),
558 read_as_text: None,
559 passed_over: None,
560 }]
561 };
562 for change in &self.view.column_changes {
564 if let crate::formats::column_types::Change::Made { from, .. } = &change.change
565 && !self.view.column_order.contains(&change.name)
566 {
567 let at = self
568 .view
569 .column_order
570 .iter()
571 .position(|c| *c == from[0])
572 .unwrap_or(self.view.column_order.len());
573 self.view.column_order.insert(at, change.name.clone());
574 }
575 }
576 self.column_changes_changed();
577 dropped
578 }
579
580 pub(super) fn forget_column_changes(&mut self) {
582 if self.view.column_changes.is_empty() && self.view.changes_dropped.is_empty() {
583 return;
584 }
585 self.view.column_changes.clear();
586 self.view.changes_dropped.clear();
587 self.view.changes_version += 1;
588 self.changes_unfit = None;
589 }
590
591 fn column_changes_changed(&mut self) {
594 self.view.changes_version += 1;
595 let (changed, _) = self.with_column_changes(self.view.base_lf.clone());
596 if let Ok(schema) = changed.clone().collect_schema() {
597 self.view.schema = schema;
598 }
599 let schema = self.view.schema.clone();
600 self.view.column_order.retain(|c| schema.contains(c));
601 for name in schema.iter_names() {
602 if !self.view.column_order.iter().any(|c| c == name.as_str()) {
603 self.view.column_order.push(name.to_string());
604 }
605 }
606 self.widths.relearn();
607 self.drop_buffer();
608 self.apply_transformations();
609 }
610
611 fn with_column_changes(
615 &self,
616 mut lf: LazyFrame,
617 ) -> (
618 LazyFrame,
619 Option<(LazyFrame, Vec<crate::formats::column_types::Typed>)>,
620 ) {
621 use crate::formats::column_types::Change;
622 if self.view.column_changes.is_empty() {
623 return (lf, None);
624 }
625 let mut schema = (*self.view.base_schema).clone();
627 let mut made = lf.clone();
628 let mut typed = Vec::new();
629 for change in &self.view.column_changes {
630 let name = PlSmallStr::from(change.name.as_str());
631 match &change.change {
632 Change::Typed(ty) => {
633 let Some(from) = schema.get(&name).cloned() else {
634 continue;
635 };
636 lf = lf.with_column(ty.expr(&change.name, &from).alias(name.clone()));
637 typed.push(crate::formats::column_types::Typed {
638 column: change.name.clone(),
639 ty: ty.clone(),
640 from,
641 });
642 schema.with_column(name, ty.dtype.clone());
643 }
644 Change::Made { from, .. } => {
645 let Some(derived) = change.derived() else {
646 continue;
647 };
648 if !from.iter().all(|f| schema.contains(f.as_str())) {
649 continue;
650 }
651 lf = lf.with_column(derived.expr().alias(name.clone()));
652 made = made.with_column(derived.expr().alias(name.clone()));
653 schema.with_column(name, DataType::Null);
654 }
655 }
656 }
657 let count = (!typed.is_empty()).then_some((made, typed));
658 (lf, count)
659 }
660
661 pub(crate) fn changes_unfit_to_count(
664 &self,
665 ) -> Option<(LazyFrame, Vec<crate::formats::column_types::Typed>, u64)> {
666 if self
667 .changes_unfit
668 .as_ref()
669 .is_some_and(|(version, _)| *version == self.view.changes_version)
670 {
671 return None;
672 }
673 let (_, count) = self.with_column_changes(self.view.base_lf.clone());
674 let (source, typed) = count?;
675 Some((source, typed, self.view.changes_version))
676 }
677
678 pub(crate) fn changes_unfit_counted(
680 &mut self,
681 version: u64,
682 unfit: &[crate::formats::column_types::Unfit],
683 ) {
684 if version == self.view.changes_version {
685 if !unfit.is_empty() {
687 self.view.notes_seen = false;
688 }
689 self.changes_unfit = Some((
690 version,
691 crate::formats::column_types::unfit_notes(unfit, "the view's column types"),
692 ));
693 }
694 }
695
696 pub(crate) fn unfit_counted(&mut self, unfit: &[crate::formats::column_types::Unfit]) {
698 if !unfit.is_empty() {
699 self.view.notes_seen = false;
700 }
701 self.unfit_notes = Some(crate::formats::column_types::unfit_notes(
702 unfit,
703 "counted over every row",
704 ));
705 }
706
707 pub(super) fn view_steps(&self) -> Vec<Step> {
709 let mut steps = self.view.base_steps.clone();
710 if !self.view.column_changes.is_empty() {
711 let said: Vec<String> = self
712 .view
713 .column_changes
714 .iter()
715 .map(crate::formats::column_types::ColumnChange::to_toml)
716 .collect();
717 steps.push(Step::Unreproducible(format!(
718 "datui typed columns as a format spec would: {}",
719 said.join("; ")
720 )));
721 }
722 if !self.view.filters.is_empty() {
723 let typed = self.typed_filters();
724 let durations: Vec<String> = typed
725 .iter()
726 .flat_map(|f| f.unscriptable_columns())
727 .collect();
728 if !durations.is_empty() {
731 steps.push(Step::Unreproducible(format!(
732 "a kept find matches {} as datui writes durations",
733 durations.join(", ")
734 )));
735 }
736 steps.push(Step::Filter(typed));
737 }
738 let left_out: Vec<String> = self
741 .view_exclusions()
742 .into_iter()
743 .map(|(_, note)| note.summary)
744 .collect();
745 if !left_out.is_empty() {
746 steps.push(Step::Unreproducible(left_out.join("; ")));
747 }
748 if !self.view.sort_columns.is_empty() {
749 steps.push(Step::Sort {
750 columns: self.view.sort_columns.clone(),
751 descending: self.view.sort_descending.clone(),
752 });
753 } else if !self.view.sort_ascending {
754 steps.push(Step::Reverse);
755 }
756 steps
757 }
758
759 pub fn python_steps(&self) -> Vec<Step> {
762 let mut steps = self.view_steps();
763 let in_order = self.view.column_order.iter().map(String::as_str).eq(self
764 .view
765 .schema
766 .iter_names()
767 .map(|s| s.as_str())
768 .filter(|s| *s != crate::formats::schema_union::DRIFT_COLUMN));
769 if !in_order {
770 steps.push(Step::Select(self.view.column_order.clone()));
771 }
772 steps
773 }
774
775 pub fn is_drilled_down(&self) -> bool {
776 self.view.drilled_down_group_index.is_some()
777 }
778
779 pub(super) fn apply_transformations(&mut self) {
782 let lf = match self.pushed_view() {
783 Some(view) => {
784 let sorted = !self.view.sort_columns.is_empty() || !self.view.sort_ascending;
785 self.view.unsorted_lf = sorted
786 .then(|| {
787 self.pushdown
788 .as_ref()
789 .and_then(|p| p.view(&self.view.filters, &[], false))
790 .map(|unsorted| unsorted.lf)
791 })
792 .flatten();
793 self.view.view_notes = Vec::new();
794 self.view.view_numbered = false;
795 self.with_column_changes(view.lf).0
796 }
797 None => self.build_view(),
798 };
799 self.invalidate_num_rows();
800 self.view.lf = lf;
801 self.restore_footer_count();
802 self.collect();
803 }
804
805 fn build_view(&mut self) -> LazyFrame {
808 let mut lf = self.with_column_changes(self.view.base_lf.clone()).0;
809 self.view.view_numbered = self.row_numbers && self.wants_view_numbers();
810 if self.view.view_numbered {
811 lf = lf.with_row_index(crate::formats::schema_union::DRIFT_COLUMN, None);
812 }
813 if let Some(e) = crate::export::python_script::filters_expr(&self.typed_filters()) {
814 lf = lf.filter(e);
815 }
816
817 let (excluded, view_notes) = self.leave_out_unread_rows(lf);
820 lf = excluded;
821 if !view_notes.is_empty() && view_notes != self.view.view_notes {
823 self.view.notes_seen = false;
824 }
825 self.view.view_notes = view_notes;
826
827 self.view.unsorted_lf =
830 (!self.view.sort_columns.is_empty() || !self.view.sort_ascending).then(|| lf.clone());
831 if !self.view.sort_columns.is_empty() {
832 lf = lf.sort_by_exprs(
833 self.view.sort_columns.iter().map(col).collect::<Vec<_>>(),
834 sort_options(self.view.sort_descending.clone()),
835 );
836 } else if !self.view.sort_ascending {
837 lf = lf.reverse();
838 }
839 lf
840 }
841
842 pub fn sort(&mut self, columns: Vec<String>, ascending: bool) {
845 let descending = vec![!ascending; columns.len()];
846 self.view.sort_ascending = ascending;
847 self.sort_by(columns, descending);
848 }
849
850 pub fn sort_by(&mut self, columns: Vec<String>, descending: Vec<bool>) {
852 debug_assert_eq!(columns.len(), descending.len());
853 if let Some(first) = descending.first() {
856 self.view.sort_ascending = !first;
857 }
858 if columns != self.view.sort_columns || descending != self.view.sort_descending {
861 self.widths.relearn();
862 }
863 self.view.sort_columns = columns;
864 self.view.sort_descending = descending;
865 self.drop_buffer();
866 self.apply_transformations();
867 }
868
869 pub fn reverse(&mut self) {
872 self.view.sort_ascending = !self.view.sort_ascending;
873 self.widths.relearn();
874 for direction in &mut self.view.sort_descending {
875 *direction = !*direction;
876 }
877 self.drop_buffer();
878 self.apply_transformations();
879 }
880
881 pub fn apply_view(
884 &mut self,
885 order: Vec<String>,
886 locked: usize,
887 filters: Vec<FilterStatement>,
888 columns: Vec<String>,
889 descending: Vec<bool>,
890 ) -> ViewChange {
891 self.set_column_order(order);
892 self.set_locked_columns(locked);
893 let change = ViewChange {
894 filters: filters != self.view.filters,
895 sort: columns != self.view.sort_columns || descending != self.view.sort_descending,
896 };
897 if change.filters || change.sort {
898 self.widths.relearn();
899 }
900 if let Some(first) = descending.first() {
901 self.view.sort_ascending = !first;
902 }
903 self.view.filters = filters;
904 self.view.sort_columns = columns;
905 self.view.sort_descending = descending;
906 self.view.start_row = 0;
907 self.drop_buffer();
908 self.apply_transformations();
909 change
910 }
911
912 pub fn filter(&mut self, filters: Vec<FilterStatement>) {
913 if filters != self.view.filters {
915 self.widths.relearn();
916 }
917 self.view.filters = filters;
918 self.view.start_row = 0;
921 self.drop_buffer();
922 self.apply_transformations();
923 }
924
925 pub fn query(&mut self, query: String) {
926 self.error = None;
927
928 let trimmed_query = query.trim();
929 if trimmed_query.is_empty() {
930 self.reset_lf_to_original();
931 self.collect();
932 return;
933 }
934
935 let source_schema = self.query_source_schema();
936 let parsed = parse_query_over(&query, Some(&source_schema))
937 .map(|parsed| parsed.past_calendar_safe(Some(&source_schema)));
938 match parsed {
939 Ok(ParsedQuery {
940 cols,
941 filter,
942 group_by: group_by_cols,
943 group_by_names: group_by_col_names,
944 distinct,
945 }) => {
946 let mut lf = self.query_source();
947 let mut schema_opt: Option<Arc<Schema>> = None;
948 let lineage = if cols.is_empty() && group_by_cols.is_empty() {
951 None
952 } else {
953 let mut kept = passed_through(&group_by_cols);
954 if cols.is_empty() {
955 kept.extend(
957 source_schema
958 .iter_names()
959 .filter(|n| !group_by_col_names.iter().any(|g| g == n.as_str()))
960 .map(|n| (n.to_string(), n.to_string())),
961 );
962 } else {
963 kept.extend(passed_through(&cols));
964 }
965 Some(Arc::new(kept))
966 };
967
968 if let Some(f) = filter {
969 lf = lf.filter(f);
970 }
971 let group_rows = lf.clone();
973
974 if !group_by_cols.is_empty() {
975 if !cols.is_empty() {
976 lf = lf.group_by(group_by_cols.clone()).agg(cols);
977 } else {
978 let agg_exprs: Vec<Expr> = source_schema
981 .iter_names()
982 .filter(|n| !group_by_col_names.iter().any(|g| g == n.as_str()))
983 .map(|n| col(n.clone()))
984 .collect();
985
986 lf = lf.group_by(group_by_cols.clone()).agg(agg_exprs);
987 }
988 let schema = match lf.collect_schema() {
991 Ok(s) => s,
992 Err(e) => {
993 self.error = Some(e);
994 return;
995 }
996 };
997 schema_opt = Some(schema.clone());
998 let sort_exprs: Vec<Expr> = schema
999 .iter_names()
1000 .take(group_by_cols.len())
1001 .map(|n| col(n.as_str()))
1002 .collect();
1003 let options = sort_options(vec![false; sort_exprs.len()]);
1004 lf = lf.sort_by_exprs(sort_exprs, options);
1005 } else if !cols.is_empty() {
1006 lf = lf.select(cols);
1007 }
1008 if distinct {
1009 lf = lf.unique_stable(None, UniqueKeepStrategy::First);
1011 }
1012
1013 let schema = match schema_opt {
1014 Some(s) => s,
1015 None => match lf.collect_schema() {
1016 Ok(s) => s,
1017 Err(e) => {
1018 self.error = Some(e);
1019 return;
1020 }
1021 },
1022 };
1023
1024 let locked = schema
1026 .iter_names()
1027 .take_while(|c| group_by_col_names.iter().any(|g| g.as_str() == c.as_str()))
1028 .count();
1029 let keys: Vec<(PlSmallStr, Expr)> =
1031 schema.iter_names().cloned().zip(group_by_cols).collect();
1032 let input = source_schema;
1034 let steps = vec![Step::Query {
1035 query: query.clone(),
1036 input: input.clone(),
1037 keys: keys.iter().map(|(name, _)| name.to_string()).collect(),
1038 }];
1039 let python_keys: Vec<Option<String>> = match crate::query::parse_nodes(&query) {
1041 Ok(mut nodes) => {
1042 nodes.resolve_division(&input);
1043 nodes
1044 .group_by
1045 .iter()
1046 .map(|key| Some(key.without_aliases().python()))
1047 .collect()
1048 }
1049 Err(_) => vec![None; keys.len()],
1050 };
1051 let python_rows = Some(vec![Step::QueryRows {
1052 query: query.clone(),
1053 input,
1054 }]);
1055 self.install_query_result(lf, schema, ActiveQuery::Dsl(query), locked, steps);
1056 self.view.lineage = lineage;
1057 if !keys.is_empty() {
1058 self.view.group_source = Some(GroupSource {
1059 rows: group_rows,
1060 keys,
1061 scratch: Vec::new(),
1062 rows_in_lists: true,
1063 python_rows,
1064 python_keys,
1065 lineage: None,
1067 });
1068 }
1069 self.forget_reshape();
1070 self.collect();
1071 if self.view.num_rows > 0 {
1073 self.view.start_row = 0;
1074 }
1075 }
1076 Err(e) => {
1077 self.error = Some(PolarsError::ComputeError(e.into()));
1079 }
1080 }
1081 }
1082
1083 pub(crate) fn query_root(&self) -> LazyFrame {
1086 if self.view.grouped.is_some() {
1087 return self.view.base_lf.clone();
1089 }
1090 Self::without_drift(
1091 self.view
1092 .reshaped_lf
1093 .clone()
1094 .unwrap_or_else(|| self.original_lf.clone()),
1095 )
1096 }
1097
1098 #[cfg(feature = "sql")]
1100 fn root_lineage(&self) -> Lineage {
1101 if self.view.grouped.is_some() {
1102 self.view.lineage.clone()
1103 } else if self.view.reshaped_lf.is_some() {
1104 self.view.reshape_lineage.clone()
1105 } else {
1106 None
1107 }
1108 }
1109
1110 #[cfg(feature = "sql")]
1112 fn query_root_steps(&self) -> Vec<Step> {
1113 if self.view.grouped.is_some() {
1114 return self.view.base_steps.clone();
1115 }
1116 match (&self.view.reshaped_lf, &self.view.reshape_steps) {
1117 (None, _) => Vec::new(),
1118 (Some(_), Some(steps)) => steps.clone(),
1119 (Some(_), None) => vec![Step::Unreproducible(
1120 "datui reshaped the data in a way it cannot write as Python".to_string(),
1121 )],
1122 }
1123 }
1124
1125 pub fn sql_query(&mut self, sql: String) {
1129 self.error = None;
1130 let trimmed = sql.trim();
1131 if trimmed.is_empty() {
1132 self.reset_lf_to_original();
1133 return;
1134 }
1135
1136 #[cfg(feature = "sql")]
1137 {
1138 use polars_sql::SQLContext;
1139 let mut ctx = SQLContext::new();
1140 let root = self.query_root();
1141 let root_steps = self.query_root_steps();
1142 ctx.register("df", root.clone());
1143 match ctx.execute(trimmed) {
1144 Ok(mut result_lf) => {
1145 crate::past_calendar::guard_plan(&mut result_lf.logical_plan);
1149 let order = ordered_by(&result_lf.logical_plan);
1152 let mut schema = match result_lf.clone().collect_schema() {
1153 Ok(s) => s,
1154 Err(e) => {
1155 self.error = Some(e);
1156 return;
1157 }
1158 };
1159 let leftover =
1160 leftover_subquery_value_columns(&mut result_lf.logical_plan, &schema);
1161 if !leftover.is_empty() {
1162 let shown = Arc::make_mut(&mut schema);
1163 for name in &leftover {
1164 shown.shift_remove(name);
1165 }
1166 result_lf = result_lf.drop(Selector::ByName {
1167 names: leftover.into(),
1168 strict: true,
1169 });
1170 }
1171 let root_lineage = self.root_lineage();
1172 let lineage = {
1173 let columns = root.clone().collect_schema().unwrap_or_default();
1174 let names: Vec<&str> = columns.iter_names().map(|n| n.as_str()).collect();
1175 let shown: Vec<&str> = schema.iter_names().map(|n| n.as_str()).collect();
1176 traced(
1177 &root_lineage,
1178 crate::query::sql_group::passed_through(trimmed, &names, &shown),
1179 )
1180 };
1181 let group_source = Self::sql_group_source(
1182 &mut ctx,
1183 trimmed,
1184 root,
1185 &root_steps,
1186 &mut result_lf,
1187 &schema,
1188 root_lineage,
1189 );
1190 if !group_source.as_ref().is_some_and(|(_, by_keys)| *by_keys) {
1193 stable_order(&mut result_lf.logical_plan);
1194 }
1195 count_subquery_values_once(&mut result_lf.logical_plan);
1196 let ordered_by = match &group_source {
1197 Some((source, true)) => schema
1198 .iter_names()
1199 .filter(|name| source.keys.iter().any(|(key, _)| key == *name))
1200 .map(|name| name.to_string())
1201 .collect(),
1202 _ => Vec::new(),
1203 };
1204 let mut steps = root_steps;
1205 steps.push(Step::Sql {
1206 sql: trimmed.to_string(),
1207 ordered_by,
1208 });
1209 let query_order = order
1210 .into_iter()
1211 .take_while(|(name, _)| schema.contains(name))
1212 .collect();
1213 self.install_query_result(result_lf, schema, ActiveQuery::Sql(sql), 0, steps);
1214 self.view.query_order = query_order;
1215 self.view.lineage = lineage;
1216 self.install_sql_group_source(group_source.map(|(source, _)| source));
1217 }
1218 Err(e) => {
1219 self.error = Some(e);
1220 }
1221 }
1222 }
1223
1224 #[cfg(not(feature = "sql"))]
1225 {
1226 self.error = Some(PolarsError::ComputeError(
1227 "SQL is not supported in this build. Rebuild with default features.".into(),
1228 ));
1229 }
1230 }
1231
1232 #[cfg(feature = "sql")]
1237 fn sql_group_source(
1238 ctx: &mut polars_sql::SQLContext,
1239 sql: &str,
1240 root: LazyFrame,
1241 root_steps: &[Step],
1242 result_lf: &mut LazyFrame,
1243 result: &Schema,
1244 lineage: Lineage,
1245 ) -> Option<(GroupSource, bool)> {
1246 use crate::query::sql_group::KeySource;
1247 let columns = root.clone().collect_schema().ok()?;
1248 let names: Vec<&str> = columns.iter_names().map(|n| n.as_str()).collect();
1249 let plan = crate::query::sql_group::plan(sql, &names, result.len())?;
1250 let mut rows = ctx.execute(&plan.source_sql).ok()?;
1251 crate::past_calendar::guard_plan(&mut rows.logical_plan);
1252 let source_schema = rows.clone().collect_schema().ok()?;
1253 let mut scratch = Vec::new();
1254 let mut keys = Vec::with_capacity(plan.keys.len());
1255 let mut python_keys = Vec::with_capacity(plan.keys.len());
1256 for key in plan.keys {
1257 let (name, dtype) = result.get_at_index(key.result_index)?;
1258 let column = match key.source {
1259 KeySource::Column(c) => PlSmallStr::from(c),
1260 KeySource::Computed(c) => {
1261 let c = PlSmallStr::from(c);
1262 scratch.push(c.clone());
1263 c
1264 }
1265 };
1266 if source_schema.get(&column) != Some(dtype) {
1268 return None;
1269 }
1270 python_keys.push(Some(format!(
1271 "pl.col({})",
1272 crate::export::python_script::py_str(&column)
1273 )));
1274 keys.push((name.clone(), col(column)));
1275 }
1276 if !plan.ordered {
1277 let by: Vec<Expr> = result
1279 .iter_names()
1280 .filter(|name| keys.iter().any(|(key, _)| key == *name))
1281 .map(|name| col(name.clone()))
1282 .collect();
1283 let options = sort_options(vec![false; by.len()]);
1284 *result_lf = result_lf.clone().sort_by_exprs(by, options);
1285 }
1286 let mut python_rows = root_steps.to_vec();
1287 python_rows.push(Step::Sql {
1288 sql: plan.source_sql.clone(),
1289 ordered_by: Vec::new(),
1290 });
1291 let source = GroupSource {
1292 rows,
1293 keys,
1294 scratch,
1295 rows_in_lists: false,
1296 python_rows: Some(python_rows),
1297 python_keys,
1298 lineage,
1300 };
1301 Some((source, !plan.ordered))
1302 }
1303
1304 #[cfg(feature = "sql")]
1307 fn install_sql_group_source(&mut self, source: Option<GroupSource>) {
1308 let Some(source) = source else {
1309 return;
1310 };
1311 self.view.locked_columns_count = self
1312 .view
1313 .schema
1314 .iter_names()
1315 .take_while(|c| source.keys.iter().any(|(k, _)| k == *c))
1316 .count();
1317 self.view.group_source = Some(source);
1318 }
1319
1320 pub fn fuzzy_search(&mut self, query: String) {
1323 self.error = None;
1324 let trimmed = query.trim();
1325 if trimmed.is_empty() {
1326 self.reset_lf_to_original();
1327 self.collect();
1328 return;
1329 }
1330 let schema = self.query_source_schema();
1333 let string_cols: Vec<String> = schema
1334 .iter()
1335 .filter(|(_, dtype)| dtype.is_string())
1336 .map(|(name, _)| name.to_string())
1337 .collect();
1338 if string_cols.is_empty() {
1339 self.error = Some(PolarsError::ComputeError(
1340 "A text match needs at least one text column".into(),
1341 ));
1342 return;
1343 }
1344 let tokens: Vec<&str> = trimmed
1345 .split_whitespace()
1346 .filter(|s| !s.is_empty())
1347 .collect();
1348 let token_exprs: Vec<Expr> = tokens
1349 .iter()
1350 .map(|token| {
1351 let pattern = fuzzy_token_regex(token);
1352 string_cols
1353 .iter()
1354 .map(|c| col(c.as_str()).str().contains(lit(pattern.as_str()), false))
1355 .reduce(|a, b| a.or(b))
1356 .unwrap()
1357 })
1358 .collect();
1359 let combined = token_exprs.into_iter().reduce(|a, b| a.and(b)).unwrap();
1360 let lf = self.query_source().filter(combined);
1361 let steps = vec![Step::Search {
1362 patterns: tokens.iter().map(|t| fuzzy_token_regex(t)).collect(),
1363 columns: string_cols.clone(),
1364 }];
1365 self.install_query_result(lf, schema, ActiveQuery::Fuzzy(query), 0, steps);
1366 self.view.lineage = None;
1368 self.forget_reshape();
1369 self.collect();
1370 }
1371}