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

Issue:
Expose Library method to clean up staged data.
**Is your feature request related to a problem? Please describe.**

I'd like to be able to clean up staged writes instead of finalizing them.

```python
df = pd.DataFrame({'column': [5,6,7]}, index=pd.date_range(start='1/1/2000', periods=3))
lib.write('staged', df, staged=True)
# I don't want to finalize, how do I clean up the staged data?
```

If I don't clean up staged data then:
- I have data that I can't see or remove
- I can't prepare any new staged data without it being mixed into the previous data on finalization

**Describe the solution you'd like**

The method exists but isn't part of the public Library API.
```
library._nvs.version_store.remove_incomplete(<symbol>)
```
Decide on naming and expose.




Discussion / hints:

Fix patch:
diff --git a/cpp/arcticdb/processing/clause.cpp b/cpp/arcticdb/processing/clause.cpp
index 81c5d0f856..da40d6d53e 100644
--- a/cpp/arcticdb/processing/clause.cpp
+++ b/cpp/arcticdb/processing/clause.cpp
@@ -553,6 +553,7 @@ Composite<ProcessingUnit> MergeClause::process(std::shared_ptr<Store> store,
 }
 
 std::optional<std::vector<Composite<ProcessingUnit>>> MergeClause::repartition(
+        ARCTICDB_UNUSED const std::shared_ptr<Store>&,
         std::vector<Composite<ProcessingUnit>> &&comps) const {
     std::vector<Composite<ProcessingUnit>> v;
     v.push_back(merge_composites_shallow(std::move(comps)));
diff --git a/cpp/arcticdb/processing/clause.hpp b/cpp/arcticdb/processing/clause.hpp
index 58b457d6d5..aec06c708f 100644
--- a/cpp/arcticdb/processing/clause.hpp
+++ b/cpp/arcticdb/processing/clause.hpp
@@ -79,8 +79,9 @@ struct IClause {
         }
 
         [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
+                [[maybe_unused]] const std::shared_ptr<Store>& store,
                 [[maybe_unused]] std::vector<Composite<ProcessingUnit>> &&comps) const {
-            return folly::poly_call<2>(*this, std::move(comps));
+            return folly::poly_call<2>(*this, store, std::move(comps));
         }
 
         [[nodiscard]] const ClauseInfo& clause_info() const { return folly::poly_call<3>(*this); };
@@ -120,7 +121,8 @@ struct PassthroughClause {
             ) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
-            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&comps) const {
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
+            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
 
@@ -159,6 +161,7 @@ struct FilterClause {
             ) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
             ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
@@ -205,7 +208,8 @@ struct ProjectClause {
     process(std::shared_ptr<Store> store, Composite<ProcessingUnit> &&p) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
-            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&comps) const {
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
+            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
 
@@ -255,6 +259,7 @@ struct PartitionClause {
     }
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
             ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&c) const {
         auto comps = std::move(c);
         std::unordered_map<size_t, Composite<ProcessingUnit>> partition_map;
@@ -330,8 +335,8 @@ struct AggregationClause {
                                                     Composite<ProcessingUnit> &&p) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
-            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&comps
-            ) const {
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
+            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
 
@@ -364,6 +369,7 @@ struct RemoveColumnPartitioningClause {
                                                     Composite<ProcessingUnit> &&p) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
             ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
@@ -393,7 +399,8 @@ struct SplitClause {
                                                     Composite<ProcessingUnit> &&procs) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
-            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&comps) const {
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
+            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
 
@@ -421,7 +428,8 @@ struct SortClause {
                                                     Composite<ProcessingUnit> &&p) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
-            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&comps) const {
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
+            ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
 
@@ -462,6 +470,7 @@ struct MergeClause {
                                                     Composite<ProcessingUnit> &&p) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
+            const std::shared_ptr<Store>& store,
             std::vector<Composite<ProcessingUnit>> &&comps) const;
 
     [[nodiscard]] const ClauseInfo& clause_info() const {
@@ -496,6 +505,7 @@ struct ColumnStatsGenerationClause {
                                                     Composite<ProcessingUnit> &&p) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
             ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
@@ -542,6 +552,7 @@ struct RowRangeClause {
                                                     Composite<ProcessingUnit> &&p) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
             ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
@@ -578,6 +589,7 @@ struct DateRangeClause {
                                                     Composite<ProcessingUnit> &&p) const;
 
     [[nodiscard]] std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
+            ARCTICDB_UNUSED const std::shared_ptr<Store>&,
             ARCTICDB_UNUSED std::vector<Composite<ProcessingUnit>> &&) const {
         return std::nullopt;
     }
diff --git a/cpp/arcticdb/vector_db/vector_db.cpp b/cpp/arcticdb/vector_db/vector_db.cpp
index ea21d72d36..d52670721f 100644
--- a/cpp/arcticdb/vector_db/vector_db.cpp
+++ b/cpp/arcticdb/vector_db/vector_db.cpp
@@ -13,12 +13,63 @@ namespace arcticdb {
             query_vector_(std::move(query_vector)),
             k_(k) {
         clause_info_.modifies_output_descriptor_ = true;
+        clause_info_.requires_repartition_ = true;
         clause_info_.follows_original_columns_order_ = false;
         clause_info_.can_combine_with_column_selection_ = false;
     }
 
-    Composite<ProcessingSegment> TopKClause::process(std::shared_ptr<Store> store,
-                                                     Composite<ProcessingSegment> &&p) const {
+    std::optional<std::vector<Composite<ProcessingUnit>>> TopKClause::repartition(
+            const std::shared_ptr<Store> &store,
+            std::vector<Composite<ProcessingUnit>> &&c) const {
+        auto comps = std::move(c);
+        TopK top_k(k_);
+        for (auto &comp: comps) {
+            comp.broadcast([&store, &top_k](auto &proc) {
+                for (const auto &slice_and_key: proc.data()) {
+                    const std::vector<std::shared_ptr<Column>> &columns = slice_and_key.segment(store).columns();
+                    for (auto &&[idx, col]: folly::enumerate(columns)) {
+                        col->type().visit_tag(
+                                [&top_k, &col = col, &slice_and_key, &store, idx = idx](auto type_desc_tag) {
+                                    using TDT = decltype(type_desc_tag);
+                                    using raw_type = typename TDT::DataTypeTag::raw_type;
+
+                                    if constexpr (is_floating_point_type(TDT::DataTypeTag::data_type)) {
+                                        top_k.try_insert(
+                                                col->template scalar_at<raw_type>(col->last_row()).value(),
+                                                col,
+                                                slice_and_key.segment(store).field(idx).name()
+                                        );
+                                        // In each column, we store the distance from the query vector in the last row.
+                                        // So col->template scalar_at<raw_type>(col->last_row()).value() gets the
+                                        // distance.
+                                        //
+                                        // We have the top-k from each segment and we want the top-k in the entire
+                                        // symbol. We therefore run top-k over the top-k from each segment again to
+                                        // obtain the top-k overall.
+                                    } else {
+                                        internal::raise<ErrorCode::E_INVALID_ARGUMENT>(
+                                                "Vectors should exclusively comprise floats; the Python layer should otherwise raise an error and prevent upsertion of vectors containing non-float components."
+                                        );
+                                    }
+                                });
+                    }
+                }
+            });
+        }
+        SegmentInMemory seg;
+        seg.descriptor().set_index(IndexDescriptor(0, IndexDescriptor::ROWCOUNT));
+        seg.set_row_id(query_vector_.size() - 1 + 1);
+        for (const auto &column: top_k.top_k_) {
+            seg.add_column(scalar_field(DataType::FLOAT64, column.name_), column.column_);
+        }
+        std::vector<Composite<ProcessingUnit>> output;
+        output.emplace_back(ProcessingUnit{std::move(seg)});
+        return output;
+    }
+
+
+    Composite<ProcessingUnit> TopKClause::process(std::shared_ptr<Store> store,
+                                                  Composite<ProcessingUnit> &&p) const {
         auto procs = std::move(p);
         TopK top_k(k_);
         // We expect a vector in each column.
@@ -29,16 +80,17 @@ namespace arcticdb {
         // Let's pretend that only the lp-norms exist. We'll use p=2 for now as a default.
 
         procs.broadcast(
-                [&store, &top_k, lp, this](const ProcessingSegment &proc) {
-                    for (const auto& slice_and_key: proc.data()) {
-                        const std::vector<std::shared_ptr<Column>>& columns = slice_and_key.segment(store).columns();
-                        for (auto&& [idx, col]: folly::enumerate(columns)) {
+                [&store, &top_k, lp, this](const ProcessingUnit &proc) {
+                    for (const auto &slice_and_key: proc.data()) {
+                        const std::vector<std::shared_ptr<Column>> &columns = slice_and_key.segment(store).columns();
+                        for (auto &&[idx, col]: folly::enumerate(columns)) {
                             internal::check<ErrorCode::E_ASSERTION_FAILURE>(
                                     static_cast<long unsigned>(col->row_count()) == this->query_vector_.size(),
                                     "Expected vector of length {}, got vector of length {}.",
                                     col->row_count(),
                                     this->query_vector_.size());
-                            col->type().visit_tag([&top_k, &slice_and_key, &store, &col=col, idx=idx, lp, this](auto type_desc_tag) {
+                            col->type().visit_tag([&top_k, &slice_and_key, &store, &col = col, idx = idx, lp, this](
+                                    auto type_desc_tag) {
                                 using TypeDescriptorTag = decltype(type_desc_tag);
                                 using RawType = typename TypeDescriptorTag::DataTypeTag::raw_type;
 
@@ -61,11 +113,10 @@ namespace arcticdb {
                                             sum_differences_exponentiated,
                                             col,
                                             slice_and_key.segment(store).field(idx).name()));
-                                }
-                                else {
+                                } else {
                                     internal::raise<ErrorCode::E_INVALID_ARGUMENT>(
                                             "Vectors should exclusively comprise floats; the Python layer should otherwise raise an error and prevent upsertion of vectors containing non-float components."
-                                            );
+                                    );
                                 }
                             });
                         }
@@ -78,7 +129,7 @@ namespace arcticdb {
         seg.descriptor().set_index(IndexDescriptor(0, IndexDescriptor::ROWCOUNT));
         seg.set_row_id(query_vector_.size() - 1 + 1);
         // -1 because of zero-indexing, +1 because we want to add the distance on at the end.
-            for (const auto &column: top_k.top_k_) {
+        for (const auto &column: top_k.top_k_) {
             // todo: replace (mem)cpying of columns' contents with holding of pointers in top_k_ and addition of pointers to seg
             // this doesn't work presently because the columns can have >1 block.
             column.column_->type().visit_tag([&column, &seg, lp, this](auto type_desc_tag) {
@@ -86,7 +137,8 @@ namespace arcticdb {
                 using RawType = typename TypeDescriptorTag::DataTypeTag::raw_type;
 
                 if constexpr (is_floating_point_type(TypeDescriptorTag::DataTypeTag::data_type)) {
-                    auto write_col = std::make_shared<Column>(column.column_->type(), this->query_vector_.size() + 1, true, false);
+                    auto write_col = std::make_shared<Column>(column.column_->type(), this->query_vector_.size() + 1,
+                                                              true, false);
 
                     auto col_data = column.column_->data();
                     auto write_ptr = reinterpret_cast<RawType *>(write_col->ptr());
@@ -98,14 +150,13 @@ namespace arcticdb {
                     *write_ptr = RawType(pow(column.distance_, 1 / static_cast<double>(lp)));
                     write_col->set_row_data(this->query_vector_.size());
                     seg.add_column(scalar_field(DataType::FLOAT64, column.name_), write_col);
-                    }
-                else {
+                } else {
                     internal::raise<ErrorCode::E_INVALID_ARGUMENT>(
                             "Vectors should exclusively comprise floats; the Python layer should otherwise raise an error and prevent upsertion of vectors containing non-float components."
                     );
                 }
             });
         }
-        return Composite{ProcessingSegment{std::move(seg)}};
+        return Composite{ProcessingUnit{std::move(seg)}};
     }
 }
\ No newline at end of file
diff --git a/cpp/arcticdb/vector_db/vector_db.hpp b/cpp/arcticdb/vector_db/vector_db.hpp
index 2094753e3d..bde47f2857 100644
--- a/cpp/arcticdb/vector_db/vector_db.hpp
+++ b/cpp/arcticdb/vector_db/vector_db.hpp
@@ -103,16 +103,20 @@ namespace arcticdb {
 
         TopKClause(std::vector<double> query_vector, uint8_t k);
 
-        [[nodiscard]] Composite<ProcessingSegment> process(
+        [[nodiscard]] Composite<ProcessingUnit> process(
                 std::shared_ptr<Store> store,
-                Composite<ProcessingSegment> &&p
+                Composite<ProcessingUnit> &&p
         ) const;
 
-        [[nodiscard]] std::optional<std::vector<Composite<ProcessingSegment>>> repartition(
-                ARCTICDB_UNUSED std::vector<Composite<ProcessingSegment>> &&) const {
-            return std::nullopt;
+        [[nodiscard]] std::vector<Composite<SliceAndKey>> structure_for_processing(
+                std::vector<SliceAndKey>& slice_and_keys, ARCTICDB_UNUSED size_t) const {
+            return structure_by_column_slice(slice_and_keys);
         }
 
+        std::optional<std::vector<Composite<ProcessingUnit>>> repartition(
+                const std::shared_ptr<Store>& store,
+                std::vector<Composite<ProcessingUnit>> &&c) const;
+
         [[nodiscard]] const ClauseInfo &clause_info() const {
             return clause_info_;
         }
diff --git a/cpp/arcticdb/version/version_core.cpp b/cpp/arcticdb/version/version_core.cpp
index 280bccb935..8f39f664d3 100644
--- a/cpp/arcticdb/version/version_core.cpp
+++ b/cpp/arcticdb/version/version_core.cpp
@@ -409,7 +409,7 @@ Composite<ProcessingUnit> process_remaining_clauses(
         std::vector<std::shared_ptr<Clause>> clauses ) { // pass by copy deliberately as we don't want to modify read_query
     while (!clauses.empty()) {
         if (clauses[0]->clause_info().requires_repartition_) {
-            std::vector<Composite<ProcessingUnit>> repartitioned_procs = clauses[0]->repartition(std::move(procs)).value();
+            std::vector<Composite<ProcessingUnit>> repartitioned_procs = clauses[0]->repartition(store, std::move(procs)).value();
             // Erasing from front of vector not ideal, but they're just shared_ptr and there shouldn't be loads of clauses
             clauses.erase(clauses.begin());
             std::vector<folly::Future<Composite<ProcessingUnit>>> fut_procs;
diff --git a/setup.py b/setup.py
index 1a86fbebbd..a5f725c361 100644
--- a/setup.py
+++ b/setup.py
@@ -26,6 +26,7 @@
 ARCTICDB_BUILD_CPP_TESTS = ARCTICDB_BUILD_CPP_TESTS != "0"
 print(f"ARCTICDB_BUILD_CPP_TESTS={ARCTICDB_BUILD_CPP_TESTS}")
 
+
 def _log_and_run(*cmd, **kwargs):
     print("Running " + " ".join(cmd))
     subprocess.check_call(cmd, **kwargs)


Test patch:
diff --git a/python/tests/unit/arcticdb/version_store/test_parallel.py b/python/tests/unit/arcticdb/version_store/test_parallel.py
index 51db1d95df..94c37e23f5 100644
--- a/python/tests/unit/arcticdb/version_store/test_parallel.py
+++ b/python/tests/unit/arcticdb/version_store/test_parallel.py
@@ -55,6 +55,7 @@ def test_remove_incomplete(lmdb_version_store):
     lib.write(sym3, df1)
     lib.remove_incomplete(sym3)
 
+
 from arcticdb.util._versions import IS_PANDAS_TWO