Skip to content

Commit 46ebd4c

Browse files
gatesngatesn
committed
Clarify docs in session accessors
Signed-off-by: Nicholas Gates <[email protected]>
1 parent ecef7da commit 46ebd4c

File tree

11 files changed

+329
-18
lines changed

11 files changed

+329
-18
lines changed

‎vortex-array/src/arrays/scalar_fn/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
mod array;
55
mod kernel;
66
mod metadata;
7+
mod rules;
78
mod vtable;
89

910
pub use array::*;

‎vortex-array/src/arrays/scalar_fn/rules.rs‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use vortex_error::VortexExpect;
5+
use vortex_error::VortexResult;
6+
use vortex_scalar::Scalar;
7+
use vortex_vector::Datum;
8+
use vortex_vector::scalar_matches_dtype;
9+
10+
use crate::Array;
11+
use crate::ArrayRef;
12+
use crate::IntoArray;
13+
use crate::arrays::AnyScalarFn;
14+
use crate::arrays::ConstantArray;
15+
use crate::arrays::ConstantVTable;
16+
use crate::arrays::ScalarFnArray;
17+
use crate::expr::ExecutionArgs;
18+
use crate::optimizer::rules::ArrayReduceRule;
19+
20+
#[derive(Debug)]
21+
pub(crate) struct ScalarFnConstantRule;
22+
impl ArrayReduceRule<AnyScalarFn> for ScalarFnConstantRule {
23+
fn matcher(&self) -> AnyScalarFn {
24+
AnyScalarFn
25+
}
26+
27+
fn reduce(&self, array: &ScalarFnArray) -> VortexResult<Option<ArrayRef>> {
28+
if !array.children.iter().all(|c| c.is::<ConstantVTable>()) {
29+
return Ok(None);
30+
}
31+
32+
let input_datums: Vec<_> = array
33+
.children
34+
.iter()
35+
.map(|c| c.as_::<ConstantVTable>().scalar().to_vector_scalar())
36+
.map(Datum::Scalar)
37+
.collect();
38+
let input_dtypes = array.children.iter().map(|c| c.dtype().clone()).collect();
39+
40+
let result = array
41+
.scalar_fn
42+
.execute(ExecutionArgs {
43+
datums: input_datums,
44+
dtypes: input_dtypes,
45+
row_count: array.len,
46+
return_dtype: array.dtype.clone(),
47+
})?
48+
.into_scalar()
49+
.vortex_expect("Scalar inputs should produce scalar output");
50+
51+
assert!(scalar_matches_dtype(&result, &array.dtype));
52+
53+
Ok(Some(
54+
ConstantArray::new(Scalar::from_vector_scalar(result, &array.dtype)?, array.len)
55+
.into_array(),
56+
))
57+
}
58+
}

‎vortex-compute/src/filter/vector/fixed_size_list.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ where
3232
type Output = FixedSizeListVector;
3333

