arrow/
tensor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Arrow Tensor Type, defined in
6//! [`format/Tensor.fbs`](https://github.com/apache/arrow/blob/master/format/Tensor.fbs).
7
8use std::marker::PhantomData;
9use std::mem;
10
11use crate::buffer::Buffer;
12use crate::datatypes::*;
13
14use crate::error::{ArrowError, Result};
15
16/// Computes the strides required assuming a row major memory layout
17fn compute_row_major_strides<T: ArrowPrimitiveType>(shape: &[usize]) -> Result<Vec<usize>> {
18    let mut remaining_bytes = mem::size_of::<T::Native>();
19
20    for i in shape {
21        if let Some(val) = remaining_bytes.checked_mul(*i) {
22            remaining_bytes = val;
23        } else {
24            return Err(ArrowError::ComputeError(
25                "overflow occurred when computing row major strides.".to_string(),
26            ));
27        }
28    }
29
30    let mut strides = Vec::<usize>::new();
31    for i in shape {
32        remaining_bytes /= *i;
33        strides.push(remaining_bytes);
34    }
35
36    Ok(strides)
37}
38
39/// Computes the strides required assuming a column major memory layout
40fn compute_column_major_strides<T: ArrowPrimitiveType>(shape: &[usize]) -> Result<Vec<usize>> {
41    let mut remaining_bytes = mem::size_of::<T::Native>();
42    let mut strides = Vec::<usize>::new();
43
44    for i in shape {
45        strides.push(remaining_bytes);
46
47        if let Some(val) = remaining_bytes.checked_mul(*i) {
48            remaining_bytes = val;
49        } else {
50            return Err(ArrowError::ComputeError(
51                "overflow occurred when computing column major strides.".to_string(),
52            ));
53        }
54    }
55
56    Ok(strides)
57}
58
59/// Tensor of primitive types
60#[derive(Debug)]
61pub struct Tensor<'a, T: ArrowPrimitiveType> {
62    data_type: DataType,
63    buffer: Buffer,
64    shape: Option<Vec<usize>>,
65    strides: Option<Vec<usize>>,
66    names: Option<Vec<&'a str>>,
67    _marker: PhantomData<T>,
68}
69
70/// [Tensor] of type [BooleanType]
71pub type BooleanTensor<'a> = Tensor<'a, BooleanType>;
72/// [Tensor] of type [Int8Type]
73pub type Date32Tensor<'a> = Tensor<'a, Date32Type>;
74/// [Tensor] of type [Int16Type]
75pub type Date64Tensor<'a> = Tensor<'a, Date64Type>;
76/// [Tensor] of type [Decimal32Type]
77pub type Decimal32Tensor<'a> = Tensor<'a, Decimal32Type>;
78/// [Tensor] of type [Decimal64Type]
79pub type Decimal64Tensor<'a> = Tensor<'a, Decimal64Type>;
80/// [Tensor] of type [Decimal128Type]
81pub type Decimal128Tensor<'a> = Tensor<'a, Decimal128Type>;
82/// [Tensor] of type [Decimal256Type]
83pub type Decimal256Tensor<'a> = Tensor<'a, Decimal256Type>;
84/// [Tensor] of type [DurationMicrosecondType]
85pub type DurationMicrosecondTensor<'a> = Tensor<'a, DurationMicrosecondType>;
86/// [Tensor] of type [DurationMillisecondType]
87pub type DurationMillisecondTensor<'a> = Tensor<'a, DurationMillisecondType>;
88/// [Tensor] of type [DurationNanosecondType]
89pub type DurationNanosecondTensor<'a> = Tensor<'a, DurationNanosecondType>;
90/// [Tensor] of type [DurationSecondType]
91pub type DurationSecondTensor<'a> = Tensor<'a, DurationSecondType>;
92/// [Tensor] of type [Float16Type]
93pub type Float16Tensor<'a> = Tensor<'a, Float16Type>;
94/// [Tensor] of type [Float32Type]
95pub type Float32Tensor<'a> = Tensor<'a, Float32Type>;
96/// [Tensor] of type [Float64Type]
97pub type Float64Tensor<'a> = Tensor<'a, Float64Type>;
98/// [Tensor] of type [Int8Type]
99pub type Int8Tensor<'a> = Tensor<'a, Int8Type>;
100/// [Tensor] of type [Int16Type]
101pub type Int16Tensor<'a> = Tensor<'a, Int16Type>;
102/// [Tensor] of type [Int32Type]
103pub type Int32Tensor<'a> = Tensor<'a, Int32Type>;
104/// [Tensor] of type [Int64Type]
105pub type Int64Tensor<'a> = Tensor<'a, Int64Type>;
106/// [Tensor] of type [IntervalDayTimeType]
107pub type IntervalDayTimeTensor<'a> = Tensor<'a, IntervalDayTimeType>;
108/// [Tensor] of type [IntervalMonthDayNanoType]
109pub type IntervalMonthDayNanoTensor<'a> = Tensor<'a, IntervalMonthDayNanoType>;
110/// [Tensor] of type [IntervalYearMonthType]
111pub type IntervalYearMonthTensor<'a> = Tensor<'a, IntervalYearMonthType>;
112/// [Tensor] of type [Time32MillisecondType]
113pub type Time32MillisecondTensor<'a> = Tensor<'a, Time32MillisecondType>;
114/// [Tensor] of type [Time32SecondType]
115pub type Time32SecondTensor<'a> = Tensor<'a, Time32SecondType>;
116/// [Tensor] of type [Time64MicrosecondType]
117pub type Time64MicrosecondTensor<'a> = Tensor<'a, Time64MicrosecondType>;
118/// [Tensor] of type [Time64NanosecondType]
119pub type Time64NanosecondTensor<'a> = Tensor<'a, Time64NanosecondType>;
120/// [Tensor] of type [TimestampMicrosecondType]
121pub type TimestampMicrosecondTensor<'a> = Tensor<'a, TimestampMicrosecondType>;
122/// [Tensor] of type [TimestampMillisecondType]
123pub type TimestampMillisecondTensor<'a> = Tensor<'a, TimestampMillisecondType>;
124/// [Tensor] of type [TimestampNanosecondType]
125pub type TimestampNanosecondTensor<'a> = Tensor<'a, TimestampNanosecondType>;
126/// [Tensor] of type [TimestampSecondType]
127pub type TimestampSecondTensor<'a> = Tensor<'a, TimestampSecondType>;
128/// [Tensor] of type [UInt8Type]
129pub type UInt8Tensor<'a> = Tensor<'a, UInt8Type>;
130/// [Tensor] of type [UInt16Type]
131pub type UInt16Tensor<'a> = Tensor<'a, UInt16Type>;
132/// [Tensor] of type [UInt32Type]
133pub type UInt32Tensor<'a> = Tensor<'a, UInt32Type>;
134/// [Tensor] of type [UInt64Type]
135pub type UInt64Tensor<'a> = Tensor<'a, UInt64Type>;
136
137impl<'a, T: ArrowPrimitiveType> Tensor<'a, T> {
138    /// Creates a new `Tensor`
139    pub fn try_new(
140        buffer: Buffer,
141        shape: Option<Vec<usize>>,
142        strides: Option<Vec<usize>>,
143        names: Option<Vec<&'a str>>,
144    ) -> Result<Self> {
145        match shape {
146            None => {
147                if buffer.len() != mem::size_of::<T::Native>() {
148                    return Err(ArrowError::InvalidArgumentError(
149                        "underlying buffer should only contain a single tensor element".to_string(),
150                    ));
151                }
152
153                if strides.is_some() {
154                    return Err(ArrowError::InvalidArgumentError(
155                        "expected None strides for tensor with no shape".to_string(),
156                    ));
157                }
158
159                if names.is_some() {
160                    return Err(ArrowError::InvalidArgumentError(
161                        "expected None names for tensor with no shape".to_string(),
162                    ));
163                }
164            }
165
166            Some(ref s) => {
167                if let Some(ref st) = strides
168                    && st.len() != s.len()
169                {
170                    return Err(ArrowError::InvalidArgumentError(
171                        "shape and stride dimensions differ".to_string(),
172                    ));
173                }
174
175                if let Some(ref n) = names
176                    && n.len() != s.len()
177                {
178                    return Err(ArrowError::InvalidArgumentError(
179                        "number of dimensions and number of dimension names differ".to_string(),
180                    ));
181                }
182
183                let total_elements: usize = s.iter().product();
184                if total_elements != (buffer.len() / mem::size_of::<T::Native>()) {
185                    return Err(ArrowError::InvalidArgumentError(
186                        "number of elements in buffer does not match dimensions".to_string(),
187                    ));
188                }
189            }
190        }
191
192        // Checking that the tensor strides used for construction are correct
193        // otherwise a row major stride is calculated and used as value for the tensor
194        let tensor_strides = {
195            if let Some(st) = strides {
196                if let Some(ref s) = shape {
197                    if compute_row_major_strides::<T>(s)? == st
198                        || compute_column_major_strides::<T>(s)? == st
199                    {
200                        Some(st)
201                    } else {
202                        return Err(ArrowError::InvalidArgumentError(
203                            "the input stride does not match the selected shape".to_string(),
204                        ));
205                    }
206                } else {
207                    Some(st)
208                }
209            } else if let Some(ref s) = shape {
210                Some(compute_row_major_strides::<T>(s)?)
211            } else {
212                None
213            }
214        };
215
216        Ok(Self {
217            data_type: T::DATA_TYPE,
218            buffer,
219            shape,
220            strides: tensor_strides,
221            names,
222            _marker: PhantomData,
223        })
224    }
225
226    /// Creates a new Tensor using row major memory layout
227    pub fn new_row_major(
228        buffer: Buffer,
229        shape: Option<Vec<usize>>,
230        names: Option<Vec<&'a str>>,
231    ) -> Result<Self> {
232        if let Some(ref s) = shape {
233            let strides = Some(compute_row_major_strides::<T>(s)?);
234
235            Self::try_new(buffer, shape, strides, names)
236        } else {
237            Err(ArrowError::InvalidArgumentError(
238                "shape required to create row major tensor".to_string(),
239            ))
240        }
241    }
242
243    /// Creates a new Tensor using column major memory layout
244    pub fn new_column_major(
245        buffer: Buffer,
246        shape: Option<Vec<usize>>,
247        names: Option<Vec<&'a str>>,
248    ) -> Result<Self> {
249        if let Some(ref s) = shape {
250            let strides = Some(compute_column_major_strides::<T>(s)?);
251
252            Self::try_new(buffer, shape, strides, names)
253        } else {
254            Err(ArrowError::InvalidArgumentError(
255                "shape required to create column major tensor".to_string(),
256            ))
257        }
258    }
259
260    /// The data type of the `Tensor`
261    pub fn data_type(&self) -> &DataType {
262        &self.data_type
263    }
264
265    /// The sizes of the dimensions
266    pub fn shape(&self) -> Option<&Vec<usize>> {
267        self.shape.as_ref()
268    }
269
270    /// Returns a reference to the underlying `Buffer`
271    pub fn data(&self) -> &Buffer {
272        &self.buffer
273    }
274
275    /// The number of bytes between elements in each dimension
276    pub fn strides(&self) -> Option<&Vec<usize>> {
277        self.strides.as_ref()
278    }
279
280    /// The names of the dimensions
281    pub fn names(&self) -> Option<&Vec<&'a str>> {
282        self.names.as_ref()
283    }
284
285    /// The number of dimensions
286    pub fn ndim(&self) -> usize {
287        match &self.shape {
288            None => 0,
289            Some(v) => v.len(),
290        }
291    }
292
293    /// The name of dimension i
294    pub fn dim_name(&self, i: usize) -> Option<&'a str> {
295        self.names.as_ref().map(|names| names[i])
296    }
297
298    /// The total number of elements in the `Tensor`
299    pub fn size(&self) -> usize {
300        match self.shape {
301            None => 0,
302            Some(ref s) => s.iter().product(),
303        }
304    }
305
306    /// Indicates if the data is laid out contiguously in memory
307    pub fn is_contiguous(&self) -> Result<bool> {
308        Ok(self.is_row_major()? || self.is_column_major()?)
309    }
310
311    /// Indicates if the memory layout row major
312    pub fn is_row_major(&self) -> Result<bool> {
313        match self.shape {
314            None => Ok(false),
315            Some(ref s) => Ok(Some(compute_row_major_strides::<T>(s)?) == self.strides),
316        }
317    }
318
319    /// Indicates if the memory layout column major
320    pub fn is_column_major(&self) -> Result<bool> {
321        match self.shape {
322            None => Ok(false),
323            Some(ref s) => Ok(Some(compute_column_major_strides::<T>(s)?) == self.strides),
324        }
325    }
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331
332    use crate::array::*;
333
334    #[test]
335    fn test_compute_row_major_strides() {
336        assert_eq!(
337            vec![48_usize, 8],
338            compute_row_major_strides::<Int64Type>(&[4_usize, 6]).unwrap()
339        );
340        assert_eq!(
341            vec![24_usize, 4],
342            compute_row_major_strides::<Int32Type>(&[4_usize, 6]).unwrap()
343        );
344        assert_eq!(
345            vec![6_usize, 1],
346            compute_row_major_strides::<Int8Type>(&[4_usize, 6]).unwrap()
347        );
348    }
349
350    #[test]
351    fn test_compute_column_major_strides() {
352        assert_eq!(
353            vec![8_usize, 32],
354            compute_column_major_strides::<Int64Type>(&[4_usize, 6]).unwrap()
355        );
356        assert_eq!(
357            vec![4_usize, 16],
358            compute_column_major_strides::<Int32Type>(&[4_usize, 6]).unwrap()
359        );
360        assert_eq!(
361            vec![1_usize, 4],
362            compute_column_major_strides::<Int8Type>(&[4_usize, 6]).unwrap()
363        );
364    }
365
366    #[test]
367    fn test_zero_dim() {
368        let buf = Buffer::from(&[1]);
369        let tensor = UInt8Tensor::try_new(buf, None, None, None).unwrap();
370        assert_eq!(0, tensor.size());
371        assert_eq!(None, tensor.shape());
372        assert_eq!(None, tensor.names());
373        assert_eq!(0, tensor.ndim());
374        assert!(!tensor.is_row_major().unwrap());
375        assert!(!tensor.is_column_major().unwrap());
376        assert!(!tensor.is_contiguous().unwrap());
377
378        let buf = Buffer::from(&[1, 2, 2, 2]);
379        let tensor = Int32Tensor::try_new(buf, None, None, None).unwrap();
380        assert_eq!(0, tensor.size());
381        assert_eq!(None, tensor.shape());
382        assert_eq!(None, tensor.names());
383        assert_eq!(0, tensor.ndim());
384        assert!(!tensor.is_row_major().unwrap());
385        assert!(!tensor.is_column_major().unwrap());
386        assert!(!tensor.is_contiguous().unwrap());
387    }
388
389    #[test]
390    fn test_tensor() {
391        let mut builder = Int32BufferBuilder::new(16);
392        for i in 0..16 {
393            builder.append(i);
394        }
395        let buf = builder.finish();
396        let tensor = Int32Tensor::try_new(buf, Some(vec![2, 8]), None, None).unwrap();
397        assert_eq!(16, tensor.size());
398        assert_eq!(Some(vec![2_usize, 8]).as_ref(), tensor.shape());
399        assert_eq!(Some(vec![32_usize, 4]).as_ref(), tensor.strides());
400        assert_eq!(2, tensor.ndim());
401        assert_eq!(None, tensor.names());
402    }
403
404    #[test]
405    fn test_new_row_major() {
406        let mut builder = Int32BufferBuilder::new(16);
407        for i in 0..16 {
408            builder.append(i);
409        }
410        let buf = builder.finish();
411        let tensor = Int32Tensor::new_row_major(buf, Some(vec![2, 8]), None).unwrap();
412        assert_eq!(16, tensor.size());
413        assert_eq!(Some(vec![2_usize, 8]).as_ref(), tensor.shape());
414        assert_eq!(Some(vec![32_usize, 4]).as_ref(), tensor.strides());
415        assert_eq!(None, tensor.names());
416        assert_eq!(2, tensor.ndim());
417        assert!(tensor.is_row_major().unwrap());
418        assert!(!tensor.is_column_major().unwrap());
419        assert!(tensor.is_contiguous().unwrap());
420    }
421
422    #[test]
423    fn test_new_column_major() {
424        let mut builder = Int32BufferBuilder::new(16);
425        for i in 0..16 {
426            builder.append(i);
427        }
428        let buf = builder.finish();
429        let tensor = Int32Tensor::new_column_major(buf, Some(vec![2, 8]), None).unwrap();
430        assert_eq!(16, tensor.size());
431        assert_eq!(Some(vec![2_usize, 8]).as_ref(), tensor.shape());
432        assert_eq!(Some(vec![4_usize, 8]).as_ref(), tensor.strides());
433        assert_eq!(None, tensor.names());
434        assert_eq!(2, tensor.ndim());
435        assert!(!tensor.is_row_major().unwrap());
436        assert!(tensor.is_column_major().unwrap());
437        assert!(tensor.is_contiguous().unwrap());
438    }
439
440    #[test]
441    fn test_with_names() {
442        let mut builder = Int64BufferBuilder::new(8);
443        for i in 0..8 {
444            builder.append(i);
445        }
446        let buf = builder.finish();
447        let names = vec!["Dim 1", "Dim 2"];
448        let tensor = Int64Tensor::new_column_major(buf, Some(vec![2, 4]), Some(names)).unwrap();
449        assert_eq!(8, tensor.size());
450        assert_eq!(Some(vec![2_usize, 4]).as_ref(), tensor.shape());
451        assert_eq!(Some(vec![8_usize, 16]).as_ref(), tensor.strides());
452        assert_eq!("Dim 1", tensor.dim_name(0).unwrap());
453        assert_eq!("Dim 2", tensor.dim_name(1).unwrap());
454        assert_eq!(2, tensor.ndim());
455        assert!(!tensor.is_row_major().unwrap());
456        assert!(tensor.is_column_major().unwrap());
457        assert!(tensor.is_contiguous().unwrap());
458    }
459
460    #[test]
461    fn test_inconsistent_strides() {
462        let mut builder = Int32BufferBuilder::new(16);
463        for i in 0..16 {
464            builder.append(i);
465        }
466        let buf = builder.finish();
467
468        let result = Int32Tensor::try_new(buf, Some(vec![2, 8]), Some(vec![2, 8, 1]), None);
469
470        if result.is_ok() {
471            panic!("shape and stride dimensions are different")
472        }
473    }
474
475    #[test]
476    fn test_inconsistent_names() {
477        let mut builder = Int32BufferBuilder::new(16);
478        for i in 0..16 {
479            builder.append(i);
480        }
481        let buf = builder.finish();
482
483        let result = Int32Tensor::try_new(
484            buf,
485            Some(vec![2, 8]),
486            Some(vec![4, 8]),
487            Some(vec!["1", "2", "3"]),
488        );
489
490        if result.is_ok() {
491            panic!("dimensions and names have different shape")
492        }
493    }
494
495    #[test]
496    fn test_incorrect_shape() {
497        let mut builder = Int32BufferBuilder::new(16);
498        for i in 0..16 {
499            builder.append(i);
500        }
501        let buf = builder.finish();
502
503        let result = Int32Tensor::try_new(buf, Some(vec![2, 6]), None, None);
504
505        if result.is_ok() {
506            panic!("number of elements does not match for the shape")
507        }
508    }
509
510    #[test]
511    fn test_incorrect_stride() {
512        let mut builder = Int32BufferBuilder::new(16);
513        for i in 0..16 {
514            builder.append(i);
515        }
516        let buf = builder.finish();
517
518        let result = Int32Tensor::try_new(buf, Some(vec![2, 8]), Some(vec![30, 4]), None);
519
520        if result.is_ok() {
521            panic!("the input stride does not match the selected shape")
522        }
523    }
524}
arrow/tensor.rs

arrow/
tensor.rs
