Skip to content

Rename functions in SelectExpr to make API clearer #4690

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

Merged
AdamGS merged 2 commits into from
Sep 19, 2025
Merged

Rename functions in SelectExpr to make API clearer #4690

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e5c7063
Rename functions in SelectExpr to make API clearer
AdamGS Sep 18, 2025
3b3d47a
Fix doctests
AdamGS Sep 18, 2025
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
32 changes: 16 additions & 16 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions vortex-dtype/src/struct_.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ impl Hash for FieldDTypeInner {

impl FieldDType {
/// Returns the concrete DType, parsing it from the underlying buffer if necessary.
#[inline]
pub fn value(&self) -> VortexResult<DType> {
self.inner.value()
}
Expand Down
99 changes: 55 additions & 44 deletions vortex-expr/src/exprs/select.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use crate::field::DisplayFieldNames;
use crate::{AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, VTable, vtable};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum SelectField {
pub enum FieldSelection {
Include(FieldNames),
Exclude(FieldNames),
}
Expand All @@ -25,13 +25,13 @@ vtable!(Select);
#[derive(Debug, Clone, Hash, Eq)]
#[allow(clippy::derived_hash_with_manual_eq)]
pub struct SelectExpr {
fields: SelectField,
selection: FieldSelection,
child: ExprRef,
}

impl PartialEq for SelectExpr {
fn eq(&self, other: &Self) -> bool {
self.fields == other.fields && self.child.eq(&other.child)
self.selection == other.selection && self.child.eq(&other.child)
}
}

Expand All @@ -52,13 +52,13 @@ impl VTable for SelectVTable {

fn metadata(expr: &Self::Expr) -> Option<Self::Metadata> {
let names = expr
.fields()
.fields()
Comment on lines -55 to -56
Copy link
Contributor Author

@AdamGS AdamGS Sep 18, 2025

Choose a reason for hiding this comment

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

Mainly to prevent this, I think this is much clearer now.

.selection()
.field_names()
.iter()
.map(|f| f.to_string())
.collect_vec();

let opts = if expr.fields().is_include() {
let opts = if expr.selection().is_include() {
Opts::Include(ProtoFieldNames { names })
} else {
Opts::Exclude(ProtoFieldNames { names })
Expand All @@ -73,7 +73,7 @@ impl VTable for SelectVTable {

fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
Ok(SelectExpr {
fields: expr.fields.clone(),
selection: expr.selection.clone(),
child: children[0].clone(),
})
}
Expand All @@ -89,10 +89,10 @@ impl VTable for SelectVTable {

let fields = match metadata.opts.as_ref() {
Some(opts) => match opts {
Opts::Include(field_names) => SelectField::Include(FieldNames::from_iter(
Opts::Include(field_names) => FieldSelection::Include(FieldNames::from_iter(
field_names.names.iter().map(|s| s.as_str()),
)),
Opts::Exclude(field_names) => SelectField::Exclude(FieldNames::from_iter(
Opts::Exclude(field_names) => FieldSelection::Exclude(FieldNames::from_iter(
field_names.names.iter().map(|s| s.as_str()),
)),
},
Expand All @@ -106,14 +106,17 @@ impl VTable for SelectVTable {
.next()
.vortex_expect("number of children validated to be one");

Ok(SelectExpr { fields, child })
Ok(SelectExpr {
selection: fields,
child,
})
}

fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
let batch = expr.child.unchecked_evaluate(scope)?.to_struct();
Ok(match &expr.fields {
SelectField::Include(f) => batch.project(f.as_ref()),
SelectField::Exclude(names) => {
Ok(match &expr.selection {
FieldSelection::Include(f) => batch.project(f.as_ref()),
FieldSelection::Exclude(names) => {
let included_names = batch
.names()
.iter()
Expand All @@ -132,9 +135,9 @@ impl VTable for SelectVTable {
.as_struct_fields_opt()
.ok_or_else(|| vortex_err!("Select child not a struct dtype"))?;

let projected = match &expr.fields {
SelectField::Include(fields) => child_struct_dtype.project(fields.as_ref())?,
SelectField::Exclude(fields) => child_struct_dtype
let projected = match &expr.selection {
FieldSelection::Include(fields) => child_struct_dtype.project(fields.as_ref())?,
FieldSelection::Exclude(fields) => child_struct_dtype
.names()
.iter()
.cloned()
Expand All @@ -154,8 +157,8 @@ impl VTable for SelectVTable {
/// # use vortex_expr::{select, root};
/// let expr = select(["name", "age"], root());
/// ```
pub fn select(fields: impl Into<FieldNames>, child: ExprRef) -> ExprRef {
SelectExpr::include_expr(fields.into(), child)
pub fn select(field_names: impl Into<FieldNames>, child: ExprRef) -> ExprRef {
SelectExpr::include_expr(field_names.into(), child)
}

/// Creates an expression that excludes specific fields from an array.
Expand All @@ -171,24 +174,27 @@ pub fn select_exclude(fields: impl Into<FieldNames>, child: ExprRef) -> ExprRef
}

impl SelectExpr {
pub fn new(fields: SelectField, child: ExprRef) -> Self {
Self { fields, child }
pub fn new(fields: FieldSelection, child: ExprRef) -> Self {
Self {
selection: fields,
child,
}
}

pub fn new_expr(fields: SelectField, child: ExprRef) -> ExprRef {
pub fn new_expr(fields: FieldSelection, child: ExprRef) -> ExprRef {
Self::new(fields, child).into_expr()
}

pub fn include_expr(columns: FieldNames, child: ExprRef) -> ExprRef {
Self::new(SelectField::Include(columns), child).into_expr()
Self::new(FieldSelection::Include(columns), child).into_expr()
}

pub fn exclude_expr(columns: FieldNames, child: ExprRef) -> ExprRef {
Self::new(SelectField::Exclude(columns), child).into_expr()
Self::new(FieldSelection::Exclude(columns), child).into_expr()
}

pub fn fields(&self) -> &SelectField {
&self.fields
pub fn selection(&self) -> &FieldSelection {
&self.selection
}

pub fn child(&self) -> &ExprRef {
Expand All @@ -200,26 +206,26 @@ impl SelectExpr {
/// For example:
/// ```rust
/// # use vortex_expr::root;
/// # use vortex_expr::{SelectExpr, SelectField};
/// # use vortex_expr::{FieldSelection, SelectExpr};
/// # use vortex_dtype::FieldNames;
/// let field_names = FieldNames::from(["a", "b", "c"]);
/// let include = SelectExpr::new(SelectField::Include(["a"].into()), root());
/// let exclude = SelectExpr::new(SelectField::Exclude(["b", "c"].into()), root());
/// let include = SelectExpr::new(FieldSelection::Include(["a"].into()), root());
/// let exclude = SelectExpr::new(FieldSelection::Exclude(["b", "c"].into()), root());
/// assert_eq!(
/// &include.as_include(&field_names).unwrap(),
/// &exclude.as_include(&field_names).unwrap()
/// );
/// ```
pub fn as_include(&self, field_names: &FieldNames) -> VortexResult<ExprRef> {
Ok(Self::new(
SelectField::Include(self.fields.as_include_names(field_names)?),
FieldSelection::Include(self.selection.as_include_names(field_names)?),
self.child.clone(),
)
.into_expr())
}
}

impl SelectField {
impl FieldSelection {
pub fn include(columns: FieldNames) -> Self {
assert_eq!(columns.iter().unique().collect_vec().len(), columns.len());
Self::Include(columns)
Expand All @@ -238,15 +244,15 @@ impl SelectField {
matches!(self, Self::Exclude(_))
}

pub fn fields(&self) -> &FieldNames {
let (SelectField::Include(fields) | SelectField::Exclude(fields)) = self;
pub fn field_names(&self) -> &FieldNames {
let (FieldSelection::Include(fields) | FieldSelection::Exclude(fields)) = self;

fields
}

pub fn as_include_names(&self, field_names: &FieldNames) -> VortexResult<FieldNames> {
if self
.fields()
.field_names()
.iter()
.any(|f| !field_names.iter().contains(f))
{
Expand All @@ -257,8 +263,8 @@ impl SelectField {
);
}
match self {
SelectField::Include(fields) => Ok(fields.clone()),
SelectField::Exclude(exc_fields) => Ok(field_names
FieldSelection::Include(fields) => Ok(fields.clone()),
FieldSelection::Exclude(exc_fields) => Ok(field_names
.iter()
.filter(|f| !exc_fields.iter().contains(f))
.cloned()
Expand All @@ -267,11 +273,11 @@ impl SelectField {
}
}

impl Display for SelectField {
impl Display for FieldSelection {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SelectField::Include(fields) => write!(f, "{{{}}}", DisplayFieldNames(fields)),
SelectField::Exclude(fields) => write!(f, "~{{{}}}", DisplayFieldNames(fields)),
FieldSelection::Include(fields) => write!(f, "{{{}}}", DisplayFieldNames(fields)),
FieldSelection::Exclude(fields) => write!(f, "~{{{}}}", DisplayFieldNames(fields)),
}
}
}
Expand All @@ -280,16 +286,21 @@ impl DisplayAs for SelectExpr {
fn fmt_as(&self, df: DisplayFormat, f: &mut std::fmt::Formatter) -> std::fmt::Result {
match df {
DisplayFormat::Compact => {
write!(f, "{}{}", self.child, self.fields)
write!(f, "{}{}", self.child, self.selection)
}
DisplayFormat::Tree => {
let field_type = if self.fields.is_include() {
let field_type = if self.selection.is_include() {
"include"
} else {
"exclude"
};

write!(f, "Select({}): {}", field_type, self.fields().fields())
write!(
f,
"Select({}): {}",
field_type,
self.selection().field_names()
)
}
}
}
Expand All @@ -310,7 +321,7 @@ mod tests {
use vortex_buffer::buffer;
use vortex_dtype::{DType, FieldName, FieldNames, Nullability};

use crate::{Scope, SelectExpr, SelectField, root, select, select_exclude, test_harness};
use crate::{FieldSelection, Scope, SelectExpr, root, select, select_exclude, test_harness};

fn test_array() -> StructArray {
StructArray::from_fields(&[
Expand Down Expand Up @@ -393,8 +404,8 @@ mod tests {
#[test]
fn test_as_include_names() {
let field_names = FieldNames::from(["a", "b", "c"]);
let include = SelectExpr::new(SelectField::Include(["a"].into()), root());
let exclude = SelectExpr::new(SelectField::Exclude(["b", "c"].into()), root());
let include = SelectExpr::new(FieldSelection::Include(["a"].into()), root());
let exclude = SelectExpr::new(FieldSelection::Exclude(["b", "c"].into()), root());
assert_eq!(
&include.as_include(&field_names).unwrap(),
&exclude.as_include(&field_names).unwrap()
Expand Down
4 changes: 2 additions & 2 deletions vortex-expr/src/transform/remove_select.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,12 +29,12 @@ fn remove_select_transformer(node: ExprRef, ctx: &DType) -> VortexResult<Transfo

let expr = pack(
select
.fields()
.selection()
.as_include_names(child_dtype.names())
.map_err(|e| {
e.with_context(format!(
"Select fields {:?} must be a subset of child fields {:?}",
select.fields(),
select.selection(),
child_dtype.names()
))
})?
Expand Down
3 changes: 2 additions & 1 deletion vortex-expr/src/traversal/references.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,8 @@ impl NodeVisitor<'_> for ReferenceCollector {
self.fields.insert(get_item.field().clone());
}
if let Some(sel) = node.as_opt::<SelectVTable>() {
self.fields.extend(sel.fields().fields().iter().cloned());
self.fields
.extend(sel.selection().field_names().iter().cloned());
}
Ok(TraversalOrder::Continue)
}
Expand Down
Loading
Loading