Bind Enumerated Data Type

Author: nguyenv

setup.py

@@ -24,7 +24,7 @@
 # - this is for builds-from-source
 # - release builds are controlled by `misc/azure-release.yml`
 # - this should be set to the current core release, not `dev`
-TILEDB_VERSION = "2.16.3"
+TILEDB_VERSION = "dev"
 
 # allow overriding w/ environment variable
 TILEDB_VERSION = os.environ.get("TILEDB_VERSION") or TILEDB_VERSION

tiledb/__init__.py

@@ -45,6 +45,7 @@
 from .dimension_label import DimLabel
 from .dimension_label_schema import DimLabelSchema
 from .domain import Domain
+from .enumeration import Enumeration
 from .filestore import Filestore
 from .filter import (
     BitShuffleFilter,

tiledb/array_schema.py

@@ -51,9 +51,15 @@ def __init__(
         allows_duplicates: bool = False,
         sparse: bool = False,
         dim_labels={},
+        enums=None,
         ctx: Ctx = None,
     ):
         super().__init__(ctx, lt.ArrayType.SPARSE if sparse else lt.ArrayType.DENSE)
+
+        if enums is not None:
+            for enum_name in enums:
+                self._add_enumeration(self._ctx, enum_name)
+
         if attrs is not None:
             for att in attrs:
                 if not isinstance(att, Attr):

tiledb/attribute.py

@@ -24,6 +24,7 @@ def __init__(
         var: bool = None,
         nullable: bool = False,
         filters: Union[FilterList, Sequence[Filter]] = None,
+        enum_label: str = None,
         ctx: Optional[Ctx] = None,
     ):
         """Class representing a TileDB array attribute.
@@ -88,6 +89,9 @@ def __init__(
         if nullable is not None:
             self._nullable = nullable
 
+        if enum_label is not None:
+            self._set_enumeration_name(self._ctx, enum_label)
+
     def __eq__(self, other):
         if not isinstance(other, Attr):
             return False
@@ -202,6 +206,10 @@ def isascii(self) -> bool:
         """
         return self._tiledb_dtype == lt.DataType.STRING_ASCII
 
+    @property
+    def enum_label(self):
+        return self._get_enumeration_name(self._ctx)
+
     def __repr__(self):
         filters_str = ""
         if self.filters:
@@ -217,11 +225,16 @@ def __repr__(self):
         else:
             attr_dtype = self.dtype
 
+        if self.enum_label is None:
+            enum_label = None
+        else:
+            enum_label = f"'{self.enum_label!s}'"
+
         # filters_str must be last with no spaces
         return (
             f"""Attr(name={repr(self.name)}, dtype='{attr_dtype!s}', """
-            f"""var={self.isvar!s}, nullable={self.isnullable!s}"""
-            f"""{filters_str})"""
+            f"""var={self.isvar!s}, nullable={self.isnullable!s}, """
+            f"""enum_label={enum_label}{filters_str})"""
         )
 
     def _repr_html_(self):

tiledb/cc/array.cc

@@ -34,6 +34,7 @@ void init_array(py::module &m) {
       .def("uri", &Array::uri)
       .def("schema", &Array::schema)
       //.def("ptr", [](Array& arr){ return py::capsule(arr.ptr()); } )
+      .def("open", (void (Array::*)(tiledb_query_type_t)) & Array::open)
       // open with encryption key
       .def("open",
            (void (Array::*)(tiledb_query_type_t, tiledb_encryption_type_t,

tiledb/cc/attribute.cc

@@ -1,4 +1,5 @@
 #include <tiledb/tiledb>
+#include <tiledb/tiledb_experimental>
 
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
@@ -39,6 +40,16 @@ py::array get_fill_value(Attribute &attr) {
   return py::array(value_type, value_num, value);
 }
 
+void set_enumeration_name(Attribute &attr, const Context &ctx,
+                          const std::string &enumeration_name) {
+  AttributeExperimental::set_enumeration_name(ctx, attr, enumeration_name);
+}
+
+std::optional<std::string> get_enumeration_name(Attribute &attr,
+                                                const Context &ctx) {
+  return AttributeExperimental::get_enumeration_name(ctx, attr);
+}
+
 void init_attribute(py::module &m) {
   py::class_<tiledb::Attribute>(m, "Attribute")
       .def(py::init<Attribute>())
@@ -73,8 +84,11 @@ void init_attribute(py::module &m) {
 
       .def_property("_fill", get_fill_value, set_fill_value)
 
+      .def("_get_enumeration_name", get_enumeration_name)
+
+      .def("_set_enumeration_name", set_enumeration_name)
+
       .def("_dump", [](Attribute &attr) { attr.dump(); });
-  ;
 }
 
 } // namespace libtiledbcpp

tiledb/cc/enumeration.cc

@@ -0,0 +1,77 @@
+#include <tiledb/tiledb>
+#include <tiledb/tiledb_experimental>
+
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include <pybind11/stl.h>
+
+#include "common.h"
+
+namespace libtiledbcpp {
+
+using namespace tiledb;
+using namespace tiledbpy::common;
+namespace py = pybind11;
+
+void init_enumeration(py::module &m) {
+  py::class_<Enumeration>(m, "Enumeration")
+      .def(py::init<Enumeration>())
+
+      .def(py::init([](const Context &ctx, const std::string &name,
+                       std::vector<std::string> &values, bool ordered,
+                       tiledb_datatype_t type) {
+        return Enumeration::create(ctx, name, values, ordered, type);
+      }))
+
+      .def(py::init([](const Context &ctx, const std::string &name,
+                       bool ordered, py::array data, py::array offsets) {
+        tiledb_datatype_t data_type;
+        try {
+          data_type = np_to_tdb_dtype(data.dtype());
+        } catch (const TileDBPyError &e) {
+          throw py::type_error(e.what());
+        }
+
+        py::buffer_info data_buffer = data.request();
+        if (data_buffer.ndim != 1)
+          throw py::type_error("Only 1D Numpy arrays can be stored as "
+                               "enumeration values");
+
+        py::size_t cell_val_num =
+            offsets.size() == 0 ? get_ncells(data.dtype()) : TILEDB_VAR_NUM;
+
+        return Enumeration::create(
+            ctx, name, data_type, cell_val_num, ordered, data.data(),
+            data.nbytes(), offsets.size() == 0 ? nullptr : offsets.data(),
+            offsets.nbytes());
+      }))
+
+      .def(py::init<const Context &, py::capsule>(), py::keep_alive<1, 2>())
+
+      .def("__capsule__",
+           [](Enumeration &enmr) {
+             return py::capsule(enmr.ptr().get(), "enmr", nullptr);
+           })
+
+      .def_property_readonly("name", &Enumeration::name)
+
+      .def_property_readonly("type", &Enumeration::type)
+
+      .def_property_readonly("cell_val_num", &Enumeration::cell_val_num)
+
+      .def_property_readonly("ordered", &Enumeration::ordered)
+
+      .def("values",
+           [](Enumeration &enmr) {
+             auto data = enmr.as_vector<std::byte>();
+             auto dtype = tdb_to_np_dtype(enmr.type(), enmr.cell_val_num());
+             return py::array(dtype, data.size() / dtype.itemsize(),
+                              data.data());
+           })
+
+      .def("str_values",
+           [](Enumeration &enmr) { return enmr.as_vector<std::string>(); });
+}
+
+} // namespace libtiledbcpp

tiledb/cc/schema.cc

@@ -278,16 +278,23 @@ void init_schema(py::module &m) {
       .def("_has_attribute", &ArraySchema::has_attribute)
 
 #if TILEDB_VERSION_MAJOR == 2 && TILEDB_VERSION_MINOR >= 15
-      .def("_has_dim_label", [](const ArraySchema &schema, const Context &ctx,
-                                const std::string &name) {
-        return ArraySchemaExperimental::has_dimension_label(ctx, schema, name);
-      });
+      .def("_has_dim_label",
+           [](const ArraySchema &schema, const Context &ctx,
+              const std::string &name) {
+             return ArraySchemaExperimental::has_dimension_label(ctx, schema,
+                                                                 name);
+           })
 #else
       .def("_has_dim_label", [](const ArraySchema &, const Context &,
                                 const std::string &) {
         return false; 
-      });
+      })
 #endif
+
+      .def("_add_enumeration", [](const ArraySchema &schema, const Context &ctx,
+                                  const Enumeration &enmr) {
+        ArraySchemaExperimental::add_enumeration(ctx, schema, enmr);
+      });
 }
 
 } // namespace libtiledbcpp

tiledb/cc/tiledbcpp.cc

@@ -19,6 +19,7 @@ void init_attribute(py::module &);
 void init_context(py::module &);
 void init_config(py::module &);
 void init_enums(py::module &);
+void init_enumeration(py::module &);
 void init_dimension_label(py::module &m);
 void init_domain(py::module &m);
 void init_file_handle(py::module &);
@@ -40,6 +41,7 @@ PYBIND11_MODULE(cc, m) {
   init_dimension_label(m);
   init_domain(m);
   init_enums(m);
+  init_enumeration(m);
   init_file_handle(m);
   init_filestore(m);
   init_filter(m);

tiledb/core.cc

@@ -316,7 +316,7 @@ class PyQuery {
   // label buffer list
   std::unordered_map<string, uint64_t> label_input_buffer_data_;
 
-  py::object pyschema_;
+  std::string uri_;
 
 public:
   tiledb_ctx_t *c_ctx_;
@@ -349,7 +349,7 @@ class PyQuery {
     domain_ =
         std::shared_ptr<tiledb::Domain>(new Domain(array_schema_->domain()));
 
-    pyschema_ = array.attr("schema");
+    uri_ = array.attr("uri").cast<std::string>();
 
     bool issparse = array_->schema().array_type() == TILEDB_SPARSE;
 
@@ -450,7 +450,7 @@ class PyQuery {
     py::object init_pyqc = cond.attr("init_query_condition");
 
     try {
-      init_pyqc(pyschema_, attrs_);
+      init_pyqc(uri_, attrs_);
     } catch (tiledb::TileDBError &e) {
       TPY_ERROR_LOC(e.what());
     } catch (py::error_already_set &e) {