ustore_key_t key {ustore_key_unknown_k};

ustore_str_view_t field {nullptr};

ustore_bytes_ptr_t ptr {nullptr};

ustore_length_t off {0};

ustore_length_t len {0};

status_t status;

ustore_key_t key = py_to_scalar<ustore_key_t>(key_py);

value_view_t val = py_to_bytes(val_py);

[[maybe_unused]] py::gil_scoped_release release;

ustore_write_t write {};

write.db = collection.db();

write.error = status.member_ptr();

write.transaction = collection.txn();

write.arena = collection.member_arena();

write.options = collection.options();

write.tasks_count = 1;

write.collections = collection.member_collection();

write.keys = &key;

write.lengths = val.member_length();

write.values = val.member_ptr();

ustore_write(&write);

status.throw_unhandled();

status_t status;

parsed_places_t parsed_places {keys_py, collection.native};

places_arg_t places = parsed_places;

parsed_contents_t parsed_contents {vals_py};

contents_arg_t contents = parsed_contents;

[[maybe_unused]] py::gil_scoped_release release;

ustore_write_t write {};

write.db = collection.db();

write.error = status.member_ptr();

write.transaction = collection.txn();

write.arena = collection.member_arena();

write.options = collection.options();

write.tasks_count = places.count;

write.collections = collection.member_collection();

write.keys = places.keys_begin.get();

write.keys_stride = places.keys_begin.stride();

write.presences = contents.presences_begin.get();

write.offsets = contents.offsets_begin.get();

write.offsets_stride = contents.offsets_begin.stride();

write.lengths = contents.lengths_begin.get();

write.lengths_stride = contents.lengths_begin.stride();

write.values = contents.contents_begin.get();

write.values_stride = contents.contents_begin.stride();

ustore_write(&write);

status.throw_unhandled();

status_t status;

parsed_places_t parsed_places {keys_py.ptr(), collection.native};

places_arg_t places = parsed_places;

value_view_t val = py_to_bytes(val_py.ptr());

[[maybe_unused]] py::gil_scoped_release release;

ustore_write_t write {};

write.db = collection.db();

write.error = status.member_ptr();

write.transaction = collection.txn();

write.arena = collection.member_arena();

write.options = collection.options();

write.tasks_count = places.count;

write.collections = collection.member_collection();

write.keys = places.keys_begin.get();

write.keys_stride = places.keys_begin.stride();

write.lengths = val.member_length();

write.values = val.member_ptr();

ustore_write(&write);

status.throw_unhandled();

status_t status;

ustore_key_t key = py_to_scalar<ustore_key_t>(key_py);

ustore_octet_t* found_presences = nullptr;

{

[[maybe_unused]] py::gil_scoped_release release;

ustore_read_t read {};

read.db = collection.db();

read.error = status.member_ptr();

read.transaction = collection.txn();

read.arena = collection.member_arena();

read.options = collection.options();

read.tasks_count = 1;

read.collections = collection.member_collection();

read.keys = &key;

read.presences = &found_presences;

ustore_read(&read);

status.throw_unhandled();

}

// Exporting Arrow objects here would be over-engineering:

// if (collection.export_into_arrow()) {

// auto shared = std::make_shared<arrow::BooleanScalar>(presences[0]);

// PyObject* obj_ptr = arrow::py::wrap_scalar(std::static_pointer_cast<arrow::Scalar>(shared));

// return py::reinterpret_steal<py::object>(obj_ptr);

// }

bits_span_t presences {found_presences};

PyObject* obj_ptr = presences[0] ? Py_True : Py_False;

return py::reinterpret_borrow<py::object>(obj_ptr);

status_t status;

ustore_key_t key = py_to_scalar<ustore_key_t>(key_py);

ustore_length_t* found_lengths = nullptr;

ustore_bytes_ptr_t found_values = nullptr;

{

[[maybe_unused]] py::gil_scoped_release release;

ustore_read_t read {};

read.db = collection.db();

read.error = status.member_ptr();

read.transaction = collection.txn();

read.arena = collection.member_arena();

read.options = collection.options();

read.tasks_count = 1;

read.collections = collection.member_collection();

read.keys = &key;

read.lengths = &found_lengths;

read.values = &found_values;

ustore_read(&read);

status.throw_unhandled();

}

// To fetch data without copies, there is a hacky way:

// https://github.com/pybind/pybind11/issues/1236#issuecomment-527730864

// But in that case we can't guarantee memory alignment, so doing a copy

// is hard to avoid in Python.

