From 9c20e63aad5b939fe259d7117b29c963848a96e0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 13 Jan 2025 15:07:29 +0000 Subject: [PATCH 1/8] wip --- vortex-expr/src/forms/nnf.rs | 2 +- vortex-expr/src/lib.rs | 3 +-- vortex-expr/src/transform/partition.rs | 2 +- vortex-expr/src/transform/simplify.rs | 2 +- vortex-expr/src/traversal/mod.rs | 4 ++++ vortex-expr/src/traversal/visitor.rs | 2 ++ vortex-layout/src/layouts/mod.rs | 1 + vortex-layout/src/layouts/visitor.rs | 21 +++++++++++++++++++++ 8 files changed, 32 insertions(+), 5 deletions(-) create mode 100644 vortex-layout/src/layouts/visitor.rs diff --git a/vortex-expr/src/forms/nnf.rs b/vortex-expr/src/forms/nnf.rs index 4ecc314db7..8292bfa607 100644 --- a/vortex-expr/src/forms/nnf.rs +++ b/vortex-expr/src/forms/nnf.rs @@ -1,6 +1,6 @@ use vortex_error::VortexResult; -use crate::traversal::{FoldChildren, FoldDown, FoldUp, FolderMut, Node as _}; +use crate::traversal::{FoldChildren, FoldDown, FoldUp, FolderMut, NodeMut}; use crate::{not, BinaryExpr, ExprRef, Not, Operator}; /// Return an equivalent expression in Negative Normal Form (NNF). diff --git a/vortex-expr/src/lib.rs b/vortex-expr/src/lib.rs index 645319f98f..facaa0f932 100644 --- a/vortex-expr/src/lib.rs +++ b/vortex-expr/src/lib.rs @@ -23,8 +23,7 @@ pub mod pruning; mod row_filter; mod select; pub mod transform; -#[allow(dead_code)] -mod traversal; +pub mod traversal; pub use binary::*; pub use column::*; diff --git a/vortex-expr/src/transform/partition.rs b/vortex-expr/src/transform/partition.rs index 08bb57a3e5..49bbb79a93 100644 --- a/vortex-expr/src/transform/partition.rs +++ b/vortex-expr/src/transform/partition.rs @@ -4,7 +4,7 @@ use vortex_array::aliases::hash_set::HashSet; use vortex_dtype::{FieldName, StructDType}; use vortex_error::{VortexExpect, VortexResult}; -use crate::traversal::{FoldChildren, FoldDown, FoldUp, Folder, FolderMut, Node}; +use crate::traversal::{FoldChildren, FoldDown, FoldUp, Folder, FolderMut, Node, NodeMut}; use crate::{get_item, ident, pack, ExprRef, GetItem, Identity, Select, SelectField}; /// Partition an expression over the fields of the scope. diff --git a/vortex-expr/src/transform/simplify.rs b/vortex-expr/src/transform/simplify.rs index 49df3532cf..eb02604af4 100644 --- a/vortex-expr/src/transform/simplify.rs +++ b/vortex-expr/src/transform/simplify.rs @@ -1,6 +1,6 @@ use vortex_error::VortexResult; -use crate::traversal::{FoldChildren, FoldUp, FolderMut, Node}; +use crate::traversal::{FoldChildren, FoldUp, FolderMut, NodeMut}; use crate::{get_item, ident, Column, ExprRef, GetItem, Pack}; pub fn simplify(e: ExprRef) -> VortexResult { diff --git a/vortex-expr/src/traversal/mod.rs b/vortex-expr/src/traversal/mod.rs index 0c4206c25c..a43ecbea02 100644 --- a/vortex-expr/src/traversal/mod.rs +++ b/vortex-expr/src/traversal/mod.rs @@ -188,7 +188,9 @@ pub trait Node: Sized { visitor: &mut V, context: V::Context, ) -> VortexResult>; +} +pub trait NodeMut: Sized { fn transform>( self, _visitor: &mut V, @@ -248,7 +250,9 @@ impl Node for ExprRef { visitor.visit_up(self, context, children) } +} +impl NodeMut for ExprRef { // A pre-order transform, with an option to ignore sub-tress (using visit_down). fn transform>( self, diff --git a/vortex-expr/src/traversal/visitor.rs b/vortex-expr/src/traversal/visitor.rs index 298589ddaa..533ce6b923 100644 --- a/vortex-expr/src/traversal/visitor.rs +++ b/vortex-expr/src/traversal/visitor.rs @@ -37,6 +37,7 @@ where } } +#[allow(dead_code)] pub fn pre_order_visit_up<'a, T: 'a + Node>( f: impl FnMut(&'a T) -> VortexResult, ) -> impl NodeVisitor<'a, NodeTy = T> { @@ -47,6 +48,7 @@ pub fn pre_order_visit_up<'a, T: 'a + Node>( } } +#[allow(dead_code)] pub fn pre_order_visit_down<'a, T: 'a + Node>( f: impl FnMut(&'a T) -> VortexResult, ) -> impl NodeVisitor<'a, NodeTy = T> { diff --git a/vortex-layout/src/layouts/mod.rs b/vortex-layout/src/layouts/mod.rs index 92c9d9c103..247530a545 100644 --- a/vortex-layout/src/layouts/mod.rs +++ b/vortex-layout/src/layouts/mod.rs @@ -2,3 +2,4 @@ pub mod chunked; pub mod flat; pub mod struct_; +mod visitor; diff --git a/vortex-layout/src/layouts/visitor.rs b/vortex-layout/src/layouts/visitor.rs new file mode 100644 index 0000000000..f4a0a31d93 --- /dev/null +++ b/vortex-layout/src/layouts/visitor.rs @@ -0,0 +1,21 @@ +use vortex_error::VortexResult; +use vortex_expr::traversal::{FoldUp, Folder, FolderMut, Node, NodeVisitor, TraversalOrder}; + +use crate::LayoutEncoding; + +impl Node for LayoutEncoding { + fn accept<'a, V: NodeVisitor<'a, NodeTy = Self>>( + &'a self, + _visitor: &mut V, + ) -> VortexResult { + todo!() + } + + fn accept_with_context<'a, V: Folder<'a, NodeTy = Self>>( + &'a self, + visitor: &mut V, + context: V::Context, + ) -> VortexResult> { + todo!() + } +} From 48474432c7650e48ed8a66ebcbdd5f8391438e12 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 13 Jan 2025 16:10:38 +0000 Subject: [PATCH 2/8] wip --- vortex-expr/src/traversal/mod.rs | 61 ++++++++++++++++++++- vortex-layout/src/layouts/chunked/reader.rs | 4 ++ vortex-layout/src/layouts/flat/reader.rs | 4 ++ vortex-layout/src/layouts/mod.rs | 1 - vortex-layout/src/layouts/struct_/reader.rs | 8 +++ vortex-layout/src/layouts/visitor.rs | 21 ------- vortex-layout/src/lib.rs | 1 + vortex-layout/src/reader.rs | 7 +++ vortex-layout/src/visitor.rs | 12 ++++ 9 files changed, 94 insertions(+), 25 deletions(-) delete mode 100644 vortex-layout/src/layouts/visitor.rs create mode 100644 vortex-layout/src/visitor.rs diff --git a/vortex-expr/src/traversal/mod.rs b/vortex-expr/src/traversal/mod.rs index a43ecbea02..a7e3cc822e 100644 --- a/vortex-expr/src/traversal/mod.rs +++ b/vortex-expr/src/traversal/mod.rs @@ -1,6 +1,8 @@ mod references; mod visitor; +use std::sync::Arc; + use itertools::Itertools; pub use references::ReferenceCollector; use vortex_error::VortexResult; @@ -185,11 +187,15 @@ pub trait Node: Sized { fn accept_with_context<'a, V: Folder<'a, NodeTy = Self>>( &'a self, - visitor: &mut V, - context: V::Context, + _visitor: &mut V, + _context: V::Context, ) -> VortexResult>; } +pub trait DynNode { + fn arc_children(&self) -> Vec<&Arc>; +} + pub trait NodeMut: Sized { fn transform>( self, @@ -203,6 +209,55 @@ pub trait NodeMut: Sized { ) -> VortexResult>; } +impl Node for Arc { + // A pre-order traversal. + fn accept<'a, V: NodeVisitor<'a, NodeTy = Arc>>( + &'a self, + visitor: &mut V, + ) -> VortexResult { + let mut ord = visitor.visit_down(self)?; + if ord == TraversalOrder::Stop { + return Ok(TraversalOrder::Stop); + } + if ord == TraversalOrder::Skip { + return Ok(TraversalOrder::Continue); + } + for child in self.arc_children() { + if ord != TraversalOrder::Continue { + return Ok(ord); + } + ord = child.accept(visitor)?; + } + if ord == TraversalOrder::Stop { + return Ok(TraversalOrder::Stop); + } + visitor.visit_up(self) + } + + fn accept_with_context<'a, V: Folder<'a, NodeTy = Self>>( + &'a self, + visitor: &mut V, + context: V::Context, + ) -> VortexResult> { + let children = match visitor.visit_down(self, context.clone())? { + FoldDown::Stop(out) => return Ok(FoldUp::Stop(out)), + FoldDown::SkipChildren => FoldChildren::Skipped, + FoldDown::Continue(child_context) => { + let mut new_children = Vec::with_capacity(self.arc_children().len()); + for child in self.arc_children() { + match child.accept_with_context(visitor, child_context.clone())? { + FoldUp::Stop(out) => return Ok(FoldUp::Stop(out)), + FoldUp::Continue(out) => new_children.push(out), + } + } + FoldChildren::Children(new_children) + } + }; + + visitor.visit_up(self, context, children) + } +} + impl Node for ExprRef { // A pre-order traversal. fn accept<'a, V: NodeVisitor<'a, NodeTy = ExprRef>>( @@ -340,7 +395,7 @@ mod tests { use vortex_error::VortexResult; use crate::traversal::visitor::pre_order_visit_down; - use crate::traversal::{MutNodeVisitor, Node, NodeVisitor, TransformResult, TraversalOrder}; + use crate::traversal::{MutNodeVisitor, NodeMut, NodeVisitor, TransformResult, TraversalOrder}; use crate::{ BinaryExpr, Column, ExprRef, FieldName, Literal, Operator, VortexExpr, VortexExprExt, }; diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index 13c5d0a0e1..e549345b2a 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -119,4 +119,8 @@ impl LayoutReader for ChunkedReader { fn layout(&self) -> &LayoutData { &self.layout } + + fn children(&self) -> VortexResult>> { + (0..self.nchunks()).map(|i| self.child(i)).collect() + } } diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 7209b2b5a1..8f76cb6985 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -44,4 +44,8 @@ impl LayoutReader for FlatReader { fn layout(&self) -> &LayoutData { &self.layout } + + fn children(&self) -> VortexResult>> { + Ok(vec![]) + } } diff --git a/vortex-layout/src/layouts/mod.rs b/vortex-layout/src/layouts/mod.rs index 247530a545..92c9d9c103 100644 --- a/vortex-layout/src/layouts/mod.rs +++ b/vortex-layout/src/layouts/mod.rs @@ -2,4 +2,3 @@ pub mod chunked; pub mod flat; pub mod struct_; -mod visitor; diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 2af0964016..acb97e4214 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -62,6 +62,10 @@ impl StructReader { .vortex_expect("Struct layout must have a struct DType, verified at construction") } + pub(crate) fn nchildren(&self) -> usize { + self.field_readers.len() + } + /// Return the child reader for the chunk. pub(crate) fn child(&self, field: &Field) -> VortexResult<&Arc> { let idx = match field { @@ -84,4 +88,8 @@ impl LayoutReader for StructReader { fn layout(&self) -> &LayoutData { &self.layout } + + fn children(&self) -> VortexResult>> { + (0..self.nchildren()).map(|idx| self.child(idx)).collect() + } } diff --git a/vortex-layout/src/layouts/visitor.rs b/vortex-layout/src/layouts/visitor.rs deleted file mode 100644 index f4a0a31d93..0000000000 --- a/vortex-layout/src/layouts/visitor.rs +++ /dev/null @@ -1,21 +0,0 @@ -use vortex_error::VortexResult; -use vortex_expr::traversal::{FoldUp, Folder, FolderMut, Node, NodeVisitor, TraversalOrder}; - -use crate::LayoutEncoding; - -impl Node for LayoutEncoding { - fn accept<'a, V: NodeVisitor<'a, NodeTy = Self>>( - &'a self, - _visitor: &mut V, - ) -> VortexResult { - todo!() - } - - fn accept_with_context<'a, V: Folder<'a, NodeTy = Self>>( - &'a self, - visitor: &mut V, - context: V::Context, - ) -> VortexResult> { - todo!() - } -} diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index d3000457ae..04b75c57a1 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -11,6 +11,7 @@ mod reader; pub use reader::*; pub mod segments; pub mod strategies; +mod visitor; /// The layout ID for a flat layout pub(crate) const FLAT_LAYOUT_ID: LayoutId = LayoutId(1); diff --git a/vortex-layout/src/reader.rs b/vortex-layout/src/reader.rs index 549ffd4968..648548d1c1 100644 --- a/vortex-layout/src/reader.rs +++ b/vortex-layout/src/reader.rs @@ -5,6 +5,7 @@ use vortex_array::stats::{Stat, StatsSet}; use vortex_array::ArrayData; use vortex_dtype::{DType, FieldPath}; use vortex_error::VortexResult; +use vortex_expr::traversal::DynNode; use vortex_expr::ExprRef; use vortex_scan::RowMask; @@ -18,12 +19,18 @@ use crate::LayoutData; pub trait LayoutReader: Send + Sync + ExprEvaluator + StatsEvaluator { /// Returns the [`LayoutData`] of this reader. fn layout(&self) -> &LayoutData; + + fn children(&self) -> VortexResult>>; } impl LayoutReader for Arc { fn layout(&self) -> &LayoutData { self.as_ref().layout() } + + fn children(&self) -> VortexResult>> { + self.as_ref().arc_children() + } } /// A trait for evaluating expressions against a [`LayoutReader`]. diff --git a/vortex-layout/src/visitor.rs b/vortex-layout/src/visitor.rs new file mode 100644 index 0000000000..373393346f --- /dev/null +++ b/vortex-layout/src/visitor.rs @@ -0,0 +1,12 @@ +use std::sync::Arc; + +use vortex_error::VortexResult; +use vortex_expr::traversal::DynNode; + +use crate::LayoutReader; + +impl DynNode for dyn LayoutReader { + fn arc_children(&self) -> VortexResult>> { + self.children() + } +} From 905268cb57e29ece76f6c173e3207ce799c4e73a Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 13 Jan 2025 17:00:40 +0000 Subject: [PATCH 3/8] export the Node trait and use it for layout reader --- vortex-expr/src/lib.rs | 8 ++- vortex-expr/src/traversal/mod.rs | 57 ++----------------- .../src/layouts/struct_/eval_expr.rs | 9 +++ vortex-layout/src/layouts/struct_/reader.rs | 8 ++- vortex-layout/src/visitor.rs | 28 ++++++++- 5 files changed, 54 insertions(+), 56 deletions(-) diff --git a/vortex-expr/src/lib.rs b/vortex-expr/src/lib.rs index facaa0f932..17a2983f37 100644 --- a/vortex-expr/src/lib.rs +++ b/vortex-expr/src/lib.rs @@ -42,7 +42,7 @@ use vortex_array::{ArrayDType as _, ArrayData, Canonical, IntoArrayData as _}; use vortex_dtype::{DType, FieldName}; use vortex_error::{VortexResult, VortexUnwrap}; -use crate::traversal::{Node, ReferenceCollector}; +use crate::traversal::{DynNode, Node, ReferenceCollector}; pub type ExprRef = Arc; @@ -92,6 +92,12 @@ impl VortexExprExt for ExprRef { } } +impl DynNode for dyn VortexExpr { + fn arc_children(&self) -> VortexResult>> { + Ok(self.children()) + } +} + /// Splits top level and operations into separate expressions pub fn split_conjunction(expr: &ExprRef) -> Vec { let mut conjunctions = vec![]; diff --git a/vortex-expr/src/traversal/mod.rs b/vortex-expr/src/traversal/mod.rs index a7e3cc822e..8470720f96 100644 --- a/vortex-expr/src/traversal/mod.rs +++ b/vortex-expr/src/traversal/mod.rs @@ -193,7 +193,7 @@ pub trait Node: Sized { } pub trait DynNode { - fn arc_children(&self) -> Vec<&Arc>; + fn arc_children(&self) -> VortexResult>>; } pub trait NodeMut: Sized { @@ -222,7 +222,7 @@ impl Node for Arc { if ord == TraversalOrder::Skip { return Ok(TraversalOrder::Continue); } - for child in self.arc_children() { + for child in self.arc_children()? { if ord != TraversalOrder::Continue { return Ok(ord); } @@ -243,57 +243,8 @@ impl Node for Arc { FoldDown::Stop(out) => return Ok(FoldUp::Stop(out)), FoldDown::SkipChildren => FoldChildren::Skipped, FoldDown::Continue(child_context) => { - let mut new_children = Vec::with_capacity(self.arc_children().len()); - for child in self.arc_children() { - match child.accept_with_context(visitor, child_context.clone())? { - FoldUp::Stop(out) => return Ok(FoldUp::Stop(out)), - FoldUp::Continue(out) => new_children.push(out), - } - } - FoldChildren::Children(new_children) - } - }; - - visitor.visit_up(self, context, children) - } -} - -impl Node for ExprRef { - // A pre-order traversal. - fn accept<'a, V: NodeVisitor<'a, NodeTy = ExprRef>>( - &'a self, - visitor: &mut V, - ) -> VortexResult { - let mut ord = visitor.visit_down(self)?; - if ord == TraversalOrder::Stop { - return Ok(TraversalOrder::Stop); - } - if ord == TraversalOrder::Skip { - return Ok(TraversalOrder::Continue); - } - for child in self.children() { - if ord != TraversalOrder::Continue { - return Ok(ord); - } - ord = child.accept(visitor)?; - } - if ord == TraversalOrder::Stop { - return Ok(TraversalOrder::Stop); - } - visitor.visit_up(self) - } - - fn accept_with_context<'a, V: Folder<'a, NodeTy = Self>>( - &'a self, - visitor: &mut V, - context: V::Context, - ) -> VortexResult> { - let children = match visitor.visit_down(self, context.clone())? { - FoldDown::Stop(out) => return Ok(FoldUp::Stop(out)), - FoldDown::SkipChildren => FoldChildren::Skipped, - FoldDown::Continue(child_context) => { - let mut new_children = Vec::with_capacity(self.children().len()); - for child in self.children() { + let mut new_children = Vec::with_capacity(self.arc_children()?.len()); + for child in self.arc_children()? { match child.accept_with_context(visitor, child_context.clone())? { FoldUp::Stop(out) => return Ok(FoldUp::Stop(out)), FoldUp::Continue(out) => new_children.push(out), diff --git a/vortex-layout/src/layouts/struct_/eval_expr.rs b/vortex-layout/src/layouts/struct_/eval_expr.rs index d40f25d852..1673050c27 100644 --- a/vortex-layout/src/layouts/struct_/eval_expr.rs +++ b/vortex-layout/src/layouts/struct_/eval_expr.rs @@ -74,6 +74,7 @@ mod tests { use crate::layouts::struct_::writer::StructLayoutWriter; use crate::segments::test::TestSegments; use crate::strategies::LayoutWriterExt; + use crate::visitor::LayoutReaderDebug; use crate::LayoutData; /// Create a chunked layout with three chunks of primitive arrays. @@ -194,4 +195,12 @@ mod tests { [4, 5].as_slice() ); } + + #[test] + fn visitor() { + let (segments, layout) = struct_layout(); + + let reader = layout.reader(segments, Default::default()).unwrap(); + println!("{:?}", LayoutReaderDebug(reader.clone())); + } } diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index acb97e4214..a2136c0382 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -75,6 +75,10 @@ impl StructReader { .ok_or_else(|| vortex_err!("Field {} not found in struct layout", n))?, Field::Index(idx) => *idx, }; + self.child_idx(idx) + } + + fn child_idx(&self, idx: usize) -> VortexResult<&Arc> { self.field_readers[idx].get_or_try_init(|| { let child_layout = self .layout @@ -90,6 +94,8 @@ impl LayoutReader for StructReader { } fn children(&self) -> VortexResult>> { - (0..self.nchildren()).map(|idx| self.child(idx)).collect() + (0..self.nchildren()) + .map(|idx| self.child_idx(idx)) + .collect() } } diff --git a/vortex-layout/src/visitor.rs b/vortex-layout/src/visitor.rs index 373393346f..d76677f8cb 100644 --- a/vortex-layout/src/visitor.rs +++ b/vortex-layout/src/visitor.rs @@ -1,7 +1,9 @@ +use std::fmt; +use std::fmt::{Debug, Formatter}; use std::sync::Arc; use vortex_error::VortexResult; -use vortex_expr::traversal::DynNode; +use vortex_expr::traversal::{DynNode, Node, NodeVisitor, TraversalOrder}; use crate::LayoutReader; @@ -10,3 +12,27 @@ impl DynNode for dyn LayoutReader { self.children() } } + +pub struct LayoutVisitor<'a, 'b> { + display: &'a mut Formatter<'b>, +} + +impl NodeVisitor<'_> for LayoutVisitor<'_, '_> { + type NodeTy = Arc; + + fn visit_down(&mut self, node: &Self::NodeTy) -> VortexResult { + writeln!(self.display, "{:?}", node.layout())?; + Ok(TraversalOrder::Continue) + } +} + +pub struct LayoutReaderDebug(pub Arc); + +impl Debug for LayoutReaderDebug { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + writeln!(f, "LayoutReader")?; + let mut vis = LayoutVisitor { display: f }; + self.0.accept(&mut vis).unwrap(); + Ok(()) + } +} From e2af2e776623292820ed2cd3431f1d2b6a49bbaf Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 13 Jan 2025 17:03:18 +0000 Subject: [PATCH 4/8] fix --- vortex-expr/src/traversal/mod.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vortex-expr/src/traversal/mod.rs b/vortex-expr/src/traversal/mod.rs index 8470720f96..8274960f4f 100644 --- a/vortex-expr/src/traversal/mod.rs +++ b/vortex-expr/src/traversal/mod.rs @@ -346,7 +346,9 @@ mod tests { use vortex_error::VortexResult; use crate::traversal::visitor::pre_order_visit_down; - use crate::traversal::{MutNodeVisitor, NodeMut, NodeVisitor, TransformResult, TraversalOrder}; + use crate::traversal::{ + MutNodeVisitor, Node, NodeMut, NodeVisitor, TransformResult, TraversalOrder, + }; use crate::{ BinaryExpr, Column, ExprRef, FieldName, Literal, Operator, VortexExpr, VortexExprExt, }; From f9efc6234281032704c5dd10d0caf13c565f9b9f Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 13 Jan 2025 17:04:29 +0000 Subject: [PATCH 5/8] fix --- vortex-expr/src/lib.rs | 2 +- vortex-expr/src/traversal/references.rs | 7 +------ 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/vortex-expr/src/lib.rs b/vortex-expr/src/lib.rs index 17a2983f37..85522cca4b 100644 --- a/vortex-expr/src/lib.rs +++ b/vortex-expr/src/lib.rs @@ -85,7 +85,7 @@ pub trait VortexExprExt { impl VortexExprExt for ExprRef { fn references(&self) -> HashSet { - let mut collector = ReferenceCollector::new(); + let mut collector = ReferenceCollector::default(); // The collector is infallible, so we can unwrap the result self.accept(&mut collector).vortex_unwrap(); collector.into_fields() diff --git a/vortex-expr/src/traversal/references.rs b/vortex-expr/src/traversal/references.rs index 885477de25..fc90b9ee49 100644 --- a/vortex-expr/src/traversal/references.rs +++ b/vortex-expr/src/traversal/references.rs @@ -5,17 +5,12 @@ use vortex_error::VortexResult; use crate::traversal::{NodeVisitor, TraversalOrder}; use crate::{Column, ExprRef, GetItem, Select}; +#[derive(Default)] pub struct ReferenceCollector { fields: HashSet, } impl ReferenceCollector { - pub fn new() -> Self { - Self { - fields: HashSet::new(), - } - } - pub fn with_set(set: HashSet) -> Self { Self { fields: set } } From c7b2d8322374fda0906c4c6773a0f01905c98212 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 13 Jan 2025 17:14:14 +0000 Subject: [PATCH 6/8] update --- vortex-layout/src/layouts/struct_/eval_expr.rs | 2 +- vortex-layout/src/visitor.rs | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-layout/src/layouts/struct_/eval_expr.rs b/vortex-layout/src/layouts/struct_/eval_expr.rs index 1673050c27..83f04b2640 100644 --- a/vortex-layout/src/layouts/struct_/eval_expr.rs +++ b/vortex-layout/src/layouts/struct_/eval_expr.rs @@ -197,7 +197,7 @@ mod tests { } #[test] - fn visitor() { + fn test_visitor() { let (segments, layout) = struct_layout(); let reader = layout.reader(segments, Default::default()).unwrap(); diff --git a/vortex-layout/src/visitor.rs b/vortex-layout/src/visitor.rs index d76677f8cb..1aa21e7159 100644 --- a/vortex-layout/src/visitor.rs +++ b/vortex-layout/src/visitor.rs @@ -21,7 +21,8 @@ impl NodeVisitor<'_> for LayoutVisitor<'_, '_> { type NodeTy = Arc; fn visit_down(&mut self, node: &Self::NodeTy) -> VortexResult { - writeln!(self.display, "{:?}", node.layout())?; + node.layout().fmt(self.display)?; + self.display.write_str("\n")?; Ok(TraversalOrder::Continue) } } From a68ada52b8a15ba83418e057ec8d1cb0dc1dbf6a Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 13 Jan 2025 17:15:16 +0000 Subject: [PATCH 7/8] update --- vortex-layout/src/layouts/struct_/eval_expr.rs | 9 --------- vortex-layout/src/visitor.rs | 6 ++++-- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/vortex-layout/src/layouts/struct_/eval_expr.rs b/vortex-layout/src/layouts/struct_/eval_expr.rs index 83f04b2640..d40f25d852 100644 --- a/vortex-layout/src/layouts/struct_/eval_expr.rs +++ b/vortex-layout/src/layouts/struct_/eval_expr.rs @@ -74,7 +74,6 @@ mod tests { use crate::layouts::struct_::writer::StructLayoutWriter; use crate::segments::test::TestSegments; use crate::strategies::LayoutWriterExt; - use crate::visitor::LayoutReaderDebug; use crate::LayoutData; /// Create a chunked layout with three chunks of primitive arrays. @@ -195,12 +194,4 @@ mod tests { [4, 5].as_slice() ); } - - #[test] - fn test_visitor() { - let (segments, layout) = struct_layout(); - - let reader = layout.reader(segments, Default::default()).unwrap(); - println!("{:?}", LayoutReaderDebug(reader.clone())); - } } diff --git a/vortex-layout/src/visitor.rs b/vortex-layout/src/visitor.rs index 1aa21e7159..ce3fe83d6c 100644 --- a/vortex-layout/src/visitor.rs +++ b/vortex-layout/src/visitor.rs @@ -2,7 +2,7 @@ use std::fmt; use std::fmt::{Debug, Formatter}; use std::sync::Arc; -use vortex_error::VortexResult; +use vortex_error::{VortexExpect, VortexResult}; use vortex_expr::traversal::{DynNode, Node, NodeVisitor, TraversalOrder}; use crate::LayoutReader; @@ -33,7 +33,9 @@ impl Debug for LayoutReaderDebug { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { writeln!(f, "LayoutReader")?; let mut vis = LayoutVisitor { display: f }; - self.0.accept(&mut vis).unwrap(); + self.0 + .accept(&mut vis) + .vortex_expect("Visitor should not fail"); Ok(()) } } From 8996048483e50092d041e5674e0676755d65f7a9 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 15 Jan 2025 11:39:32 +0000 Subject: [PATCH 8/8] update --- vortex-expr/src/transform/remove_select.rs | 2 +- vortex-layout/src/layouts/struct_/reader.rs | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-expr/src/transform/remove_select.rs b/vortex-expr/src/transform/remove_select.rs index 94e4b6a8d5..08ef17d189 100644 --- a/vortex-expr/src/transform/remove_select.rs +++ b/vortex-expr/src/transform/remove_select.rs @@ -2,7 +2,7 @@ use itertools::Itertools; use vortex_dtype::DType; use vortex_error::{vortex_err, VortexResult}; -use crate::traversal::{MutNodeVisitor, Node, TransformResult}; +use crate::traversal::{MutNodeVisitor, NodeMut, TransformResult}; use crate::{get_item, pack, ExprRef, Select}; /// Select is a useful expression, however it can be defined in terms of get_item & pack, diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index f27f0fd30e..bc8af3d55b 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -80,7 +80,8 @@ impl StructReader { .field_lookup .as_ref() .and_then(|lookup| lookup.get(name).copied()) - .or_else(|| self.struct_dtype().find_name(name)); + .or_else(|| self.struct_dtype().find_name(name)) + .ok_or_else(|| vortex_err!("StructReader::Field {} not found", name))?; self.child_idx(idx) }