#include <chrono>
#include <cmath>
#include <cstring>
#include <functional>
#include <future>
#include <iostream>
#include <map>
#include <memory>
#include <string>
#include <vector>

#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>
#include "npbuffer.h"
#include "util.h"

#include <tiledb/tiledb.h>               // C
#include <tiledb/tiledb_experimental.h>  // C
#include <tiledb/tiledb>                 // C++
#include <tiledb/tiledb_experimental>    // C++

// clang-format off
// do not re-order these headers
#include "py_arrowio"
#include "py_arrow_io_impl.h"
// clang-format on

#if defined(TILEDB_SERIALIZATION)
#include <tiledb/tiledb_serialization.h>  // C
#endif

#include "../external/tsl/robin_map.h"

#if !defined(NDEBUG)
// #include "debug.cc"
#endif

#include "query_condition.cc"

namespace tiledbpy {

using namespace tiledb;
namespace py = pybind11;
using namespace pybind11::literals;

using TimerType = std::chrono::duration<double>;
struct StatsInfo {
    std::map<std::string, TimerType> counters;
};

bool config_has_key(tiledb::Config config, std::string key) {
    return config.contains(key);
}

struct PAPair {
    int64_t get_array() {
        if (!exported_) {
            TPY_ERROR_LOC("Cannot export uninitialized array!");
        }
        return (int64_t)&array_;
    };
    int64_t get_schema() {
        if (!exported_) {
            TPY_ERROR_LOC("Cannot export uninitialized schema!");
        }
        return (int64_t)&schema_;
    }

    ArrowSchema schema_;
    ArrowArray array_;
    bool exported_ = false;
};

// global stats counters
static std::unique_ptr<StatsInfo> g_stats;

// forward declaration
py::dtype tiledb_dtype(tiledb_datatype_t type, uint32_t cell_val_num);

struct BufferInfo {
    BufferInfo(
        std::string name,
        size_t data_nbytes,
        tiledb_datatype_t data_type,
        uint32_t cell_val_num,
        size_t offsets_num,
        size_t validity_num,
        bool isvar = false,
        bool isnullable = false)

        : name(name)
        , type(data_type)
        , cell_val_num(cell_val_num)
        , isvar(isvar)
        , isnullable(isnullable) {
        try {
            dtype = tiledb_dtype(data_type, cell_val_num);
            elem_nbytes = tiledb_datatype_size(type);
            data = py::array(py::dtype("uint8"), data_nbytes);
            offsets = py::array_t<uint64_t>(offsets_num);
            validity = py::array_t<uint8_t>(validity_num);
        } catch (py::error_already_set& e) {
            TPY_ERROR_LOC(e.what())
        }
        // TODO use memset here for zero'd buffers in debug mode
    }

    string name;
    tiledb_datatype_t type;
    py::dtype dtype;
    size_t elem_nbytes = 1;
    uint64_t data_vals_read = 0;
    uint32_t cell_val_num;
    uint64_t offsets_read = 0;
    uint64_t validity_vals_read = 0;
    bool isvar;
    bool isnullable;

    py::array data;
    py::array_t<uint64_t> offsets;
    py::array_t<uint8_t> validity;
};

struct BufferHolder {
    py::object data;
    py::object offsets;
    py::object validity;

    BufferHolder(py::object d, py::object o, py::object v)
        : data(d)
        , offsets(o)
        , validity(v) {
    }

    static void free_buffer_holder(BufferHolder* self) {
        delete self;
    }
};

py::dtype tiledb_dtype(tiledb_datatype_t type, uint32_t cell_val_num) {
    if (cell_val_num == 1) {
        auto np = py::module::import("numpy");
        auto datetime64 = np.attr("datetime64");
        switch (type) {
            case TILEDB_INT32:
                return py::dtype("int32");
            case TILEDB_INT64:
                return py::dtype("int64");
            case TILEDB_FLOAT32:
                return py::dtype("float32");
            case TILEDB_FLOAT64:
                return py::dtype("float64");
            case TILEDB_INT8:
                return py::dtype("int8");
            case TILEDB_UINT8:
                return py::dtype("uint8");
            case TILEDB_INT16:
                return py::dtype("int16");
            case TILEDB_UINT16:
                return py::dtype("uint16");
            case TILEDB_UINT32:
                return py::dtype("uint32");
            case TILEDB_UINT64:
                return py::dtype("uint64");
            case TILEDB_STRING_ASCII:
                return py::dtype("S1");
            case TILEDB_STRING_UTF8:
                return py::dtype("U1");
            case TILEDB_STRING_UTF16:
            case TILEDB_STRING_UTF32:
                TPY_ERROR_LOC(
                    "Unimplemented UTF16 or UTF32 string conversion!");
            case TILEDB_STRING_UCS2:
            case TILEDB_STRING_UCS4:
                TPY_ERROR_LOC("Unimplemented UCS2 or UCS4 string conversion!");
            case TILEDB_CHAR:
                return py::dtype("S1");
            case TILEDB_DATETIME_YEAR:
                return py::dtype("M8[Y]");
            case TILEDB_DATETIME_MONTH:
                return py::dtype("M8[M]");
            case TILEDB_DATETIME_WEEK:
                return py::dtype("M8[W]");
            case TILEDB_DATETIME_DAY:
                return py::dtype("M8[D]");
            case TILEDB_DATETIME_HR:
                return py::dtype("M8[h]");
            case TILEDB_DATETIME_MIN:
                return py::dtype("M8[m]");
            case TILEDB_DATETIME_SEC:
                return py::dtype("M8[s]");
            case TILEDB_DATETIME_MS:
                return py::dtype("M8[ms]");
            case TILEDB_DATETIME_US:
                return py::dtype("M8[us]");
            case TILEDB_DATETIME_NS:
                return py::dtype("M8[ns]");
            case TILEDB_DATETIME_PS:
                return py::dtype("M8[ps]");
            case TILEDB_DATETIME_FS:
                return py::dtype("M8[fs]");
            case TILEDB_DATETIME_AS:
                return py::dtype("M8[as]");

            /* duration types map to timedelta */
            case TILEDB_TIME_HR:
                return py::dtype("m8[h]");
            case TILEDB_TIME_MIN:
                return py::dtype("m8[m]");
            case TILEDB_TIME_SEC:
                return py::dtype("m8[s]");
            case TILEDB_TIME_MS:
                return py::dtype("m8[ms]");
            case TILEDB_TIME_US:
                return py::dtype("m8[us]");
            case TILEDB_TIME_NS:
                return py::dtype("m8[ns]");
            case TILEDB_TIME_PS:
                return py::dtype("m8[ps]");
            case TILEDB_TIME_FS:
                return py::dtype("m8[fs]");
            case TILEDB_TIME_AS:
                return py::dtype("m8[as]");
            case TILEDB_BLOB:
                return py::dtype("byte");
            case TILEDB_BOOL:
                return py::dtype("bool");
#if TILEDB_VERSION_MAJOR >= 2 && TILEDB_VERSION_MINOR >= 21
            case TILEDB_GEOM_WKB:
                return py::dtype("byte");
            case TILEDB_GEOM_WKT:
                return py::dtype("S");
#endif

            case TILEDB_ANY:
                break;
        }
    } else if (cell_val_num == 2 && type == TILEDB_FLOAT32) {
        return py::dtype("complex64");
    } else if (cell_val_num == 2 && type == TILEDB_FLOAT64) {
        return py::dtype("complex128");
    } else if (
        type == TILEDB_CHAR || type == TILEDB_STRING_UTF8 ||
        type == TILEDB_STRING_ASCII) {
        std::string base_str;
        switch (type) {
            case TILEDB_CHAR:
            case TILEDB_STRING_ASCII:
                base_str = "|S";
                break;
            case TILEDB_STRING_UTF8:
                base_str = "|U";
                break;
            default:
                TPY_ERROR_LOC("internal error: unhandled string type");
        }
        if (cell_val_num < TILEDB_VAR_NUM) {
            base_str = base_str + std::to_string(cell_val_num);
        }
        return py::dtype(base_str);
    } else if (cell_val_num == TILEDB_VAR_NUM) {
        return tiledb_dtype(type, 1);
    } else if (cell_val_num > 1) {
        py::dtype base_dtype = tiledb_dtype(type, 1);
        py::tuple rec_elem = py::make_tuple("", base_dtype);
        py::list rec_list;
        for (size_t i = 0; i < cell_val_num; i++)
            rec_list.append(rec_elem);
        auto np = py::module::import("numpy");
        // note: we call the 'dtype' constructor b/c py::dtype does not accept
        // list
        auto np_dtype = np.attr("dtype");
        return np_dtype(rec_list);
    }

    TPY_ERROR_LOC(
        "tiledb datatype not understood ('" + tiledb::impl::type_to_str(type) +
        "', cell_val_num: " + std::to_string(cell_val_num) + ")");
}

py::array_t<uint8_t> uint8_bool_to_uint8_bitmap(
    py::array_t<uint8_t> validity_array) {
    // TODO profile, probably replace; avoid inplace reassignment
    auto np = py::module::import("numpy");
    auto packbits = np.attr("packbits");
    auto tmp = packbits(validity_array, "bitorder"_a = "little");
    return tmp;
}

uint64_t count_zeros(py::array_t<uint8_t> a) {
    uint64_t count = 0;
    for (Py_ssize_t idx = 0; idx < a.size(); idx++)
        count += (a.data()[idx] == 0) ? 1 : 0;
    return count;
}

class PyAgg {
    using ByteBuffer = py::array_t<uint8_t>;
    using AggToBufferMap = std::map<std::string, ByteBuffer>;
    using AttrToAggsMap = std::map<std::string, AggToBufferMap>;

