|
1 | 1 | // SPDX-License-Identifier: Apache-2.0 |
2 | 2 | // SPDX-FileCopyrightText: Copyright the Vortex contributors |
3 | 3 |
|
| 4 | +use vortex_array::ArrayRef; |
| 5 | +use vortex_array::ExecutionCtx; |
4 | 6 | use vortex_array::IntoArray; |
5 | 7 | use vortex_array::ToCanonical; |
6 | | -use vortex_array::compute::IsSortedKernel; |
7 | | -use vortex_array::compute::IsSortedKernelAdapter; |
8 | | -use vortex_array::compute::is_sorted; |
9 | | -use vortex_array::compute::is_strict_sorted; |
10 | | -use vortex_array::register_kernel; |
| 8 | +use vortex_array::aggregate_fn::AggregateFnRef; |
| 9 | +use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; |
| 10 | +use vortex_array::aggregate_fn::fns::is_sorted::is_sorted; |
| 11 | +use vortex_array::aggregate_fn::fns::is_sorted::is_strict_sorted; |
| 12 | +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; |
| 13 | +use vortex_array::scalar::Scalar; |
11 | 14 | use vortex_error::VortexResult; |
12 | 15 |
|
13 | 16 | use crate::FoR; |
14 | | -use crate::FoRArray; |
15 | 17 |
|
16 | 18 | /// FoR can express sortedness directly on its encoded form. |
17 | 19 | /// |
@@ -71,92 +73,107 @@ use crate::FoRArray; |
71 | 73 | /// Addition is order-preserving, so all the wrapped values preserve their order and they're all |
72 | 74 | /// represented as unsigned values larger than 127 so they also preserve their order with the |
73 | 75 | /// unwrapped values. |
74 | | -impl IsSortedKernel for FoR { |
75 | | - fn is_sorted(&self, array: &FoRArray) -> VortexResult<Option<bool>> { |
76 | | - let encoded = array.encoded().to_primitive(); |
77 | | - is_sorted( |
78 | | - &encoded |
79 | | - .reinterpret_cast(encoded.ptype().to_unsigned()) |
80 | | - .into_array(), |
81 | | - ) |
82 | | - } |
| 76 | +#[derive(Debug)] |
| 77 | +pub(crate) struct FoRIsSortedKernel; |
| 78 | + |
| 79 | +impl DynAggregateKernel for FoRIsSortedKernel { |
| 80 | + fn aggregate( |
| 81 | + &self, |
| 82 | + aggregate_fn: &AggregateFnRef, |
| 83 | + batch: &ArrayRef, |
| 84 | + ctx: &mut ExecutionCtx, |
| 85 | + ) -> VortexResult<Option<Scalar>> { |
| 86 | + let Some(options) = aggregate_fn.as_opt::<IsSorted>() else { |
| 87 | + return Ok(None); |
| 88 | + }; |
| 89 | + |
| 90 | + let Some(array) = batch.as_opt::<FoR>() else { |
| 91 | + return Ok(None); |
| 92 | + }; |
83 | 93 |
|
84 | | - fn is_strict_sorted(&self, array: &FoRArray) -> VortexResult<Option<bool>> { |
85 | 94 | let encoded = array.encoded().to_primitive(); |
86 | | - is_strict_sorted( |
87 | | - &encoded |
88 | | - .reinterpret_cast(encoded.ptype().to_unsigned()) |
89 | | - .into_array(), |
90 | | - ) |
| 95 | + let unsigned_array = encoded |
| 96 | + .reinterpret_cast(encoded.ptype().to_unsigned()) |
| 97 | + .into_array(); |
| 98 | + |
| 99 | + let result = if options.strict { |
| 100 | + is_strict_sorted(&unsigned_array, ctx)? |
| 101 | + } else { |
| 102 | + is_sorted(&unsigned_array, ctx)? |
| 103 | + }; |
| 104 | + |
| 105 | + Ok(Some(IsSorted::make_partial(batch, result, options.strict)?)) |
91 | 106 | } |
92 | 107 | } |
93 | 108 |
|
94 | | -register_kernel!(IsSortedKernelAdapter(FoR).lift()); |
95 | | - |
96 | 109 | #[cfg(test)] |
97 | 110 | mod test { |
98 | 111 | use vortex_array::IntoArray; |
| 112 | + use vortex_array::LEGACY_SESSION; |
| 113 | + use vortex_array::VortexSessionExecute; |
| 114 | + use vortex_array::aggregate_fn::fns::is_sorted::is_sorted; |
99 | 115 | use vortex_array::arrays::PrimitiveArray; |
100 | | - use vortex_array::compute::is_sorted; |
101 | 116 | use vortex_array::validity::Validity; |
102 | 117 | use vortex_buffer::buffer; |
103 | 118 |
|
104 | 119 | use crate::FoRArray; |
105 | 120 |
|
106 | 121 | #[test] |
107 | 122 | fn test_sorted() { |
| 123 | + let mut ctx = LEGACY_SESSION.create_execution_ctx(); |
| 124 | + |
108 | 125 | let a = PrimitiveArray::new(buffer![-1, 0, i8::MAX], Validity::NonNullable); |
109 | 126 | let b = FoRArray::encode(a).unwrap(); |
110 | 127 | assert!( |
111 | | - is_sorted(&b.clone().into_array()).unwrap().unwrap(), |
| 128 | + is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), |
112 | 129 | "{}", |
113 | 130 | b.encoded().display_values() |
114 | 131 | ); |
115 | 132 |
|
116 | 133 | let a = PrimitiveArray::new(buffer![i8::MIN, 0, i8::MAX], Validity::NonNullable); |
117 | 134 | let b = FoRArray::encode(a).unwrap(); |
118 | 135 | assert!( |
119 | | - is_sorted(&b.clone().into_array()).unwrap().unwrap(), |
| 136 | + is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), |
120 | 137 | "{}", |
121 | 138 | b.encoded().display_values() |
122 | 139 | ); |
123 | 140 |
|
124 | 141 | let a = PrimitiveArray::new(buffer![i8::MIN, 0, 30, 127], Validity::NonNullable); |
125 | 142 | let b = FoRArray::encode(a).unwrap(); |
126 | 143 | assert!( |
127 | | - is_sorted(&b.clone().into_array()).unwrap().unwrap(), |
| 144 | + is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), |
128 | 145 | "{}", |
129 | 146 | b.encoded().display_values() |
130 | 147 | ); |
131 | 148 |
|
132 | 149 | let a = PrimitiveArray::new(buffer![i8::MIN, -3, -1], Validity::NonNullable); |
133 | 150 | let b = FoRArray::encode(a).unwrap(); |
134 | 151 | assert!( |
135 | | - is_sorted(&b.clone().into_array()).unwrap().unwrap(), |
| 152 | + is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), |
136 | 153 | "{}", |
137 | 154 | b.encoded().display_values() |
138 | 155 | ); |
139 | 156 |
|
140 | 157 | let a = PrimitiveArray::new(buffer![-10, -3, -1], Validity::NonNullable); |
141 | 158 | let b = FoRArray::encode(a).unwrap(); |
142 | 159 | assert!( |
143 | | - is_sorted(&b.clone().into_array()).unwrap().unwrap(), |
| 160 | + is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), |
144 | 161 | "{}", |
145 | 162 | b.encoded().display_values() |
146 | 163 | ); |
147 | 164 |
|
148 | 165 | let a = PrimitiveArray::new(buffer![-10, -11, -1], Validity::NonNullable); |
149 | 166 | let b = FoRArray::encode(a).unwrap(); |
150 | 167 | assert!( |
151 | | - !is_sorted(&b.clone().into_array()).unwrap().unwrap(), |
| 168 | + !is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), |
152 | 169 | "{}", |
153 | 170 | b.encoded().display_values() |
154 | 171 | ); |
155 | 172 |
|
156 | 173 | let a = PrimitiveArray::new(buffer![-10, i8::MIN, -1], Validity::NonNullable); |
157 | 174 | let b = FoRArray::encode(a).unwrap(); |
158 | 175 | assert!( |
159 | | - !is_sorted(&b.clone().into_array()).unwrap().unwrap(), |
| 176 | + !is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), |
160 | 177 | "{}", |
161 | 178 | b.encoded().display_values() |
162 | 179 | ); |
|
0 commit comments