16035628905

Committed 02 Jul 2025 08:56PM UTC coverage: 77.845% (+0.02%) from 77.829%

Build # 16035628905

Build Type

push

github

Committed by

web-flow

Commit Message

fix: teach StructFieldExpressionSplitter the scope of its root expr (#3743)

Simplifications that need the scope (e.g. merge) need the _right_ scope.
The scope for the root of a split expression is a structure with one
field per partition.

Signed-off-by: Daniel King <dan@spiraldb.com>

Coverage Stats

43 of 47 new or added lines in 1 file covered. (91.49%)

43280 of 55598 relevant lines covered (77.84%)

55727.81 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

94.84

/vortex-expr/src/transform/partition.rs

use std::fmt::{Display, Formatter};
use std::hash::{BuildHasher, Hash, Hasher};
use std::sync::LazyLock;

use itertools::Itertools;
use vortex_dtype::{DType, FieldName, FieldNames, Nullability, StructFields};
use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
use vortex_utils::aliases::hash_map::{DefaultHashBuilder, HashMap};

use crate::transform::immediate_access::{FieldAccesses, immediate_scope_accesses};
use crate::transform::simplify_typed::simplify_typed;
use crate::traversal::{FoldDown, FoldUp, FolderMut, MutNodeVisitor, Node, TransformResult};
use crate::{ExprRef, GetItem, ScopeDType, get_item, is_root, pack, root};

static SPLITTER_RANDOM_STATE: LazyLock<DefaultHashBuilder> =
    LazyLock::new(DefaultHashBuilder::default);

/// Partition an expression over the fields of the scope.
///
/// This returns a partitioned expression that can be push-down over each field of the scope.
/// The results of each partition evaluation can then be recombined to reproduce the result of
/// the original expression.
///
/// ## Note
///
/// This function currently respects the validity of each field in the scope, but the not validity
/// of the scope itself. The fix would be for the returned `PartitionedExpr` to include a partition
/// expression for computing the validity, or to include that expression as part of the root.
///
/// See <https://github.com/vortex-data/vortex/issues/1907>.
///
// TODO(ngates): document the behaviour of conflicting `Field::Index` and `Field::Name`.
pub fn partition(expr: ExprRef, dtype: &DType) -> VortexResult<PartitionedExpr> {
    if !matches!(dtype, DType::Struct(..)) {
        vortex_bail!("Expected a struct dtype, got {:?}", dtype);
    }
    StructFieldExpressionSplitter::split(expr, dtype)
}

// TODO(joe): replace with let expressions.
/// The result of partitioning an expression.
#[derive(Debug)]
pub struct PartitionedExpr {
    /// The root expression used to re-assemble the results.
    pub root: ExprRef,
    /// The partitions of the expression.
    pub partitions: Box<[ExprRef]>,
    /// The field names for the partitions
    pub partition_names: FieldNames,
    /// The return DTypes of each partition.
    pub partition_dtypes: Box<[DType]>,
}

impl Display for PartitionedExpr {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "root: {} {{{}}}",
            self.root,
            self.partition_names
                .iter()
                .zip(self.partitions.iter())
                .map(|(name, partition)| format!("{name}: {partition}"))
                .join(", ")
        )
    }
}

impl PartitionedExpr {
    /// Return the partition for a given field, if it exists.
    pub fn find_partition(&self, field: &FieldName) -> Option<&ExprRef> {
        self.partition_names
            .iter()
            .position(|name| name == field)
            .map(|idx| &self.partitions[idx])
    }
}

#[derive(Debug)]
struct StructFieldExpressionSplitter<'a> {
    sub_expressions: HashMap<FieldName, Vec<ExprRef>>,
    accesses: &'a FieldAccesses<'a>,
    scope_dtype: &'a StructFields,
}

impl<'a> StructFieldExpressionSplitter<'a> {
    fn new(accesses: &'a FieldAccesses<'a>, scope_dtype: &'a StructFields) -> Self {
        Self {
            sub_expressions: HashMap::new(),
            accesses,
            scope_dtype,
        }
    }

    pub(crate) fn field_idx_name(field: &FieldName, idx: usize) -> FieldName {
        let mut hasher = SPLITTER_RANDOM_STATE.build_hasher();
        field.hash(&mut hasher);
        idx.hash(&mut hasher);
        hasher.finish().to_string().into()
    }