// https://github.com/pybind/pybind11/blob/a05bc3d2359d12840ef2329d68f613f1a7df9c5d/include/pybind11/pytypes.h#L1474

// https://docs.python.org/3/c-api/bytes.html

// https://github.com/python/cpython/blob/main/Objects/bytesobject.c

value_view_t val {found_values, found_lengths[0]};

PyObject* obj_ptr = val ? PyBytes_FromStringAndSize(val.c_str(), val.size()) : Py_None;

return py::reinterpret_borrow<py::object>(obj_ptr);

status_t status;

ustore_octet_t* found_presences = nullptr;

parsed_places_t parsed_places {keys_py, collection.native};

places_arg_t places = parsed_places;

{

[[maybe_unused]] py::gil_scoped_release release;

ustore_read_t read {};

read.db = collection.db();

read.error = status.member_ptr();

read.transaction = collection.txn();

read.arena = collection.member_arena();

read.options = collection.options();

read.tasks_count = places.count;

read.collections = collection.member_collection();

read.keys = places.keys_begin.get();

read.keys_stride = places.keys_begin.stride();

read.presences = &found_presences;

ustore_read(&read);

status.throw_unhandled();

}

bits_span_t presences {found_presences};

PyObject* tuple_ptr = PyTuple_New(places.size());

for (std::size_t i = 0; i != places.size(); ++i) {

PyObject* obj_ptr = presences[i] ? Py_True : Py_False;

PyTuple_SetItem(tuple_ptr, i, obj_ptr);

}

return py::reinterpret_steal<py::object>(tuple_ptr);

status_t status;

ustore_octet_t* found_presences = nullptr;

ustore_length_t* found_offsets = nullptr;

ustore_length_t* found_lengths = nullptr;

ustore_bytes_ptr_t found_values = nullptr;

bool const export_arrow = collection.export_into_arrow();

parsed_places_t parsed_places {keys_py, collection.native};

places_arg_t places = parsed_places;

{

[[maybe_unused]] py::gil_scoped_release release;

ustore_read_t read {};

read.db = collection.db();

read.error = status.member_ptr();

read.transaction = collection.txn();

read.arena = collection.member_arena();

read.options = collection.options();

read.tasks_count = places.count;

read.collections = collection.member_collection();

read.collections_stride = 0;

read.keys = places.keys_begin.get();

read.keys_stride = places.keys_begin.stride();

read.presences = export_arrow ? &found_presences : nullptr;

read.offsets = &found_offsets;

read.lengths = !export_arrow ? &found_lengths : nullptr;

read.values = &found_values;

ustore_read(&read);

status.throw_unhandled();

}

if (export_arrow) {

auto shared_length = static_cast<int64_t>(places.count);

auto shared_offsets = std::make_shared<arrow::Buffer>( //

reinterpret_cast<uint8_t*>(found_offsets),

(shared_length + 1) * sizeof(ustore_length_t));

auto shared_data = std::make_shared<arrow::Buffer>( //

reinterpret_cast<uint8_t*>(found_values),

static_cast<int64_t>(found_offsets[places.count]));

auto shared_bitmap = std::make_shared<arrow::Buffer>( //

reinterpret_cast<uint8_t*>(found_presences),

divide_round_up<int64_t>(shared_length, CHAR_BIT));

auto shared = std::make_shared<arrow::BinaryArray>(shared_length, shared_offsets, shared_data, shared_bitmap);

PyObject* obj_ptr = arrow::py::wrap_array(std::static_pointer_cast<arrow::Array>(shared));

return py::reinterpret_steal<py::object>(obj_ptr);

}

else {

embedded_blobs_t bins {places.size(), found_offsets, found_lengths, found_values};

PyObject* tuple_ptr = PyTuple_New(places.size());

for (std::size_t i = 0; i != places.size(); ++i) {

value_view_t val = bins[i];

PyObject* obj_ptr = val ? PyBytes_FromStringAndSize(val.c_str(), val.size()) : Py_None;

PyTuple_SetItem(tuple_ptr, i, obj_ptr);

}

return py::reinterpret_steal<py::object>(tuple_ptr);

}

auto is_single = !PySequence_Check(key_py.ptr());

auto func = is_single ? &has_one_binary<collection_at> : &has_many_binaries<collection_at>;

return func(collection, key_py.ptr());

auto is_single = !PySequence_Check(key_py.ptr());

auto func = is_single ? &read_one_binary : &read_many_binaries;

return func(collection, key_py.ptr());

auto is_single = !PySequence_Check(key_py.ptr());

