Parametrize all pipeline tests to run with arrow output format

Author: IvoDD

python/arcticdb/util/test.py

@@ -256,6 +256,11 @@ def assert_frame_equal_rebuild_index_first(expected: pd.DataFrame, actual: pd.Da
 
 def convert_arrow_to_pandas_and_remove_categoricals(table):
     new_table = stringify_dictionary_encoded_columns(table)
+    for i, name in enumerate(new_table.column_names):
+        if pa.types.is_integer(new_table.column(i).type):
+            new_col = new_table.column(i).fill_null(0)
+            print(new_col)
+            new_table = new_table.set_column(i, name, new_col)
     return new_table.to_pandas()
 
 def assert_frame_equal_with_arrow(left, right, **kwargs):

python/arcticdb/version_store/_store.py

@@ -88,6 +88,8 @@
 from packaging.version import Version
 import arcticdb_ext as ae
 
+from arcticdb.util.arrow import stringify_dictionary_encoded_columns
+
 IS_WINDOWS = sys.platform == "win32"
 
 
@@ -348,12 +350,17 @@ def _initialize(self, library, env, lib_cfg, custom_normalizer, open_mode, nativ
         self._open_mode = open_mode
         self._native_cfg = native_cfg
         self._runtime_options=runtime_options
+        self._test_convert_arrow_back_to_pandas = False
 
     def set_output_format(self, output_format: Union[OutputFormat, str]):
         if self._runtime_options is None:
             self._runtime_options = RuntimeOptions()
         self._runtime_options.set_output_format(output_format)
 
+    def _set_output_format_for_pipeline_tests(self, output_format):
+        self.set_output_format(output_format)
+        self._test_convert_arrow_back_to_pandas = True
+
     @classmethod
     def create_store_from_lib_config(cls, lib_cfg, env, open_mode=OpenMode.DELETE, native_cfg=None):
         lib = cls.create_lib_from_lib_config(lib_cfg, env, open_mode, native_cfg)
@@ -723,6 +730,9 @@ def _resolve_dynamic_strings(self, kwargs):
                     "Windows only supports dynamic_strings=True, using dynamic strings despite configuration or kwarg"
                 )
             dynamic_strings = True
+        if self._test_convert_arrow_back_to_pandas:
+            # TODO: Hackery, maybe better to skip
+            dynamic_strings = True
         return dynamic_strings
 
     last_mismatch_msg: Optional[str] = None
@@ -2400,6 +2410,14 @@ def _adapt_read_res(self, read_result: ReadResult) -> VersionedItem:
                 record_batches.append(pa.RecordBatch._import_from_c(record_batch.array(), record_batch.schema()))
             table = pa.Table.from_batches(record_batches)
             data = self._arrow_normalizer.denormalize(table, read_result.norm)
+            if self._test_convert_arrow_back_to_pandas:
+                # TODO: Deduplicate with convert_arrow_to_pandas_and_remove_categoricals
+                data = stringify_dictionary_encoded_columns(data)
+                for i, name in enumerate(data.column_names):
+                    if pa.types.is_integer(data.column(i).type):
+                        new_col = data.column(i).fill_null(0)
+                        data = data.set_column(i, name, new_col)
+                data = data.to_pandas()
         else:
             data = self._normalizer.denormalize(read_result.frame_data, read_result.norm)
             if read_result.norm.HasField("custom"):

python/tests/hypothesis/arcticdb/test_aggregation_hypothesis.py

@@ -23,7 +23,7 @@
 )
 
 
-pytestmark = pytest.mark.pipeline
+pytestmark = pytest.mark.pipeline # Covered
 
 
 @use_of_function_scoped_fixtures_in_hypothesis_checked
@@ -36,9 +36,10 @@
         ],
     ),
 )
-def test_aggregation_numeric(lmdb_version_store_v1, df):
+def test_aggregation_numeric(lmdb_version_store_v1, any_output_format, df):
     assume(not df.empty)
     lib = lmdb_version_store_v1
