Repository: man-group/ArcticDB (language: python)

Issue:
Raise a descriptive exception when trying to filter on sparse data
Short term solution until it all works properly as covered by https://github.com/man-group/ArcticDB/issues/1404


Discussion / hints:

Fix patch:
diff --git a/cpp/arcticdb/column_store/memory_segment.hpp b/cpp/arcticdb/column_store/memory_segment.hpp
index 561feaee0f..73e6d587c7 100644
--- a/cpp/arcticdb/column_store/memory_segment.hpp
+++ b/cpp/arcticdb/column_store/memory_segment.hpp
@@ -385,10 +385,10 @@ class SegmentInMemory {
         impl_->set_string_pool(string_pool);
     }
 
-    SegmentInMemory filter(const util::BitSet& filter_bitset,
+    SegmentInMemory filter(util::BitSet&& filter_bitset,
                            bool filter_down_stringpool=false,
                            bool validate=false) const{
-        return SegmentInMemory(impl_->filter(filter_bitset, filter_down_stringpool, validate));
+        return SegmentInMemory(impl_->filter(std::move(filter_bitset), filter_down_stringpool, validate));
     }
 
     /// @see SegmentInMemoryImpl::truncate
diff --git a/cpp/arcticdb/column_store/memory_segment_impl.cpp b/cpp/arcticdb/column_store/memory_segment_impl.cpp
index 0bb0fbe763..8f8acf55df 100644
--- a/cpp/arcticdb/column_store/memory_segment_impl.cpp
+++ b/cpp/arcticdb/column_store/memory_segment_impl.cpp
@@ -168,9 +168,10 @@ void SegmentInMemoryImpl::drop_column(std::string_view name) {
     column_map_->erase(name);
 }
 
-std::shared_ptr<SegmentInMemoryImpl> SegmentInMemoryImpl::filter(const util::BitSet& filter_bitset,
+std::shared_ptr<SegmentInMemoryImpl> SegmentInMemoryImpl::filter(util::BitSet&& filter_bitset,
                                                    bool filter_down_stringpool,
                                                    bool validate) const {
+    filter_bitset.resize(row_count());
     bool is_input_sparse = is_sparse();
     auto num_values = filter_bitset.count();
     if(num_values == 0)
@@ -210,18 +211,19 @@ std::shared_ptr<SegmentInMemoryImpl> SegmentInMemoryImpl::filter(const util::Bit
                 } else {
                     bitset_including_sparse.resize((*column)->row_count());
                 }
-                if (bitset_including_sparse.count() == 0) {
-                    // No values are set in the sparse column, skip it
-                    return;
-                }
                 output_col_idx = output->add_column(field(column.index), bitset_including_sparse.count(), true);
                 final_bitset = &bitset_including_sparse;
             } else {
                 final_bitset = &filter_bitset;
             }
             auto& output_col = output->column(position_t(output_col_idx));
-            if (sparse_map)
+            if (sparse_map) {
                 output_col.opt_sparse_map() = std::make_optional<util::BitSet>();
+                if (final_bitset->count() == 0) {
+                    // No values are set in the sparse column, no more work to do
+                    return;
+                }
+            }
             auto output_ptr = reinterpret_cast<RawType*>(output_col.ptr());
             auto input_data =  (*column)->data();
 
@@ -585,7 +587,7 @@ std::vector<std::shared_ptr<SegmentInMemoryImpl>> SegmentInMemoryImpl::split(siz
         util::BitSetSizeType end = std::min(start + rows, total_rows);
         // set_range is close interval on [left, right]
         bitset.set_range(start, end - 1, true);
-        output.emplace_back(filter(bitset));
+        output.emplace_back(filter(std::move(bitset)));
     }
     return output;
 }
diff --git a/cpp/arcticdb/column_store/memory_segment_impl.hpp b/cpp/arcticdb/column_store/memory_segment_impl.hpp
index b09b726d56..9747be880b 100644
--- a/cpp/arcticdb/column_store/memory_segment_impl.hpp
+++ b/cpp/arcticdb/column_store/memory_segment_impl.hpp
@@ -758,7 +758,7 @@ class SegmentInMemoryImpl {
 
     std::shared_ptr<SegmentInMemoryImpl> get_output_segment(size_t num_values, bool pre_allocate=true) const;
 
-    std::shared_ptr<SegmentInMemoryImpl> filter(const util::BitSet& filter_bitset,
+    std::shared_ptr<SegmentInMemoryImpl> filter(util::BitSet&& filter_bitset,
                                                 bool filter_down_stringpool=false,
                                                 bool validate=false) const;
 
diff --git a/cpp/arcticdb/pipeline/filter_segment.hpp b/cpp/arcticdb/pipeline/filter_segment.hpp
index 142d132bd0..d9069830aa 100644
--- a/cpp/arcticdb/pipeline/filter_segment.hpp
+++ b/cpp/arcticdb/pipeline/filter_segment.hpp
@@ -13,10 +13,10 @@
 namespace arcticdb {
 
 inline SegmentInMemory filter_segment(const SegmentInMemory& input,
-                                      const util::BitSet& filter_bitset,
+                                      util::BitSet&& filter_bitset,
                                       bool filter_down_stringpool=false,
                                       bool validate=false) {
-    return input.filter(filter_bitset, filter_down_stringpool, validate);
+    return input.filter(std::move(filter_bitset), filter_down_stringpool, validate);
 }
 
 inline std::vector<SegmentInMemory> partition_segment(const SegmentInMemory& input,
diff --git a/cpp/arcticdb/processing/clause.cpp b/cpp/arcticdb/processing/clause.cpp
index 5d38a013c4..430fbe8a54 100644
--- a/cpp/arcticdb/processing/clause.cpp
+++ b/cpp/arcticdb/processing/clause.cpp
@@ -267,9 +267,9 @@ Composite<EntityIds> FilterClause::process(
         proc.set_expression_context(expression_context_);
         auto variant_data = proc.get(expression_context_->root_node_name_);
         util::variant_match(variant_data,
-                            [&proc, &output, this](const util::BitSet& bitset) {
+                            [&proc, &output, this](util::BitSet& bitset) {
                                 if (bitset.count() > 0) {
-                                    proc.apply_filter(bitset, optimisation_);
+                                    proc.apply_filter(std::move(bitset), optimisation_);
                                     output.push_back(push_entities(component_manager_, std::move(proc)));
                                 } else {
                                     log::version().debug("Filter returned empty result");
diff --git a/cpp/arcticdb/processing/processing_unit.cpp b/cpp/arcticdb/processing/processing_unit.cpp
index f0f98c4eed..2430022b6c 100644
--- a/cpp/arcticdb/processing/processing_unit.cpp
+++ b/cpp/arcticdb/processing/processing_unit.cpp
@@ -10,7 +10,7 @@
 namespace arcticdb {
 
 void ProcessingUnit::apply_filter(
-    const util::BitSet& bitset,
+    util::BitSet&& bitset,
     PipelineOptimisation optimisation) {
     internal::check<ErrorCode::E_ASSERTION_FAILURE>(segments_.has_value() && row_ranges_.has_value() && col_ranges_.has_value(),
                                                     "ProcessingUnit::apply_filter requires all of segments, row_ranges, and col_ranges to be present");
@@ -18,7 +18,7 @@ void ProcessingUnit::apply_filter(
 
     for (auto&& [idx, segment]: folly::enumerate(*segments_)) {
         auto seg = filter_segment(*segment,
-                                  bitset,
+                                  std::move(bitset),
                                   filter_down_stringpool);
         auto num_rows = seg.is_null() ? 0 : seg.row_count();
         row_ranges_->at(idx) = std::make_shared<pipelines::RowRange>(row_ranges_->at(idx)->first, row_ranges_->at(idx)->first + num_rows);
diff --git a/cpp/arcticdb/processing/processing_unit.hpp b/cpp/arcticdb/processing/processing_unit.hpp
index b9a499bd51..eb7f6494d6 100644
--- a/cpp/arcticdb/processing/processing_unit.hpp
+++ b/cpp/arcticdb/processing/processing_unit.hpp
@@ -93,7 +93,7 @@ namespace arcticdb {
             return *this;
         }
 
-        void apply_filter(const util::BitSet& bitset, PipelineOptimisation optimisation);
+        void apply_filter(util::BitSet&& bitset, PipelineOptimisation optimisation);
 
         void truncate(size_t start_row, size_t end_row);
 
diff --git a/cpp/arcticdb/version/version_core.cpp b/cpp/arcticdb/version/version_core.cpp
index b31c957826..53eb7e2835 100644
--- a/cpp/arcticdb/version/version_core.cpp
+++ b/cpp/arcticdb/version/version_core.cpp
@@ -864,6 +864,8 @@ void copy_frame_data_to_buffer(const SegmentInMemory& destination, size_t target
         return;
     }
     auto& src_column = source.column(static_cast<position_t>(source_index));
+    schema::check<ErrorCode::E_UNSUPPORTED_COLUMN_TYPE>(!src_column.opt_sparse_map().has_value(),
+                                                   "QueryBuilder not yet supported with sparse data");
     auto& dst_column = destination.column(static_cast<position_t>(target_index));
     auto& buffer = dst_column.data().buffer();
     auto dst_rawtype_size = sizeof_datatype(dst_column.type());


Test patch:
diff --git a/cpp/arcticdb/column_store/test/test_memory_segment.cpp b/cpp/arcticdb/column_store/test/test_memory_segment.cpp
index 3416242e75..d49e02168c 100644
--- a/cpp/arcticdb/column_store/test/test_memory_segment.cpp
+++ b/cpp/arcticdb/column_store/test/test_memory_segment.cpp
@@ -481,7 +481,7 @@ TEST(MemSegment, Filter) {
         filter_bitset.set_bit(retained_row);
     }
 
-    auto filtered_seg = seg.filter(filter_bitset);
+    auto filtered_seg = seg.filter(std::move(filter_bitset));
 
     for (auto&& [idx, row]: folly::enumerate(filtered_seg)) {
         ASSERT_EQ(static_cast<int64_t>(retained_rows[idx]), row.scalar_at<int64_t>(0));
diff --git a/python/tests/unit/arcticdb/version_store/test_query_builder.py b/python/tests/unit/arcticdb/version_store/test_query_builder.py
index bab3019b06..2e505c71e9 100644
--- a/python/tests/unit/arcticdb/version_store/test_query_builder.py
+++ b/python/tests/unit/arcticdb/version_store/test_query_builder.py
@@ -11,6 +11,7 @@
 
 from arcticdb.version_store.processing import QueryBuilder
 from arcticdb.util.test import assert_frame_equal
+from arcticdb_ext.exceptions import SchemaException
 
 
 def test_reuse_querybuilder(lmdb_version_store_tiny_segment):
@@ -403,3 +404,64 @@ def test_querybuilder_pickling():
     import pickle
 
     assert pickle.loads(pickle.dumps(q)) == q
+
+
+# Remove the following test and replace with more extensive ones once this issue is fixed:
+# https://github.com/man-group/ArcticDB/issues/1404
+def test_query_builder_sparse(lmdb_version_store):
+    lib = lmdb_version_store
+    sym = "test_filter_sparse"
+    df = pd.DataFrame(
+        {
+            "sparse_col": [0.0, np.nan, 0.0],
+            "dense_col": [0.0, 1.0, 2.0]
+         },
+        index=pd.date_range("2024-01-01", periods=3)
+    )
+    lib.write(sym, df, sparsify_floats=True)
+
+    # Filters
+    # These 2 filters exercise different code paths
+    q = QueryBuilder()
+    q = q[q["sparse_col"].isnull()]
+    with pytest.raises(SchemaException):
+        lib.read(sym, query_builder=q)
+    q = QueryBuilder()
+    q = q[q["sparse_col"] == np.float64(0.0)]
+    with pytest.raises(SchemaException):
+        lib.read(sym, query_builder=q)
+
+    # Projections
+    q = QueryBuilder().apply("projected_col", q["sparse_col"] + 1)
+    with pytest.raises(SchemaException):
+        lib.read(sym, query_builder=q)
+
+    # Groupbys
+    q = QueryBuilder().groupby("sparse_col").agg({"dense_col": "sum"})
+    with pytest.raises(SchemaException):
+        lib.read(sym, query_builder=q)
+
+    # Aggregations
+    q = QueryBuilder().groupby("dense_col").agg({"sparse_col": "sum"})
+    with pytest.raises(SchemaException):
+        lib.read(sym, query_builder=q)
+
+    # Date range
+    q = QueryBuilder().date_range((pd.Timestamp("2024-01-01"), pd.Timestamp("2024-01-02")))
+    with pytest.raises(SchemaException):
+        lib.read(sym, query_builder=q)
+
+    # Head
+    q = QueryBuilder()._head(2)
+    with pytest.raises(SchemaException):
+        lib.read(sym, query_builder=q)
+
+    # Tail
+    q = QueryBuilder()._tail(2)
+    with pytest.raises(SchemaException):
+        lib.read(sym, query_builder=q)
+
+    # Row range
+    q = QueryBuilder()._row_range((1, 2))
+    with pytest.raises(SchemaException):
+        lib.read(sym, query_builder=q)