use arrow::datatypes::DataType;
use datafusion_common::{utils::take_function_args, Result, ScalarValue};
use datafusion_expr::{ColumnarValue, Documentation, ScalarFunctionArgs};
use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
use datafusion_macros::user_doc;
use std::any::Any;
#[user_doc(
doc_section(label = "Other Functions"),
description = "Returns the name of the underlying [Arrow data type](https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html) of the expression.",
syntax_example = "arrow_typeof(expression)",
sql_example = r#"```sql
> select arrow_typeof('foo'), arrow_typeof(1);
+---------------------------+------------------------+
| arrow_typeof(Utf8("foo")) | arrow_typeof(Int64(1)) |
+---------------------------+------------------------+
| Utf8 | Int64 |
+---------------------------+------------------------+
```
"#,
argument(
name = "expression",
description = "Expression to evaluate. The expression can be a constant, column, or function, and any combination of operators."
)
)]
#[derive(Debug)]
pub struct ArrowTypeOfFunc {
signature: Signature,
}
impl Default for ArrowTypeOfFunc {
fn default() -> Self {
Self::new()
}
}
impl ArrowTypeOfFunc {
pub fn new() -> Self {
Self {
signature: Signature::any(1, Volatility::Immutable),
}
}
}
impl ScalarUDFImpl for ArrowTypeOfFunc {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"arrow_typeof"
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(DataType::Utf8)
}
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
let [arg] = take_function_args(self.name(), args.args)?;
let input_data_type = arg.data_type();
Ok(ColumnarValue::Scalar(ScalarValue::from(format!(
"{input_data_type}"
))))
}
fn documentation(&self) -> Option<&Documentation> {
self.doc()
}
}