auto func = is_single ? &write_one_binary<collection_at> : &write_many_binaries<collection_at>;

return func(collection, key_py.ptr(), val_py.ptr());

auto is_single = !PySequence_Check(key_py.ptr());

auto func = is_single ? &write_one_binary<collection_at> : &write_many_binaries<collection_at>;

return func(collection, key_py.ptr(), Py_None);

status_t status;

ustore_size_t step = sizeof(py_bin_req_t);

std::vector<py_bin_req_t> keys;

keys.reserve(PyDict_Size(dict_py.ptr()));

std::size_t key_idx = 0;

py_scan_dict(dict_py.ptr(), [&](PyObject* key_obj, PyObject* val_obj) {

auto val = py_to_bytes(val_obj);

py_bin_req_t& req = keys[key_idx];

req.key = py_to_scalar<ustore_key_t>(key_obj);

req.ptr = ustore_bytes_ptr_t(val.begin());

req.len = static_cast<ustore_length_t>(val.size());

++key_idx;

});

[[maybe_unused]] py::gil_scoped_release release;

ustore_write_t write {};

write.db = collection.db();

write.error = status.member_ptr();

write.transaction = collection.txn();

write.arena = collection.member_arena();

write.options = collection.options();

write.tasks_count = static_cast<ustore_size_t>(keys.size());

write.collections = collection.member_collection();

write.keys = &keys[0].key;

write.keys_stride = step;

write.offsets = &keys[0].off;

write.offsets_stride = step;

write.lengths = &keys[0].len;

write.lengths_stride = step;

write.values = &keys[0].ptr;

write.values_stride = step;

ustore_write(&write);

status.throw_unhandled();

py_collection_gt<collection_at>& collection,

ustore_key_t min_key,

