avantgardnerio commented on code in PR #2294: URL: https://github.com/apache/datafusion-ballista/pull/2294#discussion_r3786414787
########## ballista/core/src/sort_key.rs: ########## @@ -0,0 +1,1465 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Sketching a single fixed-width `ORDER BY` key. +//! +//! [`crate::sort_key::SortKeyCodec`] is the ordering spec for one key — +//! its type, its direction, where its NULLs go — and encodes values to an +//! order-preserving `u64` and back. [`crate::sort_key::SortKeySketch`] +//! pairs that with a [`crate::kll::KllSketch`] over the encoded values and +//! a count of the NULLs, and answers quantiles over the whole population. +//! +//! Consumers want the sketch, not the codec. It is the type that knows how +//! to merge two observations, how a NULL run shifts a quantile, and what +//! goes on the wire. +//! +//! # Why an integer key +//! +//! The sketch needs `T: Ord`, and the obvious candidates for an `ORDER BY` +//! column are a type-specific wrapper (`OrderedFloat<f64>` and friends) or +//! the arrow row format. Both were measured against this encoding in +//! `benchmarks/benches/quantile_sketch.rs`; at n=1M, ratios to the +//! incumbent T-Digest are 1.23× for this encoding, 1.80× for +//! `OrderedFloat<f64>`, and 3.75× for arrow-row bytes held inline. +//! +//! Ingest cost is dominated by the `sort_unstable` inside KLL's compaction, +//! so the comparator is what matters: a `u64` compare is one instruction, +//! where a float total order is bit manipulation plus branches and +//! arrow-row pays ~25 ns/row to encode in the first place. Collapsing every +//! fixed-width type to a plain integer therefore wins on speed as well as +//! on uniformity. +//! +//! It also keeps sort direction out of the type system. `DESC` is a +//! bitwise NOT of the key rather than a second `Ord` implementation, so one +//! sketch type serves both directions instead of one per combination. +//! +//! # NULLs are out of band +//! +//! Encoding skips NULLs entirely and `SortKeySketch` counts them instead. +//! A NULL has no position among the values, only a side, and `nulls_first` +//! / `nulls_last` says which — that is one bit of plan-time information, +//! not something the key needs to carry. Keeping it out leaves the key 8 +//! bytes wide rather than 16, which the same benchmark measured at 1.23× +//! versus 1.58×. +//! +//! The cost is that a rank over the population is no longer a rank over +//! the values: the NULL run has to be stepped over first. That remap lives +//! in [`crate::sort_key::SortKeySketch::quantile`] and nowhere else. +//! Spread across call sites it would be reimplemented per consumer, and +//! getting it wrong skews every cut without failing anything. +//! +//! # The row format is the rulebook +//! +//! Arrow's row format is the only complete statement of what a SQL +//! `ORDER BY` means: it folds the column type, `nulls_first`, and +//! `descending` into a single memcmp order, and arrow's own sort agrees +//! with it. So it defines the answer, and anything faster is only allowed +//! to be an implementation of that answer. +//! +//! This encoding is exactly that. Its float transform is the same one +//! `arrow_row::fixed` applies, and both reduce to `total_cmp`, which is +//! what `ArrowNativeTypeOp::compare` uses. The test +//! `integer_keys_order_identically_to_arrow_row` pins the agreement on a +//! fixture containing ±NaN, ±0.0 and both infinities, so a divergence fails +//! a test rather than surfacing as misrouted rows. +//! +//! Following the rulebook is also what makes NaN a non-event. NaN has a +//! defined place in `total_cmp` — beyond the infinity of its own sign — so +//! it becomes an ordinary key, at the top or bottom of the `u64` range. +//! Comparisons against it behave, and this module contains no NaN handling +//! whatsoever. Code that compares raw `f64` instead has to special-case it, +//! because `partial_cmp` answers "no" to every question a router asks. +//! +//! # Exactness +//! +//! Every encoding here is a bijection on its type's value range, so a +//! quantile drawn from the sketch converts back to the precise value it +//! came from — not an approximation of it. That is what lets a +//! `Timestamp(Nanosecond)` cut stay nanosecond-exact; casting through +//! `f64` would round it to a 256 ns grid at 2020s epoch magnitudes, since +//! those sit above `f64`'s 2^53 integer limit. +//! +//! Where that exactness is worth something is narrower than it looks, and +//! worth stating so nobody over-claims it. It is not the quantiles: those +//! carry the sketch's own rank error, which on a uniform 1M stream is +//! ~0.2%, and 0.2% of a partition covering one day is about three minutes. +//! A 256 ns rounding is nine orders of magnitude beneath that. Any +//! argument resting on quantile precision is noise. +//! +//! It is the extremes. `min` and `max` are exact by construction, tracked +//! outside the compactor so no coin flip can move them, and `cut_partitions` +//! routes shuffle files on exactly those two values. There the error bars +//! are zero, so anything a cast rounds away is error introduced where none +//! existed. Keys compare with `Ord` over every element, which has no value +//! it silently ignores. +//! +//! The other case is a narrow spread at a large magnitude, since float +//! precision is relative: a partition spanning a day is unaffected, one +//! spanning 100 µs at 2020s epoch nanos is past the point where the cast +//! costs more than the sketch does. +//! +//! # Coverage +//! +//! Signed and unsigned integers, `Float32`/`Float64`, and the temporal +//! types that are `i32` or `i64` underneath (`Date`, `Time`, `Timestamp`, +//! `Duration`). [`crate::sort_key::SortKeyCodec::try_new`] returns `None` +//! for anything else +//! — `Decimal128` and wider don't fit in `u64`, `Interval` has no total +//! order, and variable-width types have no fixed encoding — leaving those +//! to the arrow-row path. + +use datafusion::arrow::array::{Array, ArrowPrimitiveType, AsArray, PrimitiveArray}; +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::{ + DataType, Date32Type, Date64Type, DurationMicrosecondType, DurationMillisecondType, + DurationNanosecondType, DurationSecondType, Float32Type, Float64Type, Int8Type, + Int16Type, Int32Type, Int64Type, Time32MillisecondType, Time32SecondType, + Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, + UInt16Type, UInt32Type, UInt64Type, +}; +use datafusion::common::{Result, ScalarValue, internal_datafusion_err}; + +use crate::kll::KllSketch; + +/// Bijection between a primitive's native value and a `u64` whose ascending +/// order matches the native ascending order. +trait SortableNative: Copy { + /// Map to the ascending `u64` key space. + fn to_key(self) -> u64; + /// Inverse of [`Self::to_key`], exact for any key that method produced. + fn from_key(key: u64) -> Self; +} + +/// Signed integers: flipping the sign bit maps the two's-complement order +/// onto unsigned order, because it slides the negative half below the +/// positive half. Narrower widths sign-extend to `i64` first, which +/// preserves order within their range. +macro_rules! impl_sortable_signed { + ($native:ty) => { + impl SortableNative for $native { + fn to_key(self) -> u64 { + (self as i64 as u64) ^ (1 << 63) + } + fn from_key(key: u64) -> Self { + (key ^ (1 << 63)) as i64 as Self + } + } + }; +} + +/// Unsigned integers are already in key order; widening preserves it. +macro_rules! impl_sortable_unsigned { + ($native:ty) => { + impl SortableNative for $native { + fn to_key(self) -> u64 { + self as u64 + } + fn from_key(key: u64) -> Self { + key as Self + } + } + }; +} + +/// IEEE-754 floats. +/// +/// This is a permutation, not a packing: 64 bits in, 64 bits out, nothing +/// compressed and nothing lost, which is why it inverts exactly. +/// +/// The layout was designed to almost sort as an integer already. +/// +/// ```text +/// 63 62 52 51 0 +/// ┌────┬─────────────────┬──────────────────────────────────────────┐ +/// │ S │ exponent │ mantissa │ +/// │ 1 │ 11 │ 52 │ +/// └────┴─────────────────┴──────────────────────────────────────────┘ +/// ^ ^ ^ +/// │ │ └── low bits +/// │ └── high bits, right below the sign +/// └── 0 = positive, 1 = negative +/// ``` +/// +/// The exponent sits *above* the mantissa deliberately. Compare two +/// same-sign floats as plain integers and the exponent dominates while the +/// mantissa breaks ties, which is exactly magnitude order. The bits already +/// sort themselves. Two things are wrong with them: +/// +/// ```text +/// as raw unsigned integers: +/// +/// 0x0000... +0.0 ─┐ +/// 0x3FF0... +1.0 │ positives: right order, stuck at the BOTTOM +/// 0x7FF0... +inf ─┘ +/// 0x8000... -0.0 ─┐ +/// 0xBFF0... -1.0 │ negatives: at the TOP, and running BACKWARDS +/// 0xFFF0... -inf ─┘ +/// +/// problem 1: a set sign bit makes negatives look huge +/// problem 2: within negatives, bigger magnitude = bigger integer +/// ``` +/// +/// One branch on the sign bit fixes both: +/// +/// ```text +/// sign bit 0 (non-negative): key = bits ^ 0x8000000000000000 +/// └─ flip only the sign bit, moving +/// them to the TOP half; the order +/// among them is untouched +/// +/// sign bit 1 (negative): key = !bits +/// └─ flip every bit: sign 1→0 moves +/// them to the BOTTOM half, and +/// inverting the rest reverses their +/// order, which is problem 2's fix +/// ``` +/// +/// What comes out the other end: +/// +/// ```text +/// value f64 bits key (u64) order +/// ─────────────────────────────────────────────────────────── +/// -NaN 0xFFF8000000000000 0x0007FFFFFFFFFFFF ▲ smallest +/// -inf 0xFFF0000000000000 0x000FFFFFFFFFFFFF │ +/// -2.0 0xC000000000000000 0x3FFFFFFFFFFFFFFF │ +/// -1.0 0xBFF0000000000000 0x400FFFFFFFFFFFFF │ +/// -0.0 0x8000000000000000 0x7FFFFFFFFFFFFFFF │ +/// +0.0 0x0000000000000000 0x8000000000000000 │ +/// +1.0 0x3FF0000000000000 0xBFF0000000000000 │ +/// +2.0 0x4000000000000000 0xC000000000000000 │ +/// +inf 0x7FF0000000000000 0xFFF0000000000000 │ +/// +NaN 0x7FF8000000000000 0xFFF8000000000000 ▼ largest +/// ``` +/// +/// That is `f64::total_cmp` order, which is what arrow sorts by. +/// +/// NaN needed no work. Its exponent is all ones with a nonzero mantissa, +/// so its pattern sits just above the infinity on its own side, which has +/// the same exponent and a zero mantissa. It lands past infinity by +/// itself. Nothing here tests for it: NaN is only awkward when compared +/// *as a float*. +/// +/// Note that all 2^64 keys are spoken for, so there is no spare slot to +/// mean NULL. That would need a 65th bit, and in practice a 16-byte key — +/// which is why NULLs are counted out of band instead. See the module +/// docs. +macro_rules! impl_sortable_float { + ($native:ty, $bits:ty, $width:expr) => { + impl SortableNative for $native { + fn to_key(self) -> u64 { + let bits = self.to_bits(); + let sign: $bits = 1 << ($width - 1); + let key: $bits = if bits & sign != 0 { !bits } else { bits ^ sign }; + // Zero-extending a narrower key preserves order, since + // every key of that width is below the widened range. + key as u64 + } + fn from_key(key: u64) -> Self { + let bits = key as $bits; + let sign = 1 << ($width - 1); + // Forward maps negatives to a cleared top bit and + // non-negatives to a set one, so the top bit selects the + // branch to undo. + let bits = if bits & sign != 0 { bits ^ sign } else { !bits }; + Self::from_bits(bits) + } + } + }; +} + +impl_sortable_signed!(i8); +impl_sortable_signed!(i16); +impl_sortable_signed!(i32); +impl_sortable_signed!(i64); +impl_sortable_unsigned!(u8); +impl_sortable_unsigned!(u16); +impl_sortable_unsigned!(u32); +impl_sortable_unsigned!(u64); +impl_sortable_float!(f32, u32, 32); +impl_sortable_float!(f64, u64, 64); + +/// Invoke `$handler!(ArrowPrimitiveType)` for the arrow type backing +/// `$data_type`, or evaluate `$fallback` when it isn't one this module +/// encodes. +/// +/// This allowlist *is* the tier boundary: every type named here gets the +/// `u64` fast path, and everything omitted falls through to arrow-row. +/// Adding a type means adding it here and nowhere else. +macro_rules! dispatch_sortable { + ($data_type:expr, $handler:ident, $fallback:expr) => { + match $data_type { + DataType::Int8 => $handler!(Int8Type), + DataType::Int16 => $handler!(Int16Type), + DataType::Int32 => $handler!(Int32Type), + DataType::Int64 => $handler!(Int64Type), + DataType::UInt8 => $handler!(UInt8Type), + DataType::UInt16 => $handler!(UInt16Type), + DataType::UInt32 => $handler!(UInt32Type), + DataType::UInt64 => $handler!(UInt64Type), + DataType::Float32 => $handler!(Float32Type), + DataType::Float64 => $handler!(Float64Type), + DataType::Date32 => $handler!(Date32Type), + DataType::Date64 => $handler!(Date64Type), + DataType::Time32(TimeUnit::Second) => $handler!(Time32SecondType), + DataType::Time32(TimeUnit::Millisecond) => $handler!(Time32MillisecondType), + DataType::Time64(TimeUnit::Microsecond) => $handler!(Time64MicrosecondType), + DataType::Time64(TimeUnit::Nanosecond) => $handler!(Time64NanosecondType), + DataType::Timestamp(TimeUnit::Second, _) => $handler!(TimestampSecondType), + DataType::Timestamp(TimeUnit::Millisecond, _) => { + $handler!(TimestampMillisecondType) + } + DataType::Timestamp(TimeUnit::Microsecond, _) => { + $handler!(TimestampMicrosecondType) + } + DataType::Timestamp(TimeUnit::Nanosecond, _) => { + $handler!(TimestampNanosecondType) + } + DataType::Duration(TimeUnit::Second) => $handler!(DurationSecondType), + DataType::Duration(TimeUnit::Millisecond) => { + $handler!(DurationMillisecondType) + } + DataType::Duration(TimeUnit::Microsecond) => { + $handler!(DurationMicrosecondType) + } + DataType::Duration(TimeUnit::Nanosecond) => $handler!(DurationNanosecondType), + _ => $fallback, + } + }; +} + +/// The complete ordering spec for one fixed-width `ORDER BY` key: its +/// type, its direction, and where its NULLs go. Encodes values to `u64` +/// and back. See the module docs for the encoding and for why NULLs are +/// handled out of band. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SortKeyCodec { + /// The column's full arrow type, retained so [`Self::decode`] can + /// rebuild a `ScalarValue` that keeps the parts the key doesn't carry + /// — a `Timestamp`'s timezone above all. + data_type: DataType, + /// `descending` inverts every key bit, reversing the sketch's ascending + /// order into the order the plan asked for. `nulls_first` never touches + /// a key, since NULLs are not encoded; it tells a sketch which end of + /// the distribution its NULL count occupies. + options: SortOptions, +} + +impl SortKeyCodec { + /// Build a codec for `data_type` under `options`, or `None` if this + /// module doesn't encode that type and the caller should fall back to + /// arrow-row. + pub fn try_new(data_type: &DataType, options: SortOptions) -> Option<Self> { + macro_rules! supported { + ($arrow_type:ty) => { + true + }; + } + let supported = dispatch_sortable!(data_type, supported, false); + supported.then(|| Self { + data_type: data_type.clone(), + options, + }) + } + + /// The arrow type this codec was built for. + pub fn data_type(&self) -> &DataType { + &self.data_type + } + + /// The sort direction and NULL placement this codec encodes for. + pub fn options(&self) -> SortOptions { + self.options + } + + /// A typed NULL of this codec's column type. What a quantile query + /// answers when the rank it asks for lands in the NULL run. + pub fn null_value(&self) -> Result<ScalarValue> { + ScalarValue::try_from(&self.data_type) + } + + /// Encode `array`'s non-NULL values in row order. + /// + /// NULLs are skipped, so the result is shorter than `array` by exactly + /// `array.null_count()` — callers that need that count read it from the + /// array. The output is ready for `KllSketch::absorb_slice`. + /// + /// Errors if `array`'s type doesn't match the one this codec was built + /// for, which would mean the routing expression changed type between + /// planning and execution. + pub fn encode(&self, array: &dyn Array) -> Result<Vec<u64>> { + if array.data_type() != &self.data_type { + return Err(internal_datafusion_err!( + "SortKeyCodec: built for {:?} but got {:?}", + self.data_type, + array.data_type() + )); + } + macro_rules! encode_as { + ($arrow_type:ty) => {{ + let typed = array.as_primitive_opt::<$arrow_type>().ok_or_else(|| { + internal_datafusion_err!( + "SortKeyCodec: {:?} array failed to downcast to its own \ + primitive type", + self.data_type + ) + })?; + Ok(self.encode_primitive(typed)) + }}; + } + dispatch_sortable!( + &self.data_type, + encode_as, + Err(internal_datafusion_err!( + "SortKeyCodec: {:?} is not encodable — try_new should have \ + returned None", + self.data_type + )) + ) + } + + /// Shared body of every [`Self::encode`] arm, monomorphized per arrow + /// type. Split out so the all-non-NULL case can read the values buffer + /// directly instead of going through the nullable iterator. + fn encode_primitive<T>(&self, array: &PrimitiveArray<T>) -> Vec<u64> + where + T: ArrowPrimitiveType, + T::Native: SortableNative, + { + let descending = self.options.descending; + let orient = move |key: u64| if descending { !key } else { key }; + if array.null_count() == 0 { + array.values().iter().map(|v| orient(v.to_key())).collect() + } else { + array.iter().flatten().map(|v| orient(v.to_key())).collect() + } + } + + /// Recover the value a key came from, as a `ScalarValue` carrying this + /// codec's full arrow type. + /// + /// Exact for any key [`Self::encode`] produced. A key from anywhere else + /// still decodes — the map is total — but to an arbitrary value of the + /// type. + pub fn decode(&self, key: u64) -> Result<ScalarValue> { + // Bitwise NOT is an involution, so the same branch undoes DESC. + let key = if self.options.descending { !key } else { key }; + macro_rules! decode_as { + ($arrow_type:ty) => {{ + let native = + <<$arrow_type as ArrowPrimitiveType>::Native as SortableNative>::from_key(key); + ScalarValue::new_primitive::<$arrow_type>(Some(native), &self.data_type) + }}; + } + dispatch_sortable!( + &self.data_type, + decode_as, + Err(internal_datafusion_err!( + "SortKeyCodec: {:?} is not decodable — try_new should have \ + returned None", + self.data_type + )) + ) + } +} + +/// KLL top-level compactor capacity. Picked for worst-case rank-error +/// parity with the T-Digest sizing it replaces (`max_size = 100`), so the +/// swap changes the sketch's cost and exactness without changing its +/// accuracy: on a uniform 1M stream, 0.0016 worst-case normalized rank +/// error against T-Digest's 0.0021. Rerun with `KLL_PARITY_CHECK=1 cargo +/// bench --bench quantile_sketch`. +const KLL_K: usize = 800; + +/// One `ORDER BY` key's observed distribution: a quantile sketch over the +/// non-NULL values, plus the count of the NULLs that have no place in it. +/// +/// Holding both together is the point. NULLs sit at one end of the order +/// rather than among the values, so any quantile over the *population* +/// has to account for the NULL run before consulting the sketch. Doing +/// that remap at call sites would mean every consumer reimplementing it, +/// and getting it wrong skews cuts silently rather than failing. So merge, +/// quantile, and the wire format all live here, once. +#[derive(Debug, Clone)] +pub struct SortKeySketch { + /// How values become keys, and which end the NULLs occupy. + codec: SortKeyCodec, + /// Quantile structure over the non-NULL values only. + sketch: KllSketch<u64>, + /// Rows whose key was NULL. Not in `sketch`, and not recoverable from + /// it. + null_count: u64, +} + +impl SortKeySketch { + /// An empty sketch for the key `codec` describes. + pub fn new(codec: SortKeyCodec) -> Self { + Self { + codec, + sketch: KllSketch::new(KLL_K), + null_count: 0, + } + } + + /// Observe every row of `array`: encode the non-NULL values into the + /// sketch and add the NULLs to the count. + /// + /// Errors if `array`'s type disagrees with the codec's, which would + /// mean the routing expression changed type between planning and + /// execution. + pub fn ingest(&mut self, array: &dyn Array) -> Result<()> { + let keys = self.codec.encode(array)?; + self.sketch.absorb_slice(&keys); + self.null_count += array.null_count() as u64; + Ok(()) + } + + /// Fold `other` into `self`. + /// + /// Errors when the two describe different keys. Merging a sketch of + /// one column into a sketch of another produces a plausible-looking + /// distribution of nothing in particular, so it is caught rather than + /// tolerated. + pub fn merge(&mut self, other: Self) -> Result<()> { + if self.codec != other.codec { + return Err(internal_datafusion_err!( + "SortKeySketch::merge: {:?} and {:?} describe different sort keys", + self.codec, + other.codec + )); + } + self.sketch.merge(other.sketch); + self.null_count += other.null_count; + Ok(()) + } + + /// Rows observed, NULLs included. + pub fn count(&self) -> u64 { + self.sketch.count() + self.null_count + } + + /// Rows observed whose key was NULL. + pub fn null_count(&self) -> u64 { + self.null_count + } + + /// The ordering spec these observations were made under. + pub fn codec(&self) -> &SortKeyCodec { + &self.codec + } + + /// The least value in sort order, or `None` when nothing was observed. + /// + /// A typed NULL when NULLs sort first and at least one was seen, since + /// then the least *row* is a NULL rather than a value. + pub fn min(&self) -> Result<Option<ScalarValue>> { + self.extreme(self.codec.options().nulls_first, self.sketch.min()) + } + + /// The greatest value in sort order, or `None` when nothing was + /// observed. A typed NULL when NULLs sort last and at least one was + /// seen. + pub fn max(&self) -> Result<Option<ScalarValue>> { + self.extreme(!self.codec.options().nulls_first, self.sketch.max()) + } + + /// Shared body of [`Self::min`] and [`Self::max`]: the extreme is a + /// NULL when the NULL run is on `nulls_are_on_this_end` and non-empty, + /// otherwise it is `value_extreme` decoded. + fn extreme( + &self, + nulls_are_on_this_end: bool, + value_extreme: Option<&u64>, + ) -> Result<Option<ScalarValue>> { + // With no values at all the run is unbounded on both sides, so the + // end it does not nominally occupy is a NULL too. Answering `None` + // there would claim nothing was observed while `count` says + // otherwise, and would hand `cut_partitions` half a range. + if self.null_count > 0 && (nulls_are_on_this_end || value_extreme.is_none()) { + return Ok(Some(self.codec.null_value()?)); + } Review Comment: > All-NULL sketch reports `min()`/`max()` as `None` Confirmed and fixed. Reproduced as a failing test first, `an_all_null_column_is_null_at_both_extremes`, which failed on exactly the case you measured: ``` assertion `left == right` failed: nulls_first=true left: None right: Some(Int64(NULL)) ``` > when `value_extreme` is `None` and `null_count > 0`, return the typed NULL regardless of which end was asked for That is the fix, with the reason inline: with no values the NULL run is unbounded on both sides, so the end it does not nominally occupy is a NULL too. `extremes_account_for_where_nulls_sort` still pins the partial case, where the value end keeps its value. -- 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. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
