refactor: tighten KLL API and query performance

CHANGELOG.md

@@ -4,9 +4,17 @@ All significant changes to this project will be documented in this file.
 
 ## Unreleased
 
+### Breaking changes
+
+* Move `SearchCriteria` from `req` to `common` and remove its `Default` implementation. Import `datasketches::common::SearchCriteria` and explicitly choose `Inclusive` or `Exclusive` for each query.
+
 ### New features
 
-* Add KLL sketches behind the `kll` feature, including rank, quantile, PMF, and CDF queries, custom item ordering, merging, and C++/Java-compatible serialization.
+* Add KLL sketches behind the `kll` feature, with rank, quantile, PMF, and CDF queries, merging, totally ordered custom item types, a `KllFloat` adapter for non-NaN floating-point values, and serialization.
+
+### Improvements
+
+* Improve truncated-input diagnostics across sketch deserializers.
 
 ### Bug fixes
 

benchmarks/Cargo.toml

@@ -23,7 +23,7 @@ edition.workspace = true
 rust-version.workspace = true
 
 [dev-dependencies]
-datasketches = { workspace = true, features = ["cpc", "req", "tdigest"] }
+datasketches = { workspace = true, features = ["cpc", "kll", "req", "tdigest"] }
 divan = { workspace = true }
 rand = { workspace = true }
 

benchmarks/kll/merge.rs

@@ -0,0 +1,38 @@
+// 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.
+
+use divan::Bencher;
+use divan::black_box;
+use divan::counter::ItemsCount;
+
+use super::support::build_sketch;
+use super::support::values;
+
+#[divan::bench]
+fn merge(bencher: Bencher) {
+    let values = values(200_000);
+    let left = build_sketch(&values[..100_000]);
+    let right = build_sketch(&values[100_000..]);
+
+    bencher
+        .counter(ItemsCount::new(values.len()))
+        .with_inputs(|| left.clone())
+        .bench_local_values(|mut left| {
+            left.merge(black_box(&right)).unwrap();
+            black_box(left)
+        });
+}

benchmarks/kll/mod.rs

@@ -0,0 +1,22 @@
+// 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.
+
+mod merge;
+mod query;
+mod serde;
+mod support;
+mod update;

benchmarks/kll/query.rs

@@ -0,0 +1,50 @@
+// 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.
+
+use datasketches::common::SearchCriteria;
+use datasketches::kll::KllFloat;
+use divan::Bencher;
+use divan::black_box;
+
+use super::support::prepared_sketch;
+
+#[divan::bench]
+fn rank(bencher: Bencher) {
+    let sketch = prepared_sketch();
+    let item = KllFloat::<f64>::new(500_000.0).unwrap();
+    bencher.bench_local(|| black_box(&sketch).rank(black_box(&item), SearchCriteria::Inclusive));
+}
+
+#[divan::bench]
+fn quantile(bencher: Bencher) {
+    let sketch = prepared_sketch();
+    bencher.bench_local(|| black_box(&sketch).quantile(black_box(0.5), SearchCriteria::Inclusive));
+}
+
+#[divan::bench]
+fn sorted_view_quantile(bencher: Bencher) {
+    let view = prepared_sketch().sorted_view();
+    bencher.bench_local(|| black_box(&view).quantile(black_box(0.5), SearchCriteria::Inclusive));
+}
+
+#[divan::bench]
+fn batch_quantiles(bencher: Bencher) {
+    let sketch = prepared_sketch();
+    let ranks = [0.01, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99];
+    bencher
+        .bench_local(|| black_box(&sketch).quantiles(black_box(&ranks), SearchCriteria::Inclusive));
+}

benchmarks/kll/serde.rs

@@ -0,0 +1,41 @@
+// 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.
+
+use datasketches::kll::KllFloat;
+use datasketches::kll::KllSketch;
+use divan::Bencher;
+use divan::black_box;
+use divan::counter::BytesCount;
+
+use super::support::prepared_sketch;
+
+#[divan::bench]
+fn serialize(bencher: Bencher) {
+    let sketch = prepared_sketch();
+    let bytes = sketch.serialize();
+    bencher
+        .counter(BytesCount::new(bytes.len()))
+        .bench_local(|| black_box(&sketch).serialize());
+}
+
+#[divan::bench]
+fn deserialize(bencher: Bencher) {
+    let bytes = prepared_sketch().serialize();
+    bencher
+        .counter(BytesCount::new(bytes.len()))
+        .bench_local(|| KllSketch::<KllFloat<f64>>::deserialize(black_box(&bytes)).unwrap());
+}

benchmarks/kll/support.rs

@@ -0,0 +1,43 @@
+// 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.
+
+use datasketches::kll::KllFloat;
+use datasketches::kll::KllSketch;
+use rand::RngExt;
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+
+pub(super) const DEFAULT_K: u16 = 200;
+
+pub(super) fn values(len: usize) -> Vec<f64> {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..len)
+        .map(|_| rng.random_range(0.0..1_000_000.0))
+        .collect()
+}
+
+pub(super) fn build_sketch(values: &[f64]) -> KllSketch<KllFloat<f64>> {
+    let mut sketch = KllSketch::new(DEFAULT_K).unwrap();
+    for &value in values {
+        sketch.update(KllFloat::<f64>::new(value).unwrap());
+    }
+    sketch
+}
+
+pub(super) fn prepared_sketch() -> KllSketch<KllFloat<f64>> {
+    build_sketch(&values(100_000))
+}

