feat: teach merge to optionally error on duplicates (#4971)

Error on duplicates is better for new users who accidentally collide
names.

---------

Signed-off-by: Daniel King <[email protected]>

Author: danking

vortex-expr/src/exprs/merge.rs

@@ -1,14 +1,18 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::fmt::Display;
 use std::hash::Hash;
 
 use itertools::Itertools as _;
 use vortex_array::arrays::StructArray;
 use vortex_array::validity::Validity;
-use vortex_array::{Array, ArrayRef, DeserializeMetadata, EmptyMetadata, IntoArray, ToCanonical};
+use vortex_array::{
+    Array, ArrayRef, DeserializeMetadata, EmptyMetadata, IntoArray as _, ToCanonical,
+};
 use vortex_dtype::{DType, FieldNames, Nullability, StructFields};
-use vortex_error::{VortexExpect as _, VortexResult, vortex_bail};
+use vortex_error::{VortexResult, vortex_bail};
+use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::display::{DisplayAs, DisplayFormat};
 use crate::{AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, VTable, vtable};
@@ -25,6 +29,32 @@ vtable!(Merge);
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct MergeExpr {
     values: Vec<ExprRef>,
+    duplicate_handling: DuplicateHandling,
+}
+
+impl MergeExpr {
+    pub fn duplicate_handling(&self) -> DuplicateHandling {
+        self.duplicate_handling
+    }
+}
+
+/// What to do when merged structs share a field name.
+#[derive(Default, Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub enum DuplicateHandling {
+    /// If two structs share a field name, take the value from the right-most struct.
+    RightMost,
+    /// If two structs share a field name, error.
+    #[default]
+    Error,
+}
+
+impl Display for DuplicateHandling {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            DuplicateHandling::Error => write!(f, "error"),
+            DuplicateHandling::RightMost => write!(f, "right-most"),
+        }
+    }
 }
 
 pub struct MergeExprEncoding;
@@ -50,8 +80,11 @@ impl VTable for MergeVTable {
         expr.values.iter().collect()
     }
 
-    fn with_children(_expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
-        Ok(MergeExpr { values: children })
+    fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
+        Ok(MergeExpr {
+            values: children,
+            duplicate_handling: expr.duplicate_handling,
+        })
     }
 
     fn build(
@@ -65,39 +98,33 @@ impl VTable for MergeVTable {
                 children
             );
         }
-        Ok(MergeExpr { values: children })
+        Ok(MergeExpr {
+            values: children,
+            duplicate_handling: DuplicateHandling::default(),
+        })
     }
 
     fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
-        let len = scope.len();
-        let value_arrays = expr
-            .values
-            .iter()
-            .map(|value_expr| value_expr.unchecked_evaluate(scope))
-            .process_results(|it| it.collect::<Vec<_>>())?;
-
         // Collect fields in order of appearance. Later fields overwrite earlier fields.
         let mut field_names = Vec::new();
         let mut arrays = Vec::new();
+        let mut duplicate_names = HashSet::<_>::new();
 
