[FEAT] Add Sparse Tensor logical type

daft/daft/__init__.pyi

@@ -993,6 +993,8 @@ class PyDataType:
     @staticmethod
     def tensor(dtype: PyDataType, shape: tuple[int, ...] | None = None) -> PyDataType: ...
     @staticmethod
+    def sparse_tensor(dtype: PyDataType, shape: tuple[int, ...] | None = None) -> PyDataType: ...
+    @staticmethod
     def python() -> PyDataType: ...
     def to_arrow(self, cast_tensor_type_for_ray: builtins.bool | None = None) -> pyarrow.DataType: ...
     def is_numeric(self) -> builtins.bool: ...

daft/datatype.py

@@ -322,6 +322,30 @@ def tensor(
                 raise ValueError("Tensor shape must be a non-empty tuple of ints, but got: ", shape)
         return cls._from_pydatatype(PyDataType.tensor(dtype._dtype, shape))
 
+    @classmethod
+    def sparse_tensor(
+        cls,
+        dtype: DataType,
+        shape: tuple[int, ...] | None = None,
+    ) -> DataType:
+        """Create a SparseTensor DataType: SparseTensor arrays implemented as 'COO Sparse Tensor' representation of n-dimensional arrays of data of the provided ``dtype`` as elements, each of the provided
+        ``shape``.
+
+        If a ``shape`` is given, each ndarray in the column will have this shape.
+
+        If ``shape`` is not given, the ndarrays in the column can have different shapes. This is much more flexible,
+        but will result in a less compact representation and may be make some operations less efficient.
+
+        Args:
+            dtype: The type of the data contained within the tensor elements.
+            shape: The shape of each SparseTensor in the column. This is ``None`` by default, which allows the shapes of
+                each tensor element to vary.
+        """
+        if shape is not None:
+            if not isinstance(shape, tuple) or not shape or any(not isinstance(n, int) for n in shape):
+                raise ValueError("SparseTensor shape must be a non-empty tuple of ints, but got: ", shape)
+        return cls._from_pydatatype(PyDataType.sparse_tensor(dtype._dtype, shape))
+
     @classmethod
     def from_arrow_type(cls, arrow_type: pa.lib.DataType) -> DataType:
         """Maps a PyArrow DataType to a Daft DataType"""

src/daft-core/src/array/growable/logical_growable.rs

@@ -74,6 +74,11 @@ impl_logical_growable!(LogicalTimeGrowable, TimeType);
 impl_logical_growable!(LogicalEmbeddingGrowable, EmbeddingType);
 impl_logical_growable!(LogicalFixedShapeImageGrowable, FixedShapeImageType);
 impl_logical_growable!(LogicalFixedShapeTensorGrowable, FixedShapeTensorType);
+impl_logical_growable!(LogicalSparseTensorGrowable, SparseTensorType);
+impl_logical_growable!(
+    LogicalFixedShapeSparseTensorGrowable,
+    FixedShapeSparseTensorType
+);
 impl_logical_growable!(LogicalImageGrowable, ImageType);
 impl_logical_growable!(LogicalDecimal128Growable, Decimal128Type);
 impl_logical_growable!(LogicalTensorGrowable, TensorType);

src/daft-core/src/array/growable/mod.rs

@@ -203,6 +203,14 @@ impl_growable_array!(
     FixedShapeTensorArray,
     logical_growable::LogicalFixedShapeTensorGrowable<'a>
 );
+impl_growable_array!(
+    SparseTensorArray,
+    logical_growable::LogicalSparseTensorGrowable<'a>
+);
+impl_growable_array!(
+    FixedShapeSparseTensorArray,
+    logical_growable::LogicalFixedShapeSparseTensorGrowable<'a>
+);
 impl_growable_array!(ImageArray, logical_growable::LogicalImageGrowable<'a>);
 impl_growable_array!(TensorArray, logical_growable::LogicalTensorGrowable<'a>);
 impl_growable_array!(