3434
fn filter(self, selection: &M) -> Self::Output {
35-
let list_size = self.list_size();
35+
let list_size = self.element_size();
3636
let filtered_validity = self.validity().filter(selection);
3737

3838
let filtered_elements = if list_size != 0 {

‎vortex-layout/src/layouts/zoned/zone_map.rs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ use vortex_array::expr::stats::Precision;
1313
use vortex_array::expr::stats::Stat;
1414
use vortex_array::expr::stats::StatsProvider;
1515
use vortex_array::mask::MaskExecutor;
16+
use vortex_array::session::ArraySessionExt;
1617
use vortex_array::stats::StatsSet;
1718
use vortex_array::validity::Validity;
1819
use vortex_dtype::DType;
@@ -152,10 +153,11 @@ impl ZoneMap {
152153
/// All zones where the predicate evaluates to `true` can be skipped entirely.
153154
pub fn prune(&self, predicate: &Expression, session: &VortexSession) -> VortexResult<Mask> {
154155
if *USE_VORTEX_OPERATORS {
155-
self.array
156-
.to_array()
157-
.apply(predicate)?
158-
.execute_mask_optimized(session)
156+
let array = self.array.to_array().apply(predicate)?;
157+
println!("PRUNE: {}", array.display_tree());
158+
let array = session.arrays().optimizer().optimize_array(&array)?;
159+
println!("PRUNE OPTIMIZED: {}", array.display_tree());
160+
array.execute_mask(session)
159161
} else {
160162
predicate
161163
.evaluate(&self.array.to_array())?

‎vortex-scalar/src/vectors.rs‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,25 @@
33

44
//! Conversion logic from this "legacy" scalar crate to Vortex Vector scalars.
55
6+
use std::ops::Deref;
67
use std::sync::Arc;
78

9+
use itertools::Itertools;
810
use vortex_buffer::Buffer;
911
use vortex_dtype::DType;
1012
use vortex_dtype::DecimalType;
13+
use vortex_dtype::NativeDecimalType;
14+
use vortex_dtype::PTypeDowncastExt;
1115
use vortex_dtype::PrecisionScale;
1216
use vortex_dtype::match_each_decimal_value_type;
1317
use vortex_dtype::match_each_native_ptype;
1418
use vortex_error::VortexExpect;
19+
use vortex_error::VortexResult;
20+
use vortex_error::vortex_ensure;
1521
use vortex_mask::Mask;
1622
use vortex_vector::VectorMut;
1723
use vortex_vector::VectorMutOps;
24+
use vortex_vector::VectorOps;
1825
use vortex_vector::binaryview::BinaryScalar;
1926
use vortex_vector::binaryview::StringScalar;
2027
use vortex_vector::bool::BoolScalar;
@@ -30,6 +37,7 @@ use vortex_vector::primitive::PVector;
3037
use vortex_vector::struct_::StructScalar;
3138
use vortex_vector::struct_::StructVector;
3239

40+
use crate::DecimalValue;
3341
use crate::Scalar;
3442

3543
impl Scalar {
@@ -158,4 +166,115 @@ impl Scalar {
158166
DType::Extension(_) => self.as_extension().storage().to_vector_scalar(),
159167
}
160168
}
169+
170+
/// Convert a `vortex-vector` [`vortex_vector::Scalar`] into a `vortex-scalar` [`Scalar`].
171+
pub fn from_vector_scalar(scalar: vortex_vector::Scalar, dtype: &DType) -> VortexResult<Self> {
172+
Ok(match dtype {
173+
DType::Null => Scalar::null(DType::Null),
174+
DType::Bool(n) => match scalar.as_bool().value() {
175+
None => {
176+
vortex_ensure!(
177+
n.is_nullable(),
178+
"Cannot create null scalar for non-nullable dtype"
179+
);
180+
Scalar::null(dtype.clone())
181+
}
182+
Some(b) => Scalar::bool(b, *n),
183+
},
184+
DType::Primitive(ptype, n) => {
185+
match_each_native_ptype!(ptype, |T| {
186+
let pscalar = scalar.into_primitive().downcast::<T>();
187+
match pscalar.value() {
188+
None => {
189+
vortex_ensure!(
190+
n.is_nullable(),
191+
"Cannot create null scalar for non-nullable dtype"
192+
);
193+
Scalar::null(dtype.clone())
194+
}
195+
Some(v) => Scalar::primitive(v, *n),
196+
}
197+
})
198+
}
199+
DType::Decimal(dec_type, n) => {
200+
let dec_scalar = scalar.into_decimal();
201+
match_each_decimal_value_type!(
202+
DecimalType::smallest_decimal_value_type(dec_type),
203+
|D| {
204+
let dscalar = <D as NativeDecimalType>::downcast(dec_scalar);
205+
match dscalar.value() {
206+
None => {
207+
vortex_ensure!(
208+
n.is_nullable(),
209+
"Cannot create null scalar for non-nullable dtype"
210+
);
211+
Scalar::null(dtype.clone())
212+
}
213+
Some(v) => Scalar::decimal(DecimalValue::from(v), dec_type.clone(), *n),
214+
}
215+
}
216+
)
217+
}
218+
DType::Utf8(n) => match scalar.as_string().value() {
219+
None => {
220+
vortex_ensure!(
221+
n.is_nullable(),
222+
"Cannot create null scalar for non-nullable dtype"
223+
);
224+
Scalar::null(dtype.clone())
225+
}
226+
Some(s) => Scalar::utf8(s.clone(), *n),
227+
},
228+
DType::Binary(n) => match scalar.as_binary().value() {
229+
None => {
230+
vortex_ensure!(
231+
n.is_nullable(),
232+
"Cannot create null scalar for non-nullable dtype"
233+
);
234+
Scalar::null(dtype.clone())
235+
}
236+
Some(b) => Scalar::binary(b.clone(), *n),
237+
},
238+
DType::List(elem_dtype, n) => {
239+
let elements = scalar.as_list().value().elements();
240+
Scalar::list(
241+
elem_dtype.clone(),
242+
(0..elements.len())
243+
.map(|idx| elements.scalar_at(idx))
244+
.map(|scalar| Scalar::from_vector_scalar(scalar, elem_dtype.deref()))
245+
.try_collect()?,
246+
*n,
247+
)
248+
}
249+
DType::FixedSizeList(elem_dtype, size, n) => {
250+
let scalar = scalar.into_fixed_size_list();
251+
let elements = scalar.value().elements();
252+
vortex_ensure!(scalar.value().element_size() == *size);
253+
254+
Scalar::fixed_size_list(
255+
DType::FixedSizeList(elem_dtype.clone(), *size, *n),
256+
(0..*size as usize)
257+
.map(|idx| {
258+
Scalar::from_vector_scalar(elements.scalar_at(idx), elem_dtype.deref())
259+
})
260+
.try_collect()?,
261+
*n,
262+
)
263+
}
264+
DType::Struct(fields, n) => Scalar::struct_(
265+
DType::Struct(fields.clone(), *n),
266+
fields
267+
.fields()
268+
.zip(scalar.into_struct().fields())
269+
.map(|(field_dtype, field_scalar)| {
270+
Scalar::from_vector_scalar(field_scalar, &field_dtype)
271+
})
272+
.try_collect()?,
273+
),
274+
DType::Extension(ext_dtype) => Scalar::extension(
275+
ext_dtype.clone(),
276+
Scalar::from_vector_scalar(scalar, ext_dtype.storage_dtype())?,
277+
),
278+
})
279+
}
161280
}

‎vortex-vector/src/decimal/scalar.rs‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4+
use vortex_dtype::DecimalTypeDowncast;
45
use vortex_dtype::DecimalTypeUpcast;
56
use vortex_dtype::NativeDecimalType;
67
use vortex_dtype::PrecisionScale;
78
use vortex_dtype::i256;
89
use vortex_error::VortexExpect;
10+
use vortex_error::vortex_panic;
911

1012
use crate::Scalar;
1113
use crate::ScalarOps;
@@ -195,3 +197,49 @@ impl DecimalTypeUpcast for DecimalScalar {
195197
DecimalScalar::D256(input)
196198
}
197199
}
200+
201+
impl DecimalTypeDowncast for DecimalScalar {
202+
type Output<T: NativeDecimalType> = DScalar<T>;
203+
204+
fn into_i8(self) -> Self::Output<i8> {
205+
if let Self::D8(v) = self {
206+
return v;
207+
}
208+
vortex_panic!("Expected DecimalScalar::D8, got {self:?}");
209+
}
210+
211+
fn into_i16(self) -> Self::Output<i16> {
212+
if let Self::D16(v) = self {
213+
return v;
214+
}
215+
vortex_panic!("Expected DecimalScalar::D16, got {self:?}");
216+
}
217+
218+
fn into_i32(self) -> Self::Output<i32> {
219+
if let Self::D32(v) = self {
220+
return v;
221+
}
222+
vortex_panic!("Expected DecimalScalar::D32, got {self:?}");
223+
}
224+
225+
fn into_i64(self) -> Self::Output<i64> {
226+
if let Self::D64(v) = self {
227+
return v;
228+
}
229+
vortex_panic!("Expected DecimalScalar::D64, got {self:?}");
230+
}
231+
232+
fn into_i128(self) -> Self::Output<i128> {
233+
if let Self::D128(v) = self {
234+
return v;
235+
}
236+
vortex_panic!("Expected DecimalScalar::D128, got {self:?}");
237+
}
238+
239+
fn into_i256(self) -> Self::Output<i256> {
240+
if let Self::D256(v) = self {
241+
return v;
242+
}
243+
vortex_panic!("Expected DecimalScalar::D256, got {self:?}");
244+
}
245+
}

‎vortex-vector/src/fixed_size_list/vector.rs‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -143,11 +143,6 @@ impl FixedSizeListVector {
143143
pub fn elements(&self) -> &Arc<Vector> {
144144
&self.elements
145145
}
146-
147-
/// Returns the size of every list in the vector.
148-
pub fn list_size(&self) -> u32 {
149-
self.list_size
150-
}
151146
}
152147

153148
impl VectorOps for FixedSizeListVector {
@@ -261,7 +256,7 @@ mod tests {
261256
let validity = Mask::new_true(2);
262257
let vec = FixedSizeListVector::new(elements.clone(), 3, validity.clone());
263258
assert_eq!(vec.len(), 2);
264-
assert_eq!(vec.list_size(), 3);
259+
assert_eq!(vec.element_size(), 3);
265260

266261
// Valid construction with try_new().
267262
let result = FixedSizeListVector::try_new(elements.clone(), 3, validity);
@@ -284,7 +279,7 @@ mod tests {
284279
assert!(result.is_ok());
285280
let vec = result.unwrap();
286281
assert_eq!(vec.len(), 5);
287-
assert_eq!(vec.list_size(), 0);
282+
assert_eq!(vec.element_size(), 0);
288283

289284
// Degenerate case with non-empty elements should fail.
290285
let result = FixedSizeListVector::try_new(elements, 0, Mask::new_true(1));
@@ -296,7 +291,7 @@ mod tests {
296291
let validity = Mask::new_true(2);
297292
let vec = unsafe { FixedSizeListVector::new_unchecked(elements, 2, validity) };
298293
assert_eq!(vec.len(), 2);
299-
assert_eq!(vec.list_size(), 2);
294+
assert_eq!(vec.element_size(), 2);
300295
}
301296

302297
#[test]
@@ -337,7 +332,7 @@ mod tests {
337332
// The error case should return the original vector.
338333
if let Err(returned_vec) = result {
339334
assert_eq!(returned_vec.len(), 2);
340-
assert_eq!(returned_vec.list_size(), 2);
335+
assert_eq!(returned_vec.element_size(), 2);
341336
}
342337
}
343338

@@ -353,7 +348,7 @@ mod tests {
353348

354349
// Test accessors.
355350
assert_eq!(vec.len(), 3);
356-
assert_eq!(vec.list_size(), 2);
351+
assert_eq!(vec.element_size(), 2);
357352
assert_eq!(vec.elements().len(), 6);
358353
assert_eq!(vec.validity().true_count(), 3);
359354

‎vortex-vector/src/fixed_size_list/vector_mut.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -371,7 +371,7 @@ mod tests {
371371
// Test freeze and accessors.
372372
let frozen = vec.freeze();
373373
assert_eq!(frozen.len(), 5);
374-
assert_eq!(frozen.list_size(), 3);
374+
assert_eq!(frozen.element_size(), 3);
375375
assert_eq!(frozen.elements().len(), 15);
376376
}
377377

@@ -503,7 +503,7 @@ mod tests {
503503
// Verify operations work correctly.
504504
let frozen = vec.freeze();
505505
assert_eq!(frozen.len(), 1);
506-
assert_eq!(frozen.list_size(), 1000);
506+
assert_eq!(frozen.element_size(), 1000);
507507
}
508508

509509
#[test]

‎vortex-vector/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ pub fn vector_matches_dtype(vector: &Vector, dtype: &DType) -> bool {
118118

119119
/// Returns true if the scalar's is compatible with the provided data type.
120120
pub fn scalar_matches_dtype(scalar: &Scalar, dtype: &DType) -> bool {
121-
if !dtype.is_nullable() && scalar.is_valid() {
121+
if !dtype.is_nullable() && scalar.is_invalid() {
122122
// Non-nullable dtype cannot have nulls in the scalar.
123123
return false;
124124
}

0 commit comments

Comments
 (0)