Skip to content

Commit 2973c55

Browse files
godnight10061godnight10061
committed
Fix scan dtype for list-of-struct projection
1 parent 014d10f commit 2973c55

File tree

7 files changed

+82
-30
lines changed

7 files changed

+82
-30
lines changed

‎vortex-array/src/expr/analysis/immediate_access.rs‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,18 +25,17 @@ pub fn annotate_scope_access(scope: &StructFields) -> impl AnnotationFn<Annotati
2525
"cannot analyse select, simplify the expression"
2626
);
2727

28-
if let Some(field_name) = expr.as_opt::<GetItem>() {
29-
if expr.child(0).is::<Root>() {
30-
return vec![field_name.clone()];
31-
}
28+
if let Some(field_name) = expr.as_opt::<GetItem>()
29+
&& expr.child(0).is::<Root>()
30+
{
31+
return vec![field_name.clone()];
3232
}
3333

34-
if expr.is::<GetItemList>() {
35-
if let Some(field_name) = expr.child(0).as_opt::<GetItem>() {
36-
if expr.child(0).child(0).is::<Root>() {
37-
return vec![field_name.clone()];
38-
}
39-
}
34+
if expr.is::<GetItemList>()
35+
&& let Some(field_name) = expr.child(0).as_opt::<GetItem>()
36+
&& expr.child(0).child(0).is::<Root>()
37+
{
38+
return vec![field_name.clone()];
4039
}
4140

4241
if expr.is::<Root>() {

‎vortex-array/src/expr/exprs/get_item.rs‎

Lines changed: 34 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -392,16 +392,33 @@ mod tests {
392392
);
393393

394394
assert_eq!(
395-
out.scalar_at(0).as_list().elements().unwrap().to_vec(),
395+
out.scalar_at(0)
396+
.unwrap()
397+
.as_list()
398+
.elements()
399+
.unwrap()
400+
.to_vec(),
396401
vec![
397402
Scalar::primitive(1i32, NonNullable),
398403
Scalar::primitive(2i32, NonNullable),
399404
]
400405
);
401-
assert!(out.scalar_at(1).as_list().elements().unwrap().is_empty());
402-
assert!(out.scalar_at(2).is_null());
406+
assert!(
407+
out.scalar_at(1)
408+
.unwrap()
409+
.as_list()
410+
.elements()
411+
.unwrap()
412+
.is_empty()
413+
);
414+
assert!(out.scalar_at(2).unwrap().is_null());
403415
assert_eq!(
404-
out.scalar_at(3).as_list().elements().unwrap().to_vec(),
416+
out.scalar_at(3)
417+
.unwrap()
418+
.as_list()
419+
.elements()
420+
.unwrap()
421+
.to_vec(),
405422
vec![Scalar::primitive(3i32, NonNullable)]
406423
);
407424
}
@@ -457,15 +474,25 @@ mod tests {
457474
);
458475

459476
assert_eq!(
460-
out.scalar_at(0).as_list().elements().unwrap().to_vec(),
477+
out.scalar_at(0)
478+
.unwrap()
479+
.as_list()
480+
.elements()
481+
.unwrap()
482+
.to_vec(),
461483
vec![
462484
Scalar::primitive(1i32, Nullability::Nullable),
463485
Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
464486
]
465487
);
466-
assert!(out.scalar_at(1).is_null());
488+
assert!(out.scalar_at(1).unwrap().is_null());
467489
assert_eq!(
468-
out.scalar_at(2).as_list().elements().unwrap().to_vec(),
490+
out.scalar_at(2)
491+
.unwrap()
492+
.as_list()
493+
.elements()
494+
.unwrap()
495+
.to_vec(),
469496
vec![
470497
Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
471498
Scalar::primitive(6i32, Nullability::Nullable),

‎vortex-array/src/expr/exprs/get_item_list.rs‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,6 @@ use vortex_error::VortexResult;
1313
use vortex_error::vortex_bail;
1414
use vortex_error::vortex_err;
1515

16-
use crate::ArrayRef;
17-
use crate::Executable;
1816
use crate::IntoArray;
1917
use crate::arrays::FixedSizeListArray;
2018
use crate::arrays::ListViewArray;
@@ -149,7 +147,7 @@ impl VTable for GetItemList {
149147
field,
150148
list.offsets().clone(),
151149
list.sizes().clone(),
152-
list.validity()?,
150+
list.validity().clone(),
153151
)?
154152
.into_array()
155153
.execute(args.ctx)
@@ -164,9 +162,14 @@ impl VTable for GetItemList {
164162
Nullability::Nullable => mask(&field, &struct_elems.validity_mask()?.not())?,
165163
};
166164

167-
FixedSizeListArray::try_new(field, list.list_size(), list.validity()?, list.len())?
168-
.into_array()
169-
.execute(args.ctx)
165+
FixedSizeListArray::try_new(
166+
field,
167+
list.list_size(),
168+
list.validity().clone(),
169+
list.len(),
170+
)?
171+
.into_array()
172+
.execute(args.ctx)
170173
}
171174
_ => Err(vortex_err!(
172175
"Expected list scope for GetItemList execution, got {}",

‎vortex-layout/src/layouts/list/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ impl VTable for ListVTable {
143143
_metadata: &<Self::Metadata as DeserializeMetadata>::Output,
144144
_segment_ids: Vec<SegmentId>,
145145
children: &dyn LayoutChildren,
146-
_ctx: ArrayContext,
146+
_ctx: &ArrayContext,
147147
) -> VortexResult<Self::Layout> {
148148
vortex_ensure!(
149149
matches!(dtype, DType::List(..) | DType::FixedSizeList(..)),

‎vortex-layout/src/layouts/list/writer.rs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,18 @@ impl LayoutStrategy for ListStrategy {
156156

157157
// validity (optional)
158158
if is_nullable {
159-
let validity = chunk.validity_mask().into_array();
159+
let validity = match chunk.validity_mask() {
160+
Ok(validity) => validity.into_array(),
161+
Err(e) => {
162+
let e: Arc<VortexError> = Arc::new(e);
163+
for tx in column_streams_tx.iter() {
164+
let _ = tx
165+
.send(Err(VortexError::from(e.clone())))
166+
.await;
167+
}
168+
break;
169+
}
170+
};
160171
let _ = column_streams_tx[0]
161172
.send(Ok((sequence_pointer.advance(), validity)))
162173
.await;

‎vortex-scan/src/gpu/gpubuilder.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,8 @@ impl<A: 'static + Send> GpuScanBuilder<A> {
6464

6565
/// The [`DType`] returned by the scan, after applying the projection.
6666
pub fn dtype(&self) -> VortexResult<DType> {
67-
self.projection.return_dtype(self.layout_reader.dtype())
67+
let projection = simplify_typed(self.projection.clone(), self.layout_reader.dtype())?;
68+
projection.return_dtype(self.layout_reader.dtype())
6869
}
6970

7071
/// Map each split of the scan. The function will be run on the spawned task.
@@ -82,7 +83,6 @@ impl<A: 'static + Send> GpuScanBuilder<A> {
8283
}
8384

8485
pub fn prepare(self) -> VortexResult<GpuScan<A>> {
85-
let dtype = self.dtype()?;
8686
let handle = self.session.handle();
8787

8888
// Spin up the root layout reader, and wrap it in a FilterLayoutReader to perform
@@ -91,6 +91,7 @@ impl<A: 'static + Send> GpuScanBuilder<A> {
9191

9292
// Normalize and simplify the expressions.
9393
let projection = simplify_typed(self.projection, layout_reader.dtype())?;
94+
let dtype = projection.return_dtype(layout_reader.dtype())?;
9495

9596
// Construct field masks and compute the row splits of the scan.
9697
let (filter_mask, projection_mask) =

‎vortex-scan/src/scan_builder.rs‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,19 @@ impl<A: 'static + Send> ScanBuilder<A> {
187187

188188
/// The [`DType`] returned by the scan, after applying the projection.
189189
pub fn dtype(&self) -> VortexResult<DType> {
190-
self.projection.return_dtype(self.layout_reader.dtype())
190+
// NOTE: `GetItem` may simplify into `GetItemList` for list-of-struct projections.
191+
// To avoid rejecting valid nested projections (like `items.a`) we must simplify before
192+
// validating the return dtype.
193+
//
194+
// Also, `row_idx` support is provided by `RowIdxLayoutReader`, so use the same reader
195+
// enrichment as `prepare`.
196+
let layout_reader = Arc::new(RowIdxLayoutReader::new(
197+
self.row_offset,
198+
self.layout_reader.clone(),
199+
self.session.clone(),
200+
));
201+
let projection = self.projection.optimize_recursive(layout_reader.dtype())?;
202+
projection.return_dtype(layout_reader.dtype())
191203
}
192204

193205
/// The session used by the scan.
@@ -220,8 +232,6 @@ impl<A: 'static + Send> ScanBuilder<A> {
220232
}
221233

222234
pub fn prepare(self) -> VortexResult<RepeatedScan<A>> {
223-
let dtype = self.dtype()?;
224-
225235
if self.filter.is_some() && self.limit.is_some() {
226236
vortex_bail!("Vortex doesn't support scans with both a filter and a limit")
227237
}
@@ -241,6 +251,7 @@ impl<A: 'static + Send> ScanBuilder<A> {
241251

242252
// Normalize and simplify the expressions.
243253
let projection = self.projection.optimize_recursive(layout_reader.dtype())?;
254+
let dtype = projection.return_dtype(layout_reader.dtype())?;
244255

245256
let filter = self
246257
.filter

0 commit comments

Comments
 (0)