WIP Arrow demo

Author: IvoDD

cpp/arcticdb/arrow/arrow_utils.cpp

@@ -71,14 +71,5 @@ std::shared_ptr<std::vector<sparrow::record_batch>> segment_to_arrow_data(Segmen
 
     return output;
 }
-ArrowReadResult create_arrow_read_result(
-    const VersionedItem& version,
-    FrameAndDescriptor&& fd) {
-    auto result = std::move(fd);
-    auto arrow_frame = ArrowOutputFrame{segment_to_arrow_data(result.frame_), names_from_segment(result.frame_)};
-
-    const auto& desc_proto = result.desc_.proto();
-    return {version, std::move(arrow_frame), desc_proto.user_meta()};
-}
 
 } // namespace arcticdb

cpp/arcticdb/arrow/arrow_utils.hpp

@@ -40,11 +40,6 @@ struct ArrowReadResult {
     arcticdb::proto::descriptors::UserDefinedMetadata user_meta_;
 };
 
-ArrowReadResult create_arrow_read_result(
-    const VersionedItem& version,
-    FrameAndDescriptor&& fd);
-
-
 std::shared_ptr<std::vector<sparrow::record_batch>> segment_to_arrow_data(SegmentInMemory& segment);
 
 std::vector<std::string> names_from_segment(const SegmentInMemory& segment);

cpp/arcticdb/entity/read_result.hpp

@@ -14,29 +14,36 @@
 #include <arcticdb/entity/frame_and_descriptor.hpp>
 #include <arcticdb/pipeline/python_output_frame.hpp>
 #include <arcticdb/util/memory_tracing.hpp>
+#include <arcticdb/arrow/arrow_utils.hpp>
 
 #include <vector>
 
 namespace arcticdb {
 
+// TODO: Rename the PythonOutputFrame to PandasOutputFrame
+using OutputFrame = std::variant<pipelines::PythonOutputFrame, ArrowOutputFrame>;
+
 struct ARCTICDB_VISIBILITY_HIDDEN ReadResult {
     ReadResult(
             const VersionedItem& versioned_item,
-            pipelines::PythonOutputFrame&& frame_data,
+            OutputFrame&& frame_data,
+            OutputFormat output_format,
             const arcticdb::proto::descriptors::NormalizationMetadata& norm_meta,
             const arcticdb::proto::descriptors::UserDefinedMetadata& user_meta,
             const arcticdb::proto::descriptors::UserDefinedMetadata& multi_key_meta,
             std::vector<entity::AtomKey>&& multi_keys) :
             item(versioned_item),
             frame_data(std::move(frame_data)),
+            output_format(output_format),
             norm_meta(norm_meta),
             user_meta(user_meta),
             multi_key_meta(multi_key_meta),
             multi_keys(std::move(multi_keys)) {
 
     }
     VersionedItem item;
-    pipelines::PythonOutputFrame frame_data;
+    OutputFrame frame_data;
+    OutputFormat output_format;
     arcticdb::proto::descriptors::NormalizationMetadata norm_meta;
     arcticdb::proto::descriptors::UserDefinedMetadata user_meta;
     arcticdb::proto::descriptors::UserDefinedMetadata multi_key_meta;
@@ -72,11 +79,18 @@ inline ReadResult create_python_read_result(
         }
     }
 
-    auto python_frame = pipelines::PythonOutputFrame{result.frame_, output_format};
+    // Is that the principal way to init a variant?
+    auto python_frame = [&]() -> OutputFrame {
+        if (output_format == OutputFormat::ARROW) {
+            return ArrowOutputFrame{segment_to_arrow_data(result.frame_), names_from_segment(result.frame_)};
+        } else {
+            return pipelines::PythonOutputFrame{result.frame_, output_format};
+        }
+    }();
     util::print_total_mem_usage(__FILE__, __LINE__, __FUNCTION__);
 
     const auto& desc_proto = result.desc_.proto();
-    return {version, std::move(python_frame), desc_proto.normalization(),
+    return {version, std::move(python_frame), output_format, desc_proto.normalization(),
             desc_proto.user_meta(), desc_proto.multi_key_meta(), std::move(result.keys_)};
 }
 

cpp/arcticdb/python/adapt_read_dataframe.hpp

@@ -17,12 +17,7 @@ inline auto adapt_read_df = [](ReadResult && ret) -> py::tuple{
     auto pynorm = python_util::pb_to_python(ret.norm_meta);
     auto pyuser_meta = python_util::pb_to_python(ret.user_meta);
     auto multi_key_meta = python_util::pb_to_python(ret.multi_key_meta);
-    return py::make_tuple(ret.item, std::move(ret.frame_data), pynorm, pyuser_meta, multi_key_meta, ret.multi_keys);
-};
-
-inline auto adapt_arrow_df = [](ArrowReadResult && ret) -> py::tuple{
-    auto pyuser_meta = python_util::pb_to_python(ret.user_meta_);
-    return py::make_tuple(ret.versioned_item_, std::move(ret.frame_), pyuser_meta);
+    return py::make_tuple(ret.item, std::move(ret.frame_data), ret.output_format, pynorm, pyuser_meta, multi_key_meta, ret.multi_keys);
 };
 
 }