+    lib._set_output_format_for_pipeline_tests(any_output_format)
     symbol = "test_aggregation_numeric"
     lib.write(symbol, df)
 
@@ -70,9 +71,10 @@ def test_aggregation_numeric(lmdb_version_store_v1, df):
         ],
     ),
 )
-def test_aggregation_strings(lmdb_version_store_v1, df):
+def test_aggregation_strings(lmdb_version_store_v1, any_output_format, df):
     assume(not df.empty)
     lib = lmdb_version_store_v1
+    lib._set_output_format_for_pipeline_tests(any_output_format)
     symbol = "test_aggregation_strings"
     lib.write(symbol, df)
 
@@ -112,12 +114,13 @@ def aggregation_dataframe_list_strategy(draw):
 @use_of_function_scoped_fixtures_in_hypothesis_checked
 @settings(deadline=None)
 @given(dfs=aggregation_dataframe_list_strategy())
-def test_aggregation_numeric_dynamic(lmdb_version_store_dynamic_schema_v1, dfs):
+def test_aggregation_numeric_dynamic(lmdb_version_store_dynamic_schema_v1, any_output_format, dfs):
     agg_column_dtypes = [df['agg_column'].dtype for df in dfs if 'agg_column' in df.columns]
     common_agg_type = functools.reduce(valid_common_type, agg_column_dtypes) if len(agg_column_dtypes) > 0 else None
     assume(any('grouping_column' in df.columns for df in dfs) and common_agg_type is not None)
 
     lib = lmdb_version_store_dynamic_schema_v1
+    lib._set_output_format_for_pipeline_tests(any_output_format)
     symbol = "test_aggregation_numeric_dynamic"
     lib.delete(symbol)
     for df in dfs:
@@ -155,9 +158,10 @@ def test_aggregation_numeric_dynamic(lmdb_version_store_dynamic_schema_v1, dfs):
         ],
     ),
 )
-def test_aggregation_strings_dynamic(lmdb_version_store_dynamic_schema_v1, df):
+def test_aggregation_strings_dynamic(lmdb_version_store_dynamic_schema_v1, any_output_format, df):
     assume(len(df) >= 3)
     lib = lmdb_version_store_dynamic_schema_v1
+    lib._set_output_format_for_pipeline_tests(any_output_format)
     symbol = "test_aggregation_strings_dynamic"
     lib.delete(symbol)
     slices = [

python/tests/hypothesis/arcticdb/test_resample.py

@@ -21,7 +21,7 @@
 MIN_DATE = np.datetime64('1969-06-01')
 MAX_DATE = np.datetime64('1970-06-01')
 
-pytestmark = pytest.mark.pipeline
+pytestmark = pytest.mark.pipeline # Covered
 
 
 @st.composite
@@ -149,8 +149,9 @@ def dynamic_schema_column_list(draw):
     origin=origin(),
     offset=offset()
 )
-def test_resample(lmdb_version_store_v1, df, rule, origin, offset):
+def test_resample(lmdb_version_store_v1, any_output_format, df, rule, origin, offset):
     lib = lmdb_version_store_v1
+    lib._set_output_format_for_pipeline_tests(any_output_format)
     sym = "sym"
     logger = get_logger()
     logger.info(f"Data frame generated has {df.shape[0]} rows")
@@ -197,9 +198,10 @@ def test_resample(lmdb_version_store_v1, df, rule, origin, offset):
     offset=offset()
 )
 @settings(deadline=None, suppress_health_check=[HealthCheck.data_too_large])
-def test_resample_dynamic_schema(lmdb_version_store_dynamic_schema_v1, df_list, rule, origin, offset):
+def test_resample_dynamic_schema(lmdb_version_store_dynamic_schema_v1, any_output_format, df_list, rule, origin, offset):
     common_column_types = compute_common_type_for_columns_in_df_list(df_list)
     lib = lmdb_version_store_dynamic_schema_v1
