feat: add intersect and except bindings for DataFrame

Author: francis-du

datafusion/tests/test_dataframe.py

@@ -212,3 +212,54 @@ def test_explain(df):
         column("a") - column("b"),
     )
     df.explain()
+
+
+def test_intersect():
+    ctx = SessionContext()
+
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
+        names=["a", "b"],
+    )
+    df_a = ctx.create_dataframe([[batch]])
+
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([3, 4, 5]), pa.array([6, 7, 8])],
+        names=["a", "b"],
+    )
+    df_b = ctx.create_dataframe([[batch]])
+
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([3]), pa.array([6])],
+        names=["a", "b"],
+    )
+    df_c = ctx.create_dataframe([[batch]]).sort(column("a").sort(ascending=True))
+
+    df_c.show()
+    df_a.intersect(df_b).sort(column("a").sort(ascending=True)).show()
+
+    assert df_c.collect() == df_a.intersect(df_b).sort(column("a").sort(ascending=True)).collect()
+
+
+def test_except_all():
+    ctx = SessionContext()
+
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
+        names=["a", "b"],
+    )
+    df_a = ctx.create_dataframe([[batch]])
+
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([3, 4, 5]), pa.array([6, 7, 8])],
+        names=["a", "b"],
+    )
+    df_b = ctx.create_dataframe([[batch]])
+
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2]), pa.array([4, 5])],
+        names=["a", "b"],
+    )
+    df_c = ctx.create_dataframe([[batch]]).sort(column("a").sort(ascending=True))
+
+    assert df_c.collect() == df_a.except_all(df_b).sort(column("a").sort(ascending=True)).collect()

src/dataframe.rs

@@ -147,7 +147,7 @@ impl PyDataFrame {
                     "The join type {} does not exist or is not implemented",
                     how
                 ))
-                .into())
+                    .into());
             }
         };
 
@@ -164,4 +164,16 @@ impl PyDataFrame {
         let batches = wait_for_future(py, df.collect())?;
         Ok(pretty::print_batches(&batches)?)
     }
+
+    /// Calculate the intersection of two `DataFrame`s.  The two `DataFrame`s must have exactly the same schema
+    fn intersect(&self, py_df: PyDataFrame) -> PyResult<Self> {
+        let new_df = self.df.intersect(py_df.df)?;
+        Ok(Self::new(new_df))
+    }
+
+    /// Calculate the exception of two `DataFrame`s.  The two `DataFrame`s must have exactly the same schema
+    fn except_all(&self, py_df: PyDataFrame) -> PyResult<Self> {
+        let new_df = self.df.except(py_df.df)?;
+        Ok(Self::new(new_df))
+    }
 }