From d65c712d4ac740c8c313ab5548d89c0426cf3892 Mon Sep 17 00:00:00 2001
From: Orson Peters <orsonpeters@gmail.com>
Date: Thu, 19 Dec 2024 17:31:22 +0100
Subject: [PATCH] perf(rust): Drop probe tables in parallel in new-streaming
 equi-join

---
 crates/polars-stream/src/nodes/joins/equi_join.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/crates/polars-stream/src/nodes/joins/equi_join.rs b/crates/polars-stream/src/nodes/joins/equi_join.rs
index c0c072d51d8b..524b3814ca46 100644
--- a/crates/polars-stream/src/nodes/joins/equi_join.rs
+++ b/crates/polars-stream/src/nodes/joins/equi_join.rs
@@ -516,6 +516,13 @@ impl ProbeState {
     }
 }
 
+impl Drop for ProbeState {
+    fn drop(&mut self) {
+        // Parallel drop as the state might be quite big.
+        self.table_per_partition.par_drain(..).for_each(drop);
+    }
+}
+
 struct EmitUnmatchedState {
     partitions: Vec<ProbeTable>,
     active_partition_idx: usize,