|
| 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 | +//! Float semantics for DataFusion. |
| 19 | +//! |
| 20 | +//! Arrow's comparison kernels work on the bit pattern of a float, |
| 21 | +//! i.e. IEEE 754 totalOrder. DataFusion deviates from that in one |
| 22 | +//! respect: `+0.0` and `-0.0` are equal, hash the same and fall into the same |
| 23 | +//! group. See [`normalize_float`], [`normalize_floats`] and |
| 24 | +//! [`normalize_float_scalar`]. Ordering still follows totalOrder, so |
| 25 | +//! [`compare_floats`] sorts `-0.0` before `+0.0`. |
| 26 | +//! |
| 27 | +//! Every comparison, hash or grouping boundary should be built from these |
| 28 | +//! helpers rather than inspecting float bits itself. Where an Arrow kernel |
| 29 | +//! compares, sorts, row-encodes or hashes the values, normalize them first. |
| 30 | +//! Where DataFusion compares values itself, use [`compare_floats`], |
| 31 | +//! [`float_lt`] or [`float_gt`]. Where it hashes them, hash [`hash_input`]. |
| 32 | +//! |
| 33 | +
|
| 34 | +use std::cmp::Ordering; |
| 35 | +use std::sync::Arc; |
| 36 | + |
| 37 | +use arrow::array::{Array, ArrayRef, AsArray, make_array}; |
| 38 | +use arrow::datatypes::{ |
| 39 | + ArrowPrimitiveType, DataType, Float16Type, Float32Type, Float64Type, |
| 40 | +}; |
| 41 | +use half::f16; |
| 42 | + |
| 43 | +use crate::ScalarValue; |
| 44 | + |
| 45 | +/// The floating point types the functions in this module are defined for: |
| 46 | +/// [`f16`], [`f32`] and [`f64`]. |
| 47 | +pub trait FloatSemantics: Copy { |
| 48 | + /// Whether [`normalize_float`] would change this value. This is a cheap |
| 49 | + /// check that lets callers skip rewriting data which is already |
| 50 | + /// normalized. |
| 51 | + fn needs_normalization(self) -> bool; |
| 52 | + |
| 53 | + /// See [`normalize_float`]. |
| 54 | + fn normalize(self) -> Self; |
| 55 | + |
| 56 | + /// See [`compare_floats`]. |
| 57 | + fn compare(self, other: Self) -> Ordering; |
| 58 | +} |
| 59 | + |
| 60 | +macro_rules! impl_float_semantics { |
| 61 | + ($(($t:ty, $bits:ty, $zero:expr)),+ $(,)?) => { |
| 62 | + $(impl FloatSemantics for $t { |
| 63 | + #[inline] |
| 64 | + fn needs_normalization(self) -> bool { |
| 65 | + // `-0.0` has only the sign bit set; no other value shares |
| 66 | + // this bit pattern. |
| 67 | + self.to_bits() == 1 << (<$bits>::BITS - 1) |
| 68 | + } |
| 69 | + |
| 70 | + #[inline] |
| 71 | + fn normalize(self) -> Self { |
| 72 | + if self.needs_normalization() { $zero } else { self } |
| 73 | + } |
| 74 | + |
| 75 | + #[inline] |
| 76 | + fn compare(self, other: Self) -> Ordering { |
| 77 | + self.total_cmp(&other) |
| 78 | + } |
| 79 | + })+ |
| 80 | + }; |
| 81 | +} |
| 82 | +impl_float_semantics!((f16, u16, f16::ZERO), (f32, u32, 0.0), (f64, u64, 0.0)); |
| 83 | + |
| 84 | +/// `-0.0` becomes `0.0`, so two values that DataFusion treats as equal also |
| 85 | +/// have the same bits. All other values, including NaN payloads, are kept as |
| 86 | +/// they are. |
| 87 | +#[inline] |
| 88 | +pub fn normalize_float<T: FloatSemantics>(v: T) -> T { |
| 89 | + v.normalize() |
| 90 | +} |
| 91 | + |
| 92 | +/// Compares two values in the order DataFusion sorts them. |
| 93 | +/// |
| 94 | +/// This is IEEE 754 totalOrder (`total_cmp`), which matches Arrow's sort |
| 95 | +/// kernels and row encoding. Note that it orders `-0.0` before `+0.0`. |
| 96 | +/// The grouped `MIN`/`MAX` accumulators start from the extremes of this |
| 97 | +/// order, so they must change together with it. |
| 98 | +#[inline] |
| 99 | +pub fn compare_floats<T: FloatSemantics>(left: T, right: T) -> Ordering { |
| 100 | + left.compare(right) |
| 101 | +} |
| 102 | + |
| 103 | +/// Whether `left` sorts before `right`, the same as |
| 104 | +/// `compare_floats(left, right).is_lt()`. |
| 105 | +#[inline] |
| 106 | +pub fn float_lt<T: FloatSemantics>(left: T, right: T) -> bool { |
| 107 | + compare_floats(left, right).is_lt() |
| 108 | +} |
| 109 | + |
| 110 | +/// Whether `left` sorts after `right`, the same as |
| 111 | +/// `compare_floats(left, right).is_gt()`. |
| 112 | +#[inline] |
| 113 | +pub fn float_gt<T: FloatSemantics>(left: T, right: T) -> bool { |
| 114 | + compare_floats(left, right).is_gt() |
| 115 | +} |
| 116 | + |
| 117 | +/// The value that DataFusion hashes in place of a float. Values that compare |
| 118 | +/// as equal hash the same, so this is the value [`normalize_float`] returns. |
| 119 | +#[inline] |
| 120 | +pub fn hash_input<T: FloatSemantics>(v: T) -> T { |
| 121 | + normalize_float(v) |
| 122 | +} |
| 123 | + |
| 124 | +/// Applies [`normalize_float`] to every float in `array`, including floats |
| 125 | +/// nested inside lists, structs, maps, unions, dictionaries and run-end |
| 126 | +/// encoded arrays. Arrays without floats are returned unchanged. |
| 127 | +/// |
| 128 | +/// Callers use this before handing float data to Arrow kernels that compare |
| 129 | +/// bit patterns, so that `+0.0 = -0.0` holds. |
| 130 | +/// |
| 131 | +/// The common case - nothing to normalize - is allocation-free: a single |
| 132 | +/// read-only scan of the values buffer (auto-vectorizable to an OR-reduction) |
| 133 | +/// decides whether to fall through to the rewriting path. Only arrays that |
| 134 | +/// contain values to normalize pay for a new buffer. |
| 135 | +pub fn normalize_floats(array: &ArrayRef) -> ArrayRef { |
| 136 | + match array.data_type() { |
| 137 | + DataType::Float16 => normalize_primitive::<Float16Type>(array), |
| 138 | + DataType::Float32 => normalize_primitive::<Float32Type>(array), |
| 139 | + DataType::Float64 => normalize_primitive::<Float64Type>(array), |
| 140 | + dt if has_float_leaf(dt) => { |
| 141 | + let data = array.to_data(); |
| 142 | + let children = data |
| 143 | + .child_data() |
| 144 | + .iter() |
| 145 | + .map(|child| normalize_floats(&make_array(child.clone())).to_data()) |
| 146 | + .collect::<Vec<_>>(); |
| 147 | + if children |
| 148 | + .iter() |
| 149 | + .zip(data.child_data()) |
| 150 | + .all(|(new, old)| new.ptr_eq(old)) |
| 151 | + { |
| 152 | + return Arc::clone(array); |
| 153 | + } |
| 154 | + make_array( |
| 155 | + data.into_builder() |
| 156 | + .child_data(children) |
| 157 | + .build() |
| 158 | + .expect("rewriting float leaves preserves the array layout"), |
| 159 | + ) |
| 160 | + } |
| 161 | + _ => Arc::clone(array), |
| 162 | + } |
| 163 | +} |
| 164 | + |
| 165 | +fn normalize_primitive<T>(array: &ArrayRef) -> ArrayRef |
| 166 | +where |
| 167 | + T: ArrowPrimitiveType, |
| 168 | + T::Native: FloatSemantics, |
| 169 | +{ |
| 170 | + let arr = array.as_primitive::<T>(); |
| 171 | + if !arr.values().iter().any(|v| v.needs_normalization()) { |
| 172 | + return Arc::clone(array); |
| 173 | + } |
| 174 | + Arc::new(arr.unary::<_, T>(normalize_float)) |
| 175 | +} |
| 176 | + |
| 177 | +/// Applies [`normalize_float`] to a `Float16`, `Float32` or `Float64` scalar. |
| 178 | +/// Other scalars are returned as is. |
| 179 | +pub fn normalize_float_scalar(scalar: ScalarValue) -> ScalarValue { |
| 180 | + match scalar { |
| 181 | + ScalarValue::Float16(Some(v)) => ScalarValue::Float16(Some(normalize_float(v))), |
| 182 | + ScalarValue::Float32(Some(v)) => ScalarValue::Float32(Some(normalize_float(v))), |
| 183 | + ScalarValue::Float64(Some(v)) => ScalarValue::Float64(Some(normalize_float(v))), |
| 184 | + other => other, |
| 185 | + } |
| 186 | +} |
| 187 | + |
| 188 | +/// Whether `data_type` is a float or contains a float at any depth of its |
| 189 | +/// nested children. |
| 190 | +pub fn has_float_leaf(data_type: &DataType) -> bool { |
| 191 | + match data_type { |
| 192 | + DataType::Float16 | DataType::Float32 | DataType::Float64 => true, |
| 193 | + DataType::List(f) |
| 194 | + | DataType::LargeList(f) |
| 195 | + | DataType::ListView(f) |
| 196 | + | DataType::LargeListView(f) |
| 197 | + | DataType::FixedSizeList(f, _) |
| 198 | + | DataType::Map(f, _) |
| 199 | + | DataType::RunEndEncoded(_, f) => has_float_leaf(f.data_type()), |
| 200 | + DataType::Struct(fields) => fields.iter().any(|f| has_float_leaf(f.data_type())), |
| 201 | + DataType::Union(fields, _) => { |
| 202 | + fields.iter().any(|(_, f)| has_float_leaf(f.data_type())) |
| 203 | + } |
| 204 | + DataType::Dictionary(_, values) => has_float_leaf(values), |
| 205 | + _ => false, |
| 206 | + } |
| 207 | +} |
| 208 | + |
| 209 | +#[cfg(test)] |
| 210 | +mod tests { |
| 211 | + use super::*; |
| 212 | + use arrow::array::{Float16Array, Float32Array, Float64Array, Int32Array, ListArray}; |
| 213 | + use arrow::datatypes::Int32Type; |
| 214 | + |
| 215 | + #[test] |
| 216 | + fn normalize_float_folds_negative_zero() { |
| 217 | + assert_eq!(normalize_float(-0.0_f64).to_bits(), 0.0_f64.to_bits()); |
| 218 | + assert_eq!(normalize_float(-0.0_f32).to_bits(), 0.0_f32.to_bits()); |
| 219 | + assert_eq!( |
| 220 | + normalize_float(f16::NEG_ZERO).to_bits(), |
| 221 | + f16::ZERO.to_bits() |
| 222 | + ); |
| 223 | + assert_eq!( |
| 224 | + hash_input(-0.0_f64).to_bits(), |
| 225 | + hash_input(0.0_f64).to_bits() |
| 226 | + ); |
| 227 | + assert!((-0.0_f64).needs_normalization()); |
| 228 | + assert!(!0.0_f64.needs_normalization()); |
| 229 | + } |
| 230 | + |
| 231 | + #[test] |
| 232 | + fn normalize_float_keeps_other_values() { |
| 233 | + let neg_nan = f64::from_bits(f64::NAN.to_bits() | (1 << 63)); |
| 234 | + let payload_nan = f64::from_bits(f64::NAN.to_bits() | 1); |
| 235 | + for v in [ |
| 236 | + 1.5, |
| 237 | + -1.5, |
| 238 | + f64::INFINITY, |
| 239 | + f64::NEG_INFINITY, |
| 240 | + f64::MIN_POSITIVE, |
| 241 | + -f64::MIN_POSITIVE, |
| 242 | + f64::NAN, |
| 243 | + neg_nan, |
| 244 | + payload_nan, |
| 245 | + ] { |
| 246 | + assert!(!v.needs_normalization()); |
| 247 | + assert_eq!(normalize_float(v).to_bits(), v.to_bits()); |
| 248 | + } |
| 249 | + } |
| 250 | + |
| 251 | + #[test] |
| 252 | + fn compare_floats_is_total_order() { |
| 253 | + let neg_nan = f64::from_bits(f64::NAN.to_bits() | (1 << 63)); |
| 254 | + assert_eq!(compare_floats(1.0_f64, 2.0), Ordering::Less); |
| 255 | + assert_eq!(compare_floats(-0.0_f64, 0.0), Ordering::Less); |
| 256 | + assert_eq!(compare_floats(f64::NAN, f64::INFINITY), Ordering::Greater); |
| 257 | + assert_eq!(compare_floats(neg_nan, f64::NEG_INFINITY), Ordering::Less); |
| 258 | + assert_eq!(compare_floats(f16::ONE, f16::ZERO), Ordering::Greater); |
| 259 | + assert!(float_lt(1.0_f64, 2.0) && !float_lt(2.0_f64, 1.0)); |
| 260 | + assert!(float_gt(f64::NAN, f64::INFINITY) && !float_gt(1.0_f64, 1.0)); |
| 261 | + } |
| 262 | + |
| 263 | + #[test] |
| 264 | + fn normalize_floats_rewrites_negative_zero() { |
| 265 | + let array: ArrayRef = |
| 266 | + Arc::new(Float64Array::from(vec![Some(-0.0), None, Some(1.0)])); |
| 267 | + let result = normalize_floats(&array); |
| 268 | + let result = result.as_primitive::<Float64Type>(); |
| 269 | + assert_eq!(result.value(0).to_bits(), 0.0_f64.to_bits()); |
| 270 | + assert!(result.is_null(1)); |
| 271 | + assert_eq!(result.value(2), 1.0); |
| 272 | + |
| 273 | + let array: ArrayRef = Arc::new(Float32Array::from(vec![-0.0_f32])); |
| 274 | + let result = normalize_floats(&array); |
| 275 | + assert_eq!( |
| 276 | + result.as_primitive::<Float32Type>().value(0).to_bits(), |
| 277 | + 0.0_f32.to_bits() |
| 278 | + ); |
| 279 | + |
| 280 | + let array: ArrayRef = Arc::new(Float16Array::from(vec![f16::NEG_ZERO])); |
| 281 | + let result = normalize_floats(&array); |
| 282 | + assert_eq!( |
| 283 | + result.as_primitive::<Float16Type>().value(0).to_bits(), |
| 284 | + f16::ZERO.to_bits() |
| 285 | + ); |
| 286 | + } |
| 287 | + |
| 288 | + #[test] |
| 289 | + fn normalize_floats_reuses_normalized_input() { |
| 290 | + let array: ArrayRef = Arc::new(Float64Array::from(vec![0.0, 1.0, f64::NAN])); |
| 291 | + assert!(Arc::ptr_eq(&normalize_floats(&array), &array)); |
| 292 | + |
| 293 | + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); |
| 294 | + assert!(Arc::ptr_eq(&normalize_floats(&array), &array)); |
| 295 | + } |
| 296 | + |
| 297 | + #[test] |
| 298 | + fn normalize_floats_rewrites_nested_floats() { |
| 299 | + let list = ListArray::from_iter_primitive::<Float64Type, _, _>(vec![ |
| 300 | + Some(vec![Some(-0.0), Some(2.0)]), |
| 301 | + None, |
| 302 | + ]); |
| 303 | + let array: ArrayRef = Arc::new(list); |
| 304 | + let result = normalize_floats(&array); |
| 305 | + assert!(!Arc::ptr_eq(&result, &array)); |
| 306 | + let values = result.as_list::<i32>().value(0); |
| 307 | + assert_eq!( |
| 308 | + values.as_primitive::<Float64Type>().value(0).to_bits(), |
| 309 | + 0.0_f64.to_bits() |
| 310 | + ); |
| 311 | + assert!(result.is_null(1)); |
| 312 | + |
| 313 | + let list = |
| 314 | + ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![Some(1)])]); |
| 315 | + let array: ArrayRef = Arc::new(list); |
| 316 | + assert!(Arc::ptr_eq(&normalize_floats(&array), &array)); |
| 317 | + } |
| 318 | + |
| 319 | + #[test] |
| 320 | + fn normalize_float_scalar_folds_negative_zero() { |
| 321 | + assert_eq!( |
| 322 | + normalize_float_scalar(ScalarValue::Float64(Some(-0.0))), |
| 323 | + ScalarValue::Float64(Some(0.0)) |
| 324 | + ); |
| 325 | + assert_eq!( |
| 326 | + normalize_float_scalar(ScalarValue::Float32(Some(-0.0))), |
| 327 | + ScalarValue::Float32(Some(0.0)) |
| 328 | + ); |
| 329 | + assert_eq!( |
| 330 | + normalize_float_scalar(ScalarValue::Float16(Some(f16::NEG_ZERO))), |
| 331 | + ScalarValue::Float16(Some(f16::ZERO)) |
| 332 | + ); |
| 333 | + assert_eq!( |
| 334 | + normalize_float_scalar(ScalarValue::Float64(None)), |
| 335 | + ScalarValue::Float64(None) |
| 336 | + ); |
| 337 | + assert_eq!( |
| 338 | + normalize_float_scalar(ScalarValue::Int32(Some(0))), |
| 339 | + ScalarValue::Int32(Some(0)) |
| 340 | + ); |
| 341 | + } |
| 342 | +} |
0 commit comments