jhorstmann commented on a change in pull request #8092:
URL: https://github.com/apache/arrow/pull/8092#discussion_r483965152
##########
File path: rust/arrow/src/array/ord.rs
##########
@@ -68,65 +80,233 @@ impl OrdArray for NullArray {
}
}
+macro_rules! float_ord_cmp {
+ ($NAME: ident, $T: ty) => {
+ #[inline]
+ fn $NAME(a: $T, b: $T) -> Ordering {
+ if a < b {
+ return Ordering::Less;
+ }
+ if a > b {
+ return Ordering::Greater;
+ }
+
+ // convert to bits with canonical pattern for NaN
+ let a = if a.is_nan() {
+ <$T>::NAN.to_bits()
+ } else {
+ a.to_bits()
+ };
+ let b = if b.is_nan() {
+ <$T>::NAN.to_bits()
+ } else {
+ b.to_bits()
+ };
+
+ if a == b {
+ // Equal or both NaN
+ Ordering::Equal
+ } else if a < b {
+ // (-0.0, 0.0) or (!NaN, NaN)
+ Ordering::Less
+ } else {
+ // (0.0, -0.0) or (NaN, !NaN)
+ Ordering::Greater
+ }
+ }
+ };
+}
+
+float_ord_cmp!(cmp_f64, f64);
+float_ord_cmp!(cmp_f32, f32);
+
+#[repr(transparent)]
+struct Float64ArrayAsOrdArray<'a>(&'a Float64Array);
+#[repr(transparent)]
+struct Float32ArrayAsOrdArray<'a>(&'a Float32Array);
+
+impl OrdArray for Float64ArrayAsOrdArray<'_> {
+ fn cmp_value(&self, i: usize, j: usize) -> Ordering {
+ let a: f64 = self.0.value(i);
+ let b: f64 = self.0.value(j);
+
+ cmp_f64(a, b)
+ }
+}
+
+impl OrdArray for Float32ArrayAsOrdArray<'_> {
+ fn cmp_value(&self, i: usize, j: usize) -> Ordering {
+ let a: f32 = self.0.value(i);
+ let b: f32 = self.0.value(j);
+
+ cmp_f32(a, b)
+ }
+}
+
+fn float32_as_ord_array<'a>(array: &'a ArrayRef) -> Box<dyn OrdArray + 'a> {
+ let float_array: &Float32Array = as_primitive_array::<Float32Type>(array);
+ Box::new(Float32ArrayAsOrdArray(float_array))
+}
+
+fn float64_as_ord_array<'a>(array: &'a ArrayRef) -> Box<dyn OrdArray + 'a> {
+ let float_array: &Float64Array = as_primitive_array::<Float64Type>(array);
+ Box::new(Float64ArrayAsOrdArray(float_array))
+}
+
+struct StringDictionaryArrayAsOrdArray<'a, T: ArrowDictionaryKeyType> {
+ dict_array: &'a DictionaryArray<T>,
+ values: StringArray,
+ keys: PrimitiveArray<T>,
+}
+
+impl<T: ArrowDictionaryKeyType> OrdArray for
StringDictionaryArrayAsOrdArray<'_, T> {
+ fn cmp_value(&self, i: usize, j: usize) -> Ordering {
+ let keys = &self.keys;
+
+ let a: T::Native = keys.value(i);
+ let b: T::Native = keys.value(j);
+
+ let dict = &self.values;
+
+ let sa = dict.value(a.to_usize().unwrap());
+ let sb = dict.value(b.to_usize().unwrap());
+
+ sa.cmp(sb)
+ }
+}
+
+fn string_dict_as_ord_array<'a, T: ArrowDictionaryKeyType>(
+ array: &'a ArrayRef,
+) -> Box<dyn OrdArray + 'a>
+where
+ T::Native: std::cmp::Ord,
+{
+ let dict_array = as_dictionary_array::<T>(array);
+ let keys = dict_array.keys_array();
+
+ let values = &dict_array.values();
+ let values = StringArray::from(values.data());
+
+ Box::new(StringDictionaryArrayAsOrdArray {
+ dict_array,
+ values,
+ keys,
+ })
+}
+
/// Convert ArrayRef to OrdArray trait object
-pub fn as_ordarray(values: &ArrayRef) -> Result<&OrdArray> {
+pub fn as_ordarray<'a>(values: &'a ArrayRef) -> Result<Box<OrdArray + 'a>> {
match values.data_type() {
- DataType::Boolean => Ok(as_boolean_array(&values)),
- DataType::Utf8 => Ok(as_string_array(&values)),
- DataType::Null => Ok(as_null_array(&values)),
- DataType::Int8 => Ok(as_primitive_array::<Int8Type>(&values)),
- DataType::Int16 => Ok(as_primitive_array::<Int16Type>(&values)),
- DataType::Int32 => Ok(as_primitive_array::<Int32Type>(&values)),
- DataType::Int64 => Ok(as_primitive_array::<Int64Type>(&values)),
- DataType::UInt8 => Ok(as_primitive_array::<UInt8Type>(&values)),
- DataType::UInt16 => Ok(as_primitive_array::<UInt16Type>(&values)),
- DataType::UInt32 => Ok(as_primitive_array::<UInt32Type>(&values)),
- DataType::UInt64 => Ok(as_primitive_array::<UInt64Type>(&values)),
- DataType::Date32(_) => Ok(as_primitive_array::<Date32Type>(&values)),
- DataType::Date64(_) => Ok(as_primitive_array::<Date64Type>(&values)),
- DataType::Time32(Second) =>
Ok(as_primitive_array::<Time32SecondType>(&values)),
- DataType::Time32(Millisecond) => {
- Ok(as_primitive_array::<Time32MillisecondType>(&values))
- }
- DataType::Time64(Microsecond) => {
- Ok(as_primitive_array::<Time64MicrosecondType>(&values))
- }
- DataType::Time64(Nanosecond) => {
- Ok(as_primitive_array::<Time64NanosecondType>(&values))
+ //DataType::Boolean => Ok(Box::new(as_boolean_array(&values))),
Review comment:
Uh, thanks for noticing. At one point a had commented out the whole
block to get the signature and borrows right
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
