Skip to content

feat[layout]: list-of-struct nested projection #6012

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
godnight10061 wants to merge 11 commits into
base: develop
Choose a base branch
from
+1,825 −29
Draft

feat[layout]: list-of-struct nested projection #6012

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
5ca3b12
Fix nested projection for list-of-struct columns
godnight10061 Jan 17, 2026
d24dced
Add unit regression for list-of-struct get_item
godnight10061 Jan 17, 2026
e4b0a5e
Fix lints in list layout and projections
godnight10061 Jan 17, 2026
8d4ff5d
Refactor list projection internals
godnight10061 Jan 20, 2026
b110e92
Improve ListReader split planning for FixedSizeList
godnight10061 Jan 21, 2026
cb2bc25
Use get_item_list for list-of-struct projection
godnight10061 Jan 22, 2026
bc0be24
Address list-of-struct projection review
godnight10061 Jan 23, 2026
d12627d
Fix scan dtype for list-of-struct projection
godnight10061 Jan 23, 2026
f964e10
Add unit regression for list-of-struct scan dtype
godnight10061 Jan 23, 2026
da77c0c
Improve ListReader split planning for List
godnight10061 Jan 24, 2026
86e36d2
Use unmasked_field_by_name in GetItemList
godnight10061 Jan 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 15 additions & 5 deletions vortex-array/src/expr/analysis/immediate_access.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use crate::expr::analysis::AnnotationFn;
use crate::expr::analysis::Annotations;
use crate::expr::descendent_annotations;
use crate::expr::exprs::get_item::GetItem;
use crate::expr::exprs::get_item_list::GetItemList;
use crate::expr::exprs::root::Root;
use crate::expr::exprs::select::Select;

Expand All @@ -24,11 +25,20 @@ pub fn annotate_scope_access(scope: &StructFields) -> impl AnnotationFn<Annotati
"cannot analyse select, simplify the expression"
);

if let Some(field_name) = expr.as_opt::<GetItem>() {
if expr.child(0).is::<Root>() {
return vec![field_name.clone()];
}
} else if expr.is::<Root>() {
if let Some(field_name) = expr.as_opt::<GetItem>()
&& expr.child(0).is::<Root>()
{
return vec![field_name.clone()];
}

if expr.is::<GetItemList>()
&& let Some(field_name) = expr.child(0).as_opt::<GetItem>()
&& expr.child(0).child(0).is::<Root>()
{
return vec![field_name.clone()];
}

if expr.is::<Root>() {
return scope.names().iter().cloned().collect();
}

Expand Down
217 changes: 208 additions & 9 deletions vortex-array/src/expr/exprs/get_item.rs
View file
Open in desktop
Copy link
Contributor

@joseph-isaacs joseph-isaacs Jan 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need these changes?

Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,11 @@ use crate::expr::Pack;
use crate::expr::ReduceCtx;
use crate::expr::ReduceNode;
use crate::expr::ReduceNodeRef;
use crate::expr::SimplifyCtx;
use crate::expr::StatsCatalog;
use crate::expr::VTable;
use crate::expr::VTableExt;
use crate::expr::exprs::get_item_list::get_item_list;
use crate::expr::exprs::root::root;
use crate::expr::lit;
use crate::expr::stats::Stat;
Expand Down Expand Up @@ -90,15 +92,7 @@ impl VTable for GetItem {
vortex_err!("Couldn't find the {} field in the input scope", field_name)
})?;

// Match here to avoid cloning the dtype if nullability doesn't need to change
if matches!(
(struct_dtype.nullability(), field_dtype.nullability()),
(Nullability::Nullable, Nullability::NonNullable)
) {
return Ok(field_dtype.with_nullability(Nullability::Nullable));
}

Ok(field_dtype)
Ok(field_dtype.union_nullability(struct_dtype.nullability()))
}

