nd_arrow_array/compute/
cf.rs

1//! CF conventions: mask fill values and unpack `scale_factor` / `add_offset`.
2//!
3//! The behavior follows netCDF4-python and xarray `decode_cf`.
4
5use std::sync::Arc;
6
7use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Scalar};
8use arrow::compute::kernels::{boolean, cmp};
9use arrow::compute::{cast as arrow_cast, nullif};
10use arrow::datatypes::DataType;
11
12use crate::compute::elementwise::{add, cast, div, fill_null, mul, round, sub};
13use crate::error::{Result, nd_err};
14use crate::{Attributes, NdArrowArray};
15
16/// CF attribute: the value that marks missing data.
17pub const FILL_VALUE: &str = "_FillValue";
18/// CF attribute: one or more values that mark missing data.
19pub const MISSING_VALUE: &str = "missing_value";
20/// CF attribute: the smallest valid value.
21pub const VALID_MIN: &str = "valid_min";
22/// CF attribute: the largest valid value.
23pub const VALID_MAX: &str = "valid_max";
24/// CF attribute: the smallest and largest valid values.
25pub const VALID_RANGE: &str = "valid_range";
26/// CF attribute: the multiplier to unpack values.
27pub const SCALE_FACTOR: &str = "scale_factor";
28/// CF attribute: the offset to unpack values.
29pub const ADD_OFFSET: &str = "add_offset";
30
31const CF_KEYS: [&str; 7] = [
32    FILL_VALUE,
33    MISSING_VALUE,
34    VALID_MIN,
35    VALID_MAX,
36    VALID_RANGE,
37    SCALE_FACTOR,
38    ADD_OFFSET,
39];
40
41/// True for the attribute keys that [`decode_cf`] uses.
42pub fn is_cf_key(key: &str) -> bool {
43    CF_KEYS.contains(&key)
44}
45
46/// The CF keys that [`decode_cf`] removed, plus the packed value type.
47/// [`encode_cf`] uses it to restore the packed form.
48#[derive(Debug, Clone, PartialEq)]
49pub struct CfEncoding {
50    /// The CF attributes that the decode removed.
51    pub attributes: Attributes,
52    /// The value type before the decode.
53    pub packed_type: DataType,
54}
55
56/// The values of an attribute: a scalar, or the items of a length-1 list.
57fn attribute_values(value: &ArrayRef) -> ArrayRef {
58    match value.data_type() {
59        DataType::List(_) => value.as_list::<i32>().value(0),
60        DataType::LargeList(_) => value.as_list::<i64>().value(0),
61        DataType::FixedSizeList(_, _) => value.as_fixed_size_list().value(0),
62        _ => value.clone(),
63    }
64}
65
66fn scalar_at(values: &ArrayRef, index: usize, to: &DataType) -> Result<Option<Scalar<ArrayRef>>> {
67    if index >= values.len() || values.is_null(index) {
68        return Ok(None);
69    }
70    Ok(Some(Scalar::new(arrow_cast(
71        values.slice(index, 1).as_ref(),
72        to,
73    )?)))
74}
75
76fn or(mask: Option<BooleanArray>, other: BooleanArray) -> Result<BooleanArray> {
77    Ok(match mask {
78        None => other,
79        Some(mask) => boolean::or(&mask, &other)?,
80    })
81}
82
83/// Mask fill and out-of-range values as nulls, then unpack scale and offset.
84///
85/// - `_FillValue` and every `missing_value` become null.
86/// - Values outside `valid_min`, `valid_max` or `valid_range` become null.
87/// - With `scale_factor` or `add_offset`, the result is `value * scale + offset`
88///   in the type of those attributes (a float type), else in `Float64`.
89///
90/// The CF keys move from the attributes to the returned [`CfEncoding`].
91pub fn decode_cf(a: &NdArrowArray) -> Result<(NdArrowArray, CfEncoding)> {
92    let attrs = a.attributes();
93    let values = a.values();
94    let data_type = a.data_type().clone();
95    let mut invalid: Option<BooleanArray> = None;
96
97    for key in [FILL_VALUE, MISSING_VALUE] {
98        if let Some(value) = attrs.get(key) {
99            let items = attribute_values(value);
100            for i in 0..items.len() {
101                if let Some(fill) = scalar_at(&items, i, &data_type)? {
102                    invalid = Some(or(invalid, cmp::eq(values, &fill)?)?);
103                }
104            }
105        }
106    }
107
108    let range = attrs.get(VALID_RANGE).map(attribute_values);
109    let low = match attrs.get(VALID_MIN) {
110        Some(v) => scalar_at(&attribute_values(v), 0, &data_type)?,
111        None => match &range {
112            Some(r) => scalar_at(r, 0, &data_type)?,
113            None => None,
114        },
115    };
116    let high = match attrs.get(VALID_MAX) {
117        Some(v) => scalar_at(&attribute_values(v), 0, &data_type)?,
118        None => match &range {
119            Some(r) => scalar_at(r, 1, &data_type)?,
120            None => None,
121        },
122    };
123    if let Some(low) = low {
124        invalid = Some(or(invalid, cmp::lt(values, &low)?)?);
125    }
126    if let Some(high) = high {
127        invalid = Some(or(invalid, cmp::gt(values, &high)?)?);
128    }
129
130    let masked: ArrayRef = match &invalid {
131        Some(mask) => nullif(values.as_ref(), mask)?,
132        None => values.clone(),
133    };
134    let mut out = NdArrowArray::try_new(masked, a.dims().clone())?;
135
136    let scale = attrs.get(SCALE_FACTOR).map(attribute_values);
137    let offset = attrs.get(ADD_OFFSET).map(attribute_values);
138    if scale.is_some() || offset.is_some() {
139        let unpacked_type = scale
140            .iter()
141            .chain(offset.iter())
142            .map(|v| v.data_type().clone())
143            .find(|t| matches!(t, DataType::Float32 | DataType::Float64))
144            .unwrap_or(DataType::Float64);
145        out = cast(&out, &unpacked_type)?;
146        if let Some(scale) = scale {
147            out = mul(&out, &scalar_array(&scale, &unpacked_type)?)?;
148        }
149        if let Some(offset) = offset {
150            out = add(&out, &scalar_array(&offset, &unpacked_type)?)?;
151        }
152    }
153
154    let mut kept = Attributes::new();
155    let mut moved = Attributes::new();
156    for (key, value) in attrs.iter() {
157        let target = if CF_KEYS.contains(&key) {
158            &mut moved
159        } else {
160            &mut kept
161        };
162        target.try_insert_array(key, value.clone())?;
163    }
164    Ok((
165        out.with_attributes(kept),
166        CfEncoding {
167            attributes: moved,
168            packed_type: data_type,
169        },
170    ))
171}
172
173fn scalar_array(values: &ArrayRef, to: &DataType) -> Result<NdArrowArray> {
174    if values.is_empty() {
175        return Err(nd_err!(InvalidAttribute, "empty CF scale or offset"));
176    }
177    NdArrowArray::try_new_scalar(arrow_cast(values.slice(0, 1).as_ref(), to)?)
178}
179
180/// Pack decoded values back with `encoding`, the inverse of [`decode_cf`].
181///
182/// Computes `(value - offset) / scale`, rounds half to even for integer
183/// packed types, casts to the packed type, fills nulls with `_FillValue` (or the
184/// first `missing_value`), and restores the CF attributes.
185pub fn encode_cf(a: &NdArrowArray, encoding: &CfEncoding) -> Result<NdArrowArray> {
186    let attrs = &encoding.attributes;
187    let mut out = a.clone();
188    let scale = attrs.get(SCALE_FACTOR).map(attribute_values);
189    let offset = attrs.get(ADD_OFFSET).map(attribute_values);
190    if let Some(offset) = &offset {
191        out = sub(&out, &scalar_array(offset, out.data_type())?)?;
192    }
193    if let Some(scale) = &scale {
194        out = div(&out, &scalar_array(scale, out.data_type())?)?;
195    }
196    let packed_is_integer = encoding.packed_type.is_integer();
197    if (scale.is_some() || offset.is_some()) && packed_is_integer {
198        out = round(&out)?;
199    }
200    out = cast(&out, &encoding.packed_type)?;
201
202    let fill = attrs
203        .get(FILL_VALUE)
204        .or_else(|| attrs.get(MISSING_VALUE))
205        .map(attribute_values)
206        .filter(|v| !v.is_empty());
207    if let Some(fill) = fill {
208        out = fill_null(&out, &Arc::new(fill.slice(0, 1)) as &ArrayRef)?;
209    }
210
211    let mut merged = a.attributes().clone();
212    for (key, value) in attrs.iter() {
213        merged.try_insert_array(key, value.clone())?;
214    }
215    Ok(out.with_attributes(merged))
216}
217
218#[cfg(test)]
219mod tests {
220    use arrow::array::{Float32Array, Int16Array};
221    use arrow::datatypes::{Float32Type, Int16Type};
222
223    use super::*;
224    use crate::Dimensions;
225
226    fn packed() -> NdArrowArray {
227        NdArrowArray::try_new(
228            Arc::new(Int16Array::from(vec![-32767, 0, 100, 200, 5000])),
229            Dimensions::try_new([("x", 5)]).unwrap(),
230        )
231        .unwrap()
232        .with_attributes(
233            Attributes::new()
234                .with(FILL_VALUE, -32767i16)
235                .with(VALID_RANGE, vec![0i16, 1000])
236                .with(SCALE_FACTOR, 0.5f32)
237                .with(ADD_OFFSET, 10.0f32)
238                .with("units", "degC"),
239        )
240    }
241
242    #[test]
243    fn decode_masks_and_unpacks() {
244        let (decoded, encoding) = decode_cf(&packed()).unwrap();
245        assert_eq!(decoded.data_type(), &DataType::Float32);
246        let values: Vec<Option<f32>> = decoded
247            .values()
248            .as_primitive::<Float32Type>()
249            .iter()
250            .collect();
251        assert_eq!(
252            values,
253            vec![None, Some(10.0), Some(60.0), Some(110.0), None]
254        );
255        assert_eq!(decoded.attributes().get_str("units"), Some("degC"));
256        assert!(!decoded.attributes().contains(SCALE_FACTOR));
257        assert_eq!(encoding.attributes.len(), 4);
258        assert_eq!(encoding.packed_type, DataType::Int16);
259    }
260
261    #[test]
262    fn encode_inverts_decode() {
263        let (decoded, encoding) = decode_cf(&packed()).unwrap();
264        let encoded = encode_cf(&decoded, &encoding).unwrap();
265        let values: Vec<i16> = encoded
266            .values()
267            .as_primitive::<Int16Type>()
268            .values()
269            .to_vec();
270        assert_eq!(values, vec![-32767, 0, 100, 200, -32767]);
271        assert!(encoded.attributes().contains(SCALE_FACTOR));
272        assert_eq!(encoded.attributes().get_str("units"), Some("degC"));
273    }
274
275    #[test]
276    fn mask_only_keeps_the_type() {
277        let a = NdArrowArray::try_new(
278            Arc::new(Float32Array::from(vec![1.0, -999.0])),
279            Dimensions::try_new([("x", 2)]).unwrap(),
280        )
281        .unwrap()
282        .with_attributes(Attributes::new().with(MISSING_VALUE, vec![-999.0f32, -998.0]));
283        let (decoded, _) = decode_cf(&a).unwrap();
284        assert_eq!(decoded.data_type(), &DataType::Float32);
285        assert_eq!(decoded.values().null_count(), 1);
286    }
287}
nd_arrow_array/compute/cf.rs

nd_arrow_array/compute/
cf.rs
