Add support for creating WKB/WKT attributes (#1912)

* Add support for WKB/WKT attributes
* Add tests

---------
Co-authored-by: Matthew Perry <[email protected]>
Co-authored-by: Agis Kounelis <[email protected]>

Author: 3 people

tiledb/attribute.py

@@ -40,7 +40,7 @@ def __init__(
         :raises tiledb.TileDBError:
         """
         dt = DataType.from_numpy(
-            np.dtype(dtype) if dtype not in ("ascii", "blob") else dtype
+            np.dtype(dtype) if dtype not in ("ascii", "blob", "wkb", "wkt") else dtype
         )
 
         # ensure that all strings are var-length
@@ -274,6 +274,16 @@ def __repr__(self):
             attr_dtype = "ascii"
         elif self._tiledb_dtype == lt.DataType.BLOB:
             attr_dtype = "blob"
+        elif (
+            hasattr(lt.DataType, "GEOM_WKB")
+            and self._tiledb_dtype == lt.DataType.GEOM_WKB
+        ):
+            attr_dtype = "wkb"
+        elif (
+            hasattr(lt.DataType, "GEOM_WKT")
+            and self._tiledb_dtype == lt.DataType.GEOM_WKT
+        ):
+            attr_dtype = "wkt"
         else:
             attr_dtype = self.dtype
 

tiledb/cc/attribute.cc

@@ -19,28 +19,42 @@ void set_fill_value(Attribute &attr, py::array value) {
 }
 
 py::array get_fill_value(Attribute &attr) {
+  // Get the fill value from the C++ API as a void* value.
   const void *value;
   uint64_t size;
-
   attr.get_fill_value(&value, &size);
 
-  auto value_num = attr.cell_val_num();
-  auto value_type = tdb_to_np_dtype(attr.type(), value_num);
-
+  // If this is a string type, we want to return each value as a single cell.
   if (is_tdb_str(attr.type())) {
-    value_type = py::dtype("|S1");
-    value_num = size;
+    auto value_type = py::dtype("|S1");
+    return py::array(value_type, size, value);
   }
 
-  // record type
+  // If this is a record type (void), return a single cell.
+  // If this is a blob-like type, we want to return each value as a single byte
+  // cell.
+  auto tdb_type = attr.type();
+  if (tdb_type == TILEDB_BLOB
+#if TILEDB_VERSION_MAJOR >= 2 && TILEDB_VERSION_MINOR >= 21
+      || tdb_type == TILEDB_GEOM_WKB || tdb_type == TILEDB_GEOM_WKT
+#endif
+  ) {
+    auto value_type = py::dtype("S");
+    return py::array(value_type, size, value);
+  }
+
+  // Get the number of values in a cell and the Python datatype.
+  auto value_num = attr.cell_val_num();
+  auto value_type = tdb_to_np_dtype(attr.type(), value_num);
+
   if (py::str(value_type.attr("kind")) == py::str("V")) {
-    value_num = 1;
+    return py::array(value_type, 1, value);
   }
 
-  // complex type - both cell values fit in a single complex element
+  // If this is a complex type both cell values fit in a single complex element.
   if (value_type == py::dtype("complex64") ||
       value_type == py::dtype("complex128")) {
-    value_num = 1;
+    return py::array(value_type, 1, value);
   }
 
   return py::array(value_type, value_num, value);

tiledb/cc/common.cc

@@ -43,7 +43,7 @@ std::unordered_map<tiledb_datatype_t, std::string> _tdb_to_np_name_dtype = {
     {TILEDB_BOOL, "bool"},
 #if TILEDB_VERSION_MAJOR >= 2 && TILEDB_VERSION_MINOR >= 21
     {TILEDB_GEOM_WKB, "byte"},
-    {TILEDB_GEOM_WKT, "S1"},
+    {TILEDB_GEOM_WKT, "S"},
 #endif
 };
 

tiledb/core.cc

@@ -218,7 +218,7 @@ py::dtype tiledb_dtype(tiledb_datatype_t type, uint32_t cell_val_num) {
     case TILEDB_GEOM_WKB:
       return py::dtype("byte");
     case TILEDB_GEOM_WKT:
-      return py::dtype("S1");
+      return py::dtype("S");
 #endif
 
     case TILEDB_ANY:

tiledb/datatypes.py

@@ -24,6 +24,12 @@ def from_numpy(cls, dtype: np.dtype) -> DataType:
         if dtype == "blob":
             return cls(np.dtype("S"), lt.DataType.BLOB, 1)
 
+        if hasattr(lt.DataType, "GEOM_WKB") and dtype == "wkb":
+            return cls(np.dtype("S"), lt.DataType.GEOM_WKB, 1)
+
+        if hasattr(lt.DataType, "GEOM_WKT") and dtype == "wkt":
+            return cls(np.dtype("S"), lt.DataType.GEOM_WKT, 1)
+
         dtype = np.dtype(dtype)
         if dtype.kind == "V":
             # fixed-size record dtypes
@@ -179,10 +185,13 @@ def uncast_tile_extent(self, tile_extent: Any) -> np.generic:
 _NUMPY_TO_TILEDB[np.dtype("complex64")] = lt.DataType.FLOAT32
 _NUMPY_TO_TILEDB[np.dtype("complex128")] = lt.DataType.FLOAT64
 
-# tiledb has STRING_ASCII and BLOB, numpy doesn't
+# tiledb has STRING_ASCII, BLOB, WKB and WKT types, numpy doesn't
 _TILEDB_TO_NUMPY = {t: n for n, t in _COMMON_DATATYPES}
 _TILEDB_TO_NUMPY[lt.DataType.STRING_ASCII] = np.dtype("S")
 _TILEDB_TO_NUMPY[lt.DataType.BLOB] = np.dtype("S")
+if hasattr(lt.DataType, "GEOM_WKB"):
+    _TILEDB_TO_NUMPY[lt.DataType.GEOM_WKB] = np.dtype("S")
+    _TILEDB_TO_NUMPY[lt.DataType.GEOM_WKT] = np.dtype("S")
 
 # pre-populate the LRU caches with all ncell=1 datatypes
 list(map(DataType.from_numpy, _NUMPY_TO_TILEDB.keys()))

tiledb/libtiledb.pxd

@@ -60,6 +60,8 @@ cdef extern from "tiledb/tiledb.h":
         TILEDB_FLOAT32
         TILEDB_FLOAT64
         TILEDB_BLOB
+        TILEDB_GEOM_WKB
+        TILEDB_GEOM_WKT
         TILEDB_CHAR
         TILEDB_INT8
         TILEDB_UINT8

tiledb/py_arrow_io_impl.h

@@ -199,6 +199,10 @@ ArrowInfo tiledb_buffer_arrow_fmt(BufferInfo bufferinfo, bool use_list = true) {
   case TILEDB_FLOAT64:
     return ArrowInfo("g");
   case TILEDB_BLOB:
+#if TILEDB_VERSION_MAJOR >= 2 && TILEDB_VERSION_MINOR >= 21
+  case TILEDB_GEOM_WKB:
+  case TILEDB_GEOM_WKT:
+#endif
     return ArrowInfo("B");
   case TILEDB_INT8:
     return ArrowInfo("c");

tiledb/tests/test_attribute.py

@@ -25,7 +25,7 @@ def test_minimal_attribute(self):
         except:
             pytest.fail(f"Could not parse attr._repr_html_(). Saw {attr._repr_html_()}")
 
-    def test_attribute(self, capfd):
+    def test_attribute_name_only(self, capfd):
         attr = tiledb.Attr("foo")
 
         attr.dump()
@@ -252,3 +252,52 @@ def test_modify_attribute_in_schema(self):
                 assert "can't set attribute" in str(exc.value)
             else:
                 assert "object has no setter" in str(exc.value)
+
+    def test_wkt_attribute(self):
+        A = np.array(
+            ["POINT (30 10)", "POLYGON ((3 1, 4 5, 2 2, 1 2, 3 1))"],
+            dtype="S",
+        )
+
+        dom = tiledb.Domain(tiledb.Dim(domain=(0, 1), tile=2))
+        att = tiledb.Attr(dtype="wkt", var=True)
+
+        schema = tiledb.ArraySchema(dom, (att,))
+
+        tiledb.DenseArray.create(self.path("foo"), schema)
+        with tiledb.DenseArray(self.path("foo"), mode="w") as T:
+            T[:] = A
+
+        # read back the data
+        with tiledb.DenseArray(self.path("foo"), mode="r") as T:
+            for i in range(2):
+                assert_array_equal(T[i], A[i])
+
+    def test_wkb_attribute(self):
+        A = np.array(
+            [
+                # representation of POINT (30 10)
+                b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00>@\x00\x00\x00\x00\x00\x00$@",
+                # representation of POLYGON ((3 1, 4 5, 2 2, 1 2, 3 1))
+                (
+                    b"\x01\x03\x00\x00\x00\x01\x00\x00\x00\x05\x00\x00\x00\x00\x00\x00\x00\x00\x00\x08@"
+                    b"\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\x10@\x00\x00\x00\x00\x00\x00\x14@"
+                    b"\x00\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00\xf0?"
+                    b"\x00\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x08@\x00\x00\x00\x00\x00\x00\xf0?"
+                ),
+            ],
+        )
+
+        dom = tiledb.Domain(tiledb.Dim(domain=(0, 1), tile=2))
+        att = tiledb.Attr(dtype="wkb", var=True)
+
+        schema = tiledb.ArraySchema(dom, (att,))
+
+        tiledb.DenseArray.create(self.path("foo"), schema)
+        with tiledb.DenseArray(self.path("foo"), mode="w") as T:
+            T[:] = A
+
+        # read back the data
+        with tiledb.DenseArray(self.path("foo"), mode="r") as T:
+            for i in range(2):
+                assert_array_equal(T[:][i].tobytes(), A[i])