Skip to content

Commit 13f3843

Browse files
authored
Introduce TypeSignatureClass::Any (#19485)
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #19438 ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> See issue. ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> Add new `TypeSignatureClass` variant `Any` and refactor `arrow_typeof` and `arrow_metadata` function signatures to use this. ## Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> Existing tests. ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> No. <!-- If there are any breaking changes to public APIs, please add the `api change` label. -->
1 parent 1d2b389 commit 13f3843

File tree

5 files changed

+42
-13
lines changed

5 files changed

+42
-13
lines changed

datafusion/expr-common/src/signature.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,8 @@ impl TypeSignature {
332332
/// arguments that can be coerced to a particular class of types.
333333
#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Hash)]
334334
pub enum TypeSignatureClass {
335+
/// Allows an arbitrary type argument without coercing the argument.
336+
Any,
335337
/// Timestamps, allowing arbitrary (or no) timezones
336338
Timestamp,
337339
/// All time types
@@ -367,6 +369,9 @@ impl TypeSignatureClass {
367369
/// documentation or error messages.
368370
fn get_example_types(&self) -> Vec<DataType> {
369371
match self {
372+
// TODO: might be too much info to return every single type here
373+
// maybe https://github.com/apache/datafusion/issues/14761 will help here?
374+
TypeSignatureClass::Any => vec![],
370375
TypeSignatureClass::Native(l) => get_data_types(l.native()),
371376
TypeSignatureClass::Timestamp => {
372377
vec![
@@ -409,6 +414,7 @@ impl TypeSignatureClass {
409414
}
410415

411416
match self {
417+
TypeSignatureClass::Any => true,
412418
TypeSignatureClass::Native(t) if t.native() == logical_type => true,
413419
TypeSignatureClass::Timestamp if logical_type.is_timestamp() => true,
414420
TypeSignatureClass::Time if logical_type.is_time() => true,
@@ -430,6 +436,7 @@ impl TypeSignatureClass {
430436
origin_type: &DataType,
431437
) -> Result<DataType> {
432438
match self {
439+
TypeSignatureClass::Any => Ok(origin_type.to_owned()),
433440
TypeSignatureClass::Native(logical_type) => {
434441
logical_type.native().default_cast_for(origin_type)
435442
}

datafusion/functions/src/core/arrow_cast.rs

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
2020
use arrow::datatypes::{DataType, Field, FieldRef};
2121
use arrow::error::ArrowError;
22+
use datafusion_common::types::logical_string;
2223
use datafusion_common::{
2324
Result, ScalarValue, arrow_datafusion_err, exec_err, internal_err,
2425
};
@@ -27,8 +28,8 @@ use std::any::Any;
2728

2829
use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
2930
use datafusion_expr::{
30-
ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs,
31-
ScalarUDFImpl, Signature, Volatility,
31+
Coercion, ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs,
32+
ScalarUDFImpl, Signature, TypeSignatureClass, Volatility,
3233
};
3334
use datafusion_macros::user_doc;
3435

@@ -102,7 +103,13 @@ impl Default for ArrowCastFunc {
102103
impl ArrowCastFunc {
103104
pub fn new() -> Self {
104105
Self {
105-
signature: Signature::any(2, Volatility::Immutable),
106+
signature: Signature::coercible(
107+
vec![
108+
Coercion::new_exact(TypeSignatureClass::Any),
109+
Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
110+
],
111+
Volatility::Immutable,
112+
),
106113
}
107114
}
108115
}

datafusion/functions/src/core/arrow_metadata.rs

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,11 @@
1717

1818
use arrow::array::{MapBuilder, StringBuilder};
1919
use arrow::datatypes::{DataType, Field, Fields};
20-
use datafusion_common::{Result, ScalarValue, exec_err};
20+
use datafusion_common::types::logical_string;
21+
use datafusion_common::{Result, ScalarValue, exec_err, internal_err};
2122
use datafusion_expr::{
22-
ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
23-
Volatility,
23+
Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
24+
TypeSignature, TypeSignatureClass, Volatility,
2425
};
2526
use datafusion_macros::user_doc;
2627
use std::any::Any;
@@ -29,7 +30,7 @@ use std::sync::Arc;
2930
#[user_doc(
3031
doc_section(label = "Other Functions"),
3132
description = "Returns the metadata of the input expression. If a key is provided, returns the value for that key. If no key is provided, returns a Map of all metadata.",
32-
syntax_example = "arrow_metadata(expression, [key])",
33+
syntax_example = "arrow_metadata(expression[, key])",
3334
sql_example = r#"```sql
3435
> select arrow_metadata(col) from table;
3536
+----------------------------+
@@ -61,7 +62,18 @@ pub struct ArrowMetadataFunc {
6162
impl ArrowMetadataFunc {
6263
pub fn new() -> Self {
6364
Self {
64-
signature: Signature::variadic_any(Volatility::Immutable),
65+
signature: Signature::one_of(
66+
vec![
67+
TypeSignature::Coercible(vec![Coercion::new_exact(
68+
TypeSignatureClass::Any,
69+
)]),
70+
TypeSignature::Coercible(vec![
71+
Coercion::new_exact(TypeSignatureClass::Any),
72+
Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
73+
]),
74+
],
75+
Volatility::Immutable,
76+
),
6577
}
6678
}
6779
}
@@ -105,7 +117,7 @@ impl ScalarUDFImpl for ArrowMetadataFunc {
105117
false,
106118
))
107119
} else {
108-
exec_err!("arrow_metadata requires 1 or 2 arguments")
120+
internal_err!("arrow_metadata requires 1 or 2 arguments")
109121
}
110122
}
111123

@@ -114,7 +126,7 @@ impl ScalarUDFImpl for ArrowMetadataFunc {
114126

115127
if args.args.len() == 2 {
116128
let key = match &args.args[1] {
117-
ColumnarValue::Scalar(ScalarValue::Utf8(Some(k))) => k,
129+
ColumnarValue::Scalar(ScalarValue::Utf8(Some(key))) => key,
118130
_ => {
119131
return exec_err!(
120132
"Second argument to arrow_metadata must be a string literal key"
@@ -142,7 +154,7 @@ impl ScalarUDFImpl for ArrowMetadataFunc {
142154
&map_array, 0,
143155
)?))
144156
} else {
145-
exec_err!("arrow_metadata requires 1 or 2 arguments")
157+
internal_err!("arrow_metadata requires 1 or 2 arguments")
146158
}
147159
}
148160
}

datafusion/sqllogictest/test_files/arrow_typeof.slt

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,9 +95,12 @@ SELECT arrow_cast('1', 'Int16')
9595
query error
9696
SELECT arrow_cast('1')
9797

98-
query error DataFusion error: Execution error: arrow_cast requires its second argument to be a non\-empty constant string
98+
query error Expect TypeSignatureClass::Native\(LogicalType\(Native\(String\), String\)\) but received NativeType::Int64, DataType: Int64
9999
SELECT arrow_cast('1', 43)
100100

101+
query error DataFusion error: Execution error: arrow_cast requires its second argument to be a non\-empty constant string
102+
SELECT arrow_cast('1', arrow_cast('Utf8', 'Utf8'))
103+
101104
query error DataFusion error: Execution error: Unsupported type 'unknown'\. Must be a supported arrow type name such as 'Int32' or 'Timestamp\(ns\)'\. Error unknown token: unknown
102105
SELECT arrow_cast('1', 'unknown')
103106

docs/source/user-guide/sql/scalar_functions.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5047,7 +5047,7 @@ arrow_cast(expression, datatype)
50475047
Returns the metadata of the input expression. If a key is provided, returns the value for that key. If no key is provided, returns a Map of all metadata.
50485048

50495049
```sql
5050-
arrow_metadata(expression, [key])
5050+
arrow_metadata(expression[, key])
50515051
```
50525052

50535053
#### Arguments

0 commit comments

Comments
 (0)