polars-plan 0.54.1

Lazy query engine for the Polars DataFrame library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

use std::fmt;
use std::sync::Arc;

use polars_utils::pl_str::PlSmallStr;

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
#[derive(Debug, Clone, Hash, PartialEq)]
pub struct Sorted {
    pub column: PlSmallStr,
    /// None -> either way / unsure
    /// Some(false) -> ascending
    /// Some(true) -> descending
    pub descending: Option<bool>,
    /// None -> either way / unsure
    /// Some(false) -> nulls (if any) at start
    /// Some(true) -> nulls (if any) at end
    pub nulls_last: Option<bool>,
}

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
#[derive(Clone, Hash, strum_macros::IntoStaticStr)]
#[strum(serialize_all = "SCREAMING_SNAKE_CASE")]
pub enum HintIR {
    Sorted(Arc<[Sorted]>),
}

impl HintIR {
    /// Removes hints based on column name and filter function. Returns false if no hints were retained.
    pub fn retain_names<F>(&mut self, mut f: F) -> bool
    where
        F: FnMut(&str) -> bool,
    {
        match self {
            Self::Sorted(s) => {
                let Some(i) = s.iter().position(|s| f(&s.column)) else {
                    return false;
                };

                *s = s
                    .get(i)
                    .into_iter()
                    .chain(s.iter().skip(1 + i).filter(|s| f(&s.column)))
                    .cloned()
                    .collect()
            },
        }

        true
    }
}

impl fmt::Display for Sorted {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let descending = match self.descending {
            None => "?",
            Some(false) => "false",
            Some(true) => "true",
        };
        let nulls_last = match self.nulls_last {
            None => "?",
            Some(false) => "false",
            Some(true) => "true",
        };

        write!(
            f,
            "'{}': {{ descending: {descending}, nulls_last: {nulls_last} }}",
            self.column,
        )
    }
}

impl fmt::Debug for HintIR {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{self}")
    }
}

impl fmt::Display for HintIR {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            HintIR::Sorted(s) => {
                write!(f, "sorted(")?;
                if let Some(fst) = s.first() {
                    fst.fmt(f)?;
                    for si in &s[1..] {
                        f.write_str(", ")?;
                        si.fmt(f)?;
                    }
                }
                write!(f, ")")
            },
        }
    }
}