feat: Support intersect all and except distinct/all in DataFrame API (#3537)

This commits adds complete support for intersect and except operations.
To support that, it's consisted of:
1. Expose new methods such as `intersect_all`, `except_distinct` and
`except_all` in python side's DataFrame API
2. Add Except operator in set_operations and convert corresponding
logical plans
3. Complete intersect all branch with similar logic with except all
4. Add a new scalar function: `list_fill(num, val)` which will create a
list of num elements with the provided `val`

The intersect/except all logic is a bit complex and it will require
list_fill and explode to correctly compute the intersected or excepted
rows.

Author: advancedxy

daft/daft/__init__.pyi

@@ -1636,6 +1636,7 @@ class LogicalPlanBuilder:
     ) -> LogicalPlanBuilder: ...
     def concat(self, other: LogicalPlanBuilder) -> LogicalPlanBuilder: ...
     def intersect(self, other: LogicalPlanBuilder, is_all: bool) -> LogicalPlanBuilder: ...
+    def except_(self, other: LogicalPlanBuilder, is_all: bool) -> LogicalPlanBuilder: ...
     def add_monotonically_increasing_id(self, column_name: str | None) -> LogicalPlanBuilder: ...
     def table_write(
         self,

daft/dataframe/dataframe.py

@@ -2557,6 +2557,94 @@ def intersect(self, other: "DataFrame") -> "DataFrame":
         builder = self._builder.intersect(other._builder)
         return DataFrame(builder)
 
+    @DataframePublicAPI
+    def intersect_all(self, other: "DataFrame") -> "DataFrame":
+        """Returns the intersection of two DataFrames, including duplicates.
+
+        Example:
+            >>> import daft
+            >>> df1 = daft.from_pydict({"a": [1, 2, 2], "b": [4, 6, 6]})
+            >>> df2 = daft.from_pydict({"a": [1, 1, 2, 2], "b": [4, 4, 6, 6]})
+            >>> df1.intersect_all(df2).sort("a").collect()
+            ╭───────┬───────╮
+            │ a     ┆ b     │
+            │ ---   ┆ ---   │
+            │ Int64 ┆ Int64 │
+            ╞═══════╪═══════╡
+            │ 1     ┆ 4     │
+            ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
+            │ 2     ┆ 6     │
+            ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
+            │ 2     ┆ 6     │
+            ╰───────┴───────╯
+            <BLANKLINE>
+            (Showing first 3 of 3 rows)
+
+        Args:
+            other (DataFrame): DataFrame to intersect with
+
+        Returns:
+            DataFrame: DataFrame with the intersection of the two DataFrames, including duplicates
+        """
+        builder = self._builder.intersect_all(other._builder)
+        return DataFrame(builder)
+
+    @DataframePublicAPI
+    def except_distinct(self, other: "DataFrame") -> "DataFrame":
+        """Returns the set difference of two DataFrames.
+
+        Example:
+            >>> import daft
+            >>> df1 = daft.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
+            >>> df2 = daft.from_pydict({"a": [1, 2, 3], "b": [4, 8, 6]})
+            >>> df1.except_distinct(df2).collect()
+            ╭───────┬───────╮
+            │ a     ┆ b     │
+            │ ---   ┆ ---   │
+            │ Int64 ┆ Int64 │
+            ╞═══════╪═══════╡
+            │ 2     ┆ 5     │
+            ╰───────┴───────╯
+            <BLANKLINE>
+            (Showing first 1 of 1 rows)
+
+        Args:
+            other (DataFrame): DataFrame to except with
+
+        Returns:
+            DataFrame: DataFrame with the set difference of the two DataFrames
+        """
+        builder = self._builder.except_distinct(other._builder)
+        return DataFrame(builder)
+
+    @DataframePublicAPI
+    def except_all(self, other: "DataFrame") -> "DataFrame":
+        """Returns the set difference of two DataFrames, considering duplicates.
+
+        Example:
+            >>> import daft
+            >>> df1 = daft.from_pydict({"a": [1, 1, 2, 2], "b": [4, 4, 6, 6]})
+            >>> df2 = daft.from_pydict({"a": [1, 2, 2], "b": [4, 6, 6]})
+            >>> df1.except_all(df2).collect()
+            ╭───────┬───────╮
+            │ a     ┆ b     │
+            │ ---   ┆ ---   │
+            │ Int64 ┆ Int64 │
+            ╞═══════╪═══════╡
+            │ 1     ┆ 4     │
+            ╰───────┴───────╯
+            <BLANKLINE>
+            (Showing first 1 of 1 rows)
+
+        Args:
+            other (DataFrame): DataFrame to except with
+
+        Returns:
+            DataFrame: DataFrame with the set difference of the two DataFrames, considering duplicates
+        """
+        builder = self._builder.except_all(other._builder)
+        return DataFrame(builder)
+
     def _materialize_results(self) -> None:
         """Materializes the results of for this DataFrame and hold a pointer to the results."""
         context = get_context()

daft/logical/builder.py

@@ -279,6 +279,18 @@ def intersect(self, other: LogicalPlanBuilder) -> LogicalPlanBuilder:
         builder = self._builder.intersect(other._builder, False)
         return LogicalPlanBuilder(builder)
 
+    def intersect_all(self, other: LogicalPlanBuilder) -> LogicalPlanBuilder:
+        builder = self._builder.intersect(other._builder, True)
+        return LogicalPlanBuilder(builder)
+
+    def except_distinct(self, other: LogicalPlanBuilder) -> LogicalPlanBuilder:
+        builder = self._builder.except_(other._builder, False)
+        return LogicalPlanBuilder(builder)
+
+    def except_all(self, other: LogicalPlanBuilder) -> LogicalPlanBuilder:
+        builder = self._builder.except_(other._builder, True)
+        return LogicalPlanBuilder(builder)
+
     def add_monotonically_increasing_id(self, column_name: str | None) -> LogicalPlanBuilder:
         builder = self._builder.add_monotonically_increasing_id(column_name)
         return LogicalPlanBuilder(builder)

src/daft-core/src/array/ops/list.rs

@@ -1,6 +1,6 @@
 use std::{iter::repeat, sync::Arc};
 
-use arrow2::offset::OffsetsBuffer;
+use arrow2::offset::{Offsets, OffsetsBuffer};
 use common_error::DaftResult;
 use indexmap::{
     map::{raw_entry_v1::RawEntryMut, RawEntryApiV1},
@@ -255,6 +255,31 @@ fn list_sort_helper_fixed_size(
         .collect()
 }
 
+fn general_list_fill_helper(element: &Series, num_array: &Int64Array) -> DaftResult<Vec<Series>> {
+    let num_iter = create_iter(num_array, element.len());
+    let mut result = Vec::with_capacity(element.len());
+    let element_data = element.as_physical()?;
+    for (row_index, num) in num_iter.enumerate() {
+        let list_arr = if element.is_valid(row_index) {
+            let mut list_growable = make_growable(
+                element.name(),
+                element.data_type(),
+                vec![&element_data],
+                false,
+                num as usize,
+            );
+            for _ in 0..num {
+                list_growable.extend(0, row_index, 1);
+            }
+            list_growable.build()?
+        } else {
+            Series::full_null(element.name(), element.data_type(), num as usize)
+        };
+        result.push(list_arr);
+    }
+    Ok(result)
+}
+
 impl ListArray {
     pub fn value_counts(&self) -> DaftResult<MapArray> {
         struct IndexRef {
@@ -625,6 +650,25 @@ impl ListArray {
             self.validity().cloned(),
         ))
     }
+
+    pub fn list_fill(elem: &Series, num_array: &Int64Array) -> DaftResult<Self> {
+        let generated = general_list_fill_helper(elem, num_array)?;
+        let generated_refs: Vec<&Series> = generated.iter().collect();
+        let lengths = generated.iter().map(|arr| arr.len());
+        let offsets = Offsets::try_from_lengths(lengths)?;
+        let flat_child = if generated_refs.is_empty() {
+            // when there's no output, we should create an empty series
+            Series::empty(elem.name(), elem.data_type())
+        } else {
+            Series::concat(&generated_refs)?
+        };
+        Ok(Self::new(
+            elem.field().to_list_field()?,
+            flat_child,
+            offsets.into(),
+            None,
+        ))
+    }
 }
 
 impl FixedSizeListArray {

src/daft-core/src/series/ops/list.rs

@@ -2,8 +2,9 @@ use common_error::{DaftError, DaftResult};
 use daft_schema::field::Field;
 
 use crate::{
+    array::ListArray,
     datatypes::{DataType, UInt64Array, Utf8Array},
-    prelude::CountMode,
+    prelude::{CountMode, Int64Array},
     series::{IntoSeries, Series},
 };
 
@@ -217,4 +218,14 @@ impl Series {
             ))),
         }
     }
+
+    /// Given a series of data T, repeat each data T with num times to create a list, returns
+    /// a series of repeated list.
+    /// # Example
+    /// ```txt
+    /// repeat([1, 2, 3], [2, 0, 1]) --> [[1, 1], [], [3]]
+    /// ```
+    pub fn list_fill(&self, num: &Int64Array) -> DaftResult<Self> {
+        ListArray::list_fill(self, num).map(|arr| arr.into_series())
+    }
 }