Skip to content

Commit b4204cb

Browse files
authored
Only check for ConstantVTable (#6002)
Signed-off-by: Nicholas Gates <[email protected]>
1 parent c83de4a commit b4204cb

File tree

19 files changed

+90
-82
lines changed

19 files changed

+90
-82
lines changed

encodings/datetime-parts/src/compute/is_constant.rs

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
use vortex_array::compute::IsConstantKernel;
55
use vortex_array::compute::IsConstantKernelAdapter;
66
use vortex_array::compute::IsConstantOpts;
7+
use vortex_array::compute::is_constant_opts;
78
use vortex_array::register_kernel;
89
use vortex_error::VortexResult;
910

@@ -14,13 +15,30 @@ impl IsConstantKernel for DateTimePartsVTable {
1415
fn is_constant(
1516
&self,
1617
array: &DateTimePartsArray,
17-
_opts: &IsConstantOpts,
18+
opts: &IsConstantOpts,
1819
) -> VortexResult<Option<bool>> {
19-
Ok(Some(
20-
array.days().is_constant()
21-
&& array.seconds().is_constant()
22-
&& array.subseconds().is_constant(),
23-
))
20+
let Some(days) = is_constant_opts(array.days().as_ref(), opts)? else {
21+
return Ok(None);
22+
};
23+
if !days {
24+
return Ok(Some(false));
25+
}
26+
27+
let Some(seconds) = is_constant_opts(array.seconds().as_ref(), opts)? else {
28+
return Ok(None);
29+
};
30+
if !seconds {
31+
return Ok(Some(false));
32+
}
33+
34+
let Some(subseconds) = is_constant_opts(array.subseconds().as_ref(), opts)? else {
35+
return Ok(None);
36+
};
37+
if !subseconds {
38+
return Ok(Some(false));
39+
}
40+
41+
Ok(Some(true))
2442
}
2543
}
2644

encodings/fastlanes/src/bitpacking/compute/between.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
use vortex_array::Array;
55
use vortex_array::ArrayRef;
66
use vortex_array::IntoArray;
7+
use vortex_array::arrays::ConstantVTable;
78
use vortex_array::compute::BetweenKernel;
89
use vortex_array::compute::BetweenKernelAdapter;
910
use vortex_array::compute::BetweenOptions;
@@ -22,7 +23,7 @@ impl BetweenKernel for BitPackedVTable {
2223
upper: &dyn Array,
2324
options: &BetweenOptions,
2425
) -> VortexResult<Option<ArrayRef>> {
25-
if !lower.is_constant() || !upper.is_constant() {
26+
if !lower.is::<ConstantVTable>() || !upper.is::<ConstantVTable>() {
2627
return Ok(None);
2728
};
2829

encodings/fastlanes/src/for/array/for_compress.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ fn compress_primitive<T: NativePType + WrappingSub + PrimInt>(
5252
#[cfg(test)]
5353
mod test {
5454
use itertools::Itertools;
55+
use vortex_array::Array;
5556
use vortex_array::ToCanonical;
5657
use vortex_array::assert_arrays_eq;
5758
use vortex_array::expr::stats::StatsProvider;
@@ -104,8 +105,8 @@ mod test {
104105
assert!(compressed.reference_scalar().dtype().is_signed_int());
105106
assert!(compressed.encoded().dtype().is_signed_int());
106107

107-
let constant = compressed.encoded().as_constant().unwrap();
108-
assert_eq!(constant, Scalar::from(0i32));
108+
let encoded = compressed.encoded().scalar_at(0);
109+
assert_eq!(encoded, Scalar::from(0i32));
109110
}
110111

111112
#[test]

encodings/runend/src/ops.rs

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@ mod tests {
4646
use vortex_array::arrays::PrimitiveArray;
4747
use vortex_array::assert_arrays_eq;
4848
use vortex_array::compute::Cost;
49+
use vortex_array::compute::IsConstantOpts;
50+
use vortex_array::compute::is_constant_opts;
4951
use vortex_buffer::buffer;
5052
use vortex_dtype::DType;
5153
use vortex_dtype::Nullability;
@@ -131,7 +133,16 @@ mod tests {
131133

132134
let sliced_array = re_array.slice(2..5);
133135

134-
assert!(sliced_array.is_constant_opts(Cost::Canonicalize))
136+
assert!(
137+
is_constant_opts(
138+
&sliced_array,
139+
&IsConstantOpts {
140+
cost: Cost::Canonicalize
141+
}
142+
)
143+
.unwrap()
144+
.unwrap_or_default()
145+
)
135146
}
136147

137148
#[test]

vortex-array/src/array/mod.rs

Lines changed: 1 addition & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,8 @@ use crate::arrays::VarBinVTable;
4545
use crate::arrays::VarBinViewVTable;
4646
use crate::builders::ArrayBuilder;
4747
use crate::compute::ComputeFn;
48-
use crate::compute::Cost;
4948
use crate::compute::InvocationArgs;
50-
use crate::compute::IsConstantOpts;
5149
use crate::compute::Output;
52-
use crate::compute::is_constant_opts;
5350
use crate::expr::stats::Precision;
5451
use crate::expr::stats::Stat;
5552
use crate::expr::stats::StatsProviderExt;
@@ -368,28 +365,8 @@ impl dyn Array + '_ {
368365
self.as_opt::<V>().is_some()
369366
}
370367

371-
pub fn is_constant(&self) -> bool {
372-
let opts = IsConstantOpts {
373-
cost: Cost::Specialized,
374-
};
375-
is_constant_opts(self, &opts)
376-
.inspect_err(|e| tracing::warn!("Failed to compute IsConstant: {e}"))
377-
.ok()
378-
.flatten()
379-
.unwrap_or_default()
380-
}
381-
382-
pub fn is_constant_opts(&self, cost: Cost) -> bool {
383-
let opts = IsConstantOpts { cost };
384-
is_constant_opts(self, &opts)
385-
.inspect_err(|e| tracing::warn!("Failed to compute IsConstant: {e}"))
386-
.ok()
387-
.flatten()
388-
.unwrap_or_default()
389-
}
390-
391368
pub fn as_constant(&self) -> Option<Scalar> {
392-
self.is_constant().then(|| self.scalar_at(0))
369+
self.as_opt::<ConstantVTable>().map(|a| a.scalar().clone())
393370
}
394371

395372
/// Total size of the array in bytes, including all children and buffers.

vortex-array/src/arrays/bool/compute/is_constant.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ mod tests {
2929
use rstest::rstest;
3030

3131
use super::*;
32+
use crate::compute::is_constant;
3233

3334
#[rstest]
3435
#[case(vec![true], true)]
@@ -40,6 +41,6 @@ mod tests {
4041
}, false)]
4142
fn test_is_constant(#[case] input: Vec<bool>, #[case] expected: bool) {
4243
let array = BoolArray::from_iter(input);
43-
assert_eq!(array.is_constant(), expected);
44+
assert_eq!(is_constant(array.as_ref()).unwrap(), Some(expected));
4445
}
4546
}

vortex-array/src/arrays/constant/compute/boolean.rs

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33

44
use vortex_dtype::DType;
55
use vortex_error::VortexResult;
6-
use vortex_error::vortex_bail;
76
use vortex_error::vortex_err;
87
use vortex_scalar::Scalar;
98

@@ -26,17 +25,15 @@ impl BooleanKernel for ConstantVTable {
2625
) -> VortexResult<Option<ArrayRef>> {
2726
// We only implement this for constant <-> constant arrays, otherwise we allow fall back
2827
// to the Arrow implementation.
29-
if !rhs.is_constant() {
28+
let Some(rhs) = rhs.as_opt::<ConstantVTable>() else {
3029
return Ok(None);
31-
}
30+
};
3231

3332
let length = lhs.len();
3433
let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
3534
let lhs = lhs.scalar().as_bool().value();
36-
let Some(rhs) = rhs.as_constant() else {
37-
vortex_bail!("Binary boolean operation requires both sides to be constant");
38-
};
3935
let rhs = rhs
36+
.scalar()
4037
.as_bool_opt()
4138
.ok_or_else(|| vortex_err!("expected rhs to be boolean"))?
4239
.value();

vortex-array/src/arrays/constant/vtable/encode.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use vortex_error::VortexResult;
66
use crate::Canonical;
77
use crate::arrays::ConstantArray;
88
use crate::arrays::ConstantVTable;
9+
use crate::compute::is_constant;
910
use crate::vtable::EncodeVTable;
1011

1112
impl EncodeVTable<ConstantVTable> for ConstantVTable {
@@ -15,7 +16,7 @@ impl EncodeVTable<ConstantVTable> for ConstantVTable {
1516
_like: Option<&ConstantArray>,
1617
) -> VortexResult<Option<ConstantArray>> {
1718
let canonical = canonical.as_ref();
18-
if canonical.is_constant() {
19+
if is_constant(canonical)?.unwrap_or_default() {
1920
let scalar = canonical.scalar_at(0);
2021
Ok(Some(ConstantArray::new(scalar, canonical.len())))
2122
} else {

vortex-array/src/arrays/dict/tests.rs

Lines changed: 0 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

44
use vortex_buffer::buffer;
5-
use vortex_scalar::Scalar;
65

76
use super::DictArray;
87
use crate::Array;
@@ -14,25 +13,6 @@ use crate::arrays::VarBinArray;
1413
use crate::assert_arrays_eq;
1514
use crate::validity::Validity;
1615

17-
#[test]
18-
fn test_slice_into_const_dict() {
19-
let dict = DictArray::try_new(
20-
PrimitiveArray::from_option_iter(vec![Some(0u32), None, Some(1)]).to_array(),
21-
PrimitiveArray::from_option_iter(vec![Some(0i32), Some(1), Some(2)]).to_array(),
22-
)
23-
.unwrap();
24-
25-
assert_eq!(
26-
Some(Scalar::new(dict.dtype().clone(), 0i32.into())),
27-
dict.slice(0..1).as_constant()
28-
);
29-
30-
assert_eq!(
31-
Some(Scalar::null(dict.dtype().clone())),
32-
dict.slice(1..2).as_constant()
33-
);
34-
}
35-
3616
#[test]
3717
fn test_scalar_at_null_code() {
3818
let dict = DictArray::try_new(

vortex-array/src/arrays/list/compute/is_constant.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use crate::arrays::ListVTable;
99
use crate::compute::IsConstantKernel;
1010
use crate::compute::IsConstantKernelAdapter;
1111
use crate::compute::IsConstantOpts;
12+
use crate::compute::is_constant;
1213
use crate::compute::numeric;
1314
use crate::register_kernel;
1415

@@ -47,7 +48,7 @@ impl IsConstantKernel for ListVTable {
4748
.slice(SMALL_ARRAY_THRESHOLD + 1..array.len() + 1);
4849
let list_lengths = numeric(&end_offsets, &start_offsets, NumericOperator::Sub)?;
4950

50-
if !list_lengths.is_constant() {
51+
if !is_constant(&list_lengths)?.unwrap_or_default() {
5152
return Ok(Some(false));
5253
}
5354
}
@@ -107,7 +108,11 @@ mod tests {
107108
.unwrap()
108109
.unwrap()
109110
);
110-
assert!(struct_of_lists.is_constant());
111+
assert!(
112+
is_constant(struct_of_lists.as_ref())
113+
.unwrap()
114+
.unwrap_or_default()
115+
);
111116
}
112117

113118
#[rstest]

0 commit comments

Comments
 (0)