Skip to content

vector_of_kll: convert array inputs to the expected dtype #73

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
leerho merged 1 commit into from
Oct 11, 2026
Merged

vector_of_kll: convert array inputs to the expected dtype #73

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions src/vector_of_kll.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
21 changes: 21 additions & 0 deletions tests/vector_of_kll_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading