Merge remote-tracking branch 'upstream/master'

Author: Maxim Morozov

cpp/arcticdb/CMakeLists.txt

@@ -969,6 +969,7 @@ if(${TEST})
             util/test/test_cursor.cpp
             util/test/test_decimal.cpp
             util/test/test_exponential_backoff.cpp
+            util/test/test_folly.cpp
             util/test/test_format_date.cpp
             util/test/test_hash.cpp
             util/test/test_id_transformation.cpp

cpp/arcticdb/async/async_store.hpp

@@ -340,14 +340,14 @@ std::vector<RemoveKeyResultType> remove_keys_sync(std::vector<entity::VariantKey
            RemoveBatchTask{std::move(keys), library_, opts}();
 }
 
-folly::Future<std::vector<VariantKey>> batch_read_compressed(
+std::vector<folly::Future<VariantKey>> batch_read_compressed(
         std::vector<std::pair<entity::VariantKey, ReadContinuation>> &&keys_and_continuations,
         const BatchReadArgs &args) override {
     util::check(!keys_and_continuations.empty(), "Unexpected empty keys/continuation vector in batch_read_compressed");
-    return folly::collect(folly::window(std::move(keys_and_continuations), [this] (auto&& key_and_continuation) {
+    return folly::window(std::move(keys_and_continuations), [this] (auto&& key_and_continuation) {
         auto [key, continuation] = std::forward<decltype(key_and_continuation)>(key_and_continuation);
         return read_and_continue(key, library_, storage::ReadKeyOpts{}, std::move(continuation));
-    }, args.batch_size_)).via(&async::io_executor());
+    }, args.batch_size_);
 }
 
 std::vector<folly::Future<pipelines::SegmentAndSlice>> batch_read_uncompressed(

cpp/arcticdb/async/batch_read_args.hpp

@@ -11,27 +11,11 @@
 
 namespace arcticdb {
 struct BatchReadArgs {
-    // The below enum controls where (IO or CPU thread pool) decoding and data processing tasks are executed.
-    // If IO is used, all work will be done in the IO thread pool which may prevent context switches - this is therefore
-    // suitable for work that is CPU light.
-    enum Scheduler {
-        IO,
-        CPU
-    };
-
-    BatchReadArgs() :
-        batch_size_(ConfigsMap::instance()->get_int("BatchRead.BatchSize", 200)),
-        scheduler_(Scheduler::CPU) {}
+    BatchReadArgs() = default;
 
     explicit BatchReadArgs(size_t batch_size) :
-        batch_size_(batch_size),
-        scheduler_(Scheduler::CPU) {}
-
-    explicit BatchReadArgs(Scheduler scheduler) :
-        batch_size_(ConfigsMap::instance()->get_int("BatchRead.BatchSize", 200)),
-        scheduler_(scheduler) { }
+        batch_size_(batch_size) {}
 
-    size_t batch_size_;
-    Scheduler scheduler_;
+    size_t batch_size_ = ConfigsMap::instance()->get_int("BatchRead.BatchSize", 200);
 };
 }

cpp/arcticdb/pipeline/read_frame.cpp

@@ -900,7 +900,8 @@ folly::Future<SegmentInMemory> fetch_data(
         }
     }
     ARCTICDB_SUBSAMPLE_DEFAULT(DoBatchReadCompressed)
-    return ssource->batch_read_compressed(std::move(keys_and_continuations), BatchReadArgs{})
+    return folly::collect(ssource->batch_read_compressed(std::move(keys_and_continuations), BatchReadArgs{}))
+    .via(&async::io_executor())
     .thenValue([frame](auto&&){ return frame; });
 }
 

cpp/arcticdb/pipeline/write_frame.cpp

