[python] Fix partition readers' naïve assumption that nnz is fast (#3915)

* add work-around for single-fragment nnz

* add raise_if_slow to nnz, and use in partitioning readers

* add raise_if_slow to nnz, and use in partitioning readers

* revert likely buggy nnz optimization

* improve fallback partitioning for ExperimentAxisQuery

* pin numcodecs

* pr fb

* pr fb

* PR fb

Author: bkmartinjr

apis/python/src/tiledbsoma/_query.py

@@ -54,6 +54,7 @@
 if TYPE_CHECKING:
     from ._experiment import Experiment
 from ._constants import SPATIAL_DISCLAIMER
+from ._exception import SOMAError
 from ._fastercsx import CompressedMatrix
 from ._measurement import Measurement
 from ._sparse_nd_array import SparseNDArray
@@ -812,7 +813,10 @@ def _read_as_csr(
 
     d0_joinids = d0_joinids_arr.to_numpy()
     d1_joinids = d1_joinids_arr.to_numpy()
-    nnz = matrix.nnz
+    try:
+        nnz: int | None = matrix._handle._handle.nnz(raise_if_slow=True)
+    except SOMAError:
+        nnz = None
 
     # if able, downcast from int64 - reduces working memory
     index_dtype = (
@@ -852,13 +856,16 @@ def _reindex(batch: pa.RecordBatch) -> pa.RecordBatch:
 
     approx_X_shape = tuple(b - a + 1 for a, b in matrix.non_empty_domain())
     # heuristically derived number (benchmarking). Thesis is that this is roughly 80% of a 1 GiB io buffer,
-    # which is the default for SOMA.
+    # which is the default for SOMA. If we have fast NNZ, use to partition based upon memory size. If we do not
+    # have fast NNZ, pick a default partition size which is large, but not so large as to rule out reasonable
+    # parallelism (in this case, calculated based on typical scRNASeq assay density of 3-6K features).
     target_point_count = 96 * 1024**2
+    fallback_row_count = 32768
     # compute partition size from array density and target point count, rounding to nearest 1024.
     partition_size = (
         max(1024 * round(approx_X_shape[0] * target_point_count / nnz / 1024), 1024)
-        if nnz > 0
-        else approx_X_shape[0]
+        if nnz is not None and nnz > 0
+        else min(fallback_row_count, approx_X_shape[0])
     )
     splits = list(
         range(
@@ -867,7 +874,7 @@ def _reindex(batch: pa.RecordBatch) -> pa.RecordBatch:
             partition_size,
         )
     )
-    if len(splits) > 0:
+    if len(splits) > 1:
         d0_joinids_splits = np.array_split(np.partition(d0_joinids, splits), splits)
         tp = matrix.context.threadpool
         tbl = pa.concat_tables(

apis/python/src/tiledbsoma/io/outgest.py

@@ -106,11 +106,14 @@ def _read_partitioned_sparse(X: SparseNDArray, d0_size: int) -> pa.Table:
     # density of matrix. Magic number determined empirically, as a tradeoff
     # between concurrency and fixed query overhead.
     tgt_point_count = 96 * 1024**2
-    nnz = X.nnz
+    try:
+        nnz: int | None = X._handle._handle.nnz(raise_if_slow=True)
+    except SOMAError:
+        nnz = None
     partition_sz = (
         max(1024 * round(d0_size * tgt_point_count / nnz / 1024), 1024)
-        if nnz > 0
-        else d0_size
+        if nnz is not None and nnz > 0
+        else d0_size  # i.e, no partitioning
     )
     partitions = [
         slice(st, min(st + partition_sz - 1, d0_size - 1))

apis/python/src/tiledbsoma/soma_array.cc

@@ -153,7 +153,11 @@ void load_soma_array(py::module& m) {
             &SOMAArray::config_options_from_schema)
         .def("context", &SOMAArray::ctx)
 
-        .def("nnz", &SOMAArray::nnz, py::call_guard<py::gil_scoped_release>())
+        .def(
+            "nnz",
+            &SOMAArray::nnz,
+            py::arg("raise_if_slow") = false,
+            py::call_guard<py::gil_scoped_release>())
 
         .def_property_readonly("uri", &SOMAArray::uri)
 

libtiledbsoma/src/soma/soma_array.cc

@@ -543,7 +543,7 @@ ArrowTable SOMAArray::_get_core_domainish(enum Domainish which_kind) {
     return ArrowTable(std::move(arrow_array), std::move(arrow_schema));
 }
 
-uint64_t SOMAArray::nnz() {
+uint64_t SOMAArray::nnz(bool raise_if_slow) {
     // Verify array is sparse
     if (schema_->array_type() != TILEDB_SPARSE) {
         throw TileDBSOMAError(
@@ -578,7 +578,7 @@ uint64_t SOMAArray::nnz() {
                          frag_ts.second <= timestamp_->second)) {
                 // fragment overlaps read timestamp range, but isn't fully
                 // contained within: fall back to count_cells to sort that out.
-                return _nnz_slow();
+                return _nnz_slow(raise_if_slow);
             }
         }
         // fall through: fragment is fully contained within the read timestamp
@@ -591,7 +591,7 @@ uint64_t SOMAArray::nnz() {
         // application's job to otherwise ensure uniqueness), then
         // sum-over-fragments is the right thing to do.
         if (!schema_->allows_dups() && frag_ts.first != frag_ts.second) {
-            return _nnz_slow();
+            return _nnz_slow(raise_if_slow);
         }
     }
 
@@ -627,7 +627,7 @@ uint64_t SOMAArray::nnz() {
             "soma_joinid or int64 soma_dim_0: using _nnz_slow",
             tiledb::impl::type_to_str(type_code),
             dim_name));
-        return _nnz_slow();
+        return _nnz_slow(raise_if_slow);
     }
 
     for (uint32_t i = 0; i < fragment_count; i++) {
@@ -668,10 +668,15 @@ uint64_t SOMAArray::nnz() {
         return total_cell_num;
     }
     // Found relevant fragments with overlap, count cells
-    return _nnz_slow();
+    return _nnz_slow(raise_if_slow);
 }
 
-uint64_t SOMAArray::_nnz_slow() {
+uint64_t SOMAArray::_nnz_slow(bool raise_if_slow) {
+    if (raise_if_slow) {
+        throw TileDBSOMAError(
+            "NNZ slow path called with 'raise_if_slow==true'");
+    }
+
     LOG_DEBUG(
         "[SOMAArray] nnz() found consolidated or overlapping fragments, "
         "counting cells...");

libtiledbsoma/src/soma/soma_array.h

@@ -659,7 +659,7 @@ class SOMAArray : public SOMAObject {
      *
      * @return uint64_t Total number of unique cells
      */
-    uint64_t nnz();
+    uint64_t nnz(bool raise_if_slow = false);
 
     /**
      * @brief Get the current capacity of each dimension.
@@ -1164,7 +1164,7 @@ class SOMAArray : public SOMAObject {
     std::shared_ptr<Array> meta_cache_arr_;
 
     // Unoptimized method for computing nnz() (issue `count_cells` query)
-    uint64_t _nnz_slow();
+    uint64_t _nnz_slow(bool raise_if_slow);
 };
 
 }  // namespace tiledbsoma