Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
 Path: blob/main/crates/polars-parquet/src/arrow/read/expr.rs
6940 views
1
use std::sync::Arc;

2


3
use arrow::array::Array;

4
use arrow::bitmap::{Bitmap, BitmapBuilder};

5
use arrow::types::AlignedBytes;

6


7
#[derive(Clone)]

8
pub enum ParquetScalar {

9
    Null,

10


11
    Boolean(bool),

12


13
    Int8(i8),

14
    Int16(i16),

15
    Int32(i32),

16
    Int64(i64),

17
    UInt8(u8),

18
    UInt16(u16),

19
    UInt32(u32),

20
    UInt64(u64),

21


22
    Float32(f32),

23
    Float64(f64),

24


25
    FixedSizeBinary(Box<[u8]>),

26


27
    String(Box<str>),

28
    Binary(Box<[u8]>),

29
}

30


31
impl ParquetScalar {

32
    pub(crate) fn is_null(&self) -> bool {

33
        matches!(self, Self::Null)

34
    }

35


36
    pub(crate) fn to_aligned_bytes<B: AlignedBytes>(&self) -> Option<B> {

37
        match self {

38
            Self::Int8(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

39
                .ok()

40
                .map(B::from_unaligned),

41
            Self::Int16(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

42
                .ok()

43
                .map(B::from_unaligned),

44
            Self::Int32(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

45
                .ok()

46
                .map(B::from_unaligned),

47
            Self::Int64(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

48
                .ok()

49
                .map(B::from_unaligned),

50
            Self::UInt8(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

51
                .ok()

52
                .map(B::from_unaligned),

53
            Self::UInt16(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

54
                .ok()

55
                .map(B::from_unaligned),

56
            Self::UInt32(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

57
                .ok()

58
                .map(B::from_unaligned),

59
            Self::UInt64(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

60
                .ok()

61
                .map(B::from_unaligned),

62
            Self::Float32(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

63
                .ok()

64
                .map(B::from_unaligned),

65
            Self::Float64(v) => <B::Unaligned>::try_from(&v.to_le_bytes())

66
                .ok()

67
                .map(B::from_unaligned),

68
            _ => None,

69
        }

70
    }

71


72
    pub(crate) fn as_str(&self) -> Option<&str> {

73
        match self {

74
            Self::String(s) => Some(s.as_ref()),

75
            _ => None,

76
        }

77
    }

78


79
    pub(crate) fn as_binary(&self) -> Option<&[u8]> {

80
        match self {

81
            Self::Binary(s) => Some(s.as_ref()),

82
            _ => None,

83
        }

84
    }

85
}

86


87
#[derive(Clone)]

88
pub enum SpecializedParquetColumnExpr {

89
    Equal(ParquetScalar),

90


91
    Between(ParquetScalar, ParquetScalar),

92


93
    EqualOneOf(Box<[ParquetScalar]>),

94


95
    StartsWith(Box<[u8]>),

96
    EndsWith(Box<[u8]>),

97
    StartEndsWith(Box<[u8]>, Box<[u8]>),

98
}

99


100
pub type ParquetColumnExprRef = Arc<dyn ParquetColumnExpr>;

101
pub trait ParquetColumnExpr: Send + Sync {

102
    fn evaluate(&self, values: &dyn Array) -> Bitmap {

103
        let mut bm = BitmapBuilder::new();

104
        self.evaluate_mut(values, &mut bm);

105
        bm.freeze()

106
    }

107
    fn evaluate_mut(&self, values: &dyn Array, bm: &mut BitmapBuilder);

108
    fn evaluate_null(&self) -> bool;

109


110
    fn as_specialized(&self) -> Option<&SpecializedParquetColumnExpr>;

111
}

112


113