Skip to content

Accept strings for the mode when finalizing staged data #2232

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions python/arcticdb/version_store/library.py
Original file line number Diff line number Diff line change
Expand Up @@ -1281,17 +1281,17 @@ def delete_staged_data(self, symbol: str):
def finalize_staged_data(
self,
symbol: str,
mode: Optional[StagedDataFinalizeMethod] = StagedDataFinalizeMethod.WRITE,
mode: Optional[Union[StagedDataFinalizeMethod, str]] = StagedDataFinalizeMethod.WRITE,
prune_previous_versions: bool = False,
metadata: Any = None,
validate_index = True,
delete_staged_data_on_failure: bool = False
) -> VersionedItem:
"""
Finalizes staged data, making it available for reads. All staged segments must be ordered and non-overlapping.
``finalize_staged_data`` is less time consuming than ``sort_and_finalize_staged_data``.
``finalize_staged_data`` is less time-consuming than ``sort_and_finalize_staged_data``.

If ``mode`` is ``StagedDataFinalizeMethod.APPEND`` the index of the first row of the new segment must be equal to or greater
If ``mode`` is ``StagedDataFinalizeMethod.APPEND`` or ``append`` the index of the first row of the new segment must be equal to or greater
than the index of the last row in the existing data.

If ``Static Schema`` is used all staged block must have matching schema (same column names, same dtype, same column ordering)
Expand All @@ -1310,9 +1310,9 @@ def finalize_staged_data(
symbol : `str`
Symbol to finalize data for.

mode : `StagedDataFinalizeMethod`, default=StagedDataFinalizeMethod.WRITE
Finalize mode. Valid options are WRITE or APPEND. Write collects the staged data and writes them to a
new version. Append collects the staged data and appends them to the latest version.
mode : Union[`StagedDataFinalizeMethod`, str], default=StagedDataFinalizeMethod.WRITE
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think lowercase write, append will match both Pandas and our QueryBuilder APIs better.

Finalize mode. Valid options are StagedDataFinalizeMethod.WRITE or StagedDataFinalizeMethod.APPEND. Write collects the staged data and writes them to a
new version. Append collects the staged data and appends them to the latest version. Also accepts "write" and "append".
prune_previous_versions: bool, default=False
Removes previous (non-snapshotted) versions from the database.
metadata : Any, default=None
Expand Down Expand Up @@ -1383,9 +1383,12 @@ def finalize_staged_data(
2024-01-03 3
2024-01-04 4
"""
if mode not in [StagedDataFinalizeMethod.APPEND, StagedDataFinalizeMethod.WRITE, "write", "append", None]:
raise ArcticInvalidApiUsageException("mode must be one of StagedDataFinalizeMethod.WRITE, StagedDataFinalizeMethod.APPEND, 'write', 'append'")

return self._nvs.compact_incomplete(
symbol,
append=mode == StagedDataFinalizeMethod.APPEND,
append=mode == StagedDataFinalizeMethod.APPEND or mode == "append",
convert_int_to_float=False,
metadata=metadata,
prune_previous_version=prune_previous_versions,
Expand Down
25 changes: 25 additions & 0 deletions python/tests/integration/arcticdb/test_finalize_staged_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,31 @@ def test_finalize_staged_data_long_scenario(basic_arctic_library):
cachedParts.verify_finalized_data(lib,symbol)


@pytest.mark.parametrize("mode" , [StagedDataFinalizeMethod.WRITE, "write", None])
def test_finalize_staged_data_mode_write(basic_arctic_library, mode):
lib = basic_arctic_library
symbol = "symbol"
df_initial = sample_dataframe('2020-1-1', [1,2,3], [4, 5, 6])
df_staged = sample_dataframe('2020-1-4', [7, 8, 9])
lib.write(symbol, df_initial)
lib.write(symbol, df_staged, staged=True)
assert_frame_equal(lib.read(symbol).data, df_initial)

lib.finalize_staged_data(symbol="symbol", mode=mode)
assert_frame_equal(lib.read(symbol).data, df_staged)


@pytest.mark.parametrize("mode" , [StagedDataFinalizeMethod.APPEND, "append"])
def test_finalize_staged_data_mode_append(basic_arctic_library, mode):
lib = basic_arctic_library
symbol = "symbol"
df_initial = sample_dataframe('2020-1-1', [1,2,3], [4, 5, 6])
df_staged = sample_dataframe('2020-1-4', [7, 8, 9], [10, 11, 12])
lib.write(symbol, df_initial)
lib.write(symbol, df_staged, staged=True)
assert_frame_equal(lib.read(symbol).data, df_initial)

lib.finalize_staged_data(symbol="symbol", mode=mode)
expected = pd.concat([df_initial, df_staged])
assert_frame_equal(lib.read(symbol).data, expected)

37 changes: 37 additions & 0 deletions python/tests/unit/arcticdb/version_store/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,12 @@
import time

from pandas import Timestamp
from unittest.mock import MagicMock
import pytest

from arcticdb.exceptions import NoSuchVersionException, NoDataFoundException
from arcticdb.util.test import distinct_timestamps
from arcticdb.version_store.library import StagedDataFinalizeMethod, ArcticInvalidApiUsageException


def test_read_descriptor(lmdb_version_store, one_col_df):
Expand Down Expand Up @@ -89,3 +91,38 @@ def test_get_num_rows_pickled(lmdb_version_store):
symbol = "test_get_num_rows_pickled"
lmdb_version_store.write(symbol, 1)
assert lmdb_version_store.get_num_rows(symbol) is None


@pytest.mark.parametrize(
"input_mode, expected_append",
[
("write", False),
("append", True),
(None, False),
(StagedDataFinalizeMethod.APPEND, True),
(StagedDataFinalizeMethod.WRITE, False),
]
)
def test_finalize_staged_data(arctic_library_lmdb, input_mode, expected_append):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good unit test but I think we should also have some integration tests that check we are actually doing the right thing.

The expected_append==None case should probably be split to its own test, rather than parameters in this one.

symbol = "sym"
arctic_library_lmdb._nvs = MagicMock()

arctic_library_lmdb.finalize_staged_data(symbol, input_mode)

default_args = {
"convert_int_to_float": False,
"metadata": None,
"prune_previous_version": False,
"validate_index": True,
"delete_staged_data_on_failure": False,
}

arctic_library_lmdb._nvs.compact_incomplete.assert_called_once_with(symbol, append=expected_append, **default_args)


@pytest.mark.parametrize("input_mode", ["wRite", "something", 3])
def test_finalize_staged_data_incorrect_args(arctic_library_lmdb, input_mode):
symbol = "sym"
arctic_library_lmdb._nvs = MagicMock()
with pytest.raises(ArcticInvalidApiUsageException):
arctic_library_lmdb.finalize_staged_data(symbol, input_mode)
Loading