/Users/andrewlamb/Software/datafusion/datafusion/expr/src/udwf.rs
Line | Count | Source (jump to first uncovered line) |
1 | | // Licensed to the Apache Software Foundation (ASF) under one |
2 | | // or more contributor license agreements. See the NOTICE file |
3 | | // distributed with this work for additional information |
4 | | // regarding copyright ownership. The ASF licenses this file |
5 | | // to you under the Apache License, Version 2.0 (the |
6 | | // "License"); you may not use this file except in compliance |
7 | | // with the License. You may obtain a copy of the License at |
8 | | // |
9 | | // http://www.apache.org/licenses/LICENSE-2.0 |
10 | | // |
11 | | // Unless required by applicable law or agreed to in writing, |
12 | | // software distributed under the License is distributed on an |
13 | | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
14 | | // KIND, either express or implied. See the License for the |
15 | | // specific language governing permissions and limitations |
16 | | // under the License. |
17 | | |
18 | | //! [`WindowUDF`]: User Defined Window Functions |
19 | | |
20 | | use arrow::compute::SortOptions; |
21 | | use std::cmp::Ordering; |
22 | | use std::hash::{DefaultHasher, Hash, Hasher}; |
23 | | use std::{ |
24 | | any::Any, |
25 | | fmt::{self, Debug, Display, Formatter}, |
26 | | sync::Arc, |
27 | | }; |
28 | | |
29 | | use arrow::datatypes::{DataType, Field}; |
30 | | |
31 | | use datafusion_common::{not_impl_err, Result}; |
32 | | use datafusion_functions_window_common::field::WindowUDFFieldArgs; |
33 | | |
34 | | use crate::expr::WindowFunction; |
35 | | use crate::{ |
36 | | function::WindowFunctionSimplification, Documentation, Expr, PartitionEvaluator, |
37 | | Signature, |
38 | | }; |
39 | | |
40 | | /// Logical representation of a user-defined window function (UDWF) |
41 | | /// A UDWF is different from a UDF in that it is stateful across batches. |
42 | | /// |
43 | | /// See the documentation on [`PartitionEvaluator`] for more details |
44 | | /// |
45 | | /// 1. For simple use cases, use [`create_udwf`] (examples in |
46 | | /// [`simple_udwf.rs`]). |
47 | | /// |
48 | | /// 2. For advanced use cases, use [`WindowUDFImpl`] which provides full API |
49 | | /// access (examples in [`advanced_udwf.rs`]). |
50 | | /// |
51 | | /// # API Note |
52 | | /// This is a separate struct from `WindowUDFImpl` to maintain backwards |
53 | | /// compatibility with the older API. |
54 | | /// |
55 | | /// [`PartitionEvaluator`]: crate::PartitionEvaluator |
56 | | /// [`create_udwf`]: crate::expr_fn::create_udwf |
57 | | /// [`simple_udwf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/simple_udwf.rs |
58 | | /// [`advanced_udwf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/advanced_udwf.rs |
59 | | #[derive(Debug, Clone, PartialOrd)] |
60 | | pub struct WindowUDF { |
61 | | inner: Arc<dyn WindowUDFImpl>, |
62 | | } |
63 | | |
64 | | /// Defines how the WindowUDF is shown to users |
65 | | impl Display for WindowUDF { |
66 | 0 | fn fmt(&self, f: &mut Formatter) -> fmt::Result { |
67 | 0 | write!(f, "{}", self.name()) |
68 | 0 | } |
69 | | } |
70 | | |
71 | | impl PartialEq for WindowUDF { |
72 | 0 | fn eq(&self, other: &Self) -> bool { |
73 | 0 | self.inner.equals(other.inner.as_ref()) |
74 | 0 | } |
75 | | } |
76 | | |
77 | | impl Eq for WindowUDF {} |
78 | | |
79 | | impl Hash for WindowUDF { |
80 | 0 | fn hash<H: Hasher>(&self, state: &mut H) { |
81 | 0 | self.inner.hash_value().hash(state) |
82 | 0 | } |
83 | | } |
84 | | |
85 | | impl WindowUDF { |
86 | | /// Create a new `WindowUDF` from a `[WindowUDFImpl]` trait object |
87 | | /// |
88 | | /// Note this is the same as using the `From` impl (`WindowUDF::from`) |
89 | 0 | pub fn new_from_impl<F>(fun: F) -> WindowUDF |
90 | 0 | where |
91 | 0 | F: WindowUDFImpl + 'static, |
92 | 0 | { |
93 | 0 | Self { |
94 | 0 | inner: Arc::new(fun), |
95 | 0 | } |
96 | 0 | } |
97 | | |
98 | | /// Return the underlying [`WindowUDFImpl`] trait object for this function |
99 | 0 | pub fn inner(&self) -> &Arc<dyn WindowUDFImpl> { |
100 | 0 | &self.inner |
101 | 0 | } |
102 | | |
103 | | /// Adds additional names that can be used to invoke this function, in |
104 | | /// addition to `name` |
105 | | /// |
106 | | /// If you implement [`WindowUDFImpl`] directly you should return aliases directly. |
107 | 0 | pub fn with_aliases(self, aliases: impl IntoIterator<Item = &'static str>) -> Self { |
108 | 0 | Self::new_from_impl(AliasedWindowUDFImpl::new(Arc::clone(&self.inner), aliases)) |
109 | 0 | } |
110 | | |
111 | | /// creates a [`Expr`] that calls the window function with default |
112 | | /// values for `order_by`, `partition_by`, `window_frame`. |
113 | | /// |
114 | | /// See [`ExprFunctionExt`] for details on setting these values. |
115 | | /// |
116 | | /// This utility allows using a user defined window function without |
117 | | /// requiring access to the registry, such as with the DataFrame API. |
118 | | /// |
119 | | /// [`ExprFunctionExt`]: crate::expr_fn::ExprFunctionExt |
120 | 0 | pub fn call(&self, args: Vec<Expr>) -> Expr { |
121 | 0 | let fun = crate::WindowFunctionDefinition::WindowUDF(Arc::new(self.clone())); |
122 | 0 |
|
123 | 0 | Expr::WindowFunction(WindowFunction::new(fun, args)) |
124 | 0 | } |
125 | | |
126 | | /// Returns this function's name |
127 | | /// |
128 | | /// See [`WindowUDFImpl::name`] for more details. |
129 | 0 | pub fn name(&self) -> &str { |
130 | 0 | self.inner.name() |
131 | 0 | } |
132 | | |
133 | | /// Returns the aliases for this function. |
134 | 0 | pub fn aliases(&self) -> &[String] { |
135 | 0 | self.inner.aliases() |
136 | 0 | } |
137 | | |
138 | | /// Returns this function's signature (what input types are accepted) |
139 | | /// |
140 | | /// See [`WindowUDFImpl::signature`] for more details. |
141 | 0 | pub fn signature(&self) -> &Signature { |
142 | 0 | self.inner.signature() |
143 | 0 | } |
144 | | |
145 | | /// Do the function rewrite |
146 | | /// |
147 | | /// See [`WindowUDFImpl::simplify`] for more details. |
148 | 0 | pub fn simplify(&self) -> Option<WindowFunctionSimplification> { |
149 | 0 | self.inner.simplify() |
150 | 0 | } |
151 | | |
152 | | /// Return a `PartitionEvaluator` for evaluating this window function |
153 | 0 | pub fn partition_evaluator_factory(&self) -> Result<Box<dyn PartitionEvaluator>> { |
154 | 0 | self.inner.partition_evaluator() |
155 | 0 | } |
156 | | |
157 | | /// Returns the field of the final result of evaluating this window function. |
158 | | /// |
159 | | /// See [`WindowUDFImpl::field`] for more details. |
160 | 0 | pub fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> { |
161 | 0 | self.inner.field(field_args) |
162 | 0 | } |
163 | | |
164 | | /// Returns custom result ordering introduced by this window function |
165 | | /// which is used to update ordering equivalences. |
166 | | /// |
167 | | /// See [`WindowUDFImpl::sort_options`] for more details. |
168 | 0 | pub fn sort_options(&self) -> Option<SortOptions> { |
169 | 0 | self.inner.sort_options() |
170 | 0 | } |
171 | | |
172 | | /// See [`WindowUDFImpl::coerce_types`] for more details. |
173 | 0 | pub fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> { |
174 | 0 | self.inner.coerce_types(arg_types) |
175 | 0 | } |
176 | | |
177 | | /// Returns the reversed user-defined window function when the |
178 | | /// order of evaluation is reversed. |
179 | | /// |
180 | | /// See [`WindowUDFImpl::reverse_expr`] for more details. |
181 | 0 | pub fn reverse_expr(&self) -> ReversedUDWF { |
182 | 0 | self.inner.reverse_expr() |
183 | 0 | } |
184 | | |
185 | | /// Returns the documentation for this Window UDF. |
186 | | /// |
187 | | /// Documentation can be accessed programmatically as well as |
188 | | /// generating publicly facing documentation. |
189 | 0 | pub fn documentation(&self) -> Option<&Documentation> { |
190 | 0 | self.inner.documentation() |
191 | 0 | } |
192 | | } |
193 | | |
194 | | impl<F> From<F> for WindowUDF |
195 | | where |
196 | | F: WindowUDFImpl + Send + Sync + 'static, |
197 | | { |
198 | 0 | fn from(fun: F) -> Self { |
199 | 0 | Self::new_from_impl(fun) |
200 | 0 | } |
201 | | } |
202 | | |
203 | | /// Trait for implementing [`WindowUDF`]. |
204 | | /// |
205 | | /// This trait exposes the full API for implementing user defined window functions and |
206 | | /// can be used to implement any function. |
207 | | /// |
208 | | /// See [`advanced_udwf.rs`] for a full example with complete implementation and |
209 | | /// [`WindowUDF`] for other available options. |
210 | | /// |
211 | | /// |
212 | | /// [`advanced_udwf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/advanced_udwf.rs |
213 | | /// # Basic Example |
214 | | /// ``` |
215 | | /// # use std::any::Any; |
216 | | /// # use std::sync::OnceLock; |
217 | | /// # use arrow::datatypes::{DataType, Field}; |
218 | | /// # use datafusion_common::{DataFusionError, plan_err, Result}; |
219 | | /// # use datafusion_expr::{col, Signature, Volatility, PartitionEvaluator, WindowFrame, ExprFunctionExt, Documentation}; |
220 | | /// # use datafusion_expr::{WindowUDFImpl, WindowUDF}; |
221 | | /// # use datafusion_expr::window_doc_sections::DOC_SECTION_ANALYTICAL; |
222 | | /// # use datafusion_functions_window_common::field::WindowUDFFieldArgs; |
223 | | /// |
224 | | /// #[derive(Debug, Clone)] |
225 | | /// struct SmoothIt { |
226 | | /// signature: Signature, |
227 | | /// } |
228 | | /// |
229 | | /// impl SmoothIt { |
230 | | /// fn new() -> Self { |
231 | | /// Self { |
232 | | /// signature: Signature::uniform(1, vec![DataType::Int32], Volatility::Immutable), |
233 | | /// } |
234 | | /// } |
235 | | /// } |
236 | | /// |
237 | | /// static DOCUMENTATION: OnceLock<Documentation> = OnceLock::new(); |
238 | | /// |
239 | | /// fn get_doc() -> &'static Documentation { |
240 | | /// DOCUMENTATION.get_or_init(|| { |
241 | | /// Documentation::builder() |
242 | | /// .with_doc_section(DOC_SECTION_ANALYTICAL) |
243 | | /// .with_description("smooths the windows") |
244 | | /// .with_syntax_example("smooth_it(2)") |
245 | | /// .with_argument("arg1", "The int32 number to smooth by") |
246 | | /// .build() |
247 | | /// .unwrap() |
248 | | /// }) |
249 | | /// } |
250 | | /// |
251 | | /// /// Implement the WindowUDFImpl trait for SmoothIt |
252 | | /// impl WindowUDFImpl for SmoothIt { |
253 | | /// fn as_any(&self) -> &dyn Any { self } |
254 | | /// fn name(&self) -> &str { "smooth_it" } |
255 | | /// fn signature(&self) -> &Signature { &self.signature } |
256 | | /// // The actual implementation would smooth the window |
257 | | /// fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> { unimplemented!() } |
258 | | /// fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> { |
259 | | /// if let Some(DataType::Int32) = field_args.get_input_type(0) { |
260 | | /// Ok(Field::new(field_args.name(), DataType::Int32, false)) |
261 | | /// } else { |
262 | | /// plan_err!("smooth_it only accepts Int32 arguments") |
263 | | /// } |
264 | | /// } |
265 | | /// fn documentation(&self) -> Option<&Documentation> { |
266 | | /// Some(get_doc()) |
267 | | /// } |
268 | | /// } |
269 | | /// |
270 | | /// // Create a new WindowUDF from the implementation |
271 | | /// let smooth_it = WindowUDF::from(SmoothIt::new()); |
272 | | /// |
273 | | /// // Call the function `add_one(col)` |
274 | | /// // smooth_it(speed) OVER (PARTITION BY car ORDER BY time ASC) |
275 | | /// let expr = smooth_it.call(vec![col("speed")]) |
276 | | /// .partition_by(vec![col("car")]) |
277 | | /// .order_by(vec![col("time").sort(true, true)]) |
278 | | /// .window_frame(WindowFrame::new(None)) |
279 | | /// .build() |
280 | | /// .unwrap(); |
281 | | /// ``` |
282 | | pub trait WindowUDFImpl: Debug + Send + Sync { |
283 | | // Note: When adding any methods (with default implementations), remember to add them also |
284 | | // into the AliasedWindowUDFImpl below! |
285 | | |
286 | | /// Returns this object as an [`Any`] trait object |
287 | | fn as_any(&self) -> &dyn Any; |
288 | | |
289 | | /// Returns this function's name |
290 | | fn name(&self) -> &str; |
291 | | |
292 | | /// Returns the function's [`Signature`] for information about what input |
293 | | /// types are accepted and the function's Volatility. |
294 | | fn signature(&self) -> &Signature; |
295 | | |
296 | | /// Invoke the function, returning the [`PartitionEvaluator`] instance |
297 | | fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>>; |
298 | | |
299 | | /// Returns any aliases (alternate names) for this function. |
300 | | /// |
301 | | /// Note: `aliases` should only include names other than [`Self::name`]. |
302 | | /// Defaults to `[]` (no aliases) |
303 | 0 | fn aliases(&self) -> &[String] { |
304 | 0 | &[] |
305 | 0 | } |
306 | | |
307 | | /// Optionally apply per-UDWF simplification / rewrite rules. |
308 | | /// |
309 | | /// This can be used to apply function specific simplification rules during |
310 | | /// optimization. The default implementation does nothing. |
311 | | /// |
312 | | /// Note that DataFusion handles simplifying arguments and "constant |
313 | | /// folding" (replacing a function call with constant arguments such as |
314 | | /// `my_add(1,2) --> 3` ). Thus, there is no need to implement such |
315 | | /// optimizations manually for specific UDFs. |
316 | | /// |
317 | | /// Example: |
318 | | /// [`simplify_udwf_expression.rs`]: <https://github.com/apache/arrow-datafusion/blob/main/datafusion-examples/examples/simplify_udwf_expression.rs> |
319 | | /// |
320 | | /// # Returns |
321 | | /// [None] if simplify is not defined or, |
322 | | /// |
323 | | /// Or, a closure with two arguments: |
324 | | /// * 'window_function': [crate::expr::WindowFunction] for which simplified has been invoked |
325 | | /// * 'info': [crate::simplify::SimplifyInfo] |
326 | 0 | fn simplify(&self) -> Option<WindowFunctionSimplification> { |
327 | 0 | None |
328 | 0 | } |
329 | | |
330 | | /// Return true if this window UDF is equal to the other. |
331 | | /// |
332 | | /// Allows customizing the equality of window UDFs. |
333 | | /// Must be consistent with [`Self::hash_value`] and follow the same rules as [`Eq`]: |
334 | | /// |
335 | | /// - reflexive: `a.equals(a)`; |
336 | | /// - symmetric: `a.equals(b)` implies `b.equals(a)`; |
337 | | /// - transitive: `a.equals(b)` and `b.equals(c)` implies `a.equals(c)`. |
338 | | /// |
339 | | /// By default, compares [`Self::name`] and [`Self::signature`]. |
340 | 0 | fn equals(&self, other: &dyn WindowUDFImpl) -> bool { |
341 | 0 | self.name() == other.name() && self.signature() == other.signature() |
342 | 0 | } |
343 | | |
344 | | /// Returns a hash value for this window UDF. |
345 | | /// |
346 | | /// Allows customizing the hash code of window UDFs. Similarly to [`Hash`] and [`Eq`], |
347 | | /// if [`Self::equals`] returns true for two UDFs, their `hash_value`s must be the same. |
348 | | /// |
349 | | /// By default, hashes [`Self::name`] and [`Self::signature`]. |
350 | 0 | fn hash_value(&self) -> u64 { |
351 | 0 | let hasher = &mut DefaultHasher::new(); |
352 | 0 | self.name().hash(hasher); |
353 | 0 | self.signature().hash(hasher); |
354 | 0 | hasher.finish() |
355 | 0 | } |
356 | | |
357 | | /// The [`Field`] of the final result of evaluating this window function. |
358 | | /// |
359 | | /// Call `field_args.name()` to get the fully qualified name for defining |
360 | | /// the [`Field`]. For a complete example see the implementation in the |
361 | | /// [Basic Example](WindowUDFImpl#basic-example) section. |
362 | | fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field>; |
363 | | |
364 | | /// Allows the window UDF to define a custom result ordering. |
365 | | /// |
366 | | /// By default, a window UDF doesn't introduce an ordering. |
367 | | /// But when specified by a window UDF this is used to update |
368 | | /// ordering equivalences. |
369 | 0 | fn sort_options(&self) -> Option<SortOptions> { |
370 | 0 | None |
371 | 0 | } |
372 | | |
373 | | /// Coerce arguments of a function call to types that the function can evaluate. |
374 | | /// |
375 | | /// This function is only called if [`WindowUDFImpl::signature`] returns [`crate::TypeSignature::UserDefined`]. Most |
376 | | /// UDWFs should return one of the other variants of `TypeSignature` which handle common |
377 | | /// cases |
378 | | /// |
379 | | /// See the [type coercion module](crate::type_coercion) |
380 | | /// documentation for more details on type coercion |
381 | | /// |
382 | | /// For example, if your function requires a floating point arguments, but the user calls |
383 | | /// it like `my_func(1::int)` (aka with `1` as an integer), coerce_types could return `[DataType::Float64]` |
384 | | /// to ensure the argument was cast to `1::double` |
385 | | /// |
386 | | /// # Parameters |
387 | | /// * `arg_types`: The argument types of the arguments this function with |
388 | | /// |
389 | | /// # Return value |
390 | | /// A Vec the same length as `arg_types`. DataFusion will `CAST` the function call |
391 | | /// arguments to these specific types. |
392 | 0 | fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> { |
393 | 0 | not_impl_err!("Function {} does not implement coerce_types", self.name()) |
394 | 0 | } |
395 | | |
396 | | /// Allows customizing the behavior of the user-defined window |
397 | | /// function when it is evaluated in reverse order. |
398 | 0 | fn reverse_expr(&self) -> ReversedUDWF { |
399 | 0 | ReversedUDWF::NotSupported |
400 | 0 | } |
401 | | |
402 | | /// Returns the documentation for this Window UDF. |
403 | | /// |
404 | | /// Documentation can be accessed programmatically as well as |
405 | | /// generating publicly facing documentation. |
406 | 0 | fn documentation(&self) -> Option<&Documentation> { |
407 | 0 | None |
408 | 0 | } |
409 | | } |
410 | | |
411 | | pub enum ReversedUDWF { |
412 | | /// The result of evaluating the user-defined window function |
413 | | /// remains identical when reversed. |
414 | | Identical, |
415 | | /// A window function which does not support evaluating the result |
416 | | /// in reverse order. |
417 | | NotSupported, |
418 | | /// Customize the user-defined window function for evaluating the |
419 | | /// result in reverse order. |
420 | | Reversed(Arc<WindowUDF>), |
421 | | } |
422 | | |
423 | | impl PartialEq for dyn WindowUDFImpl { |
424 | 0 | fn eq(&self, other: &Self) -> bool { |
425 | 0 | self.equals(other) |
426 | 0 | } |
427 | | } |
428 | | |
429 | | impl PartialOrd for dyn WindowUDFImpl { |
430 | 0 | fn partial_cmp(&self, other: &Self) -> Option<Ordering> { |
431 | 0 | match self.name().partial_cmp(other.name()) { |
432 | 0 | Some(Ordering::Equal) => self.signature().partial_cmp(other.signature()), |
433 | 0 | cmp => cmp, |
434 | | } |
435 | 0 | } |
436 | | } |
437 | | |
438 | | /// WindowUDF that adds an alias to the underlying function. It is better to |
439 | | /// implement [`WindowUDFImpl`], which supports aliases, directly if possible. |
440 | | #[derive(Debug)] |
441 | | struct AliasedWindowUDFImpl { |
442 | | inner: Arc<dyn WindowUDFImpl>, |
443 | | aliases: Vec<String>, |
444 | | } |
445 | | |
446 | | impl AliasedWindowUDFImpl { |
447 | 0 | pub fn new( |
448 | 0 | inner: Arc<dyn WindowUDFImpl>, |
449 | 0 | new_aliases: impl IntoIterator<Item = &'static str>, |
450 | 0 | ) -> Self { |
451 | 0 | let mut aliases = inner.aliases().to_vec(); |
452 | 0 | aliases.extend(new_aliases.into_iter().map(|s| s.to_string())); |
453 | 0 |
|
454 | 0 | Self { inner, aliases } |
455 | 0 | } |
456 | | } |
457 | | |
458 | | impl WindowUDFImpl for AliasedWindowUDFImpl { |
459 | 0 | fn as_any(&self) -> &dyn Any { |
460 | 0 | self |
461 | 0 | } |
462 | | |
463 | 0 | fn name(&self) -> &str { |
464 | 0 | self.inner.name() |
465 | 0 | } |
466 | | |
467 | 0 | fn signature(&self) -> &Signature { |
468 | 0 | self.inner.signature() |
469 | 0 | } |
470 | | |
471 | 0 | fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> { |
472 | 0 | self.inner.partition_evaluator() |
473 | 0 | } |
474 | | |
475 | 0 | fn aliases(&self) -> &[String] { |
476 | 0 | &self.aliases |
477 | 0 | } |
478 | | |
479 | 0 | fn simplify(&self) -> Option<WindowFunctionSimplification> { |
480 | 0 | self.inner.simplify() |
481 | 0 | } |
482 | | |
483 | 0 | fn equals(&self, other: &dyn WindowUDFImpl) -> bool { |
484 | 0 | if let Some(other) = other.as_any().downcast_ref::<AliasedWindowUDFImpl>() { |
485 | 0 | self.inner.equals(other.inner.as_ref()) && self.aliases == other.aliases |
486 | | } else { |
487 | 0 | false |
488 | | } |
489 | 0 | } |
490 | | |
491 | 0 | fn hash_value(&self) -> u64 { |
492 | 0 | let hasher = &mut DefaultHasher::new(); |
493 | 0 | self.inner.hash_value().hash(hasher); |
494 | 0 | self.aliases.hash(hasher); |
495 | 0 | hasher.finish() |
496 | 0 | } |
497 | | |
498 | 0 | fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> { |
499 | 0 | self.inner.field(field_args) |
500 | 0 | } |
501 | | |
502 | 0 | fn sort_options(&self) -> Option<SortOptions> { |
503 | 0 | self.inner.sort_options() |
504 | 0 | } |
505 | | |
506 | 0 | fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> { |
507 | 0 | self.inner.coerce_types(arg_types) |
508 | 0 | } |
509 | | |
510 | 0 | fn documentation(&self) -> Option<&Documentation> { |
511 | 0 | self.inner.documentation() |
512 | 0 | } |
513 | | } |
514 | | |
515 | | // Window UDF doc sections for use in public documentation |
516 | | pub mod window_doc_sections { |
517 | | use crate::DocSection; |
518 | | |
519 | | pub fn doc_sections() -> Vec<DocSection> { |
520 | | vec![ |
521 | | DOC_SECTION_AGGREGATE, |
522 | | DOC_SECTION_RANKING, |
523 | | DOC_SECTION_ANALYTICAL, |
524 | | ] |
525 | | } |
526 | | |
527 | | pub const DOC_SECTION_AGGREGATE: DocSection = DocSection { |
528 | | include: true, |
529 | | label: "Aggregate Functions", |
530 | | description: Some("All aggregate functions can be used as window functions."), |
531 | | }; |
532 | | |
533 | | pub const DOC_SECTION_RANKING: DocSection = DocSection { |
534 | | include: true, |
535 | | label: "Ranking Functions", |
536 | | description: None, |
537 | | }; |
538 | | |
539 | | pub const DOC_SECTION_ANALYTICAL: DocSection = DocSection { |
540 | | include: true, |
541 | | label: "Analytical Functions", |
542 | | description: None, |
543 | | }; |
544 | | } |
545 | | |
546 | | #[cfg(test)] |
547 | | mod test { |
548 | | use crate::{PartitionEvaluator, WindowUDF, WindowUDFImpl}; |
549 | | use arrow::datatypes::{DataType, Field}; |
550 | | use datafusion_common::Result; |
551 | | use datafusion_expr_common::signature::{Signature, Volatility}; |
552 | | use datafusion_functions_window_common::field::WindowUDFFieldArgs; |
553 | | use std::any::Any; |
554 | | use std::cmp::Ordering; |
555 | | |
556 | | #[derive(Debug, Clone)] |
557 | | struct AWindowUDF { |
558 | | signature: Signature, |
559 | | } |
560 | | |
561 | | impl AWindowUDF { |
562 | | fn new() -> Self { |
563 | | Self { |
564 | | signature: Signature::uniform( |
565 | | 1, |
566 | | vec![DataType::Int32], |
567 | | Volatility::Immutable, |
568 | | ), |
569 | | } |
570 | | } |
571 | | } |
572 | | |
573 | | /// Implement the WindowUDFImpl trait for AddOne |
574 | | impl WindowUDFImpl for AWindowUDF { |
575 | | fn as_any(&self) -> &dyn Any { |
576 | | self |
577 | | } |
578 | | fn name(&self) -> &str { |
579 | | "a" |
580 | | } |
581 | | fn signature(&self) -> &Signature { |
582 | | &self.signature |
583 | | } |
584 | | fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> { |
585 | | unimplemented!() |
586 | | } |
587 | | fn field(&self, _field_args: WindowUDFFieldArgs) -> Result<Field> { |
588 | | unimplemented!() |
589 | | } |
590 | | } |
591 | | |
592 | | #[derive(Debug, Clone)] |
593 | | struct BWindowUDF { |
594 | | signature: Signature, |
595 | | } |
596 | | |
597 | | impl BWindowUDF { |
598 | | fn new() -> Self { |
599 | | Self { |
600 | | signature: Signature::uniform( |
601 | | 1, |
602 | | vec![DataType::Int32], |
603 | | Volatility::Immutable, |
604 | | ), |
605 | | } |
606 | | } |
607 | | } |
608 | | |
609 | | /// Implement the WindowUDFImpl trait for AddOne |
610 | | impl WindowUDFImpl for BWindowUDF { |
611 | | fn as_any(&self) -> &dyn Any { |
612 | | self |
613 | | } |
614 | | fn name(&self) -> &str { |
615 | | "b" |
616 | | } |
617 | | fn signature(&self) -> &Signature { |
618 | | &self.signature |
619 | | } |
620 | | fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> { |
621 | | unimplemented!() |
622 | | } |
623 | | fn field(&self, _field_args: WindowUDFFieldArgs) -> Result<Field> { |
624 | | unimplemented!() |
625 | | } |
626 | | } |
627 | | |
628 | | #[test] |
629 | | fn test_partial_ord() { |
630 | | let a1 = WindowUDF::from(AWindowUDF::new()); |
631 | | let a2 = WindowUDF::from(AWindowUDF::new()); |
632 | | assert_eq!(a1.partial_cmp(&a2), Some(Ordering::Equal)); |
633 | | |
634 | | let b1 = WindowUDF::from(BWindowUDF::new()); |
635 | | assert!(a1 < b1); |
636 | | assert!(!(a1 == b1)); |
637 | | } |
638 | | } |