This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new f220d72472 feat(arrow-schema): expose FixedShapeTensor logical shape
(#11249)
f220d72472 is described below
commit f220d72472e652e555a2d73d4a9cad48c78d7397
Author: yongster <[email protected]>
AuthorDate: Fri Oct 2 02:05:42 2026 +0800
feat(arrow-schema): expose FixedShapeTensor logical shape (#11249)
# Which issue does this PR close?
- Closes #11248.
# Rationale for this change
`FixedShapeTensor` stores a physical shape and an optional permutation,
but does not expose the logical shape. The extension defines logical
dimension `i` as physical dimension `permutations[i]`. Physical shape
`[100, 200, 500]` with permutation `[2, 0, 1]` is logical shape `[500,
100, 200]`.
This is independent of #11247, which exposes the physical shape. The two
accessors answer different questions.
# What changes are included in this PR?
- `FixedShapeTensor::logical_shape(&self) -> Vec<usize>`
- `FixedShapeTensorMetadata::logical_shape(&self) -> Vec<usize>`
- Without a permutation, the result is the physical shape
- A doc example for `[100, 200, 500]` / `[2, 0, 1]`, plus tests for no
permutation, that permutation, and an identity permutation
Serialization and validation are unchanged. This does not rename the
metadata key `permutations`.
# Are these changes tested?
Yes.
- `cargo test -p arrow-schema --features canonical_extension_types --lib
logical_shape`
- `cargo test -p arrow-schema --doc --features canonical_extension_types
logical_shape`
- `cargo clippy -p arrow-schema --features canonical_extension_types
--all-targets --no-deps -- -D warnings`
# Are there any user-facing changes?
Additive API only, behind the existing `canonical_extension_types`
feature. No breaking change.
AI assistance: the accessor and tests were drafted with AI and checked
against the permutation rule already documented on `FixedShapeTensor`.
Co-authored-by: yang3.xie <[email protected]>
Co-authored-by: Andrew Lamb <[email protected]>
---
.../src/extension/canonical/fixed_shape_tensor.rs | 55 ++++++++++++++++++++++
1 file changed, 55 insertions(+)
diff --git a/arrow-schema/src/extension/canonical/fixed_shape_tensor.rs
b/arrow-schema/src/extension/canonical/fixed_shape_tensor.rs
index bf85d53fe5..8e8f71b922 100644
--- a/arrow-schema/src/extension/canonical/fixed_shape_tensor.rs
+++ b/arrow-schema/src/extension/canonical/fixed_shape_tensor.rs
@@ -146,6 +146,30 @@ impl FixedShapeTensor {
pub fn permutations(&self) -> Option<&[usize]> {
self.metadata.permutations()
}
+
+ /// Returns the logical shape.
+ ///
+ /// Logical dimension `i` has the size of physical dimension
+ /// `permutations[i]`. Without a permutation, this is the physical shape.
+ ///
+ /// Physical shape `[100, 200, 500]` and permutation `[2, 0, 1]` have
+ /// logical shape `[500, 100, 200]`.
+ ///
+ /// ```
+ /// # use arrow_schema::extension::FixedShapeTensor;
+ /// # use arrow_schema::DataType;
+ /// let tensor = FixedShapeTensor::try_new(
+ /// DataType::Float32,
+ /// [100, 200, 500],
+ /// None,
+ /// Some(vec![2, 0, 1]),
+ /// )
+ /// .unwrap();
+ /// assert_eq!(tensor.logical_shape(), vec![500, 100, 200]);
+ /// ```
+ pub fn logical_shape(&self) -> Vec<usize> {
+ self.metadata.logical_shape()
+ }
}
/// Extension type metadata for [`FixedShapeTensor`].
@@ -377,6 +401,16 @@ impl FixedShapeTensorMetadata {
pub fn permutations(&self) -> Option<&[usize]> {
self.permutations.as_ref().map(AsRef::as_ref)
}
+
+ /// Returns the logical shape.
+ ///
+ /// See [`FixedShapeTensor::logical_shape`].
+ pub fn logical_shape(&self) -> Vec<usize> {
+ match &self.permutations {
+ Some(permutation) => permutation.iter().map(|&index|
self.shape[index]).collect(),
+ None => self.shape.clone(),
+ }
+ }
}
impl ExtensionType for FixedShapeTensor {
@@ -494,6 +528,27 @@ mod tests {
Ok(())
}
+ #[test]
+ fn logical_shape_follows_permutation() -> Result<(), ArrowError> {
+ let physical = FixedShapeTensor::try_new(DataType::Float32, [100, 200,
500], None, None)?;
+ assert_eq!(physical.logical_shape(), vec![100, 200, 500]);
+ assert_eq!(physical.metadata().logical_shape(), vec![100, 200, 500]);
+
+ let permuted = FixedShapeTensor::try_new(
+ DataType::Float32,
+ [100, 200, 500],
+ Some(vec!["C".to_owned(), "H".to_owned(), "W".to_owned()]),
+ Some(vec![2, 0, 1]),
+ )?;
+ assert_eq!(permuted.logical_shape(), vec![500, 100, 200]);
+ assert_eq!(permuted.metadata().logical_shape(), vec![500, 100, 200]);
+
+ let identity =
+ FixedShapeTensor::try_new(DataType::Float32, [2, 6], None,
Some(vec![0, 1]))?;
+ assert_eq!(identity.logical_shape(), vec![2, 6]);
+ Ok(())
+ }
+
#[test]
fn shape_distinguishes_equal_sizes() -> Result<(), ArrowError> {
let wide = FixedShapeTensor::try_new(DataType::Float32, [2, 6], None,
None)?;
