feat: expose image attribute as expression (#4848)

Author: Jay-ju

.pre-commit-config.yaml

@@ -85,7 +85,7 @@ repos:
     language: system
     types: [rust]
     pass_filenames: false
-    args: [--workspace, --all-targets]
+    args: [--workspace, --all-targets, -v]
 
   - id: cargo-check-all-features
     name: cargo check (all features)

Cargo.lock

@@ -74,7 +74,7 @@ def refresh_logger() -> None:
     from_pylist,
     from_ray_dataset,
 )
-from daft.daft import ImageFormat, ImageMode, ResourceRequest
+from daft.daft import ImageFormat, ImageMode, ImageProperty, ResourceRequest
 from daft.dataframe import DataFrame
 from daft.schema import Schema
 from daft.datatype import DataType, TimeUnit
@@ -161,6 +161,7 @@ def refresh_logger() -> None:
     "Identifier",
     "ImageFormat",
     "ImageMode",
+    "ImageProperty",
     "ResourceRequest",
     "Schema",
     "Series",

daft/__init__.py

@@ -77,6 +77,17 @@ class ImageMode(Enum):
         """
         ...
 
+class ImageProperty(Enum):
+    """Supported image properties for Daft's image type."""
+
+    Height = 1
+    Width = 2
+    Channel = 3
+    Mode = 4
+
+    @staticmethod
+    def from_property_string(attr: str) -> ImageProperty: ...
+
 class PyWindowBoundary:
     """Represents a window frame boundary in window functions."""
 

daft/daft/__init__.pyi

@@ -20,6 +20,7 @@
     CountMode,
     ImageFormat,
     ImageMode,
+    ImageProperty,
     ResourceRequest,
     initialize_udfs,
     resolved_col,
@@ -5224,6 +5225,60 @@ def to_mode(self, mode: str | ImageMode) -> Expression:
         f = native.get_function_from_registry("to_mode")
         return Expression._from_pyexpr(f(self._expr, mode=image_mode))
 
+    def attribute(self, name: Literal["width", "height", "channel", "mode"] | ImageProperty) -> Expression:
+        """Get a property of the image, such as 'width', 'height', 'channel', or 'mode'.
+
+        Args:
+            name (str): The name of the property to retrieve.
+
+        Returns:
+            Expression: An Expression representing the requested property.
+        """
+        if isinstance(name, str):
+            name = ImageProperty.from_property_string(name)
+        f = native.get_function_from_registry("image_attribute")
+        return Expression._from_pyexpr(f(self._expr, lit(name)._expr))
+
+    def width(self) -> Expression:
+        """Gets the width of an image in pixels.
+
+        Example:
+            >>> # Create a dataframe with an image column
+            >>> df = ...  # doctest: +SKIP
+            >>> df = df.with_column("width", df["images"].image.width())  # doctest: +SKIP
+        """
+        return self.attribute("width")
+
+    def height(self) -> Expression:
+        """Gets the height of an image in pixels.
+
+        Example:
+            >>> # Create a dataframe with an image column
+            >>> df = ...  # doctest: +SKIP
+            >>> df = df.with_column("height", df["images"].image.height())  # doctest: +SKIP
+        """
+        return self.attribute("height")
+
+    def channel(self) -> Expression:
+        """Gets the number of channels in an image.
+
+        Example:
+            >>> # Create a dataframe with an image column
+            >>> df = ...  # doctest: +SKIP
+            >>> df = df.with_column("channel", df["images"].image.channel())  # doctest: +SKIP
+        """
+        return self.attribute("channel")
+
+    def mode(self) -> Expression:
+        """Gets the mode of an image as a string.
+
+        Example:
+            >>> # Create a dataframe with an image column
+            >>> df = ...  # doctest: +SKIP
+            >>> df = df.with_column("mode", df["images"].image.mode())  # doctest: +SKIP
+        """
+        return self.attribute("mode")
+
 
 class ExpressionPartitioningNamespace(ExpressionNamespace):
     """The following methods are available under the `expr.partition` attribute."""

daft/expressions/expressions.py

