Merge branch 'master' into remove-mmap-buffer

man-group · Mar 25, 2024 · 8c8da59 · 8c8da59
2 parents 60e3988 + e3453ad
commit 8c8da59
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 16 deletions.
diff --git a/python/arcticdb/version_store/_store.py b/python/arcticdb/version_store/_store.py
@@ -1237,10 +1237,11 @@ def _batch_write_internal(
         )
         norm_failure_options_msg = kwargs.get("norm_failure_options_msg", self.norm_failure_options_msg_write)
 
+        # metadata_vector used to be type-hinted as an Iterable, so handle this case in case anyone is relying on it
         if metadata_vector is None:
-            metadata_itr = itertools.repeat(None)
+            metadata_vector = len(symbols) * [None]
         else:
-            metadata_itr = iter(metadata_vector)
+            metadata_vector = list(metadata_vector)
 
         for idx in range(len(symbols)):
             _handle_categorical_columns(symbols[idx], data_vector[idx], False)
@@ -1249,7 +1250,7 @@ def _batch_write_internal(
             self._try_normalize(
                 symbols[idx],
                 data_vector[idx],
-                next(metadata_itr),
+                metadata_vector[idx],
                 pickle_on_failure,
                 dynamic_strings,
                 None,
@@ -1264,13 +1265,11 @@ def _batch_write_internal(
             symbols, items, norm_metas, udms, prune_previous_version, validate_index, throw_on_error
         )
         write_results = []
-        if metadata_vector is not None:
-            metadata_itr = iter(metadata_vector)
-        for result in cxx_versioned_items:
+        for idx, result in enumerate(cxx_versioned_items):
             if isinstance(result, DataError):
                 write_results.append(result)
             else:
-                write_results.append(self._convert_thin_cxx_item_to_python(result, next(metadata_itr)))
+                write_results.append(self._convert_thin_cxx_item_to_python(result, metadata_vector[idx]))
         return write_results
 
     def _batch_write_metadata_to_versioned_items(
@@ -1399,10 +1398,11 @@ def _batch_append_to_versioned_items(
         )
         dynamic_strings = self._resolve_dynamic_strings(kwargs)
 
+        # metadata_vector used to be type-hinted as an Iterable, so handle this case in case anyone is relying on it
         if metadata_vector is None:
-            metadata_itr = itertools.repeat(None)
+            metadata_vector = len(symbols) * [None]
         else:
-            metadata_itr = iter(metadata_vector)
+            metadata_vector = list(metadata_vector)
 
         for idx in range(len(symbols)):
             _handle_categorical_columns(symbols[idx], data_vector[idx])
@@ -1411,7 +1411,7 @@ def _batch_append_to_versioned_items(
             self._try_normalize(
                 symbols[idx],
                 data_vector[idx],
-                next(metadata_itr),
+                metadata_vector[idx],
                 False,
                 dynamic_strings,
                 None,
@@ -1436,13 +1436,11 @@ def _batch_append_to_versioned_items(
             throw_on_error,
         )
         append_results = []
-        if metadata_vector is not None:
-            metadata_itr = iter(metadata_vector)
-        for result in cxx_versioned_items:
+        for idx, result in enumerate(cxx_versioned_items):
             if isinstance(result, DataError):
                 append_results.append(result)
             else:
-                append_results.append(self._convert_thin_cxx_item_to_python(result, next(metadata_itr)))
+                append_results.append(self._convert_thin_cxx_item_to_python(result, metadata_vector[idx]))
         return append_results
 
     def batch_restore_version(

diff --git a/python/tests/integration/arcticdb/version_store/test_basic_version_store.py b/python/tests/integration/arcticdb/version_store/test_basic_version_store.py
@@ -1342,7 +1342,7 @@ def test_batch_write(basic_store_tombstone_and_sync_passive):
         vitem = lmdb_version_store.read(sym)
         sequential_results[vitem.symbol] = vitem
     lmdb_version_store.batch_write(
-        list(multi_data.keys()), list(multi_data.values()), metadata_vector=list(metadata.values())
+        list(multi_data.keys()), list(multi_data.values()), metadata_vector=(metadata[sym] for sym in multi_data)
     )
     batch_results = lmdb_version_store.batch_read(list(multi_data.keys()))
     assert len(batch_results) == len(sequential_results)
@@ -2122,7 +2122,7 @@ def test_batch_append(basic_store_tombstone, three_col_df):
     metadata = {"sym1": {"key1": "val1"}, "sym2": None, "sym3": None}
 
     lmdb_version_store.batch_write(
-        list(multi_data.keys()), list(multi_data.values()), metadata_vector=list(metadata.values())
+        list(multi_data.keys()), list(multi_data.values()), metadata_vector=(metadata[sym] for sym in multi_data)
     )
 
     multi_append = {"sym1": three_col_df(10), "sym2": three_col_df(11), "sym3": three_col_df(12)}