ustore_length_t count_limit) {

status_t status;

ustore_length_t* found_lengths = nullptr;

ustore_key_t* found_keys = nullptr;

bool const export_arrow = collection.export_into_arrow();

ustore_scan_t scan {};

scan.db = collection.db();

scan.error = status.member_ptr();

scan.transaction = collection.txn();

scan.arena = collection.member_arena();

scan.options = collection.options();

scan.tasks_count = 1;

scan.collections = collection.member_collection();

scan.start_keys = &min_key;

scan.count_limits = &count_limit;

scan.counts = &found_lengths;

scan.keys = &found_keys;

ustore_scan(&scan);

status.throw_unhandled();

if (export_arrow) {

auto shared_length = static_cast<int64_t>(found_lengths[0]);

auto shared_data = std::make_shared<arrow::Buffer>( //

reinterpret_cast<uint8_t*>(found_keys),

shared_length * sizeof(ustore_key_t));

static_assert(std::is_same_v<ustore_key_t, int64_t>, "Change the following line!");

auto shared = std::make_shared<arrow::NumericArray<arrow::Int64Type>>(shared_length, shared_data);

PyObject* obj_ptr = arrow::py::wrap_array(std::static_pointer_cast<arrow::Array>(shared));

return py::reinterpret_steal<py::object>(obj_ptr);

}

else

return py::array_t<ustore_key_t>(found_lengths[0], found_keys);

ustore_database_t db_ptr,

ustore_transaction_t txn_ptr,

ustore_collection_t collection_ptr,

managed_arena_t& arena,

py::handle keys_arr,

py::handle values_arr,

py::handle values_lengths_arr,

std::uint8_t padding_char = 0) {

// Check if we are receiving protocol buffers

PyObject* keys_obj = keys_arr.ptr();

PyObject* values_obj = values_arr.ptr();

PyObject* values_lengths_obj = values_lengths_arr.ptr();

if (!PyObject_CheckBuffer(keys_obj) | !PyObject_CheckBuffer(values_obj) | !PyObject_CheckBuffer(values_lengths_obj))

throw std::invalid_argument("All arguments must implement the buffer protocol");

// Take buffer protocol handles

// Flags can be: https://docs.python.org/3/c-api/buffer.html#readonly-format

auto output_flags = PyBUF_WRITABLE | PyBUF_ANY_CONTIGUOUS | PyBUF_STRIDED;

py_received_buffer_t keys, values, values_lengths;

keys.initialized = PyObject_GetBuffer(keys_obj, &keys.py, PyBUF_ANY_CONTIGUOUS) == 0;

values.initialized = PyObject_GetBuffer(values_obj, &values.py, output_flags) == 0;

values_lengths.initialized = PyObject_GetBuffer(values_lengths_obj, &values_lengths.py, output_flags) == 0;

if (!keys.initialized | !values.initialized | !values_lengths.initialized)

throw std::invalid_argument("Couldn't obtain buffer overviews");

if (!values.py.shape | !values_lengths.py.shape | !values.py.strides | !values_lengths.py.strides)

throw std::invalid_argument("Outputs shape wasn't inferred");

// Validate the format of `keys`

if (keys.py.itemsize != sizeof(ustore_key_t))

throw std::invalid_argument("Keys type mismatch");

if (keys.py.ndim != 1 || !PyBuffer_IsContiguous(&keys.py, 'A'))

throw std::invalid_argument("Keys must be placed in a continuous 1 dimensional array");

if (keys.py.strides[0] != sizeof(ustore_key_t))

throw std::invalid_argument("Keys can't be strided");

ustore_size_t const tasks_count = static_cast<ustore_size_t>(keys.py.len / keys.py.itemsize);

ustore_key_t const* keys_ptr = reinterpret_cast<ustore_key_t const*>(keys.py.buf);

// Validate the format of `values`

if (values.py.ndim != 2)

throw std::invalid_argument("Output tensor must have rank 2");

if (values.py.itemsize != sizeof(std::uint8_t))

throw std::invalid_argument("Output tensor must have single-byte entries");

if ((values.py.shape[0] <= 0) || values.py.shape[1] <= 0)

throw std::invalid_argument("Output tensor sides can't be zero");

if ((values.py.strides[0] <= 0) || values.py.strides[1] <= 0)

throw std::invalid_argument("Output tensor strides can't be negative");

if (tasks_count != static_cast<ustore_size_t>(values.py.shape[0]))

throw std::invalid_argument("Number of input keys and output slots doesn't match");

auto outputs_bytes = reinterpret_cast<std::uint8_t*>(values.py.buf);

auto outputs_bytes_stride = static_cast<std::size_t>(values.py.strides[0]);

auto output_bytes_cap = static_cast<ustore_length_t>(values.py.shape[1]);

// Validate the format of `values_lengths`

if (values_lengths.py.ndim != 1)

throw std::invalid_argument("Lengths tensor must have rank 1");

if (values_lengths.py.itemsize != sizeof(ustore_length_t))

throw std::invalid_argument("Lengths tensor must have 4-byte entries");

if (values_lengths.py.shape[0] <= 0)

throw std::invalid_argument("Lengths tensor sides can't be zero");

if (values_lengths.py.strides[0] <= 0)

throw std::invalid_argument("Lengths tensor strides can't be negative");

if (tasks_count != static_cast<ustore_size_t>(values_lengths.py.shape[0]))

throw std::invalid_argument("Number of input keys and output slots doesn't match");

auto outputs_lengths_bytes = reinterpret_cast<std::uint8_t*>(values_lengths.py.buf);

auto outputs_lengths_bytes_stride = static_cast<std::size_t>(values_lengths.py.strides[0]);

// Perform the read

[[maybe_unused]] py::gil_scoped_release release;

status_t status;

ustore_length_t* found_lengths = nullptr;

ustore_bytes_ptr_t found_values = nullptr;

ustore_options_t options = ustore_options_default_k;

ustore_read_t read {};

read.db = db_ptr;

read.error = status.internal_cptr();

read.transaction = txn_ptr;

read.arena = arena.internal_cptr();

read.options = options;

read.tasks_count = tasks_count;

read.collections = &collection_ptr;

read.keys = keys_ptr;

read.keys_stride = sizeof(ustore_key_t);

read.lengths = &found_lengths;

read.values = &found_values;

ustore_read(&read);

status.throw_unhandled();

// Export the data into the matrix

taped_values_view_t inputs {found_lengths, found_values, tasks_count};

tape_iterator_t input_it = inputs.begin();

for (ustore_size_t i = 0; i != tasks_count; ++i, ++input_it) {

value_view_t input = *input_it;

auto input_bytes = reinterpret_cast<std::uint8_t const*>(input.begin());

auto input_length = static_cast<ustore_length_t const>(input.size());

std::uint8_t* output_bytes = outputs_bytes + outputs_bytes_stride * i;

ustore_length_t& output_length =

*reinterpret_cast<ustore_length_t*>(outputs_lengths_bytes + outputs_lengths_bytes_stride * i);

std::size_t count_copy = std::min(output_bytes_cap, input_length);

std::size_t count_pads = output_bytes_cap - count_copy;

std::memcpy(output_bytes, input_bytes, count_copy);

std::memset(output_bytes + count_copy, padding_char, count_pads);

output_length = count_copy;

}