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

polars_lazy/frame/
cached_arenas.rs
