Remove redundant memory alloc for Non Shuffled Scenario (#38)

* Remove redundant memory alloc

* Cleanup redundant typecasting

* Set default batch size

Author: AaditSaluja

src/tiledbsoma_ml/_io_batch_iterable.py

@@ -55,27 +55,36 @@ def __iter__(self) -> Iterator[IOBatch]:
         """Emit |IOBatch|'s."""
         # Because obs/var IDs have been partitioned/split/shuffled upstream of this class, this RNG does not need to be
         # identical across sub-processes, but seeding is supported anyway, for testing/reproducibility.
-        shuffle_rng = np.random.default_rng(self.seed)
         X = self.X
         context = X.context
+
+        # only build rng if we shuffle
+        shuffle_rng = np.random.default_rng(self.seed) if self.shuffle else None
+
         obs_column_names = (
             list(self.obs_column_names)
             if "soma_joinid" in self.obs_column_names
             else ["soma_joinid", *self.obs_column_names]
         )
         # NOTE: `.astype("int64")` works around the `np.int64` singleton failing reference-equality after cross-process
         # SerDes.
-        var_joinids = self.var_joinids.astype("int64")
+        var_joinids = np.asarray(
+            self.var_joinids, dtype=np.int64
+        )  # as array only typecasts if needed
         var_indexer = IntIndexer(var_joinids, context=context)
 
         for obs_coords in self.io_batch_ids:
             st_time = time.perf_counter()
-            obs_shuffled_coords = (
+
+            if shuffle_rng is None:
+                obs_order = np.fromiter(
+                    obs_coords, dtype=np.int64, count=len(obs_coords)
+                )
+            else:
                 np.array(obs_coords)
-                if not self.shuffle
-                else shuffle_rng.permuted(obs_coords)
-            )
-            obs_indexer = IntIndexer(obs_shuffled_coords, context=context)
+                obs_order = shuffle_rng.permuted(obs_coords)
+
+            obs_indexer = IntIndexer(obs_order, context=context)
             logger.debug(
                 f"Retrieving next SOMA IO batch of length {len(obs_coords)}..."
             )
@@ -123,14 +132,14 @@ def make_io_buffer(
                 .concat()
                 .to_pandas()
                 .set_index("soma_joinid")
-                .reindex(obs_shuffled_coords, copy=False)
+                .reindex(obs_order, copy=False)
                 .reset_index()  # demote "soma_joinid" to a column
                 [self.obs_column_names]
             )  # fmt: on
 
             X_io_batch = CSR_IO_Buffer.merge(tuple(_io_buf_iter))
 
-            del obs_indexer, obs_coords, obs_shuffled_coords, _io_buf_iter
+            del obs_indexer, obs_coords, obs_order, _io_buf_iter
             gc.collect()
 
             tm = time.perf_counter() - st_time