avantgardnerio commented on code in PR #2294: URL: https://github.com/apache/datafusion-ballista/pull/2294#discussion_r3786725902
########## ballista/core/src/sort_key.rs: ########## @@ -0,0 +1,1588 @@ +// 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. + /// + /// The match is exact `DataType` equality, so a codec built for + /// `Timestamp(ns, Some("UTC"))` rejects an array tagged `Some("+00:00")` + /// even though the two encode to identical keys. Strict because the + /// codec's type is what [`Self::decode`] rebuilds a `ScalarValue` from, + /// so accepting an array of a type the codec doesn't carry would + /// relabel every cut it later produces. Review Comment: > **`encode` requires exact `DataType` equality** Documented rather than changed. The doc previously said the types must "match", which reads as a semantic type change, and the `data_type` field doc right above says the timezone is deliberately *not* in the key, so the rejection looked like a bug. It now states that the check is `DataType` equality and gives the reason: the codec's type is what `decode` rebuilds a `ScalarValue` from, so accepting an array of a type the codec doesn't carry would relabel every cut it produces. > Strictness is defensible; worth knowing when the consumer lands. Agreed, and left open. Codec and array come from the same routing expression today, so they cannot diverge. The case that would force this is the consumer building the codec in the scheduler from the table schema while the array comes from a scan whose file metadata tags the zone differently. If that turns up, the principled relaxation is comparing the dispatch arm rather than the full type, since the key only depends on the primitive. Deciding it there, where there is something to test against. ########## ballista/core/src/kll.rs: ########## @@ -202,6 +202,39 @@ fn level_capacity(k: usize, num_levels: usize, height: usize) -> usize { raw.max(MIN_LEVEL_WIDTH) } +/// Summarizes rather than dumping the compactor stack, which holds on the +/// order of `3k` items and would bury whatever else a caller was printing. +impl<T: Ord + Clone> std::fmt::Debug for KllSketch<T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("KllSketch") + .field("k", &self.k) + .field("count", &self.count()) + .field("levels", &self.levels.len()) + .field("retained", &self.levels.iter().map(Vec::len).sum::<usize>()) + .finish() + } +} + +/// Clone copies the compactor stack and the tracked extremes, and gives +/// the copy a fresh PRNG rather than duplicating the original's position. +/// +/// Written by hand because `StdRng` is not `Clone`. Reseeding is the +/// correct behaviour anyway: two sketches sharing a coin-flip sequence +/// would correlate their compaction decisions, and every quantile the +/// clone can answer is already determined by the state that was copied. +impl<T: Ord + Clone> Clone for KllSketch<T> { Review Comment: > **`Clone` for `KllSketch` reseeds from `rand::random()`**, so cloning is non-deterministic Agreed on all three points, no change. The reasoning is already on the impl, including the decorrelation argument: two sketches sharing a coin-flip sequence would correlate their compaction decisions, and everything the clone can answer is already fixed by the state that was copied. ########## ballista/core/src/kll.rs: ########## @@ -477,15 +539,37 @@ impl<T: Ord + Clone> KllSketch<T> { if q == 1.0 { return self.max.as_ref(); } - let total_weight: u64 = self - .levels - .iter() - .enumerate() - .map(|(h, level)| (1u64 << h) * level.len() as u64) - .sum(); + let total_weight = self.count(); if total_weight == 0 { return None; } + self.at_rank((q * total_weight as f64) as u64) + } + + /// Return the item at `rank`, counting cumulative weight from the + /// smallest item up. `None` if the sketch is empty. + /// + /// Semantics: the smallest retained item whose cumulative weight is at + /// least `rank`. Ranks at or beyond the ends give the tracked extremes, + /// which bypass the compactor so coin-flip history can't move them. + /// + /// This is the primitive [`Self::quantile`] is expressed in, and the + /// one to prefer whenever the caller already knows the rank it wants. + /// Converting a known rank into a fraction and back loses it: for 99 + /// items, rank 59 becomes `59/99`, and multiplying that back by 99 + /// yields `58.999…`, which truncates to 58. Callers that adjust a rank + /// — stepping over a run of NULLs, say — must stay in integers. + pub fn at_rank(&self, rank: u64) -> Option<&T> { + let total_weight = self.count(); + if total_weight == 0 { + return None; + } + if rank == 0 { + return self.min.as_ref(); + } + if rank >= total_weight { + return self.max.as_ref(); + } Review Comment: > `at_rank` allocates and sorts the ~3k retained pairs per call, so `cuts(P)` is O(P · m log m) Correct, and left as is. The bench module doc carries the reason: at N=1M, P=64, K=64 cuts, merge and quantile are 3+ orders of magnitude cheaper than ingest, so this is not where the time goes. If a consumer ever calls `cuts` per batch rather than per stage that changes, and the fix is to sort the pairs once per `cuts` call instead of once per cut. -- 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]