   private:
    Context ctx_;
    std::shared_ptr<tiledb::Array> array_;
    std::shared_ptr<tiledb::Query> query_;
    AttrToAggsMap result_buffers_;
    AttrToAggsMap validity_buffers_;

    py::dict original_input_;
    std::vector<std::string> attrs_;

   public:
    PyAgg() = delete;

    PyAgg(
        const Context& ctx,
        py::object py_array,
        py::object py_layout,
        py::dict attr_to_aggs_input)
        : ctx_(ctx)
        , original_input_(attr_to_aggs_input) {
        tiledb_array_t* c_array_ = (py::capsule)py_array.attr("__capsule__")();

        // We never own this pointer; pass own=false
        array_ = std::make_shared<tiledb::Array>(ctx_, c_array_, false);
        query_ = std::make_shared<tiledb::Query>(ctx_, *array_, TILEDB_READ);

        bool issparse = array_->schema().array_type() == TILEDB_SPARSE;
        tiledb_layout_t layout = (tiledb_layout_t)py_layout.cast<int32_t>();
        if (!issparse && layout == TILEDB_UNORDERED) {
            TPY_ERROR_LOC(
                "TILEDB_UNORDERED read is not supported for dense arrays")
        }
        query_->set_layout(layout);

        // Iterate through the requested attributes
        for (auto attr_to_aggs : attr_to_aggs_input) {
            auto attr_name = attr_to_aggs.first.cast<std::string>();
            auto aggs = attr_to_aggs.second.cast<std::vector<std::string>>();

            tiledb::Attribute attr = array_->schema().attribute(attr_name);
            attrs_.push_back(attr_name);

            // For non-nullable attributes, applying max and min to the empty
            // set is undefined. To check for this, we need to also run the
            // count aggregate to make sure count != 0
            bool requested_max = std::find(aggs.begin(), aggs.end(), "max") !=
                                 aggs.end();
            bool requested_min = std::find(aggs.begin(), aggs.end(), "min") !=
                                 aggs.end();
            if (!attr.nullable() && (requested_max || requested_min)) {
                // If the user already also requested count, then we don't need
                // to request it again
                if (std::find(aggs.begin(), aggs.end(), "count") ==
                    aggs.end()) {
                    aggs.push_back("count");
                }
            }

            // Iterate through the aggreate operations to apply on the given
            // attribute
            for (auto agg_name : aggs) {
                _apply_agg_operator_to_attr(agg_name, attr_name);

                // Set the result data buffers
                auto* res_buf = &result_buffers_[attr_name][agg_name];
                if ("count" == agg_name || "null_count" == agg_name ||
                    "mean" == agg_name) {
                    // count and null_count use uint64 and mean uses float64
                    *res_buf = py::array(py::dtype("uint8"), 8);
                } else {
                    // max, min, and sum use the dtype of the attribute
                    py::dtype dt(tiledb_dtype(attr.type(), attr.cell_size()));
                    *res_buf = py::array(py::dtype("uint8"), dt.itemsize());
                }
                query_->set_data_buffer(
                    attr_name + agg_name, (void*)res_buf->data(), 1);

                if (attr.nullable()) {
                    // For nullable attributes, if the input set for the
                    // aggregation contains all NULL values, we will not get an
                    // aggregate value back as this operation is undefined. We
                    // need to check the validity buffer beforehand to see if we
                    // had a valid result
                    if (!("count" == agg_name || "null_count" == agg_name)) {
                        auto* val_buf = &validity_buffers_[attr.name()]
                                                          [agg_name];
                        *val_buf = py::array(py::dtype("uint8"), 1);
                        query_->set_validity_buffer(
                            attr_name + agg_name, (uint8_t*)val_buf->data(), 1);
                    }
                }
            }
        }
    }

    void _apply_agg_operator_to_attr(
        const std::string& op_label, const std::string& attr_name) {
        using AggregateFunc = std::function<ChannelOperation(
            const Query&, const std::string&)>;

        std::unordered_map<std::string, AggregateFunc> label_to_agg_func = {
            {"sum", QueryExperimental::create_unary_aggregate<SumOperator>},
            {"min", QueryExperimental::create_unary_aggregate<MinOperator>},
            {"max", QueryExperimental::create_unary_aggregate<MaxOperator>},
            {"mean", QueryExperimental::create_unary_aggregate<MeanOperator>},
            {"null_count",
             QueryExperimental::create_unary_aggregate<NullCountOperator>},
        };

        QueryChannel default_channel = QueryExperimental::get_default_channel(
            *query_);

        if (label_to_agg_func.find(op_label) != label_to_agg_func.end()) {
            AggregateFunc create_unary_aggregate = label_to_agg_func.at(
                op_label);
            ChannelOperation op = create_unary_aggregate(*query_, attr_name);
            default_channel.apply_aggregate(attr_name + op_label, op);
        } else if ("count" == op_label) {
            default_channel.apply_aggregate(
                attr_name + op_label, CountOperation());
        } else {
            TPY_ERROR_LOC(
                "Invalid channel operation " + op_label +
                " passed to apply_aggregate.");
        }
    }

    py::dict get_aggregate() {
        query_->submit();

        // Cast the results to the correct dtype and output this as a Python
        // dict
        py::dict output;
        for (auto attr_to_agg : original_input_) {
            // Be clear in our variable names for strings as py::dict uses
            // py::str keys whereas std::map uses std::string keys
            std::string attr_cpp_name = attr_to_agg.first.cast<string>();

            py::str attr_py_name(attr_cpp_name);
            output[attr_py_name] = py::dict();

            tiledb::Attribute attr = array_->schema().attribute(attr_cpp_name);

            for (auto agg_py_name : original_input_[attr_py_name]) {
                std::string agg_cpp_name = agg_py_name.cast<string>();

                if (_is_invalid(attr, agg_cpp_name)) {
                    output[attr_py_name][agg_py_name] = _is_integer_dtype(
                                                            attr) ?
                                                            py::none() :
                                                            py::cast(NAN);
                } else {
                    output[attr_py_name][agg_py_name] = _set_result(
                        attr, agg_cpp_name);
                }
            }
        }
        return output;
    }

    bool _is_invalid(tiledb::Attribute attr, std::string agg_name) {
        if (attr.nullable()) {
            if ("count" == agg_name || "null_count" == agg_name)
                return false;

            // For nullable attributes, check if the validity buffer returned
            // false
            const void* val_buf = validity_buffers_[attr.name()][agg_name]
                                      .data();
            return *((uint8_t*)(val_buf)) == 0;
        } else {
            // For non-nullable attributes, max and min are undefined for the
            // empty set, so we must check the count == 0
            if ("max" == agg_name || "min" == agg_name) {
                const void* count_buf = result_buffers_[attr.name()]["count"]
                                            .data();
                return *((uint64_t*)(count_buf)) == 0;
            }
            return false;
        }
    }

