more itemizations in assay-matrix ingestor

apis/python/src/tiledbsc/assay_matrix.py

@@ -53,12 +53,14 @@ def __init__(
  * For reading from an already-populated SOMA, we wish to avoid cache-coherency issues.
  """
  super().__init__(uri=uri, name=name, parent=parent)
+ s0 = self.timing_start("__init__", "total")
 
  self.row_dim_name = row_dim_name
  self.col_dim_name = col_dim_name
  self.attr_name = "value"
  self.row_dataframe = row_dataframe
  self.col_dataframe = col_dataframe
+ self.timing_end(s0)
 
  # ----------------------------------------------------------------
  def shape(self) -> Tuple[int, int]:
@@ -69,12 +71,15 @@ def shape(self) -> Tuple[int, int]:
 
  Note: currently implemented via data scan -- will be optimized for TileDB core 2.10.
  """
+ s1 = self.timing_start("shape", "total")
  with self._open():
  # These TileDB arrays are string-dimensioned sparse arrays so there is no '.shape'.
  # Instead we compute it ourselves. See also:
  num_rows = self.row_dataframe.shape()[0]
  num_cols = self.col_dataframe.shape()[0]
- return (num_rows, num_cols)
+ retval = (num_rows, num_cols)
+ self.timing_end(s1)
+ return retval
 
  # ----------------------------------------------------------------
  def dim_select(
@@ -89,7 +94,13 @@ def dim_select(
  Either or both of the ID lists may be `None`, meaning, do not subselect along
  that dimension. If both ID lists are `None`, the entire matrix is returned.
  """
+ s0 = self.timing_start("dim_select", "open")
+
+ s1 = self.timing_start("dim_select", "open")
  with tiledb.open(self.uri, ctx=self._ctx) as A:
+ self.timing_end(s1)
+
+ s2 = self.timing_start("dim_select", "tiledb_query")
  query = A.query(return_arrow=return_arrow)
  if obs_ids is None:
  if var_ids is None:
@@ -101,8 +112,14 @@ def dim_select(
  df = query.df[obs_ids, :]
  else:
  df = query.df[obs_ids, var_ids]
+ self.timing_end(s2)
+
+ s3 = self.timing_start("dim_select", "set_index")
  if not return_arrow:
  df.set_index([self.row_dim_name, self.col_dim_name], inplace=True)
+ self.timing_end(s3)
+
+ self.timing_end(s0)
  return df
 
  # ----------------------------------------------------------------
@@ -126,15 +143,21 @@ def csr(
  """
  Like `.df()` but returns results in `scipy.sparse.csr_matrix` format.
  """
- return self._csr_or_csc("csr", obs_ids, var_ids)
+ s0 = self.timing_start("csr", "total")
+ retval = self._csr_or_csc("csr", obs_ids, var_ids)
+ self.timing_end(s0)
+ return retval
 
  def csc(
  self, obs_ids: Optional[Ids] = None, var_ids: Optional[Ids] = None
  ) -> sp.csc_matrix:
  """
  Like `.df()` but returns results in `scipy.sparse.csc_matrix` format.
  """
- return self._csr_or_csc("csc", obs_ids, var_ids)
+ s0 = self.timing_start("csc", "total")
+ retval = self._csr_or_csc("csc", obs_ids, var_ids)
+ self.timing_end(s0)
+ return retval
 
  def _csr_or_csc(
  self,
@@ -168,6 +191,7 @@ def from_matrix_and_dim_values(
  `scipy.sparse.csr_matrix`, `scipy.sparse.csc_matrix`, `numpy.ndarray`, etc.
  For ingest from `AnnData`, these should be `ann.obs_names` and `ann.var_names`.
  """
+ s0 = self.timing_start("from_matrix_and_dim_values", "total")
 
  s = util.get_start_stamp()
  log_io(
@@ -206,11 +230,14 @@ def from_matrix_and_dim_values(
  util.format_elapsed(s, f"{self._indent}FINISH WRITING {self.uri}"),
  )
 
+ self.timing_end(s0)
+
  # ----------------------------------------------------------------
  def _create_empty_array(self, matrix_dtype: np.dtype) -> None:
  """
  Create a TileDB 2D sparse array with string dimensions and a single attribute.
  """
+ s0 = self.timing_start("_create_empty_array", "total")
 
  level = self._soma_options.string_dim_zstd_level
  dom = tiledb.Domain(
@@ -253,6 +280,7 @@ def _create_empty_array(self, matrix_dtype: np.dtype) -> None:
  )
 
  tiledb.Array.create(self.uri, sch, ctx=self._ctx)
+ self.timing_end(s0)
 
  # ----------------------------------------------------------------
  def ingest_data_whole(
@@ -269,6 +297,7 @@ def ingest_data_whole(
  :param row_names: List of row names.
  :param col_names: List of column names.
  """
+ s0 = self.timing_start("ingest_data_whole", "total")
 
  assert len(row_names) == matrix.shape[0]
  assert len(col_names) == matrix.shape[1]
@@ -279,6 +308,7 @@ def ingest_data_whole(
 
  with tiledb.open(self.uri, mode="w", ctx=self._ctx) as A:
  A[d0, d1] = mat_coo.data
+ self.timing_end(s0)
 
  # ----------------------------------------------------------------
  # Example: suppose this 4x3 is to be written in two chunks of two rows each
@@ -326,7 +356,9 @@ def ingest_data_rows_chunked(
  :param row_names: List of row names.
  :param col_names: List of column names.
  """
+ s0 = self.timing_start("ingest_data_rows_chunked", "total")
 
+ s1 = self.timing_start("ingest_data_rows_chunked", "sortprep")
  assert len(row_names) == matrix.shape[0]
  assert len(col_names) == matrix.shape[1]
 
@@ -346,13 +378,20 @@ def ingest_data_rows_chunked(
  f"{self._indent}START ingest_data_rows_chunked",
  )
 
+ self.timing_end(s1)
+
  eta_tracker = util.ETATracker()
+ s2 = self.timing_start("ingest_data_rows_chunked", "open")
  with tiledb.open(self.uri, mode="w", ctx=self._ctx) as A:
+ self.timing_end(s2)
+
  nrow = len(sorted_row_names)
 
  i = 0
  while i < nrow:
  t1 = time.time()
+
+ s3 = self.timing_start("ingest_data_rows_chunked", "chunkprep")
  # Find a number of CSR rows which will result in a desired nnz for the chunk.
  chunk_size = util._find_csr_chunk_size(
  matrix, permutation, i, self._soma_options.goal_chunk_nnz
@@ -365,6 +404,7 @@ def ingest_data_rows_chunked(
  # Write the chunk-COO to TileDB.
  d0 = sorted_row_names[chunk_coo.row + i]
  d1 = col_names[chunk_coo.col]
+ self.timing_end(s3)
 
  if len(d0) == 0:
  i = i2
@@ -390,7 +430,9 @@ def ingest_data_rows_chunked(
  )
 
  # Write a TileDB fragment
+ s4 = self.timing_start("ingest_data_rows_chunked", "tiledb-write")
  A[d0, d1] = chunk_coo.data
+ self.timing_end(s4)
 
  t2 = time.time()
  chunk_seconds = t2 - t1
@@ -413,6 +455,7 @@ def ingest_data_rows_chunked(
  f"{self._indent}FINISH __ingest_coo_data_string_dims_rows_chunked",
  ),
  )
+ self.timing_end(s0)
 
  # This method is very similar to ingest_data_rows_chunked. The code is largely repeated,
  # and this is intentional. The algorithm here is non-trivial (among the most non-trivial
@@ -432,6 +475,7 @@ def ingest_data_cols_chunked(
  :param row_names: List of row names.
  :param col_names: List of column names.
  """
+ s0 = self.timing_start("ingest_data_cols_chunked", "total")
 
  assert len(row_names) == matrix.shape[0]
  assert len(col_names) == matrix.shape[1]
@@ -519,6 +563,7 @@ def ingest_data_cols_chunked(
  f"{self._indent}FINISH __ingest_coo_data_string_dims_rows_chunked",
  ),
  )
+ self.timing_end(s0)
 
  # This method is very similar to ingest_data_rows_chunked. The code is largely repeated,
  # and this is intentional. The algorithm here is non-trivial (among the most non-trivial
@@ -538,6 +583,7 @@ def ingest_data_dense_rows_chunked(
  :param row_names: List of row names.
  :param col_names: List of column names.
  """
+ s0 = self.timing_start("ingest_data_dense_rows_chunked", "total")
 
  assert len(row_names) == matrix.shape[0]
  assert len(col_names) == matrix.shape[1]
@@ -627,6 +673,7 @@ def ingest_data_dense_rows_chunked(
  f"{self._indent}FINISH __ingest_coo_data_string_dims_dense_rows_chunked",
  ),
  )
+ self.timing_end(s0)
 
  # ----------------------------------------------------------------
  def to_csr_matrix(self, row_labels: Labels, col_labels: Labels) -> sp.csr_matrix:
@@ -638,6 +685,7 @@ def to_csr_matrix(self, row_labels: Labels, col_labels: Labels) -> sp.csr_matrix
  be in the same order as they were in any anndata object which was used to create the
  TileDB storage.
  """
+ s0 = self.timing_start("to_csr_matrix", "total")
 
  s = util.get_start_stamp()
  log_io(None, f"{self._indent}START read {self.uri}")
@@ -649,4 +697,5 @@ def to_csr_matrix(self, row_labels: Labels, col_labels: Labels) -> sp.csr_matrix
  util.format_elapsed(s, f"{self._indent}FINISH read {self.uri}"),
  )
 
+ self.timing_end(s0)
  return csr