    fn split(expr: ExprRef, dtype: &DType) -> VortexResult<PartitionedExpr> {
        let scope_dtype = match dtype {
            DType::Struct(scope_dtype, _) => scope_dtype,
            _ => vortex_bail!("Expected a struct dtype, got {:?}", dtype),
        };

        let field_accesses = immediate_scope_accesses(&expr, scope_dtype)?;

        let mut splitter = StructFieldExpressionSplitter::new(&field_accesses, scope_dtype);

        let split = expr
            .clone()
            .transform_with_context(&mut splitter, ())?
            .result();

        let mut remove_accesses: Vec<FieldName> = Vec::new();

        // Create partitions which can be passed to layout fields
        let mut partitions = Vec::with_capacity(splitter.sub_expressions.len());
        let mut partition_names = Vec::with_capacity(splitter.sub_expressions.len());
        let mut partition_dtypes = Vec::with_capacity(splitter.sub_expressions.len());
        for (name, exprs) in splitter.sub_expressions.into_iter() {
            // If there is a single expr then we don't need to `pack` this, and we must update
            // the root expr removing this access.
            let expr = if exprs.len() == 1 {
                remove_accesses.push(Self::field_idx_name(&name, 0));
                exprs.first().vortex_expect("exprs is non-empty").clone()
            } else {
                pack(
                    exprs
                        .into_iter()
                        .enumerate()
                        .map(|(idx, expr)| (Self::field_idx_name(&name, idx), expr)),
                    Nullability::NonNullable,
                )
            };

            let field_dtype = scope_dtype
                .field(&name)
                .ok_or_else(|| vortex_err!("Missing field {name}"))?;
            let field_ctx = ScopeDType::new(field_dtype);
            let expr = simplify_typed(expr.clone(), &field_ctx)?;
            let expr_dtype = expr.return_dtype(&field_ctx)?;

            partitions.push(expr);
            partition_names.push(name);
            partition_dtypes.push(expr_dtype);
        }

        let expression_access_counts = field_accesses.get(&expr).map(|ac| ac.len());
        // Ensure that there are not more accesses than partitions, we missed something
        assert!(expression_access_counts.unwrap_or(0) <= partitions.len());
        // Ensure that there are as many partitions as there are accesses/fields in the scope,
        // this will affect performance, not correctness.
        debug_assert_eq!(expression_access_counts.unwrap_or(0), partitions.len());

        let split = split
            .transform(&mut ReplaceAccessesWithChild(remove_accesses))?
            .into_inner();

        let ctx = ScopeDType::new(DType::Struct(
            StructFields::new(
                FieldNames::from(partition_names.clone()),
                partition_dtypes.clone(),
            ),
            Nullability::NonNullable,
        ));

        Ok(PartitionedExpr {
            root: simplify_typed(split, &ctx)?,
            partitions: partitions.into_boxed_slice(),
            partition_names: partition_names.into(),
            partition_dtypes: partition_dtypes.into_boxed_slice(),
        })
    }
}

impl FolderMut for StructFieldExpressionSplitter<'_> {
    type NodeTy = ExprRef;
    type Out = ExprRef;
    type Context = ();

    fn visit_down(
        &mut self,
        node: &Self::NodeTy,
        _context: Self::Context,
    ) -> VortexResult<FoldDown<ExprRef, Self::Context>> {
        // If this expression only accesses a single field, then we can skip the children
        let access = self.accesses.get(node);
        if access.as_ref().is_some_and(|a| a.len() == 1) {
            let field_name = access
                .vortex_expect("access is non-empty")
                .iter()
                .next()
                .vortex_expect("expected one field");

            let sub_exprs = self.sub_expressions.entry(field_name.clone()).or_default();
            let idx = sub_exprs.len();

            // Need to replace get_item(f, ident) with ident, making the expr relative to the child.
            let replaced = node
                .clone()
                .transform(&mut ScopeStepIntoFieldExpr(field_name.clone()))?;
            sub_exprs.push(replaced.into_inner());

            let access = get_item(
                Self::field_idx_name(field_name, idx),
                get_item(field_name.clone(), root()),
            );

            return Ok(FoldDown::SkipChildren(access));
        };

        // If the expression is an identity, then we need to partition it into the fields of the scope.
        if is_root(node) {
            let field_names = self.scope_dtype.names();

            let mut elements = Vec::with_capacity(field_names.len());

            for field_name in field_names.iter() {
                let sub_exprs = self
                    .sub_expressions
                    .entry(field_name.clone())
                    .or_insert_with(Vec::new);

                let idx = sub_exprs.len();

                sub_exprs.push(root());

                elements.push((
                    field_name.clone(),
                    // Partitions are packed into a struct of field name -> occurrence idx -> array
                    get_item(
                        Self::field_idx_name(field_name, idx),
                        get_item(field_name.clone(), root()),
                    ),
                ));
            }

            return Ok(FoldDown::SkipChildren(pack(
                elements,
                Nullability::NonNullable,
            )));
        }

        // Otherwise, continue traversing.
        Ok(FoldDown::Continue(()))
    }

    fn visit_up(
        &mut self,
        node: Self::NodeTy,
        _context: Self::Context,
        children: Vec<Self::Out>,
    ) -> VortexResult<FoldUp<Self::Out>> {
        Ok(FoldUp::Continue(node.replacing_children(children)))
    }
}

