Skip to content

Commit cbea7ad

Browse files
jorgecarleitaonevi-me
authored andcommitted
ARROW-11354: [Rust] Speed-up cast of dates and times (2-4x)
This PR improves the performance of certain time / date casts by using the brand new API proposed in #9235 . That API allows for a very fast execution of unary and infalible operations on primitive arrays, and this PR leverages that for cast operations that require a simple mathematical operation. ``` Switched to branch 'fast_unary' Compiling arrow v3.0.0-SNAPSHOT (/Users/jorgecarleitao/projects/arrow/rust/arrow) Finished bench [optimized] target(s) in 1m 06s Running /Users/jorgecarleitao/projects/arrow/rust/target/release/deps/cast_kernels-25ee76597a8b997b Gnuplot not found, using plotters backend cast date64 to date32 512 time: [1.1668 us 1.1706 us 1.1749 us] change: [-83.347% -83.248% -83.144%] (p = 0.00 < 0.05) Performance has improved. Found 5 outliers among 100 measurements (5.00%) 3 (3.00%) high mild 2 (2.00%) high severe cast date32 to date64 512 time: [899.73 ns 930.58 ns 971.56 ns] change: [-86.799% -86.520% -86.190%] (p = 0.00 < 0.05) Performance has improved. Found 10 outliers among 100 measurements (10.00%) 3 (3.00%) high mild 7 (7.00%) high severe cast time32s to time32ms 512 time: [728.73 ns 732.33 ns 735.90 ns] change: [-54.503% -54.201% -53.917%] (p = 0.00 < 0.05) Performance has improved. Found 2 outliers among 100 measurements (2.00%) 1 (1.00%) high mild 1 (1.00%) high severe cast time64ns to time32s 512 time: [4.8374 us 4.8447 us 4.8526 us] change: [-57.022% -56.791% -56.587%] (p = 0.00 < 0.05) Performance has improved. Found 8 outliers among 100 measurements (8.00%) 6 (6.00%) high mild 2 (2.00%) high severe ``` Closes #9297 from jorgecarleitao/fast_unary Authored-by: Jorge C. Leitao <jorgecarleitao@gmail.com> Signed-off-by: Neville Dipale <nevilledips@gmail.com>
1 parent 2696951 commit cbea7ad

File tree

3 files changed

+124
-61
lines changed

3 files changed

+124
-61
lines changed
Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
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+
18+
//! Defines kernels suitable to perform operations to primitive arrays.
19+
20+
use crate::array::{Array, ArrayData, PrimitiveArray};
21+
use crate::buffer::Buffer;
22+
use crate::datatypes::ArrowPrimitiveType;
23+
24+
#[inline]
25+
fn into_primitive_array_data<I: ArrowPrimitiveType, O: ArrowPrimitiveType>(
26+
array: &PrimitiveArray<I>,
27+
buffer: Buffer,
28+
) -> ArrayData {
29+
ArrayData::new(
30+
O::DATA_TYPE,
31+
array.len(),
32+
None,
33+
array.data_ref().null_buffer().cloned(),
34+
0,
35+
vec![buffer],
36+
vec![],
37+
)
38+
}
39+
40+
/// Applies an unary and infalible function to a primitive array.
41+
/// This is the fastest way to perform an operation on a primitive array when
42+
/// the benefits of a vectorized operation outweights the cost of branching nulls and non-nulls.
43+
/// # Implementation
44+
/// This will apply the function for all values, including those on null slots.
45+
/// This implies that the operation must be infalible for any value of the corresponding type
46+
/// or this function may panic.
47+
/// # Example
48+
/// ```rust
49+
/// # use arrow::array::Int32Array;
50+
/// # use arrow::datatypes::Int32Type;
51+
/// # use arrow::compute::kernels::arity::unary;
52+
/// # fn main() {
53+
/// let array = Int32Array::from(vec![Some(5), Some(7), None]);
54+
/// let c = unary::<_, _, Int32Type>(&array, |x| x * 2 + 1);
55+
/// assert_eq!(c, Int32Array::from(vec![Some(11), Some(15), None]));
56+
/// # }
57+
/// ```
58+
pub fn unary<I, F, O>(array: &PrimitiveArray<I>, op: F) -> PrimitiveArray<O>
59+
where
60+
I: ArrowPrimitiveType,
61+
O: ArrowPrimitiveType,
62+
F: Fn(I::Native) -> O::Native,
63+
{
64+
let values = array.values().iter().map(|v| op(*v));
65+
// JUSTIFICATION
66+
// Benefit
67+
// ~60% speedup
68+
// Soundness
69+
// `values` is an iterator with a known size because arrays are sized.
70+
let buffer = unsafe { Buffer::from_trusted_len_iter(values) };
71+
72+
let data = into_primitive_array_data::<_, O>(array, buffer);
73+
PrimitiveArray::<O>::from(std::sync::Arc::new(data))
74+
}

