polars_lazy/frame/
cached_arenas.rs

1use super::*;
2
3pub(crate) struct CachedArena {
4    lp_arena: Arena<IR>,
5    expr_arena: Arena<AExpr>,
6}
7
8impl LazyFrame {
9    pub fn set_cached_arena(&self, lp_arena: Arena<IR>, expr_arena: Arena<AExpr>) {
10        let mut cached = self.cached_arena.lock().unwrap();
11        *cached = Some(CachedArena {
12            lp_arena,
13            expr_arena,
14        });
15    }
16
17    pub fn schema_with_arenas(
18        &mut self,
19        lp_arena: &mut Arena<IR>,
20        expr_arena: &mut Arena<AExpr>,
21    ) -> PolarsResult<SchemaRef> {
22        let node = to_alp(
23            self.logical_plan.clone(),
24            expr_arena,
25            lp_arena,
26            &mut OptFlags::schema_only(),
27        )?;
28
29        let schema = lp_arena.get(node).schema(lp_arena).into_owned();
30        // Cache the logical plan so that next schema call is cheap.
31        self.logical_plan = DslPlan::IR {
32            node: Some(node),
33            dsl: Arc::new(self.logical_plan.clone()),
34            version: lp_arena.version(),
35        };
36        Ok(schema)
37    }
38
39    /// Get a handle to the schema — a map from column names to data types — of the current
40    /// `LazyFrame` computation.
41    ///
42    /// Returns an `Err` if the logical plan has already encountered an error (i.e., if
43    /// `self.collect()` would fail), `Ok` otherwise.
44    pub fn collect_schema(&mut self) -> PolarsResult<SchemaRef> {
45        let mut cached_arenas = self.cached_arena.lock().unwrap();
46
47        match &mut *cached_arenas {
48            None => {
49                let mut lp_arena = Default::default();
50                let mut expr_arena = Default::default();
51                // Code duplication because of bchk. :(
52                let node = to_alp(
53                    self.logical_plan.clone(),
54                    &mut expr_arena,
55                    &mut lp_arena,
56                    &mut OptFlags::schema_only(),
57                )?;
58
59                let schema = lp_arena.get(node).schema(&lp_arena).into_owned();
60                // Cache the logical plan so that next schema call is cheap.
61                self.logical_plan = DslPlan::IR {
62                    node: Some(node),
63                    dsl: Arc::new(self.logical_plan.clone()),
64                    version: lp_arena.version(),
65                };
66                *cached_arenas = Some(CachedArena {
67                    lp_arena,
68                    expr_arena,
69                });
70
71                Ok(schema)
72            },
73            Some(arenas) => {
74                match self.logical_plan {
75                    // We have got arenas and don't need to convert the DSL.
76                    DslPlan::IR {
77                        node: Some(node), ..
78                    } => Ok(arenas
79                        .lp_arena
80                        .get(node)
81                        .schema(&arenas.lp_arena)
82                        .into_owned()),
83                    _ => {
84                        // We have got arenas, but still need to convert (parts) of the DSL.
85                        // Code duplication because of bchk. :(
86                        let node = to_alp(
87                            self.logical_plan.clone(),
88                            &mut arenas.expr_arena,
89                            &mut arenas.lp_arena,
90                            &mut OptFlags::schema_only(),
91                        )?;
92
93                        let schema = arenas
94                            .lp_arena
95                            .get(node)
96                            .schema(&arenas.lp_arena)
97                            .into_owned();
98                        // Cache the logical plan so that next schema call is cheap.
99                        self.logical_plan = DslPlan::IR {
100                            node: Some(node),
101                            dsl: Arc::new(self.logical_plan.clone()),
102                            version: arenas.lp_arena.version(),
103                        };
104                        Ok(schema)
105                    },
106                }
107            },
108        }
109    }
110
111    pub(super) fn get_arenas(&mut self) -> (Arena<IR>, Arena<AExpr>) {
112        match self.cached_arena.lock().unwrap().as_mut() {
113            Some(arenas) => (arenas.lp_arena.clone(), arenas.expr_arena.clone()),
114            None => (Arena::with_capacity(16), Arena::with_capacity(16)),
115        }
116    }
117}
polars_lazy/frame/cached_arenas.rs

polars_lazy/frame/
cached_arenas.rs
