fix(columnar): estimate BlockwiseLinear's residuals per block, not per column

columnar/src/column_values/u64_based/blockwise_linear.rs

@@ -1,4 +1,5 @@
 use std::io::Write;
+use std::num::NonZeroU64;
 use std::sync::Arc;
 use std::{io, iter};
 
@@ -8,6 +9,7 @@ use tantivy_bitpacker::{BitPacker, BitUnpacker, compute_num_bits};
 
 use crate::MonotonicallyMappableToU64;
 use crate::column_values::u64_based::line::Line;
+use crate::column_values::u64_based::stats_collector::compute_gcd;
 use crate::column_values::u64_based::{ColumnCodec, ColumnCodecEstimator, ColumnStats};
 use crate::column_values::{ColumnValues, VecColumn};
 
@@ -44,25 +46,61 @@ fn compute_num_blocks(num_vals: u32) -> u32 {
 
 pub struct BlockwiseLinearEstimator {
     block: Vec<u64>,
-    values_num_bytes: u64,
+    values_num_bits: u64,
+    gcd_blocks: Vec<(u64, u64, u32)>,
     meta_num_bytes: u64,
 }
 
 impl Default for BlockwiseLinearEstimator {
     fn default() -> Self {
         Self {
             block: Vec::with_capacity(BLOCK_SIZE as usize),
-            values_num_bytes: 0u64,
+            values_num_bits: 0u64,
+            gcd_blocks: Vec::new(),
             meta_num_bytes: 0u64,
         }
     }
 }
 
 impl BlockwiseLinearEstimator {
+    fn block_min_and_gcd(&self) -> (u64, NonZeroU64) {
+        let Some((&first_val, rest)) = self.block.split_first() else {
+            return (0u64, NonZeroU64::MIN);
+        };
+        let mut block_min = first_val;
+        let mut block_gcd: Option<NonZeroU64> = None;
+        for &buffer_val in rest {
+            block_min = block_min.min(buffer_val);
+            if block_gcd.map(NonZeroU64::get) == Some(1) {
+                continue;
+            }
+            let Some(non_zero_diff) = NonZeroU64::new(buffer_val.abs_diff(first_val)) else {
+                continue;
+            };
+            block_gcd = Some(match block_gcd {
+                Some(gcd) => compute_gcd(non_zero_diff, gcd),
+                None => non_zero_diff,
+            });
+        }
+        (block_min, block_gcd.unwrap_or(NonZeroU64::MIN))
+    }
+
     fn flush_block_estimate(&mut self) {
         if self.block.is_empty() {
             return;
         }
+        let (block_min, block_gcd) = self.block_min_and_gcd();
+        if block_gcd.get() > 1 {
+            let divider = DividerU64::divide_by(block_gcd.get());
+            for buffer_val in self.block.iter_mut() {
+                *buffer_val = divider.divide(*buffer_val - block_min);
+            }
+        } else {
+            for buffer_val in self.block.iter_mut() {
+                *buffer_val -= block_min;
+            }
+        }
+
         let column = VecColumn::from(std::mem::take(&mut self.block));
         let line = Line::train(&column);
         self.block = column.into();
@@ -73,8 +111,12 @@ impl BlockwiseLinearEstimator {
             let val = buffer_val.wrapping_sub(interpolated_val);
             max_value = val.max(max_value);
         }
-        let bit_width = compute_num_bits(max_value) as usize;
-        self.values_num_bytes += (bit_width * self.block.len() + 7) as u64 / 8;
+        let num_rows = self.block.len() as u32;
+        if block_gcd.get() > 1 {
+            self.gcd_blocks.push((max_value, block_gcd.get(), num_rows));
+        } else {
+            self.values_num_bits += compute_num_bits(max_value) as u64 * u64::from(num_rows);
+        }
         self.meta_num_bytes += 1 + line.num_bytes();
     }
 }
@@ -88,13 +130,18 @@ impl ColumnCodecEstimator for BlockwiseLinearEstimator {
         }
     }
     fn estimate(&self, stats: &ColumnStats) -> Option<u64> {
-        let mut estimate = 4 + stats.num_bytes() + self.meta_num_bytes + self.values_num_bytes;
-        if stats.gcd.get() > 1 {
-            let estimate_gain_from_gcd =
-                (stats.gcd.get() as f32).log2().floor() * stats.num_rows as f32 / 8.0f32;
-            estimate = estimate.saturating_sub(estimate_gain_from_gcd as u64);
-        }
-        Some(estimate)
+        let gcd = stats.gcd.get();
+        let values_num_bits: u64 = self.values_num_bits
+            + self
+                .gcd_blocks
+                .iter()
+                .map(|&(max_residual, block_gcd, num_rows)| {
+                    let scale = (block_gcd / gcd).max(1);
+                    let bit_width = compute_num_bits(max_residual.saturating_mul(scale)) as u64;
+                    bit_width * u64::from(num_rows)
+                })
+                .sum::<u64>();
+        Some(4 + stats.num_bytes() + self.meta_num_bytes + values_num_bits.div_ceil(8))
     }
 
     fn finalize(&mut self) {

columnar/src/column_values/u64_based/stats_collector.rs

@@ -8,7 +8,7 @@ use crate::column_values::ColumnStats;
 /// Compute the gcd of two non null numbers.
 ///
 /// It is recommended, but not required, to feed values such that `large >= small`.
-fn compute_gcd(mut large: NonZeroU64, mut small: NonZeroU64) -> NonZeroU64 {
+pub(crate) fn compute_gcd(mut large: NonZeroU64, mut small: NonZeroU64) -> NonZeroU64 {
     loop {
         let rem: u64 = large.get() % small;
         if let Some(new_small) = NonZeroU64::new(rem) {

columnar/src/column_values/u64_based/tests.rs

@@ -322,6 +322,96 @@ fn estimation_test_bad_interpolation_case_monotonically_increasing() {
     assert_le!(bitpacked_estimation, linear_interpol_estimation);
 }
 
+#[test]
+fn estimation_blockwise_linear_accounts_for_the_gcd_per_block() {
+    let data: Vec<u64> = (0..20_000u64)
+        .map(|i| 1_700_000_000 + (i / 2_000) * 3_600)
+        .collect();
+
+    let mut stats_collector = StatsCollector::default();
+    let mut estimator = CodecType::BlockwiseLinear.estimator();
+    for &val in &data {
+        stats_collector.collect(val);
+        estimator.collect(val);
+    }
+    estimator.finalize();
+    let estimated = estimator.estimate(&stats_collector.stats()).unwrap();
+
+    let mut buffer = Vec::new();
+    serialize_u64_based_column_values(&&data[..], &[CodecType::BlockwiseLinear], &mut buffer)
+        .unwrap();
+    let actual = buffer.len() as u64;
+
+    assert!(
+        estimated * 10 >= actual * 9 && estimated * 9 <= actual * 10,
+        "estimated {estimated} bytes, wrote {actual}"
+    );
+}
+
+#[test]
+fn test_selection_does_not_pick_a_much_larger_codec() {
+    let n = 100_000u64;
+    let mix = |i: u64| i.wrapping_mul(0x9E37_79B9_7F4A_7C15);
+    let shapes: Vec<(&str, Vec<u64>)> = vec![
+        (
+            "fixed_cadence",
+            (0..n).map(|i| 1_700_000_000_000 + i * 250).collect(),
+        ),
+        (
+            "noisy_ramp",
+            (0..n)
+                .map(|i| 1_700_000_000_000 + i * 250 + mix(i) % 4_000)
+                .collect(),
+        ),
+        ("uniform_narrow", (0..n).map(|i| mix(i) % 1_000).collect()),
+        (
+            "plateaus",
+            (0..n).map(|i| (i / 500) * 1_000 + mix(i) % 8).collect(),
+        ),
+        ("two_values", {
+            let mut vals = vec![42_000u64; n as usize];
+            for val in vals.iter_mut().step_by(997) {
+                *val = u64::MAX / 2;
+            }
+            vals
+        }),
+        (
+            "shuffled_cadence",
+            (0..n)
+                .map(|i| 1_700_000_000_000 + (mix(i) % n) * 250)
+                .collect(),
+        ),
+    ];
+
+    let codec_sets: [&[CodecType]; 2] = [
+        &ALL_U64_CODEC_TYPES,
+        &[CodecType::Bitpacked, CodecType::BlockwiseLinear],
+    ];
+    for (name, vals) in shapes {
+        for codec_types in codec_sets {
+            let smallest = codec_types
+                .iter()
+                .map(|&codec_type| {
+                    let mut buffer = Vec::new();
+                    serialize_u64_based_column_values(&&vals[..], &[codec_type], &mut buffer)
+                        .unwrap();
+                    buffer.len()
+                })
+                .min()
+                .unwrap();
+            let mut chosen = Vec::new();
+            serialize_u64_based_column_values(&&vals[..], codec_types, &mut chosen).unwrap();
+            assert!(
+                chosen.len() * 100 <= smallest * 105,
+                "{name} over {codec_types:?}: selection wrote {} bytes ({:?}) when {smallest} \
+                 were available",
+                chosen.len(),
+                CodecType::try_from_code(chosen[0]).unwrap(),
+            );
+        }
+    }
+}
+
 #[test]
 fn test_fast_field_codec_type_to_code() {
     let mut count_codec = 0;