struct ScopeStepIntoFieldExpr(FieldName);

impl MutNodeVisitor for ScopeStepIntoFieldExpr {
    type NodeTy = ExprRef;

    fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<TransformResult<ExprRef>> {
        if is_root(&node) {
            Ok(TransformResult::yes(pack(
                [(self.0.clone(), root())],
                Nullability::NonNullable,
            )))
        } else {
            Ok(TransformResult::no(node))
        }
    }
}

pub(crate) struct ReplaceAccessesWithChild(Vec<FieldName>);

impl ReplaceAccessesWithChild {
    pub(crate) fn new(field_names: Vec<FieldName>) -> Self {
        Self(field_names)
    }
}

impl MutNodeVisitor for ReplaceAccessesWithChild {
    type NodeTy = ExprRef;

    fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<TransformResult<ExprRef>> {
        if let Some(item) = node.as_any().downcast_ref::<GetItem>() {
            if self.0.contains(item.field()) {
                return Ok(TransformResult::yes(item.child().clone()));
            }
        }
        Ok(TransformResult::no(node))
    }
}

#[cfg(test)]
mod tests {

    use vortex_dtype::Nullability::NonNullable;
    use vortex_dtype::PType::I32;
    use vortex_dtype::{DType, StructFields};
    use vortex_utils::aliases::hash_set::HashSet;

    use super::*;
    use crate::transform::simplify::simplify;
    use crate::transform::simplify_typed::simplify_typed;
    use crate::{Pack, and, col, get_item, lit, merge, pack, root, select};

    fn dtype() -> DType {
        DType::Struct(
            StructFields::from_iter([
                (
                    "a",
                    DType::Struct(
                        StructFields::from_iter([("a", I32.into()), ("b", DType::from(I32))]),
                        NonNullable,
                    ),
                ),
                ("b", I32.into()),
                ("c", I32.into()),
            ]),
            NonNullable,
        )
    }

    #[test]
    fn test_expr_top_level_ref() {
        let dtype = dtype();

        let expr = root();

        let split = StructFieldExpressionSplitter::split(expr, &dtype);

        assert!(split.is_ok());

        let partitioned = split.unwrap();

        assert!(partitioned.root.as_any().is::<Pack>());
        // Have a single top level pack with all fields in dtype
        assert_eq!(
            partitioned.partitions.len(),
            dtype.as_struct().unwrap().names().len()
        )
    }

    #[test]
    fn test_expr_top_level_ref_get_item_and_split() {
        let dtype = dtype();

        let expr = get_item("b", get_item("a", root()));

        let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();
        let split_a = partitioned.find_partition(&"a".into());
        assert!(split_a.is_some());
        let split_a = split_a.unwrap();

        assert_eq!(&partitioned.root, &get_item("a", root()));
        assert_eq!(&simplify(split_a.clone()).unwrap(), &get_item("b", root()));
    }

    #[test]
    fn test_expr_top_level_ref_get_item_and_split_pack() {
        let dtype = dtype();

        let expr = pack(
            [
                ("a", get_item("a", get_item("a", root()))),
                ("b", get_item("b", get_item("a", root()))),
                ("c", get_item("c", root())),
            ],
            NonNullable,
        );
        let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();

        let split_a = partitioned.find_partition(&"a".into()).unwrap();
        assert_eq!(
            &simplify(split_a.clone()).unwrap(),
            &pack(
                [
                    (
                        StructFieldExpressionSplitter::field_idx_name(&"a".into(), 0),
                        get_item("a", root())
                    ),
                    (
                        StructFieldExpressionSplitter::field_idx_name(&"a".into(), 1),
                        get_item("b", root())
                    )
                ],
                NonNullable
            )
        );
        let split_c = partitioned.find_partition(&"c".into()).unwrap();
        assert_eq!(&simplify(split_c.clone()).unwrap(), &root())
    }

    #[test]
    fn test_expr_top_level_ref_get_item_add() {
        let dtype = dtype();

        let expr = and(get_item("b", get_item("a", root())), lit(1));
        let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();

        // Whole expr is a single split
        assert_eq!(partitioned.partitions.len(), 1);
    }

