This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion-python.git
The following commit(s) were added to refs/heads/main by this push:
new c2768d8 Add random missing bindings (#522)
c2768d8 is described below
commit c2768d8f9d558f3a5e042f1e934c02f945ede797
Author: Jeremy Dyer <[email protected]>
AuthorDate: Wed Oct 18 11:40:26 2023 -0400
Add random missing bindings (#522)
---
src/expr/aggregate.rs | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
src/expr/join.rs | 14 ++++++++++++++
src/expr/sort.rs | 4 ++++
3 files changed, 67 insertions(+)
diff --git a/src/expr/aggregate.rs b/src/expr/aggregate.rs
index c3de967..5ebf8c6 100644
--- a/src/expr/aggregate.rs
+++ b/src/expr/aggregate.rs
@@ -16,12 +16,15 @@
// under the License.
use datafusion_common::DataFusionError;
+use datafusion_expr::expr::{AggregateFunction, AggregateUDF, Alias};
use datafusion_expr::logical_plan::Aggregate;
+use datafusion_expr::Expr;
use pyo3::prelude::*;
use std::fmt::{self, Display, Formatter};
use super::logical_node::LogicalNode;
use crate::common::df_schema::PyDFSchema;
+use crate::errors::py_type_err;
use crate::expr::PyExpr;
use crate::sql::logical::PyLogicalPlan;
@@ -84,6 +87,24 @@ impl PyAggregate {
.collect())
}
+ /// Returns the inner Aggregate Expr(s)
+ pub fn agg_expressions(&self) -> PyResult<Vec<PyExpr>> {
+ Ok(self
+ .aggregate
+ .aggr_expr
+ .iter()
+ .map(|e| PyExpr::from(e.clone()))
+ .collect())
+ }
+
+ pub fn agg_func_name(&self, expr: PyExpr) -> PyResult<String> {
+ Self::_agg_func_name(&expr.expr)
+ }
+
+ pub fn aggregation_arguments(&self, expr: PyExpr) -> PyResult<Vec<PyExpr>>
{
+ self._aggregation_arguments(&expr.expr)
+ }
+
// Retrieves the input `LogicalPlan` to this `Aggregate` node
fn input(&self) -> PyResult<Vec<PyLogicalPlan>> {
Ok(Self::inputs(self))
@@ -99,6 +120,34 @@ impl PyAggregate {
}
}
+impl PyAggregate {
+ #[allow(clippy::only_used_in_recursion)]
+ fn _aggregation_arguments(&self, expr: &Expr) -> PyResult<Vec<PyExpr>> {
+ match expr {
+ // TODO: This Alias logic seems to be returning some strange
results that we should investigate
+ Expr::Alias(Alias { expr, .. }) =>
self._aggregation_arguments(expr.as_ref()),
+ Expr::AggregateFunction(AggregateFunction { fun: _, args, .. })
+ | Expr::AggregateUDF(AggregateUDF { fun: _, args, .. }) => {
+ Ok(args.iter().map(|e| PyExpr::from(e.clone())).collect())
+ }
+ _ => Err(py_type_err(
+ "Encountered a non Aggregate type in aggregation_arguments",
+ )),
+ }
+ }
+
+ fn _agg_func_name(expr: &Expr) -> PyResult<String> {
+ match expr {
+ Expr::Alias(Alias { expr, .. }) =>
Self::_agg_func_name(expr.as_ref()),
+ Expr::AggregateFunction(AggregateFunction { fun, .. }) =>
Ok(fun.to_string()),
+ Expr::AggregateUDF(AggregateUDF { fun, .. }) =>
Ok(fun.name.clone()),
+ _ => Err(py_type_err(
+ "Encountered a non Aggregate type in agg_func_name",
+ )),
+ }
+ }
+}
+
impl LogicalNode for PyAggregate {
fn inputs(&self) -> Vec<PyLogicalPlan> {
vec![PyLogicalPlan::from((*self.aggregate.input).clone())]
diff --git a/src/expr/join.rs b/src/expr/join.rs
index 8016629..a53ddd3 100644
--- a/src/expr/join.rs
+++ b/src/expr/join.rs
@@ -46,6 +46,10 @@ impl PyJoinType {
pub fn is_outer(&self) -> bool {
self.join_type.is_outer()
}
+
+ fn __repr__(&self) -> PyResult<String> {
+ Ok(format!("{}", self.join_type))
+ }
}
impl Display for PyJoinType {
@@ -72,6 +76,16 @@ impl From<PyJoinConstraint> for JoinConstraint {
}
}
+#[pymethods]
+impl PyJoinConstraint {
+ fn __repr__(&self) -> PyResult<String> {
+ match self.join_constraint {
+ JoinConstraint::On => Ok("On".to_string()),
+ JoinConstraint::Using => Ok("Using".to_string()),
+ }
+ }
+}
+
#[pyclass(name = "Join", module = "datafusion.expr", subclass)]
#[derive(Clone)]
pub struct PyJoin {
diff --git a/src/expr/sort.rs b/src/expr/sort.rs
index 8843c63..f9f9e58 100644
--- a/src/expr/sort.rs
+++ b/src/expr/sort.rs
@@ -72,6 +72,10 @@ impl PySort {
.collect())
}
+ fn get_fetch_val(&self) -> PyResult<Option<usize>> {
+ Ok(self.sort.fetch)
+ }
+
/// Retrieves the input `LogicalPlan` to this `Sort` node
fn input(&self) -> PyResult<Vec<PyLogicalPlan>> {
Ok(Self::inputs(self))
