common_function/scalars/math/
clamp.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::{self, Display};
16use std::sync::Arc;
17
18use datafusion::arrow::array::{Array, ArrayRef, AsArray, PrimitiveArray};
19use datafusion::arrow::datatypes::DataType as ArrowDataType;
20use datafusion::logical_expr::{ColumnarValue, Volatility};
21use datafusion_common::{DataFusionError, ScalarValue, utils};
22use datafusion_expr::{ScalarFunctionArgs, Signature};
23
24use crate::function::Function;
25use crate::helper::NUMERICS;
26
27#[derive(Clone, Debug)]
28pub struct ClampFunction {
29    signature: Signature,
30}
31
32impl Default for ClampFunction {
33    fn default() -> Self {
34        Self {
35            // input, min, max
36            signature: Signature::uniform(3, NUMERICS.to_vec(), Volatility::Immutable),
37        }
38    }
39}
40
41const CLAMP_NAME: &str = "clamp";
42
43impl Function for ClampFunction {
44    fn name(&self) -> &str {
45        CLAMP_NAME
46    }
47
48    fn return_type(
49        &self,
50        input_types: &[ArrowDataType],
51    ) -> datafusion_common::Result<ArrowDataType> {
52        // Type check is done by `signature`
53        Ok(input_types[0].clone())
54    }
55
56    fn signature(&self) -> &Signature {
57        &self.signature
58    }
59
60    fn invoke_with_args(
61        &self,
62        args: ScalarFunctionArgs,
63    ) -> datafusion_common::Result<ColumnarValue> {
64        let [col, min, max] = utils::take_function_args(self.name(), args.args)?;
65        clamp_impl(col, min, max)
66    }
67}
68
69impl Display for ClampFunction {
70    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
71        write!(f, "{}", CLAMP_NAME.to_ascii_uppercase())
72    }
73}
74
75fn clamp_impl(
76    col: ColumnarValue,
77    min: ColumnarValue,
78    max: ColumnarValue,
79) -> datafusion_common::Result<ColumnarValue> {
80    if col.data_type() != min.data_type() || min.data_type() != max.data_type() {
81        return Err(DataFusionError::Execution(format!(
82            "argument data types mismatch: {}, {}, {}",
83            col.data_type(),
84            min.data_type(),
85            max.data_type(),
86        )));
87    }
88
89    macro_rules! with_match_numerics_types {
90        ($data_type:expr, | $_:tt $T:ident | $body:tt) => {{
91            macro_rules! __with_ty__ {
92                ( $_ $T:ident ) => {
93                    $body
94                };
95            }
96
97            use datafusion::arrow::datatypes::{
98                Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type,
99                UInt16Type, UInt32Type, UInt64Type,
100            };
101
102            match $data_type {
103                ArrowDataType::Int8 => Ok(__with_ty__! { Int8Type }),
104                ArrowDataType::Int16 => Ok(__with_ty__! { Int16Type }),
105                ArrowDataType::Int32 => Ok(__with_ty__! { Int32Type }),
106                ArrowDataType::Int64 => Ok(__with_ty__! { Int64Type }),
107                ArrowDataType::UInt8 => Ok(__with_ty__! { UInt8Type }),
108                ArrowDataType::UInt16 => Ok(__with_ty__! { UInt16Type }),
109                ArrowDataType::UInt32 => Ok(__with_ty__! { UInt32Type }),
110                ArrowDataType::UInt64 => Ok(__with_ty__! { UInt64Type }),
111                ArrowDataType::Float32 => Ok(__with_ty__! { Float32Type }),
112                ArrowDataType::Float64 => Ok(__with_ty__! { Float64Type }),
113                _ => Err(DataFusionError::Execution(format!(
114                    "unsupported numeric data type: '{}'",
115                    $data_type
116                ))),
117            }
118        }};
119    }
120
121    macro_rules! clamp {
122        ($v: ident, $min: ident, $max: ident) => {
123            if $v < $min {
124                $min
125            } else if $v > $max {
126                $max
127            } else {
128                $v
129            }
130        };
131    }
132
133    match (col, min, max) {
134        (ColumnarValue::Scalar(col), ColumnarValue::Scalar(min), ColumnarValue::Scalar(max)) => {
135            if min > max {
136                return Err(DataFusionError::Execution(format!(
137                    "min '{}' > max '{}'",
138                    min, max
139                )));
140            }
141            Ok(ColumnarValue::Scalar(clamp!(col, min, max)))
142        }
143
144        (ColumnarValue::Array(col), ColumnarValue::Array(min), ColumnarValue::Array(max)) => {
145            if col.len() != min.len() || col.len() != max.len() {
146                return Err(DataFusionError::Internal(
147                    "arguments not of same length".to_string(),
148                ));
149            }
150            let result = with_match_numerics_types!(
151                col.data_type(),
152                |$S| {
153                    let col = col.as_primitive::<$S>();
154                    let min = min.as_primitive::<$S>();
155                    let max = max.as_primitive::<$S>();
156                    Arc::new(PrimitiveArray::<$S>::from(
157                        (0..col.len())
158                            .map(|i| {
159                                let v = col.is_valid(i).then(|| col.value(i));
160                                // Index safety: checked above, all have same length.
161                                let min = min.is_valid(i).then(|| min.value(i));
162                                let max = max.is_valid(i).then(|| max.value(i));
163                                Ok(match (v, min, max) {
164                                    (Some(v), Some(min), Some(max)) => {
165                                        if min > max {
166                                            return Err(DataFusionError::Execution(format!(
167                                                "min '{}' > max '{}'",
168                                                min, max
169                                            )));
170                                        }
171                                        Some(clamp!(v, min, max))
172                                    },
173                                    _ => None,
174                                })
175                            })
176                            .collect::<datafusion_common::Result<Vec<_>>>()?,
177                        )
178                    ) as ArrayRef
179                }
180            )?;
181            Ok(ColumnarValue::Array(result))
182        }
183
184        (ColumnarValue::Array(col), ColumnarValue::Scalar(min), ColumnarValue::Scalar(max)) => {
185            if min.is_null() || max.is_null() {
186                return Err(DataFusionError::Execution(
187                    "argument 'min' or 'max' is null".to_string(),
188                ));
189            }
190            let min = min.to_array()?;
191            let max = max.to_array()?;
192            let result = with_match_numerics_types!(
193                col.data_type(),
194                |$S| {
195                    let col = col.as_primitive::<$S>();
196                    // Index safety: checked above, both are not nulls.
197                    let min = min.as_primitive::<$S>().value(0);
198                    let max = max.as_primitive::<$S>().value(0);
199                    if min > max {
200                        return Err(DataFusionError::Execution(format!(
201                            "min '{}' > max '{}'",
202                            min, max
203                        )));
204                    }
205                    Arc::new(PrimitiveArray::<$S>::from(
206                        (0..col.len())
207                            .map(|x| {
208                                col.is_valid(x).then(|| {
209                                    let v = col.value(x);
210                                    clamp!(v, min, max)
211                                })
212                            })
213                            .collect::<Vec<_>>(),
214                        )
215                    ) as ArrayRef
216                }
217            )?;
218            Ok(ColumnarValue::Array(result))
219        }
220        _ => Err(DataFusionError::Internal(
221            "argument column types mismatch".to_string(),
222        )),
223    }
224}
225
226#[derive(Clone, Debug)]
227pub struct ClampMinFunction {
228    signature: Signature,
229}
230
231impl Default for ClampMinFunction {
232    fn default() -> Self {
233        Self {
234            // input, min
235            signature: Signature::uniform(2, NUMERICS.to_vec(), Volatility::Immutable),
236        }
237    }
238}
239
240const CLAMP_MIN_NAME: &str = "clamp_min";
241
242impl Function for ClampMinFunction {
243    fn name(&self) -> &str {
244        CLAMP_MIN_NAME
245    }
246
247    fn return_type(
248        &self,
249        input_types: &[ArrowDataType],
250    ) -> datafusion_common::Result<ArrowDataType> {
251        Ok(input_types[0].clone())
252    }
253
254    fn signature(&self) -> &Signature {
255        &self.signature
256    }
257
258    fn invoke_with_args(
259        &self,
260        args: ScalarFunctionArgs,
261    ) -> datafusion_common::Result<ColumnarValue> {
262        let [col, min] = utils::take_function_args(self.name(), args.args)?;
263
264        let Some(max) = ScalarValue::max(&min.data_type()) else {
265            return Err(DataFusionError::Internal(format!(
266                "cannot find a max value for numeric data type {}",
267                min.data_type()
268            )));
269        };
270        clamp_impl(col, min, ColumnarValue::Scalar(max))
271    }
272}
273
274impl Display for ClampMinFunction {
275    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
276        write!(f, "{}", CLAMP_MIN_NAME.to_ascii_uppercase())
277    }
278}
279
280#[derive(Clone, Debug)]
281pub struct ClampMaxFunction {
282    signature: Signature,
283}
284
285impl Default for ClampMaxFunction {
286    fn default() -> Self {
287        Self {
288            // input, max
289            signature: Signature::uniform(2, NUMERICS.to_vec(), Volatility::Immutable),
290        }
291    }
292}
293
294const CLAMP_MAX_NAME: &str = "clamp_max";
295
296impl Function for ClampMaxFunction {
297    fn name(&self) -> &str {
298        CLAMP_MAX_NAME
299    }
300
301    fn return_type(
302        &self,
303        input_types: &[ArrowDataType],
304    ) -> datafusion_common::Result<ArrowDataType> {
305        Ok(input_types[0].clone())
306    }
307
308    fn signature(&self) -> &Signature {
309        &self.signature
310    }
311
312    fn invoke_with_args(
313        &self,
314        args: ScalarFunctionArgs,
315    ) -> datafusion_common::Result<ColumnarValue> {
316        let [col, max] = utils::take_function_args(self.name(), args.args)?;
317
318        let Some(min) = ScalarValue::min(&max.data_type()) else {
319            return Err(DataFusionError::Internal(format!(
320                "cannot find a min value for numeric data type {}",
321                max.data_type()
322            )));
323        };
324        clamp_impl(col, ColumnarValue::Scalar(min), max)
325    }
326}
327
328impl Display for ClampMaxFunction {
329    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
330        write!(f, "{}", CLAMP_MAX_NAME.to_ascii_uppercase())
331    }
332}
333
334#[cfg(test)]
335mod test {
336
337    use std::sync::Arc;
338
339    use arrow_schema::Field;
340    use datafusion_common::config::ConfigOptions;
341    use datatypes::arrow::array::{
342        ArrayRef, Decimal128Array, Float64Array, Int64Array, UInt64Array,
343    };
344    use datatypes::arrow_array::StringArray;
345
346    use super::*;
347
348    macro_rules! impl_test_eval {
349        ($func: ty) => {
350            impl $func {
351                fn test_eval(
352                    &self,
353                    args: Vec<ColumnarValue>,
354                    number_rows: usize,
355                ) -> datafusion_common::Result<ArrayRef> {
356                    let input_type = args[0].data_type();
357                    self.invoke_with_args(ScalarFunctionArgs {
358                        args,
359                        arg_fields: vec![],
360                        number_rows,
361                        return_field: Arc::new(Field::new("x", input_type, false)),
362                        config_options: Arc::new(ConfigOptions::new()),
363                    })
364                    .and_then(|v| ColumnarValue::values_to_arrays(&[v]).map_err(Into::into))
365                    .map(|mut a| a.remove(0))
366                }
367            }
368        };
369    }
370
371    impl_test_eval!(ClampFunction);
372    impl_test_eval!(ClampMinFunction);
373    impl_test_eval!(ClampMaxFunction);
374
375    fn decimal_array(values: Vec<i128>) -> ColumnarValue {
376        ColumnarValue::Array(Arc::new(
377            Decimal128Array::from(values)
378                .with_precision_and_scale(10, 2)
379                .unwrap(),
380        ))
381    }
382
383    fn decimal_scalar(value: i128) -> ColumnarValue {
384        ColumnarValue::Scalar(ScalarValue::Decimal128(Some(value), 10, 2))
385    }
386
387    #[allow(deprecated)]
388    fn evaluate_decimal(
389        function: &dyn Function,
390        args: Vec<ColumnarValue>,
391    ) -> datafusion_common::Result<ArrayRef> {
392        let input_types = args
393            .iter()
394            .map(ColumnarValue::data_type)
395            .collect::<Vec<_>>();
396        let planned_types = datafusion_expr::type_coercion::functions::data_types(
397            function.name(),
398            &input_types,
399            function.signature(),
400        )?;
401        let args = args
402            .into_iter()
403            .zip(planned_types)
404            .map(|(arg, planned_type)| arg.cast_to(&planned_type, None))
405            .collect::<datafusion_common::Result<Vec<_>>>()?;
406        function
407            .invoke_with_args(ScalarFunctionArgs {
408                args,
409                arg_fields: vec![],
410                number_rows: 3,
411                return_field: Arc::new(Field::new("x", ArrowDataType::Float64, false)),
412                config_options: Arc::new(ConfigOptions::new()),
413            })
414            .and_then(|value| value.to_array(3))
415    }
416
417    #[test]
418    fn clamp_decimal_coercion_executes_as_float64() {
419        let test_cases = [
420            (
421                Box::new(ClampFunction::default()) as Box<dyn Function>,
422                vec![
423                    decimal_array(vec![100, 300, 500]),
424                    decimal_scalar(200),
425                    decimal_scalar(400),
426                ],
427                vec![2.0, 3.0, 4.0],
428            ),
429            (
430                Box::new(ClampMinFunction::default()),
431                vec![decimal_array(vec![100, 300, 500]), decimal_scalar(200)],
432                vec![2.0, 3.0, 5.0],
433            ),
434            (
435                Box::new(ClampMaxFunction::default()),
436                vec![decimal_array(vec![100, 300, 500]), decimal_scalar(200)],
437                vec![1.0, 2.0, 2.0],
438            ),
439        ];
440
441        for (function, args, expected) in test_cases {
442            let result = evaluate_decimal(function.as_ref(), args).unwrap();
443            let expected: ArrayRef = Arc::new(Float64Array::from(expected));
444            assert_eq!(expected.as_ref(), result.as_ref());
445        }
446    }
447
448    #[test]
449    fn clamp_i64() {
450        let inputs = [
451            (
452                vec![Some(-3), Some(-2), Some(-1), Some(0), Some(1), Some(2)],
453                -1i64,
454                10i64,
455                vec![Some(-1), Some(-1), Some(-1), Some(0), Some(1), Some(2)],
456            ),
457            (
458                vec![Some(-3), Some(-2), Some(-1), Some(0), Some(1), Some(2)],
459                0i64,
460                0i64,
461                vec![Some(0), Some(0), Some(0), Some(0), Some(0), Some(0)],
462            ),
463            (
464                vec![Some(-3), None, Some(-1), None, None, Some(2)],
465                -2i64,
466                1i64,
467                vec![Some(-2), None, Some(-1), None, None, Some(1)],
468            ),
469            (
470                vec![None, None, None, None, None],
471                0i64,
472                1i64,
473                vec![None, None, None, None, None],
474            ),
475        ];
476
477        let func = ClampFunction::default();
478        for (in_data, min, max, expected) in inputs {
479            let number_rows = in_data.len();
480            let args = vec![
481                ColumnarValue::Array(Arc::new(Int64Array::from(in_data))),
482                ColumnarValue::Scalar(min.into()),
483                ColumnarValue::Scalar(max.into()),
484            ];
485            let result = func.test_eval(args, number_rows).unwrap();
486            let expected: ArrayRef = Arc::new(Int64Array::from(expected));
487            assert_eq!(expected.as_ref(), result.as_ref());
488        }
489    }
490
491    #[test]
492    fn clamp_u64() {
493        let inputs = [
494            (
495                vec![Some(0), Some(1), Some(2), Some(3), Some(4), Some(5)],
496                1u64,
497                3u64,
498                vec![Some(1), Some(1), Some(2), Some(3), Some(3), Some(3)],
499            ),
500            (
501                vec![Some(0), Some(1), Some(2), Some(3), Some(4), Some(5)],
502                0u64,
503                0u64,
504                vec![Some(0), Some(0), Some(0), Some(0), Some(0), Some(0)],
505            ),
506            (
507                vec![Some(0), None, Some(2), None, None, Some(5)],
508                1u64,
509                3u64,
510                vec![Some(1), None, Some(2), None, None, Some(3)],
511            ),
512            (
513                vec![None, None, None, None, None],
514                0u64,
515                1u64,
516                vec![None, None, None, None, None],
517            ),
518        ];
519
520        let func = ClampFunction::default();
521        for (in_data, min, max, expected) in inputs {
522            let number_rows = in_data.len();
523            let args = vec![
524                ColumnarValue::Array(Arc::new(UInt64Array::from(in_data))),
525                ColumnarValue::Scalar(min.into()),
526                ColumnarValue::Scalar(max.into()),
527            ];
528            let result = func.test_eval(args, number_rows).unwrap();
529            let expected: ArrayRef = Arc::new(UInt64Array::from(expected));
530            assert_eq!(expected.as_ref(), result.as_ref());
531        }
532    }
533
534    #[test]
535    fn clamp_f64() {
536        let inputs = [
537            (
538                vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)],
539                -1.0,
540                10.0,
541                vec![Some(-1.0), Some(-1.0), Some(-1.0), Some(0.0), Some(1.0)],
542            ),
543            (
544                vec![Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)],
545                0.0,
546                0.0,
547                vec![Some(0.0), Some(0.0), Some(0.0), Some(0.0)],
548            ),
549            (
550                vec![Some(-3.0), None, Some(-1.0), None, None, Some(2.0)],
551                -2.0,
552                1.0,
553                vec![Some(-2.0), None, Some(-1.0), None, None, Some(1.0)],
554            ),
555            (
556                vec![None, None, None, None, None],
557                0.0,
558                1.0,
559                vec![None, None, None, None, None],
560            ),
561        ];
562
563        let func = ClampFunction::default();
564        for (in_data, min, max, expected) in inputs {
565            let number_rows = in_data.len();
566            let args = vec![
567                ColumnarValue::Array(Arc::new(Float64Array::from(in_data))),
568                ColumnarValue::Scalar(min.into()),
569                ColumnarValue::Scalar(max.into()),
570            ];
571            let result = func.test_eval(args, number_rows).unwrap();
572            let expected: ArrayRef = Arc::new(Float64Array::from(expected));
573            assert_eq!(expected.as_ref(), result.as_ref());
574        }
575    }
576
577    #[test]
578    fn clamp_invalid_min_max() {
579        let input = vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)];
580        let min = 10.0;
581        let max = -1.0;
582
583        let func = ClampFunction::default();
584        let number_rows = input.len();
585        let args = vec![
586            ColumnarValue::Array(Arc::new(Float64Array::from(input))),
587            ColumnarValue::Scalar(min.into()),
588            ColumnarValue::Scalar(max.into()),
589        ];
590        let result = func.test_eval(args, number_rows);
591        assert!(result.is_err());
592    }
593
594    #[test]
595    fn clamp_type_not_match() {
596        let input = vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)];
597        let min = -1i64;
598        let max = 10u64;
599
600        let func = ClampFunction::default();
601        let number_rows = input.len();
602        let args = vec![
603            ColumnarValue::Array(Arc::new(Float64Array::from(input))),
604            ColumnarValue::Scalar(min.into()),
605            ColumnarValue::Scalar(max.into()),
606        ];
607        let result = func.test_eval(args, number_rows);
608        assert!(result.is_err());
609    }
610
611    #[test]
612    fn clamp_min_is_not_scalar() {
613        let input = vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)];
614        let min = -10.0;
615        let max = 1.0;
616
617        let func = ClampFunction::default();
618        let number_rows = input.len();
619        let args = vec![
620            ColumnarValue::Array(Arc::new(Float64Array::from(input))),
621            ColumnarValue::Array(Arc::new(Float64Array::from(vec![min, max]))),
622            ColumnarValue::Array(Arc::new(Float64Array::from(vec![max, min]))),
623        ];
624        let result = func.test_eval(args, number_rows);
625        assert!(result.is_err());
626    }
627
628    #[test]
629    fn clamp_no_max() {
630        let input = vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)];
631        let min = -10.0;
632
633        let func = ClampFunction::default();
634        let number_rows = input.len();
635        let args = vec![
636            ColumnarValue::Array(Arc::new(Float64Array::from(input))),
637            ColumnarValue::Scalar(min.into()),
638        ];
639        let result = func.test_eval(args, number_rows);
640        assert!(result.is_err());
641    }
642
643    #[test]
644    fn clamp_on_string() {
645        let input = vec![Some("foo"), Some("foo"), Some("foo"), Some("foo")];
646
647        let func = ClampFunction::default();
648        let number_rows = input.len();
649        let args = vec![
650            ColumnarValue::Array(Arc::new(StringArray::from(input))),
651            ColumnarValue::Scalar("bar".into()),
652            ColumnarValue::Scalar("baz".into()),
653        ];
654        let result = func.test_eval(args, number_rows);
655        assert!(result.is_err());
656    }
657
658    #[test]
659    fn clamp_min_i64() {
660        let inputs = [
661            (
662                vec![Some(-3), Some(-2), Some(-1), Some(0), Some(1), Some(2)],
663                -1i64,
664                vec![Some(-1), Some(-1), Some(-1), Some(0), Some(1), Some(2)],
665            ),
666            (
667                vec![Some(-3), None, Some(-1), None, None, Some(2)],
668                -2i64,
669                vec![Some(-2), None, Some(-1), None, None, Some(2)],
670            ),
671        ];
672
673        let func = ClampMinFunction::default();
674        for (in_data, min, expected) in inputs {
675            let number_rows = in_data.len();
676            let args = vec![
677                ColumnarValue::Array(Arc::new(Int64Array::from(in_data))),
678                ColumnarValue::Scalar(min.into()),
679            ];
680            let result = func.test_eval(args, number_rows).unwrap();
681            let expected: ArrayRef = Arc::new(Int64Array::from(expected));
682            assert_eq!(expected.as_ref(), result.as_ref());
683        }
684    }
685
686    #[test]
687    fn clamp_max_i64() {
688        let inputs = [
689            (
690                vec![Some(-3), Some(-2), Some(-1), Some(0), Some(1), Some(2)],
691                1i64,
692                vec![Some(-3), Some(-2), Some(-1), Some(0), Some(1), Some(1)],
693            ),
694            (
695                vec![Some(-3), None, Some(-1), None, None, Some(2)],
696                0i64,
697                vec![Some(-3), None, Some(-1), None, None, Some(0)],
698            ),
699        ];
700
701        let func = ClampMaxFunction::default();
702        for (in_data, max, expected) in inputs {
703            let number_rows = in_data.len();
704            let args = vec![
705                ColumnarValue::Array(Arc::new(Int64Array::from(in_data))),
706                ColumnarValue::Scalar(max.into()),
707            ];
708            let result = func.test_eval(args, number_rows).unwrap();
709            let expected: ArrayRef = Arc::new(Int64Array::from(expected));
710            assert_eq!(expected.as_ref(), result.as_ref());
711        }
712    }
713
714    #[test]
715    fn clamp_min_f64() {
716        let inputs = [(
717            vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)],
718            -1.0,
719            vec![Some(-1.0), Some(-1.0), Some(-1.0), Some(0.0), Some(1.0)],
720        )];
721
722        let func = ClampMinFunction::default();
723        for (in_data, min, expected) in inputs {
724            let number_rows = in_data.len();
725            let args = vec![
726                ColumnarValue::Array(Arc::new(Float64Array::from(in_data))),
727                ColumnarValue::Scalar(min.into()),
728            ];
729            let result = func.test_eval(args, number_rows).unwrap();
730            let expected: ArrayRef = Arc::new(Float64Array::from(expected));
731            assert_eq!(expected.as_ref(), result.as_ref());
732        }
733    }
734
735    #[test]
736    fn clamp_max_f64() {
737        let inputs = [(
738            vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)],
739            0.0,
740            vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(0.0)],
741        )];
742
743        let func = ClampMaxFunction::default();
744        for (in_data, max, expected) in inputs {
745            let number_rows = in_data.len();
746            let args = vec![
747                ColumnarValue::Array(Arc::new(Float64Array::from(in_data))),
748                ColumnarValue::Scalar(max.into()),
749            ];
750            let result = func.test_eval(args, number_rows).unwrap();
751            let expected: ArrayRef = Arc::new(Float64Array::from(expected));
752            assert_eq!(expected.as_ref(), result.as_ref());
753        }
754    }
755
756    #[test]
757    fn clamp_min_type_not_match() {
758        let input = vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)];
759        let min = -1i64;
760
761        let func = ClampMinFunction::default();
762        let number_rows = input.len();
763        let args = vec![
764            ColumnarValue::Array(Arc::new(Float64Array::from(input))),
765            ColumnarValue::Scalar(min.into()),
766        ];
767        let result = func.test_eval(args, number_rows);
768        assert!(result.is_err());
769    }
770
771    #[test]
772    fn clamp_max_type_not_match() {
773        let input = vec![Some(-3.0), Some(-2.0), Some(-1.0), Some(0.0), Some(1.0)];
774        let max = 1i64;
775
776        let func = ClampMaxFunction::default();
777        let number_rows = input.len();
778        let args = vec![
779            ColumnarValue::Array(Arc::new(Float64Array::from(input))),
780            ColumnarValue::Scalar(max.into()),
781        ];
782        let result = func.test_eval(args, number_rows);
783        assert!(result.is_err());
784    }
785}
common_function/scalars/math/clamp.rs

common_function/scalars/math/
clamp.rs
