1use super::*;
5
6pub struct ViewRollback {
12 root_generation: u64,
14 counted: Option<CountedRows>,
16 table_state: TableState,
17 termcol_index: usize,
18 view: View,
19}
20
21impl ViewRollback {
22 pub fn count_landed(
26 &mut self,
27 len_generation: u64,
28 rows: usize,
29 file_row_groups: Option<&[Vec<usize>]>,
30 ) -> bool {
31 let ours = len_generation == self.view.len_generation;
32 if ours {
33 self.counted = Some(CountedRows {
34 rows,
35 file_row_groups: file_row_groups.map(<[_]>::to_vec),
36 });
37 }
38 ours
39 }
40}
41
42pub(super) struct CountedRows {
45 rows: usize,
46 file_row_groups: Option<Vec<Vec<usize>>>,
47}
48
49impl DataTableState {
50 pub fn get_filters(&self) -> &[FilterStatement] {
54 match &self.view.grouped {
55 Some(view) => &view.filters,
56 None => &self.view.filters,
57 }
58 }
59
60 pub fn get_sort_columns(&self) -> &[String] {
61 match &self.view.grouped {
62 Some(view) => &view.sort_columns,
63 None => &self.view.sort_columns,
64 }
65 }
66
67 pub fn get_sort_ascending(&self) -> bool {
68 match &self.view.grouped {
69 Some(view) => view.sort_ascending,
70 None => self.view.sort_ascending,
71 }
72 }
73
74 pub fn get_sort_descending(&self) -> &[bool] {
75 match &self.view.grouped {
76 Some(view) => &view.sort_descending,
77 None => &self.view.sort_descending,
78 }
79 }
80
81 pub fn view_filters(&self) -> &[FilterStatement] {
84 &self.view.filters
85 }
86
87 pub fn view_sort_columns(&self) -> &[String] {
88 &self.view.sort_columns
89 }
90
91 pub fn view_sort_ascending(&self) -> bool {
92 self.view.sort_ascending
93 }
94
95 pub fn view_sort_descending(&self) -> &[bool] {
96 &self.view.sort_descending
97 }
98
99 pub(crate) fn header_sort(&self) -> (Vec<String>, Vec<bool>) {
102 if self.view.sort_columns.is_empty() && self.view.grouped.is_none() {
103 self.view.query_order.iter().cloned().unzip()
104 } else {
105 (
106 self.view.sort_columns.clone(),
107 self.view.sort_descending.clone(),
108 )
109 }
110 }
111
112 #[cfg(test)]
114 pub(crate) fn reshaped_lf_clone(&self) -> Option<LazyFrame> {
115 self.view.reshaped_lf.clone()
116 }
117
118 pub fn get_column_order(&self) -> &[String] {
119 &self.view.column_order
120 }
121
122 pub fn is_at_defaults(&self) -> bool {
126 self.sampled.is_none()
127 && self.view.column_changes.is_empty()
128 && self.view.active_query.is_empty()
129 && self.view.active_sql_query.is_empty()
130 && self.view.active_fuzzy_query.is_empty()
131 && self.view.filters.is_empty()
132 && self.view.sort_columns.is_empty()
133 && self.view.last_pivot_spec.is_none()
134 && self.view.last_melt_spec.is_none()
135 && self.locked_columns_count() == 0
136 && self.view.column_order.iter().map(String::as_str).eq(self
137 .view
138 .schema
139 .iter_names()
140 .map(|s| s.as_str()))
141 }
142
143 pub fn get_active_query(&self) -> &str {
144 &self.view.active_query
145 }
146
147 pub fn get_active_sql_query(&self) -> &str {
148 &self.view.active_sql_query
149 }
150
151 pub fn check_plan(&self) -> PolarsResult<()> {
155 self.view
156 .lf
157 .clone()
158 .select(self.binary_stub_exprs())
159 .collect_schema()
160 .map(|_| ())
161 }
162
163 pub fn rollback_point(&self) -> ViewRollback {
165 ViewRollback {
166 root_generation: self.root_generation,
167 counted: None,
168 table_state: self.table_state,
169 termcol_index: self.termcol_index,
170 view: self.view.clone(),
171 }
172 }
173
174 pub fn roll_back(&mut self, saved: ViewRollback) {
179 if saved.root_generation != self.root_generation {
180 self.return_to_root();
181 return;
182 }
183 self.widths.keep_learned();
184 self.view = saved.view;
185 self.table_state = saved.table_state;
186 self.termcol_index = saved.termcol_index;
187 self.clear_column_moves();
188 self.reveal_cursor = true;
189 self.error = None;
190 if let Some(counted) = saved.counted {
192 self.take_count(counted.rows, counted.file_row_groups.as_deref());
193 }
194 }
195
196 pub fn try_transition<T, E>(
200 &mut self,
201 steps: impl FnOnce(&mut Self) -> std::result::Result<T, E>,
202 ) -> std::result::Result<(T, ViewRollback), E> {
203 let saved = self.rollback_point();
204 self.error = None;
205 match self.deferred(steps) {
206 Ok(value) => Ok((value, saved)),
207 Err(e) => {
208 self.roll_back(saved);
209 Err(e)
210 }
211 }
212 }
213
214 pub fn deferred<R>(&mut self, steps: impl FnOnce(&mut Self) -> R) -> R {
217 let deferred = std::mem::replace(&mut self.defer_collect, true);
218 let result = steps(self);
219 self.defer_collect = deferred;
220 result
221 }
222
223 pub fn count_landed(
226 &mut self,
227 len_generation: u64,
228 rows: usize,
229 file_row_groups: Option<&[Vec<usize>]>,
230 ) -> bool {
231 let current = len_generation == self.view.len_generation;
232 if current {
233 self.take_count(rows, file_row_groups);
234 }
235 current
236 }
237
238 #[cfg(test)]
241 pub(crate) fn set_provisional_rows(&mut self, n: usize) {
242 self.view.num_rows = n;
243 }
244
245 fn take_count(&mut self, rows: usize, file_row_groups: Option<&[Vec<usize>]>) {
248 match file_row_groups {
249 Some(groups) => self.record_file_row_groups(groups),
250 None => self.set_num_rows(rows),
251 }
252 }
253
254 pub fn follow(&self) -> Option<&crate::loading::follow::Follow> {
256 self.follow.as_ref()
257 }
258
259 pub fn follow_mut(&mut self) -> Option<&mut crate::loading::follow::Follow> {
260 self.follow.as_mut()
261 }
262
263 pub(crate) fn join_followed_fields(
267 &mut self,
268 fields: &[Field],
269 ) -> std::result::Result<bool, ()> {
270 if !self.scan_is_the_root() {
271 return Err(());
272 }
273 let (Some(follow), Some(format)) = (self.follow.as_ref(), self.read_as) else {
274 return Ok(false);
275 };
276 let (path, rows) = (follow.path().to_path_buf(), follow.shown());
277 let Some(mut lf) =
278 crate::loading::follow::widen(&self.original_lf, &path, format, fields, rows)
279 else {
280 return Ok(false);
281 };
282 let Ok(schema) = lf.collect_schema() else {
283 return Ok(false);
284 };
285 let known: std::collections::HashSet<&str> =
286 self.view.column_order.iter().map(String::as_str).collect();
287 let joining: Vec<String> = schema
288 .iter_names()
289 .map(|name| name.to_string())
290 .filter(|name| !known.contains(name.as_str()))
291 .collect();
292 drop(known);
293 self.view.column_order.extend(joining);
294 self.replace_root(lf, schema);
295 if self.is_pristine() {
296 self.set_num_rows(rows);
298 }
299 self.deferred(Self::apply_transformations);
302 Ok(true)
303 }
304
305 pub fn start_following(&mut self, follow: crate::loading::follow::Follow) {
307 self.follow = Some(follow);
308 }
309
310 pub fn stop_following(&mut self) {
312 if let Some(mut follow) = self.follow.take() {
313 follow.end();
314 }
315 }
316
317 pub(crate) fn aim_at_end(&mut self) {
320 if self.view.num_rows_valid && self.visible_rows > 0 {
321 self.view.start_row = self.view.num_rows.saturating_sub(self.visible_rows);
322 }
323 }
324
325 pub fn on_last_row(&self) -> bool {
327 self.view.num_rows_valid
328 && (self.view.num_rows == 0
329 || self.view.start_row + self.table_state.selected().unwrap_or(0) + 1
330 >= self.view.num_rows)
331 }
332
333 pub(super) fn each_frame(&mut self, mut f: impl FnMut(&mut LazyFrame)) {
335 f(&mut self.original_lf);
336 f(&mut self.view.base_lf);
337 f(&mut self.view.lf);
338 if let Some(lf) = self.view.unsorted_lf.as_mut() {
339 f(lf);
340 }
341 if let Some(lf) = self.view.reshaped_lf.as_mut() {
342 f(lf);
343 }
344 if let Some(source) = self.view.group_source.as_mut() {
345 f(&mut source.rows);
346 }
347 if let Some(grouped) = self.view.grouped.as_mut() {
348 f(&mut grouped.lf);
349 f(&mut grouped.base_lf);
350 if let Some(source) = grouped.group_source.as_mut() {
351 f(&mut source.rows);
352 }
353 }
354 }
355
356 pub(crate) fn follow_to(&mut self, rows: usize, restarted: bool) -> bool {
360 let Some(path) = self.follow.as_ref().map(|f| f.path().to_path_buf()) else {
361 return true;
362 };
363 let rows_stand = !restarted
364 && self.view.sort_columns.is_empty()
365 && self.view.sort_ascending
366 && self.scan_is_the_root();
367 let known = self.known_before_follow(&path, restarted);
368 self.each_frame(|lf| crate::loading::follow::bound(lf, &path, rows));
369 self.invalidate_num_rows();
370 self.follow_known = known.map(|known| (self.view.len_generation, known));
371 if self.is_pristine() {
372 self.set_num_rows(rows);
374 } else if self.scan_is_the_root() {
375 self.pristine_rows = Some(rows);
376 }
377 if restarted {
378 self.view.start_row = 0;
379 self.table_state.select(Some(0));
380 }
381 if !rows_stand {
382 self.drop_buffer();
383 }
384 rows_stand
385 }
386
387 fn known_before_follow(&mut self, path: &Path, restarted: bool) -> Option<Vec<(usize, usize)>> {
391 if restarted || self.is_pristine() || !self.scan_is_the_root() {
392 return None;
393 }
394 let mut known = self
395 .follow_known
396 .take()
397 .filter(|(generation, _)| *generation == self.view.len_generation)
398 .map(|(_, known)| known);
399 if self.view.num_rows_valid
400 && let Some(row) = crate::loading::follow::bound_of(&self.view.lf, path)
401 {
402 let known = known.get_or_insert_with(Vec::new);
403 if let [.., before, last] = known.as_slice()
405 && last.1 - before.1 < crate::loading::follow::MARK_ROWS as usize
406 {
407 known.pop();
408 }
409 if known.last().is_none_or(|&(_, at)| at < row) {
410 known.push((self.view.num_rows, row));
411 }
412 }
413 known
414 }
415
416 pub(crate) fn read_followed_through(&mut self, file: &std::fs::File) {
419 let Some(path) = self.follow.as_ref().map(|f| f.path().to_path_buf()) else {
420 return;
421 };
422 self.each_frame(|lf| crate::loading::follow::read_through(lf, &path, file));
423 }
424
425 pub fn lf(&self) -> &LazyFrame {
428 &self.view.lf
429 }
430
431 pub fn schema(&self) -> &Arc<Schema> {
433 &self.view.schema
434 }
435
436 pub fn num_rows(&self) -> usize {
439 self.view.num_rows
440 }
441
442 pub fn error(&self) -> Option<&PolarsError> {
444 self.error.as_ref()
445 }
446
447 pub fn dismiss_error(&mut self) {
449 self.error = None;
450 }
451
452 pub fn start_row(&self) -> usize {
454 self.view.start_row
455 }
456
457 pub fn partition_columns(&self) -> Option<&[String]> {
459 self.partition_columns.as_deref()
460 }
461
462 pub fn polars_streaming(&self) -> bool {
464 self.polars_streaming
465 }
466}