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& items, char order); void merge(const vector_of_kll_sketches& other); + // The array alternative is typed so that nanobind converts other dtypes + // (e.g. float32 ranks, int64 indices) instead of reinterpreting the buffer. template - using ArrInputType = std::variant, nb::list, V>; + using ArrInputType = std::variant>, nb::list, V>; // returns a single sketch combining all data in the array kll_sketch collapse(ArrInputType& isk) const; @@ -194,8 +196,8 @@ auto vector_of_kll_sketches::make_ndarray(size_t rows, size_t cols) const template template auto vector_of_kll_sketches::input_to_vec(ArrInputType& input) const -> Array1D { - if (std::holds_alternative>(input)) { - nb::ndarray<> arr = std::get>(input); + if (std::holds_alternative>>(input)) { + nb::ndarray> arr = std::get>>(input); return Array1D(arr); } else if (std::holds_alternative(input)) { nb::list list = std::get(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()