-        for value_array in value_arrays.iter() {
+        for expr in expr.values.iter() {
             // TODO(marko): When nullable, we need to merge struct validity into field validity.
-            if value_array.dtype().is_nullable() {
-                todo!("merge nullable structs");
+            let array = expr.unchecked_evaluate(scope)?;
+            if array.dtype().is_nullable() {
+                vortex_bail!("merge expects non-nullable input");
             }
-            if !value_array.dtype().is_struct() {
-                vortex_bail!("merge expects non-nullable struct input");
+            if !array.dtype().is_struct() {
+                vortex_bail!("merge expects struct input");
             }
+            let array = array.to_struct();
 
-            let struct_array = value_array.to_struct();
-
-            for (field_name, array) in struct_array
-                .names()
-                .iter()
-                .zip_eq(struct_array.fields().iter().cloned())
-            {
+            for (field_name, array) in array.names().iter().zip_eq(array.fields().iter().cloned()) {
                 // Update or insert field.
                 if let Some(idx) = field_names.iter().position(|name| name == field_name) {
+                    duplicate_names.insert(field_name.clone());
                     arrays[idx] = array;
                 } else {
                     field_names.push(field_name.clone());
@@ -106,8 +133,16 @@ impl VTable for MergeVTable {
             }
         }
 
+        if expr.duplicate_handling == DuplicateHandling::Error && !duplicate_names.is_empty() {
+            vortex_bail!(
+                "merge: duplicate fields in children: {}",
+                duplicate_names.into_iter().format(", ")
+            )
+        }
+
         // TODO(DK): When children are allowed to be nullable, this needs to change.
         let validity = Validity::NonNullable;
+        let len = scope.len();
         Ok(
             StructArray::try_new(FieldNames::from(field_names), arrays, len, validity)?
                 .into_array(),
@@ -117,27 +152,23 @@ impl VTable for MergeVTable {
     fn return_dtype(expr: &Self::Expr, scope: &DType) -> VortexResult<DType> {
         let mut field_names = Vec::new();
         let mut arrays = Vec::new();
+        let mut merge_nullability = Nullability::NonNullable;
+        let mut duplicate_names = HashSet::<_>::new();
 
-        let mut nullability = Nullability::NonNullable;
-
-        for value in expr.values.iter() {
-            let dtype = value.return_dtype(scope)?;
-            if !dtype.is_struct() {
+        for expr in expr.values.iter() {
+            let dtype = expr.return_dtype(scope)?;
+            let Some(fields) = dtype.as_struct_fields_opt() else {
                 vortex_bail!("merge expects struct input");
-            }
+            };
             if dtype.is_nullable() {
                 vortex_bail!("merge expects non-nullable input");
             }
-            nullability |= dtype.nullability();
 
-            let struct_dtype = dtype
-                .as_struct_fields_opt()
-                .vortex_expect("merge expects struct input");
+            merge_nullability |= dtype.nullability();
 
-            for i in 0..struct_dtype.nfields() {
-                let field_name = struct_dtype.field_name(i).vortex_expect("never OOB");
-                let field_dtype = struct_dtype.field_by_index(i).vortex_expect("never OOB");
+            for (field_name, field_dtype) in fields.names().iter().zip_eq(fields.fields()) {
                 if let Some(idx) = field_names.iter().position(|name| name == field_name) {
+                    duplicate_names.insert(field_name.clone());
                     arrays[idx] = field_dtype;
                 } else {
                     field_names.push(field_name.clone());
@@ -146,21 +177,42 @@ impl VTable for MergeVTable {
             }
         }
 
+        if expr.duplicate_handling == DuplicateHandling::Error && !duplicate_names.is_empty() {
+            vortex_bail!(
+                "merge: duplicate fields in children: {}",
+                duplicate_names.into_iter().format(", ")
+            )
+        }
+
         Ok(DType::Struct(
             StructFields::new(FieldNames::from(field_names), arrays),
-            nullability,
+            merge_nullability,
         ))
     }
 }
 
 impl MergeExpr {
     pub fn new(values: Vec<ExprRef>) -> Self {
-        MergeExpr { values }
+        MergeExpr {
+            values,
+            duplicate_handling: DuplicateHandling::default(),
+        }
     }
 
     pub fn new_expr(values: Vec<ExprRef>) -> ExprRef {
         Self::new(values).into_expr()
     }
+
+    pub fn new_opts(values: Vec<ExprRef>, duplicate_handling: DuplicateHandling) -> Self {
+        MergeExpr {
+            values,
+            duplicate_handling,
+        }
+    }
+
+    pub fn new_expr_opts(values: Vec<ExprRef>, duplicate_handling: DuplicateHandling) -> ExprRef {
+        Self::new_opts(values, duplicate_handling).into_expr()
+    }
 }
 
 /// Creates an expression that merges struct expressions into a single struct.
@@ -178,11 +230,24 @@ pub fn merge(elements: impl IntoIterator<Item = impl Into<ExprRef>>) -> ExprRef
     MergeExpr::new(values).into_expr()
 }
 
+pub fn merge_opts(
+    elements: impl IntoIterator<Item = impl Into<ExprRef>>,
+    duplicate_handling: DuplicateHandling,
+) -> ExprRef {
+    let values = elements.into_iter().map(|value| value.into()).collect_vec();
+    MergeExpr::new_opts(values, duplicate_handling).into_expr()
+}
+
 impl DisplayAs for MergeExpr {
     fn fmt_as(&self, df: DisplayFormat, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match df {
             DisplayFormat::Compact => {
-                write!(f, "merge({})", self.values.iter().format(", "),)
+                write!(
+                    f,
+                    "merge[{}]({})",
+                    self.duplicate_handling,
+                    self.values.iter().format(", "),
+                )
             }
             DisplayFormat::Tree => {
                 write!(f, "Merge")
@@ -200,7 +265,7 @@ mod tests {
     use vortex_buffer::buffer;
     use vortex_error::{VortexResult, vortex_bail};
 
-    use crate::{MergeExpr, Scope, get_item, merge, root};
+    use crate::{DuplicateHandling, MergeExpr, Scope, get_item, merge, root};
 
     fn primitive_field(array: &dyn Array, field_path: &[&str]) -> VortexResult<PrimitiveArray> {
         let mut field_path = field_path.iter();
@@ -217,12 +282,15 @@ mod tests {
     }
 
     #[test]
-    pub fn test_merge() {
-        let expr = MergeExpr::new(vec![
-            get_item("0", root()),
-            get_item("1", root()),
-            get_item("2", root()),
-        ]);
+    pub fn test_merge_right_most() {
+        let expr = MergeExpr::new_opts(
+            vec![
+                get_item("0", root()),
+                get_item("1", root()),
+                get_item("2", root()),
+            ],
+            DuplicateHandling::RightMost,
+        );
 
         let test_array = StructArray::from_fields(&[
             (
@@ -294,6 +362,52 @@ mod tests {
         );
     }
 
+    #[test]
+    #[should_panic(expected = "merge: duplicate fields in children")]
+    pub fn test_merge_error_on_dupe_return_dtype() {
+        let expr = MergeExpr::new_opts(
+            vec![get_item("0", root()), get_item("1", root())],
+            DuplicateHandling::Error,
+        );
+        let test_array = StructArray::try_from_iter([
+            (
+                "0",
+                StructArray::try_from_iter([("a", buffer![1]), ("b", buffer![1])]).unwrap(),
+            ),
+            (
+                "1",
+                StructArray::try_from_iter([("c", buffer![1]), ("b", buffer![1])]).unwrap(),
+            ),
+        ])
+        .unwrap()
+        .into_array();
+
+        expr.return_dtype(test_array.dtype()).unwrap();
+    }
+
+    #[test]
+    #[should_panic(expected = "merge: duplicate fields in children")]
+    pub fn test_merge_error_on_dupe_evaluate() {
+        let expr = MergeExpr::new_opts(
+            vec![get_item("0", root()), get_item("1", root())],
+            DuplicateHandling::Error,
+        );
+        let test_array = StructArray::try_from_iter([
+            (
+                "0",
+                StructArray::try_from_iter([("a", buffer![1]), ("b", buffer![1])]).unwrap(),
+            ),
+            (
+                "1",
+                StructArray::try_from_iter([("c", buffer![1]), ("b", buffer![1])]).unwrap(),
+            ),
+        ])
+        .unwrap()
+        .into_array();
+
+        expr.evaluate(&Scope::new(test_array)).unwrap();
+    }
+
     #[test]
     pub fn test_empty_merge() {
         let expr = MergeExpr::new(Vec::new());
@@ -310,7 +424,10 @@ mod tests {
     pub fn test_nested_merge() {
         // Nested structs are not merged!
 
-        let expr = MergeExpr::new(vec![get_item("0", root()), get_item("1", root())]);
+        let expr = MergeExpr::new_opts(
+            vec![get_item("0", root()), get_item("1", root())],
+            DuplicateHandling::RightMost,
+        );
 
         let test_array = StructArray::from_fields(&[
             (
@@ -396,9 +513,9 @@ mod tests {
     #[test]
     pub fn test_display() {
         let expr = merge([get_item("struct1", root()), get_item("struct2", root())]);
-        assert_eq!(expr.to_string(), "merge($.struct1, $.struct2)");
+        assert_eq!(expr.to_string(), "merge[error]($.struct1, $.struct2)");
 
         let expr2 = MergeExpr::new(vec![get_item("a", root())]);
-        assert_eq!(expr2.to_string(), "merge($.a)");
+        assert_eq!(expr2.to_string(), "merge[error]($.a)");
     }
 }