Inline get_buffer_size_generator and rename some names

Author: gsakkis

Diff for: tiledb/ml/readers/_tensor_gen.py

@@ -49,7 +49,7 @@ def iter_tensors(self, buffer_slice: slice) -> Iterator[np.ndarray]:
             yield buf_array[buffer_slice]
 
 
-class SparseTileDBTensorGenerator(TileDBTensorGenerator[Tensor]):
+class TileDBSparseTensorGenerator(TileDBTensorGenerator[Tensor]):
     def __init__(self, array: tiledb.Array, attrs: Sequence[str]) -> None:
         schema = array.schema
         if schema.ndim != 2:
@@ -96,8 +96,8 @@ def _tensor_from_coo(
         """Convert a scipy.sparse.coo_matrix to a Tensor"""
 
 
-DenseTensor = TypeVar("DenseTensor")
-SparseTensor = TypeVar("SparseTensor")
+DT = TypeVar("DT")
+ST = TypeVar("ST")
 
 
 def tensor_generator(
@@ -109,13 +109,9 @@ def tensor_generator(
     y_attrs: Sequence[str],
     start_offset: int = 0,
     stop_offset: int = 0,
-    dense_tensor_generator_cls: Type[
-        TileDBTensorGenerator[DenseTensor]
-    ] = TileDBNumpyGenerator,
-    sparse_tensor_generator_cls: Type[
-        TileDBTensorGenerator[SparseTensor]
-    ] = SparseTileDBTensorGenerator,
-) -> Iterator[Sequence[Union[DenseTensor, SparseTensor]]]:
+    dense_generator_cls: Type[TileDBTensorGenerator[DT]] = TileDBNumpyGenerator,
+    sparse_generator_cls: Type[TileDBTensorGenerator[ST]] = TileDBSparseTensorGenerator,
+) -> Iterator[Sequence[Union[DT, ST]]]:
     """
     Generator for batches of tensors.
 
@@ -130,20 +126,19 @@ def tensor_generator(
     :param y_attrs: Attribute names of y_array.
     :param start_offset: Start row offset; defaults to 0.
     :param stop_offset: Stop row offset; defaults to number of rows.
-    :param dense_tensor_generator_cls: Dense tensor generator type.
-    :param sparse_tensor_generator_cls: Sparse tensor generator type.
+    :param dense_generator_cls: Dense tensor generator type.
+    :param sparse_generator_cls: Sparse tensor generator type.
     """
-
-    def get_buffer_size_generator(
-        array: tiledb.Array, attrs: Sequence[str]
-    ) -> Union[TileDBTensorGenerator[DenseTensor], TileDBTensorGenerator[SparseTensor]]:
-        if array.schema.sparse:
-            return sparse_tensor_generator_cls(array, attrs)
-        else:
-            return dense_tensor_generator_cls(array, attrs)
-
-    x_gen = get_buffer_size_generator(x_array, x_attrs)
-    y_gen = get_buffer_size_generator(y_array, y_attrs)
+    x_gen: Union[TileDBTensorGenerator[DT], TileDBTensorGenerator[ST]] = (
+        sparse_generator_cls(x_array, x_attrs)
+        if x_array.schema.sparse
+        else dense_generator_cls(x_array, x_attrs)
+    )
+    y_gen: Union[TileDBTensorGenerator[DT], TileDBTensorGenerator[ST]] = (
+        sparse_generator_cls(y_array, y_attrs)
+        if y_array.schema.sparse
+        else dense_generator_cls(y_array, y_attrs)
+    )
     if not stop_offset:
         stop_offset = x_array.shape[0]
     for batch in iter_batches(x_buffer_size, y_buffer_size, start_offset, stop_offset):

Diff for: tiledb/ml/readers/pytorch.py

@@ -11,7 +11,7 @@
 import tiledb
 
 from ._buffer_utils import get_attr_names, get_buffer_size
-from ._tensor_gen import SparseTileDBTensorGenerator, tensor_generator
+from ._tensor_gen import TileDBSparseTensorGenerator, tensor_generator
 
 
 def PyTorchTileDBDataLoader(
@@ -92,7 +92,7 @@ def __init__(
             y_buffer_size=get_buffer_size(y_array, y_attrs, buffer_bytes),
             x_attrs=x_attrs,
             y_attrs=y_attrs,
-            sparse_tensor_generator_cls=PyTorchSparseTileDBTensorGenerator,
+            sparse_generator_cls=PyTorchSparseTensorGenerator,
         )
 
     def __iter__(self) -> Iterator[Sequence[torch.Tensor]]:
@@ -166,7 +166,7 @@ def iter_shuffled(iterable: Iterable[T], buffer_size: int) -> Iterator[T]:
         yield buffer.pop()
 
 
-class PyTorchSparseTileDBTensorGenerator(SparseTileDBTensorGenerator[torch.Tensor]):
+class PyTorchSparseTensorGenerator(TileDBSparseTensorGenerator[torch.Tensor]):
     @staticmethod
     def _tensor_from_coo(
         data: np.ndarray,

Diff for: tiledb/ml/readers/tensorflow.py

@@ -9,7 +9,7 @@
 import tiledb
 
 from ._buffer_utils import get_attr_names, get_buffer_size
-from ._tensor_gen import SparseTileDBTensorGenerator, tensor_generator
+from ._tensor_gen import TileDBSparseTensorGenerator, tensor_generator
 
 # TODO: We have to track the following issues:
 # - https://github.com/tensorflow/tensorflow/issues/47532
@@ -56,7 +56,7 @@ def TensorflowTileDBDataset(
             y_buffer_size=get_buffer_size(y_array, y_attrs, buffer_bytes),
             x_attrs=x_attrs,
             y_attrs=y_attrs,
-            sparse_tensor_generator_cls=TensorflowSparseTileDBTensorGenerator,
+            sparse_generator_cls=TensorflowSparseTensorGenerator,
         ),
         output_signature=(
             *_iter_tensor_specs(x_array.schema, x_attrs),
@@ -77,9 +77,7 @@ def _iter_tensor_specs(
         yield cls(shape=(None, *schema.shape[1:]), dtype=schema.attr(attr).dtype)
 
 
-class TensorflowSparseTileDBTensorGenerator(
-    SparseTileDBTensorGenerator[tf.SparseTensor]
-):
+class TensorflowSparseTensorGenerator(TileDBSparseTensorGenerator[tf.SparseTensor]):
     @staticmethod
     def _tensor_from_coo(
         data: np.ndarray,