From 4025ecbebe9118655366dff8d12f6a6131105389 Mon Sep 17 00:00:00 2001
From: Lee Rhodes <leerho@users.noreply.github.com>
Date: Sat, 10 Oct 2026 14:54:23 -0700
Subject: [PATCH] vector_of_kll: convert array inputs to the expected dtype

The query methods took array arguments as an untyped nb::ndarray<> and
input_to_vec() wrapped the buffer as the expected element type without
checking its dtype. A float32 ranks array was read as float64, a float64
values array was read as float32 by the floats sketches, and an int64
index array (numpy's default int) was read as int32, all silently.

Type the array alternative of ArrInputType as nb::ndarray<V, nb::ndim<1>>
so nanobind converts the input to the expected dtype. This fixes
get_quantiles, get_ranks, get_pmf, get_cdf, and the isk index argument
of every method that takes one.

Fixes #63.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/vector_of_kll.cpp       |  8 +++++---
 tests/vector_of_kll_test.py | 21 +++++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/vector_of_kll.cpp b/src/vector_of_kll.cpp
index 70139152..0f4fd64a 100644
--- a/src/vector_of_kll.cpp
+++ b/src/vector_of_kll.cpp
@@ -65,8 +65,10 @@ class vector_of_kll_sketches {
     void update(nb::ndarray<T>& items, char order);
     void merge(const vector_of_kll_sketches<T>& other);
 
+    // The array alternative is typed so that nanobind converts other dtypes
+    // (e.g. float32 ranks, int64 indices) instead of reinterpreting the buffer.
     template<typename V>
-    using ArrInputType = std::variant<nb::ndarray<>, nb::list, V>;
+    using ArrInputType = std::variant<nb::ndarray<V, nb::ndim<1>>, nb::list, V>;
 
     // returns a single sketch combining all data in the array
     kll_sketch<T, C> collapse(ArrInputType<int>& isk) const;
@@ -194,8 +196,8 @@ auto vector_of_kll_sketches<T, C>::make_ndarray(size_t rows, size_t cols) const
 template<typename T, typename C>
 template<typename TT>
 auto vector_of_kll_sketches<T, C>::input_to_vec(ArrInputType<TT>& input) const -> Array1D<TT> {
-  if (std::holds_alternative<nb::ndarray<>>(input)) {
-    nb::ndarray<> arr = std::get<nb::ndarray<>>(input);
+  if (std::holds_alternative<nb::ndarray<TT, nb::ndim<1>>>(input)) {
+    nb::ndarray<TT, nb::ndim<1>> arr = std::get<nb::ndarray<TT, nb::ndim<1>>>(input);
     return Array1D<TT>(arr);
   } else if (std::holds_alternative<nb::list>(input)) {
     nb::list list = std::get<nb::list>(input);
diff --git a/tests/vector_of_kll_test.py b/tests/vector_of_kll_test.py
index 51b2f0b2..adaae849 100644
--- a/tests/vector_of_kll_test.py
+++ b/tests/vector_of_kll_test.py
@@ -97,6 +97,27 @@ def test_kll_ints_sketches(self):
       kll = vector_of_kll_ints_sketches(k, d)
       self.assertTrue(np.all(kll.is_empty()))
 
+    def test_kll_input_dtypes(self):
+      # array inputs must give the same results as lists, whatever their dtype (issue #63)
+      kll = vector_of_kll_floats_sketches(200, 2)
+      kll.update(np.column_stack([np.arange(1, 101), np.arange(101, 201)]).astype(np.float32))
+
+      ranks = [0.0625, 0.9375] # exactly representable in float32
+      expected = kll.get_quantiles(ranks)
+      for dtype in (np.float32, np.float64):
+        np.testing.assert_array_equal(kll.get_quantiles(np.array(ranks, dtype=dtype)), expected)
+
+      values = [25.5, 150.5]
+      for fn in (kll.get_ranks, kll.get_pmf, kll.get_cdf):
+        expected = fn(values)
+        for dtype in (np.float32, np.float64):
+          np.testing.assert_array_equal(fn(np.array(values, dtype=dtype)), expected)
+
+      # sketch indices, including numpy's default int64
+      expected = kll.get_quantiles(0.5, [1, 0])
+      for dtype in (np.int32, np.int64):
+        np.testing.assert_array_equal(kll.get_quantiles(0.5, np.array([1, 0], dtype=dtype)), expected)
+
     def test_kll_2Dupdates(self):
       # 1D case tested in the first example
       # 2D case will follow same idea, but focusing on update()