alopex_dataframe/lazy/
lazyframe.rs

1use std::path::Path;
2
3use crate::lazy::{LogicalPlan, Optimizer, ProjectionKind};
4use crate::ops::{FillNull, JoinKeys, JoinType, SortOptions};
5use crate::{DataFrame, Expr, Result};
6
7/// A lazily-evaluated query backed by a `LogicalPlan`.
8#[derive(Debug, Clone)]
9pub struct LazyFrame {
10    plan: LogicalPlan,
11}
12
13impl LazyFrame {
14    /// Create a `LazyFrame` that scans an in-memory `DataFrame`.
15    pub fn from_dataframe(df: DataFrame) -> Self {
16        Self {
17            plan: LogicalPlan::DataFrameScan { df },
18        }
19    }
20
21    /// Build a CSV scan plan (no file I/O is performed until `collect()`).
22    pub fn scan_csv(path: impl AsRef<Path>) -> Result<Self> {
23        Ok(Self {
24            plan: LogicalPlan::CsvScan {
25                path: path.as_ref().to_path_buf(),
26                predicate: None,
27                projection: None,
28            },
29        })
30    }
31
32    /// Build a Parquet scan plan (no file I/O is performed until `collect()`).
33    pub fn scan_parquet(path: impl AsRef<Path>) -> Result<Self> {
34        Ok(Self {
35            plan: LogicalPlan::ParquetScan {
36                path: path.as_ref().to_path_buf(),
37                predicate: None,
38                projection: None,
39            },
40        })
41    }
42
43    /// Add a projection (`select`) node to the logical plan.
44    pub fn select(self, exprs: Vec<Expr>) -> Self {
45        Self {
46            plan: LogicalPlan::Projection {
47                input: Box::new(self.plan),
48                exprs,
49                kind: ProjectionKind::Select,
50            },
51        }
52    }
53
54    /// Add a filter node to the logical plan.
55    pub fn filter(self, predicate: Expr) -> Self {
56        Self {
57            plan: LogicalPlan::Filter {
58                input: Box::new(self.plan),
59                predicate,
60            },
61        }
62    }
63
64    /// Add a projection (`with_columns`) node to the logical plan.
65    pub fn with_columns(self, exprs: Vec<Expr>) -> Self {
66        Self {
67            plan: LogicalPlan::Projection {
68                input: Box::new(self.plan),
69                exprs,
70                kind: ProjectionKind::WithColumns,
71            },
72        }
73    }
74
75    /// Start a group-by on this `LazyFrame`.
76    pub fn group_by(self, by: Vec<Expr>) -> LazyGroupBy {
77        LazyGroupBy {
78            plan: self.plan,
79            by,
80        }
81    }
82
83    /// Join with another `LazyFrame` using provided join keys.
84    pub fn join<K: Into<JoinKeys>>(self, other: LazyFrame, keys: K, how: JoinType) -> Self {
85        let keys = keys.into();
86        Self {
87            plan: LogicalPlan::Join {
88                left: Box::new(self.plan),
89                right: Box::new(other.plan),
90                keys,
91                how,
92            },
93        }
94    }
95
96    /// Sort by one or more columns.
97    pub fn sort(self, mut options: SortOptions) -> Self {
98        options.nulls_last = true;
99        options.stable = true;
100        Self {
101            plan: LogicalPlan::Sort {
102                input: Box::new(self.plan),
103                options,
104            },
105        }
106    }
107
108    /// Return the first `n` rows.
109    pub fn head(self, n: usize) -> Self {
110        Self {
111            plan: LogicalPlan::Slice {
112                input: Box::new(self.plan),
113                offset: 0,
114                len: n,
115                from_end: false,
116            },
117        }
118    }
119
120    /// Return the last `n` rows.
121    pub fn tail(self, n: usize) -> Self {
122        Self {
123            plan: LogicalPlan::Slice {
124                input: Box::new(self.plan),
125                offset: 0,
126                len: n,
127                from_end: true,
128            },
129        }
130    }
131
132    /// Remove duplicate rows.
133    pub fn unique(self, subset: Option<Vec<String>>) -> Self {
134        Self {
135            plan: LogicalPlan::Unique {
136                input: Box::new(self.plan),
137                subset,
138            },
139        }
140    }
141
142    /// Fill null values using a scalar or strategy.
143    pub fn fill_null<T: Into<FillNull>>(self, fill: T) -> Self {
144        Self {
145            plan: LogicalPlan::FillNull {
146                input: Box::new(self.plan),
147                fill: fill.into(),
148            },
149        }
150    }
151
152    /// Drop rows containing null values.
153    pub fn drop_nulls(self, subset: Option<Vec<String>>) -> Self {
154        Self {
155            plan: LogicalPlan::DropNulls {
156                input: Box::new(self.plan),
157                subset,
158            },
159        }
160    }
161
162    /// Count null values per column.
163    pub fn null_count(self) -> Self {
164        Self {
165            plan: LogicalPlan::NullCount {
166                input: Box::new(self.plan),
167            },
168        }
169    }
170
171    /// Explode one `List<Utf8>` column into multiple rows.
172    pub fn explode(self, column: impl Into<String>) -> Self {
173        Self {
174            plan: LogicalPlan::Explode {
175                input: Box::new(self.plan),
176                column: column.into(),
177            },
178        }
179    }
180
181    /// Implode UTF-8 columns into one row of `List<Utf8>` columns.
182    pub fn implode(self) -> Self {
183        Self {
184            plan: LogicalPlan::Implode {
185                input: Box::new(self.plan),
186            },
187        }
188    }
189
190    /// Optimize, compile, and execute this `LazyFrame` into an eager `DataFrame`.
191    pub fn collect(self) -> Result<DataFrame> {
192        let optimized = Optimizer::optimize(&self.plan);
193        let physical = crate::physical::compile(&optimized)?;
194        let batches = crate::physical::Executor::execute(physical)?;
195        DataFrame::from_batches(batches)
196    }
197
198    /// Render the logical plan as a human-readable string.
199    ///
200    /// If `optimized` is `true`, includes optimizer rewrites such as pushdowns.
201    pub fn explain(self, optimized: bool) -> String {
202        if optimized {
203            Optimizer::optimize(&self.plan).display()
204        } else {
205            self.plan.display()
206        }
207    }
208}
209
210/// Group-by builder for `LazyFrame`.
211#[derive(Debug, Clone)]
212pub struct LazyGroupBy {
213    by: Vec<Expr>,
214    plan: LogicalPlan,
215}
216
217impl LazyGroupBy {
218    /// Add an aggregate node to the logical plan.
219    pub fn agg(self, aggs: Vec<Expr>) -> LazyFrame {
220        LazyFrame {
221            plan: LogicalPlan::Aggregate {
222                input: Box::new(self.plan),
223                group_by: self.by,
224                aggs,
225            },
226        }
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use std::sync::Arc;
233
234    use arrow::array::{ArrayRef, Int64Array};
235
236    use super::LazyFrame;
237    use crate::expr::{col, lit};
238    use crate::{DataFrame, Series};
239
240    fn df() -> DataFrame {
241        let a: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));
242        let b: ArrayRef = Arc::new(Int64Array::from(vec![10, 20, 30]));
243        DataFrame::new(vec![
244            Series::from_arrow("a", vec![a]).unwrap(),
245            Series::from_arrow("b", vec![b]).unwrap(),
246        ])
247        .unwrap()
248    }
249
250    #[test]
251    fn explain_builds_plan_without_io() {
252        let lf = LazyFrame::scan_csv("test.csv").unwrap();
253        let s = lf.explain(false);
254        assert!(s.contains("scan[csv"));
255    }
256
257    #[test]
258    fn collect_executes_filter_and_select_on_dataframe_scan() {
259        let lf = LazyFrame::from_dataframe(df())
260            .filter(col("a").gt(lit(1_i64)))
261            .select(vec![col("b").alias("bb")]);
262        let out = lf.collect().unwrap();
263        assert_eq!(out.height(), 2);
264        let bb = out.column("bb").unwrap();
265        assert_eq!(bb.len(), 2);
266    }
267
268    #[test]
269    fn group_by_agg_executes_sum_and_count() {
270        let lf = LazyFrame::from_dataframe(df())
271            .group_by(vec![col("a")])
272            .agg(vec![
273                col("b").sum().alias("sum_b"),
274                col("b").count().alias("cnt_b"),
275            ]);
276        let out = lf.collect().unwrap();
277        assert_eq!(out.width(), 3);
278        assert_eq!(out.column("a").unwrap().len(), 3);
279        assert_eq!(out.column("sum_b").unwrap().len(), 3);
280        assert_eq!(out.column("cnt_b").unwrap().len(), 3);
281    }
282}
alopex_dataframe/lazy/lazyframe.rs

alopex_dataframe/lazy/
lazyframe.rs
