From 7886e29b271cdfd947eaead5cc22ae009672b40c Mon Sep 17 00:00:00 2001
From: Andrew Lamb <andrew@nerdnetworks.org>
Date: Wed, 17 Jul 2024 15:56:55 -0400
Subject: [PATCH] Optimize `single_column`

---
 datafusion/core/src/physical_optimizer/pruning.rs | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/datafusion/core/src/physical_optimizer/pruning.rs b/datafusion/core/src/physical_optimizer/pruning.rs
index 0deb698bc6eb..3c18e53497fd 100644
--- a/datafusion/core/src/physical_optimizer/pruning.rs
+++ b/datafusion/core/src/physical_optimizer/pruning.rs
@@ -747,10 +747,13 @@ impl RequiredColumns {
     /// * `a > 5 OR b < 10` returns `None`
     /// * `true` returns None
     pub(crate) fn single_column(&self) -> Option<&phys_expr::Column> {
-        let cols = self.iter().map(|(c, _s, _f)| c).collect::<HashSet<_>>();
-
-        if cols.len() == 1 {
-            cols.iter().next().copied()
+        if self.columns.windows(2).all(|w| {
+            // check if all columns are the same (ignoring statistics and field)
+            let c1 = &w[0].0;
+            let c2 = &w[1].0;
+            c1 == c2
+        }) {
+            self.columns.first().map(|r| &r.0)
         } else {
             None
         }