rust/arrow/src/compute/kernels/cast.rs

Lines changed: 49 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ use std::str;
3939
use std::sync::Arc;
4040

4141
use crate::compute::kernels::arithmetic::{divide, multiply};
42+
use crate::compute::kernels::arity::unary;
4243
use crate::datatypes::*;
4344
use crate::error::{ArrowError, Result};
4445
use crate::{array::*, compute::take};
@@ -569,45 +570,43 @@ pub fn cast(array: &ArrayRef, to_type: &DataType) -> Result<ArrayRef> {
569570
(Time64(_), Int64) => cast_array_data::<Int64Type>(array, to_type.clone()),
570571
(Date32(DateUnit::Day), Date64(DateUnit::Millisecond)) => {
571572
let date_array = array.as_any().downcast_ref::<Date32Array>().unwrap();
572-
let mut b = Date64Builder::new(array.len());
573-
for i in 0..array.len() {
574-
if array.is_null(i) {
575-
b.append_null()?;
576-
} else {
577-
b.append_value(date_array.value(i) as i64 * MILLISECONDS_IN_DAY)?;
578-
}
579-
}
580573

581-
Ok(Arc::new(b.finish()) as ArrayRef)
574+
let values =
575+
unary::<_, _, Date64Type>(date_array, |x| x as i64 * MILLISECONDS_IN_DAY);
576+
577+
Ok(Arc::new(values) as ArrayRef)
582578
}
583579
(Date64(DateUnit::Millisecond), Date32(DateUnit::Day)) => {
584580
let date_array = array.as_any().downcast_ref::<Date64Array>().unwrap();
585-
let mut b = Date32Builder::new(array.len());
586-
for i in 0..array.len() {
587-
if array.is_null(i) {
588-
b.append_null()?;
589-
} else {
590-
b.append_value((date_array.value(i) / MILLISECONDS_IN_DAY) as i32)?;
591-
}
592-
}
593581

594-
Ok(Arc::new(b.finish()) as ArrayRef)
582+
let values = unary::<_, _, Date32Type>(date_array, |x| {
583+
(x / MILLISECONDS_IN_DAY) as i32
584+
});
585+
586+
Ok(Arc::new(values) as ArrayRef)
595587
}
596588
(Time32(TimeUnit::Second), Time32(TimeUnit::Millisecond)) => {
597-
let time_array = Time32MillisecondArray::from(array.data());
598-
let mult =
599-
Time32MillisecondArray::from(vec![MILLISECONDS as i32; array.len()]);
600-
let time32_ms = multiply(&time_array, &mult)?;
589+
let time_array = array.as_any().downcast_ref::<Time32SecondArray>().unwrap();
601590

602-
Ok(Arc::new(time32_ms) as ArrayRef)
591+
let values = unary::<_, _, Time32MillisecondType>(time_array, |x| {
592+
x * MILLISECONDS as i32
593+
});
594+
595+
Ok(Arc::new(values) as ArrayRef)
603596
}
604597
(Time32(TimeUnit::Millisecond), Time32(TimeUnit::Second)) => {
605-
let time_array = Time32SecondArray::from(array.data());
606-
let divisor = Time32SecondArray::from(vec![MILLISECONDS as i32; array.len()]);
607-
let time32_s = divide(&time_array, &divisor)?;
598+
let time_array = array
599+
.as_any()
600+
.downcast_ref::<Time32MillisecondArray>()
601+
.unwrap();
602+
603+
let values = unary::<_, _, Time32SecondType>(time_array, |x| {
604+
x / (MILLISECONDS as i32)
605+
});
608606

609-
Ok(Arc::new(time32_s) as ArrayRef)
607+
Ok(Arc::new(values) as ArrayRef)
610608
}
609+
//(Time32(TimeUnit::Second), Time64(_)) => {},
611610
(Time32(from_unit), Time64(to_unit)) => {
612611
let time_array = Int32Array::from(array.data());
613612
// note: (numeric_cast + SIMD multiply) is faster than (cast & multiply)
@@ -632,18 +631,24 @@ pub fn cast(array: &ArrayRef, to_type: &DataType) -> Result<ArrayRef> {
632631
}
633632
}
634633
(Time64(TimeUnit::Microsecond), Time64(TimeUnit::Nanosecond)) => {
635-
let time_array = Time64NanosecondArray::from(array.data());
636-
let mult = Time64NanosecondArray::from(vec![MILLISECONDS; array.len()]);
637-
let time64_ns = multiply(&time_array, &mult)?;
634+
let time_array = array
635+
.as_any()
636+
.downcast_ref::<Time64MicrosecondArray>()
637+
.unwrap();
638638

639-
Ok(Arc::new(time64_ns) as ArrayRef)
639+
let values =
640+
unary::<_, _, Time64NanosecondType>(time_array, |x| x * MILLISECONDS);
641+
Ok(Arc::new(values) as ArrayRef)
640642
}
641643
(Time64(TimeUnit::Nanosecond), Time64(TimeUnit::Microsecond)) => {
642-
let time_array = Time64MicrosecondArray::from(array.data());
643-
let divisor = Time64MicrosecondArray::from(vec![MILLISECONDS; array.len()]);
644-
let time64_us = divide(&time_array, &divisor)?;
644+
let time_array = array
645+
.as_any()
646+
.downcast_ref::<Time64NanosecondArray>()
647+
.unwrap();
645648

646-
Ok(Arc::new(time64_us) as ArrayRef)
649+
let values =
650+
unary::<_, _, Time64MicrosecondType>(time_array, |x| x / MILLISECONDS);
651+
Ok(Arc::new(values) as ArrayRef)
647652
}
648653
(Time64(from_unit), Time32(to_unit)) => {
649654
let time_array = Int64Array::from(array.data());
@@ -652,33 +657,16 @@ pub fn cast(array: &ArrayRef, to_type: &DataType) -> Result<ArrayRef> {
652657
let divisor = from_size / to_size;
653658
match to_unit {
654659
TimeUnit::Second => {
655-
let mut b = Time32SecondBuilder::new(array.len());
656-
for i in 0..array.len() {
657-
if array.is_null(i) {
658-
b.append_null()?;
659-
} else {
660-
b.append_value(
661-
(time_array.value(i) as i64 / divisor) as i32,
662-
)?;
663-
}
664-
}
665-
666-
Ok(Arc::new(b.finish()) as ArrayRef)
660+
let values = unary::<_, _, Time32SecondType>(&time_array, |x| {
661+
(x as i64 / divisor) as i32
662+
});
663+
Ok(Arc::new(values) as ArrayRef)
667664
}
668665
TimeUnit::Millisecond => {
669-
// currently can't dedup this builder [ARROW-4164]
670-
let mut b = Time32MillisecondBuilder::new(array.len());
671-
for i in 0..array.len() {
672-
if array.is_null(i) {
673-
b.append_null()?;
674-
} else {
675-
b.append_value(
676-
(time_array.value(i) as i64 / divisor) as i32,
677-
)?;
678-
}
679-
}
680-
681-
Ok(Arc::new(b.finish()) as ArrayRef)
666+
let values = unary::<_, _, Time32MillisecondType>(&time_array, |x| {
667+
(x as i64 / divisor) as i32
668+
});
669+
Ok(Arc::new(values) as ArrayRef)
682670
}
683671
_ => unreachable!("array type not supported"),
684672
}
@@ -806,7 +794,7 @@ pub fn cast(array: &ArrayRef, to_type: &DataType) -> Result<ArrayRef> {
806794
}
807795

808796
/// Get the time unit as a multiple of a second
809-
fn time_unit_multiple(unit: &TimeUnit) -> i64 {
797+
const fn time_unit_multiple(unit: &TimeUnit) -> i64 {
810798
match unit {
811799
TimeUnit::Second => 1,
812800
TimeUnit::Millisecond => MILLISECONDS,

rust/arrow/src/compute/kernels/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
2020
pub mod aggregate;
2121
pub mod arithmetic;
22+
pub mod arity;
2223
pub mod boolean;
2324
pub mod cast;
2425
pub mod comparison;

0 commit comments

Comments
 (0)