Support recursive normalizer in write in v2 API

Author: phoebusm

build_tooling/format.py

@@ -31,30 +31,38 @@ def install_tools():
     return black or clang
 
 
-def lint_python(in_place: bool):
+def lint_python(in_place: bool, specific_file: str):
     try:
         import black
         assert black.__version__ == black_version
     except ImportError:
         raise RuntimeError("black not installed. Run this script with --install-tools then try again")
 
+    if specific_file:
+        path = specific_file
+    else:
+        path = "python/"
+
     if in_place:
-        return subprocess.run(["black", "-l", "120", "python/"]).returncode
+        return subprocess.run(["black", "-l", "120", path]).returncode
     else:
-        return subprocess.run(["black", "-l", "120", "--check", "python/"]).returncode
+        return subprocess.run(["black", "-l", "120", "--check", path]).returncode
 
 
-def lint_cpp(in_place: bool):
+def lint_cpp(in_place: bool, specific_file: str):
     try:
         import clang_format
     except ImportError:
         raise RuntimeError("clang-format not installed. Run this script with --install-tools then try again")
 
     files = []
-    root = pathlib.Path("cpp", "arcticdb")
-    for e in ("*.cpp", "*.hpp"):
-        for f in root.rglob(e):
-            files.append(str(f))
+    if specific_file:
+        files.append(specific_file)
+    else:
+        root = pathlib.Path("cpp", "arcticdb")
+        for e in ("*.cpp", "*.hpp"):
+            for f in root.rglob(e):
+                files.append(str(f))
 
     args = ["clang-format"]
     if in_place:
@@ -69,11 +77,11 @@ def lint_cpp(in_place: bool):
     return subprocess.run(args).returncode
 
 
-def main(type: str, in_place: bool):
+def main(type: str, in_place: bool, specific_file: str):
     if type == "python":
-        return lint_python(in_place)
+        return lint_python(in_place, specific_file)
     elif type == "cpp":
-        return lint_cpp(in_place)
+        return lint_cpp(in_place, specific_file)
     else:
         return lint_python(in_place) or lint_cpp(in_place)
 
@@ -105,6 +113,11 @@ def main(type: str, in_place: bool):
         action='store_true',
         help="Apply linting rules to your working copy. Changes files."
     )
+    parser.add_argument(
+        "-f", 
+        "--file", 
+        help="Apply linting rules to a specific file."
+    )
     args = parser.parse_args()
 
     if args.install_tools:
@@ -118,10 +131,13 @@ def main(type: str, in_place: bool):
         raise RuntimeError("Must specify --type")
     if args.type not in ("python", "cpp", "all"):
         raise RuntimeError("Invalid --type")
+    if args.type == "all" and args.file:
+        raise RuntimeError("Cannot specify file type when specifying a single file for formatting")
 
     return_code = main(
         type=args.type,
         in_place=args.in_place,
+        specific_file=args.file,
     )
 
     sys.exit(return_code)

python/arcticdb/adapters/arctic_library_adapter.py

