Merge when the target symbol does not exist

Author: vasil-pashov

python/tests/unit/arcticdb/version_store/test_merge.py

@@ -14,7 +14,7 @@
 from arcticdb.version_store import VersionedItem
 from arcticdb_ext.storage import KeyType
 import numpy as np
-from arcticdb.exceptions import SchemaException, UserInputException, SortingException
+from arcticdb.exceptions import SchemaException, UserInputException, SortingException, StorageException
 from typing import NamedTuple
 
 
@@ -569,6 +569,66 @@ def test_merge_update_target_is_empty(self, lmdb_library, monkeypatch, merge_met
         assert_vit_equals_except_data(merge_vit, read_vit)
         assert_frame_equal(read_vit.data, expected)
 
+    @pytest.mark.parametrize(
+        "source",
+        (
+            pd.DataFrame([], index=pd.DatetimeIndex([])),
+            pd.DataFrame({"a": []}, index=pd.DatetimeIndex([])),
+            pd.DataFrame({"a": [1]}, index=pd.DatetimeIndex([pd.Timestamp("2024-01-01")])),
+        ),
+    )
+    @pytest.mark.parametrize("upsert", (True, False))
+    def test_update_on_match_target_symbol_does_not_exist(self, lmdb_library, monkeypatch, source, upsert):
+        # Model non-existing target after Library.update
+        # There is an upsert parameter to control whether to create the index. If upsert=False exception is thrown.
+        # Since we're doing a merge on non-existing data, I think it's logical to assume that nothing matches. If
+        # upsert=True, we will create an empty dataframe with the same schema of the source.
+        lib = lmdb_library
+
+        if not upsert:
+            monkeypatch.setattr(lib.__class__, "merge", raise_wrapper(StorageException), raising=False)
+            with pytest.raises(StorageException):
+                lib.merge(
+                    "sym", source, strategy=MergeStrategy(MergeAction.UPDATE, MergeAction.DO_NOTHING), upsert=upsert
+                )
+        else:
+            import datetime
+
+            monkeypatch.setattr(
+                lib.__class__,
+                "merge",
+                lambda *args, **kwargs: VersionedItem(
+                    symbol="sym",
+                    library=lmdb_library.name,
+                    data=None,
+                    version=0,
+                    metadata=None,
+                    host=lmdb_library._nvs.env,
+                    timestamp=pd.Timestamp(datetime.datetime.now()),
+                ),
+                raising=False,
+            )
+            merge_vit = lib.merge(
+                "sym", source, strategy=MergeStrategy(MergeAction.UPDATE, MergeAction.DO_NOTHING), upsert=upsert
+            )
+            expected = pd.DataFrame({"a": []}, index=pd.DatetimeIndex([]))
+            monkeypatch.setattr(
+                lib,
+                "read",
+                lambda *args, **kwargs: VersionedItem(
+                    symbol=merge_vit.symbol,
+                    library=merge_vit.library,
+                    data=expected,
+                    version=merge_vit.version,
+                    metadata=merge_vit.metadata,
+                    host=merge_vit.host,
+                    timestamp=merge_vit.timestamp,
+                ),
+            )
+            read_vit = lib.read("sym")
+            assert_vit_equals_except_data(merge_vit, read_vit)
+            assert_frame_equal(read_vit.data, expected)
+
     # ================================================================================================
     # ================================= TEST INSERT NOT MATCHED ======================================
     # ================================================================================================
@@ -991,6 +1051,63 @@ def test_update_and_insert_with_slicing(self, lmdb_library_factory, monkeypatch,
         assert len(lt.find_keys_for_symbol(KeyType.TABLE_INDEX, "sym")) == 2
         assert len(lt.find_keys_for_symbol(KeyType.VERSION, "sym")) == 2
 
+    @pytest.mark.parametrize(
+        "source",
+        (
+            pd.DataFrame([], index=pd.DatetimeIndex([])),
+            pd.DataFrame({"a": []}, index=pd.DatetimeIndex([])),
+            pd.DataFrame({"a": [1]}, index=pd.DatetimeIndex([pd.Timestamp("2024-01-01")])),
+        ),
+    )
+    @pytest.mark.parametrize("upsert", (True, False))
+    @pytest.mark.parametrize("strategy", (MergeStrategy("update", "insert"), MergeStrategy("do_nothing", "insert")))
+    def test_insert_when_target_symbol_does_not_exist(self, lmdb_library, monkeypatch, source, upsert, strategy):
+        # Model non-existing target after Library.update
+        # There is an upsert parameter to control whether to create the index. If upsert=False exception is thrown.
+        # Since we're doing a merge on non-existing data, I think it's logical to assume that nothing matches. No
+        # updates are performed. The insert operation will insert the data into the newly created target symbol.
+        lib = lmdb_library
+
+        if not upsert:
+            monkeypatch.setattr(lib.__class__, "merge", raise_wrapper(StorageException), raising=False)
+            with pytest.raises(StorageException):
+                lib.merge("sym", source, strategy=strategy, upsert=upsert)
+        else:
+            import datetime
+
+            monkeypatch.setattr(
+                lib.__class__,
+                "merge",
+                lambda *args, **kwargs: VersionedItem(
+                    symbol="sym",
+                    library=lmdb_library.name,
+                    data=None,
+                    version=0,
+                    metadata=None,
+                    host=lmdb_library._nvs.env,
+                    timestamp=pd.Timestamp(datetime.datetime.now()),
+                ),
+                raising=False,
+            )
+            merge_vit = lib.merge("sym", source, strategy=strategy, upsert=upsert)
+            expected = source
+            monkeypatch.setattr(
+                lib,
+                "read",
+                lambda *args, **kwargs: VersionedItem(
+                    symbol=merge_vit.symbol,
+                    library=merge_vit.library,
+                    data=expected,
+                    version=merge_vit.version,
+                    metadata=merge_vit.metadata,
+                    host=merge_vit.host,
+                    timestamp=merge_vit.timestamp,
+                ),
+            )
+            read_vit = lib.read("sym")
+            assert_vit_equals_except_data(merge_vit, read_vit)
+            assert_frame_equal(read_vit.data, expected)
+
     def test_update_and_insert_on_index_and_column(self, lmdb_library, monkeypatch):
         lib = lmdb_library
         target = pd.DataFrame(