@@ -95,8 +95,14 @@ std::tuple<stream::StreamSink::PartialKey, SegmentInMemory, FrameSlice> WriteToS
             auto& tensor = frame_->field_tensors[slice_.absolute_field_col(col)];
             auto opt_error = aggregator_set_data(
                 fd.type(),
-                tensor, agg, abs_col, rows_to_write, offset_in_frame, slice_num_for_column_,
-                regular_slice_size, sparsify_floats_);
+                tensor,
+                agg,
+                abs_col,
+                rows_to_write,
+                offset_in_frame,
+                slice_num_for_column_,
+                regular_slice_size,
+                sparsify_floats_);
             if (opt_error.has_value()) {
                 opt_error->raise(fd.name(), offset_in_frame);
             }

cpp/arcticdb/python/python_to_tensor_frame.cpp

@@ -145,14 +145,10 @@ NativeTensor obj_to_tensor(PyObject *ptr, bool empty_types) {
         // wide type always is 64bits
         val_bytes = 8;
 
-        // If Numpy has type 'O' then get_value_type above will return type 'BYTES'
-        // If there is no value, and we can't deduce a type then leave it that way,
-        // otherwise try to work out whether it was a bytes (string) type or unicode
         if (!is_fixed_string_type(val_type) && element_count > 0) {
             auto none = py::none{};
             auto obj = reinterpret_cast<PyObject **>(arr->data);
-            bool empty = false;
-            bool all_nans = false;
+            bool empty_string_placeholder = false;
             PyObject *sample = *obj;
             PyObject** current_object = obj;
             // Arctic allows both None and NaN to represent a string with no value. We have 3 options:
@@ -163,31 +159,28 @@ NativeTensor obj_to_tensor(PyObject *ptr, bool empty_types) {
             // * In case there is at least one actual string we can sample it and decide the type of the column segment
             //      based on it
             // Note: ValueType::ASCII_DYNAMIC was used when Python 2 was supported. It is no longer supported, and
-            //  we're not expected to enter that branch.
+            // we're not expected to enter that branch.
             if (sample == none.ptr() || is_py_nan(sample)) {
-                empty = true;
-                all_nans = true;
+                empty_string_placeholder = true;
                 util::check(c_style, "Non contiguous columns with first element as None not supported yet.");
                 const auto* end = obj + size;
                 while(current_object < end) {
-
-                    if(*current_object == none.ptr()) {
-                        all_nans = false;
-                    } else if(is_py_nan(*current_object)) {
-                        empty = false;
-                    } else {
-                        all_nans = false;
-                        empty = false;
+                    if(!(is_py_nan(*current_object) || *current_object == none.ptr())) {
+                        empty_string_placeholder = false;
                         break;
                     }
                     ++current_object;
                 }
                 if(current_object != end)
                     sample = *current_object;
             }
-            if (empty && kind == 'O') {
+            // Column full of NaN values is interpreted differently based on the kind. If kind is object "O" the column
+            // is assigned a string type if kind is float "f" the column is assigned a float type. This is done in
+            // order to preserve a legacy behavior of ArcticDB allowing to use both NaN and None as a placeholder for
+            // missing string values.
+            if (empty_string_placeholder && kind == 'O') {
                 val_type = empty_types ? ValueType::EMPTY : ValueType::UTF_DYNAMIC;
-            } else if(all_nans || is_unicode(sample)){
+            } else if(is_unicode(sample)) {
                 val_type = ValueType::UTF_DYNAMIC;
             } else if (PYBIND11_BYTES_CHECK(sample)) {
                 val_type = ValueType::ASCII_DYNAMIC;

cpp/arcticdb/storage/test/in_memory_store.hpp

@@ -53,7 +53,7 @@ class InMemoryStore : public Store {
         throw std::runtime_error("Not implemented for tests");
     }
 
-    folly::Future<std::vector<VariantKey>> batch_read_compressed(
+    std::vector<folly::Future<VariantKey>> batch_read_compressed(
             std::vector<std::pair<entity::VariantKey, ReadContinuation>>&&,
             const BatchReadArgs&) override {
         throw std::runtime_error("Not implemented for tests");

cpp/arcticdb/stream/stream_source.hpp

@@ -55,15 +55,13 @@ struct StreamSource {
 
     using ReadContinuation = folly::Function<VariantKey(storage::KeySegmentPair &&)>;
 
-    virtual folly::Future<std::vector<VariantKey>> batch_read_compressed(
+    virtual std::vector<folly::Future<VariantKey>> batch_read_compressed(
         std::vector<std::pair<entity::VariantKey, ReadContinuation>> &&ks,
-        const BatchReadArgs &args) = 0;
+        const BatchReadArgs& args) = 0;
 
     [[nodiscard]] virtual std::vector<folly::Future<bool>> batch_key_exists(
             const std::vector<entity::VariantKey> &keys) = 0;
 
-    using DecodeContinuation = folly::Function<folly::Unit(SegmentInMemory &&)>;
-
     virtual std::vector<folly::Future<pipelines::SegmentAndSlice>> batch_read_uncompressed(
         std::vector<pipelines::RangesAndKey>&& ranges_and_keys,
         std::shared_ptr<std::unordered_set<std::string>> columns_to_decode) = 0;

cpp/arcticdb/toolbox/storage_mover.hpp

@@ -122,14 +122,18 @@ struct BatchCopier {
                                 ++skipped_;
                             }
                         }
-                        auto collected_kvs = source_store_->batch_read_compressed(std::move(keys_to_copy), BatchReadArgs{}).get();
-                        if (!collected_kvs.empty()) {
+
+                        size_t n_keys = keys_to_copy.size();
+                        auto collected_kvs = folly::collect(source_store_->batch_read_compressed(std::move(keys_to_copy), BatchReadArgs{}))
+                            .via(&async::io_executor())
+                            .get();
+                        if (n_keys > 0) {
                             const size_t bytes_being_copied = std::accumulate(segments_ptr->begin(), segments_ptr->end(), size_t{0}, [] (size_t a, const storage::KeySegmentPair& ks) {
                                 return a + ks.segment().size();
                             });
                             target_store_->batch_write_compressed(*segments_ptr.release()).get();
                             bytes_moved_.fetch_add(bytes_being_copied, std::memory_order_relaxed);
-                            objects_moved_.fetch_add(collected_kvs.size(), std::memory_order_relaxed);
+                            objects_moved_.fetch_add(n_keys, std::memory_order_relaxed);
                         }
                         ++count_;
                         if (count_.compare_exchange_strong(logging_frequency, 0)) {
@@ -177,8 +181,11 @@ struct BatchCopier {
             }
         );
         keys_.clear();
-        auto collected_kvs = source_store_->batch_read_compressed(std::move(keys_to_copy), BatchReadArgs{}).get();
-        if (!collected_kvs.empty()) {
+        size_t n_keys = keys_to_copy.size();
+        auto collected_kvs = folly::collect(source_store_->batch_read_compressed(std::move(keys_to_copy), BatchReadArgs{}))
+            .via(&async::io_executor())
+            .get();
+        if (n_keys > 0) {
             bytes_moved_ += std::accumulate(segments.begin(), segments.end(), size_t{0}, [] (size_t a, const storage::KeySegmentPair& ks) {
                 return a + ks.segment().size();
             });
@@ -523,7 +530,9 @@ class ARCTICDB_VISIBILITY_HIDDEN StorageMover {
             }
 
             total_copied += copied;
-            auto keys = source_store_->batch_read_compressed(std::move(keys_to_copy), BatchReadArgs{}).get();
+            [[maybe_unused]] auto keys = folly::collect(source_store_->batch_read_compressed(std::move(keys_to_copy), BatchReadArgs{}))
+                .via(&async::io_executor())
+                .get();
             std::erase_if(segments, [](const auto& segment) { return variant_key_type(segment.variant_key()) == KeyType::UNDEFINED; });
             util::check(keys.size() == segments.size(), "Keys and segments size mismatch, maybe due to parallel deletes");
             write_futs.push_back(target_store_->batch_write_compressed(std::move(segments)));

cpp/arcticdb/util/test/test_folly.cpp

@@ -0,0 +1,49 @@
+/* Copyright 2025 Man Group Operations Limited
+ *
+ * Use of this software is governed by the Business Source License 1.1 included in the file licenses/BSL.txt.
+ *
+ * As of the Change Date specified in that file, in accordance with the Business Source License, use of this software will be governed by the Apache License, version 2.0.
+ */
+
+#include <folly/futures/Future.h>
+#include <gtest/gtest.h>
+
+/**
+ * We rely on folly::window continuing to process the collection even after one of the futures throws.
+ *
+ * Since this is not clearly tested in Folly itself, we have this regression test in case they break it in future.
+ *
+ * At time of writing, the symbol size calculation APIs `test_symbol_sizes.py` rely on this.
+ */
+TEST(Window, ContinuesOnException) {
+    using namespace folly;
+    std::vector<int> ints(1000);
+    for (int i = 0; i < 1000; i++) {
+        ints.push_back(i);
+    }
+
+    std::vector<Promise<int>> ps(1000);
+
+    auto res = reduce(
+        window(
+            ints,
+            [&ps](int i) {
+                if (i % 4 == 0) {
+                    throw std::runtime_error("exception should not kill process");
+                }
+                return ps[i].getFuture();
+            },
+            2),
+        0,
+        [](int sum, const Try<int>& b) {
+            sum += b.hasException<std::exception>() ? 0 : 1;
+            return sum;
+        });
+
+    for (auto& p : ps) {
+        p.setValue(0);
+    }
+
+    // 100 / 4 = 250 of the futures threw, 750 completed successfully
+    EXPECT_EQ(750, std::move(res).get());
+}

cpp/arcticdb/version/local_versioned_engine.cpp

@@ -1702,15 +1702,35 @@ timestamp LocalVersionedEngine::latest_timestamp(const std::string& symbol) {
 }
 
 struct AtomicKeySizesInfo {
-    size_t count;
+    size_t count{0};
     std::atomic<size_t> compressed_size;
     std::atomic<size_t> uncompressed_size;
 };
 
+using KeySizeCalculators = std::vector<std::pair<VariantKey, stream::StreamSource::ReadContinuation>>;
+
+static void read_ignoring_key_not_found(Store& store, KeySizeCalculators&& calculators) {
+    if (calculators.empty()) {
+        return;
+    }
+
+    auto read_futures = store.batch_read_compressed(std::move(calculators), BatchReadArgs{});
+    std::vector<folly::Future<folly::Unit>> res;
+    res.reserve(read_futures.size());
+    for (auto&& fut : std::move(read_futures)) {
+        // Ignore some exceptions, someone might be deleting while we scan
+        res.push_back(std::move(fut)
+                          .thenValue([](auto&&) {return folly::Unit{};})
+                          .thenError(folly::tag_t<storage::KeyNotFoundException>{}, [](auto&&) { return folly::Unit{}; }));
+    }
+
+    folly::collect(res).get();
+}
+
 std::unordered_map<KeyType, KeySizesInfo> LocalVersionedEngine::scan_object_sizes() {
     std::unordered_map<KeyType, AtomicKeySizesInfo> sizes;
     foreach_key_type([&store=store(), &sizes=sizes](KeyType key_type) {
-        std::vector<std::pair<VariantKey, stream::StreamSource::ReadContinuation>> key_size_calculators;
+        KeySizeCalculators key_size_calculators;
         store->iterate_type(key_type, [&key_size_calculators, &sizes, &key_type](const VariantKey&& k) {
             auto& sizes_info = sizes[key_type];
             ++sizes_info.count;
@@ -1723,10 +1743,7 @@ std::unordered_map<KeyType, KeySizesInfo> LocalVersionedEngine::scan_object_size
             });
         });
 
-        if(!key_size_calculators.empty()) {
-            store->batch_read_compressed(std::move(key_size_calculators), BatchReadArgs{}).get();
-        }
-
+        read_ignoring_key_not_found(*store, std::move(key_size_calculators));
     });
 
     std::unordered_map<KeyType, KeySizesInfo> result;
@@ -1742,10 +1759,10 @@ std::unordered_map<StreamId, std::unordered_map<KeyType, KeySizesInfo>> LocalVer
     auto streams = symbol_list().get_symbols(store());
 
     foreach_key_type([&store=store(), &sizes, &mutex](KeyType key_type) {
-        std::vector<std::pair<VariantKey, stream::StreamSource::ReadContinuation>> keys;
+        KeySizeCalculators key_size_calculators;
 
-        store->iterate_type(key_type, [&keys, &mutex, &sizes, key_type](const VariantKey&& k){
-            keys.emplace_back(std::forward<const VariantKey>(k), [key_type, &sizes, &mutex] (auto&& ks) {
+        store->iterate_type(key_type, [&key_size_calculators, &mutex, &sizes, key_type](const VariantKey&& k){
+            key_size_calculators.emplace_back(std::forward<const VariantKey>(k), [key_type, &sizes, &mutex] (auto&& ks) {
                 auto key_seg = std::forward<decltype(ks)>(ks);
                 auto variant_key = key_seg.variant_key();
                 auto stream_id = variant_key_id(variant_key);
@@ -1766,9 +1783,7 @@ std::unordered_map<StreamId, std::unordered_map<KeyType, KeySizesInfo>> LocalVer
 
         });
 
-        if (!keys.empty()) {
-            store->batch_read_compressed(std::move(keys), BatchReadArgs{}).get();
-        }
+        read_ignoring_key_not_found(*store, std::move(key_size_calculators));
     });
 
     return sizes;

python/tests/hypothesis/arcticdb/test_sort_merge.py

@@ -127,6 +127,7 @@ def get_append_keys(lib, sym):
 @given(df_list=generate_dataframes(COLUMN_DESCRIPTIONS))
 def test_sort_merge_static_schema_write(lmdb_library, df_list):
     lib = lmdb_library
+    lib._nvs.version_store.clear()
     sym = "test_sort_merge_static_schema_write"
     for df in df_list:
         lib.write(sym, df, staged=True, validate_index=False)
@@ -150,6 +151,7 @@ def test_sort_merge_static_schema_write(lmdb_library, df_list):
 @given(df_list=generate_dataframes(COLUMN_DESCRIPTIONS), initial_df=generate_single_dataframe(COLUMN_DESCRIPTIONS, min_size=1, allow_nat_in_index=False))
 def test_sort_merge_static_schema_append(lmdb_library, df_list, initial_df):
     lib = lmdb_library
+    lib._nvs.version_store.clear()
     sym = "test_sort_merge_static_schema_append"
     initial_df.sort_index(inplace=True)
     lib.write(sym, initial_df)
@@ -179,6 +181,7 @@ def test_sort_merge_static_schema_append(lmdb_library, df_list, initial_df):
 @given(df_list=generate_dataframes(COLUMN_DESCRIPTIONS))
 def test_sort_merge_dynamic_schema_write(lmdb_library_dynamic_schema, df_list):
     lib = lmdb_library_dynamic_schema
+    lib._nvs.version_store.clear()
     sym = "test_sort_merge_dynamic_schema_write"
     for df in df_list:
         lib.write(sym, df, staged=True, validate_index=False)
@@ -203,6 +206,7 @@ def test_sort_merge_dynamic_schema_write(lmdb_library_dynamic_schema, df_list):
 @given(df_list=generate_dataframes(COLUMN_DESCRIPTIONS), initial_df=generate_single_dataframe(COLUMN_DESCRIPTIONS, min_size=1, allow_nat_in_index=False))
 def test_sort_merge_dynamic_schema_append(lmdb_library_dynamic_schema, df_list, initial_df):    
     lib = lmdb_library_dynamic_schema
+    lib._nvs.version_store.clear()
     sym = "test_sort_merge_dynamic_schema_append"
     initial_df.sort_index(inplace=True)
     lib.write(sym, initial_df)