Add extra offset element to buffer size in order to avoid incomplete retry

[sc-skip]

Author: ihnorton

HISTORY.md

@@ -2,7 +2,7 @@
 
 ## Bug Fixes
 * Fix default validity for write to nullable attribute [#994]((https://github.com/TileDB-Inc/TileDB-Py/pull/994)
-
+* Reduce query time for dense var-length arrays by including extra offset element in initial buffer allocation [#1005](https://github.com/TileDB-Inc/TileDB-Py/pull/1005)
 ## API Changes
 * Addition of `ArraySchema.version` to get version of array schema [#949](https://github.com/TileDB-Inc/TileDB-Py/pull/949)
 * Deprecate `coords_filters` from `ArraySchema` [#993](https://github.com/TileDB-Inc/TileDB-Py/pull/993)

tiledb/core.cc

@@ -825,6 +825,11 @@ class PyQuery {
       } 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_

tiledb/tests/test_fixes.py

@@ -87,3 +87,39 @@ def slice_array(a: tiledb.Array, indexer, selection, expected):
             # Important: must get each result here or else assertion
             # failures or exceptions will disappear.
             list(map(lambda x: x.result(), futures))
+
+    # skip, does not currently work, because we cannot force use
+    # of the memory estimate
+    @pytest.mark.skip
+    def test_sc16301_arrow_extra_estimate_dense(self):
+        """
+        Test that dense query of array with var-length attribute completes
+        in one try. We are currently adding an extra element to the offset
+        estimate from libtiledb, in order to avoid an unnecessary pair of
+        query resubmits when the offsets won't fit in the estimated buffer.
+        """
+
+        uri = self.path("test_sc16301_arrow_extra_estimate_dense")
+
+        dim1 = tiledb.Dim(name="d1", dtype="int64", domain=(1, 3))
+        att = tiledb.Attr(name="a1", dtype="<U0", var=True)
+
+        schema = tiledb.ArraySchema(
+            domain=tiledb.Domain(dim1),
+            attrs=(att,),
+            sparse=False,
+            allows_duplicates=False,
+        )
+        tiledb.Array.create(uri, schema)
+
+        with tiledb.open(uri, "w") as A:
+            A[:] = np.array(["aaa", "bb", "c"])
+
+        with tiledb.open(uri) as A:
+            tiledb.stats_enable()
+            r = A[:]
+            assert (
+                """"Context.StorageManager.Query.Reader.loop_num": 1"""
+                in tiledb.stats_dump(print_out=False)
+            )
+            tiledb.stats_disable()