    bool _is_integer_dtype(tiledb::Attribute attr) {
        switch (attr.type()) {
            case TILEDB_INT8:
            case TILEDB_INT16:
            case TILEDB_UINT8:
            case TILEDB_INT32:
            case TILEDB_INT64:
            case TILEDB_UINT16:
            case TILEDB_UINT32:
            case TILEDB_UINT64:
                return true;
            default:
                return false;
        }
    }

    py::object _set_result(tiledb::Attribute attr, std::string agg_name) {
        const void* agg_buf = result_buffers_[attr.name()][agg_name].data();

        if ("mean" == agg_name)
            return py::cast(*((double*)agg_buf));

        if ("count" == agg_name || "null_count" == agg_name)
            return py::cast(*((uint64_t*)agg_buf));

        switch (attr.type()) {
            case TILEDB_FLOAT32:
                return py::cast(
                    "sum" == agg_name ? *((double*)agg_buf) :
                                        *((float*)agg_buf));
            case TILEDB_FLOAT64:
                return py::cast(*((double*)agg_buf));
            case TILEDB_INT8:
                return py::cast(*((int8_t*)agg_buf));
            case TILEDB_UINT8:
                return py::cast(*((uint8_t*)agg_buf));
            case TILEDB_INT16:
                return py::cast(*((int16_t*)agg_buf));
            case TILEDB_UINT16:
                return py::cast(*((uint16_t*)agg_buf));
            case TILEDB_UINT32:
                return py::cast(*((uint32_t*)agg_buf));
            case TILEDB_INT32:
                return py::cast(*((int32_t*)agg_buf));
            case TILEDB_INT64:
                return py::cast(*((int64_t*)agg_buf));
            case TILEDB_UINT64:
                return py::cast(*((uint64_t*)agg_buf));
            default:
                TPY_ERROR_LOC(
                    "[_cast_agg_result] Invalid tiledb dtype for aggregation "
                    "result")
        }
    }

    void set_subarray(py::object py_subarray) {
        query_->set_subarray(*py_subarray.cast<tiledb::Subarray*>());
    }

    void set_cond(py::object cond) {
        py::object init_pyqc = cond.attr("init_query_condition");

        try {
            init_pyqc(array_->uri(), attrs_, ctx_);
        } catch (tiledb::TileDBError& e) {
            TPY_ERROR_LOC(e.what());
        } catch (py::error_already_set& e) {
            TPY_ERROR_LOC(e.what());
        }
        auto pyqc = (cond.attr("c_obj")).cast<PyQueryCondition>();
        auto qc = pyqc.ptr().get();
        query_->set_condition(*qc);
    }
};

class PyQuery {
   private:
    Context ctx_;
    std::shared_ptr<tiledb::Domain> domain_;
    std::shared_ptr<tiledb::ArraySchema> array_schema_;
    std::shared_ptr<tiledb::Array> array_;
    std::shared_ptr<tiledb::Query> query_;
    std::vector<std::string> attrs_;
    std::vector<std::string> dims_;
    std::map<std::string, BufferInfo> buffers_;
    std::vector<std::string> buffers_order_;

    bool deduplicate_ = true;
    bool use_arrow_ = false;
    // initialize the query buffers with exactly `init_buffer_bytes`
    // rather than the estimated result size. for incomplete testing.
    bool exact_init_bytes_ = false;
    uint64_t init_buffer_bytes_ = DEFAULT_INIT_BUFFER_BYTES;
    uint64_t alloc_max_bytes_ = DEFAULT_ALLOC_MAX_BYTES;
    tiledb_layout_t layout_ = TILEDB_ROW_MAJOR;

    // label buffer list
    unordered_map<string, uint64_t> label_input_buffer_data_;

   public:
    tiledb_ctx_t* c_ctx_;
    tiledb_array_t* c_array_;
    bool preload_metadata_ = false;
    bool return_incomplete_ = false;
    size_t retries_ = 0;

   public:
    PyQuery() = delete;

    PyQuery(
        const Context& ctx,
        py::object array,
        py::iterable attrs,
        py::iterable dims,
        py::object py_layout,
        py::object use_arrow)
        : ctx_(ctx) {
        init_config();
        // initialize arrow argument from user, if provided
        // call after init_config
        if (!use_arrow.is(py::none())) {
            use_arrow_ = py::cast<bool>(use_arrow);
        }

        tiledb_array_t* c_array_ = (py::capsule)array.attr("__capsule__")();

        // we never own this pointer, pass own=false
        array_ = std::make_shared<tiledb::Array>(ctx_, c_array_, false);

        array_schema_ = std::make_shared<tiledb::ArraySchema>(array_->schema());

        domain_ = std::make_shared<tiledb::Domain>(array_schema_->domain());

        bool issparse = array_->schema().array_type() == TILEDB_SPARSE;

        std::string mode = py::str(array.attr("mode"));
        auto query_mode = TILEDB_READ;
        if (mode == "r") {
            query_mode = TILEDB_READ;
        } else if (mode == "d") {
            query_mode = TILEDB_DELETE;
        } else {
            throw std::invalid_argument("Invalid query mode: " + mode);
        }

        // initialize the dims that we are asked to read
        for (auto d : dims) {
            dims_.push_back(d.cast<string>());
        }

        // initialize the attrs that we are asked to read
        for (auto a : attrs) {
            attrs_.push_back(a.cast<string>());
        }

        if (query_mode == TILEDB_READ) {
            py::object pre_buffers = array.attr("_buffers");
            if (!pre_buffers.is(py::none())) {
                py::dict pre_buffers_dict = pre_buffers.cast<py::dict>();

                // iterate over (key, value) pairs
                for (std::pair<py::handle, py::handle> b : pre_buffers_dict) {
                    py::str name = b.first.cast<py::str>();

                    // unpack value tuple of (data, offsets)
                    auto bfrs = b.second
                                    .cast<std::pair<py::handle, py::handle>>();
                    auto data_array = bfrs.first.cast<py::array>();
                    auto offsets_array = bfrs.second.cast<py::array>();

                    import_buffer(name, data_array, offsets_array);
                }
            }
        }

        query_ = std::make_shared<tiledb::Query>(ctx_, *array_, query_mode);
        //        [](Query* p){} /* note: no deleter*/);

        if (query_mode == TILEDB_READ) {
            layout_ = (tiledb_layout_t)py_layout.cast<int32_t>();
            if (!issparse && layout_ == TILEDB_UNORDERED) {
                TPY_ERROR_LOC(
                    "TILEDB_UNORDERED read is not supported for dense arrays")
            }
            query_->set_layout(layout_);
        }

        if (use_arrow_) {
            // enable arrow mode in the Query
            auto tmp_config = ctx_.config();
            tmp_config.set("sm.var_offsets.bitsize", "64");
            tmp_config.set("sm.var_offsets.mode", "elements");
            tmp_config.set("sm.var_offsets.extra_element", "true");
            ctx_.handle_error(tiledb_query_set_config(
                ctx_.ptr().get(), query_->ptr().get(), tmp_config.ptr().get()));
        }
    }

    void set_subarray(py::object py_subarray) {
        query_->set_subarray(*py_subarray.cast<tiledb::Subarray*>());
    }

#if defined(TILEDB_SERIALIZATION)
    void set_serialized_query(py::buffer serialized_query) {
        int rc;
        tiledb_query_t* c_query;
        tiledb_buffer_t* c_buffer;
        tiledb_ctx_t* c_ctx = ctx_.ptr().get();

        rc = tiledb_buffer_alloc(c_ctx, &c_buffer);
        if (rc == TILEDB_ERR)
            TPY_ERROR_LOC("Could not allocate c_buffer.");

        py::buffer_info buffer_info = serialized_query.request();
        rc = tiledb_buffer_set_data(
            c_ctx, c_buffer, buffer_info.ptr, buffer_info.shape[0]);
        if (rc == TILEDB_ERR)
            TPY_ERROR_LOC("Could not set c_buffer.");

        c_query = query_.get()->ptr().get();
        rc = tiledb_deserialize_query(
            c_ctx, c_buffer, TILEDB_CAPNP, 0, c_query);
        if (rc == TILEDB_ERR)
            TPY_ERROR_LOC("Could not deserialize query.");
    }
#endif

