proof_of_sql_planner/postprocessing/
select_postprocessing.rs

1use super::{evaluate_expr, PostprocessingResult, PostprocessingStep};
2use ahash::AHasher;
3use alloc::vec::Vec;
4use core::hash::BuildHasherDefault;
5use datafusion::logical_expr::Expr;
6use indexmap::IndexMap;
7use proof_of_sql::base::{
8    database::{OwnedColumn, OwnedTable},
9    scalar::Scalar,
10};
11use sqlparser::ast::Ident;
12
13/// The select expression used to select, reorder, and apply alias transformations
14#[derive(Debug, Clone, PartialEq)]
15pub struct SelectPostprocessing {
16    /// The expressions we select
17    exprs: Vec<Expr>,
18}
19
20impl SelectPostprocessing {
21    /// Create a new `SelectPostprocessing` node.
22    #[must_use]
23    pub fn new(exprs: Vec<Expr>) -> Self {
24        Self { exprs }
25    }
26}
27
28impl<S: Scalar> PostprocessingStep<S> for SelectPostprocessing {
29    /// Apply the select transformation to the given `OwnedTable`.
30    fn apply(&self, owned_table: OwnedTable<S>) -> PostprocessingResult<OwnedTable<S>> {
31        let cols: IndexMap<Ident, OwnedColumn<S>, BuildHasherDefault<AHasher>> = self
32            .exprs
33            .iter()
34            .map(|expr| -> PostprocessingResult<(Ident, OwnedColumn<S>)> {
35                let result_column = evaluate_expr(&owned_table, expr)?;
36                Ok((expr.display_name()?.as_str().into(), result_column))
37            })
38            .collect::<PostprocessingResult<_>>()?;
39        Ok(OwnedTable::try_new(cols)?)
40    }
41}