avantgardnerio commented on code in PR #2294: URL: https://github.com/apache/datafusion-ballista/pull/2294#discussion_r3786551052
########## ballista/core/src/sort_key.rs: ########## @@ -0,0 +1,1569 @@ +// 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()?)); + } + value_extreme.map(|key| self.codec.decode(*key)).transpose() + } + + /// The value at the `q`-quantile of everything observed, NULLs + /// included. `q` is clamped to `[0, 1]`. + /// + /// `Ok(None)` means nothing was observed at all. `Ok(Some(null))` means + /// the rank `q` asks for lands inside the NULL run, which is a real + /// answer: a cut there says the partition below it holds only NULLs. + /// + /// # The remap + /// + /// NULLs occupy a contiguous run at one end, so a rank over the + /// population maps onto a rank over the values by subtracting the run + /// when it sits below, and by nothing when it sits above: + /// + /// ```text + /// nulls_first nulls_last + /// ┌────────┬─────┐ ┌─────┬────────┐ + /// │ NULLs │ vals│ │ vals│ NULLs │ + /// └────────┴─────┘ └─────┴────────┘ + /// 0 n N 0 v N + /// + /// rank = (q · N) as integer + /// rank <= n -> NULL rank <= v -> values, at rank + /// else -> values, else -> NULL + /// at rank - n + /// ``` + /// + /// The subtraction stays in integers on purpose. Rescaling the rank + /// into a fraction of the value run and letting the sketch multiply it + /// back loses it: with 99 values, rank 59 becomes `59/99`, and `59/99 · + /// 99` is `58.999…`, which truncates to 58. That off-by-one was caught + /// by `quantiles_match_the_nulls_in_key_oracle`. + pub fn quantile(&self, q: f64) -> Result<Option<ScalarValue>> { + let total = self.count(); + if total == 0 { + return Ok(None); + } + // No NULL run to step over, so the population rank *is* the value + // rank and the sketch can answer directly. Also the only path that + // reaches the sketch's exact-extreme handling at q = 0 and q = 1. + if self.null_count == 0 { + return self + .sketch + .quantile(q) + .map(|key| self.codec.decode(*key)) + .transpose(); + } + let values = self.sketch.count(); + if values == 0 { + return Ok(Some(self.codec.null_value()?)); + } + + // Same rank arithmetic the sketch itself would do, so that + // stepping over the NULL run is the only difference between this + // answer and a sketch of the whole population. + let rank = (q.clamp(0.0, 1.0) * total as f64) as u64; + let value_rank = if self.codec.options().nulls_first { + if rank <= self.null_count { + return Ok(Some(self.codec.null_value()?)); + } + rank - self.null_count + } else { + if rank > values { + return Ok(Some(self.codec.null_value()?)); + } + rank + }; + self.sketch + .at_rank(value_rank) + .map(|key| self.codec.decode(*key)) + .transpose() + } + + /// The `partitions - 1` boundaries that split everything observed into + /// `partitions` equally-sized runs, in sort order. + /// + /// Empty when `partitions < 2` or nothing was observed. Entries may + /// repeat where one value dominates, and may be typed NULLs where the + /// NULL run spans a boundary; both are faithful answers about a skewed + /// distribution rather than errors. + pub fn cuts(&self, partitions: usize) -> Result<Vec<ScalarValue>> { + if partitions < 2 { + return Ok(Vec::new()); + } + (1..partitions) + .map(|cut| self.quantile(cut as f64 / partitions as f64)) + .collect::<Result<Vec<_>>>() + .map(|cuts| cuts.into_iter().flatten().collect()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{ + Float64Array, Int64Array, TimestampNanosecondArray, UInt64Array, + }; + use std::sync::Arc; + + /// Encoding must be monotone: sorting the keys must give the same + /// permutation as sorting the values. Checked over a set chosen to + /// straddle every boundary the encoding has to get right — sign + /// changes, zero, and the extremes of the type. + #[test] + fn ascending_keys_follow_value_order() { + let values = vec![ + i64::MIN, + i64::MIN + 1, + -1_000_000, + -1, + 0, + 1, + 1_000_000, + i64::MAX - 1, + i64::MAX, + ]; + let codec = + SortKeyCodec::try_new(&DataType::Int64, sort_options(false, false)).unwrap(); + let keys = codec.encode(&Int64Array::from(values.clone())).unwrap(); + assert!( + keys.windows(2).all(|w| w[0] < w[1]), + "keys must be strictly increasing for strictly increasing values: {keys:?}" + ); + } + + /// `DESC` must invert that order while staying a bijection. + #[test] + fn descending_keys_reverse_value_order() { + let values = vec![-5i64, 0, 5, 100]; + let codec = + SortKeyCodec::try_new(&DataType::Int64, sort_options(true, false)).unwrap(); + let keys = codec.encode(&Int64Array::from(values.clone())).unwrap(); + assert!( + keys.windows(2).all(|w| w[0] > w[1]), + "DESC keys must strictly decrease for increasing values: {keys:?}" + ); + for (key, value) in keys.iter().zip(&values) { + assert_eq!( + codec.decode(*key).unwrap(), + ScalarValue::Int64(Some(*value)) + ); + } + } + + /// Floats order by `total_cmp`, which puts `-0.0` below `+0.0` and + /// sorts NaN to the ends by sign. Infinities are ordinary values. + #[test] + fn float_keys_follow_total_order() { + let values = vec![ + f64::NEG_INFINITY, + -1.5, + -0.0, + 0.0, + f64::MIN_POSITIVE, + 1.5, + f64::INFINITY, + ]; + let codec = SortKeyCodec::try_new(&DataType::Float64, sort_options(false, false)) + .unwrap(); + let keys = codec.encode(&Float64Array::from(values.clone())).unwrap(); + assert!( + keys.windows(2).all(|w| w[0] < w[1]), + "float keys must respect total_cmp order: {keys:?}" + ); + // Round-trip preserves the sign of zero, which `==` on f64 would + // not distinguish. + assert_eq!( + codec.decode(keys[2]).unwrap(), + ScalarValue::Float64(Some(-0.0)) + ); + assert!( + matches!(codec.decode(keys[2]).unwrap(), ScalarValue::Float64(Some(v)) if v.is_sign_negative()) + ); + } + + /// NaN and the infinities have defined places in `total_cmp`: a NaN + /// sits beyond the infinity of its own sign, and the infinities are + /// otherwise ordinary values. Arrow sorts floats by exactly this + /// predicate (`ArrowNativeTypeOp::compare` delegates to `total_cmp`), + /// so the key has to agree — otherwise the sketch would order its + /// input differently from the `SortExec` that produced it. + /// + /// Unlike T-Digest, which interpolates centroid means and can turn + /// `Inf - Inf` into NaN, the sketch only ever compares keys, so an + /// infinity is exactly as safe here as any other value. + #[test] + fn nan_and_infinities_take_their_total_cmp_positions() { + let values = vec![ + -f64::NAN, + f64::NEG_INFINITY, + -1.0, + 0.0, + 1.0, + f64::INFINITY, + f64::NAN, + ]; + assert!( + values.windows(2).all(|w| w[0].total_cmp(&w[1]).is_lt()), + "test premise: the fixture is in total_cmp order" + ); + + let codec = SortKeyCodec::try_new(&DataType::Float64, sort_options(false, false)) + .unwrap(); + let keys = codec.encode(&Float64Array::from(values.clone())).unwrap(); + assert!( + keys.windows(2).all(|w| w[0] < w[1]), + "keys must reproduce total_cmp order: {keys:?}" + ); + + for (key, value) in keys.iter().zip(&values) { + match codec.decode(*key).unwrap() { + ScalarValue::Float64(Some(decoded)) => assert_eq!( + decoded.total_cmp(value), + std::cmp::Ordering::Equal, + "decode must round-trip {value} bit-exactly, including NaN sign" + ), + other => panic!("expected Float64, got {other:?}"), + } + } + } + + /// The two encoders have to agree. A column's routing must not depend + /// on whether it happened to qualify for the integer fast path or fell + /// through to arrow-row, so the orders they induce must be identical — + /// including on the float edge cases. + #[test] + fn integer_keys_order_identically_to_arrow_row() { + use datafusion::arrow::row::{RowConverter, SortField}; + + let values = vec![ + -f64::NAN, + f64::NEG_INFINITY, + -7.25, + -0.0, + 0.0, + f64::MIN_POSITIVE, + 7.25, + f64::INFINITY, + f64::NAN, + ]; + let array = Float64Array::from(values.clone()); + + let codec = SortKeyCodec::try_new(&DataType::Float64, sort_options(false, false)) + .unwrap(); + let keys = codec.encode(&array).unwrap(); + + let converter = + RowConverter::new(vec![SortField::new(DataType::Float64)]).unwrap(); + let rows = converter + .convert_columns(&[Arc::new(array) as Arc<dyn Array>]) + .unwrap(); + + let by_key = argsort(&keys); + let by_row: Vec<usize> = { + let mut order: Vec<usize> = (0..values.len()).collect(); + order.sort_by_key(|&i| rows.row(i).as_ref().to_vec()); + order + }; + assert_eq!( + by_key, by_row, + "integer keys and arrow-row bytes must induce the same permutation" + ); + } + + /// `SortOptions` without the struct-literal noise at every call site. + fn sort_options(descending: bool, nulls_first: bool) -> SortOptions { + SortOptions { + descending, + nulls_first, + } + } + + /// Ingest keys through a sketch, one `absorb_slice` per batch. + fn sketch_batches(codec: &SortKeyCodec, batches: Vec<Vec<f64>>) -> KllSketch<u64> { + let mut sketch = KllSketch::<u64>::new(64); + for batch in batches { + let keys = codec.encode(&Float64Array::from(batch)).unwrap(); + sketch.absorb_slice(&keys); + } + sketch + } + + /// The extremes a sketch reports are what decides which shuffle files + /// overlap which output partition, so losing one loses rows. They have + /// to survive ingest even when a batch *begins and ends* with NaN, + /// which is what a `total_cmp` sort does to any batch containing one. + /// + /// This is the case DataFusion's T-Digest gets wrong. It reads a + /// batch's extremes from `first()` and `last()` alone and folds them in + /// with `f64::min`/`f64::max`, which discard NaN rather than order it. + /// Fed the same two batches, it reports the stream's range as + /// `[1.0, 5.0]` and loses both infinities. Comparing keys with `Ord` + /// over every element has no such blind spot. + #[test] + fn extremes_survive_a_batch_whose_ends_are_nan() { + let codec = SortKeyCodec::try_new(&DataType::Float64, sort_options(false, false)) + .unwrap(); + let sketch = sketch_batches( + &codec, + vec![ + vec![1.0, 5.0], + vec![f64::NAN, -f64::NAN, f64::INFINITY, f64::NEG_INFINITY], + ], + ); + + // Per total_cmp, -NaN is below every value and +NaN above every + // value, so those are this stream's true extremes. + let min = codec.decode(*sketch.min().unwrap()).unwrap(); + let max = codec.decode(*sketch.max().unwrap()).unwrap(); + assert!( + matches!(min, ScalarValue::Float64(Some(v)) if v.is_nan() && v.is_sign_negative()), + "min must be -NaN, got {min:?}" + ); + assert!( + matches!(max, ScalarValue::Float64(Some(v)) if v.is_nan() && v.is_sign_positive()), + "max must be +NaN, got {max:?}" + ); + } + + /// Same shape without NaN, which is the plainer data-loss case: a + /// stream spanning both infinities must report spanning both + /// infinities. T-Digest reports `[1.0, 5.0]` here once a NaN has + /// appeared in the same batch as the infinities. + #[test] + fn infinities_are_reported_as_the_extremes_they_are() { + let codec = SortKeyCodec::try_new(&DataType::Float64, sort_options(false, false)) + .unwrap(); + let sketch = sketch_batches( + &codec, + vec![vec![1.0, 5.0], vec![f64::INFINITY, f64::NEG_INFINITY]], + ); + assert_eq!( + codec.decode(*sketch.min().unwrap()).unwrap(), + ScalarValue::Float64(Some(f64::NEG_INFINITY)) + ); + assert_eq!( + codec.decode(*sketch.max().unwrap()).unwrap(), + ScalarValue::Float64(Some(f64::INFINITY)) + ); + } + + /// Batch arrival order must not change the answer. T-Digest's is order + /// dependent: the first batch assigns extremes directly, later batches + /// fold theirs in with NaN-dropping comparisons, so the same data gives + /// different bounds depending on which batch landed first. + #[test] + fn extremes_do_not_depend_on_batch_order() { + let codec = SortKeyCodec::try_new(&DataType::Float64, sort_options(false, false)) + .unwrap(); + let nan_batch = vec![f64::NAN, -f64::NAN, f64::INFINITY, f64::NEG_INFINITY]; + let plain_batch = vec![1.0, 5.0]; + + let nan_first = + sketch_batches(&codec, vec![nan_batch.clone(), plain_batch.clone()]); + let nan_second = sketch_batches(&codec, vec![plain_batch, nan_batch]); + + assert_eq!(nan_first.min(), nan_second.min(), "min is order dependent"); + assert_eq!(nan_first.max(), nan_second.max(), "max is order dependent"); + } + + /// Build a sketch over `values`, ingested as one batch. + fn int_sketch(values: Vec<Option<i64>>, options: SortOptions) -> SortKeySketch { + let codec = SortKeyCodec::try_new(&DataType::Int64, options).unwrap(); + let mut sketch = SortKeySketch::new(codec); + sketch.ingest(&Int64Array::from(values)).unwrap(); + sketch + } + + /// NULLs are counted, not sketched, so they must show up in the total + /// without disturbing the values. + #[test] + fn nulls_are_counted_beside_the_values() { + let sketch = int_sketch( + vec![Some(10), None, Some(20), None, None], + sort_options(false, true), + ); + assert_eq!(sketch.count(), 5, "population is values plus NULLs"); + assert_eq!(sketch.null_count(), 3); + } + + /// With no NULLs the population rank is the value rank, so the remap + /// must be a no-op. This is the case the general path's `rank <= n` + /// test would get wrong at q = 0, which is why it has its own branch. + #[test] + fn quantiles_are_unshifted_when_nothing_is_null() { + let values: Vec<Option<i64>> = (1..=100).map(Some).collect(); + let sketch = int_sketch(values, sort_options(false, true)); + assert_eq!( + sketch.quantile(0.0).unwrap(), + Some(ScalarValue::Int64(Some(1))), + "q=0 must be the smallest value, not a NULL" + ); + assert_eq!( + sketch.quantile(1.0).unwrap(), + Some(ScalarValue::Int64(Some(100))) + ); + } + + /// Half NULL, NULLs first. The bottom half of the population is the + /// NULL run, so every quantile below the midpoint is a NULL and the + /// top half compresses the whole value range into `q > 0.5`. A + /// consumer that skipped the remap would report the value median at + /// q=0.5 and skew every cut. + #[test] + fn nulls_first_shifts_the_value_range_upward() { + let mut values: Vec<Option<i64>> = (1..=50).map(Some).collect(); + values.extend(std::iter::repeat_n(None, 50)); + let sketch = int_sketch(values, sort_options(false, true)); + assert_eq!(sketch.count(), 100); + assert_eq!(sketch.null_count(), 50); + + let null = ScalarValue::Int64(None); + assert_eq!(sketch.quantile(0.0).unwrap(), Some(null.clone())); + assert_eq!(sketch.quantile(0.25).unwrap(), Some(null.clone())); + assert_eq!( + sketch.quantile(0.5).unwrap(), + Some(null), + "the NULL run reaches exactly the midpoint" + ); + // Past the run, rank 0.75·100 = 75 is value 25 of 50. + assert_eq!( + sketch.quantile(0.75).unwrap(), + Some(ScalarValue::Int64(Some(25))) + ); + assert_eq!( + sketch.quantile(1.0).unwrap(), + Some(ScalarValue::Int64(Some(50))) + ); + } + + /// Same data, NULLs last: the values now occupy the bottom half and + /// the NULL run the top. + #[test] + fn nulls_last_leaves_the_value_range_at_the_bottom() { + let mut values: Vec<Option<i64>> = (1..=50).map(Some).collect(); + values.extend(std::iter::repeat_n(None, 50)); + let sketch = int_sketch(values, sort_options(false, false)); + + assert_eq!( + sketch.quantile(0.25).unwrap(), + Some(ScalarValue::Int64(Some(25))), + "rank 25 of 100 is value 25 of 50" + ); + assert_eq!( + sketch.quantile(0.5).unwrap(), + Some(ScalarValue::Int64(Some(50))), + "the value run ends exactly at the midpoint" + ); + assert_eq!( + sketch.quantile(0.75).unwrap(), + Some(ScalarValue::Int64(None)) + ); + assert_eq!( + sketch.quantile(1.0).unwrap(), + Some(ScalarValue::Int64(None)) + ); + } + + /// All NULL: every quantile is a NULL, and nothing consults the empty + /// value sketch. + #[test] + fn an_all_null_column_answers_null_everywhere() { + let sketch = int_sketch(vec![None; 8], sort_options(false, true)); + for q in [0.0, 0.5, 1.0] { + assert_eq!( + sketch.quantile(q).unwrap(), + Some(ScalarValue::Int64(None)), + "q={q}" + ); + } + assert_eq!(sketch.count(), 8); + } + + /// With no values at all the NULL run is unbounded on both sides, so + /// both extremes are NULL whichever end the run nominally occupies. + /// + /// The end it does not occupy has no value to fall back to, and + /// answering `None` there would claim nothing was observed while + /// `count` says otherwise. `cut_partitions` routes files on the pair, + /// so half a range routes half the rows. + #[test] + fn an_all_null_column_is_null_at_both_extremes() { + for nulls_first in [true, false] { + let sketch = int_sketch(vec![None; 8], sort_options(false, nulls_first)); + assert_eq!( + sketch.min().unwrap(), + Some(ScalarValue::Int64(None)), + "nulls_first={nulls_first}" + ); + assert_eq!( + sketch.max().unwrap(), + Some(ScalarValue::Int64(None)), + "nulls_first={nulls_first}" + ); + } + } + + /// Nothing observed at all is distinct from observing NULLs. + #[test] + fn an_empty_sketch_answers_none_not_null() { + let codec = + SortKeyCodec::try_new(&DataType::Int64, sort_options(false, true)).unwrap(); + let sketch = SortKeySketch::new(codec); + assert_eq!(sketch.count(), 0); + assert_eq!(sketch.quantile(0.5).unwrap(), None); + assert_eq!(sketch.min().unwrap(), None); + assert!(sketch.cuts(4).unwrap().is_empty()); + } + + /// The extremes follow NULL placement: with NULLs first the least row + /// is a NULL, and the greatest is still the largest value. + #[test] + fn extremes_account_for_where_nulls_sort() { + let values = vec![Some(10), None, Some(20)]; + let first = int_sketch(values.clone(), sort_options(false, true)); + assert_eq!(first.min().unwrap(), Some(ScalarValue::Int64(None))); + assert_eq!(first.max().unwrap(), Some(ScalarValue::Int64(Some(20)))); + + let last = int_sketch(values, sort_options(false, false)); + assert_eq!(last.min().unwrap(), Some(ScalarValue::Int64(Some(10)))); + assert_eq!(last.max().unwrap(), Some(ScalarValue::Int64(None))); + } + + /// Merging folds both the value sketches and the NULL counts. Losing + /// the second sketch's NULLs would silently shift every quantile. + #[test] + fn merge_folds_values_and_null_counts() { + let options = sort_options(false, true); + let mut left = int_sketch(vec![Some(1), Some(2), None], options); + let right = int_sketch(vec![Some(3), Some(4), None, None], options); + left.merge(right).unwrap(); + + assert_eq!(left.count(), 7, "3 + 4 rows"); + assert_eq!(left.null_count(), 3, "1 + 2 NULLs"); + assert_eq!(left.max().unwrap(), Some(ScalarValue::Int64(Some(4)))); + assert_eq!(left.min().unwrap(), Some(ScalarValue::Int64(None))); + } + + /// Sketches of different keys must not fold together. The result would + /// look like a perfectly ordinary distribution of nothing in + /// particular. + #[test] + fn merge_rejects_a_different_sort_key() { + let options = sort_options(false, true); + let mut ints = int_sketch(vec![Some(1)], options); + + let float_codec = SortKeyCodec::try_new(&DataType::Float64, options).unwrap(); + let floats = SortKeySketch::new(float_codec); + let err = ints + .merge(floats) + .expect_err("merging Int64 with Float64 must be refused"); + assert!( + err.to_string().contains("different sort keys"), + "got: {err}" + ); + + // Direction is part of the key's identity too: the same column + // sketched ASC and DESC has keys running opposite ways. + let mut ascending = int_sketch(vec![Some(1)], sort_options(false, true)); + let descending = int_sketch(vec![Some(1)], sort_options(true, true)); + assert!( + ascending.merge(descending).is_err(), + "ASC and DESC keys are inverses; folding them is meaningless" + ); + } + + /// `cuts` splits the population, NULL run included, so a mostly-NULL + /// column spends its low cuts inside that run and only crosses into + /// the values once the run is behind it. + /// + /// 60 NULLs then values 1..=40, NULLs first. Quartile ranks are 25, 50 + /// and 75 of 100; the first two sit inside the 60-row NULL run, and the + /// third is 15 rows past it, which is value 15 of 40. + #[test] + fn cuts_split_the_population_including_nulls() { + let mut values: Vec<Option<i64>> = vec![None; 60]; + values.extend((1..=40).map(Some)); + let sketch = int_sketch(values, sort_options(false, true)); + + let cuts = sketch.cuts(4).unwrap(); + assert_eq!(cuts.len(), 3, "K-1 cuts for K partitions"); + assert_eq!( + cuts[0], + ScalarValue::Int64(None), + "the first quarter is entirely NULL" + ); + assert_eq!(cuts[1], ScalarValue::Int64(None), "so is the second"); + assert_eq!( + cuts[2], + ScalarValue::Int64(Some(15)), + "the third crosses out of the run" + ); + } + + /// A NULL run ending exactly on a cut still belongs to the NULL side: + /// the rank it occupies is the run's last row, not the values' first. + /// 75 NULLs of 100 rows puts the 0.75 cut precisely there. + #[test] + fn a_cut_landing_on_the_end_of_the_null_run_is_null() { + let mut values: Vec<Option<i64>> = vec![None; 75]; + values.extend((1..=25).map(Some)); + let sketch = int_sketch(values, sort_options(false, true)); + assert_eq!( + sketch.quantile(0.75).unwrap(), + Some(ScalarValue::Int64(None)), + "rank 75 of 100 is the last NULL, so three of four partitions \ + hold only NULLs" + ); + } + + /// The design priced and dropped in `benchmarks/benches/ + /// quantile_sketch.rs` as `kll_norm_u128`, kept here as an oracle. + /// + /// It puts NULLs *inside* the key, tagged above the value bits, so the + /// sketch sees the whole population in order and a quantile is a + /// straight lookup with no arithmetic. That makes it the independent + /// check on [`SortKeySketch::quantile`], whose whole job is to + /// reproduce this answer while keeping the key 8 bytes wide. + /// + /// Value encoding is shared, so a disagreement can only come from NULL + /// placement or the rank remap. + struct NullsInKeyOracle { + sketch: KllSketch<u128>, + codec: SortKeyCodec, + null_tag: u128, + } + + impl NullsInKeyOracle { + fn new(values: &[Option<i64>], options: SortOptions) -> Self { + let codec = SortKeyCodec::try_new(&DataType::Int64, options).unwrap(); + // The tag alone decides which side the NULL run sits on. + let (null_tag, value_tag) = if options.nulls_first { + (0u128, 1u128) + } else { + (1u128, 0u128) + }; + let keys: Vec<u128> = values + .iter() + .map(|value| match value { + None => null_tag << 64, + Some(value) => { + let key = + codec.encode(&Int64Array::from(vec![*value])).unwrap()[0]; + (value_tag << 64) | key as u128 + } + }) + .collect(); + let mut sketch = KllSketch::<u128>::new(KLL_K); + sketch.absorb_slice(&keys); + Self { + sketch, + codec, + null_tag, + } + } + + fn quantile(&self, q: f64) -> Option<ScalarValue> { + let key = self.sketch.quantile(q)?; + if (key >> 64) == self.null_tag { + Some(self.codec.null_value().unwrap()) + } else { + Some(self.codec.decode(*key as u64).unwrap()) + } + } + } + + /// Differential test: the shipped out-of-band-NULL sketch must answer + /// exactly what the nulls-in-the-key oracle answers, across every NULL + /// fraction and both placements. + /// + /// Both sketches are exact here — the fixtures are far below `KLL_K`, + /// so nothing compacts — which means any disagreement is the rank + /// remap and not sketch error. + #[test] + fn quantiles_match_the_nulls_in_key_oracle() { + // Deliberately includes 0 and every-row-NULL, plus fractions whose + // ranks land exactly on the run boundary for some q. + let null_counts = [0usize, 1, 25, 50, 60, 75, 99, 100]; + let quantile_probes: Vec<f64> = (0..=20).map(|step| step as f64 / 20.0).collect(); + + for nulls in null_counts { + let mut values: Vec<Option<i64>> = vec![None; nulls]; + values.extend((1..=(100 - nulls) as i64).map(Some)); + for nulls_first in [true, false] { + let options = sort_options(false, nulls_first); + let oracle = NullsInKeyOracle::new(&values, options); + let sketch = int_sketch(values.clone(), options); + + assert_eq!( + sketch.count(), + 100, + "nulls={nulls} nulls_first={nulls_first}: population size" + ); + for &q in &quantile_probes { + assert_eq!( + sketch.quantile(q).unwrap(), + oracle.quantile(q), + "nulls={nulls} nulls_first={nulls_first} q={q}: \ + out-of-band NULLs disagreed with nulls-in-key" + ); + } + } + } + } + + /// The same agreement has to hold under `DESC`, where the value keys + /// are inverted but the NULL run still sits where `nulls_first` says. + #[test] + fn descending_quantiles_match_the_nulls_in_key_oracle() { + let mut values: Vec<Option<i64>> = vec![None; 30]; + values.extend((1..=70).map(Some)); + for nulls_first in [true, false] { + let options = sort_options(true, nulls_first); + let oracle = NullsInKeyOracle::new(&values, options); + let sketch = int_sketch(values.clone(), options); + for step in 0..=20 { + let q = step as f64 / 20.0; + assert_eq!( + sketch.quantile(q).unwrap(), + oracle.quantile(q), + "DESC nulls_first={nulls_first} q={q}" + ); + } + } + } + + /// Indices of `keys` in ascending key order. + fn argsort(keys: &[u64]) -> Vec<usize> { + let mut order: Vec<usize> = (0..keys.len()).collect(); + order.sort_by_key(|&i| keys[i]); + order + } + + /// The worked table in `impl_sortable_float`'s docs, asserted. Hex + /// written by hand into a comment rots silently; pinning it here makes + /// a drift fail instead. + #[test] + fn float_key_table_in_docs_is_accurate() { + let table: [(f64, u64); 10] = [ + (-f64::NAN, 0x0007FFFFFFFFFFFF), + (f64::NEG_INFINITY, 0x000FFFFFFFFFFFFF), + (-2.0, 0x3FFFFFFFFFFFFFFF), + (-1.0, 0x400FFFFFFFFFFFFF), + (-0.0, 0x7FFFFFFFFFFFFFFF), + (0.0, 0x8000000000000000), + (1.0, 0xBFF0000000000000), + (2.0, 0xC000000000000000), + (f64::INFINITY, 0xFFF0000000000000), + (f64::NAN, 0xFFF8000000000000), + ]; + for (value, expected) in table { + assert_eq!( + value.to_key(), + expected, + "{value} should key to {expected:#018X}, got {:#018X}", + value.to_key() + ); + } + assert!( + table.windows(2).all(|w| w[0].1 < w[1].1), + "the table is written in ascending value order, so its keys \ + must ascend too" + ); + } + + /// Unsigned keys must not be shifted the way signed ones are, or the + /// top half of the range would wrap below the bottom. + #[test] + fn unsigned_keys_span_full_range_in_order() { + let values = vec![0u64, 1, 1 << 62, u64::MAX - 1, u64::MAX]; + let codec = + SortKeyCodec::try_new(&DataType::UInt64, sort_options(false, false)).unwrap(); + let keys = codec.encode(&UInt64Array::from(values.clone())).unwrap(); + assert!(keys.windows(2).all(|w| w[0] < w[1]), "{keys:?}"); + for (key, value) in keys.iter().zip(&values) { + assert_eq!( + codec.decode(*key).unwrap(), + ScalarValue::UInt64(Some(*value)) + ); + } + } + + /// The motivating case: nanosecond timestamps round-trip exactly at + /// magnitudes where an `f64` cast would quantize them. The timezone + /// rides on the `DataType` rather than the key, so it has to survive + /// the decode as well. + #[test] + fn nanosecond_timestamps_round_trip_exactly_with_timezone() { + let data_type = DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())); + // 2026-08-12T00:00:00Z in nanos, then the next three nanoseconds. + // Above 2^53, so an f64 cast would collapse them onto one value. + let base = 1_786_233_600_000_000_000i64; + let values = vec![base, base + 1, base + 2, base + 3]; + assert!( + base as f64 as i64 != base + 1, + "test premise: f64 must be unable to separate adjacent nanos here" + ); + + let array = TimestampNanosecondArray::from(values.clone()) + .with_timezone("UTC".to_string()); + let codec = + SortKeyCodec::try_new(&data_type, sort_options(false, false)).unwrap(); + let keys = codec.encode(&array).unwrap(); + assert!( + keys.windows(2).all(|w| w[0] < w[1]), + "adjacent nanoseconds must stay distinct and ordered: {keys:?}" + ); + for (key, value) in keys.iter().zip(&values) { + assert_eq!( + codec.decode(*key).unwrap(), + ScalarValue::TimestampNanosecond(Some(*value), Some("UTC".into())), + "decode must preserve both the instant and the timezone" + ); + } + } + + /// Encode `ascending` under both directions, asserting the keys run + /// the way the direction asks and that each one decodes back to the + /// value it came from. + /// + /// `ascending` must be in `total_cmp` order and free of duplicates, + /// since the keys are checked for strict monotonicity. + fn check_dispatched<T: ArrowPrimitiveType>(ascending: Vec<T::Native>) { + let array = PrimitiveArray::<T>::from_iter_values(ascending.iter().copied()); + let data_type = array.data_type().clone(); + for descending in [false, true] { + let codec = SortKeyCodec::try_new(&data_type, sort_options(descending, true)) + .unwrap_or_else(|| { + panic!("{data_type:?} is dispatched but try_new declined it") + }); + let keys = codec.encode(&array).unwrap(); + assert_eq!(keys.len(), ascending.len(), "{data_type:?}"); + + for pair in keys.windows(2) { + let [lo, hi] = pair else { + unreachable!("windows(2) yields pairs") + }; + let ordered = if descending { lo > hi } else { lo < hi }; + assert!( + ordered, + "{data_type:?} descending={descending}: {lo:#018x} then {hi:#018x}" + ); + } + + for (key, value) in keys.iter().zip(&ascending) { + let expected = + ScalarValue::new_primitive::<T>(Some(*value), &data_type).unwrap(); + assert_eq!( + codec.decode(*key).unwrap(), + expected, + "{data_type:?} descending={descending}" + ); + } + } + } + + /// Every type `dispatch_sortable!` admits, end to end: the allowlist + /// accepts it, the keys run the way the direction asks, and the decode + /// arm hands back the value that went in. + /// + /// The narrow widths are their own code rather than a special case of + /// the 64-bit ones. `impl_sortable_float!(f32, u32, 32)` zero-extends a + /// 32-bit key that DESC then inverts across all 64 bits, and the narrow + /// signed arms sign-extend on the way out and truncate on the way back. + #[test] + fn every_dispatched_type_encodes_in_order_and_decodes_back() { + check_dispatched::<Int8Type>(vec![i8::MIN, -1, 0, 1, i8::MAX]); + check_dispatched::<Int16Type>(vec![i16::MIN, -1, 0, 1, i16::MAX]); + check_dispatched::<Int32Type>(vec![i32::MIN, -1, 0, 1, i32::MAX]); + check_dispatched::<Int64Type>(vec![i64::MIN, -1, 0, 1, i64::MAX]); Review Comment: > 20 of 24 dispatched types have no round-trip test Fixed. `every_dispatched_type_encodes_in_order_and_decodes_back` walks one fixture per dispatched type: `try_new` accepts it, keys are strictly monotone in the direction asked for, and every key decodes back to the value it came from. Float fixtures carry both NaNs, both infinities and both zeros. > That leaves the entire `impl_sortable_float!(f32, u32, 32)` instantiation unpinned Mutation-tested rather than trusted green. Changing it to `impl_sortable_float!(f32, u32, 31)` fails the new test at `Float32 descending=false: 0x00000000ff800000 then 0x00000000c0000000`, and fails nothing else in the file. The arm was correct; it had no coverage. -- 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]
