feat: add `arrow_cast` function to support supports arbitrary arrow types #5166

alamb · 2023-02-02T22:07:58Z

Which issue does this PR close?

Rationale for this change

This function is important to be able to test thing such as DictionaryArray via sql (datafusion-cli as well as sqllogictests). It also will help control output into parquet files more precisely, for example.

See #5016 for more details

What changes are included in this PR?

Add special handling for arrow_cast in sql planner
code to convert from string to an DataType (opposite of data_type.to_string())
sqllogicbased tests
Documentation of the new function

Are these changes tested?

Yes

Are there any user-facing changes?

a new arrow_cast function

alamb · 2023-02-02T22:08:19Z

datafusion/core/tests/sqllogictests/test_files/arrow_typeof.slt

-# SELECT arrow_typeof('1')
-# ----
-# Utf8
+query T


I just uncommented out these tests -- I am not sure why they were commented out 🤷

alamb · 2023-02-02T22:08:50Z

datafusion/sql/src/expr/function.rs

        }
+
+        // Special case arrow_cast
+        if &name == ARROW_CAST_NAME {


This is the main difference -- if all other function resolution fails, try to resolve it as arrow_cast before erroring

alamb · 2023-02-02T22:09:26Z

datafusion/sql/src/expr/function.rs

-        match self.schema_provider.get_function_meta(&name) {
-            Some(fm) => {
-                let args = self.function_args_to_expr(function.args, schema)?;
+        if let Some(fm) = self.schema_provider.get_function_meta(&name) {


This change was just to reduce the indent level (as adding support for arrow_cast in the existing structure would have added another level)

alamb · 2023-02-02T22:09:57Z

datafusion/sql/tests/integration_test.rs

@@ -2497,7 +2506,7 @@ impl ContextProvider for MockContextProvider {
    }

    fn get_function_meta(&self, _name: &str) -> Option<Arc<ScalarUDF>> {
-        unimplemented!()
+        None


otherwise the test context panic's before it gets to seeing arrow_cast

alamb · 2023-02-02T22:10:39Z

datafusion/sql/src/expr/cast.rs

+    // do the actual lookup to the appropriate data type
+    let data_type = parse_data_type(&data_type_string)?;
+
+    return arg0.cast_to(&data_type, schema);


the actual implementation simply calls the existing Expr::Cast

alamb · 2023-03-03T18:05:05Z

datafusion/core/tests/sqllogictests/test_files/arrow_typeof.slt

+
+statement ok
+create table foo as select
+  arrow_cast(to_timestamp('2020-01-02 01:01:11.1234567890Z'), 'Timestamp(Second, None)') as col_ts_s,


This is kind of cool -- it shows how to convert to arbitrary Arrow timestamp types.

I wonder if we should deprecate / remove to_timestamp_millis, to_timstamp_nanos, ... 🤔 Any thoughts @waitingkuo ?

this is great! thank you @alamb

i wonder if to_timestamp_micros with two parameters (as mentioned in #5398) is needed. If not, I think it's a great timing to deprecate them and align the current to_timestamp to postgrseql's (1. return type 2. accept the 2nd argument).

i wonder if to_timestamp_micros with two parameters (as mentioned in #5398) is needed

I think to_timestamp with two parameters (specifically with user defined formatting) is still a feature gap after this PR.

What I was thinking is that the to_timestamp_millis, to_timestamp_secs, etc variants whose only difference is the output arrow timestamp precision (TimeUnit) are not needed as arrow_cast can be used to convert between those precisions

alamb · 2023-03-03T18:06:34Z

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

+
+## Supported Arrow Types
+
+The following types are supported by the `arrow_typeof` function:


I am pretty stoked about the ability to cast to types that have no direct SQL mapping (for IOx it is dictionary types) but the different timestamp and durations are also quite relevant I think

alamb · 2023-03-03T18:08:33Z

datafusion/sql/src/expr/arrow_cast.rs

+    arg0.cast_to(&data_type, schema)
+}
+
+/// Parses `str` into a `DataType`.


Here is the main entrypoint. This uses a Parser very much like the one in https://github.com/sqlparser-rs/sqlparser-rs because I was familiar with that structure and I figured that other contributors to DataFusion might be too

I am not sure what reviewers think of proposing putting this in the arrow-rs crate -- it feels like a better fit there to me

filed apache/arrow-rs#3821

alamb · 2023-03-03T18:09:38Z

datafusion/sql/src/expr/arrow_cast.rs

+                    ),
+                ),
+            ),
+            // TODO support more structured types (List, LargeList, Struct, Union, Map, RunEndEncoded, etc)


List, Struct, etc is not important for my use cases (yet) so I didn't add support for the yet, but I think all the necessary patterns are present. Someone who cares just needs to implement them

alamb · 2023-03-04T21:16:24Z

This PR is now ready for review

waitingkuo

hi @alamb thank you, this feature is great.

I'll be happy to implement the timezone parsing after this pr merged

datafusion/sql/src/expr/arrow_cast.rs

Co-authored-by: Wei-Ting Kuo <waitingkuo0527@gmail.com>

alamb · 2023-03-07T21:50:35Z

@comphead @Jefffrey or @ygf11 might one of you have time to help review this pull request?

alamb · 2023-03-07T21:50:58Z

(I would like to use it to write tests)

comphead

Great job @alamb

comphead · 2023-03-07T22:10:39Z

datafusion/core/tests/sqllogictests/test_files/arrow_typeof.slt

+  arrow_typeof(col_u64),
+  arrow_typeof(col_f32),
+  arrow_typeof(col_f64)
+  FROM foo;


nit: you can remove from foo

When I did that I got the error:

arrow_typeof(col_f64);" Error: query failed: DataFusion error: Schema error: No field named 'col_i8'. [SQL] SELECT arrow_typeof(col_i8), arrow_typeof(col_i16), arrow_typeof(col_i32), arrow_typeof(col_i64), arrow_typeof(col_u8), arrow_typeof(col_u16), arrow_typeof(col_u32), arrow_typeof(col_u64), arrow_typeof(col_f32), arrow_typeof(col_f64);

The point of this test was that the values inserted into the foo table do indeed have the proper types. I will add some comments to make this clearer.

comphead · 2023-03-07T22:13:53Z

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

+You can cast a SQL expression to a specific Arrow type using the `arrow_cast` function
+For example, to cast the output of `now()` to a `Timestamp` with second precision rather:
+
+```sql


what if select arrow_cast(now(), 'Timestamp(Second, Some("+00:00"))')

comphead · 2023-03-07T22:14:18Z

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

+| `Timestamp(Second, None)`                                   |
+| `Timestamp(Millisecond, None)`                              |
+| `Timestamp(Microsecond, None)`                              |
+| `Timestamp(Nanosecond, None)`                               |


No support for timestamptz?

Sadly, it wasn't important for my initial usecase -- @waitingkuo has nicely offered to implement it https://github.com/apache/arrow-datafusion/pull/5166/files#r1125621584 as a follow on. I will file appropriate tickets.

comphead · 2023-03-07T22:23:15Z

datafusion/core/tests/sqllogictests/test_files/arrow_typeof.slt

+
+query TTTT
+SELECT
+  arrow_typeof(col_ts_s),


Would be also great to have tests select arrow_type_of(arrow_cast(.....))

Added in a4f2753

tustvold

Just some relatively minor comments, LGTM

tustvold · 2023-03-08T15:51:34Z

datafusion/sql/src/expr/arrow_cast.rs

+
+    /// parse the characters in val starting at pos, until the next
+    /// `,`, `(`, or `)` or end of line
+    fn parse_word(&mut self) -> Result<Token> {


Do we need to copy here, could we instead use something like str::split_once perhaps?

I will try this

I spent some time messing around with using split_once -- it was better to avoid copying but I think it made the code quite ugly.

I eventually went with a compromise approach of reusing the String in 84901bf which saves an allocation on each token.

tustvold · 2023-03-08T15:52:35Z

datafusion/sql/src/expr/arrow_cast.rs

+/// * Token::Rparen,
+struct Tokenizer<'a> {
+    val: &'a str,
+    chars: Peekable<Chars<'a>>,


Unless I am mistaken, all tokens are ASCII, and so I think this could use bytes directly without needing to worry about UTF-8 shenanigans...

I may be able to use bytes directly, though then I found it challenging to create str from those bytes without revalidating each time the data is uf8. Maybe using bytes directly would be a good performance optimization in the future 🤔

Given this function is called once per call to arrow_cast I don't think it is performance critical (yet)

tustvold · 2023-03-08T16:10:22Z

datafusion/sql/src/expr/arrow_cast.rs

+/// assert_eq!(data_type, DataType::Int32);
+/// ```
+///
+/// TODO file a ticket about bringing this into arrow possibly


I think having FromStr and Display implementations for DataType would be very compelling 👍

apache/arrow-rs#3821

tustvold · 2023-03-08T16:16:14Z

datafusion/sql/src/expr/arrow_cast.rs

+    /// convert data_type to a string, and then parse it as a type
+    /// verifying it is the same
+    fn round_trip(data_type: DataType) {
+        let data_type_string = data_type.to_string();


I like that this is round-tripping the Display output 👍

Jefffrey · 2023-03-08T19:59:46Z

Just did some quick testing, noticed a minor quirk:

❯ select arrow_cast(to_timestamp('2020-01-02 01:01:11.1234567890Z'), '          Timestamp        (Nanosecond,      None               )');
+------------------------------------------------------+
| totimestamp(Utf8("2020-01-02 01:01:11.1234567890Z")) |
+------------------------------------------------------+
| 2020-01-02T01:01:11.123456789                        |
+------------------------------------------------------+
1 row in set. Query took 0.003 seconds.
❯ select arrow_cast(to_timestamp('2020-01-02 01:01:11.1234567890Z'), '          Timestamp        (Nanosecond,      None               ) ');
Error during planning: Unsupported type '          Timestamp        (Nanosecond,      None               ) '. Must be a supported arrow type name such as 'Int32' or 'Timestamp(Nanosecond, None)'. Error checking trailing content after parsing 'Timestamp(Nanosecond, None)'
❯

The parser is very tolerant of whitespace in the constant string containing the type unless the whitespace is at the very end, unsure if this is intended

alamb · 2023-03-08T20:39:52Z

Just did some quick testing, noticed a minor quirk:

Great catch. Thank you @Jefffrey -- It was not intended. Fixed in ff5d72b

ursabot · 2023-03-08T22:42:54Z

Benchmark runs are scheduled for baseline = 84530a2 and contender = e46924d. e46924d is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on ec2-t3-xlarge-us-east-2] ec2-t3-xlarge-us-east-2
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on test-mac-arm] test-mac-arm
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on ursa-i9-9960x] ursa-i9-9960x
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on ursa-thinkcentre-m75q] ursa-thinkcentre-m75q
Buildkite builds:
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java

