nd_arrow_array/
extension.rs

1//! The `nd_arrow.array` Arrow extension type.
2//!
3//! An nd column is an Arrow `Struct` with this storage layout. Each row holds
4//! one complete N-dimensional array:
5//!
6//! ```text
7//! Struct {
8//!   values:    List<T>,        // flat C-order values
9//!   dim_names: List<Utf8>,     // one name per axis
10//!   dim_sizes: List<UInt32>,   // one size per axis
11//!   attrs:     Struct { .. },  // optional: one nullable child per attribute key
12//! }
13//! ```
14//!
15//! The field carries `ARROW:extension:name = nd_arrow.array` and
16//! `ARROW:extension:metadata = {"version":1}`, so the intent survives IPC,
17//! Parquet and any engine that keeps field metadata.
18//!
19//! # Rules
20//!
21//! `FORMAT.md` in the repository is the full, language-neutral specification.
22//! This section is a summary.
23//!
24//! - The field order is fixed: `values`, `dim_names`, `dim_sizes`, then `attrs`
25//!   if present.
26//! - A column without attributes leaves out the `attrs` field, so no storage
27//!   format needs a struct with zero fields. Readers also accept an `attrs`
28//!   struct with zero fields and treat it as no attributes.
29//! - `dim_names` and `dim_sizes` have the same length. This length is the rank.
30//! - The length of `values` is the product of `dim_sizes`.
31//! - Rank 0 (empty dimension lists) is a scalar with one value.
32//! - Dimension names are not empty and are unique in a row.
33//! - Validity of elements is in the null buffer of `values`. There is no fill value.
34//! - An attribute can be any Arrow type. The `attrs` struct holds the union of
35//!   the keys of all rows. A null child means that the row does not have that key.
36//! - A null row decodes to a rank-0 null scalar.
37//!
38//! # Example
39//!
40//! A `(time=2, lat=3)` array of `Float64` with `units = "K"`:
41//!
42//! ```text
43//! values:    [0.0, 0.1, 0.2, 1.0, 1.1, 1.2]
44//! dim_names: ["time", "lat"]
45//! dim_sizes: [2, 3]
46//! attrs:     { units: "K" }
47//! ```
48//!
49//! # Global attributes
50//!
51//! An encoded record batch (see [`encode_batches`](crate::encode_batches))
52//! stores the global attributes of each nd batch in one extra column,
53//! [`GLOBAL_ATTRIBUTES_FIELD`], tagged with [`GLOBAL_ATTRIBUTES_EXTENSION`].
54//! It has one row per nd batch and the union of keys, with the same null rule
55//! as `attrs`. The column is left out when no batch has global attributes.
56
57use std::collections::HashMap;
58use std::sync::Arc;
59
60use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema};
61use arrow::error::ArrowError;
62use arrow_schema::extension::{
63    EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
64};
65use serde::{Deserialize, Serialize};
66
67use crate::error::{Result, nd_err};
68
69/// Storage field with the flat values.
70pub const VALUES_FIELD: &str = "values";
71/// Storage field with the dimension names.
72pub const DIM_NAMES_FIELD: &str = "dim_names";
73/// Storage field with the dimension sizes.
74pub const DIM_SIZES_FIELD: &str = "dim_sizes";
75/// Optional storage field with the attributes.
76pub const ATTRS_FIELD: &str = "attrs";
77
78/// Name of the list item field in `values`, `dim_names` and `dim_sizes`.
79pub const ITEM_FIELD: &str = "item";
80
81/// The encoding version this crate writes and reads.
82pub const ENCODING_VERSION: u32 = 1;
83
84/// Extension metadata, serialized as JSON in `ARROW:extension:metadata`.
85#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
86pub struct NdArrayMetadata {
87    /// The encoding version.
88    pub version: u32,
89}
90
91impl Default for NdArrayMetadata {
92    fn default() -> Self {
93        Self {
94            version: ENCODING_VERSION,
95        }
96    }
97}
98
99/// The nd extension type: the element field of the values, plus the attribute fields.
100#[derive(Debug, Clone, PartialEq, Eq)]
101pub struct NdArrayType {
102    value_field: FieldRef,
103    attrs: Fields,
104    metadata: NdArrayMetadata,
105}
106
107impl NdArrayType {
108    /// An nd type with nullable values of `value_type` and the given attribute fields.
109    pub fn new(value_type: DataType, attrs: Fields) -> Self {
110        Self::with_value_field(Arc::new(Field::new(ITEM_FIELD, value_type, true)), attrs)
111    }
112
113    /// An nd type with a custom item field for the values list,
114    /// for example to keep nested extension metadata on the elements.
115    pub fn with_value_field(value_field: FieldRef, attrs: Fields) -> Self {
116        let attrs = attrs
117            .iter()
118            .map(|f| {
119                if f.is_nullable() {
120                    f.clone()
121                } else {
122                    Arc::new(f.as_ref().clone().with_nullable(true))
123                }
124            })
125            .collect();
126        Self {
127            value_field,
128            attrs,
129            metadata: NdArrayMetadata::default(),
130        }
131    }
132
133    /// The item field of the `values` list.
134    pub fn value_field(&self) -> &FieldRef {
135        &self.value_field
136    }
137
138    /// The Arrow type of the array values.
139    pub fn value_type(&self) -> &DataType {
140        self.value_field.data_type()
141    }
142
143    /// One field per attribute key. Empty when there are no attributes.
144    pub fn attribute_fields(&self) -> &Fields {
145        &self.attrs
146    }
147
148    /// The storage type of this extension type.
149    pub fn storage_type(&self) -> DataType {
150        DataType::Struct(self.storage_fields())
151    }
152
153    /// True when the storage has an `attrs` field.
154    pub fn has_attributes(&self) -> bool {
155        !self.attrs.is_empty()
156    }
157
158    pub(crate) fn storage_fields(&self) -> Fields {
159        let mut fields = vec![
160            Field::new(
161                VALUES_FIELD,
162                DataType::List(self.value_field.clone()),
163                false,
164            ),
165            Field::new(DIM_NAMES_FIELD, dim_names_type(), false),
166            Field::new(DIM_SIZES_FIELD, dim_sizes_type(), false),
167        ];
168        if self.has_attributes() {
169            fields.push(Field::new(
170                ATTRS_FIELD,
171                DataType::Struct(self.attrs.clone()),
172                false,
173            ));
174        }
175        fields.into()
176    }
177
178    /// Parse the storage layout. This ignores field metadata.
179    /// An `attrs` struct with zero fields reads as no attributes.
180    pub fn try_from_storage_type(data_type: &DataType) -> Result<Self> {
181        let DataType::Struct(fields) = data_type else {
182            return Err(nd_err!(
183                InvalidEncoding,
184                "storage type must be a Struct, got {data_type}"
185            ));
186        };
187        let names: Vec<&str> = fields.iter().map(|f| f.name().as_str()).collect();
188        let expected = [VALUES_FIELD, DIM_NAMES_FIELD, DIM_SIZES_FIELD, ATTRS_FIELD];
189        if !(names == expected[..3] || names == expected) {
190            return Err(nd_err!(
191                InvalidEncoding,
192                "storage fields must be {expected:?}, got {names:?}"
193            ));
194        }
195
196        let DataType::List(value_field) = fields[0].data_type() else {
197            return Err(nd_err!(
198                InvalidEncoding,
199                "'{VALUES_FIELD}' must be a List, got {}",
200                fields[0].data_type()
201            ));
202        };
203        expect_list_of(&fields[1], &DataType::Utf8)?;
204        expect_list_of(&fields[2], &DataType::UInt32)?;
205        let attrs = match fields.get(3).map(|f| f.data_type()) {
206            None => Fields::empty(),
207            Some(DataType::Struct(attrs)) => attrs.clone(),
208            Some(other) => {
209                return Err(nd_err!(
210                    InvalidEncoding,
211                    "'{ATTRS_FIELD}' must be a Struct, got {other}"
212                ));
213            }
214        };
215
216        Ok(Self::with_value_field(value_field.clone(), attrs))
217    }
218
219    /// A field named `name` with this storage type and the extension tag.
220    pub fn to_field(&self, name: impl Into<String>, nullable: bool) -> Field {
221        Field::new(name, self.storage_type(), nullable).with_extension_type(self.clone())
222    }
223}
224
225fn dim_names_type() -> DataType {
226    DataType::List(Arc::new(Field::new(ITEM_FIELD, DataType::Utf8, false)))
227}
228
229fn dim_sizes_type() -> DataType {
230    DataType::List(Arc::new(Field::new(ITEM_FIELD, DataType::UInt32, false)))
231}
232
233fn expect_list_of(field: &Field, item_type: &DataType) -> Result<()> {
234    match field.data_type() {
235        DataType::List(item) if item.data_type() == item_type => Ok(()),
236        other => Err(nd_err!(
237            InvalidEncoding,
238            "'{}' must be List<{item_type}>, got {other}",
239            field.name()
240        )),
241    }
242}
243
244impl ExtensionType for NdArrayType {
245    const NAME: &'static str = "nd_arrow.array";
246
247    type Metadata = NdArrayMetadata;
248
249    fn metadata(&self) -> &Self::Metadata {
250        &self.metadata
251    }
252
253    fn serialize_metadata(&self) -> Option<String> {
254        Some(serde_json::to_string(&self.metadata).expect("metadata serializes"))
255    }
256
257    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
258        let metadata = match metadata.map(str::trim) {
259            None | Some("") | Some("{}") => NdArrayMetadata::default(),
260            Some(json) => serde_json::from_str(json).map_err(|e| {
261                ArrowError::InvalidArgumentError(format!("invalid nd extension metadata: {e}"))
262            })?,
263        };
264        if metadata.version != ENCODING_VERSION {
265            return Err(ArrowError::InvalidArgumentError(format!(
266                "unsupported nd encoding version {}, expected {ENCODING_VERSION}",
267                metadata.version
268            )));
269        }
270        Ok(metadata)
271    }
272
273    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
274        let parsed = Self::try_from_storage_type(data_type)?;
275        if parsed.value_field != self.value_field || parsed.attrs != self.attrs {
276            return Err(ArrowError::InvalidArgumentError(format!(
277                "nd extension type does not match storage type {data_type}"
278            )));
279        }
280        Ok(())
281    }
282
283    fn try_new(data_type: &DataType, metadata: Self::Metadata) -> Result<Self, ArrowError> {
284        let mut parsed = Self::try_from_storage_type(data_type)?;
285        parsed.metadata = metadata;
286        Ok(parsed)
287    }
288}
289
290/// Extension name of the column that holds the global attributes of each nd batch.
291pub const GLOBAL_ATTRIBUTES_EXTENSION: &str = "nd_arrow.global_attributes";
292
293/// Name of the global attributes column in an encoded record batch.
294pub const GLOBAL_ATTRIBUTES_FIELD: &str = "__global_attrs__";
295
296/// The global attributes column: a `Struct` with one nullable child per key,
297/// tagged with [`GLOBAL_ATTRIBUTES_EXTENSION`].
298pub fn global_attributes_field(fields: Fields) -> Field {
299    Field::new(GLOBAL_ATTRIBUTES_FIELD, DataType::Struct(fields), false).with_metadata(
300        HashMap::from([
301            (
302                EXTENSION_TYPE_NAME_KEY.to_string(),
303                GLOBAL_ATTRIBUTES_EXTENSION.to_string(),
304            ),
305            (
306                EXTENSION_TYPE_METADATA_KEY.to_string(),
307                serde_json::to_string(&NdArrayMetadata::default()).expect("metadata serializes"),
308            ),
309        ]),
310    )
311}
312
313/// True when `field` is a global attributes column.
314pub fn is_global_attributes_field(field: &Field) -> bool {
315    field.extension_type_name() == Some(GLOBAL_ATTRIBUTES_EXTENSION)
316        && matches!(field.data_type(), DataType::Struct(_))
317}
318
319/// True when `field` carries the nd extension tag.
320pub fn is_nd_field(field: &Field) -> bool {
321    field.extension_type_name() == Some(NdArrayType::NAME)
322}
323
324/// The nd field for a logical field: same name and metadata, values of the
325/// logical type, no attributes. The field is nullable.
326pub fn encoded_field(logical: &Field) -> Field {
327    let field = NdArrayType::new(logical.data_type().clone(), Fields::empty())
328        .to_field(logical.name(), true);
329    with_base_metadata(field, logical.metadata())
330}
331
332/// The logical field of an nd field: the element type of the values,
333/// with the field metadata minus the extension keys.
334pub fn logical_field(encoded: &Field) -> Result<Field> {
335    let ext: NdArrayType = encoded.try_extension_type()?;
336    let mut metadata = encoded.metadata().clone();
337    metadata.remove(EXTENSION_TYPE_NAME_KEY);
338    metadata.remove(EXTENSION_TYPE_METADATA_KEY);
339    Ok(Field::new(
340        encoded.name(),
341        ext.value_type().clone(),
342        encoded.is_nullable(),
343    )
344    .with_metadata(metadata))
345}
346
347/// Every field of `logical` as an nd field, see [`encoded_field`].
348pub fn encoded_schema(logical: &Schema) -> Schema {
349    let fields: Vec<Field> = logical.fields().iter().map(|f| encoded_field(f)).collect();
350    Schema::new_with_metadata(fields, logical.metadata().clone())
351}
352
353/// Every field of `encoded` as its logical field, see [`logical_field`].
354pub fn logical_schema(encoded: &Schema) -> Result<Schema> {
355    let fields = encoded
356        .fields()
357        .iter()
358        .map(|f| logical_field(f))
359        .collect::<Result<Vec<_>>>()?;
360    Ok(Schema::new_with_metadata(
361        fields,
362        encoded.metadata().clone(),
363    ))
364}
365
366/// `field` with `base` metadata added under its own (extension) keys.
367pub(crate) fn with_base_metadata(field: Field, base: &HashMap<String, String>) -> Field {
368    let mut metadata = base.clone();
369    metadata.extend(field.metadata().clone());
370    field.with_metadata(metadata)
371}
372
373#[cfg(test)]
374mod tests {
375    use super::*;
376
377    fn sample() -> NdArrayType {
378        NdArrayType::new(
379            DataType::Float64,
380            Fields::from(vec![Field::new("units", DataType::Utf8, true)]),
381        )
382    }
383
384    #[test]
385    fn storage_type_round_trips() {
386        let ext = sample();
387        let parsed = NdArrayType::try_from_storage_type(&ext.storage_type()).unwrap();
388        assert_eq!(parsed, ext);
389    }
390
391    #[test]
392    fn field_is_tagged_and_parses_back() {
393        let field = sample().to_field("sst", true);
394        assert!(is_nd_field(&field));
395        let ext: NdArrayType = field.try_extension_type().unwrap();
396        assert_eq!(ext, sample());
397        assert_eq!(
398            field.metadata().get(EXTENSION_TYPE_METADATA_KEY).unwrap(),
399            r#"{"version":1}"#
400        );
401    }
402
403    #[test]
404    fn logical_and_encoded_fields_round_trip() {
405        let logical = Field::new("sst", DataType::Float32, true)
406            .with_metadata(HashMap::from([("k".to_string(), "v".to_string())]));
407        let encoded = encoded_field(&logical);
408        assert!(is_nd_field(&encoded));
409        assert_eq!(encoded.metadata().get("k").unwrap(), "v");
410        assert_eq!(logical_field(&encoded).unwrap(), logical);
411    }
412
413    #[test]
414    fn invalid_storage_rejected() {
415        assert!(NdArrayType::try_from_storage_type(&DataType::Float64).is_err());
416        let wrong_order = DataType::Struct(Fields::from(vec![
417            Field::new(DIM_NAMES_FIELD, dim_names_type(), false),
418            Field::new(
419                VALUES_FIELD,
420                DataType::new_list(DataType::Int32, true),
421                false,
422            ),
423            Field::new(DIM_SIZES_FIELD, dim_sizes_type(), false),
424            Field::new(ATTRS_FIELD, DataType::Struct(Fields::empty()), false),
425        ]));
426        assert!(NdArrayType::try_from_storage_type(&wrong_order).is_err());
427    }
428
429    #[test]
430    fn unknown_version_rejected() {
431        assert!(NdArrayType::deserialize_metadata(Some(r#"{"version":2}"#)).is_err());
432        assert!(NdArrayType::deserialize_metadata(None).is_ok());
433    }
434}
nd_arrow_array/extension.rs

nd_arrow_array/
extension.rs
