Preserve UDAF distincit flag in serde (#259)

coralogix · Aug 21, 2024 · 7f3eac7 · 7f3eac7
1 parent 6465608
commit 7f3eac7
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 1 deletion.
diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -522,6 +522,7 @@ message AggregateUDFExprNode {
   repeated LogicalExprNode args = 2;
   LogicalExprNode filter = 3;
   repeated LogicalExprNode order_by = 4;
+  bool distinct = 5;
   optional bytes fun_definition = 6;
 }
 

diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs
diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs
diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs
@@ -658,7 +658,7 @@ pub fn parse_expr(
             Ok(Expr::AggregateFunction(expr::AggregateFunction::new_udf(
                 agg_fn,
                 parse_exprs(&pb.args, registry, codec)?,
-                false,
+                pb.distinct,
                 parse_optional_expr(pb.filter.as_deref(), registry, codec)?.map(Box::new),
                 parse_vec_expr(&pb.order_by, registry, codec)?,
                 None,

diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs
@@ -502,6 +502,7 @@ pub fn serialize_expr(
                                 Some(e) => serialize_exprs(e, codec)?,
                                 None => vec![],
                             },
+                            distinct: *distinct,
                             fun_definition: (!buf.is_empty()).then_some(buf),
                         },
                     ))),