    #[test]
    fn test_expr_top_level_ref_get_item_add_cannot_split() {
        let dtype = dtype();

        let expr = and(get_item("b", get_item("a", root())), get_item("b", root()));
        let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();

        // One for id.a and id.b
        assert_eq!(partitioned.partitions.len(), 2);
    }

    // Test that typed_simplify removes select and partition precise
    #[test]
    fn test_expr_partition_many_occurrences_of_field() {
        let dtype = dtype();

        let expr = and(
            get_item("b", get_item("a", root())),
            select(vec!["a".into(), "b".into()], root()),
        );
        let expr = simplify_typed(expr, &ScopeDType::new(dtype.clone())).unwrap();
        let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();

        // One for id.a and id.b
        assert_eq!(partitioned.partitions.len(), 2);

        // This fetches [].$c which is unused, however a previous optimisation should replace select
        // with get_item and pack removing this field.
        assert_eq!(
            &partitioned.root,
            &and(
                get_item(
                    StructFieldExpressionSplitter::field_idx_name(&"a".into(), 0),
                    get_item("a", root())
                ),
                pack(
                    [
                        (
                            "a",
                            get_item(
                                StructFieldExpressionSplitter::field_idx_name(&"a".into(), 1),
                                get_item("a", root())
                            )
                        ),
                        ("b", get_item("b", root()))
                    ],
                    NonNullable
                )
            )
        )
    }

    #[test]
    fn test_expr_merge() {
        let dtype = dtype();

        let expr = merge(
            [col("a"), pack([("b", col("b"))], NonNullable)],
            NonNullable,
        );

        let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();
        let expected = pack(
            [
                ("a", get_item("a", col("a"))),
                ("b", get_item("b", col("b"))),
            ],
            NonNullable,
        );
        assert_eq!(
            &partitioned.root, &expected,
            "{} {}",
            partitioned.root, expected
        );
        let expected = [root(), pack([("b", root())], NonNullable)]
            .into_iter()
            .collect::<HashSet<_>>();
        assert_eq!(
            &partitioned
                .partitions
                .clone()
                .into_iter()
                .collect::<HashSet<_>>(),
            &expected,
            "{} {}",
            partitioned.partitions.iter().join(";"),
            expected.iter().join(";")
        );
    }
}