    void set_cond(py::object cond) {
        py::object init_pyqc = cond.attr("init_query_condition");

        try {
            init_pyqc(array_->uri(), attrs_, ctx_);
        } catch (tiledb::TileDBError& e) {
            TPY_ERROR_LOC(e.what());
        } catch (py::error_already_set& e) {
            TPY_ERROR_LOC(e.what());
        }
        auto pyqc = (cond.attr("c_obj")).cast<PyQueryCondition>();
        auto qc = pyqc.ptr().get();
        query_->set_condition(*qc);
    }

    bool is_dimension(std::string name) {
        return domain_->has_dimension(name);
    }

    bool is_attribute(std::string name) {
        return array_schema_->has_attribute(name);
    }

    bool is_dimension_label(std::string name) {
        return ArraySchemaExperimental::has_dimension_label(
            ctx_, *array_schema_, name);
    }

    bool is_var(std::string name) {
        if (is_dimension(name)) {
            auto dim = domain_->dimension(name);
            return dim.cell_val_num() == TILEDB_VAR_NUM;
        } else if (is_attribute(name)) {
            auto attr = array_schema_->attribute(name);
            return attr.cell_val_num() == TILEDB_VAR_NUM;
        } else if (is_dimension_label(name)) {
            auto dim_label = ArraySchemaExperimental::dimension_label(
                ctx_, *array_schema_, name);
            return dim_label.label_cell_val_num() == TILEDB_VAR_NUM;
        } else {
            TPY_ERROR_LOC(
                "Unknown buffer type for is_var check (expected attribute "
                "or dimension)")
        }
    }

    bool is_nullable(std::string name) {
        if (is_dimension(name) || is_dimension_label(name)) {
            return false;
        }

        auto attr = array_schema_->attribute(name);
        return attr.nullable();
    }

    std::pair<tiledb_datatype_t, uint32_t> buffer_type(std::string name) {
        tiledb_datatype_t type;
        uint32_t cell_val_num;
        if (is_dimension(name)) {
            type = domain_->dimension(name).type();
            cell_val_num = domain_->dimension(name).cell_val_num();
        } else if (is_attribute(name)) {
            type = array_schema_->attribute(name).type();
            cell_val_num = array_schema_->attribute(name).cell_val_num();
        } else if (is_dimension_label(name)) {
            auto dim_label = ArraySchemaExperimental::dimension_label(
                ctx_, *array_schema_, name);
            type = dim_label.label_type();
            cell_val_num = dim_label.label_cell_val_num();
        } else {
            TPY_ERROR_LOC("Unknown buffer '" + name + "'");
        }
        return {type, cell_val_num};
    }

    uint32_t buffer_ncells(std::string name) {
        if (is_dimension(name)) {
            return domain_->dimension(name).cell_val_num();
        } else if (is_attribute(name)) {
            return array_schema_->attribute(name).cell_val_num();
        }
        TPY_ERROR_LOC("Unknown buffer '" + name + "' for buffer_ncells");
    }

    py::dtype buffer_dtype(std::string name) {
        try {
            auto t = buffer_type(name);
            return tiledb_dtype(t.first, t.second);
        } catch (TileDBError& e) {
            (void)e;
            return py::none();
        }
    }

    bool is_sparse() {
        return array_->schema().array_type() == TILEDB_SPARSE;
    }

    void import_buffer(std::string name, py::array data, py::array offsets) {
        tiledb_datatype_t type;
        uint32_t cell_val_num;
        std::tie(type, cell_val_num) = buffer_type(name);
        auto dtype = tiledb_dtype(type, cell_val_num);

        buffers_order_.push_back(name);
        // set nbytes and noffsets=0 here to avoid allocation; buffers set below
        auto buffer_info = BufferInfo(
            name,
            0,
            type,
            cell_val_num,
            0,
            0,  // TODO
            is_var(name),
            is_nullable(name));
        buffer_info.data = data;
        buffer_info.offsets = offsets;
        buffers_.insert({name, buffer_info});
    }

    void alloc_buffer(std::string name) {
        tiledb_datatype_t type;
        uint32_t cell_val_num;
        uint64_t cell_nbytes;
        bool var;
        bool nullable;
        uint64_t buf_nbytes = 0;
        uint64_t offsets_num = 0;
        uint64_t validity_num = 0;
        bool dense = array_schema_->array_type() == TILEDB_DENSE;
        if (is_dimension_label(name)) {
            auto dim_label = ArraySchemaExperimental::dimension_label(
                ctx_, *array_schema_, name);
            type = dim_label.label_type();
            cell_val_num = dim_label.label_cell_val_num();
            var = cell_val_num == TILEDB_VAR_NUM;
            nullable = false;

            cell_nbytes = tiledb_datatype_size(type);
            uint64_t ncells = label_input_buffer_data_[name];

            if (!var) {
                cell_nbytes *= cell_val_num;
            } else {
                offsets_num = ncells;
            }
            buf_nbytes = ncells * cell_nbytes;
        } else {
            std::tie(type, cell_val_num) = buffer_type(name);
            cell_nbytes = tiledb_datatype_size(type);
            if (cell_val_num != TILEDB_VAR_NUM) {
                cell_nbytes *= cell_val_num;
            }
            var = is_var(name);
            nullable = is_nullable(name);

            if (retries_ < 1 && dense) {
                // we must not call after submitting
                if (nullable && var) {
                    auto sizes = query_->est_result_size_var_nullable(name);
                    offsets_num = sizes[0];
                    buf_nbytes = sizes[1];
                    validity_num = sizes[2] / sizeof(uint8_t);
                } else if (nullable && !var) {
                    auto sizes = query_->est_result_size_nullable(name);
                    buf_nbytes = sizes[0];
                    validity_num = sizes[1] / sizeof(uint8_t);
                } else if (!nullable && var) {
                    auto size_pair = query_->est_result_size_var(name);
                    buf_nbytes = size_pair[0];
                    offsets_num = size_pair[1];
                } else {  // !nullable && !var
                    buf_nbytes = query_->est_result_size(name);
                }

                // Add extra offset to estimate in order to avoid incomplete
                // resubmit libtiledb 2.7.* does not include extra element in
                // estimate. Remove this section after resolution of SC-16301.
                offsets_num += (var && use_arrow_) ? 1 : 0;
            }
        }

        // - for sparse arrays: don't try to allocate more than alloc_max_bytes_
        // - for dense arrays: the estimate should be exact, so don't cap
        if (is_sparse() && buf_nbytes > alloc_max_bytes_) {
            buf_nbytes = alloc_max_bytes_;
        }
        // use max to avoid overflowing to zero in the multiplication, in case
        // the
        //   estimate is too large
        if (max(validity_num, validity_num * sizeof(uint8_t)) >
            alloc_max_bytes_) {
            validity_num = alloc_max_bytes_ / sizeof(uint8_t);
        }
        if (max(offsets_num, offsets_num * sizeof(uint64_t)) >
            alloc_max_bytes_) {
            offsets_num = alloc_max_bytes_ / sizeof(uint64_t);
        }

        // use init_buffer_bytes configuration option if the
        // estimate is smaller
        if ((var || is_sparse()) &&
            (buf_nbytes < init_buffer_bytes_ || exact_init_bytes_)) {
            buf_nbytes = init_buffer_bytes_;
            offsets_num = init_buffer_bytes_ / sizeof(uint64_t);
            validity_num = init_buffer_bytes_ / cell_nbytes;
        }

        buffers_order_.push_back(name);
        buffers_.insert(
            {name,
             BufferInfo(
                 name,
                 buf_nbytes,
                 type,
                 cell_val_num,
                 offsets_num,
                 validity_num,
                 var,
                 nullable)});
    }

    void add_label_buffer(std::string& label_name, uint64_t ncells) {
        label_input_buffer_data_[label_name] = ncells;
    }

    py::object get_buffers() {
        py::dict rmap;
        for (auto& bp : buffers_) {
            py::list result;
            const auto name = bp.first;
            const BufferInfo b = bp.second;
            result.append(b.data);
            result.append(b.offsets);
            rmap[py::str{name}] = result;
        }
        return std::move(rmap);
    }

