feat: add register_udaf_fun funcation for dataframe (apache#49)

francis-du · web-flow · commit 0f048ad1bade · 2022-09-07T21:33:54.000-06:00
diff --git a/src/context.rs b/src/context.rs
@@ -34,6 +34,7 @@ use crate::catalog::{PyCatalog, PyTable};
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
 use crate::errors::DataFusionError;
+use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
 use crate::utils::wait_for_future;
 
@@ -227,6 +228,11 @@ impl PySessionContext {
         Ok(())
     }
 
+    fn register_udaf(&mut self, udaf: PyAggregateUDF) -> PyResult<()> {
+        self.ctx.register_udaf(udaf.function);
+        Ok(())
+    }
+
     #[args(name = "\"datafusion\"")]
     fn catalog(&self, name: &str) -> PyResult<PyCatalog> {
         match self.ctx.catalog(name) {
