|
22 | 22 | #include <arcticdb/python/adapt_read_dataframe.hpp>
|
23 | 23 | #include <arcticdb/version/schema_checks.hpp>
|
24 | 24 | #include <arcticdb/util/pybind_mutex.hpp>
|
25 |
| - |
26 |
| -#include "python/python_handler_data.hpp" |
| 25 | +#include <arcticdb/python/python_handler_data.hpp> |
27 | 26 |
|
28 | 27 |
|
29 | 28 | namespace arcticdb::version_store {
|
@@ -233,7 +232,10 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
|
233 | 232 | version.def("write_dataframe_to_file", &write_dataframe_to_file);
|
234 | 233 | version.def("read_dataframe_from_file",
|
235 | 234 | [] (StreamId sid, std::string path, std::shared_ptr<ReadQuery>& read_query, const ReadOptions& read_options){
|
236 |
| - return adapt_read_df(read_dataframe_from_file(sid, path, read_query, read_options)); |
| 235 | + const OutputFormat output_format = read_options.output_format(); |
| 236 | + auto handler_data = TypeHandlerRegistry::instance()->get_handler_data(output_format); |
| 237 | + std::pair<std::any&, OutputFormat> handler{handler_data, output_format}; |
| 238 | + return adapt_read_df(read_dataframe_from_file(sid, path, read_query, read_options, handler_data), &handler); |
237 | 239 | });
|
238 | 240 |
|
239 | 241 | using FrameDataWrapper = arcticdb::pipelines::FrameDataWrapper;
|
@@ -572,7 +574,9 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
|
572 | 574 | py::call_guard<SingleThreadMutexHolder>(), "Drop column stats")
|
573 | 575 | .def("read_column_stats_version",
|
574 | 576 | [&](PythonVersionStore& v, StreamId sid, const VersionQuery& version_query){
|
575 |
| - return adapt_read_df(v.read_column_stats_version(sid, version_query)); |
| 577 | + auto handler_data = TypeHandlerRegistry::instance()->get_handler_data(OutputFormat::PANDAS); |
| 578 | + std::pair<std::any&, OutputFormat> handler{handler_data, OutputFormat::PANDAS}; |
| 579 | + return adapt_read_df(v.read_column_stats_version(sid, version_query, handler_data), &handler); |
576 | 580 | },
|
577 | 581 | py::call_guard<SingleThreadMutexHolder>(), "Read the column stats")
|
578 | 582 | .def("get_column_stats_info_version",
|
@@ -679,17 +683,19 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
|
679 | 683 | &PythonVersionStore::write_dataframe_specific_version,
|
680 | 684 | py::call_guard<SingleThreadMutexHolder>(), "Write a specific version of this dataframe to the store")
|
681 | 685 | .def("read_dataframe_version",
|
682 |
| - [&](PythonVersionStore& v, StreamId sid, const VersionQuery& version_query, const std::shared_ptr<ReadQuery>& read_query, const ReadOptions& read_options) { |
683 |
| - auto handler_data = TypeHandlerRegistry::instance()->get_handler_data(read_options.output_format()); |
684 |
| - py::list result = adapt_read_df(v.read_dataframe_version(sid, version_query, read_query, read_options, handler_data)); |
685 |
| - apply_global_refcounts(handler_data, read_options.output_format()); |
686 |
| - return result; |
| 686 | + [&](PythonVersionStore& v, StreamId sid, const VersionQuery& version_query, const std::shared_ptr<ReadQuery>& read_query, const ReadOptions& read_options) { |
| 687 | + const OutputFormat output_format = read_options.output_format(); |
| 688 | + auto handler_data = TypeHandlerRegistry::instance()->get_handler_data(output_format); |
| 689 | + std::pair<std::any&, OutputFormat> handler{handler_data, output_format}; |
| 690 | + return adapt_read_df(v.read_dataframe_version(sid, version_query, read_query, read_options, handler_data), &handler); |
687 | 691 | },
|
688 | 692 | py::call_guard<SingleThreadMutexHolder>(),
|
689 | 693 | "Read the specified version of the dataframe from the store")
|
690 | 694 | .def("read_index",
|
691 | 695 | [&](PythonVersionStore& v, StreamId sid, const VersionQuery& version_query){
|
692 |
| - return adapt_read_df(v.read_index(sid, version_query)); |
| 696 | + auto handler_data = TypeHandlerRegistry::instance()->get_handler_data(OutputFormat::PANDAS); |
| 697 | + std::pair<std::any&, OutputFormat> handler{handler_data, OutputFormat::PANDAS}; |
| 698 | + return adapt_read_df(v.read_index(sid, version_query, handler_data), &handler); |
693 | 699 | },
|
694 | 700 | py::call_guard<SingleThreadMutexHolder>(), "Read the most recent dataframe from the store")
|
695 | 701 | .def("get_update_time",
|
@@ -747,7 +753,7 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
|
747 | 753 | tsd_proto.multi_key_meta(),
|
748 | 754 | std::vector<entity::AtomKey>{}
|
749 | 755 | };
|
750 |
| - return adapt_read_df(std::move(res)); }, |
| 756 | + return adapt_read_df(std::move(res), nullptr); }, |
751 | 757 | py::call_guard<SingleThreadMutexHolder>(), "Restore a previous version of a symbol.")
|
752 | 758 | .def("check_ref_key",
|
753 | 759 | &PythonVersionStore::check_ref_key,
|
@@ -782,15 +788,18 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
|
782 | 788 | const std::vector<VersionQuery>& version_queries,
|
783 | 789 | std::vector<std::shared_ptr<ReadQuery>>& read_queries,
|
784 | 790 | const ReadOptions& read_options){
|
| 791 | + const OutputFormat output_format = read_options.output_format(); |
785 | 792 | auto handler_data = TypeHandlerRegistry::instance()->get_handler_data(read_options.output_format());
|
786 |
| - py::list result = python_util::adapt_read_dfs(v.batch_read(stream_ids, version_queries, read_queries, read_options)); |
787 |
| - apply_global_refcounts(handler_data, read_options.output_format()); |
788 |
| - return result; |
| 793 | + std::pair<std::any&, OutputFormat> handler{handler_data, output_format}; |
| 794 | + return python_util::adapt_read_dfs(v.batch_read(stream_ids, version_queries, read_queries, read_options, handler_data), &handler); |
789 | 795 | },
|
790 | 796 | py::call_guard<SingleThreadMutexHolder>(), "Read a dataframe from the store")
|
791 | 797 | .def("batch_read_keys",
|
792 | 798 | [&](PythonVersionStore& v, std::vector<AtomKey> atom_keys) {
|
793 |
| - return python_util::adapt_read_dfs(frame_to_read_result(v.batch_read_keys(atom_keys))); |
| 799 | + constexpr OutputFormat output_format = OutputFormat::PANDAS; |
| 800 | + auto handler_data = TypeHandlerRegistry::instance()->get_handler_data(output_format); |
| 801 | + std::pair<std::any&, OutputFormat> handler{handler_data, output_format}; |
| 802 | + return python_util::adapt_read_dfs(frame_to_read_result(v.batch_read_keys(atom_keys, handler_data)), &handler); |
794 | 803 | },
|
795 | 804 | py::call_guard<SingleThreadMutexHolder>(), "Read a specific version of a dataframe from the store")
|
796 | 805 | .def("batch_write",
|
@@ -821,7 +830,7 @@ void register_bindings(py::module &version, py::exception<arcticdb::ArcticExcept
|
821 | 830 | tsd_proto.user_meta(),
|
822 | 831 | tsd_proto.multi_key_meta(), {}};
|
823 | 832 |
|
824 |
| - output.emplace_back(adapt_read_df(std::move(res))); |
| 833 | + output.emplace_back(adapt_read_df(std::move(res), nullptr)); |
825 | 834 | }
|
826 | 835 | return output;
|
827 | 836 | },
|
|
0 commit comments