    void set_buffers() {
        for (auto bp : buffers_) {
            auto name = bp.first;
            const BufferInfo b = bp.second;

            size_t offsets_read = b.offsets_read;
            size_t data_vals_read = b.data_vals_read;
            size_t validity_vals_read = b.validity_vals_read;

            void* data_ptr = (void*)((char*)b.data.data() +
                                     (data_vals_read * b.elem_nbytes));
            uint64_t data_nelem = (b.data.size() -
                                   (data_vals_read * b.elem_nbytes)) /
                                  b.elem_nbytes;

            // Experimental version of API call is needed to support
            // type-checking on dimension label buffers.
            QueryExperimental::set_data_buffer(
                *query_, b.name, data_ptr, data_nelem);

            if (b.isvar) {
                size_t offsets_size = b.offsets.size() - offsets_read;
                uint64_t* offsets_ptr = (uint64_t*)b.offsets.data() +
                                        offsets_read;

                query_->set_offsets_buffer(
                    b.name, (uint64_t*)(offsets_ptr), offsets_size);
            }
            if (b.isnullable) {
                uint64_t validity_size = b.validity.size() - validity_vals_read;
                uint8_t* validity_ptr = (uint8_t*)b.validity.data() +
                                        validity_vals_read;

                query_->set_validity_buffer(
                    b.name, validity_ptr, validity_size);
            }
        }
    }

    void update_read_elem_num() {
        auto result_elements =
            QueryExperimental::result_buffer_elements_nullable_labels(*query_);

        for (const auto& read_info : result_elements) {
            auto name = read_info.first;
            uint64_t offset_elem_num = 0, data_vals_num = 0,
                     validity_elem_num = 0;
            std::tie(
                offset_elem_num,
                data_vals_num,
                validity_elem_num) = read_info.second;

            BufferInfo& buf = buffers_.at(name);

            // TODO if we ever support per-attribute read offset bitsize
            // then need to handle here. Currently this is hard-coded to
            // 64-bit to match query config.
            auto offset_ptr = buf.offsets.mutable_data();

            if (buf.isvar) {
                if (offset_elem_num > 0) {
                    // account for 'sm.var_offsets.extra_element'
                    offset_elem_num -= (use_arrow_) ? 1 : 0;
                }

                if (buf.offsets_read > 0) {
                    if (offset_ptr[buf.offsets_read] == 0) {
                        auto last_size = (buf.data_vals_read * buf.elem_nbytes);

                        for (uint64_t i = 0; i < offset_elem_num; i++) {
                            offset_ptr[buf.offsets_read + i] += last_size;
                        }
                    }
                }
            }

            buf.data_vals_read += data_vals_num;
            buf.offsets_read += offset_elem_num;
            buf.validity_vals_read += validity_elem_num;

            if ((Py_ssize_t)(buf.data_vals_read * buf.elem_nbytes) >
                (Py_ssize_t)buf.data.size()) {
                throw TileDBError(
                    "After read query, data buffer out of bounds: " + name +
                    " (" +
                    std::to_string(buf.data_vals_read * buf.elem_nbytes) +
                    " > " + std::to_string(buf.data.size()) + ")");
            }
            if ((Py_ssize_t)buf.offsets_read > buf.offsets.size()) {
                throw TileDBError(
                    "After read query, offsets buffer out of bounds: " + name +
                    " (" + std::to_string(buf.offsets_read) + " > " +
                    std::to_string(buf.offsets.size()) + ")");
            }
            if ((Py_ssize_t)buf.validity_vals_read > buf.validity.size()) {
                throw TileDBError(
                    "After read query, validity buffer out of bounds: " + name +
                    " (" + std::to_string(buf.validity_vals_read) + " > " +
                    std::to_string(buf.validity.size()) + ")");
            }
        }
    }

    void reset_read_elem_num() {
        for (auto& bp : buffers_) {
            auto& buf = bp.second;

            buf.offsets_read = 0;
            buf.data_vals_read = 0;
            buf.validity_vals_read = 0;
        }
    }

    uint64_t get_max_retries() {
        // should make this a templated getter for any key w/ default
        std::string tmp_str;
        size_t max_retries;
        try {
            tmp_str = ctx_.config().get("py.max_incomplete_retries");
            max_retries = std::stoull(tmp_str);
        } catch (const std::invalid_argument& e) {
            (void)e;
            throw TileDBError(
                "Failed to convert 'py.max_incomplete_retries' to uint64_t ('" +
                tmp_str + "')");
        } catch (tiledb::TileDBError& e) {
            (void)e;
            max_retries = 100;
        }
        return max_retries;
    }

    void resubmit_read() {
        tiledb_query_status_details_t status_details;
        tiledb_query_get_status_details(
            ctx_.ptr().get(), query_.get()->ptr().get(), &status_details);

        if (status_details.incomplete_reason ==
            TILEDB_REASON_USER_BUFFER_SIZE) {
            auto start_incomplete_buffer_update =
                std::chrono::high_resolution_clock::now();
            for (auto& bp : buffers_) {
                auto& buf = bp.second;

                // Check if values buffer should be resized
                if ((buf.data_vals_read == 0) ||
                    (int64_t)(buf.data_vals_read * buf.elem_nbytes) >
                        (buf.data.nbytes() + 1) / 2) {
                    size_t new_size = buf.data.size() * 2;
                    buf.data.resize({new_size}, false);
                }

                // Check if offset buffer should be resized
                if ((buf.isvar && buf.offsets_read == 0) ||
                    ((int64_t)(buf.offsets_read * sizeof(uint64_t)) >
                     (buf.offsets.nbytes() + 1) / 2)) {
                    size_t new_offsets_size = buf.offsets.size() * 2;
                    buf.offsets.resize({new_offsets_size}, false);
                }

                // Check if validity buffer should be resized
                if ((buf.isnullable && buf.validity_vals_read == 0) ||
                    ((int64_t)(buf.validity_vals_read * sizeof(uint8_t)) >
                     (buf.validity.nbytes() + 1) / 2)) {
                    size_t new_validity_size = buf.validity.size() * 2;
                    buf.validity.resize({new_validity_size}, false);
                }
            }

            // note: this block confuses lldb. continues from here unless bp set
            // after block.
            set_buffers();

            if (g_stats) {
                auto now = std::chrono::high_resolution_clock::now();
                g_stats.get()
                    ->counters["py.read_query_incomplete_buffer_resize_time"] +=
                    now - start_incomplete_buffer_update;
            }
        }

        {
            py::gil_scoped_release release;
            query_->submit();
        }

        if (query_->query_status() == Query::Status::UNINITIALIZED) {
            TPY_ERROR_LOC(
                "Unexpected state: Query::Submit returned uninitialized");
        }

        update_read_elem_num();

        return;
    }

    void resize_output_buffers() {
        // resize the output buffers to match the final read total
        // the higher level code uses the size of the buffers to
        // determine how much to unpack, but we may have over-allocated

        // account for the extra element at the end of offsets in arrow mode
        size_t arrow_offset_size = use_arrow_ ? 1 : 0;

        for (auto& bp : buffers_) {
            auto name = bp.first;
            auto& buf = bp.second;

            Py_ssize_t final_data_nbytes = buf.data_vals_read * buf.elem_nbytes;
            Py_ssize_t final_offsets_count = buf.offsets_read +
                                             arrow_offset_size;
            Py_ssize_t final_validity_count = buf.validity_vals_read;

            assert(final_data_nbytes <= buf.data.size());
            assert(
                final_offsets_count <=
                (Py_ssize_t)(buf.offsets.size() + arrow_offset_size));

            buf.data.resize({final_data_nbytes});
            buf.offsets.resize({final_offsets_count});
            buf.validity.resize({final_validity_count});

            if (use_arrow_) {
                if (retries_ > 0) {
                    // we need to write the final size to the final offset slot
                    // because core doesn't track size between incomplete
                    // submits
                    buf.offsets
                        .mutable_data()[buf.offsets_read] = final_data_nbytes;
                }

                // reset bytes-read so that set_buffers uses the full buffer
                // size
                buf.data_vals_read = 0;
                buf.offsets_read = 0;
                buf.validity_vals_read = 0;
            }
        }
        if (use_arrow_) {
            // this is a very light hack:
            // call set_buffers here to reset the buffers to the *full*
            // buffer in case there were incomplete queries. without this call,
            // the active tiledb::Query only knows about the buffer ptr/size
            // for the *last* submit loop, so we don't get full result set.
            // ArrowAdapter gets the buffer sizes from tiledb::Query.
            set_buffers();
        }
    }

