vortex-datafusion 0.72.0

Apache Datafusion integration for Vortex
Documentation 
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use datafusion_common::ColumnStatistics;
use datafusion_common::stats::Precision;
use vortex::array::stats::StatsSet;
use vortex::dtype::DType;
use vortex::dtype::Nullability;
use vortex::dtype::PType;
use vortex::error::VortexExpect;
use vortex::error::VortexResult;
use vortex::expr::stats::Precision as VortexPrecision;
use vortex::expr::stats::Stat;
use vortex::scalar::Scalar;

use crate::PrecisionExt;
use crate::convert::TryToDataFusion;

/// Convert a stats set for an array with the given dtype.
pub(crate) fn stats_set_to_df(
    stats_set: &StatsSet,
    dtype: &DType,
) -> VortexResult<ColumnStatistics> {
    // Update the total size in bytes.
    let column_size = stats_set.get_as::<usize>(Stat::UncompressedSizeInBytes, &PType::U64.into());

    // TODO(connor): There's a lot that can go wrong here, should probably handle this
    // more gracefully...
    // Find the min statistic.
    let min = stats_set.get(Stat::Min).and_then(|pstat_val| {
        pstat_val
            .map(|stat_val| {
                Scalar::try_new(
                    Stat::Min
                        .dtype(dtype)
                        .vortex_expect("must have a valid dtype"),
                    Some(stat_val),
                )
                .vortex_expect("`Stat::Min` somehow had an incompatible `DType`")
                .try_to_df()
                .ok()
            })
            .transpose()
    });

    // Find the max statistic.
    let max = stats_set.get(Stat::Max).and_then(|pstat_val| {
        pstat_val
            .map(|stat_val| {
                Scalar::try_new(
                    Stat::Max
                        .dtype(dtype)
                        .vortex_expect("must have a valid dtype"),
                    Some(stat_val),
                )
                .vortex_expect("`Stat::Max` somehow had an incompatible `DType`")
                .try_to_df()
                .ok()
            })
            .transpose()
    });

    // Find the sum statistic
    let sum = stats_set.get(Stat::Sum).and_then(|pstat_val| {
        pstat_val
            .map(|stat_val| {
                Scalar::try_new(
                    Stat::Sum
                        .dtype(dtype)
                        .vortex_expect("must have a valid dtype"),
                    Some(stat_val),
                )
                .vortex_expect("`Stat::Sum` somehow had an incompatible `DType`")
                .try_to_df()
                .ok()
            })
            .transpose()
    });

    let null_count = stats_set.get_as::<usize>(Stat::NullCount, &PType::U64.into());

    Ok(ColumnStatistics {
        null_count: null_count.to_df(),
        min_value: min.to_df(),
        max_value: max.to_df(),
        sum_value: sum.to_df(),
        distinct_count: is_constant_to_distinct_count(
            stats_set.get_as::<bool>(Stat::IsConstant, &DType::Bool(Nullability::NonNullable)),
        ),
        byte_size: column_size.to_df(),
    })
}

pub(crate) fn is_constant_to_distinct_count(
    is_constant: Option<VortexPrecision<bool>>,
) -> Precision<usize> {
    match is_constant.and_then(VortexPrecision::as_exact) {
        Some(true) => Precision::Exact(1),
        Some(false) | None => Precision::Absent,
    }
}

#[cfg(test)]
mod tests {
    use vortex::expr::stats::Precision as VortexPrecision;

    use super::*;

    #[test]
    fn is_constant_false_does_not_imply_one_distinct_value() -> VortexResult<()> {
        let false_constant = StatsSet::of(Stat::IsConstant, VortexPrecision::exact(false));
        let false_stats = stats_set_to_df(&false_constant, &DType::Bool(Nullability::NonNullable))?;

        assert_eq!(false_stats.distinct_count, Precision::Absent);

        let true_constant = StatsSet::of(Stat::IsConstant, VortexPrecision::exact(true));
        let true_stats = stats_set_to_df(&true_constant, &DType::Bool(Nullability::NonNullable))?;

        assert_eq!(true_stats.distinct_count, Precision::Exact(1));

        Ok(())
    }
}