Somewhat dynamic sampling (#2847)

Instead of hard-coding 640 samples (1% of 64K chunk), we try and sample
1% or at least 640 values.

Author: AdamGS

vortex-btrblocks/src/float/stats.rs

@@ -70,7 +70,7 @@ impl CompressorStats for FloatStats {
         &self.src
     }
 
-    fn sample_opts(&self, sample_size: u16, sample_count: u16, opts: GenerateStatsOptions) -> Self {
+    fn sample_opts(&self, sample_size: u32, sample_count: u32, opts: GenerateStatsOptions) -> Self {
         let sampled = sample(self.src.clone(), sample_size, sample_count)
             .to_primitive()
             .vortex_expect("primitive");

vortex-btrblocks/src/integer/stats.rs

@@ -142,7 +142,7 @@ impl CompressorStats for IntegerStats {
         &self.src
     }
 
-    fn sample_opts(&self, sample_size: u16, sample_count: u16, opts: GenerateStatsOptions) -> Self {
+    fn sample_opts(&self, sample_size: u32, sample_count: u32, opts: GenerateStatsOptions) -> Self {
         let sampled = sample(self.src.clone(), sample_size, sample_count)
             .to_primitive()
             .vortex_expect("primitive");

vortex-btrblocks/src/lib.rs

@@ -9,7 +9,7 @@ use vortex_array::variants::{ExtensionArrayTrait, PrimitiveArrayTrait, StructArr
 use vortex_array::{Array, ArrayRef, Canonical};
 use vortex_dtype::datetime::TemporalMetadata;
 use vortex_dtype::{DType, Nullability};
-use vortex_error::{VortexExpect, VortexResult};
+use vortex_error::{VortexExpect, VortexResult, VortexUnwrap};
 
 pub use crate::float::FloatCompressor;
 pub use crate::integer::IntCompressor;
@@ -39,6 +39,8 @@ impl Default for GenerateStatsOptions {
     }
 }
 
+const SAMPLE_SIZE: u32 = 64;
+
 /// Stats for the compressor.
 pub trait CompressorStats: Debug + Clone {
     type ArrayType: Array;
@@ -52,11 +54,11 @@ pub trait CompressorStats: Debug + Clone {
 
     fn source(&self) -> &Self::ArrayType;
 
-    fn sample(&self, sample_size: u16, sample_count: u16) -> Self {
+    fn sample(&self, sample_size: u32, sample_count: u32) -> Self {
         self.sample_opts(sample_size, sample_count, GenerateStatsOptions::default())
     }
 
-    fn sample_opts(&self, sample_size: u16, sample_count: u16, opts: GenerateStatsOptions) -> Self;
+    fn sample_opts(&self, sample_size: u32, sample_count: u32, opts: GenerateStatsOptions) -> Self;
 }
 
 /// Top-level compression scheme trait.
@@ -127,7 +129,16 @@ pub fn estimate_compression_ratio_with_sampling<T: Scheme + ?Sized>(
     let sample = if is_sample {
         stats.clone()
     } else {
-        stats.sample(64, 10)
+        // We want to sample about 1% of data
+        let source_len = stats.source().len();
+
+        // We want to sample about 1% of data, while keeping a minimal sample of 640 values.
+        let sample_count = usize::max(
+            (source_len / 100) / usize::try_from(SAMPLE_SIZE).vortex_unwrap(),
+            10,
+        );
+
+        stats.sample(SAMPLE_SIZE, sample_count.try_into().vortex_unwrap())
     };
 
     let after = compressor

vortex-btrblocks/src/sample.rs

@@ -5,7 +5,7 @@ use vortex_array::compute::slice;
 use vortex_array::{Array, ArrayRef};
 use vortex_error::VortexExpect;
 
-pub(crate) fn sample<T: Array + Clone>(input: T, sample_size: u16, sample_count: u16) -> ArrayRef {
+pub(crate) fn sample<T: Array + Clone>(input: T, sample_size: u32, sample_count: u32) -> ArrayRef {
     if input.len() <= (sample_size as usize) * (sample_count as usize) {
         return input.to_array();
     }
@@ -31,8 +31,8 @@ pub(crate) fn sample<T: Array + Clone>(input: T, sample_size: u16, sample_count:
 
 pub fn stratified_slices(
     length: usize,
-    sample_size: u16,
-    sample_count: u16,
+    sample_size: u32,
+    sample_count: u32,
     rng: &mut StdRng,
 ) -> Vec<(usize, usize)> {
     let total_num_samples: usize = (sample_count as usize) * (sample_size as usize);
@@ -62,7 +62,7 @@ pub fn stratified_slices(
 
 /// Split a range of array indices into as-equal-as-possible slices. If the provided `num_partitions` doesn't
 /// evenly divide into `length`, then the first `(length % num_partitions)` slices will have an extra element.
-pub fn partition_indices(length: usize, num_partitions: u16) -> Vec<(usize, usize)> {
+pub fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> {
     let num_long_parts = length % num_partitions as usize;
     let short_step = length / num_partitions as usize;
     let long_step = short_step + 1;

vortex-btrblocks/src/string.rs

@@ -68,7 +68,7 @@ impl CompressorStats for StringStats {
         &self.src
     }
 
-    fn sample_opts(&self, sample_size: u16, sample_count: u16, opts: GenerateStatsOptions) -> Self {
+    fn sample_opts(&self, sample_size: u32, sample_count: u32, opts: GenerateStatsOptions) -> Self {
         let sampled = sample(self.src.clone(), sample_size, sample_count)
             .to_varbinview()
             .vortex_expect("varbinview");