    void allocate_buffers() {
        // allocate buffers for dims
        //   - we want to return dims first, if any requested
        for (size_t dim_idx = 0; dim_idx < domain_->ndim(); dim_idx++) {
            auto dim = domain_->dimension(dim_idx);
            if ((std::find(dims_.begin(), dims_.end(), dim.name()) ==
                 dims_.end()) &&
                // we need to also check if this is an attr for
                // backward-compatibility
                (std::find(attrs_.begin(), attrs_.end(), dim.name()) ==
                 attrs_.end())) {
                continue;
            }
            alloc_buffer(dim.name());
        }

        // allocate buffers for label dimensions
        for (const auto& label_data : label_input_buffer_data_) {
            alloc_buffer(label_data.first);
        }

        // allocate buffers for attributes
        //   - schema.attributes() is unordered, but we need to return ordered
        //   results
        for (size_t attr_idx = 0; attr_idx < array_schema_->attribute_num();
             attr_idx++) {
            auto attr = array_schema_->attribute(attr_idx);
            if (std::find(attrs_.begin(), attrs_.end(), attr.name()) ==
                attrs_.end()) {
                continue;
            }
            alloc_buffer(attr.name());
        }
    }

    void submit_read() {
        if (retries_ > 0 &&
            query_->query_status() == tiledb::Query::Status::INCOMPLETE) {
            buffers_.clear();
            assert(buffers_.size() == 0);

            buffers_order_.clear();
            // reset_read_elem_num();
        } else if (buffers_.size() != 0) {
            // we have externally imported buffers
            return;
        }

        // start time
        auto start = std::chrono::high_resolution_clock::now();

        // Initiate a metadata API request to make libtiledb fetch the
        // metadata ahead of time. In some queries we know we will always
        // access metadata, so initiating this call saves time when loading
        // from remote arrays because metadata i/o is lazy in core.
        /*
        // This section is disabled pending final disposition of SC-11720
        // This call is not currently safe as of TileDB 2.5, and has caused
        // reproducible deadlocks with the subesequent calls to
        // Array::est_result_sizes from the main thread.
        //
        std::future<uint64_t> metadata_num_preload;
        if (preload_metadata_) {
          metadata_num_preload = std::async(
              std::launch::async, [this]() { return array_->metadata_num(); });
        }
        */

        allocate_buffers();

        // set the buffers on the Query
        set_buffers();

        size_t max_retries = get_max_retries();

        auto start_submit = std::chrono::high_resolution_clock::now();
        {
            py::gil_scoped_release release;
            query_->submit();
        }

        if (query_->query_status() == Query::Status::UNINITIALIZED) {
            TPY_ERROR_LOC(
                "Unexpected state: Query::Submit returned uninitialized");
        }

        if (g_stats) {
            auto now = std::chrono::high_resolution_clock::now();
            g_stats.get()->counters["py.core_read_query_initial_submit_time"] +=
                now - start_submit;
        }

        // update the BufferInfo read-counts to match the query results read
        update_read_elem_num();

        // fetch the result of the metadata get task
        // this will block if not yet completed
        /*
        // disabled, see comment above
        if (preload_metadata_) {
          metadata_num_preload.get();
        }
        */

        auto incomplete_start = std::chrono::high_resolution_clock::now();

        // TODO: would be nice to have a callback here for custom realloc
        // strategy
        while (!return_incomplete_ &&
               query_->query_status() == Query::Status::INCOMPLETE) {
            if (++retries_ > max_retries)
                TPY_ERROR_LOC(
                    "Exceeded maximum retries ('py.max_incomplete_retries': '" +
                    std::to_string(max_retries) + "')");

            resubmit_read();
        }

        if (g_stats && retries_ > 0) {
            auto now = std::chrono::high_resolution_clock::now();
            g_stats.get()
                ->counters["py.core_read_query_incomplete_retry_time"] +=
                now - incomplete_start;
        }

        // update TileDB-Py stat counter
        if (g_stats) {
            auto now = std::chrono::high_resolution_clock::now();
            g_stats.get()->counters["py.core_read_query_total_time"] += now -
                                                                        start;
        }

        if (g_stats) {
            g_stats.get()->counters["py.query_retries_count"] += TimerType(
                retries_);
        }

        resize_output_buffers();

        if (return_incomplete_) {
            // increment in case we submit again
            retries_++;
        }
    }

    py::array unpack_buffer(
        std::string name, py::array buf, py::array_t<uint64_t> off) {
        auto start = std::chrono::high_resolution_clock::now();

        if (off.size() < 1)
            TPY_ERROR_LOC(
                std::string("Unexpected empty offsets array ('") + name + "')");

        auto dtype = buffer_dtype(name);
        bool is_unicode = dtype.is(py::dtype("U"));
        bool is_str = dtype.is(py::dtype("S"));
        if (is_unicode || is_str) {
            dtype = py::dtype("O");
        }

        // Hashmap for string deduplication
        // fastest so far:
        typedef tsl::robin_map<size_t, uint64_t> MapType;
        MapType map;
        std::vector<py::object> object_v;
        if (is_unicode) {
            map.reserve(size_t(off.size() / 10) + 1);
        }

        auto result_array = py::array(py::dtype("O"), off.size());
        auto result_p = (py::object*)result_array.mutable_data();
        uint64_t last = 0;
        uint64_t cur = 0;
        size_t size = 0;
        uint64_t create = 0;

        auto off_data = off.data();
        last = off_data[0];  // initial should always be 0
        for (auto i = 1; i < off.size() + 1; i++) {
            if (i == off.size())
                cur = buf.nbytes();
            else {
                cur = off_data[i];
            }

            size = cur - last;

            py::object o;
            auto data_ptr = (char*)buf.data() + last;
            if (is_unicode)
                if (size == 0 || (data_ptr[0] == '\0' && size == 1)) {
                    o = py::str("");
                } else {
                    if (!deduplicate_) {
                        o = py::str(data_ptr, size);
                    } else {
                        auto v = std::string_view{data_ptr, size};
                        auto h = std::hash<std::string_view>()(v);
                        auto needle = map.find(h);
                        if (needle == map.end()) {
                            o = py::str(data_ptr, size);
                            map.insert(needle, {h, create});
                            object_v.push_back(o);
                            create++;
                        } else {
                            auto idx = needle->second;
                            o = object_v[idx];
                        }
                    }
                }
            else if (is_str)
                if (size == 0 || (data_ptr[0] == '\0' && size == 1)) {
                    o = py::bytes("");
                } else {
                    o = py::bytes(data_ptr, size);
                }
            else {
                o = py::array(py::dtype("uint8"), size, data_ptr);
                o.attr("dtype") = dtype;
            }

            result_p[i - 1] = o;
            last = cur;
        }

        if (g_stats) {
            auto now = std::chrono::high_resolution_clock::now();
            g_stats.get()->counters["py.buffer_conversion_time"] += now - start;
        }

        return result_array;
    }

    void submit_write() {
    }

    void submit() {
        if (array_->query_type() == TILEDB_READ) {
            submit_read();
        } else if (array_->query_type() == TILEDB_WRITE) {
            submit_write();
        } else if (array_->query_type() == TILEDB_DELETE) {
            query_->submit();
            if (query_->query_status() == Query::Status::UNINITIALIZED) {
                TPY_ERROR_LOC(
                    "Unexpected state: Query::Submit returned uninitialized");
            }
        } else {
            TPY_ERROR_LOC("Unknown query type!")
        }
    }

    py::dict results() {
        py::dict results;
        for (auto& buffer_name : buffers_order_) {
            auto bp = buffers_.at(buffer_name);
            results[py::str(buffer_name)] = py::make_tuple(
                bp.data, bp.offsets, bp.validity);
        }
        return results;
    }

    std::unique_ptr<PAPair> buffer_to_pa(std::string name) {
        if (query_->query_status() != tiledb::Query::Status::COMPLETE)
            TPY_ERROR_LOC("Cannot convert buffers unless Query is complete");

        tiledb::arrow::ArrowAdapter adapter(&ctx_, query_.get());

        std::unique_ptr<PAPair> pa_pair(new PAPair());

        adapter.export_buffer(
            name.c_str(),
            &(pa_pair->array_),
            &(pa_pair->schema_),
            nullptr,
            nullptr);
        pa_pair->exported_ = true;

        return pa_pair;
    }