benchmarks/kll/update.rs

@@ -0,0 +1,31 @@
+// 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.
+
+use divan::Bencher;
+use divan::black_box;
+use divan::counter::ItemsCount;
+
+use super::support::build_sketch;
+use super::support::values;
+
+#[divan::bench(args = [1_000, 10_000, 100_000])]
+fn update(bencher: Bencher, len: usize) {
+    let values = values(len);
+    bencher
+        .counter(ItemsCount::new(len))
+        .bench_local(|| build_sketch(black_box(&values)));
+}

benchmarks/main.rs

@@ -21,6 +21,7 @@ use divan::AllocProfiler;
 static ALLOC: AllocProfiler = AllocProfiler::system();
 
 mod cpc;
+mod kll;
 mod req;
 mod tdigest;
 

benchmarks/req/query.rs

@@ -15,8 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use datasketches::common::SearchCriteria;
 use datasketches::req::ReqFloat;
-use datasketches::req::SearchCriteria;
 use divan::Bencher;
 use divan::black_box;
 

datasketches/src/bloom/sketch.rs

@@ -495,11 +495,12 @@ impl BloomFilter {
                 .checked_add(1)
                 .and_then(|words| words.checked_mul(size_of::<u64>()))
                 .ok_or_else(|| Error::deserial("Bloom filter payload length overflows"))?;
-            if payload_bytes > cursor.remaining().len() {
-                return Err(Error::insufficient_data(format!(
-                    "Bloom filter payload requires {payload_bytes} bytes, got {}",
-                    cursor.remaining().len()
-                )));
+            let available_bytes = cursor.remaining().len();
+            if available_bytes < payload_bytes {
+                return Err(Error::insufficient_data_of(
+                    "Bloom filter payload",
+                    format_args!("expected {payload_bytes} bytes, got {available_bytes}"),
+                ));
             }
         }
         let mut bit_array = vec![0u64; num_words].into_boxed_slice();

datasketches/src/codec/assert.rs

@@ -21,7 +21,7 @@ use std::ops::RangeBounds;
 use crate::error::Error;
 
 pub fn insufficient_data(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error {
-    move |_| Error::insufficient_data(tag)
+    move |error| Error::insufficient_data_of(tag, error)
 }
 
 pub fn ensure_serial_version_is(expected: u8, actual: u8) -> Result<(), Error> {

datasketches/src/common/mod.rs

@@ -19,8 +19,10 @@
 
 mod num_std_dev;
 mod resize;
+mod search_criteria;
 pub use self::num_std_dev::NumStdDev;
 pub use self::resize::ResizeFactor;
+pub use self::search_criteria::SearchCriteria;
 
 #[cfg(any(feature = "cpc", feature = "hll"))]
 pub(crate) mod inv_pow2;

datasketches/src/common/search_criteria.rs

@@ -0,0 +1,25 @@
+// 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.
+
+/// Selects the rank definition used by rank, quantile, PMF, and CDF queries.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SearchCriteria {
+    /// Define rank as the fraction of values less than or equal to the boundary.
+    Inclusive,
+    /// Define rank as the fraction of values strictly less than the boundary.
+    Exclusive,
+}

datasketches/src/countmin/sketch.rs

@@ -430,11 +430,12 @@ impl<T: CountMinValue> CountMinSketch<T> {
             let payload_bytes = payload_values
                 .checked_mul(LONG_SIZE_BYTES)
                 .ok_or_else(|| Error::deserial("CountMin payload size overflows"))?;
-            if payload_bytes > cursor.remaining().len() {
-                return Err(Error::insufficient_data(format!(
-                    "CountMin payload requires {payload_bytes} bytes, got {}",
-                    cursor.remaining().len()
-                )));
+            let available_bytes = cursor.remaining().len();
+            if available_bytes < payload_bytes {
+                return Err(Error::insufficient_data_of(
+                    "CountMin payload",
+                    format_args!("expected {payload_bytes} bytes, got {available_bytes}"),
+                ));
             }
         }
 

datasketches/src/cpc/sketch.rs

@@ -775,10 +775,14 @@ impl CpcSketch {
         let payload_bytes = window_data_bytes
             .checked_add(table_data_bytes)
             .ok_or_else(|| Error::deserial("CPC payload length overflows"))?;
-        let payload = cursor
-            .remaining()
-            .get(..payload_bytes)
-            .ok_or_else(|| Error::deserial("insufficient data for CPC compressed payload"))?;
+        let available_bytes = cursor.remaining().len();
+        if available_bytes < payload_bytes {
+            return Err(Error::insufficient_data_of(
+                "CPC compressed payload",
+                format_args!("expected {payload_bytes} bytes, got {available_bytes}"),
+            ));
+        }
+        let payload = &cursor.remaining()[..payload_bytes];
         let (window_data, table_data) = payload.split_at(window_data_bytes);
         let (table, window) = match flavor {
             Flavor::Empty => (PairTable::new(2, lg_k + 6), vec![]),