Merge branch 'master' into vasil.pashov/nans-bugfix

Author: vasil-pashov

cpp/arcticdb/version/version_core.cpp

@@ -500,21 +500,25 @@ VersionedItem update_impl(
 folly::Future<ReadVersionOutput> read_multi_key(
     const std::shared_ptr<Store>& store,
     const SegmentInMemory& index_key_seg,
-    std::any& handler_data) {
+    std::any& handler_data,
+    AtomKey&& key
+    ) {
     std::vector<AtomKey> keys;
+    keys.reserve(index_key_seg.row_count());
     for (size_t idx = 0; idx < index_key_seg.row_count(); idx++) {
-        keys.emplace_back(stream::read_key_row(index_key_seg, static_cast<ssize_t>(idx)));
+        keys.emplace_back(read_key_row(index_key_seg, static_cast<ssize_t>(idx)));
     }
 
     AtomKey dup{keys[0]};
     VersionedItem versioned_item{std::move(dup)};
     TimeseriesDescriptor multi_key_desc{index_key_seg.index_descriptor()};
 
     return read_frame_for_version(store, versioned_item, std::make_shared<ReadQuery>(), ReadOptions{}, handler_data)
-    .thenValue([multi_key_desc=std::move(multi_key_desc), keys=std::move(keys)](ReadVersionOutput&& read_version_output) mutable {
+    .thenValue([multi_key_desc=std::move(multi_key_desc), keys=std::move(keys), key=std::move(key)](ReadVersionOutput&& read_version_output) mutable {
         multi_key_desc.mutable_proto().mutable_normalization()->CopyFrom(read_version_output.frame_and_descriptor_.desc_.proto().normalization());
         read_version_output.frame_and_descriptor_.desc_ = std::move(multi_key_desc);
         read_version_output.frame_and_descriptor_.keys_ = std::move(keys);
+        read_version_output.versioned_item_ = VersionedItem(std::move(key));
         return std::move(read_version_output);
     });
 }
@@ -2046,7 +2050,7 @@ folly::Future<ReadVersionOutput> read_frame_for_version(
 
     if(pipeline_context->multi_key_) {
         check_multi_key_is_not_index_only(*pipeline_context, *read_query);
-        return read_multi_key(store, *pipeline_context->multi_key_, handler_data);
+        return read_multi_key(store, *pipeline_context->multi_key_, handler_data, std::move(res_versioned_item.key_));
     }
 
     if(opt_false(read_options.incompletes())) {
@@ -2067,7 +2071,7 @@ folly::Future<ReadVersionOutput> read_frame_for_version(
     ARCTICDB_DEBUG(log::version(), "Fetching data to frame");
 
     DecodePathData shared_data;
-    return version_store::do_direct_read_or_process(store, read_query, read_options, pipeline_context, shared_data, handler_data)
+    return do_direct_read_or_process(store, read_query, read_options, pipeline_context, shared_data, handler_data)
     .thenValue([res_versioned_item, pipeline_context, read_options, &handler_data, read_query, shared_data](auto&& frame) mutable {
         ARCTICDB_DEBUG(log::version(), "Reduce and fix columns");
         return reduce_and_fix_columns(pipeline_context, frame, read_options, handler_data)

cpp/arcticdb/version/version_store_api.cpp

@@ -571,11 +571,16 @@ VersionedItem PythonVersionStore::write_versioned_composite_data(
     ARCTICDB_DEBUG(log::version(), "write_versioned_composite_data for stream_id: {} , version_id = {}", stream_id, version_id);
     // TODO: Assuming each sub key is always going to have the same version attached to it.
     std::vector<VersionId> version_ids;
+    version_ids.reserve(sub_keys.size());
+
     std::vector<py::object> user_metas;
+    user_metas.reserve(sub_keys.size());
+
+    std::vector<std::shared_ptr<DeDupMap>> de_dup_maps;
+    de_dup_maps.reserve(sub_keys.size());
 
     auto write_options = get_write_options();
     auto de_dup_map = get_de_dup_map(stream_id, maybe_prev, write_options);
-    std::vector<std::shared_ptr<DeDupMap>> de_dup_maps;
     for (auto i = 0u; i < sub_keys.size(); ++i) {
         version_ids.emplace_back(version_id);
         user_metas.emplace_back(py::none());

cpp/arcticdb/version/version_utils.cpp

@@ -20,6 +20,47 @@ using namespace arcticdb::entity;
 using namespace arcticdb::stream;
 
 
+VariantKey write_multi_index_entry(
+    std::shared_ptr<StreamSink> store,
+    std::vector<AtomKey> &keys,
+    const StreamId &stream_id,
+    const py::object &metastruct,
+    const py::object &user_meta,
+    VersionId version_id
+) {
+    ARCTICDB_SAMPLE(WriteJournalEntry, 0)
+    ARCTICDB_DEBUG(log::version(), "Version map writing multi key");
+    folly::Future<VariantKey> multi_key_fut = folly::Future<VariantKey>::makeEmpty();
+
+    IndexAggregator<RowCountIndex> multi_index_agg(stream_id, [&multi_key_fut, &store, version_id, stream_id](auto &&segment) {
+        multi_key_fut = store->write(KeyType::MULTI_KEY,
+                                     version_id,  // version_id
+                                     stream_id,
+                                     NumericIndex{0},  // start_index
+                                     NumericIndex{0},  // end_index
+                                     std::forward<decltype(segment)>(segment)).wait();
+    });
+
+    for (auto &key : keys) {
+        multi_index_agg.add_key(to_atom(key));
+    }
+    TimeseriesDescriptor timeseries_descriptor;
+
+    if (!metastruct.is_none()) {
+        arcticdb::proto::descriptors::UserDefinedMetadata multi_key_proto;
+        python_util::pb_from_python(metastruct, multi_key_proto);
+        timeseries_descriptor.set_multi_key_metadata(std::move(multi_key_proto));
+    }
+    if (!user_meta.is_none()) {
+        arcticdb::proto::descriptors::UserDefinedMetadata user_meta_proto;
+        python_util::pb_from_python(user_meta, user_meta_proto);
+        timeseries_descriptor.set_user_metadata(std::move(user_meta_proto));
+    }
+    multi_index_agg.set_timeseries_descriptor(timeseries_descriptor);
+    multi_index_agg.commit();
+    return multi_key_fut.wait().value();
+}
+
 std::unordered_map<StreamId, size_t> get_num_version_entries(const std::shared_ptr<Store>& store, size_t batch_size)  {
     std::unordered_map<StreamId, size_t> output;
     size_t max_blocks = ConfigsMap::instance()->get_int("VersionMap.MaxVersionBlocks", 5);

cpp/arcticdb/version/version_utils.hpp

@@ -22,46 +22,14 @@
 
 namespace arcticdb {
 
-inline entity::VariantKey write_multi_index_entry(
+VariantKey write_multi_index_entry(
     std::shared_ptr<StreamSink> store,
     std::vector<AtomKey> &keys,
     const StreamId &stream_id,
     const py::object &metastruct,
     const py::object &user_meta,
     VersionId version_id
-) {
-    ARCTICDB_SAMPLE(WriteJournalEntry, 0)
-    ARCTICDB_DEBUG(log::version(), "Version map writing multi key");
-    folly::Future<VariantKey> multi_key_fut = folly::Future<VariantKey>::makeEmpty();
-
-    IndexAggregator<RowCountIndex> multi_index_agg(stream_id, [&multi_key_fut, &store, version_id, stream_id](auto &&segment) {
-        multi_key_fut = store->write(KeyType::MULTI_KEY,
-                                     version_id,  // version_id
-                                     stream_id,
-                                     NumericIndex{0},  // start_index
-                                     NumericIndex{0},  // end_index
-                                     std::forward<decltype(segment)>(segment)).wait();
-    });
-
-    for (auto &key : keys) {
-        multi_index_agg.add_key(to_atom(key));
-    }
-    TimeseriesDescriptor timeseries_descriptor;
-
-    if (!metastruct.is_none()) {
-        arcticdb::proto::descriptors::UserDefinedMetadata multi_key_proto;
-        python_util::pb_from_python(metastruct, multi_key_proto);
-        timeseries_descriptor.set_multi_key_metadata(std::move(multi_key_proto));
-    }
-    if (!user_meta.is_none()) {
-        arcticdb::proto::descriptors::UserDefinedMetadata user_meta_proto;
-        python_util::pb_from_python(user_meta, user_meta_proto);
-        timeseries_descriptor.set_user_metadata(std::move(user_meta_proto));
-    }
-    multi_index_agg.set_timeseries_descriptor(timeseries_descriptor);
-    multi_index_agg.commit();
-    return multi_key_fut.wait().value();
-}
+);
 
 inline std::optional<AtomKey> read_segment_with_keys(
     const SegmentInMemory &seg,

python/arcticdb/util/test.py

@@ -901,3 +901,19 @@ def generic_resample_test(
         assert_dfs_approximate(expected, received)
     else:
         assert_frame_equal(expected, received, check_dtype=False)
+
+def equals(x, y):
+    if isinstance(x, tuple) or isinstance(x, list):
+        assert len(x) == len(y)
+        for vx, vy in zip(x, y):
+            equals(vx, vy)
+    elif isinstance(x, dict):
+        assert isinstance(y, dict)
+        assert set(x.keys()) == set(y.keys())
+        for k in x.keys():
+            equals(x[k], y[k])
+    elif isinstance(x, np.ndarray):
+        assert isinstance(y, np.ndarray)
+        assert np.allclose(x, y)
+    else:
+        assert x == y

python/arcticdb/version_store/_store.py

@@ -433,6 +433,7 @@ def _try_flatten_and_write_composite_object(
                     metadata=metadata,
                     data=None,
                     host=self.env,
+                    timestamp=vit_composite.timestamp
                 )
 
     def _write_options(self):
@@ -2153,7 +2154,6 @@ def _adapt_read_res(self, read_result: ReadResult) -> VersionedItem:
         data = self._normalizer.denormalize(frame_data, read_result.norm)
         if read_result.norm.HasField("custom"):
             data = self._custom_normalizer.denormalize(data, read_result.norm.custom)
-
         return VersionedItem(
             symbol=read_result.version.symbol,
             library=self._library.library_path,

python/tests/compat/arcticdb/test_compatibility.py

@@ -217,4 +217,4 @@ def test_compat_read_incomplete(old_venv_and_arctic_uri, lib_name):
         assert_frame_equal(read_df, df[["float_col", "str_col"]])
 
         read_df = curr.lib._nvs.read(sym, date_range=(pd.Timestamp(2024, 1, 5), pd.Timestamp(2024, 1, 9)), incomplete=True, columns=["float_col", "str_col"]).data
-        assert_frame_equal(read_df, df[["float_col", "str_col"]].iloc[4:9])
+        assert_frame_equal(read_df, df[["float_col", "str_col"]].iloc[4:9])