    py::object buffers_to_pa_table() {
        using namespace pybind11::literals;

        auto pa = py::module::import("pyarrow");
        auto pa_array_import = pa.attr("Array").attr("_import_from_c");

        tiledb::arrow::ArrowAdapter adapter(&ctx_, query_.get());

        py::list names;
        py::list results;
        for (auto& buffer_name : buffers_order_) {
            BufferInfo& buffer_info = buffers_.at(buffer_name);

            auto buffer_holder = new BufferHolder(
                buffer_info.data, buffer_info.validity, buffer_info.offsets);

            ArrowArray c_pa_array;
            ArrowSchema c_pa_schema;
            adapter.export_buffer(
                buffer_name.c_str(),
                static_cast<void*>(&c_pa_array),
                static_cast<void*>(&c_pa_schema),
                (tiledb::arrow::ArrowAdapter::release_cb)
                    BufferHolder::free_buffer_holder,
                buffer_holder);

            if (is_nullable(buffer_name)) {
                // count zeros before converting to bitmap
                c_pa_array.null_count = count_zeros(buffer_info.validity);
                // convert to bitmap
                buffer_info.validity = uint8_bool_to_uint8_bitmap(
                    buffer_info.validity);
                c_pa_array.buffers[0] = buffer_info.validity.data();
                c_pa_array.n_buffers = is_var(buffer_name) ? 3 : 2;
                c_pa_schema.flags |= ARROW_FLAG_NULLABLE;
            } else if (!is_var(buffer_name)) {
                // reset the number of buffers for non-nullable data
                c_pa_array.n_buffers = 2;
            }

            // work around for SC-11522: metadata field must be set to nullptr
            c_pa_schema.metadata = nullptr;
            py::object pa_array = pa_array_import(
                py::int_((ptrdiff_t)&c_pa_array),
                py::int_((ptrdiff_t)&c_pa_schema));

            results.append(pa_array);
            names.append(buffer_name);
        }

        auto pa_table = pa.attr("Table").attr("from_arrays")(
            results, "names"_a = names);
        return pa_table;
    }

    py::object is_incomplete() {
        if (!query_) {
            throw TileDBPyError("Internal error: PyQuery not initialized!");
        }
        return py::cast<bool>(
            query_->query_status() == tiledb::Query::Status::INCOMPLETE);
    }

    py::object estimated_result_sizes() {
        // vector of names to estimate
        std::vector<std::string> estim_names;

        for (size_t dim_idx = 0; dim_idx < domain_->ndim(); dim_idx++) {
            auto dim = domain_->dimension(dim_idx);
            if ((std::find(dims_.begin(), dims_.end(), dim.name()) ==
                 dims_.end()) &&
                // we need to also check if this is an attr for
                // backward-compatibility
                (std::find(attrs_.begin(), attrs_.end(), dim.name()) ==
                 attrs_.end())) {
                continue;
            }
            estim_names.push_back(dim.name());
        }

        // iterate by idx: schema.attributes() is unordered, but we need to
        //                 return ordered results
        for (size_t attr_idx = 0; attr_idx < array_schema_->attribute_num();
             attr_idx++) {
            auto attr = array_schema_->attribute(attr_idx);
            if (std::find(attrs_.begin(), attrs_.end(), attr.name()) ==
                attrs_.end()) {
                continue;
            }
            estim_names.push_back(attr.name());
        }

        py::dict results;

        for (auto const& name : estim_names) {
            size_t est_offsets = 0, est_data_bytes = 0;

            if (is_var(name)) {
                query_->est_result_size_var(name);
                auto est_sizes = query_->est_result_size_var(name);
                est_offsets = std::get<0>(est_sizes);
                est_data_bytes = std::get<1>(est_sizes);
            } else {
                est_data_bytes = query_->est_result_size(name);
            }
            results[py::str(name)] = py::make_tuple(
                est_offsets, est_data_bytes);
        }

        return std::move(results);
    }

    py::array _test_array() {
        py::array_t<uint8_t> a;
        a.resize({10});

        a.resize({20});
        return std::move(a);
    }

    uint64_t _test_init_buffer_bytes() {
        // test helper to get the configured init_buffer_bytes
        return init_buffer_bytes_;
    }
    uint64_t _test_alloc_max_bytes() {
        // test helper to get the configured init_buffer_bytes
        return alloc_max_bytes_;
    }

    std::string get_stats() {
        return query_->stats();
    }