1	use std::fmt::{Display, Formatter};
2	use std::hash::{BuildHasher, Hash, Hasher};
3	use std::sync::LazyLock;
4
5	use itertools::Itertools;
6	use vortex_dtype::{DType, FieldName, FieldNames, Nullability, StructFields};
7	use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
8	use vortex_utils::aliases::hash_map::{DefaultHashBuilder, HashMap};
9
10	use crate::transform::immediate_access::{FieldAccesses, immediate_scope_accesses};
11	use crate::transform::simplify_typed::simplify_typed;
12	use crate::traversal::{FoldDown, FoldUp, FolderMut, MutNodeVisitor, Node, TransformResult};
13	use crate::{ExprRef, GetItem, ScopeDType, get_item, is_root, pack, root};
14
15	static SPLITTER_RANDOM_STATE: LazyLock<DefaultHashBuilder> =
16	LazyLock::new(DefaultHashBuilder::default);
17
18	/// Partition an expression over the fields of the scope.
19	///
20	/// This returns a partitioned expression that can be push-down over each field of the scope.
21	/// The results of each partition evaluation can then be recombined to reproduce the result of
22	/// the original expression.
23	///
24	/// ## Note
25	///
26	/// This function currently respects the validity of each field in the scope, but the not validity
27	/// of the scope itself. The fix would be for the returned `PartitionedExpr` to include a partition
28	/// expression for computing the validity, or to include that expression as part of the root.
29	///
30	/// See <https://github.com/vortex-data/vortex/issues/1907>.
31	///
32	// TODO(ngates): document the behaviour of conflicting `Field::Index` and `Field::Name`.
33	pub fn partition(expr: ExprRef, dtype: &DType) -> VortexResult<PartitionedExpr> {	1,218✔
34	if !matches!(dtype, DType::Struct(..)) {	1,218✔
35	vortex_bail!("Expected a struct dtype, got {:?}", dtype);	×
36	}	1,218✔
37	StructFieldExpressionSplitter::split(expr, dtype)	1,218✔
38	}	1,218✔
39
40	// TODO(joe): replace with let expressions.
41	/// The result of partitioning an expression.
42	#[derive(Debug)]
43	pub struct PartitionedExpr {
44	/// The root expression used to re-assemble the results.
45	pub root: ExprRef,
46	/// The partitions of the expression.
47	pub partitions: Box<[ExprRef]>,
48	/// The field names for the partitions
49	pub partition_names: FieldNames,
50	/// The return DTypes of each partition.
51	pub partition_dtypes: Box<[DType]>,
52	}
53
54	impl Display for PartitionedExpr {
55	fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {	×
56	write!(	×
57	f,	×
58	"root: {} {{{}}}",	×
59	self.root,	×
60	self.partition_names	×
61	.iter()	×
62	.zip(self.partitions.iter())	×
63	.map(\|(name, partition)\| format!("{name}: {partition}"))	×
64	.join(", ")	×
65	)	×
66	}	×
67	}
68
69	impl PartitionedExpr {
70	/// Return the partition for a given field, if it exists.
71	pub fn find_partition(&self, field: &FieldName) -> Option<&ExprRef> {	3✔
72	self.partition_names	3✔
73	.iter()	3✔
74	.position(\|name\| name == field)	4✔
75	.map(\|idx\| &self.partitions[idx])	3✔
76	}	3✔
77	}
78
79	#[derive(Debug)]
80	struct StructFieldExpressionSplitter<'a> {
81	sub_expressions: HashMap<FieldName, Vec<ExprRef>>,
82	accesses: &'a FieldAccesses<'a>,
83	scope_dtype: &'a StructFields,
84	}
85
86	impl<'a> StructFieldExpressionSplitter<'a> {
87	fn new(accesses: &'a FieldAccesses<'a>, scope_dtype: &'a StructFields) -> Self {	1,225✔
88	Self {	1,225✔
89	sub_expressions: HashMap::new(),	1,225✔
90	accesses,	1,225✔
91	scope_dtype,	1,225✔
92	}	1,225✔
93	}	1,225✔
94
95	pub(crate) fn field_idx_name(field: &FieldName, idx: usize) -> FieldName {	3,184✔
96	let mut hasher = SPLITTER_RANDOM_STATE.build_hasher();	3,184✔
97	field.hash(&mut hasher);	3,184✔
98	idx.hash(&mut hasher);	3,184✔
99	hasher.finish().to_string().into()	3,184✔
100	}	3,184✔
101
102	fn split(expr: ExprRef, dtype: &DType) -> VortexResult<PartitionedExpr> {	1,225✔
103	let scope_dtype = match dtype {	1,225✔
104	DType::Struct(scope_dtype, _) => scope_dtype,	1,225✔
105	_ => vortex_bail!("Expected a struct dtype, got {:?}", dtype),	×
106	};
107
108	let field_accesses = immediate_scope_accesses(&expr, scope_dtype)?;	1,225✔
109
110	let mut splitter = StructFieldExpressionSplitter::new(&field_accesses, scope_dtype);	1,225✔
111
112	let split = expr	1,225✔
113	.clone()	1,225✔
114	.transform_with_context(&mut splitter, ())?	1,225✔
115	.result();	1,225✔
116		1,225✔
117	let mut remove_accesses: Vec<FieldName> = Vec::new();	1,225✔
118		1,225✔
119	// Create partitions which can be passed to layout fields	1,225✔
120	let mut partitions = Vec::with_capacity(splitter.sub_expressions.len());	1,225✔
121	let mut partition_names = Vec::with_capacity(splitter.sub_expressions.len());	1,225✔
122	let mut partition_dtypes = Vec::with_capacity(splitter.sub_expressions.len());	1,225✔
123	for (name, exprs) in splitter.sub_expressions.into_iter() {	1,567✔
124	// If there is a single expr then we don't need to `pack` this, and we must update
125	// the root expr removing this access.
126	let expr = if exprs.len() == 1 {	1,567✔
127	remove_accesses.push(Self::field_idx_name(&name, 0));	1,544✔
128	exprs.first().vortex_expect("exprs is non-empty").clone()	1,544✔
129	} else {
130	pack(	23✔
131	exprs	23✔
132	.into_iter()	23✔
133	.enumerate()	23✔
134	.map(\|(idx, expr)\| (Self::field_idx_name(&name, idx), expr)),	46✔
135	Nullability::NonNullable,	23✔
136	)	23✔
137	};
138
139	let field_dtype = scope_dtype	1,567✔
140	.field(&name)	1,567✔
141	.ok_or_else(\|\| vortex_err!("Missing field {name}"))?;	1,567✔
142	let field_ctx = ScopeDType::new(field_dtype);	1,567✔
143	let expr = simplify_typed(expr.clone(), &field_ctx)?;	1,567✔
144	let expr_dtype = expr.return_dtype(&field_ctx)?;	1,567✔
145
146	partitions.push(expr);	1,567✔
147	partition_names.push(name);	1,567✔
148	partition_dtypes.push(expr_dtype);	1,567✔
149	}
150
151	let expression_access_counts = field_accesses.get(&expr).map(\|ac\| ac.len());	1,225✔
152	// Ensure that there are not more accesses than partitions, we missed something	1,225✔
153	assert!(expression_access_counts.unwrap_or(0) <= partitions.len());	1,225✔
154	// Ensure that there are as many partitions as there are accesses/fields in the scope,
155	// this will affect performance, not correctness.
156	debug_assert_eq!(expression_access_counts.unwrap_or(0), partitions.len());	1,225✔
157
158	let split = split	1,225✔
159	.transform(&mut ReplaceAccessesWithChild(remove_accesses))?	1,225✔
160	.into_inner();	1,225✔
161		1,225✔
162	let ctx = ScopeDType::new(DType::Struct(	1,225✔
163	StructFields::new(	1,225✔
164	FieldNames::from(partition_names.clone()),	1,225✔
165	partition_dtypes.clone(),	1,225✔
166	),	1,225✔
167	Nullability::NonNullable,	1,225✔
168	));	1,225✔
169		1,225✔
170	Ok(PartitionedExpr {	1,225✔
171	root: simplify_typed(split, &ctx)?,	1,225✔
172	partitions: partitions.into_boxed_slice(),	1,225✔
173	partition_names: partition_names.into(),	1,225✔
174	partition_dtypes: partition_dtypes.into_boxed_slice(),	1,225✔
175	})
176	}	1,225✔
177	}
178
179	impl FolderMut for StructFieldExpressionSplitter<'_> {
180	type NodeTy = ExprRef;
181	type Out = ExprRef;
182	type Context = ();
183
184	fn visit_down(	1,656✔
185	&mut self,	1,656✔
186	node: &Self::NodeTy,	1,656✔
187	_context: Self::Context,	1,656✔
188	) -> VortexResult<FoldDown<ExprRef, Self::Context>> {	1,656✔
189	// If this expression only accesses a single field, then we can skip the children	1,656✔
190	let access = self.accesses.get(node);	1,656✔
191	if access.as_ref().is_some_and(\|a\| a.len() == 1) {	1,656✔
192	let field_name = access	1,062✔
193	.vortex_expect("access is non-empty")	1,062✔
194	.iter()	1,062✔
195	.next()	1,062✔
196	.vortex_expect("expected one field");	1,062✔
197		1,062✔
198	let sub_exprs = self.sub_expressions.entry(field_name.clone()).or_default();	1,062✔
199	let idx = sub_exprs.len();	1,062✔
200
201	// Need to replace get_item(f, ident) with ident, making the expr relative to the child.
202	let replaced = node	1,062✔
203	.clone()	1,062✔
204	.transform(&mut ScopeStepIntoFieldExpr(field_name.clone()))?;	1,062✔
205	sub_exprs.push(replaced.into_inner());	1,062✔
206		1,062✔
207	let access = get_item(	1,062✔
208	Self::field_idx_name(field_name, idx),	1,062✔
209	get_item(field_name.clone(), root()),	1,062✔
210	);	1,062✔
211		1,062✔
212	return Ok(FoldDown::SkipChildren(access));	1,062✔
213	};	594✔
214		594✔
215	// If the expression is an identity, then we need to partition it into the fields of the scope.	594✔
216	if is_root(node) {	594✔
217	let field_names = self.scope_dtype.names();	211✔
218		211✔
219	let mut elements = Vec::with_capacity(field_names.len());	211✔
220
221	for field_name in field_names.iter() {	528✔
222	let sub_exprs = self	528✔
223	.sub_expressions	528✔
224	.entry(field_name.clone())	528✔
225	.or_insert_with(Vec::new);	528✔
226		528✔
227	let idx = sub_exprs.len();	528✔
228		528✔
229	sub_exprs.push(root());	528✔
230		528✔
231	elements.push((	528✔
232	field_name.clone(),	528✔
233	// Partitions are packed into a struct of field name -> occurrence idx -> array	528✔
234	get_item(	528✔
235	Self::field_idx_name(field_name, idx),	528✔
236	get_item(field_name.clone(), root()),	528✔
237	),	528✔
238	));	528✔
239	}	528✔
240
241	return Ok(FoldDown::SkipChildren(pack(	211✔
242	elements,	211✔
243	Nullability::NonNullable,	211✔
244	)));	211✔
245	}	383✔
246		383✔
247	// Otherwise, continue traversing.	383✔
248	Ok(FoldDown::Continue(()))	383✔
249	}	1,656✔
250
251	fn visit_up(	383✔
252	&mut self,	383✔
253	node: Self::NodeTy,	383✔
254	_context: Self::Context,	383✔
255	children: Vec<Self::Out>,	383✔
256	) -> VortexResult<FoldUp<Self::Out>> {	383✔
257	Ok(FoldUp::Continue(node.replacing_children(children)))	383✔
258	}	383✔
259	}
260
261	struct ScopeStepIntoFieldExpr(FieldName);
262
263	impl MutNodeVisitor for ScopeStepIntoFieldExpr {
264	type NodeTy = ExprRef;
265
266	fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<TransformResult<ExprRef>> {	3,288✔
267	if is_root(&node) {	3,288✔
268	Ok(TransformResult::yes(pack(	1,083✔
269	[(self.0.clone(), root())],	1,083✔
270	Nullability::NonNullable,	1,083✔
271	)))	1,083✔
272	} else {
273	Ok(TransformResult::no(node))	2,205✔
274	}
275	}	3,288✔
276	}
277
278	pub(crate) struct ReplaceAccessesWithChild(Vec<FieldName>);
279
280	impl ReplaceAccessesWithChild {
281	pub(crate) fn new(field_names: Vec<FieldName>) -> Self {	67✔
282	Self(field_names)	67✔
283	}	67✔
284	}
285
286	impl MutNodeVisitor for ReplaceAccessesWithChild {
287	type NodeTy = ExprRef;
288
289	fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<TransformResult<ExprRef>> {	5,638✔
290	if let Some(item) = node.as_any().downcast_ref::<GetItem>() {	5,638✔
291	if self.0.contains(item.field()) {	3,294✔
292	return Ok(TransformResult::yes(item.child().clone()));	1,612✔
293	}	1,682✔
294	}	2,344✔
295	Ok(TransformResult::no(node))	4,026✔
296	}	5,638✔
297	}
298
299	#[cfg(test)]
300	mod tests {
301
302	use vortex_dtype::Nullability::NonNullable;
303	use vortex_dtype::PType::I32;
304	use vortex_dtype::{DType, StructFields};
305	use vortex_utils::aliases::hash_set::HashSet;
306
307	use super::*;
308	use crate::transform::simplify::simplify;
309	use crate::transform::simplify_typed::simplify_typed;
310	use crate::{Pack, and, col, get_item, lit, merge, pack, root, select};
311
312	fn dtype() -> DType {	7✔
313	DType::Struct(	7✔
314	StructFields::from_iter([	7✔
315	(	7✔
316	"a",	7✔
317	DType::Struct(	7✔
318	StructFields::from_iter([("a", I32.into()), ("b", DType::from(I32))]),	7✔
319	NonNullable,	7✔
320	),	7✔
321	),	7✔
322	("b", I32.into()),	7✔
323	("c", I32.into()),	7✔
324	]),	7✔
325	NonNullable,	7✔
326	)	7✔
327	}	7✔
328
329	#[test]
330	fn test_expr_top_level_ref() {	1✔
331	let dtype = dtype();	1✔
332		1✔
333	let expr = root();	1✔
334		1✔
335	let split = StructFieldExpressionSplitter::split(expr, &dtype);	1✔
336		1✔
337	assert!(split.is_ok());	1✔
338
339	let partitioned = split.unwrap();	1✔
340		1✔
341	assert!(partitioned.root.as_any().is::<Pack>());	1✔
342	// Have a single top level pack with all fields in dtype
343	assert_eq!(	1✔
344	partitioned.partitions.len(),	1✔
345	dtype.as_struct().unwrap().names().len()	1✔
346	)	1✔
347	}	1✔
348
349	#[test]
350	fn test_expr_top_level_ref_get_item_and_split() {	1✔
351	let dtype = dtype();	1✔
352		1✔
353	let expr = get_item("b", get_item("a", root()));	1✔
354		1✔
355	let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();	1✔
356	let split_a = partitioned.find_partition(&"a".into());	1✔
357	assert!(split_a.is_some());	1✔
358	let split_a = split_a.unwrap();	1✔
359		1✔
360	assert_eq!(&partitioned.root, &get_item("a", root()));	1✔
361	assert_eq!(&simplify(split_a.clone()).unwrap(), &get_item("b", root()));	1✔
362	}	1✔
363
364	#[test]
365	fn test_expr_top_level_ref_get_item_and_split_pack() {	1✔
366	let dtype = dtype();	1✔
367		1✔
368	let expr = pack(	1✔
369	[	1✔
370	("a", get_item("a", get_item("a", root()))),	1✔
371	("b", get_item("b", get_item("a", root()))),	1✔
372	("c", get_item("c", root())),	1✔
373	],	1✔
374	NonNullable,	1✔
375	);	1✔
376	let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();	1✔
377		1✔
378	let split_a = partitioned.find_partition(&"a".into()).unwrap();	1✔
379	assert_eq!(	1✔
380	&simplify(split_a.clone()).unwrap(),	1✔
381	&pack(	1✔
382	[	1✔
383	(	1✔
384	StructFieldExpressionSplitter::field_idx_name(&"a".into(), 0),	1✔
385	get_item("a", root())	1✔
386	),	1✔
387	(	1✔
388	StructFieldExpressionSplitter::field_idx_name(&"a".into(), 1),	1✔
389	get_item("b", root())	1✔
390	)	1✔
391	],	1✔
392	NonNullable	1✔
393	)	1✔
394	);	1✔
395	let split_c = partitioned.find_partition(&"c".into()).unwrap();	1✔
396	assert_eq!(&simplify(split_c.clone()).unwrap(), &root())	1✔
397	}	1✔
398
399	#[test]
400	fn test_expr_top_level_ref_get_item_add() {	1✔
401	let dtype = dtype();	1✔
402		1✔
403	let expr = and(get_item("b", get_item("a", root())), lit(1));	1✔
404	let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();	1✔
405		1✔
406	// Whole expr is a single split	1✔
407	assert_eq!(partitioned.partitions.len(), 1);	1✔
408	}	1✔
409
410	#[test]
411	fn test_expr_top_level_ref_get_item_add_cannot_split() {	1✔
412	let dtype = dtype();	1✔
413		1✔
414	let expr = and(get_item("b", get_item("a", root())), get_item("b", root()));	1✔
415	let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();	1✔
416		1✔
417	// One for id.a and id.b	1✔
418	assert_eq!(partitioned.partitions.len(), 2);	1✔
419	}	1✔
420
421	// Test that typed_simplify removes select and partition precise
422	#[test]
423	fn test_expr_partition_many_occurrences_of_field() {	1✔
424	let dtype = dtype();	1✔
425		1✔
426	let expr = and(	1✔
427	get_item("b", get_item("a", root())),	1✔
428	select(vec!["a".into(), "b".into()], root()),	1✔
429	);	1✔
430	let expr = simplify_typed(expr, &ScopeDType::new(dtype.clone())).unwrap();	1✔
431	let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();	1✔
432		1✔
433	// One for id.a and id.b	1✔
434	assert_eq!(partitioned.partitions.len(), 2);	1✔
435
436	// This fetches [].$c which is unused, however a previous optimisation should replace select
437	// with get_item and pack removing this field.
438	assert_eq!(	1✔
439	&partitioned.root,	1✔
440	&and(	1✔
441	get_item(	1✔
442	StructFieldExpressionSplitter::field_idx_name(&"a".into(), 0),	1✔
443	get_item("a", root())	1✔
444	),	1✔
445	pack(	1✔
446	[	1✔
447	(	1✔
448	"a",	1✔
449	get_item(	1✔
450	StructFieldExpressionSplitter::field_idx_name(&"a".into(), 1),	1✔
451	get_item("a", root())	1✔
452	)	1✔
453	),	1✔
454	("b", get_item("b", root()))	1✔
455	],	1✔
456	NonNullable	1✔
457	)	1✔
458	)	1✔
459	)	1✔
460	}	1✔
461
462	#[test]
463	fn test_expr_merge() {	1✔
464	let dtype = dtype();	1✔
465		1✔
466	let expr = merge(	1✔
467	[col("a"), pack([("b", col("b"))], NonNullable)],	1✔
468	NonNullable,	1✔
469	);	1✔
470		1✔
471	let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();	1✔
472	let expected = pack(	1✔
473	[	1✔
474	("a", get_item("a", col("a"))),	1✔
475	("b", get_item("b", col("b"))),	1✔
476	],	1✔
477	NonNullable,	1✔
478	);	1✔
479	assert_eq!(	1✔
480	&partitioned.root, &expected,	1✔
NEW 481	"{} {}",	×
482	partitioned.root, expected
483	);
484	let expected = [root(), pack([("b", root())], NonNullable)]	1✔
485	.into_iter()	1✔
486	.collect::<HashSet<_>>();	1✔
487	assert_eq!(	1✔
488	&partitioned	1✔
489	.partitions	1✔
490	.clone()	1✔
491	.into_iter()	1✔
492	.collect::<HashSet<_>>(),	1✔
493	&expected,	1✔
NEW 494	"{} {}",	×
NEW 495	partitioned.partitions.iter().join(";"),	×
NEW 496	expected.iter().join(";")	×
497	);
498	}	1✔
499	}

vortex-data / vortex / 16035628905

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous