datafusion_spark/function/datetime/
add_months.rs

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
18use std::ops::Add;
19use std::sync::Arc;
20
21use arrow::datatypes::{DataType, Field, FieldRef, IntervalUnit};
22use datafusion_common::utils::take_function_args;
23use datafusion_common::{Result, internal_err};
24use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext};
25use datafusion_expr::{
26    ColumnarValue, Expr, ExprSchemable, ReturnFieldArgs, ScalarFunctionArgs,
27    ScalarUDFImpl, Signature, Volatility,
28};
29
30/// <https://spark.apache.org/docs/latest/api/sql/index.html#add_months>
31#[derive(Debug, PartialEq, Eq, Hash)]
32pub struct SparkAddMonths {
33    signature: Signature,
34}
35
36impl Default for SparkAddMonths {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42impl SparkAddMonths {
43    pub fn new() -> Self {
44        Self {
45            signature: Signature::exact(
46                vec![DataType::Date32, DataType::Int32],
47                Volatility::Immutable,
48            ),
49        }
50    }
51}
52
53impl ScalarUDFImpl for SparkAddMonths {
54    fn name(&self) -> &str {
55        "add_months"
56    }
57
58    fn signature(&self) -> &Signature {
59        &self.signature
60    }
61
62    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
63        internal_err!("return_field_from_args should be used instead")
64    }
65
66    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
67        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
68
69        Ok(Arc::new(Field::new(
70            self.name(),
71            DataType::Date32,
72            nullable,
73        )))
74    }
75
76    fn simplify(
77        &self,
78        args: Vec<Expr>,
79        info: &SimplifyContext,
80    ) -> Result<ExprSimplifyResult> {
81        let [date_arg, months_arg] = take_function_args("add_months", args)?;
82        let interval = months_arg
83            .cast_to(&DataType::Interval(IntervalUnit::YearMonth), info.schema())?;
84        Ok(ExprSimplifyResult::Simplified(date_arg.add(interval)))
85    }
86
87    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
88        internal_err!("invoke should not be called on a simplified add_months() function")
89    }
90}
datafusion_spark/function/datetime/add_months.rs

datafusion_spark/function/datetime/
add_months.rs
