feat: expose image attribute as expression

Author: Jay-ju

daft/expressions/expressions.py

@@ -5153,6 +5153,18 @@ def to_mode(self, mode: str | ImageMode) -> Expression:
         image_mode = lit(mode)._expr
         f = native.get_function_from_registry("to_mode")
         return Expression._from_pyexpr(f(self._expr, mode=image_mode))
+    
+    def attribute(self, name: str) -> Expression:
+        """Get an attribute of the image, such as 'width', 'height', or 'mode'.
+
+        Args:
+            name (str): The name of the attribute to retrieve.
+
+        Returns:
+            Expression: An Expression representing the requested attribute.
+        """
+        f = native.get_function_from_registry("image_attribute")
+        return Expression._from_pyexpr(f(self._expr, lit(name)._expr))
 
 
 class ExpressionPartitioningNamespace(ExpressionNamespace):

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

@@ -56,23 +56,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-image/src/functions/attribute.rs

@@ -0,0 +1,51 @@
+use common_error::{DaftError, DaftResult};
+use daft_core::prelude::*;
+use daft_dsl::{
+    functions::{FunctionArgs, ScalarUDF},
+    ExprRef,
+};
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
+pub struct ImageAttribute;
+
+#[derive(FunctionArgs)]
+struct ImageAttributeArgs<T> {
+    input: T,
+    attr: String,
+}
+
+#[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, attr } = 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)"
+    }
+}

src/daft-image/src/functions/mod.rs

@@ -5,6 +5,7 @@ pub mod decode;
 pub mod encode;
 pub mod resize;
 pub mod to_mode;
+pub mod attribute;
 
 pub struct ImageFunctions;
 
@@ -15,5 +16,6 @@ impl FunctionModule for ImageFunctions {
         parent.add_fn(encode::ImageEncode);
         parent.add_fn(resize::ImageResize);
         parent.add_fn(to_mode::ImageToMode);
+        parent.add_fn(attribute::ImageAttribute);
     }
 }

src/daft-image/src/ops.rs

@@ -38,6 +38,7 @@ pub trait ImageOps {
     fn to_mode(&self, mode: ImageMode) -> DaftResult<Self>
     where
         Self: Sized;
+    fn attribute(&self, attr: &str) -> DaftResult<DataArray<UInt32Type>>;
 }
 
 pub(crate) fn image_array_from_img_buffers(
@@ -194,8 +195,22 @@ impl ImageOps for ImageArray {
             .collect();
         image_array_from_img_buffers(self.name(), &buffers, Some(mode))
     }
+
+    fn attribute(&self, attr: &str) -> DaftResult<DataArray<UInt32Type>> {
+        match attr.to_lowercase().as_str() {
+            "height" | "heights" => Ok(self.heights().clone().rename(self.name())),
+            "width" | "widths" => Ok(self.widths().clone().rename(self.name())),
+            "channel" | "channels" => Ok(self.channels().clone().cast(&DataType::UInt32)?.u32()?.clone().rename(self.name())),
+            "mode" | "modes" => Ok(self.modes().clone().cast(&DataType::UInt32)?.u32()?.clone().rename(self.name())),
+            _ => Err(DaftError::ValueError(format!(
+                "Unsupported image attribute: {}, available: [heights, widths, channels, modes]",
+                attr
+            ))),
+        }
+    }
 }
 
+
 impl ImageOps for FixedShapeImageArray {
     fn encode(&self, image_format: ImageFormat) -> DaftResult<BinaryArray> {
         encode_images(self, image_format)
@@ -254,6 +269,36 @@ impl ImageOps for FixedShapeImageArray {
         };
         fixed_image_array_from_img_buffers(self.name(), &buffers, &mode, *height, *width)
     }
+
+    fn attribute(&self, attr: &str) -> DaftResult<DataArray<UInt32Type>> {
+        let (height, width) = match self.data_type() {
+            DataType::FixedShapeImage(_, h, w) => (h, w),
+            _ => unreachable!("Should be FixedShapeImage type"),
+        };
+
+        match attr.to_lowercase().as_str() {
+            "height" | "heights" => Ok(UInt32Array::from((
+                self.name(),
+                vec![*height; self.len()].as_slice(),
+            ))),
+            "width" | "widths" => Ok(UInt32Array::from((
+                self.name(),
+                vec![*width; self.len()].as_slice(),
+            ))),
+            "channel" | "channels" => Ok(UInt32Array::from((
+                self.name(),
+                vec![self.image_mode().num_channels() as u32; self.len()].as_slice(),
+            ))),
+            "mode" | "modes" => Ok(UInt32Array::from((
+                self.name(),
+                vec![(*self.image_mode() as u8) as u32; self.len()].as_slice(),
+            ))),
+            _ => Err(DaftError::ValueError(format!(
+                "Unsupported image attribute: {}, available: [heights, widths, channels, modes]",
+                attr
+            ))),
+        }
+    }
 }
 
 impl AsImageObj for ImageArray {

src/daft-image/src/series.rs

@@ -197,3 +197,30 @@ pub fn to_mode(s: &Series, mode: ImageMode) -> DaftResult<Series> {
         ))),
     }
 }
+
+
+/// Get metadata attributes from image series
+///
+/// # Arguments
+/// * `s` - Input Series containing image data  
+/// * `attr` - Attribute name to retrieve ("height", "width", "channels", "mode")
+///
+/// # Returns
+/// Series of UInt32 values containing requested attribute
+pub fn attribute(s: &Series, attr: &str) -> DaftResult<Series> {
+    match s.data_type() {
+        DataType::Image(_) => {
+            let array = s.downcast::<ImageArray>()?;
+            Ok(array.attribute(attr)?.into_series())
+        }
+        DataType::FixedShapeImage(..) => {
+            let array = s.downcast::<FixedShapeImageArray>()?;
+            Ok(array.attribute(attr)?.into_series())
+        }
+        dt => Err(DaftError::ValueError(format!(
+            "datatype: {} does not support Image attributes. Occurred while processing Series: {}",
+            dt,
+            s.name()
+        ))),
+    }
+}

tests/recordbatch/image/test_attribute.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import pytest
+
+import daft
+from daft import DataType
+from daft.recordbatch import MicroPartition
+from tests.recordbatch.image.conftest import MODE_TO_NUM_CHANNELS
+
+
+@pytest.mark.parametrize("attr", ["width", "height", "channels", "mode"])
+def test_image_attribute_mixed_shape(mixed_shape_data_fixture, attr):
+    table = daft.from_pydict({"images": mixed_shape_data_fixture})
+    table = table.with_column(attr, daft.col("images").image.attribute(attr))
+    values = table.to_pydict()[attr]
+    
+    image_dtype = mixed_shape_data_fixture.datatype()
+    mode = image_dtype.image_mode
+    mode_name = str(mode).split(".")[-1]
+
+    if attr == "width":
+        assert all(x in (3, 4) for x in values if x is not None)
+    elif attr == "height":
+        assert all(x in (2, 3) for x in values if x is not None)
+    elif attr == "channels":
+        expected_channels = MODE_TO_NUM_CHANNELS[mode_name]
+        assert all(x == expected_channels for x in values if x is not None)
+    elif attr == "mode":
+        expected_mode = MODE_TO_NUM_CHANNELS[mode_name]
+        assert all(x == expected_mode for x in values if x is not None)
+
+