github-actions bot added core Core DataFusion crate sql SQL Planner sqllogictest SQL Logic Tests (.slt) labels Feb 2, 2023

alamb commented Feb 2, 2023

View reviewed changes

alamb force-pushed the alamb/arrow_cast branch from e43c546 to 3c6f139 Compare March 3, 2023 16:14

Add arrow_cast function

52e5c58

alamb force-pushed the alamb/arrow_cast branch from 3c6f139 to 52e5c58 Compare March 3, 2023 17:58

prettier

a58f00a

alamb commented Mar 3, 2023

View reviewed changes

alamb mentioned this pull request Mar 3, 2023

upper (and other string functions) don't support String Dictionary types: Internal error: The "upper" function can only accept strings. #5471

Closed

alamb marked this pull request as ready for review March 3, 2023 18:16

alamb mentioned this pull request Mar 4, 2023

Use _dyn and _scalar_dyn for operations like, ilike, nlike, and nilike to support Dictionary and LargeStringArray #5227

Closed

alamb mentioned this pull request Mar 4, 2023

Add casting from arbitrary arrow types #5016

Closed

waitingkuo reviewed Mar 5, 2023

View reviewed changes

alamb and others added 2 commits March 5, 2023 06:56

Update datafusion/sql/src/expr/arrow_cast.rs