   private:
    void init_config() {
        // get config parameters
        std::string tmp_str;
        if (config_has_key(ctx_.config(), "py.init_buffer_bytes")) {
            tmp_str = ctx_.config().get("py.init_buffer_bytes");
            try {
                init_buffer_bytes_ = std::stoull(tmp_str);
            } catch (const std::invalid_argument& e) {
                (void)e;
                throw std::invalid_argument(
                    "Failed to convert 'py.init_buffer_bytes' to uint64_t ('" +
                    tmp_str + "')");
            }
        }

        if (config_has_key(ctx_.config(), "py.alloc_max_bytes")) {
            tmp_str = ctx_.config().get("py.alloc_max_bytes");
            try {
                alloc_max_bytes_ = std::stoull(tmp_str);
            } catch (const std::invalid_argument& e) {
                (void)e;
                throw std::invalid_argument(
                    "Failed to convert 'py.alloc_max_bytes' to uint64_t ('" +
                    tmp_str + "')");
            }
            if (alloc_max_bytes_ < pow(1024, 2)) {
                throw std::invalid_argument(
                    "Invalid parameter: 'py.alloc_max_bytes' "
                    "must be >= 1 MB (1024 ** 2 bytes)");
            };
        }

        if (config_has_key(ctx_.config(), "py.deduplicate")) {
            tmp_str = ctx_.config().get("py.deduplicate");
            if (tmp_str == "true") {
                deduplicate_ = true;
            } else if (tmp_str == "false") {
                deduplicate_ = false;
            } else {
                throw std::invalid_argument(
                    "Failed to convert configuration 'py.deduplicate' to bool "
                    "('" +
                    tmp_str + "')");
            }
        }

        if (config_has_key(ctx_.config(), "py.exact_init_buffer_bytes")) {
            tmp_str = ctx_.config().get("py.exact_init_buffer_bytes");
            if (tmp_str == "true") {
                exact_init_bytes_ = true;
            } else if (tmp_str == "false") {
                exact_init_bytes_ = false;
            } else {
                throw std::invalid_argument(
                    "Failed to convert configuration "
                    "'py.exact_init_buffer_bytes' to bool ('" +
                    tmp_str + "')");
            }
        }

        if (config_has_key(ctx_.config(), "py.use_arrow")) {
            tmp_str = ctx_.config().get("py.use_arrow");
            if (tmp_str == "True") {
                use_arrow_ = true;
            } else if (tmp_str == "False") {
                use_arrow_ = false;
            } else {
                throw std::invalid_argument(
                    "Failed to convert configuration 'py.use_arrow' to bool "
                    "('" +
                    tmp_str + "')");
            }
        }
    }

};  // namespace tiledbpy

void init_stats() {
    g_stats.reset(new StatsInfo());

    auto stats_counters = g_stats.get()->counters;
    stats_counters["py.core_read_query_initial_submit_time"] = TimerType();
    stats_counters["py.core_read_query_total_time"] = TimerType();
    stats_counters["py.core_read_query_incomplete_retry_time"] = TimerType();
    stats_counters["py.buffer_conversion_time"] = TimerType();
    stats_counters["py.read_query_incomplete_buffer_resize_time"] = TimerType();
    stats_counters["py.query_retries_count"] = TimerType();
}

void disable_stats() {
    g_stats.reset(nullptr);
}

bool stats_enabled() {
    return (bool)g_stats;
}

void increment_stat(std::string key, double value) {
    auto& stats_counters = g_stats.get()->counters;

    if (stats_counters.count(key) == 0)
        stats_counters[key] = TimerType();

    auto& timer = stats_counters[key];
    auto incr = std::chrono::duration<double>(value);
    timer += incr;
}

bool use_stats() {
    return (bool)g_stats;
}

py::object get_stats() {
    if (!g_stats) {
        TPY_ERROR_LOC("Stats counters are not uninitialized!")
    }

    auto& stats_counters = g_stats.get()->counters;

    py::dict res;
    for (auto iter = stats_counters.begin(); iter != stats_counters.end();
         ++iter) {
        auto val = std::chrono::duration<double>(iter->second);
        res[py::str(iter->first)] = py::float_(val.count());
    }
    return std::move(res);
}

py::object python_internal_stats(bool dict = false) {
    if (!g_stats) {
        TPY_ERROR_LOC("Stats counters are not uninitialized!")
    }

    auto counters = g_stats.get()->counters;
    auto rq_time = counters["py.core_read_query_initial_submit_time"].count();

    if (dict) {
        py::dict stats;

        // core.cc is only tracking read time right now; don't output if we
        // have no query submission time
        if (rq_time == 0) {
            return std::move(stats);
        }

        for (auto& stat : counters) {
            stats[py::str(stat.first)] = stat.second.count();
        }

        return std::move(stats);
    } else {
        std::ostringstream os;

        // core.cc is only tracking read time right now; don't output if we
        // have no query submission time
        if (rq_time == 0) {
            return py::str(os.str());
        }

        os << std::endl;
        os << "==== Python Stats ====" << std::endl << std::endl;

        for (auto& stat : counters) {
            os << "  " << stat.first << " : " << stat.second.count()
               << std::endl;
        }

        return py::str(os.str());
    }
}

py::str as_built_dump() {
    tiledb_string_t* s;
    int rc = tiledb_as_built_dump(&s);
    if (rc != TILEDB_OK) {
        TPY_ERROR_LOC("Could not dump as built.");
    }
    const char* data_ptr;
    py::size_t length;

    tiledb_string_view(s, &data_ptr, &length);
    py::str res(data_ptr, length);
    tiledb_string_free(&s);

    return res;
}

auto walk_callback =
    [](const char* path_ptr, tiledb_object_t obj, void* pyfunc) {
        std::string my_path(path_ptr);
        std::string my_objtype;
        if (obj == TILEDB_GROUP) {
            my_objtype = "group";
        } else if (obj == TILEDB_ARRAY) {
            my_objtype = "array";
        } else {
            return 0;
        }
        try {
            py::function func = py::cast<py::function>((PyObject*)pyfunc);
            func(my_path, my_objtype);
        } catch (py::stop_iteration) {
            return 0;
        }
        return 1;
    };

void init_core(py::module& m) {
    init_query_condition(m);

    auto pq =
        py::class_<PyQuery>(m, "PyQuery")
            .def(py::init<
                 const Context&,
                 py::object,
                 py::iterable,
                 py::iterable,
                 py::object,
                 py::object>())
            .def("add_label_buffer", &PyQuery::add_label_buffer)
            .def("buffer_dtype", &PyQuery::buffer_dtype)
            .def("results", &PyQuery::results)
            .def("set_subarray", &PyQuery::set_subarray)
            .def("set_cond", &PyQuery::set_cond)
#if defined(TILEDB_SERIALIZATION)
            .def("set_serialized_query", &PyQuery::set_serialized_query)
#endif
            .def("submit", &PyQuery::submit)
            .def("unpack_buffer", &PyQuery::unpack_buffer)
            .def("estimated_result_sizes", &PyQuery::estimated_result_sizes)
            .def("get_stats", &PyQuery::get_stats)
            .def("_allocate_buffers", &PyQuery::allocate_buffers)
            .def("_get_buffers", &PyQuery::get_buffers)
            .def("_buffer_to_pa", &PyQuery::buffer_to_pa)
            .def("_buffers_to_pa_table", &PyQuery::buffers_to_pa_table)
            .def("_test_array", &PyQuery::_test_array)
            .def(
                "_test_err",
                [](py::object self, std::string s) { throw TileDBPyError(s); })
            .def_readwrite("_preload_metadata", &PyQuery::preload_metadata_)
            .def_readwrite("_return_incomplete", &PyQuery::return_incomplete_)
            // properties
            .def_property_readonly("is_incomplete", &PyQuery::is_incomplete)
            .def_property_readonly(
                "_test_init_buffer_bytes", &PyQuery::_test_init_buffer_bytes)
            .def_property_readonly(
                "_test_alloc_max_bytes", &PyQuery::_test_alloc_max_bytes)
            .def_readonly("retries", &PyQuery::retries_);

    py::class_<PyAgg>(m, "PyAgg")
        .def(
            py::init<const Context&, py::object, py::object, py::dict>(),
            "ctx"_a,
            "py_array"_a,
            "py_layout"_a,
            "attr_to_aggs_input"_a)
        .def("set_subarray", &PyAgg::set_subarray)
        .def("set_cond", &PyAgg::set_cond)
        .def("get_aggregate", &PyAgg::get_aggregate);

    py::class_<PAPair>(m, "PAPair")
        .def(py::init())
        .def("get_array", &PAPair::get_array)
        .def("get_schema", &PAPair::get_schema);

    m.def("array_to_buffer", &convert_np);

    m.def("init_stats", []() {
        Stats::enable();
        init_stats();
    });
    m.def("disable_stats", []() {
        Stats::disable();
        disable_stats();
    });
    m.def("stats_enabled", &stats_enabled);
    m.def("reset_stats", []() {
        Stats::reset();
        init_stats();
    });
    m.def("stats_raw_dump_str", []() {
        std::string out;
        Stats::raw_dump(&out);
        return out;
    });
    m.def("stats_dump_str", []() {
        std::string out;
        Stats::dump(&out);
        return out;
    });
    m.def(
        "python_internal_stats",
        &python_internal_stats,
        py::arg("dict") = false);
    m.def("increment_stat", &increment_stat);
    m.def("get_stats", &get_stats);
    m.def("use_stats", &use_stats);
    m.def("datatype_size", &tiledb_datatype_size);
    m.def("as_built_dump", &as_built_dump);
    m.def(
        "object_type",
        [](const std::string& uri, const Context& ctx) -> py::object {
            tiledb_object_t res;
            ctx.handle_error(
                tiledb_object_type(ctx.ptr().get(), uri.c_str(), &res));
            if (res == TILEDB_ARRAY) {
                return py::str("array");
            } else if (res == TILEDB_GROUP) {
                return py::str("group");
            }
            return py::none();
        });
    m.def(
        "ls",
        [](const std::string& path, py::function func, const Context& ctx) {
            ctx.handle_error(tiledb_object_ls(
                ctx.ptr().get(),
                path.c_str(),
                walk_callback,
                (void*)func.ptr()));
        });
    m.def(
        "walk",
        [](const std::string path,
           py::function func,
           const std::string order,
           const Context& ctx) {
            tiledb_walk_order_t walk_order;
            if (order == "postorder") {
                walk_order = TILEDB_POSTORDER;
            } else if (order == "preorder") {
                walk_order = TILEDB_PREORDER;
            } else {
                throw TileDBError("unknown walk order " + order);
            }
            ctx.handle_error(tiledb_object_walk(
                ctx.ptr().get(),
                path.c_str(),
                walk_order,
                walk_callback,
                (void*)func.ptr()));
        });
    m.def("remove", &Object::remove);
    m.def("move", &Object::move);

    /*
     We need to make sure C++ TileDBError is translated to a correctly-typed py
     error. Note that using py::exception(..., "TileDBError") creates a new
     exception in the *readquery* module, so we must import to reference.
  */
    static auto tiledb_py_error = (py::object)py::module::import("tiledb").attr(
        "TileDBError");

    py::register_exception_translator([](std::exception_ptr p) {
        try {
            if (p)
                std::rethrow_exception(p);
        } catch (const TileDBPyError& e) {
            PyErr_SetString(tiledb_py_error.ptr(), e.what());
        } catch (const tiledb::TileDBError& e) {
            PyErr_SetString(tiledb_py_error.ptr(), e.what());
        } catch (py::builtin_exception& e) {
            // just forward the error
            throw;
            //} catch (std::runtime_error &e) {
            //  std::cout << "unexpected runtime_error: " << e.what() <<
            //  std::endl;
        }
    });
};
};  // namespace tiledbpy