+    lib._set_output_format_for_pipeline_tests(any_output_format)
     lib.version_store.clear()
     sym = "sym"
     agg = {f"{name}_{op}": (name, op) for name in common_column_types for op in ALL_AGGREGATIONS}

python/tests/integration/arcticdb/version_store/test_basic_version_store.py

@@ -822,10 +822,11 @@ def test_range_index(basic_store, sym):
     assert_equal(expected, vit.data)
 
 
-@pytest.mark.pipeline
+@pytest.mark.pipeline # Covered
 @pytest.mark.parametrize("use_date_range_clause", [True, False])
 @pytest.mark.storage
-def test_date_range(basic_store, use_date_range_clause):
+def test_date_range(basic_store, use_date_range_clause, any_output_format):
+    basic_store._set_output_format_for_pipeline_tests(any_output_format)
     initial_timestamp = pd.Timestamp("2019-01-01")
     df = pd.DataFrame(data=np.arange(100), index=pd.date_range(initial_timestamp, periods=100))
     sym = "date_test"
@@ -871,10 +872,11 @@ def test_date_range(basic_store, use_date_range_clause):
     assert data_closed[data_closed.columns[0]][-1] == end_offset
 
 
-@pytest.mark.pipeline
+@pytest.mark.pipeline # Covered
 @pytest.mark.parametrize("use_date_range_clause", [True, False])
 @pytest.mark.storage
-def test_date_range_none(basic_store, use_date_range_clause):
+def test_date_range_none(basic_store, use_date_range_clause, any_output_format):
+    basic_store._set_output_format_for_pipeline_tests(any_output_format)
     sym = "date_test2"
     rows = 100
     initial_timestamp = pd.Timestamp("2019-01-01")
@@ -891,10 +893,11 @@ def test_date_range_none(basic_store, use_date_range_clause):
     assert len(data) == rows
 
 
-@pytest.mark.pipeline
+@pytest.mark.pipeline # Covered
 @pytest.mark.parametrize("use_date_range_clause", [True, False])
 @pytest.mark.storage
-def test_date_range_start_equals_end(basic_store, use_date_range_clause):
+def test_date_range_start_equals_end(basic_store, use_date_range_clause, any_output_format):
+    basic_store._set_output_format_for_pipeline_tests(any_output_format)
     sym = "date_test2"
     rows = 100
     initial_timestamp = pd.Timestamp("2019-01-01")
@@ -914,11 +917,12 @@ def test_date_range_start_equals_end(basic_store, use_date_range_clause):
     assert data[data.columns[0]][0] == start_offset
 
 
-@pytest.mark.pipeline
+@pytest.mark.pipeline # Covered
 @pytest.mark.parametrize("use_date_range_clause", [True, False])
 @pytest.mark.storage
-def test_date_range_row_sliced(basic_store_tiny_segment, use_date_range_clause):
+def test_date_range_row_sliced(basic_store_tiny_segment, use_date_range_clause, any_output_format):
     lib = basic_store_tiny_segment
+    lib._set_output_format_for_pipeline_tests(any_output_format)
     sym = "test_date_range_row_sliced"
     # basic_store_tiny_segment produces 2x2 segments
     num_rows = 6
@@ -2719,11 +2723,12 @@ def test_batch_append_with_throw_exception(basic_store, three_col_df):
         )
 
 
-@pytest.mark.pipeline
+@pytest.mark.pipeline # Covered
 @pytest.mark.parametrize("use_date_range_clause", [True, False])
 @pytest.mark.storage
-def test_batch_read_date_range(basic_store_tombstone_and_sync_passive, use_date_range_clause):
+def test_batch_read_date_range(basic_store_tombstone_and_sync_passive, use_date_range_clause, any_output_format):
     lmdb_version_store = basic_store_tombstone_and_sync_passive
+    lmdb_version_store._set_output_format_for_pipeline_tests(any_output_format)
     symbols = []
     for i in range(5):
         symbols.append("sym_{}".format(i))