fix: Improve logic for determining when an UnpackOrDeepCopy is needed (#2142)

Author: andygrove

native/core/src/execution/operators/copy.rs

@@ -91,6 +91,14 @@ impl CopyExec {
             mode,
         }
     }
+
+    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
+        &self.input
+    }
+
+    pub fn mode(&self) -> &CopyMode {
+        &self.mode
+    }
 }
 
 impl DisplayAs for CopyExec {

native/core/src/execution/planner.rs

@@ -2610,10 +2610,23 @@ impl From<ExpressionError> for DataFusionError {
 /// modification. This is used to determine if we need to copy the input batch to avoid
 /// data corruption from reusing the input batch.
 fn can_reuse_input_batch(op: &Arc<dyn ExecutionPlan>) -> bool {
-    if op.as_any().is::<ProjectionExec>() || op.as_any().is::<LocalLimitExec>() {
-        can_reuse_input_batch(op.children()[0])
+    if op.as_any().is::<ScanExec>() {
+        // JVM side can return arrow buffers to the pool
+        // Also, native_comet scan reuses mutable buffers
+        true
+    } else if op.as_any().is::<CopyExec>() {
+        let copy_exec = op.as_any().downcast_ref::<CopyExec>().unwrap();
+        copy_exec.mode() == &CopyMode::UnpackOrClone && can_reuse_input_batch(copy_exec.input())
+    } else if op.as_any().is::<CometFilterExec>() {
+        // CometFilterExec guarantees that all arrays have been copied
+        false
     } else {
-        op.as_any().is::<ScanExec>()
+        for child in op.children() {
+            if can_reuse_input_batch(child) {
+                return true;
+            }
+        }
+        false
     }
 }