78965db

Co-authored-by: Wei-Ting Kuo <waitingkuo0527@gmail.com>

Apply suggestions from code review

3df081f

Co-authored-by: Wei-Ting Kuo <waitingkuo0527@gmail.com>

comphead approved these changes Mar 7, 2023

View reviewed changes

tustvold approved these changes Mar 8, 2023

View reviewed changes

alamb added 6 commits March 8, 2023 13:59

Clarify intent of tests

d3b17e8

Add more error tests

7fa282b

Merge remote-tracking branch 'apache/main' into alamb/arrow_cast

739d1c0

More tests

a4f2753

fix test

b3defe4

reuse buffer to avoid an allocation per word

84901bf

alamb mentioned this pull request Mar 8, 2023

Implement FromStr for DataType / Parse DataType description apache/arrow-rs#3821

Closed

2 tasks

add ticket link

1e6b648

alamb added 2 commits March 8, 2023 15:38

allow trailing whitespace, add tests for whitespace

ff5d72b

Merge remote-tracking branch 'apache/main' into alamb/arrow_cast

6c5c55e

alamb merged commit e46924d into apache:main Mar 8, 2023

alamb mentioned this pull request Mar 12, 2023

Improve performance of COUNT (distinct x) for dictionary columns #258

Open

andygrove added the enhancement New feature or request label Mar 12, 2023

alamb mentioned this pull request Mar 13, 2023

Minor: Port more aggregate tests to sqllogictests #5574

Merged

waitingkuo mentioned this pull request Apr 7, 2023

allow ARROW_CAST to cast to TimestampTz #5913

Closed

alamb deleted the alamb/arrow_cast branch October 18, 2024 20:46


		## Supported Arrow Types

		The following types are supported by the `arrow_typeof` function:

feat: add arrow_cast function to support supports arbitrary arrow types #5166

feat: add arrow_cast function to support supports arbitrary arrow types #5166

Conversation

alamb commented Feb 2, 2023 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented Mar 4, 2023

waitingkuo left a comment • edited Loading

Choose a reason for hiding this comment

alamb commented Mar 7, 2023

alamb commented Mar 7, 2023

comphead left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tustvold left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Jefffrey commented Mar 8, 2023

alamb commented Mar 8, 2023

ursabot commented Mar 8, 2023

feat: add `arrow_cast` function to support supports arbitrary arrow types #5166

feat: add `arrow_cast` function to support supports arbitrary arrow types #5166

alamb commented Feb 2, 2023 •

edited

Loading

waitingkuo left a comment •

edited

Loading