Add extension hooks for encoding and decoding UDAFs and UDWFs (#253)

* Add extension hooks for encoding and decoding UDAFs and UDWFs * Add tests for encoding and decoding UDAFs
coralogix · Jul 17, 2024 · c0f40e1 · c0f40e1
1 parent 55cdadf
commit c0f40e1
Show file tree

Hide file tree

Showing 13 changed files with 640 additions and 315 deletions.
diff --git a/datafusion/physical-expr-common/src/aggregate/mod.rs b/datafusion/physical-expr-common/src/aggregate/mod.rs
@@ -255,6 +255,11 @@ impl AggregateFunctionExpr {
     pub fn is_distinct(&self) -> bool {
         self.is_distinct
     }
+
+    /// Return if the aggregation ignores nulls
+    pub fn ignore_nulls(&self) -> bool {
+        self.ignore_nulls
+    }
 }
 
 impl AggregateExpr for AggregateFunctionExpr {

diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -164,7 +164,7 @@ message CreateExternalTableNode {
   map<string, string> options = 8;
   datafusion_common.Constraints constraints = 12;
   map<string, LogicalExprNode> column_defaults = 13;
- }
+}
 
 message PrepareNode {
   string name = 1;
@@ -249,30 +249,30 @@ message DistinctOnNode {
 }
 
 message CopyToNode {
-    LogicalPlanNode input = 1;
-    string output_url = 2;
-    oneof format_options {
-      datafusion_common.CsvOptions csv = 8;
-      datafusion_common.JsonOptions json = 9;
-      datafusion_common.TableParquetOptions parquet = 10;
-      datafusion_common.AvroOptions avro = 11;
-      datafusion_common.ArrowOptions arrow = 12;
-    }
-    repeated string partition_by = 7;
+  LogicalPlanNode input = 1;
+  string output_url = 2;
+  oneof format_options {
+    datafusion_common.CsvOptions csv = 8;
+    datafusion_common.JsonOptions json = 9;
+    datafusion_common.TableParquetOptions parquet = 10;
+    datafusion_common.AvroOptions avro = 11;
+    datafusion_common.ArrowOptions arrow = 12;
+  }
+  repeated string partition_by = 7;
 }
 
 message UnnestNode {
-    LogicalPlanNode input = 1;
-    repeated datafusion_common.Column exec_columns = 2;
-    repeated uint64 list_type_columns = 3;
-    repeated uint64 struct_type_columns = 4;
-    repeated uint64 dependency_indices = 5;
-    datafusion_common.DfSchema schema = 6;
-    UnnestOptions options = 7;
+  LogicalPlanNode input = 1;
+  repeated datafusion_common.Column exec_columns = 2;
+  repeated uint64 list_type_columns = 3;
+  repeated uint64 struct_type_columns = 4;
+  repeated uint64 dependency_indices = 5;
+  datafusion_common.DfSchema schema = 6;
+  UnnestOptions options = 7;
 }
 
 message UnnestOptions {
-    bool preserve_nulls = 1;
+  bool preserve_nulls = 1;
 }
 
 message UnionNode {
@@ -522,6 +522,7 @@ message AggregateUDFExprNode {
   repeated LogicalExprNode args = 2;
   LogicalExprNode filter = 3;
   repeated LogicalExprNode order_by = 4;
+  optional bytes fun_definition = 6;
 }
 
 message ScalarUDFExprNode {
@@ -556,6 +557,7 @@ message WindowExprNode {
   repeated LogicalExprNode order_by = 6;
   // repeated LogicalExprNode filter = 7;
   WindowFrame window_frame = 8;
+  optional bytes fun_definition = 10;
 }
 
 message BetweenNode {
@@ -856,6 +858,8 @@ message PhysicalAggregateExprNode {
   repeated PhysicalExprNode expr = 2;
   repeated PhysicalSortExprNode ordering_req = 5;
   bool distinct = 3;
+  bool ignore_nulls = 6;
+  optional bytes fun_definition = 7;
 }
 
 message PhysicalWindowExprNode {
@@ -869,6 +873,7 @@ message PhysicalWindowExprNode {
   repeated PhysicalSortExprNode order_by = 6;
   WindowFrame window_frame = 7;
   string name = 8;
+  optional bytes fun_definition = 9;
 }
 
 message PhysicalIsNull {