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::physical::budget::StreamOptions;
6use crate::physical::{BatchOpenContext, BudgetedMaterializedExecutor, DataFrameStream};
7use crate::{DataFrame, DataFrameError, Expr, Result};
8
9/// A lazily-evaluated query backed by a `LogicalPlan`.
10#[derive(Debug, Clone)]
11pub struct LazyFrame {
12    plan: LogicalPlan,
13}
14
15impl LazyFrame {
16    /// Create a `LazyFrame` that scans an in-memory `DataFrame`.
17    pub fn from_dataframe(df: DataFrame) -> Self {
18        Self {
19            plan: LogicalPlan::DataFrameScan { df },
20        }
21    }
22
23    /// Build a CSV scan plan (no file I/O is performed until `collect()`).
24    pub fn scan_csv(path: impl AsRef<Path>) -> Result<Self> {
25        Ok(Self {
26            plan: LogicalPlan::CsvScan {
27                path: path.as_ref().to_path_buf(),
28                predicate: None,
29                projection: None,
30            },
31        })
32    }
33
34    /// Build a Parquet scan plan (no file I/O is performed until `collect()`).
35    pub fn scan_parquet(path: impl AsRef<Path>) -> Result<Self> {
36        Ok(Self {
37            plan: LogicalPlan::ParquetScan {
38                path: path.as_ref().to_path_buf(),
39                predicate: None,
40                projection: None,
41            },
42        })
43    }
44
45    /// Construct strict vertical concatenation from two or more lazy inputs.
46    ///
47    /// When every schema is already known, mismatch is reported at plan construction. For lazy
48    /// file scans whose schema is intentionally unavailable until bounded open, the streaming
49    /// executor preflights every child before its first output and rejects any mismatch then.
50    pub fn concat(inputs: Vec<LazyFrame>) -> Result<Self> {
51        let mut plans = Vec::with_capacity(inputs.len());
52        let mut known_schemas = Vec::with_capacity(inputs.len());
53        for input in inputs {
54            known_schemas.push(known_plan_schema(&input.plan));
55            plans.push(input.plan);
56        }
57        if known_schemas.iter().all(Option::is_some) {
58            let plans = plans
59                .into_iter()
60                .zip(known_schemas)
61                .map(|(plan, schema)| (plan, schema.expect("all schemas were checked")))
62                .collect();
63            return Ok(Self {
64                plan: LogicalPlan::concat(plans)?,
65            });
66        }
67
68        let mut first_known: Option<arrow::datatypes::SchemaRef> = None;
69        for (index, schema) in known_schemas.iter().enumerate() {
70            let Some(schema) = schema else {
71                continue;
72            };
73            if let Some(first) = &first_known {
74                if schema.as_ref() != first.as_ref() {
75                    return Err(DataFrameError::schema_mismatch(format!(
76                        "concat_schema_mismatch: known input {index} differs from another known input"
77                    )));
78                }
79            } else {
80                first_known = Some(schema.clone());
81            }
82        }
83        Ok(Self {
84            plan: LogicalPlan::concat_deferred(plans)?,
85        })
86    }
87
88    /// Add a projection (`select`) node to the logical plan.
89    pub fn select(self, exprs: Vec<Expr>) -> Self {
90        Self {
91            plan: LogicalPlan::Projection {
92                input: Box::new(self.plan),
93                exprs,
94                kind: ProjectionKind::Select,
95            },
96        }
97    }
98
99    /// Add a filter node to the logical plan.
100    pub fn filter(self, predicate: Expr) -> Self {
101        Self {
102            plan: LogicalPlan::Filter {
103                input: Box::new(self.plan),
104                predicate,
105            },
106        }
107    }
108
109    /// Add a projection (`with_columns`) node to the logical plan.
110    pub fn with_columns(self, exprs: Vec<Expr>) -> Self {
111        Self {
112            plan: LogicalPlan::Projection {
113                input: Box::new(self.plan),
114                exprs,
115                kind: ProjectionKind::WithColumns,
116            },
117        }
118    }
119
120    /// Start a group-by on this `LazyFrame`.
121    pub fn group_by(self, by: Vec<Expr>) -> LazyGroupBy {
122        LazyGroupBy {
123            plan: self.plan,
124            by,
125        }
126    }
127
128    /// Join with another `LazyFrame` using provided join keys.
129    pub fn join<K: Into<JoinKeys>>(self, other: LazyFrame, keys: K, how: JoinType) -> Self {
130        let keys = keys.into();
131        Self {
132            plan: LogicalPlan::Join {
133                left: Box::new(self.plan),
134                right: Box::new(other.plan),
135                keys,
136                how,
137            },
138        }
139    }
140
141    /// Sort by one or more columns.
142    pub fn sort(self, mut options: SortOptions) -> Self {
143        options.nulls_last = true;
144        options.stable = true;
145        Self {
146            plan: LogicalPlan::Sort {
147                input: Box::new(self.plan),
148                options,
149            },
150        }
151    }
152
153    /// Return the first `n` rows.
154    pub fn head(self, n: usize) -> Self {
155        Self {
156            plan: LogicalPlan::Slice {
157                input: Box::new(self.plan),
158                offset: 0,
159                len: n,
160                from_end: false,
161            },
162        }
163    }
164
165    /// Return the last `n` rows.
166    pub fn tail(self, n: usize) -> Self {
167        Self {
168            plan: LogicalPlan::Slice {
169                input: Box::new(self.plan),
170                offset: 0,
171                len: n,
172                from_end: true,
173            },
174        }
175    }
176
177    /// Remove duplicate rows.
178    pub fn unique(self, subset: Option<Vec<String>>) -> Self {
179        Self {
180            plan: LogicalPlan::Unique {
181                input: Box::new(self.plan),
182                subset,
183            },
184        }
185    }
186
187    /// Fill null values using a scalar or strategy.
188    pub fn fill_null<T: Into<FillNull>>(self, fill: T) -> Self {
189        Self {
190            plan: LogicalPlan::FillNull {
191                input: Box::new(self.plan),
192                fill: fill.into(),
193            },
194        }
195    }
196
197    /// Drop rows containing null values.
198    pub fn drop_nulls(self, subset: Option<Vec<String>>) -> Self {
199        Self {
200            plan: LogicalPlan::DropNulls {
201                input: Box::new(self.plan),
202                subset,
203            },
204        }
205    }
206
207    /// Count null values per column.
208    pub fn null_count(self) -> Self {
209        Self {
210            plan: LogicalPlan::NullCount {
211                input: Box::new(self.plan),
212            },
213        }
214    }
215
216    /// Explode one `List<Utf8>` column into multiple rows.
217    pub fn explode(self, column: impl Into<String>) -> Self {
218        Self {
219            plan: LogicalPlan::Explode {
220                input: Box::new(self.plan),
221                column: column.into(),
222            },
223        }
224    }
225
226    /// Implode UTF-8 columns into one row of `List<Utf8>` columns.
227    pub fn implode(self) -> Self {
228        Self {
229            plan: LogicalPlan::Implode {
230                input: Box::new(self.plan),
231            },
232        }
233    }
234
235    /// Optimize, compile, and execute this `LazyFrame` into an eager `DataFrame`.
236    pub fn collect(self) -> Result<DataFrame> {
237        let optimized = Optimizer::optimize(&self.plan);
238        let physical = crate::physical::compile(&optimized)?;
239        let batches = crate::physical::Executor::execute(physical)?;
240        DataFrame::from_batches(batches)
241    }
242
243    /// Materialize a supported bounded stream while reserving every retained output batch.
244    ///
245    /// This deliberately uses the streaming source/operator path and never wraps the eager
246    /// executor with accounting after allocation.
247    pub fn collect_with_options(self, options: StreamOptions) -> Result<DataFrame> {
248        let mut stream = self.collect_streaming(options)?;
249        BudgetedMaterializedExecutor::new(stream.budget().clone()).collect_stream(&mut stream)
250    }
251
252    /// Compile this plan for incremental bounded execution.
253    ///
254    /// The compiler performs eligibility analysis before any source is opened. Until a source's
255    /// v0.8 `BatchSourceFactory` is installed, it returns `StreamingUnsupported` rather than
256    /// calling the legacy eager executor.
257    pub fn collect_streaming(self, options: StreamOptions) -> Result<DataFrameStream> {
258        let optimized = Optimizer::optimize(&self.plan);
259        let physical = crate::physical::compile(&optimized)?;
260        let compiled = crate::physical::compile_streaming(&physical, options)?;
261        let context = BatchOpenContext::new(options);
262        let source = crate::physical::open_streaming_plan(compiled, context.clone())?;
263        Ok(DataFrameStream::from_opened_source(source, context.budget))
264    }
265
266    /// Render the logical plan as a human-readable string.
267    ///
268    /// If `optimized` is `true`, includes optimizer rewrites such as pushdowns.
269    pub fn explain(self, optimized: bool) -> String {
270        if optimized {
271            Optimizer::optimize(&self.plan).display()
272        } else {
273            self.plan.display()
274        }
275    }
276}
277
278fn known_plan_schema(plan: &LogicalPlan) -> Option<arrow::datatypes::SchemaRef> {
279    match plan {
280        LogicalPlan::DataFrameScan { df } => Some(df.schema()),
281        LogicalPlan::Concat { schema, .. } => schema.clone(),
282        LogicalPlan::Filter { input, .. }
283        | LogicalPlan::Sort { input, .. }
284        | LogicalPlan::Slice { input, .. }
285        | LogicalPlan::Unique { input, .. }
286        | LogicalPlan::FillNull { input, .. }
287        | LogicalPlan::DropNulls { input, .. } => known_plan_schema(input),
288        LogicalPlan::CsvScan { .. }
289        | LogicalPlan::ParquetScan { .. }
290        | LogicalPlan::Projection { .. }
291        | LogicalPlan::Aggregate { .. }
292        | LogicalPlan::Join { .. }
293        | LogicalPlan::NullCount { .. }
294        | LogicalPlan::Explode { .. }
295        | LogicalPlan::Implode { .. } => None,
296    }
297}
298
299/// Group-by builder for `LazyFrame`.
300#[derive(Debug, Clone)]
301pub struct LazyGroupBy {
302    by: Vec<Expr>,
303    plan: LogicalPlan,
304}
305
306impl LazyGroupBy {
307    /// Add an aggregate node to the logical plan.
308    pub fn agg(self, aggs: Vec<Expr>) -> LazyFrame {
309        LazyFrame {
310            plan: LogicalPlan::Aggregate {
311                input: Box::new(self.plan),
312                group_by: self.by,
313                aggs,
314            },
315        }
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use std::num::NonZeroUsize;
322    use std::sync::Arc;
323
324    use arrow::array::{Array, ArrayRef, Int64Array, StringArray};
325
326    use super::LazyFrame;
327    use crate::expr::{col, concat_str, lit, ConcatStrNullBehavior};
328    use crate::physical::budget::StreamOptions;
329    use crate::{DataFrame, Series};
330
331    fn df() -> DataFrame {
332        let a: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));
333        let b: ArrayRef = Arc::new(Int64Array::from(vec![10, 20, 30]));
334        DataFrame::new(vec![
335            Series::from_arrow("a", vec![a]).unwrap(),
336            Series::from_arrow("b", vec![b]).unwrap(),
337        ])
338        .unwrap()
339    }
340
341    #[test]
342    fn explain_builds_plan_without_io() {
343        let lf = LazyFrame::scan_csv("test.csv").unwrap();
344        let s = lf.explain(false);
345        assert!(s.contains("scan[csv"));
346    }
347
348    #[test]
349    fn collect_executes_filter_and_select_on_dataframe_scan() {
350        let lf = LazyFrame::from_dataframe(df())
351            .filter(col("a").gt(lit(1_i64)))
352            .select(vec![col("b").alias("bb")]);
353        let out = lf.collect().unwrap();
354        assert_eq!(out.height(), 2);
355        let bb = out.column("bb").unwrap();
356        assert_eq!(bb.len(), 2);
357    }
358
359    #[test]
360    fn group_by_agg_executes_sum_and_count() {
361        let lf = LazyFrame::from_dataframe(df())
362            .group_by(vec![col("a")])
363            .agg(vec![
364                col("b").sum().alias("sum_b"),
365                col("b").count().alias("cnt_b"),
366            ]);
367        let out = lf.collect().unwrap();
368        assert_eq!(out.width(), 3);
369        assert_eq!(out.column("a").unwrap().len(), 3);
370        assert_eq!(out.column("sum_b").unwrap().len(), 3);
371        assert_eq!(out.column("cnt_b").unwrap().len(), 3);
372    }
373
374    #[test]
375    fn collect_streaming_opens_the_bounded_csv_source_without_eager_collect() {
376        let dir = tempfile::tempdir().unwrap();
377        let path = dir.path().join("stream.csv");
378        std::fs::write(&path, "a\n1\n2\n").unwrap();
379        let options = StreamOptions::new(
380            16 * 1024,
381            NonZeroUsize::new(1).unwrap(),
382            NonZeroUsize::new(1).unwrap(),
383        );
384
385        let mut stream = LazyFrame::scan_csv(&path)
386            .unwrap()
387            .collect_streaming(options)
388            .unwrap();
389        assert_eq!(stream.next_batch().unwrap().unwrap().height(), 1);
390        assert_eq!(stream.next_batch().unwrap().unwrap().height(), 1);
391        assert!(stream.next_batch().unwrap().is_none());
392    }
393
394    #[test]
395    fn collect_streaming_opens_the_bounded_parquet_source_without_eager_collect() {
396        let dir = tempfile::tempdir().unwrap();
397        let path = dir.path().join("stream.parquet");
398        let schema = Arc::new(arrow::datatypes::Schema::new(vec![
399            arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Int64, true),
400        ]));
401        let batch = arrow::record_batch::RecordBatch::try_new(
402            schema,
403            vec![Arc::new(Int64Array::from(vec![1_i64, 2])) as ArrayRef],
404        )
405        .unwrap();
406        crate::io::write_parquet(&path, &DataFrame::from_batches(vec![batch]).unwrap()).unwrap();
407        let options = StreamOptions::new(
408            64 * 1024,
409            NonZeroUsize::new(1).unwrap(),
410            NonZeroUsize::new(1).unwrap(),
411        );
412
413        let mut stream = LazyFrame::scan_parquet(&path)
414            .unwrap()
415            .collect_streaming(options)
416            .unwrap();
417        assert_eq!(stream.next_batch().unwrap().unwrap().height(), 1);
418        assert_eq!(stream.next_batch().unwrap().unwrap().height(), 1);
419        assert!(stream.next_batch().unwrap().is_none());
420    }
421
422    #[test]
423    fn lazy_concat_preflights_deferred_csv_schemas_then_preserves_input_order() {
424        let dir = tempfile::tempdir().unwrap();
425        let first = dir.path().join("first.csv");
426        let second = dir.path().join("second.csv");
427        std::fs::write(&first, "a\n1\n2\n").unwrap();
428        std::fs::write(&second, "a\n3\n").unwrap();
429        let options = StreamOptions::new(
430            16 * 1024,
431            NonZeroUsize::new(1).unwrap(),
432            NonZeroUsize::new(1).unwrap(),
433        );
434
435        let mut stream = LazyFrame::concat(vec![
436            LazyFrame::scan_csv(&first).unwrap(),
437            LazyFrame::scan_csv(&second).unwrap(),
438        ])
439        .unwrap()
440        .collect_streaming(options)
441        .unwrap();
442        let mut values = Vec::new();
443        while let Some(batch) = stream.next_batch().unwrap() {
444            let batch = batch.to_arrow().pop().unwrap();
445            let values_array = batch
446                .column(0)
447                .as_any()
448                .downcast_ref::<Int64Array>()
449                .unwrap();
450            values.extend((0..values_array.len()).map(|index| values_array.value(index)));
451        }
452        assert_eq!(values, vec![1, 2, 3]);
453    }
454
455    #[test]
456    fn deferred_concat_schema_mismatch_fails_before_a_stream_is_returned() {
457        let dir = tempfile::tempdir().unwrap();
458        let first = dir.path().join("first.csv");
459        let second = dir.path().join("second.csv");
460        std::fs::write(&first, "a\n1\n").unwrap();
461        std::fs::write(&second, "a\nnot-a-number\n").unwrap();
462        let options = StreamOptions::new(
463            16 * 1024,
464            NonZeroUsize::new(1).unwrap(),
465            NonZeroUsize::new(1).unwrap(),
466        );
467
468        let result = LazyFrame::concat(vec![
469            LazyFrame::scan_csv(&first).unwrap(),
470            LazyFrame::scan_csv(&second).unwrap(),
471        ])
472        .unwrap()
473        .collect_streaming(options);
474        let err = match result {
475            Ok(_) => panic!("mismatched deferred concat must fail before returning a stream"),
476            Err(error) => error,
477        };
478        assert!(err.to_string().contains("concat_schema_mismatch"));
479    }
480
481    #[test]
482    fn streaming_select_filter_concat_str_and_bounded_collect_share_batch_semantics() {
483        let dir = tempfile::tempdir().unwrap();
484        let path = dir.path().join("operators.csv");
485        std::fs::write(&path, "a,left,right\n1,x,y\n2,m,n\n3,p,q\n").unwrap();
486        // This test asserts streaming and bounded collection semantics rather
487        // than a particular reservation threshold. Dedicated budget tests
488        // exercise rejection at constrained limits.
489        let options = StreamOptions::default();
490        let repeated = col("a").add(lit(1_i64));
491        let label = concat_str(
492            vec![col("left"), col("right")],
493            "-",
494            ConcatStrNullBehavior::Propagate,
495        )
496        .unwrap()
497        .alias("label");
498        let plan = LazyFrame::scan_csv(&path)
499            .unwrap()
500            .filter(col("a").gt(lit(1_i64)))
501            .select(vec![
502                repeated.clone().alias("next"),
503                repeated.alias("next_again"),
504                label,
505            ]);
506
507        let mut stream = plan.clone().collect_streaming(options).unwrap();
508        let mut rows = Vec::new();
509        while let Some(dataframe) = stream.next_batch().unwrap() {
510            let batch = dataframe.to_arrow().pop().unwrap();
511            let next = batch
512                .column(0)
513                .as_any()
514                .downcast_ref::<Int64Array>()
515                .unwrap();
516            let repeated = batch
517                .column(1)
518                .as_any()
519                .downcast_ref::<Int64Array>()
520                .unwrap();
521            let label = batch
522                .column(2)
523                .as_any()
524                .downcast_ref::<StringArray>()
525                .unwrap();
526            for index in 0..batch.num_rows() {
527                rows.push((
528                    next.value(index),
529                    repeated.value(index),
530                    label.value(index).to_string(),
531                ));
532            }
533        }
534        assert_eq!(rows, vec![(3, 3, "m-n".into()), (4, 4, "p-q".into())]);
535
536        let materialized = plan.collect_with_options(options).unwrap();
537        assert_eq!(materialized.height(), 2);
538        assert_eq!(materialized.schema().fields()[2].name(), "label");
539    }
540
541    #[test]
542    fn streaming_with_columns_and_forward_head_preserve_schema_and_order() {
543        let dir = tempfile::tempdir().unwrap();
544        let path = dir.path().join("with-columns.csv");
545        std::fs::write(&path, "a\n1\n2\n3\n").unwrap();
546        let options = StreamOptions::new(
547            16 * 1024,
548            NonZeroUsize::new(1).unwrap(),
549            NonZeroUsize::new(1).unwrap(),
550        );
551
552        let mut stream = LazyFrame::scan_csv(&path)
553            .unwrap()
554            .with_columns(vec![col("a").add(lit(10_i64)).alias("next")])
555            .head(2)
556            .collect_streaming(options)
557            .unwrap();
558        let mut values = Vec::new();
559        while let Some(dataframe) = stream.next_batch().unwrap() {
560            let batch = dataframe.to_arrow().pop().unwrap();
561            assert_eq!(batch.schema().fields()[1].name(), "next");
562            let next = batch
563                .column(1)
564                .as_any()
565                .downcast_ref::<Int64Array>()
566                .unwrap();
567            values.extend((0..next.len()).map(|index| next.value(index)));
568        }
569        assert_eq!(values, vec![11, 12]);
570    }
571
572    #[test]
573    fn known_lazy_concat_schema_mismatch_fails_during_plan_construction() {
574        let first = df().lazy();
575        let incompatible = DataFrame::new(vec![Series::from_arrow(
576            "other",
577            vec![Arc::new(Int64Array::from(vec![1])) as ArrayRef],
578        )
579        .unwrap()])
580        .unwrap()
581        .lazy();
582
583        let err = LazyFrame::concat(vec![first, incompatible]).unwrap_err();
584        assert!(err.to_string().contains("concat_schema_mismatch"));
585    }
586}
alopex_dataframe/lazy/lazyframe.rs

alopex_dataframe/lazy/
lazyframe.rs
