[python] Use new shape in one more spot

single-cell-data · Nov 12, 2024 · efc15c5 · efc15c5
1 parent afe5b69
commit efc15c5
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/apis/python/src/tiledbsoma/io/outgest.py b/apis/python/src/tiledbsoma/io/outgest.py
@@ -443,11 +443,14 @@ def _extract_obsm_or_varm(
     num_cols = width_configs.get(element_name, None)
 
     if num_cols is None:
-        try:
-            used_shape = soma_nd_array.used_shape()
-            num_cols = used_shape[1][1] + 1
-        except SOMAError:
-            pass  # We tried; moving on to next option
+        if soma_nd_array.tiledbsoma_has_upgraded_shape:
+            num_cols = soma_nd_array.shape[1]
+        else:
+            try:
+                used_shape = soma_nd_array.used_shape()
+                num_cols = used_shape[1][1] + 1
+            except SOMAError:
+                pass  # We tried; moving on to next option
 
     if num_cols is None:
         num_rows_times_width, coo_column_count = matrix_tbl.shape