feat: expose image attribute as expression

Jay-ju · Jay-ju · commit 07ea9c5ca647 · 2025-08-14T19:03:19.000+08:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/daft/__init__.py b/daft/__init__.py
@@ -69,7 +69,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
@@ -148,6 +148,7 @@ def refresh_logger() -> None:
     "Identifier",
     "ImageFormat",
     "ImageMode",
+    "ImageProperty",
     "ResourceRequest",
     "Schema",
     "Series",
diff --git a/daft/daft/__init__.pyi b/daft/daft/__init__.pyi
@@ -76,6 +76,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."""
 
diff --git a/daft/expressions/expressions.py b/daft/expressions/expressions.py
@@ -21,6 +21,7 @@
     CountMode,
     ImageFormat,
     ImageMode,
+    ImageProperty,
     ResourceRequest,
     initialize_udfs,
     resolved_col,
@@ -5217,6 +5218,20 @@ 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))
+
 
 class ExpressionPartitioningNamespace(ExpressionNamespace):
     """The following methods are available under the `expr.partition` attribute."""
diff --git a/src/daft-core/src/array/image_array.rs b/src/daft-core/src/array/image_array.rs
@@ -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(
diff --git a/src/daft-core/src/lit/conversions.rs b/src/daft-core/src/lit/conversions.rs
@@ -8,7 +8,7 @@ use super::{deserializer::LiteralDeserializer, FromLiteral, Literal};
 use crate::python::{PyDataType, PyTimeUnit};
 use crate::{
     datatypes::IntervalValue,
-    prelude::{CountMode, DataType, ImageFormat, ImageMode, TimeUnit},
+    prelude::{CountMode, DataType, ImageFormat, ImageMode, ImageProperty, TimeUnit},
     series::Series,
 };
 
@@ -219,7 +219,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);
diff --git a/src/daft-core/src/prelude.rs b/src/daft-core/src/prelude.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};
 
diff --git a/src/daft-image/Cargo.toml b/src/daft-image/Cargo.toml
@@ -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}
diff --git a/src/daft-image/src/functions/attribute.rs b/src/daft-image/src/functions/attribute.rs
@@ -0,0 +1,52 @@
+use common_error::{DaftError, DaftResult};
+use daft_core::prelude::*;
+use daft_dsl::{
+    functions::{FunctionArgs, ScalarUDF},
+    ExprRef,
+};
+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)"
+    }
+}
diff --git a/src/daft-image/src/functions/mod.rs b/src/daft-image/src/functions/mod.rs
@@ -1,5 +1,6 @@
 use daft_dsl::functions::FunctionModule;
 
+pub mod attribute;
 pub mod crop;
 pub mod decode;
 pub mod encode;
@@ -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);
     }
 }
diff --git a/src/daft-image/src/ops.rs b/src/daft-image/src/ops.rs
@@ -11,6 +11,7 @@ use daft_core::{
     datatypes::prelude::*,
     prelude::ImageArray,
 };
+use daft_schema::image_property::ImageProperty;
 
 use crate::{iters::ImageBufferIter, CountingWriter};
 
@@ -31,6 +32,7 @@ pub trait ImageOps {
     fn to_mode(&self, mode: ImageMode) -> DaftResult<Self>
     where
         Self: Sized;
+    fn attribute(&self, attr: ImageProperty) -> DaftResult<DataArray<UInt32Type>>;
 }
 
 impl ImageOps for ImageArray {
@@ -75,6 +77,27 @@ impl ImageOps for ImageArray {
             .collect();
         image_array_from_img_buffers(self.name(), &buffers, Some(mode))
     }
+
+    fn attribute(&self, attr: ImageProperty) -> DaftResult<DataArray<UInt32Type>> {
+        match attr {
+            ImageProperty::Height => Ok(self.heights().clone().rename(self.name())),
+            ImageProperty::Width => Ok(self.widths().clone().rename(self.name())),
+            ImageProperty::Channel => Ok(self
+                .channels()
+                .clone()
+                .cast(&DataType::UInt32)?
+                .u32()?
+                .clone()
+                .rename(self.name())),
+            ImageProperty::Mode => Ok(self
+                .modes()
+                .clone()
+                .cast(&DataType::UInt32)?
+                .u32()?
+                .clone()
+                .rename(self.name())),
+        }
+    }
 }
 
 impl ImageOps for FixedShapeImageArray {
@@ -135,6 +158,32 @@ impl ImageOps for FixedShapeImageArray {
         };
         fixed_image_array_from_img_buffers(self.name(), &buffers, &mode, *height, *width)
     }
+
+    fn attribute(&self, attr: ImageProperty) -> DaftResult<DataArray<UInt32Type>> {
+        let (height, width) = match self.data_type() {
+            DataType::FixedShapeImage(_, h, w) => (h, w),
+            _ => unreachable!("Should be FixedShapeImage type"),
+        };
+
+        match attr {
+            ImageProperty::Height => Ok(UInt32Array::from((
+                self.name(),
+                vec![*height; self.len()].as_slice(),
+            ))),
+            ImageProperty::Width => Ok(UInt32Array::from((
+                self.name(),
+                vec![*width; self.len()].as_slice(),
+            ))),
+            ImageProperty::Channel => Ok(UInt32Array::from((
+                self.name(),
+                vec![self.image_mode().num_channels() as u32; self.len()].as_slice(),
+            ))),
+            ImageProperty::Mode => Ok(UInt32Array::from((
+                self.name(),
+                vec![(*self.image_mode() as u8) as u32; self.len()].as_slice(),
+            ))),
+        }
+    }
 }
 
 fn encode_images<Arr: AsImageObj>(
diff --git a/src/daft-image/src/series.rs b/src/daft-image/src/series.rs
@@ -1,6 +1,7 @@
 use common_error::{DaftError, DaftResult};
 use common_image::CowImage;
 use daft_core::{array::ops::image::image_array_from_img_buffers, prelude::*};
+use daft_schema::image_property::ImageProperty;
 
 use crate::ops::ImageOps;
 fn image_decode_impl(
@@ -195,3 +196,29 @@ 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", "channel", "mode")
+///
+/// # Returns
+/// Series of UInt32 values containing requested attribute
+pub fn attribute(s: &Series, attr: ImageProperty) -> 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()
+        ))),
+    }
+}
diff --git a/src/daft-ir/src/lib.rs b/src/daft-ir/src/lib.rs
@@ -239,6 +239,7 @@ pub mod schema {
     pub use daft_schema::time_unit::TimeUnit;
     pub use daft_schema::image_format::ImageFormat;
     pub use daft_schema::image_mode::ImageMode;
+    pub use daft_schema::image_property::ImageProperty;
 }
 
 /// Python exports for daft_ir.
diff --git a/src/daft-schema/src/image_property.rs b/src/daft-schema/src/image_property.rs
diff --git a/src/daft-schema/src/lib.rs b/src/daft-schema/src/lib.rs
diff --git a/src/daft-schema/src/prelude.rs b/src/daft-schema/src/prelude.rs
diff --git a/src/daft-schema/src/python/mod.rs b/src/daft-schema/src/python/mod.rs
diff --git a/tests/recordbatch/image/test_attribute.py b/tests/recordbatch/image/test_attribute.py

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,6 @@`
`1`	`1`	`use daft_dsl::functions::FunctionModule;`
`2`	`2`
	`3`	`+pub mod attribute;`
`3`	`4`	`pub mod crop;`
`4`	`5`	`pub mod decode;`
`5`	`6`	`pub mod encode;`
`@@ -15,5 +16,6 @@ impl FunctionModule for ImageFunctions {`
`15`	`16`	`parent.add_fn(encode::ImageEncode);`
`16`	`17`	`parent.add_fn(resize::ImageResize);`
`17`	`18`	`parent.add_fn(to_mode::ImageToMode);`
	`19`	`+ parent.add_fn(attribute::ImageAttribute);`
`18`	`20`	`}`
`19`	`21`	`}`