fn execute(
Expand All @@ -120,6 +114,30 @@ impl VTable for GetItem {
.execute(args.ctx)
}

fn simplify(
&self,
field_name: &FieldName,
expr: &Expression,
ctx: &dyn SimplifyCtx,
) -> VortexResult<Option<Expression>> {
let child = expr.child(0);
let child_dtype = ctx.return_dtype(child)?;

let element_dtype = match child_dtype {
DType::List(element_dtype, _) => Some(element_dtype),
DType::FixedSizeList(element_dtype, ..) => Some(element_dtype),
_ => None,
};

if let Some(element_dtype) = element_dtype
&& element_dtype.as_struct_fields_opt().is_some()
{
Ok(Some(get_item_list(field_name.clone(), child.clone())))
} else {
Ok(None)
}
}

fn reduce(
&self,
field_name: &FieldName,
Expand Down Expand Up @@ -240,6 +258,8 @@ pub fn get_item(field: impl Into<FieldName>, child: Expression) -> Expression {

#[cfg(test)]
mod tests {
use std::sync::Arc;

use vortex_buffer::buffer;
use vortex_dtype::DType;
use vortex_dtype::FieldNames;
Expand All @@ -251,9 +271,12 @@ mod tests {

use crate::Array;
use crate::IntoArray;
use crate::arrays::FixedSizeListArray;
use crate::arrays::ListArray;
use crate::arrays::StructArray;
use crate::expr::exprs::binary::checked_add;
use crate::expr::exprs::get_item::get_item;
use crate::expr::exprs::get_item_list::get_item_list;
use crate::expr::exprs::literal::lit;
use crate::expr::exprs::pack::pack;
use crate::expr::exprs::root::root;
Expand Down Expand Up @@ -301,6 +324,182 @@ mod tests {
);
}

#[test]
fn get_item_list_of_struct() {
let element_dtype = Arc::new(DType::Struct(
[
("a", DType::Primitive(PType::I32, NonNullable)),
("b", DType::Utf8(NonNullable)),
]
.into_iter()
.collect(),
NonNullable,
));

let row_count = 4;
let items = ListArray::from_iter_opt_slow::<u32, _, _>(
[
Some(vec![
Scalar::struct_(
(*element_dtype).clone(),
vec![
Scalar::primitive(1i32, NonNullable),
Scalar::utf8("x", NonNullable),
],
),
Scalar::struct_(
(*element_dtype).clone(),
vec![
Scalar::primitive(2i32, NonNullable),
Scalar::utf8("y", NonNullable),
],
),
]),
Some(Vec::new()),
None,
Some(vec![Scalar::struct_(
(*element_dtype).clone(),
vec![
Scalar::primitive(3i32, NonNullable),
Scalar::utf8("z", NonNullable),
],
)]),
],
element_dtype,
)
.unwrap();

let ids = buffer![0i32, 1, 2, 3].into_array();

let data = StructArray::new(
FieldNames::from(["id", "items"]),
vec![ids, items],
row_count,
Validity::NonNullable,
)
.into_array();

// Regression for nested field projection on list-of-struct: `items.a`.
let projection = get_item_list("a", get_item("items", root()));
let out = data.apply(&projection).expect("apply");

assert_eq!(
out.dtype(),
&DType::List(
Arc::new(DType::Primitive(PType::I32, NonNullable)),
Nullability::Nullable
)
);

assert_eq!(
out.scalar_at(0)
.unwrap()
.as_list()
.elements()
.unwrap()
.to_vec(),
vec![
Scalar::primitive(1i32, NonNullable),
Scalar::primitive(2i32, NonNullable),
]
);
assert!(
out.scalar_at(1)
.unwrap()
.as_list()
.elements()
.unwrap()
.is_empty()
);
assert!(out.scalar_at(2).unwrap().is_null());
assert_eq!(
out.scalar_at(3)
.unwrap()
.as_list()
.elements()
.unwrap()
.to_vec(),
vec![Scalar::primitive(3i32, NonNullable)]
);
}

#[test]
fn get_item_fixed_size_list_of_struct() {
let n_lists: usize = 3;
let list_size: u32 = 2;
let n_elements = n_lists * list_size as usize;

let struct_elems = StructArray::try_new(
FieldNames::from(["a", "b"]),
vec![
buffer![1i32, 2, 3, 4, 5, 6].into_array(),
buffer![10i64, 20, 30, 40, 50, 60].into_array(),
],
n_elements,
Validity::from_iter([true, false, true, true, false, true]),
)
.unwrap()
.into_array();

let items = FixedSizeListArray::try_new(
struct_elems,
list_size,
Validity::from_iter([true, false, true]),
n_lists,
)
.unwrap()
.into_array();

let ids = buffer![0i32, 1, 2].into_array();

let data = StructArray::new(
FieldNames::from(["id", "items"]),
vec![ids, items],
n_lists,
Validity::NonNullable,
)
.into_array();

// FixedSizeList-of-struct projection: `items.a`, including struct-level nulls inside the list.
let projection = get_item_list("a", get_item("items", root()));
let out = data.apply(&projection).expect("apply");

assert_eq!(
out.dtype(),
&DType::FixedSizeList(
Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)),
list_size,
Nullability::Nullable
)
);

assert_eq!(
out.scalar_at(0)
.unwrap()
.as_list()
.elements()
.unwrap()
.to_vec(),
vec![
Scalar::primitive(1i32, Nullability::Nullable),
Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
]
);
assert!(out.scalar_at(1).unwrap().is_null());
assert_eq!(
out.scalar_at(2)
.unwrap()
.as_list()
.elements()
.unwrap()
.to_vec(),
vec![
Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
Scalar::primitive(6i32, Nullability::Nullable),
]
);
}

#[test]
fn test_pack_get_item_rule() {
// Create: pack(a: lit(1), b: lit(2)).get_item("b")
Expand Down
Loading