alopex_dataframe/lazy/
logical_plan.rs

1use std::path::PathBuf;
2
3use crate::ops::{FillNull, JoinKeys, JoinType, SortOptions};
4use crate::{DataFrame, Expr};
5
6/// How a projection node should be interpreted.
7#[derive(Debug, Clone)]
8pub enum ProjectionKind {
9    /// Select columns/expressions, producing a new schema.
10    Select,
11    /// Add or overwrite columns, preserving existing columns.
12    WithColumns,
13}
14
15/// Logical query plan nodes for `LazyFrame`.
16#[derive(Debug, Clone)]
17pub enum LogicalPlan {
18    /// Scan an in-memory `DataFrame`.
19    DataFrameScan { df: DataFrame },
20    /// Scan a CSV file (predicate/projection may be pushed down).
21    CsvScan {
22        path: PathBuf,
23        predicate: Option<Expr>,
24        projection: Option<Vec<String>>,
25    },
26    /// Scan a Parquet file (predicate/projection may be pushed down).
27    ParquetScan {
28        path: PathBuf,
29        predicate: Option<Expr>,
30        projection: Option<Vec<String>>,
31    },
32    /// Projection node (select or with_columns).
33    Projection {
34        input: Box<LogicalPlan>,
35        exprs: Vec<Expr>,
36        kind: ProjectionKind,
37    },
38    /// Filter node.
39    Filter {
40        input: Box<LogicalPlan>,
41        predicate: Expr,
42    },
43    /// Aggregate node (group keys and aggregations).
44    Aggregate {
45        input: Box<LogicalPlan>,
46        group_by: Vec<Expr>,
47        aggs: Vec<Expr>,
48    },
49    /// Join two inputs.
50    Join {
51        left: Box<LogicalPlan>,
52        right: Box<LogicalPlan>,
53        keys: JoinKeys,
54        how: JoinType,
55    },
56    /// Sort input rows.
57    Sort {
58        input: Box<LogicalPlan>,
59        options: SortOptions,
60    },
61    /// Slice rows (used for head/tail).
62    Slice {
63        input: Box<LogicalPlan>,
64        offset: usize,
65        len: usize,
66        from_end: bool,
67    },
68    /// Remove duplicate rows.
69    Unique {
70        input: Box<LogicalPlan>,
71        subset: Option<Vec<String>>,
72    },
73    /// Fill nulls using a scalar or strategy.
74    FillNull {
75        input: Box<LogicalPlan>,
76        fill: FillNull,
77    },
78    /// Drop rows containing nulls.
79    DropNulls {
80        input: Box<LogicalPlan>,
81        subset: Option<Vec<String>>,
82    },
83    /// Count nulls per column.
84    NullCount { input: Box<LogicalPlan> },
85    /// Explode one list column.
86    Explode {
87        input: Box<LogicalPlan>,
88        column: String,
89    },
90    /// Implode columns into one row of list columns.
91    Implode { input: Box<LogicalPlan> },
92}
93
94impl LogicalPlan {
95    /// Render this plan as a readable string (used by `explain()` and tests).
96    pub fn display(&self) -> String {
97        let mut out = String::new();
98        self.fmt_into(&mut out, 0);
99        out
100    }
101
102    fn fmt_into(&self, out: &mut String, indent: usize) {
103        let pad = "  ".repeat(indent);
104        match self {
105            LogicalPlan::DataFrameScan { .. } => {
106                out.push_str(&format!("{pad}scan[dataframe]\n"));
107            }
108            LogicalPlan::CsvScan {
109                path,
110                predicate,
111                projection,
112            } => {
113                out.push_str(&format!("{pad}scan[csv path='{}']", path.display()));
114                if let Some(projection) = projection {
115                    out.push_str(&format!(" projection={:?}", projection));
116                }
117                if let Some(predicate) = predicate {
118                    out.push_str(&format!(" filters=[{}]", fmt_expr(predicate)));
119                }
120                out.push('\n');
121            }
122            LogicalPlan::ParquetScan {
123                path,
124                predicate,
125                projection,
126            } => {
127                out.push_str(&format!("{pad}scan[parquet path='{}']", path.display()));
128                if let Some(projection) = projection {
129                    out.push_str(&format!(" projection={:?}", projection));
130                }
131                if let Some(predicate) = predicate {
132                    out.push_str(&format!(" filters=[{}]", fmt_expr(predicate)));
133                }
134                out.push('\n');
135            }
136            LogicalPlan::Projection { input, exprs, kind } => {
137                let label = match kind {
138                    ProjectionKind::Select => "project",
139                    ProjectionKind::WithColumns => "with_columns",
140                };
141                out.push_str(&format!(
142                    "{pad}{label} [{}]\n",
143                    exprs.iter().map(fmt_expr).collect::<Vec<_>>().join(", ")
144                ));
145                input.fmt_into(out, indent + 1);
146            }
147            LogicalPlan::Filter { input, predicate } => {
148                out.push_str(&format!("{pad}filter [{}]\n", fmt_expr(predicate)));
149                input.fmt_into(out, indent + 1);
150            }
151            LogicalPlan::Aggregate {
152                input,
153                group_by,
154                aggs,
155            } => {
156                out.push_str(&format!(
157                    "{pad}aggregate by=[{}] aggs=[{}]\n",
158                    group_by.iter().map(fmt_expr).collect::<Vec<_>>().join(", "),
159                    aggs.iter().map(fmt_expr).collect::<Vec<_>>().join(", ")
160                ));
161                input.fmt_into(out, indent + 1);
162            }
163            LogicalPlan::Join {
164                left,
165                right,
166                keys,
167                how,
168            } => {
169                out.push_str(&format!(
170                    "{pad}join how={how:?} keys={}\n",
171                    fmt_join_keys(keys)
172                ));
173                left.fmt_into(out, indent + 1);
174                right.fmt_into(out, indent + 1);
175            }
176            LogicalPlan::Sort { input, options } => {
177                out.push_str(&format!(
178                    "{pad}sort by={:?} desc={:?} nulls_last={} stable={}\n",
179                    options.by, options.descending, options.nulls_last, options.stable
180                ));
181                input.fmt_into(out, indent + 1);
182            }
183            LogicalPlan::Slice {
184                input,
185                offset,
186                len,
187                from_end,
188            } => {
189                out.push_str(&format!(
190                    "{pad}slice offset={offset} len={len} from_end={from_end}\n"
191                ));
192                input.fmt_into(out, indent + 1);
193            }
194            LogicalPlan::Unique { input, subset } => {
195                out.push_str(&format!("{pad}unique subset={subset:?}\n"));
196                input.fmt_into(out, indent + 1);
197            }
198            LogicalPlan::FillNull { input, fill } => {
199                out.push_str(&format!("{pad}fill_null {}\n", fmt_fill_null(fill)));
200                input.fmt_into(out, indent + 1);
201            }
202            LogicalPlan::DropNulls { input, subset } => {
203                out.push_str(&format!("{pad}drop_nulls subset={subset:?}\n"));
204                input.fmt_into(out, indent + 1);
205            }
206            LogicalPlan::NullCount { input } => {
207                out.push_str(&format!("{pad}null_count\n"));
208                input.fmt_into(out, indent + 1);
209            }
210            LogicalPlan::Explode { input, column } => {
211                out.push_str(&format!("{pad}explode column={column}\n"));
212                input.fmt_into(out, indent + 1);
213            }
214            LogicalPlan::Implode { input } => {
215                out.push_str(&format!("{pad}implode\n"));
216                input.fmt_into(out, indent + 1);
217            }
218        }
219    }
220}
221
222fn fmt_join_keys(keys: &JoinKeys) -> String {
223    match keys {
224        JoinKeys::On(cols) => format!("on={cols:?}"),
225        JoinKeys::LeftRight { left_on, right_on } => {
226            format!("left_on={left_on:?} right_on={right_on:?}")
227        }
228    }
229}
230
231fn fmt_fill_null(fill: &FillNull) -> String {
232    match fill {
233        FillNull::Value(value) => format!("value={value:?}"),
234        FillNull::Strategy(strategy) => format!("strategy={strategy:?}"),
235    }
236}
237
238fn fmt_expr(expr: &Expr) -> String {
239    use crate::expr::{
240        AggFunc, DatetimeFunction, Expr as E, ExprFunction, ListFunction, Operator, Scalar,
241        StringFunction, UnaryOperator,
242    };
243
244    match expr {
245        E::Column(name) => format!("col({name})"),
246        E::Literal(Scalar::Null) => "lit(null)".to_string(),
247        E::Literal(Scalar::Boolean(v)) => format!("lit({v})"),
248        E::Literal(Scalar::Int64(v)) => format!("lit({v})"),
249        E::Literal(Scalar::Float64(v)) => format!("lit({v})"),
250        E::Literal(Scalar::Utf8(v)) => format!("lit({v:?})"),
251        E::Wildcard => "*".to_string(),
252        E::Alias { expr, name } => format!("{} as {name}", fmt_expr(expr)),
253        E::UnaryOp {
254            op: UnaryOperator::Not,
255            expr,
256        } => format!("not({})", fmt_expr(expr)),
257        E::BinaryOp { left, op, right } => {
258            let op_s = match op {
259                Operator::Add => "+",
260                Operator::Sub => "-",
261                Operator::Mul => "*",
262                Operator::Div => "/",
263                Operator::Eq => "==",
264                Operator::Neq => "!=",
265                Operator::Gt => ">",
266                Operator::Lt => "<",
267                Operator::Ge => ">=",
268                Operator::Le => "<=",
269                Operator::And => "and",
270                Operator::Or => "or",
271            };
272            format!("({} {op_s} {})", fmt_expr(left), fmt_expr(right))
273        }
274        E::Agg { func, expr } => {
275            let f = match func {
276                AggFunc::Sum => "sum",
277                AggFunc::Mean => "mean",
278                AggFunc::Count => "count",
279                AggFunc::Min => "min",
280                AggFunc::Max => "max",
281            };
282            format!("{f}({})", fmt_expr(expr))
283        }
284        E::Function { input, function } => {
285            let f = match function {
286                ExprFunction::String(function) => match function {
287                    StringFunction::ToLowercase => "str.to_lowercase".to_string(),
288                    StringFunction::ToUppercase => "str.to_uppercase".to_string(),
289                    StringFunction::Contains { pattern } => {
290                        format!("str.contains({pattern:?})")
291                    }
292                    StringFunction::Replace {
293                        pattern,
294                        replacement,
295                    } => format!("str.replace({pattern:?}, {replacement:?})"),
296                    StringFunction::StripChars { chars } => {
297                        format!("str.strip_chars({chars:?})")
298                    }
299                    StringFunction::Split { separator } => {
300                        format!("str.split({separator:?})")
301                    }
302                    StringFunction::LenChars => "str.len_chars".to_string(),
303                    StringFunction::Extract {
304                        pattern,
305                        capture_group,
306                    } => format!("str.extract({pattern:?}, {capture_group})"),
307                },
308                ExprFunction::Datetime(function) => match function {
309                    DatetimeFunction::Year => "dt.year".to_string(),
310                    DatetimeFunction::Month => "dt.month".to_string(),
311                    DatetimeFunction::Day => "dt.day".to_string(),
312                    DatetimeFunction::Weekday => "dt.weekday".to_string(),
313                    DatetimeFunction::ToString => "dt.to_string".to_string(),
314                    DatetimeFunction::ConvertTimeZone {
315                        from_offset,
316                        to_offset,
317                    } => format!("dt.convert_time_zone({from_offset:?}, {to_offset:?})"),
318                },
319                ExprFunction::List(function) => match function {
320                    ListFunction::Join {
321                        separator,
322                        null_value,
323                    } => format!("list.join({separator:?}, {null_value:?})"),
324                    ListFunction::Len => "list.len".to_string(),
325                    ListFunction::Contains { value } => {
326                        format!("list.contains({value:?})")
327                    }
328                },
329            };
330            format!("{f}({})", fmt_expr(input))
331        }
332    }
333}
334
335#[cfg(test)]
336mod tests {
337    use super::{LogicalPlan, ProjectionKind};
338    use crate::expr::{col, lit};
339
340    #[test]
341    fn display_is_readable_and_stable() {
342        let plan = LogicalPlan::Filter {
343            input: Box::new(LogicalPlan::Projection {
344                input: Box::new(LogicalPlan::CsvScan {
345                    path: "data.csv".into(),
346                    predicate: None,
347                    projection: Some(vec!["a".to_string(), "b".to_string()]),
348                }),
349                exprs: vec![col("a"), col("b").alias("bb")],
350                kind: ProjectionKind::Select,
351            }),
352            predicate: col("a").gt(lit(1_i64)),
353        };
354
355        let s = plan.display();
356        assert!(s.contains("scan[csv"));
357        assert!(s.contains("project"));
358        assert!(s.contains("filter"));
359        assert!(s.contains("col(a)"));
360    }
361}
alopex_dataframe/lazy/logical_plan.rs

alopex_dataframe/lazy/
logical_plan.rs
