Fix Arrow unit test and add test coverage for 64-bit offset data. (#5236)

This PR fixes the Arrow unit test, which was only running for
`col_sizes=0`. It also tests the `ArrowAdapter` against 64-bit offset
data for completeness.

[sc-52393]

---
TYPE: NO_HISTORY
DESC: Fix Arrow unit test and add test coverage for 64-bit offset data.

---------

Co-authored-by: KiterLuc <[email protected]>

Author: kounelisagis

test/src/unit-arrow.cc

@@ -46,7 +46,7 @@ namespace py = pybind11;
 namespace {
 struct CPPArrayFx {
  public:
-  CPPArrayFx(std::string uri, const uint64_t col_size)
+  CPPArrayFx(std::string uri, const uint64_t col_size, const uint8_t offset)
       : vfs(ctx)
       , uri(uri) {
     if (vfs.is_dir(uri))
@@ -59,41 +59,47 @@ struct CPPArrayFx {
     domain.add_dimensions(d1);
 
     std::vector<Attribute> attrs;
-    attrs.insert(
-        attrs.end(),
-        {Attribute::create<int8_t>(ctx, "int8"),
-         Attribute::create<int16_t>(ctx, "int16"),
-         Attribute::create<int32_t>(ctx, "int32"),
-         Attribute::create<int64_t>(ctx, "int64"),
-
-         Attribute::create<uint8_t>(ctx, "uint8"),
-         Attribute::create<uint16_t>(ctx, "uint16"),
-         Attribute::create<uint32_t>(ctx, "uint32"),
-         Attribute::create<uint64_t>(ctx, "uint64"),
-
-         Attribute::create<float>(ctx, "float32"),
-         Attribute::create<double>(ctx, "float64")});
-
-    // must be constructed manually to get TILEDB_STRING_UTF8 type
-    {
-      auto str_attr = Attribute(ctx, "utf_string1", TILEDB_STRING_UTF8);
+    if (offset == 64) {
+      auto str_attr = Attribute(ctx, "utf_big_string", TILEDB_STRING_UTF8);
       str_attr.set_cell_val_num(TILEDB_VAR_NUM);
       attrs.push_back(str_attr);
-    }
-    {
-      auto str_attr = Attribute(ctx, "utf_string2", TILEDB_STRING_UTF8);
-      str_attr.set_cell_val_num(TILEDB_VAR_NUM);
-      attrs.push_back(str_attr);
-    }
-    {
-      auto str_attr = Attribute(ctx, "tiledb_char", TILEDB_CHAR);
-      str_attr.set_cell_val_num(TILEDB_VAR_NUM);
-      attrs.push_back(str_attr);
-    }
+    } else if (offset == 32) {
+      attrs.insert(
+          attrs.end(),
+          {Attribute::create<int8_t>(ctx, "int8"),
+           Attribute::create<int16_t>(ctx, "int16"),
+           Attribute::create<int32_t>(ctx, "int32"),
+           Attribute::create<int64_t>(ctx, "int64"),
+
+           Attribute::create<uint8_t>(ctx, "uint8"),
+           Attribute::create<uint16_t>(ctx, "uint16"),
+           Attribute::create<uint32_t>(ctx, "uint32"),
+           Attribute::create<uint64_t>(ctx, "uint64"),
+
+           Attribute::create<float>(ctx, "float32"),
+           Attribute::create<double>(ctx, "float64")});
+
+      // must be constructed manually to get TILEDB_STRING_UTF8 type
+      {
+        auto str_attr = Attribute(ctx, "utf_string1", TILEDB_STRING_UTF8);
+        str_attr.set_cell_val_num(TILEDB_VAR_NUM);
+        attrs.push_back(str_attr);
+      }
+      {
+        auto str_attr = Attribute(ctx, "utf_string2", TILEDB_STRING_UTF8);
+        str_attr.set_cell_val_num(TILEDB_VAR_NUM);
+        attrs.push_back(str_attr);
+      }
+      {
+        auto str_attr = Attribute(ctx, "tiledb_char", TILEDB_CHAR);
+        str_attr.set_cell_val_num(TILEDB_VAR_NUM);
+        attrs.push_back(str_attr);
+      }
 
-    // must be constructed manually to get TILEDB_DATETIME_NS type
-    auto datetimens_attr = Attribute(ctx, "datetime_ns", TILEDB_DATETIME_NS);
-    attrs.push_back(datetimens_attr);
+      // must be constructed manually to get TILEDB_DATETIME_NS type
+      auto datetimens_attr = Attribute(ctx, "datetime_ns", TILEDB_DATETIME_NS);
+      attrs.push_back(datetimens_attr);
+    }
 
     FilterList filters(ctx);
     filters.add_filter({ctx, TILEDB_FILTER_LZ4});
@@ -217,9 +223,11 @@ void allocate_query_buffers(tiledb::Query* const query) {
 
 };  // namespace
 
-void test_for_column_size(size_t col_size) {
-  std::string uri("test_arrow_io_" + std::to_string(col_size));
-  CPPArrayFx _fx(uri, col_size);
+void test_for_column_size(const size_t col_size, const uint8_t offset) {
+  std::string uri(
+      "test_arrow_io_" + std::to_string(col_size) + "_" +
+      std::to_string(offset));
+  CPPArrayFx _fx(uri, col_size, offset);
 
   py::object py_data_source;
   py::object py_data_arrays;
@@ -234,7 +242,8 @@ void test_for_column_size(size_t col_size) {
   unit_arrow = py::module::import("unit_arrow");
 
   // this class generates random test data for each attribute
-  auto h_data_source = unit_arrow.attr("DataFactory");
+  auto class_name = "DataFactory" + std::to_string(offset);
+  auto h_data_source = unit_arrow.attr(class_name.c_str());
   py_data_source = h_data_source(py::int_(col_size));
   py_data_names = py_data_source.attr("names");
   py_data_arrays = py_data_source.attr("arrays");
@@ -248,7 +257,7 @@ void test_for_column_size(size_t col_size) {
      * Test write
      */
     Config config;
-    config["sm.var_offsets.bitsize"] = 32;
+    config["sm.var_offsets.bitsize"] = offset;
     config["sm.var_offsets.mode"] = "elements";
     config["sm.var_offsets.extra_element"] = "true";
     Context ctx(config);
@@ -303,14 +312,12 @@ void test_for_column_size(size_t col_size) {
   // However, there is an unexplained crash due to an early destructor
   // when both brace scopes are converted to SECTIONs.
   // SECTION("Test reading data back via ArrowAdapter into pyarrow arrays")
-
-  // test both bitsize read modes
-  for (auto bitsize : {32, 64}) {
+  {
     /*
      * Test read
      */
     Config config;
-    config["sm.var_offsets.bitsize"] = bitsize;
+    config["sm.var_offsets.bitsize"] = 64;
     config["sm.var_offsets.mode"] = "elements";
     config["sm.var_offsets.extra_element"] = "true";
     Context ctx(config);
@@ -437,8 +444,9 @@ TEST_CASE("Arrow IO integration tests", "[arrow]") {
 #endif
 
   // do not use catch2 GENERATE here: it causes bad things to happen w/ python
-  uint64_t col_sizes[] = {0};  //,1,2,3,4,11,103};
+  uint64_t col_sizes[] = {0, 1, 2, 3, 4, 11, 103};
   for (auto sz : col_sizes) {
-    test_for_column_size(sz);
+    test_for_column_size(sz, 32);
+    test_for_column_size(sz, 64);
   }
 }

test/src/unit_arrow.py

@@ -11,16 +11,10 @@
 #          Data generators                                                   #
 # ************************************************************************** #
 
-# python 2 vs 3 compatibility
-if sys.hexversion >= 0x3000000:
-    getchr = chr
-else:
-    getchr = unichr
-
 def gen_chr(max, printable=False):
     while True:
         # TODO we exclude 0x0 here because the key API does not embedded NULL
-        s = getchr(random.randrange(1, max))
+        s = chr(random.randrange(1, max))
         if printable and not s.isprintable():
             continue
         if len(s) > 0:
@@ -59,10 +53,10 @@ def rand_ascii_bytes(size=5, printable=False):
     return b''.join([gen_chr(127, printable).encode('utf-8') for _ in range(0,size)])
 
 # ************************************************************************** #
-#           Test class                                                       #
+#           Test classes                                                     #
 # ************************************************************************** #
 
-class DataFactory():
+class DataFactory32():
   def __init__(self, col_size):
     self.results = {}
     self.col_size = col_size
@@ -81,35 +75,65 @@ def create(self):
     for dt in (np.int8, np.uint8, np.int16, np.uint16, np.int32, np.uint32, np.int64, np.uint64):
         key = np.dtype(dt).name
         dtinfo = np.iinfo(dt)
-        self.data[key] = np.random.randint(dtinfo.min, dtinfo.max, size=col_size, dtype=dt)
+        self.data[key] = pa.array(np.random.randint(dtinfo.min, dtinfo.max, size=col_size, dtype=dt))
 
     for dt in (np.float32, np.float64):
         key = np.dtype(dt).name
-        self.data[key] = np.random.rand(col_size).astype(dt)
+        self.data[key] = pa.array(np.random.rand(col_size).astype(dt))
 
     # var-len (strings)
-    self.data['tiledb_char'] = np.array([rand_ascii_bytes(np.random.randint(1,100))
-                                       for _ in range(col_size)]).astype("S1")
-    self.data['utf_string1'] = np.array([rand_utf8(np.random.randint(1, 100))
+    self.data['tiledb_char'] = pa.array(np.array([rand_ascii_bytes(np.random.randint(1,100))
+                                       for _ in range(col_size)]).astype("S1"))
+    utf_strings = np.array([rand_utf8(np.random.randint(1, 100))
                                        for _ in range(col_size)]).astype("U0")
+    self.data['utf_string1'] = pa.array(utf_strings)
 
-    # another version with some important cells set to empty
-    self.data['utf_string2'] = np.array([rand_utf8(np.random.randint(0, 100))
-                                       for _ in range(col_size)]).astype("U0")
+    # another version with some cells set to empty
+    utf_strings[np.random.randint(0, col_size, size=col_size//2)] = ''
+    self.data['utf_string2'] = pa.array(utf_strings)
+
+    self.data['datetime_ns'] = pa.array(rand_datetime64_array(col_size))
+
+    ##########################################################################
+
+    self.arrays = list(self.data.values())
+    self.names = list(self.data.keys())
+
+  def import_result(self, name, c_array, c_schema):
+    self.results[name] = pa.Array._import_from_c(c_array, c_schema)
+
+  def check(self):
+    for key,val in self.data.items():
+      assert (key in self.results), "Expected key '{}' not found in results!".format(key)
+
+      res_val = self.results[key]
+      assert_array_equal(val, res_val)
 
-    utf_string2 = self.data['utf_string2']
-    for i in range(len(utf_string2)):
-        self.data['utf_string2'][i] = ''
-    range_start = len(utf_string2) - 1
-    range_end = len(utf_string2) % 3
-    for i in range(range_start, range_end, -1):
-        self.data['utf_string2'][i] = ''
+    return True
+
+
+class DataFactory64():
+  def __init__(self, col_size):
+    self.results = {}
+    self.col_size = col_size
+    self.create()
+
+  def __len__(self):
+    if not self.data:
+      raise ValueError("Uninitialized data")
+    return len(self.data)
+
+  def create(self):
+    # generate test data for all columns
+    col_size = self.col_size
+    self.data = {}
 
-    self.data['datetime_ns'] = rand_datetime64_array(col_size)
+    self.data['utf_big_string'] = pa.array([rand_utf8(np.random.randint(1, 100))
+                                       for _ in range(col_size)], type=pa.large_utf8())
 
     ##########################################################################
 
-    self.arrays = [pa.array(val) for val in self.data.values()]
+    self.arrays = list(self.data.values())
     self.names = list(self.data.keys())
 
   def import_result(self, name, c_array, c_schema):
@@ -122,4 +146,4 @@ def check(self):
       res_val = self.results[key]
       assert_array_equal(val, res_val)
 
-    return True
+    return True

tiledb/sm/cpp_api/arrow_io_impl.h

@@ -759,8 +759,6 @@ void ArrowExporter::export_(
   if (bufferinfo.is_var) {
     buffers = {nullptr, bufferinfo.offsets, bufferinfo.data};
   } else {
-    cpp_schema = new CPPArrowSchema(
-        name, arrow_fmt.fmt_, std::nullopt, arrow_flags, {}, {});
     buffers = {nullptr, bufferinfo.data};
   }
   cpp_schema->export_ptr(schema);