vortex_array/compute/
is_constant.rs1use vortex_error::{VortexExpect as _, VortexResult, vortex_bail};
2
3use crate::arrays::{ConstantArray, NullArray};
4use crate::stats::{Precision, Stat, StatsProviderExt};
5use crate::{Array, ArrayExt, Encoding};
6
7pub trait IsConstantFn<A> {
8    fn is_constant(&self, array: A, opts: &IsConstantOpts) -> VortexResult<Option<bool>>;
15}
16
17impl<E: Encoding> IsConstantFn<&dyn Array> for E
18where
19    E: for<'a> IsConstantFn<&'a E::Array>,
20{
21    fn is_constant(&self, array: &dyn Array, opts: &IsConstantOpts) -> VortexResult<Option<bool>> {
22        let array_ref = array
23            .as_any()
24            .downcast_ref::<E::Array>()
25            .vortex_expect("Failed to downcast array");
26        IsConstantFn::is_constant(self, array_ref, opts)
27    }
28}
29
30#[derive(Clone)]
32pub struct IsConstantOpts {
33    pub canonicalize: bool,
35}
36
37impl Default for IsConstantOpts {
38    fn default() -> Self {
39        Self { canonicalize: true }
40    }
41}
42
43pub fn is_constant(array: &dyn Array) -> VortexResult<bool> {
54    let opts = IsConstantOpts::default();
55    is_constant_opts(array, &opts)
56}
57
58pub fn is_constant_opts(array: &dyn Array, opts: &IsConstantOpts) -> VortexResult<bool> {
62    if let Some(Precision::Exact(value)) = array.statistics().get_as::<bool>(Stat::IsConstant) {
64        return Ok(value);
65    }
66
67    let is_constant = is_constant_impl(array, opts)?;
68
69    if let Some(is_constant) = is_constant {
70        array
71            .statistics()
72            .set(Stat::IsConstant, Precision::Exact(is_constant.into()));
73    }
74
75    Ok(is_constant.unwrap_or_default())
76}
77
78fn is_constant_impl(array: &dyn Array, opts: &IsConstantOpts) -> VortexResult<Option<bool>> {
79    match array.len() {
80        0 => return Ok(Some(false)),
82        1 => return Ok(Some(true)),
84        _ => {}
85    }
86
87    if array.as_opt::<ConstantArray>().is_some() || array.as_opt::<NullArray>().is_some() {
89        return Ok(Some(true));
90    }
91
92    let all_invalid = array.all_invalid()?;
93    if all_invalid {
94        return Ok(Some(true));
95    }
96
97    let all_valid = array.all_valid()?;
98
99    if !all_valid && !all_invalid {
101        return Ok(Some(false));
102    }
103
104    let min = array
106        .statistics()
107        .get_scalar(Stat::Min, array.dtype())
108        .and_then(|p| p.as_exact());
109    let max = array
110        .statistics()
111        .get_scalar(Stat::Max, array.dtype())
112        .and_then(|p| p.as_exact());
113
114    if let Some((min, max)) = min.zip(max) {
115        if min == max {
116            return Ok(Some(true));
117        }
118    }
119
120    assert!(
121        all_valid,
122        "All values must be valid as an invariant of the VTable."
123    );
124    let is_constant = if let Some(vtable_fn) = array.vtable().is_constant_fn() {
125        vtable_fn.is_constant(array, opts)?
126    } else {
127        log::debug!(
128            "No is_constant implementation found for {}",
129            array.encoding()
130        );
131
132        if opts.canonicalize {
133            let array = array.to_canonical()?;
134
135            if let Some(is_constant_fn) = array.as_ref().vtable().is_constant_fn() {
136                is_constant_fn.is_constant(array.as_ref(), opts)?
137            } else {
138                vortex_bail!(
139                    "No is_constant function for canonical array: {}",
140                    array.as_ref().encoding(),
141                )
142            }
143        } else {
144            None
145        }
146    };
147
148    Ok(is_constant)
149}