cpp/arcticdb/python/python_utils.hpp

@@ -251,7 +251,7 @@ inline py::list adapt_read_dfs(std::vector<std::variant<ReadResult, DataError>>&
                 auto pynorm = python_util::pb_to_python(read_result.norm_meta);
                 auto pyuser_meta = python_util::pb_to_python(read_result.user_meta);
                 auto multi_key_meta = python_util::pb_to_python(read_result.multi_key_meta);
-                lst.append(py::make_tuple(read_result.item, std::move(read_result.frame_data), pynorm, pyuser_meta, multi_key_meta,
+                lst.append(py::make_tuple(read_result.item, std::move(read_result.frame_data), read_result.output_format, pynorm, pyuser_meta, multi_key_meta,
                                           read_result.multi_keys));
             },
             [&lst] (DataError& data_error) {

cpp/arcticdb/version/python_bindings.cpp

@@ -692,16 +692,9 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
              [&](PythonVersionStore& v,  StreamId sid, const VersionQuery& version_query, const std::shared_ptr<ReadQuery>& read_query, const ReadOptions& read_options) {
                 auto handler_data = TypeHandlerRegistry::instance()->get_handler_data(read_options.output_format());
                 return adapt_read_df(v.read_dataframe_version(sid, version_query, read_query, read_options, handler_data));
-              },
+             },
              py::call_guard<SingleThreadMutexHolder>(),
              "Read the specified version of the dataframe from the store")
-            .def("read_dataframe_version_arrow",
-            [&](PythonVersionStore& v,  StreamId sid, const VersionQuery& version_query, const std::shared_ptr<ReadQuery>& read_query, const ReadOptions& read_options) {
-                auto handler_data = TypeHandlerRegistry::instance()->get_handler_data(read_options.output_format());
-                return adapt_arrow_df(v.read_dataframe_version_arrow(sid, version_query, read_query, read_options, handler_data));
-            },
-            py::call_guard<SingleThreadMutexHolder>(),
-            "Read the specified version of the dataframe from the store")
         .def("read_index",
              [&](PythonVersionStore& v, StreamId sid, const VersionQuery& version_query){
                  return adapt_read_df(v.read_index(sid, version_query));
@@ -757,6 +750,7 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
                     vit,
                     PythonOutputFrame{
                         SegmentInMemory{tsd.as_stream_descriptor()},  read_options.output_format()},
+                        read_options.output_format(),
                         tsd_proto.normalization(),
                         tsd_proto.user_meta(),
                         tsd_proto.multi_key_meta(),
@@ -827,6 +821,7 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
                      const auto& tsd_proto = tsd.proto();
                      ReadResult res{vit, PythonOutputFrame{
                          SegmentInMemory{tsd.as_stream_descriptor()}, read_options.output_format()},
+                                    read_options.output_format(),
                                     tsd_proto.normalization(),
                                     tsd_proto.user_meta(),
                                     tsd_proto.multi_key_meta(), {}};

cpp/arcticdb/version/version_store_api.cpp

@@ -868,18 +868,6 @@ ReadResult PythonVersionStore::read_dataframe_version(
     return create_python_read_result(opt_version_and_frame.versioned_item_, read_options.output_format(), std::move(opt_version_and_frame.frame_and_descriptor_));
 }
 
-ArrowReadResult PythonVersionStore::read_dataframe_version_arrow(
-    const StreamId &stream_id,
-    const VersionQuery& version_query,
-    const std::shared_ptr<ReadQuery>& read_query,
-    const ReadOptions& read_options,
-    std::any& handler_data) {
-    ARCTICDB_RUNTIME_DEBUG(log::version(), "Command: read_dataframe_version_arrow");
-    util::check(read_options.output_format() == OutputFormat::ARROW, "Expected arrow format in read_dataframe_version_arrow");
-    auto opt_version_and_frame = read_dataframe_version_internal(stream_id, version_query, read_query, read_options, handler_data);
-    return create_arrow_read_result(opt_version_and_frame.versioned_item_, std::move(opt_version_and_frame.frame_and_descriptor_));
-}
-
 namespace {
 
 std::vector<SnapshotVariantKey> ARCTICDB_UNUSED iterate_snapshot_tombstones (

cpp/arcticdb/version/version_store_api.hpp

@@ -167,13 +167,6 @@ class PythonVersionStore : public LocalVersionedEngine {
         const ReadOptions& read_options,
         std::any& handler_data);
 
-    ArrowReadResult read_dataframe_version_arrow(
-        const StreamId &stream_id,
-        const VersionQuery& version_query,
-        const std::shared_ptr<ReadQuery>& read_query,
-        const ReadOptions& read_options,
-        std::any& handler_data);
-
     VersionedItem sort_merge(
             const StreamId& stream_id,
             const py::object& user_meta,

python/arcticdb/__init__.py

@@ -8,7 +8,7 @@
 from arcticdb.version_store._store import VersionedItem
 import arcticdb.version_store.library as library
 from arcticdb.tools import set_config_from_env_vars
-from arcticdb_ext.version_store import DataError, VersionRequestType
+from arcticdb_ext.version_store import DataError, VersionRequestType, OutputFormat
 from arcticdb_ext.exceptions import ErrorCode, ErrorCategory
 from arcticdb.version_store.library import (
     WritePayload,

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
 from arcticdb_ext.storage import LibraryManager
 from arcticdb.exceptions import LibraryNotFound, MismatchingLibraryOptions
 from arcticdb.version_store.library import ArcticInvalidApiUsageException, Library
@@ -90,7 +90,7 @@ def __init__(self, uri: str, encoding_version: EncodingVersion = DEFAULT_ENCODIN
         self._library_manager = LibraryManager(self._library_adapter.config_library)
         self._uri = uri
 
-    def __getitem__(self, name: str) -> Library:
+    def _get_library(self, name:str, runtime_options:Optional[RuntimeOptions]):
         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)
@@ -100,20 +100,24 @@ def __getitem__(self, name: str) -> Library:
             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) -> Library:
+        self._get_library(name, None)
+
     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, runtime_options: Optional[RuntimeOptions] = None
     ) -> Library:
         """
         Returns the library named ``name``.
@@ -136,6 +140,9 @@ def get_library(
             match these.
             Unused if create_if_missing is False.
 
+        runtime_options: Optional[RuntimeOptions], default = None
+            TODO
+
         Examples
         --------
         >>> arctic = adb.Arctic('s3://MY_ENDPOINT:MY_BUCKET')
@@ -152,7 +159,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, runtime_options)
             if create_if_missing and library_options:
                 if library_options.encoding_version is None:
                     library_options.encoding_version = self._encoding_version

python/arcticdb/options.py

@@ -11,6 +11,7 @@
 
 from arcticdb.encoding_version import EncodingVersion
 from arcticdb_ext.storage import ModifiableLibraryOption, ModifiableEnterpriseLibraryOption
+from arcticdb_ext.version_store import OutputFormat
 
 
 DEFAULT_ENCODING_VERSION = EncodingVersion.V1
@@ -146,6 +147,15 @@ def __repr__(self):
         )
 
 
+class RuntimeOptions:
+    """
+    TODO: Docs
+    """
+
+    def __init__(self, output_format=OutputFormat.PANDAS):
+        self.output_format=output_format
+
+
 class EnterpriseLibraryOptions:
     """
     Configuration options for ArcticDB libraries, that should only be used when you are using the ArcticDB enterprise