perf: speed up INTERSECT ALL and EXCEPT ALL on mostly distinct rows

datafusion/core/src/dataframe/mod.rs

@@ -2123,7 +2123,13 @@ impl DataFrame {
     pub fn intersect(self, dataframe: DataFrame) -> Result<DataFrame> {
         let left_plan = self.plan;
         let right_plan = dataframe.plan;
-        let plan = LogicalPlanBuilder::intersect(left_plan, right_plan, true)?;
+        let Some((count, range)) = self.session_state.set_operation_functions() else {
+            return plan_err!(
+                "DataFrame::intersect requires the built-in count and range functions to be available"
+            );
+        };
+        let plan =
+            LogicalPlanBuilder::intersect_all(left_plan, right_plan, &count, &range)?;
         Ok(DataFrame {
             session_state: self.session_state,
             plan,
@@ -2201,7 +2207,12 @@ impl DataFrame {
     pub fn except(self, dataframe: DataFrame) -> Result<DataFrame> {
         let left_plan = self.plan;
         let right_plan = dataframe.plan;
-        let plan = LogicalPlanBuilder::except(left_plan, right_plan, true)?;
+        let Some((count, range)) = self.session_state.set_operation_functions() else {
+            return plan_err!(
+                "DataFrame::except requires the built-in count and range functions to be available"
+            );
+        };
+        let plan = LogicalPlanBuilder::except_all(left_plan, right_plan, &count, &range)?;
         Ok(DataFrame {
             session_state: self.session_state,
             plan,

datafusion/core/src/execution/session_state.rs

@@ -69,6 +69,10 @@ use datafusion_expr::simplify::SimplifyContext;
 use datafusion_expr::{
     AggregateUDF, Explain, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
 };
+#[cfg(feature = "nested_expressions")]
+use datafusion_functions_aggregate::count::count_udaf;
+#[cfg(feature = "nested_expressions")]
+use datafusion_functions_nested::range::range_udf;
 use datafusion_optimizer::simplify_expressions::ExprSimplifier;
 use datafusion_optimizer::{
     Analyzer, AnalyzerRule, Optimizer, OptimizerConfig, OptimizerRule,
@@ -1058,6 +1062,24 @@ impl SessionState {
         &self.inner.aggregate_functions
     }
 
+    /// Return the built-in functions used by `INTERSECT ALL` and `EXCEPT ALL`.
+    ///
+    /// Returns `None` when `count` or `range` is unavailable or
+    /// `nested_expressions` is disabled.
+    pub fn set_operation_functions(&self) -> Option<(Arc<AggregateUDF>, Arc<ScalarUDF>)> {
+        self.aggregate_functions().get("count")?;
+        self.scalar_functions().get("range")?;
+
+        #[cfg(feature = "nested_expressions")]
+        {
+            Some((count_udaf(), range_udf()))
+        }
+        #[cfg(not(feature = "nested_expressions"))]
+        {
+            None
+        }
+    }
+
     /// Return reference to window functions
     pub fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>> {
         &self.inner.window_functions
@@ -2192,6 +2214,10 @@ impl ContextProvider for SessionContextProvider<'_> {
         self.state.aggregate_functions().get(name).cloned()
     }
 
+    fn get_set_operation_functions(&self) -> Option<(Arc<AggregateUDF>, Arc<ScalarUDF>)> {
+        self.state.set_operation_functions()
+    }
+
     fn get_window_meta(&self, name: &str) -> Option<Arc<WindowUDF>> {
         self.state.window_functions().get(name).cloned()
     }

datafusion/core/tests/dataframe/mod.rs

@@ -33,7 +33,7 @@ use arrow::datatypes::{
 use arrow::error::ArrowError;
 use arrow::util::pretty::pretty_format_batches;
 use arrow_schema::{SortOptions, TimeUnit};
-use datafusion::{assert_batches_eq, dataframe};
+use datafusion::{assert_batches_eq, assert_batches_sorted_eq, dataframe};
 use datafusion_common::metadata::FieldMetadata;
 use datafusion_expr::select_expr::SelectExpr;
 use datafusion_functions_aggregate::count::{count_all, count_all_window};
@@ -1849,31 +1849,90 @@ async fn sendable() {
 
 #[tokio::test]
 async fn intersect() -> Result<()> {
-    let df = test_table().await?.select_columns(&["c1", "c3"])?;
-    let d2 = df.clone();
-    let plan = df.intersect(d2)?;
-    let result = plan.logical_plan().clone();
-    let expected = create_plan(
-        "SELECT c1, c3 FROM aggregate_test_100
-            INTERSECT ALL SELECT c1, c3 FROM aggregate_test_100",
-    )
-    .await?;
-    assert_same_plan(&result, &expected);
+    let ctx = SessionContext::new();
+    let left = ctx
+        .sql(
+            "SELECT * FROM VALUES (NULL, 'n'), (NULL, 'n'), (1, 'x'), (1, 'x'), (2, 'y')",
+        )
+        .await?;
+    let right = ctx
+        .sql("SELECT * FROM VALUES (NULL, 'n'), (1, 'x'), (1, 'x'), (1, 'x'), (3, 'z')")
+        .await?;
+    let result = left.intersect(right)?.collect().await?;
+    assert_batches_sorted_eq!(
+        [
+            "+---------+---------+",
+            "| column1 | column2 |",
+            "+---------+---------+",
+            "|         | n       |",
+            "| 1       | x       |",
+            "| 1       | x       |",
+            "+---------+---------+",
+        ],
+        &result
+    );
     Ok(())
 }
 
 #[tokio::test]
 async fn except() -> Result<()> {
-    let df = test_table().await?.select_columns(&["c1", "c3"])?;
-    let d2 = df.clone();
-    let plan = df.except(d2)?;
-    let result = plan.logical_plan().clone();
-    let expected = create_plan(
-        "SELECT c1, c3 FROM aggregate_test_100
-            EXCEPT ALL SELECT c1, c3 FROM aggregate_test_100",
-    )
-    .await?;
-    assert_same_plan(&result, &expected);
+    let ctx = SessionContext::new();
+    let left = ctx
+        .sql(
+            "SELECT * FROM VALUES (NULL, 'n'), (NULL, 'n'), (1, 'x'), (1, 'x'), (2, 'y')",
+        )
+        .await?;
+    let right = ctx
+        .sql("SELECT * FROM VALUES (NULL, 'n'), (1, 'x'), (1, 'x'), (1, 'x'), (3, 'z')")
+        .await?;
+    let result = left.except(right)?.collect().await?;
+    assert_batches_sorted_eq!(
+        [
+            "+---------+---------+",
+            "| column1 | column2 |",
+            "+---------+---------+",
+            "|         | n       |",
+            "| 2       | y       |",
+            "+---------+---------+",
+        ],
+        &result
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn intersect_and_except_count_rows_in_the_type_they_are_compared_in() -> Result<()>
+{
+    // '1' and '01' are both the integer 1, so the strings are two copies of the
+    // row that the integers repeat twice, and the result is an integer column.
+    let ctx = SessionContext::new();
+    let integers = ctx.sql("SELECT * FROM VALUES (1), (1)").await?;
+    let strings = ctx.sql("SELECT * FROM VALUES ('1'), ('01')").await?;
+
+    let result = integers.clone().except(strings.clone())?.collect().await?;
+    assert_eq!(
+        result.iter().map(|batch| batch.num_rows()).sum::<usize>(),
+        0
+    );
+    let result = strings.clone().except(integers.clone())?.collect().await?;
+    assert_eq!(
+        result.iter().map(|batch| batch.num_rows()).sum::<usize>(),
+        0
+    );
+
+    let intersection = integers.intersect(strings)?;
+    assert_eq!(intersection.schema().field(0).data_type(), &DataType::Int64);
+    assert_batches_sorted_eq!(
+        [
+            "+---------+",
+            "| column1 |",
+            "+---------+",
+            "| 1       |",
+            "| 1       |",
+            "+---------+",
+        ],
+        &intersection.collect().await?
+    );
     Ok(())
 }
 

datafusion/core/tests/set_operations_no_nested.rs

@@ -0,0 +1,43 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#![cfg(all(feature = "sql", not(feature = "nested_expressions")))]
+
+use datafusion::prelude::SessionContext;
+
+#[tokio::test]
+async fn set_operations_require_nested_expressions() -> datafusion::error::Result<()> {
+    let ctx = SessionContext::new();
+    let err = ctx
+        .sql("SELECT 1 INTERSECT ALL SELECT 1")
+        .await
+        .unwrap_err();
+    assert!(
+        err.to_string()
+            .contains("built-in count and range functions to be available"),
+        "{err}"
+    );
+
+    let dataframe = ctx.sql("SELECT 1").await?;
+    let err = dataframe.clone().intersect(dataframe).unwrap_err();
+    assert!(
+        err.to_string()
+            .contains("built-in count and range functions to be available"),
+        "{err}"
+    );
+    Ok(())
+}

datafusion/core/tests/sql/mod.rs

@@ -70,6 +70,7 @@ mod path_partition;
 mod runtime_config;
 pub mod select;
 mod set_comparison;
+mod set_operations;
 mod sql_api;
 mod tpcds_planning;
 mod union_comparison;