Make read_batch take a global arrow_string_format_per_column

to be used as an override if the read request doesn't specify it.

It is also used to set up per column formatting for `read_batch_and_join`

Author: IvoDD

python/arcticdb/version_store/_store.py

@@ -1218,6 +1218,7 @@ def batch_read(
         query_builder: Optional[Union[QueryBuilder, List[QueryBuilder]]] = None,
         columns: Optional[List[List[str]]] = None,
         arrow_string_format_default: Optional[Union[ArrowOutputStringFormat, "pa.DataType"]] = None,
+        arrow_string_format_per_column: Optional[Dict[str, Union[ArrowOutputStringFormat, "pa.DataType"]]] = None,
         per_symbol_arrow_string_format_default: Optional[
             List[Optional[Union[ArrowOutputStringFormat, "pa.DataType"]]]
         ] = None,
@@ -1257,6 +1258,11 @@ def batch_read(
             If using `output_format=EXPERIMENTAL_ARROW` it sets the output format of string columns for arrow.
             See documentation of `ArrowOutputStringFormat` for more information on the different options.
             It serves as the default for the entire batch.
+        arrow_string_format_per_column: Optional[Dict[str, Union[ArrowOutputStringFormat, "pa.DataType"]]], default=None,
+            If using `output_format=EXPERIMENTAL_ARROW` it sets the output format of string columns for arrow.
+            See documentation of `ArrowOutputStringFormat` for more information on the different options.
+            It defines the setting per column. It is applied to all symbols which don't have a
+            `per_symbol_arrow_string_format_per_column` set.
         per_symbol_arrow_string_format_default: Optional[List[Optional[Union[ArrowOutputStringFormat, "pa.DataType"]]]], default=None,
             If using `output_format=EXPERIMENTAL_ARROW` it sets the output format of string columns for arrow.
             See documentation of `ArrowOutputStringFormat` for more information on the different options.
@@ -1292,6 +1298,7 @@ def batch_read(
             query_builder=query_builder,
             throw_on_error=throw_on_error,
             arrow_string_format_default=arrow_string_format_default,
+            arrow_string_format_per_column=arrow_string_format_per_column,
             per_symbol_arrow_string_format_default=per_symbol_arrow_string_format_default,
             per_symbol_arrow_string_format_per_column=per_symbol_arrow_string_format_per_column,
             **kwargs,
@@ -1312,6 +1319,7 @@ def _batch_read_to_versioned_items(
         query_builder,
         throw_on_error,
         arrow_string_format_default,
+        arrow_string_format_per_column,
         per_symbol_arrow_string_format_default,
         per_symbol_arrow_string_format_per_column,
         **kwargs,
@@ -1327,6 +1335,7 @@ def _batch_read_to_versioned_items(
             len(symbols),
             throw_on_error,
             arrow_string_format_default,
+            arrow_string_format_per_column,
             per_symbol_arrow_string_format_default,
             per_symbol_arrow_string_format_per_column,
             **kwargs,
@@ -2136,6 +2145,7 @@ def _get_batch_read_options(
         num_symbols,
         batch_throw_on_error,
         global_arrow_string_format_default=None,
+        global_arrow_string_format_per_column=None,
         per_symbol_arrow_string_format_default=None,
         per_symbol_arrow_string_format_per_column=None,
         **kwargs,
@@ -2159,15 +2169,17 @@ def _get_batch_read_options(
             )
         for idx in range(num_symbols):
             arrow_string_format_default = global_arrow_string_format_default
-            arrow_string_format_per_column = None
+            arrow_string_format_per_column = global_arrow_string_format_per_column
 
             if per_symbol_arrow_string_format_default is not None:
                 arrow_string_format_default = (
                     per_symbol_arrow_string_format_default[idx] or global_arrow_string_format_default
                 )
 
             if per_symbol_arrow_string_format_per_column is not None:
-                arrow_string_format_per_column = per_symbol_arrow_string_format_per_column[idx]
+                arrow_string_format_per_column = (
+                    per_symbol_arrow_string_format_per_column[idx] or global_arrow_string_format_per_column
+                )
 
             read_options_per_symbol.append(
                 self._get_read_options(

python/arcticdb/version_store/library.py

@@ -538,6 +538,8 @@ class LazyDataFrameCollection(QueryBuilder):
     def __init__(
         self,
         lazy_dataframes: List[LazyDataFrame],
+        arrow_string_format_default: Optional[Union[ArrowOutputStringFormat, "pa.DataType"]] = None,
+        arrow_string_format_per_column: Optional[Dict[str, Union[ArrowOutputStringFormat, "pa.DataType"]]] = None,
     ):
         """
         Gather a list of `LazyDataFrame`s into a single object that can be collected together.
@@ -563,6 +565,8 @@ def __init__(
         )
         super().__init__()
         self._lazy_dataframes = lazy_dataframes
+        self._arrow_string_format_default = arrow_string_format_default
+        self._arrow_string_format_per_column = arrow_string_format_per_column
         if len(self._lazy_dataframes):
             self._lib = self._lazy_dataframes[0].lib
             self._output_format = self._lazy_dataframes[0].read_request.output_format
@@ -588,7 +592,12 @@ def collect(self) -> List[Union[VersionedItem, DataError]]:
         """
         if not len(self._lazy_dataframes):
             return []
-        return self._lib.read_batch(self._read_requests(), output_format=self._output_format)
+        return self._lib.read_batch(
+            self._read_requests(),
+            output_format=self._output_format,
+            arrow_string_format_default=self._arrow_string_format_default,
+            arrow_string_format_per_column=self._arrow_string_format_per_column,
+        )
 
     def _read_requests(self) -> List[ReadRequest]:
         # Combines queries for individual LazyDataFrames with the global query associated with this
@@ -647,33 +656,6 @@ def __init__(
         super().__init__()
         self._lazy_dataframes = lazy_dataframes
         self.then(join)
-        self.arrow_string_format_default = None
-        self.arrow_string_format_per_column = {}
-        for lf in self._lazy_dataframes._lazy_dataframes:
-            self.arrow_string_format_default = (
-                self.arrow_string_format_default or lf.read_request.arrow_string_format_default
-            )
-            check(
-                lf.read_request.arrow_string_format_default is None
-                or self.arrow_string_format_default == lf.read_request.arrow_string_format_default,
-                "Lazy frames from collection cannot be combined for join because they have incompatible arrow_string_format_default values {} and {}",
-                self.arrow_string_format_default,
-                lf.read_request.arrow_string_format_default,
-            )
-            if lf.read_request.arrow_string_format_per_column is not None:
-                common_cols = (
-                    self.arrow_string_format_per_column.keys() & lf.read_request.arrow_string_format_per_column.keys()
-                )
-                for common_col in common_cols:
-                    check(
-                        self.arrow_string_format_per_column[common_col]
-                        == lf.read_request.arrow_string_format_per_column[common_col],
-                        "Lazy frames from collection cannot be combined for join because they have incompatible arrow_string_format_per_column values {} and {} for column {}",
-                        self.arrow_string_format_per_column[common_col],
-                        lf.read_request.arrow_string_format_per_column[common_col],
-                        common_col,
-                    )
-                self.arrow_string_format_per_column.update(lf.read_request.arrow_string_format_per_column)
 
     def collect(self) -> VersionedItemWithJoin:
         """
@@ -693,8 +675,8 @@ def collect(self) -> VersionedItemWithJoin:
                 self._lazy_dataframes._read_requests(),
                 self,
                 output_format=self._lazy_dataframes._output_format,
-                arrow_string_format_default=self.arrow_string_format_default,
-                arrow_string_format_per_column=self.arrow_string_format_per_column,
+                arrow_string_format_default=self._lazy_dataframes._arrow_string_format_default,
+                arrow_string_format_per_column=self._lazy_dataframes._arrow_string_format_per_column,
             )
 
     def __str__(self) -> str:
@@ -2079,6 +2061,7 @@ def read_batch(
         lazy: bool = False,
         output_format: Optional[Union[OutputFormat, str]] = None,
         arrow_string_format_default: Optional[Union[ArrowOutputStringFormat, "pa.DataType"]] = None,
+        arrow_string_format_per_column: Optional[Dict[str, Union[ArrowOutputStringFormat, "pa.DataType"]]] = None,
     ) -> Union[List[Union[VersionedItem, DataError]], LazyDataFrameCollection]:
         """
         Reads multiple symbols.
@@ -2107,6 +2090,10 @@ def read_batch(
             It serves as the default for the entire batch. The string format settings inside the `ReadRequest`s will
             override this batch level setting.
 
+        arrow_string_format_per_column: Optional[Dict[str, Union[ArrowOutputStringFormat, "pa.DataType"]]], default=None,
+            Provides per column name overrides for `arrow_string_format_default`. It is only applied to symbols which
+            don't have a `arrow_string_format_per_column` set in their `ReadRequest`.
+
         Returns
         -------
         Union[List[Union[VersionedItem, DataError]], LazyDataFrameCollection]
@@ -2221,14 +2208,18 @@ def handle_symbol(s_):
                             columns=columns[idx],
                             query_builder=q,
                             output_format=output_format,
-                            arrow_string_format_default=(
-                                per_symbol_arrow_string_format_default[idx] or arrow_string_format_default
-                            ),
-                            arrow_string_format_per_column=per_symbol_arrow_string_format_per_column[idx],
+                            arrow_string_format_default=per_symbol_arrow_string_format_default[idx]
+                            or arrow_string_format_default,
+                            arrow_string_format_per_column=per_symbol_arrow_string_format_per_column[idx]
+                            or arrow_string_format_per_column,
                         ),
                     )
                 )
-            return LazyDataFrameCollection(lazy_dataframes)
+            return LazyDataFrameCollection(
+                lazy_dataframes,
+                arrow_string_format_default=arrow_string_format_default,
+                arrow_string_format_per_column=arrow_string_format_per_column,
+            )
         else:
             return self._nvs._batch_read_to_versioned_items(
                 symbol_strings,
@@ -2242,6 +2233,7 @@ def handle_symbol(s_):
                 iterate_snapshots_if_tombstoned=False,
                 output_format=output_format,
                 arrow_string_format_default=arrow_string_format_default,
+                arrow_string_format_per_column=arrow_string_format_per_column,
                 per_symbol_arrow_string_format_default=per_symbol_arrow_string_format_default,
                 per_symbol_arrow_string_format_per_column=per_symbol_arrow_string_format_per_column,
             )

python/tests/unit/arcticdb/test_arrow_api.py

@@ -342,10 +342,14 @@ def test_read_batch_strings(lmdb_storage, lib_name, lazy, batch_default):
         ReadRequest(
             symbol=sym_2,
             arrow_string_format_default=ArrowOutputStringFormat.LARGE_STRING,
-            arrow_string_format_per_column={"col_2": ArrowOutputStringFormat.CATEGORICAL},
         ),
     ]
-    batch_result = lib.read_batch(read_requests, arrow_string_format_default=batch_default, lazy=lazy)
+    batch_result = lib.read_batch(
+        read_requests,
+        arrow_string_format_default=batch_default,
+        arrow_string_format_per_column={"col_2": ArrowOutputStringFormat.CATEGORICAL},
+        lazy=lazy,
+    )
     if lazy:
         batch_result = batch_result.collect()
     table_1 = batch_result[0].data
@@ -354,49 +358,32 @@ def test_read_batch_strings(lmdb_storage, lib_name, lazy, batch_default):
     assert_frame_equal_with_arrow(table_1, df_1)
     table_2 = batch_result[1].data
     assert table_2.schema.field(0).type == pa.large_string()  # per symbol default
-    assert table_2.schema.field(1).type == pa.dictionary(pa.int32(), pa.large_string())  # per_column override
+    assert table_2.schema.field(1).type == pa.dictionary(pa.int32(), pa.large_string())  # global per_column
     assert_frame_equal_with_arrow(table_2, df_2)
 
 
-@pytest.mark.parametrize("default_1", [None, ArrowOutputStringFormat.SMALL_STRING])
-@pytest.mark.parametrize(
-    "default_2", [None, ArrowOutputStringFormat.LARGE_STRING, ArrowOutputStringFormat.SMALL_STRING]
-)
-@pytest.mark.parametrize("per_column_1", [ArrowOutputStringFormat.CATEGORICAL, ArrowOutputStringFormat.LARGE_STRING])
-def test_read_batch_and_join_strings(lmdb_storage, lib_name, default_1, default_2, per_column_1):
+@pytest.mark.parametrize("default", [None, ArrowOutputStringFormat.SMALL_STRING])
+@pytest.mark.parametrize("per_column", [None, ArrowOutputStringFormat.CATEGORICAL])
+def test_read_batch_and_join_strings(lmdb_storage, lib_name, default, per_column):
     ac = lmdb_storage.create_arctic(output_format=OutputFormat.EXPERIMENTAL_ARROW)
     lib = ac.create_library(lib_name, library_options=LibraryOptions(dynamic_schema=True))
     sym_1, sym_2 = "sym_1", "sym_2"
     df_1 = pd.DataFrame({"col_1": ["a", "a", "bb"], "col_2": ["x", "y", "z"]})
     df_2 = pd.DataFrame({"col_2": ["a", "aa", "aaa"], "col_3": ["a", "a", "a"]})
     lib.write_batch([WritePayload(sym_1, df_1), WritePayload(sym_2, df_2)])
 
-    read_requests = [
-        ReadRequest(
-            symbol=sym_1, arrow_string_format_default=default_1, arrow_string_format_per_column={"col_2": per_column_1}
-        ),
-        ReadRequest(
-            symbol=sym_2,
-            arrow_string_format_default=default_2,
-            arrow_string_format_per_column={
-                "col_2": ArrowOutputStringFormat.CATEGORICAL,
-                "col_3": ArrowOutputStringFormat.LARGE_STRING,
-            },
-        ),
-    ]
-    lazy_dfs = lib.read_batch(read_requests, lazy=True)
-
-    has_mismatch_default = default_1 is not None and default_2 is not None and default_1 != default_2
-    has_mismatch_per_column = per_column_1 != ArrowOutputStringFormat.CATEGORICAL
-    should_raise = has_mismatch_default or has_mismatch_per_column
-    if should_raise:
-        with pytest.raises(ArcticNativeException):
-            lazy_with_join = concat(lazy_dfs)
-    else:
-        lazy_with_join = concat(lazy_dfs)
-        result = lazy_with_join.collect().data
-        assert result.schema.field(0).type == default_1 or default_2 or pa.large_string()
-        assert result.schema.field(1).type == pa.dictionary(pa.int32(), pa.large_string())
-        assert result.schema.field(2).type == pa.large_string()
-        expected_df = pd.concat([df_1, df_2]).reset_index(drop=True)
-        assert_frame_equal_with_arrow(expected_df, result)
+    arrow_string_format_per_column = {"col_1": per_column, "col_3": per_column} if per_column is not None else None
+    lazy_dfs = lib.read_batch(
+        [sym_1, sym_2],
+        arrow_string_format_default=default,
+        arrow_string_format_per_column=arrow_string_format_per_column,
+        lazy=True,
+    )
+
+    lazy_with_join = concat(lazy_dfs)
+    result = lazy_with_join.collect().data
+    assert result.schema.field(0).type == per_column or default or pa.large_string()
+    assert result.schema.field(1).type == default or pa.large_string()
+    assert result.schema.field(2).type == per_column or default or pa.large_string()
+    expected_df = pd.concat([df_1, df_2]).reset_index(drop=True)
+    assert_frame_equal_with_arrow(expected_df, result)

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

@@ -1135,22 +1135,24 @@ def test_arrow_read_batch_with_strings(lmdb_version_store_arrow):
     lib.batch_write([sym_1, sym_2], [df_1, df_2])
 
     arrow_string_format_default = ArrowOutputStringFormat.SMALL_STRING
-    per_symbol_arrow_string_format_default = [None, ArrowOutputStringFormat.LARGE_STRING]
+    arrow_string_format_per_column = {"col_1": ArrowOutputStringFormat.CATEGORICAL}
+    per_symbol_arrow_string_format_default = [ArrowOutputStringFormat.LARGE_STRING, None]
     per_symbol_arrow_string_format_per_column = [
-        {"col_1": ArrowOutputStringFormat.CATEGORICAL},
+        None,  # First item will use the global arrow_string_format_per_column
         {"col_2": ArrowOutputStringFormat.CATEGORICAL},
     ]
     batch_result = lib.batch_read(
         [sym_1, sym_2],
         arrow_string_format_default=arrow_string_format_default,
+        arrow_string_format_per_column=arrow_string_format_per_column,
         per_symbol_arrow_string_format_default=per_symbol_arrow_string_format_default,
         per_symbol_arrow_string_format_per_column=per_symbol_arrow_string_format_per_column,
     )
     table_1 = batch_result[sym_1].data
-    assert table_1.schema.field(0).type == pa.dictionary(pa.int32(), pa.large_string())  # per_column override
-    assert table_1.schema.field(1).type == pa.string()  # global default for all symbols
+    assert table_1.schema.field(0).type == pa.dictionary(pa.int32(), pa.large_string())  # global per_column
+    assert table_1.schema.field(1).type == pa.large_string()  # per symbol default
     assert_frame_equal_with_arrow(table_1, df_1)
     table_2 = batch_result[sym_2].data
-    assert table_2.schema.field(0).type == pa.large_string()  # per symbol default
+    assert table_2.schema.field(0).type == pa.string()  # global default for all symbols
     assert table_2.schema.field(1).type == pa.dictionary(pa.int32(), pa.large_string())  # per_column override
     assert_frame_equal_with_arrow(table_2, df_2)