diff --git a/daft/daft.pyi b/daft/daft.pyi index 78d30c8d26..8b792adc21 100644 --- a/daft/daft.pyi +++ b/daft/daft.pyi @@ -1003,10 +1003,13 @@ class PyDataType: def is_numeric(self) -> builtins.bool: ... def is_image(self) -> builtins.bool: ... def is_fixed_shape_image(self) -> builtins.bool: ... + def is_list(self) -> builtins.bool: ... def is_tensor(self) -> builtins.bool: ... def is_fixed_shape_tensor(self) -> builtins.bool: ... def is_map(self) -> builtins.bool: ... def is_logical(self) -> builtins.bool: ... + def is_boolean(self) -> builtins.bool: ... + def is_string(self) -> builtins.bool: ... def is_temporal(self) -> builtins.bool: ... def is_equal(self, other: Any) -> builtins.bool: ... @staticmethod diff --git a/daft/datatype.py b/daft/datatype.py index 4077a7b0e3..741bf438d1 100644 --- a/daft/datatype.py +++ b/daft/datatype.py @@ -492,6 +492,15 @@ def _is_fixed_shape_image_type(self) -> builtins.bool: def _is_numeric_type(self) -> builtins.bool: return self._dtype.is_numeric() + def _is_list(self) -> builtins.bool: + return self._dtype.is_list() + + def _is_boolean(self) -> builtins.bool: + return self._dtype.is_boolean() + + def _is_string(self) -> builtins.bool: + return self._dtype.is_string() + def _is_map(self) -> builtins.bool: return self._dtype.is_map() diff --git a/src/daft-core/src/datatypes/dtype.rs b/src/daft-core/src/datatypes/dtype.rs index dbf40d7551..140ffd45d6 100644 --- a/src/daft-core/src/datatypes/dtype.rs +++ b/src/daft-core/src/datatypes/dtype.rs @@ -369,6 +369,21 @@ impl DataType { matches!(self, DataType::Map(..)) } + #[inline] + pub fn is_list(&self) -> bool { + matches!(self, DataType::List(..)) + } + + #[inline] + pub fn is_string(&self) -> bool { + matches!(self, DataType::Utf8) + } + + #[inline] + pub fn is_boolean(&self) -> bool { + matches!(self, DataType::Boolean) + } + #[inline] pub fn is_null(&self) -> bool { match self { diff --git a/src/daft-core/src/python/datatype.rs b/src/daft-core/src/python/datatype.rs index 082605b739..c553a78eb0 100644 --- a/src/daft-core/src/python/datatype.rs +++ b/src/daft-core/src/python/datatype.rs @@ -379,6 +379,18 @@ impl PyDataType { Ok(self.dtype.is_map()) } + pub fn is_list(&self) -> PyResult { + Ok(self.dtype.is_list()) + } + + pub fn is_boolean(&self) -> PyResult { + Ok(self.dtype.is_boolean()) + } + + pub fn is_string(&self) -> PyResult { + Ok(self.dtype.is_string()) + } + pub fn is_logical(&self) -> PyResult { Ok(self.dtype.is_logical()) }