@@ -36,23 +36,39 @@ impl ImageArray {
     }
 
     pub fn channel_array(&self) -> &arrow2::array::UInt16Array {
-        let array = self.physical.children.get(Self::IMAGE_CHANNEL_IDX).unwrap();
-        array.u16().unwrap().as_arrow()
+        self.channels().as_arrow()
     }
 
     pub fn height_array(&self) -> &arrow2::array::UInt32Array {
-        let array = self.physical.children.get(Self::IMAGE_HEIGHT_IDX).unwrap();
-        array.u32().unwrap().as_arrow()
+        self.heights().as_arrow()
     }
 
     pub fn width_array(&self) -> &arrow2::array::UInt32Array {
-        let array = self.physical.children.get(Self::IMAGE_WIDTH_IDX).unwrap();
-        array.u32().unwrap().as_arrow()
+        self.widths().as_arrow()
     }
 
     pub fn mode_array(&self) -> &arrow2::array::UInt8Array {
+        self.modes().as_arrow()
+    }
+
+    pub fn channels(&self) -> &DataArray<UInt16Type> {
+        let array = self.physical.children.get(Self::IMAGE_CHANNEL_IDX).unwrap();
+        array.u16().unwrap()
+    }
+
+    pub fn heights(&self) -> &DataArray<UInt32Type> {
+        let array = self.physical.children.get(Self::IMAGE_HEIGHT_IDX).unwrap();
+        array.u32().unwrap()
+    }
+
+    pub fn widths(&self) -> &DataArray<UInt32Type> {
+        let array = self.physical.children.get(Self::IMAGE_WIDTH_IDX).unwrap();
+        array.u32().unwrap()
+    }
+
+    pub fn modes(&self) -> &DataArray<UInt8Type> {
         let array = self.physical.children.get(Self::IMAGE_MODE_IDX).unwrap();
-        array.u8().unwrap().as_arrow()
+        array.u8().unwrap()
     }
 
     pub fn from_list_array(

src/daft-core/src/array/image_array.rs

@@ -9,7 +9,7 @@ use super::{FromLiteral, Literal, deserializer::LiteralDeserializer};
 use crate::python::{PyDataType, PyTimeUnit};
 use crate::{
     datatypes::IntervalValue,
-    prelude::{CountMode, DataType, ImageFormat, ImageMode, TimeUnit},
+    prelude::{CountMode, DataType, ImageFormat, ImageMode, ImageProperty, TimeUnit},
     series::Series,
 };
 
@@ -223,7 +223,8 @@ impl_float_fromliteral!(f32);
 impl_float_fromliteral!(f64);
 impl_pyobj_fromliteral!(IOConfig, common_io_config::python::IOConfig);
 impl_pyobj_fromliteral!(ImageMode, ImageMode);
-impl_pyobj_fromliteral!(ImageFormat, ImageFormat);
+impl_pyobj_fromliteral!(ImageProperty, ImageProperty);
 impl_pyobj_fromliteral!(CountMode, CountMode);
 impl_pyobj_fromliteral!(TimeUnit, PyTimeUnit);
 impl_pyobj_fromliteral!(DataType, PyDataType);
+impl_pyobj_fromliteral!(ImageFormat, ImageFormat);

src/daft-core/src/lit/conversions.rs

@@ -4,6 +4,7 @@
 
 // Re-export arrow2 bitmap
 pub use arrow2::bitmap;
+pub use daft_schema::image_property::ImageProperty;
 // Re-export core series structures
 pub use daft_schema::schema::{Schema, SchemaRef};
 

src/daft-core/src/prelude.rs

@@ -5,6 +5,7 @@ common-error = {path = "../common/error", default-features = false}
 common-image = {workspace = true}
 daft-core = {path = "../daft-core", default-features = false}
 daft-dsl = {path = "../daft-dsl", default-features = false}
+daft-schema = {path = "../daft-schema", default-features = false}
 log = {workspace = true}
 serde = {workspace = true}
 typetag = {workspace = true}

src/daft-image/Cargo.toml

@@ -0,0 +1,52 @@
+use common_error::{DaftError, DaftResult};
+use daft_core::prelude::*;
+use daft_dsl::{
+    ExprRef,
+    functions::{FunctionArgs, ScalarUDF},
+};
+use daft_schema::image_property::ImageProperty;
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
+pub struct ImageAttribute;
+
+#[derive(FunctionArgs)]
+struct ImageAttributeArgs<T> {
+    input: T,
+    attr: ImageProperty,
+}
+
+#[typetag::serde]
+impl ScalarUDF for ImageAttribute {
+    fn call(&self, inputs: FunctionArgs<Series>) -> DaftResult<Series> {
+        let ImageAttributeArgs { input, attr } = inputs.try_into()?;
+        crate::series::attribute(&input, attr)
+    }
+
+    fn name(&self) -> &'static str {
+        "image_attribute"
+    }
+
+    fn get_return_field(
+        &self,
+        inputs: FunctionArgs<ExprRef>,
+        schema: &Schema,
+    ) -> DaftResult<Field> {
+        let ImageAttributeArgs { input, .. } = inputs.try_into()?;
+
+        let input_field = input.to_field(schema)?;
+        match input_field.dtype {
+            DataType::Image(_) | DataType::FixedShapeImage(..) => {
+                Ok(Field::new(input_field.name, DataType::UInt32))
+            }
+            _ => Err(DaftError::TypeError(format!(
+                "Image attribute can only be retrieved from ImageArrays, got {}",
+                input_field.dtype
+            ))),
+        }
+    }
+
+    fn docstring(&self) -> &'static str {
+        "Extracts metadata attributes from image series (height/width/channels/mode)"
+    }
+}