Arrow read support (experimental release) (#2489)

#### Reference Issues/PRs
<!--Example: Fixes #1234. See also #3456.-->

#### What does this implement or fix?
Adds a frontend API for using Arrow. The API looks like below:
```
from arcticdb import Arctic, OutputFormat
ac = adb.Arctic(uri, output_format=OutputFormat.EXPERIMENTAL_ARROW) # Sets a runtime output format option so all read operations return arrow tables.
lib = ac["lib_name"]
lib.read(sym).data # This will return pyarrow.Table
lib.read(sym, output_format=OutputFormat.PANDAS) # We can also override the output format for any specific read operation
```

All read operations `read`, `read_batch`, `read_batch_and_join` and
their `lazy` equivalents adhere to the `output_format` argument.

The changes in this PR are:
- Change the `_output_format` argument to `output_format`
- Separates the internal C++ `OutputFormat` from the python one. For
python we use a `StrEnum` which allows users to pass both the enum and
the string value
- Adds a `RuntimeOptions` class stored inside the `Arctic` and
`NativeVersionStore` instances. `RuntimeOptions` contains only the
`output_format` currently but will later on include things like
`arrow_string_column_encoding` and other layout configurations.
- Allow passing `output_format` in all V2 APIs and make it work for
`lazy=True` cases
- Run all query builder tests also with `output_format=ARROW` 
- Clean up `test_arrow` and `test_arrow_normalization` to adhere to new
API
- Add tests for all user facing arrow APIs in `test_arrow_api.py` 
- Fixes an issue where `read_batch_and_join` didn't respect the input
`ReadOptions`
- Introduces the `arcticdb.dependencies` for handling optional
dependencies.

#### Any other comments?
Writing tests for the optional dependencies is difficult. So I ran a
manual test:

In a venv with pyarrow:
```
>>> import arcticdb as adb
>>> ac = adb.Arctic("lmdb:///tmp/test-arrow")
>>> lib = ac["test"]
>>> lib.read("test", output_format="pandas").data
   x
0  5
1  6
2  7
>>> lib.read("test", output_format="experimental_arrow").data
pyarrow.Table
x: int64
----
x: [[5,6,7]]
```

And in a venv without pyarrow:
```
>>> import arcticdb as adb
>>> ac = adb.Arctic("lmdb:///tmp/test-arrow")
>>> lib = ac["test"]
>>> lib.read("test", output_format="pandas").data
   x
0  5
1  6
2  7
>>> lib.read("test", output_format="experimental_arrow").data
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/ivo/source/read_as_arrow/python/arcticdb/version_store/library.py", line 1887, in read
    return self._nvs.read(
           ^^^^^^^^^^^^^^^
  File "/home/ivo/source/read_as_arrow/python/arcticdb/version_store/_store.py", line 2063, in read
    version_query, read_options, read_query = self._get_queries(
                                              ^^^^^^^^^^^^^^^^^^
  File "/home/ivo/source/read_as_arrow/python/arcticdb/version_store/_store.py", line 1970, in _get_queries
    read_options = self._get_read_options(**kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ivo/source/read_as_arrow/python/arcticdb/version_store/_store.py", line 1958, in _get_read_options
    output_format_to_internal(
  File "/home/ivo/source/read_as_arrow/python/arcticdb/options.py", line 162, in output_format_to_internal
    raise ModuleNotFoundError("ArcticDB's pyarrow optional dependency missing but is required to use arrow output format.")
ModuleNotFoundError: ArcticDB's pyarrow optional dependency missing but is required to use arrow output format.
```

#### Checklist

<details>
  <summary>
   Checklist for code changes...
  </summary>
 
- [ ] Have you updated the relevant docstrings, documentation and
copyright notice?
- [ ] Is this contribution tested against [all ArcticDB's
features](../docs/mkdocs/docs/technical/contributing.md)?
- [ ] Do all exceptions introduced raise appropriate [error
messages](https://docs.arcticdb.io/error_messages/)?
 - [ ] Are API changes highlighted in the PR description?
- [ ] Is the PR labelled as enhancement or bug so it appears in
autogenerated release notes?
</details>

<!--
Thanks for contributing a Pull Request to ArcticDB! Please ensure you
have taken a look at:
- ArcticDB's Code of Conduct:
https://github.com/man-group/ArcticDB/blob/master/CODE_OF_CONDUCT.md
- ArcticDB's Contribution Licensing:
https://github.com/man-group/ArcticDB/blob/master/docs/mkdocs/docs/technical/contributing.md#contribution-licensing
-->

Author: IvoDD

cpp/arcticdb/pipeline/read_options.hpp

@@ -90,5 +90,9 @@ struct ReadOptions {
     [[nodiscard]] OutputFormat output_format() const {
         return data_->output_format_;
     }
+
+    [[nodiscard]] ReadOptions clone() const {
+        return ReadOptions(std::make_shared<ReadOptionsData>(*data_));
+    }
 };
 } //namespace arcticdb

cpp/arcticdb/version/local_versioned_engine.cpp

@@ -1293,22 +1293,22 @@ MultiSymbolReadOutput LocalVersionedEngine::batch_read_and_join_internal(
     }
     auto clauses_ptr = std::make_shared<std::vector<std::shared_ptr<Clause>>>(std::move(clauses));
     return folly::collect(symbol_processing_result_futs).via(&async::io_executor())
-    .thenValueInline([this, &handler_data, clauses_ptr, component_manager](std::vector<SymbolProcessingResult>&& symbol_processing_results) mutable {
+    .thenValueInline([this, &handler_data, clauses_ptr, component_manager, read_options](std::vector<SymbolProcessingResult>&& symbol_processing_results) mutable {
         auto [input_schemas, entity_ids, res_versioned_items, res_metadatas] = unpack_symbol_processing_results(std::move(symbol_processing_results));
         auto pipeline_context = setup_join_pipeline_context(std::move(input_schemas), *clauses_ptr);
         return schedule_remaining_iterations(std::move(entity_ids), clauses_ptr)
         .thenValueInline([component_manager](std::vector<EntityId>&& processed_entity_ids) {
             auto proc = gather_entities<std::shared_ptr<SegmentInMemory>, std::shared_ptr<RowRange>, std::shared_ptr<ColRange>>(*component_manager, std::move(processed_entity_ids));
             return collect_segments(std::move(proc));
         })
-        .thenValueInline([store=store(), &handler_data, pipeline_context](std::vector<SliceAndKey>&& slice_and_keys) mutable {
-            return prepare_output_frame(std::move(slice_and_keys), pipeline_context, store, ReadOptions{}, handler_data);
+        .thenValueInline([store=store(), &handler_data, pipeline_context, read_options](std::vector<SliceAndKey>&& slice_and_keys) mutable {
+            return prepare_output_frame(std::move(slice_and_keys), pipeline_context, store, read_options, handler_data);
         })
-        .thenValueInline([&handler_data, pipeline_context, res_versioned_items, res_metadatas](SegmentInMemory&& frame) mutable {
+        .thenValueInline([&handler_data, pipeline_context, res_versioned_items, res_metadatas, read_options](SegmentInMemory&& frame) mutable {
             // Needed to force our usual backfilling behaviour when columns have been outer-joined and some are not present in all input symbols
-            ReadOptions read_options;
-            read_options.set_dynamic_schema(true);
-            return reduce_and_fix_columns(pipeline_context, frame, read_options, handler_data)
+            ReadOptions read_options_with_dynamic_schema = read_options.clone();
+            read_options_with_dynamic_schema.set_dynamic_schema(true);
+            return reduce_and_fix_columns(pipeline_context, frame, read_options_with_dynamic_schema, handler_data)
             .thenValueInline([pipeline_context, frame, res_versioned_items, res_metadatas](auto&&) mutable {
                 return MultiSymbolReadOutput{
                     std::move(*res_versioned_items),

cpp/arcticdb/version/python_bindings.cpp

@@ -230,7 +230,7 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
         .def("set_timestamp", &VersionQuery::set_timestamp)
         .def("set_version", &VersionQuery::set_version);
 
-    py::enum_<OutputFormat>(version, "OutputFormat")
+    py::enum_<OutputFormat>(version, "InternalOutputFormat")
         .value("PANDAS", OutputFormat::PANDAS)
         .value("ARROW", OutputFormat::ARROW);
 

python/arcticdb/__init__.py

@@ -5,7 +5,7 @@
 import sys as _sys
 
 from arcticdb.arctic import Arctic
-from arcticdb.options import LibraryOptions
+from arcticdb.options import LibraryOptions, OutputFormat, RuntimeOptions
 from arcticdb.version_store.processing import QueryBuilder, where
 from arcticdb.version_store._store import VersionedItem
 import arcticdb.version_store.library as library

python/arcticdb/arctic.py

@@ -8,7 +8,7 @@
 import logging
 from typing import List, Optional, Any, Union
 
-from arcticdb.options import DEFAULT_ENCODING_VERSION, LibraryOptions, EnterpriseLibraryOptions
+from arcticdb.options import DEFAULT_ENCODING_VERSION, LibraryOptions, EnterpriseLibraryOptions, RuntimeOptions, OutputFormat
 from arcticdb_ext.storage import LibraryManager
 from arcticdb.exceptions import LibraryNotFound, MismatchingLibraryOptions
 from arcticdb.version_store.library import ArcticInvalidApiUsageException, Library
@@ -46,7 +46,7 @@ class Arctic:
     # It is set by the LmdbStorageFixture
     _accessed_libs: Optional[List[NativeVersionStore]] = None
 
-    def __init__(self, uri: str, encoding_version: EncodingVersion = DEFAULT_ENCODING_VERSION):
+    def __init__(self, uri: str, encoding_version: EncodingVersion = DEFAULT_ENCODING_VERSION, output_format: Union[OutputFormat, str] = OutputFormat.PANDAS):
         """
         Initializes a top-level Arctic library management instance.
 
@@ -59,10 +59,19 @@ def __init__(self, uri: str, encoding_version: EncodingVersion = DEFAULT_ENCODIN
             URI specifying the backing store used to access, configure, and create Arctic libraries.
             For more details about the parameters, please refer to the [Arctic URI Documentation](./arctic_uri.md).
 
-        encoding_version: EncodingVersion, default DEFAULT_ENCODING_VERSION
+        encoding_version: EncodingVersion, default = EncodingVersion.V1
             When creating new libraries with this Arctic instance, the default encoding version to use.
             Can be overridden by specifying the encoding version in the LibraryOptions argument to create_library.
 
+        output_format: Union[OutputFormat, str], default = OutputFormat.PANDAS
+            Controls the default output format of all operations returning a dataframe.
+            The default behavior (OutputFormat.PANDAS) is to return `pandas.DataFrame`s or `pandas.Series` backed by
+            numpy arrays.
+            OutputFormat.EXPERIMENTAL_ARROW will return all dataframes as `pyarrow.Table`s. The arrow API is still
+            experimental and the arrow layout might change in a minor release.
+            Accepts the OutputFormat as either OutputFormat enum values or as case-insensitive strings like "pandas"
+            and "experimental_arrow".
+
         Examples
         --------
 
@@ -89,31 +98,46 @@ def __init__(self, uri: str, encoding_version: EncodingVersion = DEFAULT_ENCODIN
         self._library_adapter: ArcticLibraryAdapter = _cls(uri, self._encoding_version)
         self._library_manager = LibraryManager(self._library_adapter.config_library)
         self._uri = uri
+        self._runtime_options = RuntimeOptions(output_format=output_format)
 
-    def __getitem__(self, name: str) -> Library:
+    def _get_library(self, name: str, output_format: Optional[Union[OutputFormat, str]] = None) -> Library:
         lib_mgr_name = self._library_adapter.get_name_for_library_manager(name)
         if not self._library_manager.has_library(lib_mgr_name):
             raise LibraryNotFound(name)
 
         storage_override = self._library_adapter.get_storage_override()
+
+        if output_format is not None:
+            runtime_options = RuntimeOptions(output_format=output_format)
+        else:
+            runtime_options = self._runtime_options
+
         lib = NativeVersionStore(
             self._library_manager.get_library(lib_mgr_name, storage_override, native_storage_config=self._library_adapter.native_config()),
             repr(self._library_adapter),
             lib_cfg=self._library_manager.get_library_config(lib_mgr_name, storage_override),
-            native_cfg=self._library_adapter.native_config()
+            native_cfg=self._library_adapter.native_config(),
+            runtime_options=runtime_options
         )
         if self._accessed_libs is not None:
             self._accessed_libs.append(lib)
         return Library(repr(self), lib)
 
+    def __getitem__(self, name: str):
+        return self._get_library(name)
+
     def __repr__(self):
         return "Arctic(config=%r)" % self._library_adapter
 
     def __contains__(self, name: str):
         return self.has_library(name)
 
     def get_library(
-        self, name: str, create_if_missing: Optional[bool] = False, library_options: Optional[LibraryOptions] = None
+        self,
+        name: str,
+        create_if_missing: Optional[bool] = False,
+        library_options: Optional[LibraryOptions] = None,
+        output_format: Optional[Union[OutputFormat, str]] = None,
     ) -> Library:
         """
         Returns the library named ``name``.
@@ -136,6 +160,11 @@ def get_library(
             match these.
             Unused if create_if_missing is False.
 
+        output_format: Optional[Union[OutputFormat, str]], default = None
+            Controls the default output format of all operations on the library returning a dataframe.
+            For more information see documentation of `Arctic.__init__`.
+            If `None` uses the output format from the Arctic instance.
+
         Examples
         --------
         >>> arctic = adb.Arctic('s3://MY_ENDPOINT:MY_BUCKET')
@@ -152,7 +181,7 @@ def get_library(
                 "In get_library, library_options must be falsey if create_if_missing is falsey"
             )
         try:
-            lib = self[name]
+            lib = self._get_library(name, output_format)
             if create_if_missing and library_options:
                 if library_options.encoding_version is None:
                     library_options.encoding_version = self._encoding_version
@@ -161,14 +190,15 @@ def get_library(
             return lib
         except LibraryNotFound as e:
             if create_if_missing:
-                return self.create_library(name, library_options)
+                return self.create_library(name, library_options, output_format)
             else:
                 raise e
 
     def create_library(self,
                        name: str,
                        library_options: Optional[LibraryOptions] = None,
-                       enterprise_library_options: Optional[EnterpriseLibraryOptions] = None) -> Library:
+                       enterprise_library_options: Optional[EnterpriseLibraryOptions] = None,
+                       output_format: Optional[Union[OutputFormat, str]] = None) -> Library:
         """
         Creates the library named ``name``.
 
@@ -192,6 +222,11 @@ def create_library(self,
             Enterprise options to use in configuring the library. Defaults if not provided are the same as documented in
             EnterpriseLibraryOptions. These options are only relevant to ArcticDB enterprise users.
 
+        output_format: Optional[Union[OutputFormat, str]], default = None
+            Controls the default output format of all operations on the library returning a dataframe.
+            For more information see documentation of `Arctic.__init__`.
+            If `None` uses the output format from the Arctic instance.
+
         Examples
         --------
         >>> arctic = adb.Arctic('s3://MY_ENDPOINT:MY_BUCKET')
@@ -214,7 +249,7 @@ def create_library(self,
         cfg = self._library_adapter.get_library_config(name, library_options, enterprise_library_options)
         lib_mgr_name = self._library_adapter.get_name_for_library_manager(name)
         self._library_manager.write_library_config(cfg, lib_mgr_name, self._library_adapter.get_masking_override())
-        return self.get_library(name)
+        return self.get_library(name, output_format=output_format)
 
     def delete_library(self, name: str) -> None:
         """

python/arcticdb/dependencies.py

@@ -0,0 +1,39 @@
+from types import ModuleType
+from typing import Any, Tuple
+from importlib import import_module
+
+
+class MissingModule(ModuleType):
+    """
+    A dummy module used to represent a missing optional dependency.
+    Raises a meaningful error message when trying to get any attribute.
+    """
+
+    def __init__(
+            self,
+            module_name: str,
+    ) -> None:
+        self._module_name = module_name
+        super().__init__(module_name)
+
+    def __getattr__(self, name: str) -> Any:
+        msg = f"ArcticDB's {self._module_name!r} optional dependency is missing but should be installed to use this feature."
+        raise ModuleNotFoundError(msg)
+
+
+def _import_optional_dependency(module_name: str) -> Tuple[ModuleType, bool]:
+    try:
+        module = import_module(module_name)
+        return module, True
+    except ImportError:
+        module = MissingModule(module_name)
+        return module, False
+
+
+pyarrow, _PYARROW_AVAILABLE = _import_optional_dependency("pyarrow")
+
+
+__all__ = [
+    "pyarrow",
+    "_PYARROW_AVAILABLE",
+]

python/arcticdb/options.py

@@ -6,11 +6,13 @@
 As of the Change Date specified in that file, in accordance with the Business Source License, use of this software will be governed by the Apache License, version 2.0.
 """
 
-from typing import Optional
+from typing import Optional, Union
 from enum import Enum
 
+from arcticdb.dependencies import _PYARROW_AVAILABLE
 from arcticdb.encoding_version import EncodingVersion
 from arcticdb_ext.storage import ModifiableLibraryOption, ModifiableEnterpriseLibraryOption
+from arcticdb_ext.version_store import InternalOutputFormat
 
 
 DEFAULT_ENCODING_VERSION = EncodingVersion.V1
@@ -146,6 +148,35 @@ def __repr__(self):
         )
 
 
+# TODO: Use enum.StrEnum when we no longer need to support python 3.9
+class OutputFormat(str, Enum):
+    PANDAS = "PANDAS"
+    EXPERIMENTAL_ARROW = "EXPERIMENTAL_ARROW"
+
+
+def output_format_to_internal(output_format: Union[OutputFormat, str]) -> InternalOutputFormat:
+    if output_format.lower() == OutputFormat.PANDAS.lower():
+        return InternalOutputFormat.PANDAS
+    elif output_format.lower() == OutputFormat.EXPERIMENTAL_ARROW.lower():
+        if not _PYARROW_AVAILABLE:
+            raise ModuleNotFoundError("ArcticDB's pyarrow optional dependency missing but is required to use arrow output format.")
+        return InternalOutputFormat.ARROW
+    else:
+        raise ValueError(f"Unknown OutputFormat: {output_format}")
+
+class RuntimeOptions:
+    def __init__(
+        self,
+        *,
+        output_format: Union[OutputFormat, str] = OutputFormat.PANDAS,
+    ):
+        self.output_format=output_format
+
+
+    def set_output_format(self, output_format: Union[OutputFormat, str]):
+        self.output_format = output_format
+
+
 class EnterpriseLibraryOptions:
     """
     Configuration options for ArcticDB libraries, that should only be used when you are using the ArcticDB enterprise

python/arcticdb/util/arrow.py

@@ -0,0 +1,18 @@
+from arcticdb.dependencies import pyarrow as pa
+
+def stringify_dictionary_encoded_columns(table, string_type=None):
+    """
+    Converts all pyarrow.Table dictionary encoded columns to strings.
+
+    ArcticDB currently returns string columns in dictionary encoded arrow arrays when using
+    OutputFormat.EXPERIMENTAL_ARROW.
+
+    Useful for testing when comparing to the source dataframe where we want regular large_string columns instead of
+    categorical columns.
+    """
+    if string_type is None:
+        string_type = pa.large_string()
+    for i, name in enumerate(table.column_names):
+        if pa.types.is_dictionary(table.column(i).type):
+            table = table.set_column(i, name, table.column(name).cast(string_type))
+    return table

python/arcticdb/util/test.py

@@ -19,6 +19,7 @@
 import time
 import attr
 from functools import wraps, reduce
+from arcticdb.dependencies import pyarrow as pa
 
 from arcticdb.util.marks import SHORTER_LOGS
 
@@ -29,6 +30,7 @@
 
 from arcticdb import QueryBuilder
 from arcticdb.util._versions import IS_PANDAS_ONE, PANDAS_VERSION, CHECK_FREQ_VERSION
+from arcticdb.util.arrow import stringify_dictionary_encoded_columns
 from arcticdb.version_store import NativeVersionStore
 from arcticdb.version_store._custom_normalizers import CustomNormalizer
 from arcticc.pb2.descriptors_pb2 import NormalizationMetadata
@@ -253,38 +255,15 @@ def assert_frame_equal_rebuild_index_first(expected: pd.DataFrame, actual: pd.Da
 
 
 def convert_arrow_to_pandas_and_remove_categoricals(table):
-    """
-    Converts a pyarrow.Table to a pandas.DataFrame and unwinds all dictionary encoded columns.
-
-    Useful for testing because ArcticDB currently returns string columns in arrow as dictionary encoded, but when
-    comparing to the source pandas dataframes we want regular string columns instead of categoricals.
-    """
-    import pyarrow as pa
-    new_columns = []
-    new_fields = []
-    metadata = table.schema.metadata
-
-    for i, col in enumerate(table.columns):
-        field = table.field(i)
-        if isinstance(field.type, pa.DictionaryType):
-            col = pa.compute.cast(col, field.type.value_type)
-            field = field.with_type(field.type.value_type)
-        new_columns.append(col)
-        new_fields.append(field)
-
-    new_table = pa.Table.from_arrays(
-        new_columns,
-        schema=pa.schema(new_fields).with_metadata(metadata)
-    )
+    new_table = stringify_dictionary_encoded_columns(table)
     return new_table.to_pandas()
 
-def assert_frame_equal_with_arrow(left, right):
-    import pyarrow as pa
+def assert_frame_equal_with_arrow(left, right, **kwargs):
     if isinstance(left, pa.Table):
         left = convert_arrow_to_pandas_and_remove_categoricals(left)
     if isinstance(right, pa.Table):
         right = convert_arrow_to_pandas_and_remove_categoricals(right)
-    assert_frame_equal(left, right)
+    assert_frame_equal(left, right, **kwargs)
 
 
 unicode_symbol = "\u00A0"  # start of latin extensions