diff --git a/cpp/heimdall_common/impl/filtered_column.hpp b/cpp/heimdall_common/impl/filtered_column.hpp index 26ef363fa0..b089675bae 100644 --- a/cpp/heimdall_common/impl/filtered_column.hpp +++ b/cpp/heimdall_common/impl/filtered_column.hpp @@ -129,7 +129,7 @@ class filtered_column : public chained_column_view, public query_core::index_hol // The incoming filter operates on our filtered space, so we need to expand it // to the source space by applying it through our index mapping auto expanded_filter = std::make_shared(); - for (int64_t i = 0; i < filter->cardinality() && i < indices_.size(); ++i) { + for (int64_t i = 0; i < indices_.size(); ++i) { if (filter->contains(i)) { auto source_idx = indices_[i]; if (source_idx < source_size) { @@ -144,15 +144,15 @@ class filtered_column : public chained_column_view, public query_core::index_hol return source_->index_holder() ->run_query(info, data, combined_filter) - .then([indices = indices_](auto&& results) { + .then([indices = indices_, filter = std::move(filter)](auto&& results) { // Convert results back to our filtered index space for (auto& result : results) { std::vector new_indices; new_indices.reserve(result.indices.size()); for (auto source_idx : result.indices) { - // Find where this source index maps to in our filtered space + // Preserve the first-match behavior, restricted to eligible view positions. for (int64_t i = 0; i < indices.size(); ++i) { - if (indices[i] == source_idx) { + if (indices[i] == source_idx && (!filter || filter->contains(i))) { new_indices.push_back(i); break; } diff --git a/python/tests/test_filtered_column_query.py b/python/tests/test_filtered_column_query.py new file mode 100644 index 0000000000..0477b4533c --- /dev/null +++ b/python/tests/test_filtered_column_query.py @@ -0,0 +1,51 @@ +"""Exercise sparse filters through real columns and nested dataset queries. + +Run with the native extension built from this checkout. A published wheel can +validate the fixture/API, but cannot validate changes to the C++ header here. +""" + +import deeplake +import pytest + + +@pytest.fixture +def reordered_view(tmp_path): + """Create an indexed column and a view with source mapping [9, 2, 7].""" + dataset = deeplake.create(str(tmp_path / "filtered-column")) + dataset.add_column("id", deeplake.types.Int32()) + dataset.add_column("rank", deeplake.types.Int32()) + dataset.add_column("embedding", deeplake.types.Embedding(size=2)) + dataset.append( + [ + { + "id": index, + "rank": {9: 0, 2: 1, 7: 2}.get(index, 10), + "embedding": [1.0, float(index)], + } + for index in range(10) + ] + ) + dataset.commit() + view = dataset.query("SELECT * WHERE rank < 3 ORDER BY rank") + assert [int(row["id"]) for row in view] == [9, 2, 7] + return view + + +@pytest.mark.parametrize("source_id", [2, 7]) +def test_sparse_filter_after_first_view_position(reordered_view, source_id): + """One selected row may occupy view position one or two, not position zero.""" + result = reordered_view.query( + f"SELECT * WHERE id = {source_id} " + f"ORDER BY COSINE_SIMILARITY(embedding, ARRAY[1.0, {source_id}.0]) " + "DESC LIMIT 1" + ) + assert [int(row["id"]) for row in result] == [source_id] + + +def test_top_k_without_additional_filter(reordered_view): + """Mapping a source hit back to the view preserves top-k ordering.""" + result = reordered_view.query( + "SELECT * ORDER BY COSINE_SIMILARITY(embedding, ARRAY[1.0, 7.0]) " + "DESC LIMIT 1" + ) + assert [int(row["id"]) for row in result] == [7]