From d24fb0eae7159a986a92c4e8be04e99321f4fc3f Mon Sep 17 00:00:00 2001 From: Marco Neumann Date: Wed, 7 Sep 2022 15:59:16 +0000 Subject: [PATCH] fix: support `u64` in `query_functions::selectors` (#5582) This this to be used for InfluxRPC and is currently panicking the prod querier from time to time. --- query_functions/src/selectors.rs | 66 +++++++++++++++++++++-- query_functions/src/selectors/internal.rs | 46 +++++++++++++++- 2 files changed, 108 insertions(+), 4 deletions(-) diff --git a/query_functions/src/selectors.rs b/query_functions/src/selectors.rs index 22d18eeee7..d32b9d43e9 100644 --- a/query_functions/src/selectors.rs +++ b/query_functions/src/selectors.rs @@ -29,8 +29,9 @@ mod internal; use internal::{ BooleanFirstSelector, BooleanLastSelector, BooleanMaxSelector, BooleanMinSelector, F64FirstSelector, F64LastSelector, F64MaxSelector, F64MinSelector, I64FirstSelector, - I64LastSelector, I64MaxSelector, I64MinSelector, Utf8FirstSelector, Utf8LastSelector, - Utf8MaxSelector, Utf8MinSelector, + I64LastSelector, I64MaxSelector, I64MinSelector, U64FirstSelector, U64LastSelector, + U64MaxSelector, U64MinSelector, Utf8FirstSelector, Utf8LastSelector, Utf8MaxSelector, + Utf8MinSelector, }; use schema::TIME_DATA_TYPE; @@ -58,6 +59,7 @@ pub fn selector_first(data_type: &DataType, output: SelectorOutput) -> Aggregate match data_type { DataType::Float64 => make_uda::(name, output), DataType::Int64 => make_uda::(name, output), + DataType::UInt64 => make_uda::(name, output), DataType::Utf8 => make_uda::(name, output), DataType::Boolean => make_uda::(name, output), _ => unimplemented!("first not supported for {:?}", data_type), @@ -88,6 +90,7 @@ pub fn selector_last(data_type: &DataType, output: SelectorOutput) -> AggregateU match data_type { DataType::Float64 => make_uda::(name, output), DataType::Int64 => make_uda::(name, output), + DataType::UInt64 => make_uda::(name, output), DataType::Utf8 => make_uda::(name, output), DataType::Boolean => make_uda::(name, output), _ => unimplemented!("last not supported for {:?}", data_type), @@ -118,6 +121,7 @@ pub fn selector_min(data_type: &DataType, output: SelectorOutput) -> AggregateUD match data_type { DataType::Float64 => make_uda::(name, output), DataType::Int64 => make_uda::(name, output), + DataType::UInt64 => make_uda::(name, output), DataType::Utf8 => make_uda::(name, output), DataType::Boolean => make_uda::(name, output), _ => unimplemented!("min not supported for {:?}", data_type), @@ -148,6 +152,7 @@ pub fn selector_max(data_type: &DataType, output: SelectorOutput) -> AggregateUD match data_type { DataType::Float64 => make_uda::(name, output), DataType::Int64 => make_uda::(name, output), + DataType::UInt64 => make_uda::(name, output), DataType::Utf8 => make_uda::(name, output), DataType::Boolean => make_uda::(name, output), _ => unimplemented!("max not supported for {:?}", data_type), @@ -300,7 +305,10 @@ where #[cfg(test)] mod test { use arrow::{ - array::{BooleanArray, Float64Array, Int64Array, StringArray, TimestampNanosecondArray}, + array::{ + BooleanArray, Float64Array, Int64Array, StringArray, TimestampNanosecondArray, + UInt64Array, + }, datatypes::{Field, Schema, SchemaRef}, record_batch::RecordBatch, util::pretty::pretty_format_batches, @@ -337,6 +345,18 @@ mod test { "+------------------------------------------+-----------------------------------------+", ], ), + ( + selector_first(&DataType::UInt64, SelectorOutput::Value), + selector_first(&DataType::UInt64, SelectorOutput::Time), + "u64_value", + vec![ + "+------------------------------------------+-----------------------------------------+", + "| selector_first_value(t.u64_value,t.time) | selector_first_time(t.u64_value,t.time) |", + "+------------------------------------------+-----------------------------------------+", + "| 20 | 1970-01-01 00:00:00.000001 |", + "+------------------------------------------+-----------------------------------------+", + ], + ), ( selector_first(&DataType::Utf8, SelectorOutput::Value), selector_first(&DataType::Utf8, SelectorOutput::Time), @@ -403,6 +423,18 @@ mod test { "+-----------------------------------------+----------------------------------------+", ], ), + ( + selector_last(&DataType::UInt64, SelectorOutput::Value), + selector_last(&DataType::UInt64, SelectorOutput::Time), + "u64_value", + vec![ + "+-----------------------------------------+----------------------------------------+", + "| selector_last_value(t.u64_value,t.time) | selector_last_time(t.u64_value,t.time) |", + "+-----------------------------------------+----------------------------------------+", + "| 30 | 1970-01-01 00:00:00.000006 |", + "+-----------------------------------------+----------------------------------------+", + ], + ), ( selector_last(&DataType::Utf8, SelectorOutput::Value), selector_last(&DataType::Utf8, SelectorOutput::Time), @@ -469,6 +501,18 @@ mod test { "+----------------------------------------+---------------------------------------+", ], ), + ( + selector_min(&DataType::UInt64, SelectorOutput::Value), + selector_min(&DataType::UInt64, SelectorOutput::Time), + "u64_value", + vec![ + "+----------------------------------------+---------------------------------------+", + "| selector_min_value(t.u64_value,t.time) | selector_min_time(t.u64_value,t.time) |", + "+----------------------------------------+---------------------------------------+", + "| 10 | 1970-01-01 00:00:00.000004 |", + "+----------------------------------------+---------------------------------------+", + ], + ), ( selector_min(&DataType::Utf8, SelectorOutput::Value), selector_min(&DataType::Utf8, SelectorOutput::Time), @@ -535,6 +579,18 @@ mod test { "+----------------------------------------+---------------------------------------+", ], ), + ( + selector_max(&DataType::UInt64, SelectorOutput::Value), + selector_max(&DataType::UInt64, SelectorOutput::Time), + "u64_value", + vec![ + "+----------------------------------------+---------------------------------------+", + "| selector_max_value(t.u64_value,t.time) | selector_max_time(t.u64_value,t.time) |", + "+----------------------------------------+---------------------------------------+", + "| 50 | 1970-01-01 00:00:00.000005 |", + "+----------------------------------------+---------------------------------------+", + ], + ), ( selector_max(&DataType::Utf8, SelectorOutput::Value), selector_max(&DataType::Utf8, SelectorOutput::Time), @@ -594,6 +650,7 @@ mod test { let schema = Arc::new(Schema::new(vec![ Field::new("f64_value", DataType::Float64, true), Field::new("i64_value", DataType::Int64, true), + Field::new("u64_value", DataType::UInt64, true), Field::new("string_value", DataType::Utf8, true), Field::new("bool_value", DataType::Boolean, true), Field::new("time", TIME_DATA_TYPE(), true), @@ -605,6 +662,7 @@ mod test { vec![ Arc::new(Float64Array::from(vec![Some(2.0), Some(4.0), None])), Arc::new(Int64Array::from(vec![Some(20), Some(40), None])), + Arc::new(UInt64Array::from(vec![Some(20), Some(40), None])), Arc::new(StringArray::from(vec![Some("two"), Some("four"), None])), Arc::new(BooleanArray::from(vec![Some(true), Some(false), None])), Arc::new(TimestampNanosecondArray::from_vec( @@ -621,6 +679,7 @@ mod test { vec![ Arc::new(Float64Array::from(vec![] as Vec>)), Arc::new(Int64Array::from(vec![] as Vec>)), + Arc::new(UInt64Array::from(vec![] as Vec>)), Arc::new(StringArray::from(vec![] as Vec>)), Arc::new(BooleanArray::from(vec![] as Vec>)), Arc::new(TimestampNanosecondArray::from_vec( @@ -638,6 +697,7 @@ mod test { vec![ Arc::new(Float64Array::from(vec![Some(1.0), Some(5.0), Some(3.0)])), Arc::new(Int64Array::from(vec![Some(10), Some(50), Some(30)])), + Arc::new(UInt64Array::from(vec![Some(10), Some(50), Some(30)])), Arc::new(StringArray::from(vec![ Some("a_one"), Some("z_five"), diff --git a/query_functions/src/selectors/internal.rs b/query_functions/src/selectors/internal.rs index 4d37f81a53..05a0e85709 100644 --- a/query_functions/src/selectors/internal.rs +++ b/query_functions/src/selectors/internal.rs @@ -11,7 +11,7 @@ use std::fmt::Debug; use arrow::{ array::{ Array, ArrayRef, BooleanArray, Float64Array, Int64Array, StringArray, - TimestampNanosecondArray, + TimestampNanosecondArray, UInt64Array, }, compute::kernels::aggregate::{ max as array_max, max_boolean as array_max_boolean, max_string as array_max_string, @@ -49,6 +49,12 @@ impl LtVal for i64 { } } +impl LtVal for u64 { + fn lt_val(&self, v: &Self) -> bool { + self < v + } +} + impl LtVal for bool { fn lt_val(&self, v: &Self) -> bool { self < v @@ -85,6 +91,12 @@ impl ToState for i64 { } } +impl ToState for u64 { + fn to_state(&self) -> Self { + *self + } +} + impl ToState for bool { fn to_state(&self) -> Self { *self @@ -578,6 +590,14 @@ make_first_selector!( array_min, ScalarValue::Int64 ); +make_first_selector!( + U64FirstSelector, + u64, + DataType::UInt64, + UInt64Array, + array_min, + ScalarValue::UInt64 +); make_first_selector!( Utf8FirstSelector, String, @@ -613,6 +633,14 @@ make_last_selector!( array_max, ScalarValue::Int64 ); +make_last_selector!( + U64LastSelector, + u64, + DataType::UInt64, + UInt64Array, + array_max, + ScalarValue::UInt64 +); make_last_selector!( Utf8LastSelector, String, @@ -648,6 +676,14 @@ make_min_selector!( array_min, ScalarValue::Int64 ); +make_min_selector!( + U64MinSelector, + u64, + DataType::UInt64, + UInt64Array, + array_min, + ScalarValue::UInt64 +); make_min_selector!( Utf8MinSelector, String, @@ -683,6 +719,14 @@ make_max_selector!( array_max, ScalarValue::Int64 ); +make_max_selector!( + U64MaxSelector, + u64, + DataType::UInt64, + UInt64Array, + array_max, + ScalarValue::UInt64 +); make_max_selector!( Utf8MaxSelector, String,