@@ -24,7 +24,6 @@ def set_library_options(
     write_options = lib_desc.version.write_options
 
     write_options.dynamic_strings = True
-    write_options.recursive_normalizers = True
     write_options.use_tombstones = True
     write_options.fast_tombstone_all = True
     lib_desc.version.symbol_list = True
@@ -38,6 +37,7 @@ def set_library_options(
     write_options.de_duplication = options.dedup
     write_options.segment_row_size = options.rows_per_segment
     write_options.column_group_size = options.columns_per_segment
+    write_options.recursive_normalizers = options.recursive_normalizers
 
     lib_desc.version.encoding_version = (
         options.encoding_version if options.encoding_version is not None else DEFAULT_ENCODING_VERSION

python/arcticdb/options.py

@@ -43,6 +43,7 @@ def __init__(
         rows_per_segment: int = 100_000,
         columns_per_segment: int = 127,
         encoding_version: Optional[EncodingVersion] = None,
+        recursive_normalizers: bool = True,
     ):
         """
         Parameters
@@ -124,12 +125,19 @@ def __init__(
         encoding_version: Optional[EncodingVersion], default None
             The encoding version to use when writing data to storage.
             v2 is faster, but still experimental, so use with caution.
+
+        recursive_normalizers: bool, default True
+            Whether to recursively normalize nested data structures when writing sequence-like or dict-like data.
+            The data structure can be nested or a mix of lists and dictionaries.
+            Note: If the leaf nodes cannot be natively normalized and must be written using write_pickle, those leaf nodes
+            will be pickled, resulting in the overall data being only partially normalized and partially pickled.
         """
         self.dynamic_schema = dynamic_schema
         self.dedup = dedup
         self.rows_per_segment = rows_per_segment
         self.columns_per_segment = columns_per_segment
         self.encoding_version = encoding_version
+        self.recursive_normalizers = recursive_normalizers
 
     def __eq__(self, right):
         return (
@@ -138,13 +146,15 @@ def __eq__(self, right):
             and self.rows_per_segment == right.rows_per_segment
             and self.columns_per_segment == right.columns_per_segment
             and self.encoding_version == right.encoding_version
+            and self.recursive_normalizers == right.recursive_normalizers
         )
 
     def __repr__(self):
         return (
             f"LibraryOptions(dynamic_schema={self.dynamic_schema}, dedup={self.dedup},"
             f" rows_per_segment={self.rows_per_segment}, columns_per_segment={self.columns_per_segment},"
-            f" encoding_version={self.encoding_version if self.encoding_version is not None else 'Default'})"
+            f" encoding_version={self.encoding_version if self.encoding_version is not None else 'Default'},"
+            f" recursive_normalizers={self.recursive_normalizers})"
         )
 
 

python/arcticdb/version_store/_store.py

@@ -612,6 +612,15 @@ def stage(
             log.warning("The data could not be normalized to an ArcticDB format and has not been written")
             return None
 
+    def _is_recursive_normalizers_enabled(self, **kwargs):
+        return resolve_defaults(
+            "recursive_normalizers",
+            self._lib_cfg.lib_desc.version.write_options,
+            global_default=False,
+            uppercase=False,
+            **kwargs,
+        )
+
     def write(
         self,
         symbol: str,
@@ -698,11 +707,9 @@ def write(
         prune_previous_version = resolve_defaults(
             "prune_previous_version", proto_cfg, global_default=False, existing_value=prune_previous_version, **kwargs
         )
-        recursive_normalizers = resolve_defaults(
-            "recursive_normalizers", proto_cfg, global_default=False, uppercase=False, **kwargs
-        )
         parallel = resolve_defaults("parallel", proto_cfg, global_default=False, uppercase=False, **kwargs)
         incomplete = resolve_defaults("incomplete", proto_cfg, global_default=False, uppercase=False, **kwargs)
+        recursive_normalizers = self._is_recursive_normalizers_enabled(**kwargs)
 
         # TODO remove me when dynamic strings is the default everywhere
         if parallel:
@@ -3147,12 +3154,8 @@ def will_item_be_pickled(self, item, recursive_normalizers: Optional[bool] = Non
         )
         if result and log_warning_message:
             proto_cfg = self._lib_cfg.lib_desc.version.write_options
-            resolved_recursive_normalizers = resolve_defaults(
-                "recursive_normalizers",
-                proto_cfg,
-                global_default=False,
-                uppercase=False,
-                **{"recursive_normalizers": recursive_normalizers},
+            resolved_recursive_normalizers = self._is_recursive_normalizers_enabled(
+                **{"recursive_normalizers": recursive_normalizers}
             )
             warning_msg = ""
             is_recursive_normalize_preferred, _, _ = self._try_flatten(item, "")

python/arcticdb/version_store/library.py

@@ -944,6 +944,7 @@ def write(
         staged=False,
         validate_index=True,
         index_column: Optional[str] = None,
+        recursive_normalizers: bool = None,
     ) -> VersionedItem:
         """
         Write ``data`` to the specified ``symbol``. If ``symbol`` already exists then a new version will be created to
@@ -992,6 +993,10 @@ def write(
         index_column: Optional[str], default=None
             Optional specification of timeseries index column if data is an Arrow table. Ignored if data is not an Arrow
             table.
+        recursive_normalizers: bool, default None
+            Whether to recursively normalize nested data structures when writing sequence-like or dict-like data.
+            If None, falls back to the corresponding setting in the library configuration.
+            The data structure can be nested or a mix of lists and dictionaries.
 
         Returns
         -------
@@ -1020,7 +1025,9 @@ def write(
         >>> w = adb.WritePayload("symbol", df, metadata={'the': 'metadata'})
         >>> lib.write(*w, staged=True)
         """
-        if not self._allowed_input_type(data):
+        if not self._nvs._is_recursive_normalizers_enabled(
+            **{"recursive_normalizers": recursive_normalizers}
+        ) and not self._allowed_input_type(data):
             raise ArcticUnsupportedDataTypeException(
                 "data is of a type that cannot be normalized. Consider using "
                 f"write_pickle instead. type(data)=[{type(data)}]"
@@ -1037,10 +1044,17 @@ def write(
             index_column=index_column,
             norm_failure_options_msg="Using write_pickle will allow the object to be written. However, many operations "
             "(such as date_range filtering and column selection) will not work on pickled data.",
+            recursive_normalizers=recursive_normalizers,
         )
 
     def write_pickle(
-        self, symbol: str, data: Any, metadata: Any = None, prune_previous_versions: bool = False, staged=False
+        self,
+        symbol: str,
+        data: Any,
+        metadata: Any = None,
+        prune_previous_versions: bool = False,
+        staged=False,
+        recursive_normalizers: bool = None,
     ) -> VersionedItem:
         """
         See `write`. This method differs from `write` only in that ``data`` can be of any type that is serialisable via
@@ -1062,6 +1076,10 @@ def write_pickle(
             See documentation on `write`.
         staged
             See documentation on `write`.
+        recursive_normalizers: bool, default None
+            See documentation on `write`.
+            If the leaf nodes cannot be natively normalized, they will be pickled,
+            resulting in the overall data being recursively normalized and partially pickled.
 
         Returns
         -------
@@ -1086,6 +1104,7 @@ def write_pickle(
             prune_previous_version=prune_previous_versions,
             pickle_on_failure=True,
             parallel=staged,
+            recursive_normalizers=recursive_normalizers,
         )
 
     @staticmethod

python/tests/conftest.py

@@ -655,12 +655,25 @@ def arctic_client_lmdb(request, encoding_version) -> Arctic:
     return ac
 
 
+@pytest.fixture
+def arctic_client_s3(request) -> Arctic:
+    storage_fixture: StorageFixture = request.getfixturevalue("s3_storage")
+    ac = storage_fixture.create_arctic()
+    return ac
+
+
 @pytest.fixture
 def arctic_library(arctic_client, lib_name) -> Generator[Library, None, None]:
     yield arctic_client.create_library(lib_name)
     arctic_client.delete_library(lib_name)
 
 
+@pytest.fixture
+def arctic_library_s3(arctic_client_s3, lib_name) -> Generator[Library, None, None]:
+    yield arctic_client_s3.create_library(lib_name)
+    arctic_client_s3.delete_library(lib_name)
+
+
 @pytest.fixture
 def arctic_library_dynamic(arctic_client, lib_name) -> Generator[Library, None, None]:
     lib_opts = LibraryOptions(dynamic_schema=True)

python/tests/integration/arcticdb/test_arctic.py

@@ -824,6 +824,22 @@ def test_s3_repr(s3_storage: S3Bucket, one_col_df, lib_name):
     assert written_vi.host == config
 
 
+def test_default_library_config(mem_storage, lib_name):
+    ac = mem_storage.create_arctic()
+    lib = ac.create_library(lib_name)
+    assert lib._nvs.lib_cfg().lib_desc.version.symbol_list == True
+
+    write_options = lib._nvs.lib_cfg().lib_desc.version.write_options
+    assert write_options.dynamic_strings == True
+    assert write_options.recursive_normalizers == True
+    assert write_options.use_tombstones == True
+    assert write_options.fast_tombstone_all == True
+    assert write_options.prune_previous_version == False
+    assert write_options.pickle_on_failure == False
+    assert write_options.snapshot_dedup == False
+    assert write_options.delayed_deletes == False
+
+
 class A:
     """A dummy user defined type that requires pickling to serialize."""
 

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

@@ -20,6 +20,8 @@
     UnsupportedKeyInDictionary,
     ArcticException as ArcticNativeException,
 )
+from arcticdb.version_store.library import ArcticUnsupportedDataTypeException
+import arcticdb.toolbox.query_stats as qs
 from arcticdb_ext.storage import KeyType
 from arcticdb_ext.version_store import NoSuchVersionException
 import arcticdb_ext.stream as adb_stream
@@ -57,6 +59,40 @@ def assert_vit_equals_except_data(left, right):
     assert left.timestamp == right.timestamp
 
 
+@pytest.mark.parametrize("recursive_normalizers", (True, False, None))
+def test_v2_api(arctic_library_s3, sym, recursive_normalizers, clear_query_stats):
+    lib = arctic_library_s3
+    data = {"a": np.arange(5), "b": pd.DataFrame({"col": [1, 2, 3]})}
+    if recursive_normalizers is None or recursive_normalizers is True:
+        with qs.query_stats():
+            lib.write(sym, data, recursive_normalizers=recursive_normalizers)
+        stats = qs.get_query_stats()
+        assert "MULTI_KEY" in stats["storage_operations"]["S3_PutObject"].keys()
+    else:
+        with pytest.raises(ArcticUnsupportedDataTypeException) as e:
+            lib.write(sym, data, recursive_normalizers=recursive_normalizers)
+
+
+@pytest.mark.parametrize("recursive_normalizers", (True, False, None))
+def test_v2_api_pickle(arctic_library_s3, sym, recursive_normalizers, clear_query_stats):
+    lib = arctic_library_s3
+    data = (
+        {
+            "a": [1, 2, 3],
+            "b": {"c": np.arange(24)},
+            "d": [AlmostAListNormalizer()],  # A random item that will be pickled
+        },
+    )
+    with qs.query_stats():
+        lib.write_pickle(sym, data, recursive_normalizers=recursive_normalizers)
+    keys = qs.get_query_stats()["storage_operations"]["S3_PutObject"].keys()
+    if recursive_normalizers is None or recursive_normalizers is True:
+        assert "MULTI_KEY" in keys
+    else:
+        assert "MULTI_KEY" not in keys
+        assert lib._nvs.is_symbol_pickled(sym) == True
+
+
 @pytest.mark.parametrize("read", (lambda lib, sym: lib.batch_read([sym])[sym], lambda lib, sym: lib.read(sym)))
 @pytest.mark.storage
 def test_recursively_written_data(basic_store, read):