questdb/ingress/
ndarr.rs

1/*******************************************************************************
2 *     ___                  _   ____  ____
3 *    / _ \ _   _  ___  ___| |_|  _ \| __ )
4 *   | | | | | | |/ _ \/ __| __| | | |  _ \
5 *   | |_| | |_| |  __/\__ \ |_| |_| | |_) |
6 *    \__\_\\__,_|\___||___/\__|____/|____/
7 *
8 *  Copyright (c) 2014-2019 Appsicle
9 *  Copyright (c) 2019-2025 QuestDB
10 *
11 *  Licensed under the Apache License, Version 2.0 (the "License");
12 *  you may not use this file except in compliance with the License.
13 *  You may obtain a copy of the License at
14 *
15 *  http://www.apache.org/licenses/LICENSE-2.0
16 *
17 *  Unless required by applicable law or agreed to in writing, software
18 *  distributed under the License is distributed on an "AS IS" BASIS,
19 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20 *  See the License for the specific language governing permissions and
21 *  limitations under the License.
22 *
23 ******************************************************************************/
24
25pub trait NdArrayView<T>
26where
27    T: ArrayElement,
28{
29    type Iter<'a>: Iterator<Item = &'a T>
30    where
31        Self: 'a,
32        T: 'a;
33
34    /// Returns the number of dimensions (rank) of the array.
35    fn ndim(&self) -> usize;
36
37    /// Returns the size of the specified dimension.
38    fn dim(&self, index: usize) -> Result<usize, Error>;
39
40    /// Return the array’s data as a slice, if it is c-major-layout.
41    /// Return `None` otherwise.
42    fn as_slice(&self) -> Option<&[T]>;
43
44    /// Return an iterator of references to the elements of the array.
45    /// Iterator element type is `&T`.
46    fn iter(&self) -> Self::Iter<'_>;
47}
48
49pub(crate) fn write_array_data<A: NdArrayView<T>, T>(
50    array: &A,
51    buf: &mut [u8],
52    expect_size: usize,
53) -> Result<(), Error>
54where
55    T: ArrayElement,
56{
57    // When working with contiguous layout, benchmark shows `copy_from_slice` has better performance than
58    // `std::ptr::copy_nonoverlapping` on both Arm(Macos) and x86(Linux) platform.
59    // This may because `copy_from_slice` benefits more from compiler.
60    if let Some(contiguous) = array.as_slice() {
61        let bytes = unsafe {
62            slice::from_raw_parts(contiguous.as_ptr() as *const u8, size_of_val(contiguous))
63        };
64
65        if bytes.len() != expect_size {
66            return Err(error::fmt!(
67                ArrayError,
68                "Array write buffer length mismatch (actual: {}, expected: {})",
69                expect_size,
70                bytes.len()
71            ));
72        }
73
74        if buf.len() < bytes.len() {
75            return Err(error::fmt!(
76                ArrayError,
77                "Buffer capacity {} < required {}",
78                buf.len(),
79                bytes.len()
80            ));
81        }
82
83        buf[..bytes.len()].copy_from_slice(bytes);
84        return Ok(());
85    }
86
87    // For non-contiguous memory layouts, direct raw pointer operations are preferred.
88    if buf.len() < expect_size {
89        return Err(error::fmt!(
90            ArrayError,
91            "Buffer capacity {} < required {}",
92            buf.len(),
93            expect_size
94        ));
95    }
96
97    let elem_size = size_of::<T>();
98    let mut total_len = 0usize;
99    for &element in array.iter() {
100        let end = total_len
101            .checked_add(elem_size)
102            .ok_or_else(|| error::fmt!(ArrayError, "Array write buffer length overflow"))?;
103        if end > expect_size {
104            return Err(error::fmt!(
105                ArrayError,
106                "Array write buffer length mismatch (actual: {}, expected: {})",
107                end,
108                expect_size
109            ));
110        }
111        unsafe {
112            std::ptr::copy_nonoverlapping(
113                &element as *const T as *const u8,
114                buf.as_mut_ptr().add(total_len),
115                elem_size,
116            )
117        }
118        total_len = end;
119    }
120    if total_len != expect_size {
121        return Err(error::fmt!(
122            ArrayError,
123            "Array write buffer length mismatch (actual: {}, expected: {})",
124            total_len,
125            expect_size
126        ));
127    }
128    Ok(())
129}
130
131pub(crate) fn check_and_get_array_bytes_size<A: NdArrayView<T>, T>(
132    array: &A,
133) -> Result<usize, Error>
134where
135    T: ArrayElement,
136{
137    let mut size = std::mem::size_of::<T>();
138    for dim_index in 0..array.ndim() {
139        let dim = array.dim(dim_index)?;
140        if dim > MAX_ARRAY_DIM_LEN {
141            return Err(error::fmt!(
142                ArrayError,
143                "dimension length out of range: dim {}, dim length {}, max length {}",
144                dim_index,
145                dim,
146                MAX_ARRAY_DIM_LEN
147            ));
148        }
149        // following dimension's length may be zero, so check the size in out of loop
150        size = size.checked_mul(dim).ok_or_else(|| {
151            error::fmt!(
152                ArrayError,
153                "Array buffer size overflow: maximum: {}",
154                MAX_ARRAY_BUFFER_SIZE
155            )
156        })?;
157    }
158
159    if size > MAX_ARRAY_BUFFER_SIZE {
160        return Err(error::fmt!(
161            ArrayError,
162            "Array buffer size too big: {}, maximum: {}",
163            size,
164            MAX_ARRAY_BUFFER_SIZE
165        ));
166    }
167    Ok(size)
168}
169
170/// Marker trait for valid array element types.
171///
172/// Implemented for primitive types that can be stored in arrays.
173/// Combines type information with data type classification.
174pub trait ArrayElement: Copy + 'static {}
175
176pub(crate) trait ArrayElementSealed {
177    /// Returns the binary format identifier for array element types compatible
178    /// with QuestDB's io.questdb.cairo.ColumnType numeric type constants.
179    fn type_tag() -> u8;
180}
181
182impl ArrayElement for f64 {}
183
184impl ArrayElementSealed for f64 {
185    fn type_tag() -> u8 {
186        10 // Double
187    }
188}
189
190impl ArrayElement for i64 {}
191
192impl ArrayElementSealed for i64 {
193    fn type_tag() -> u8 {
194        6 // Long
195    }
196}
197
198/// impl NdArrayView for one dimension vector
199impl<T: ArrayElement> NdArrayView<T> for Vec<T> {
200    type Iter<'a>
201        = std::slice::Iter<'a, T>
202    where
203        T: 'a;
204
205    fn ndim(&self) -> usize {
206        1
207    }
208
209    fn dim(&self, idx: usize) -> Result<usize, Error> {
210        if idx == 0 {
211            Ok(self.len())
212        } else {
213            Err(error::fmt!(
214                ArrayError,
215                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
216                idx,
217                1
218            ))
219        }
220    }
221
222    fn as_slice(&self) -> Option<&[T]> {
223        Some(self.as_slice())
224    }
225
226    fn iter(&self) -> Self::Iter<'_> {
227        self.as_slice().iter()
228    }
229}
230
231/// impl NdArrayView for one dimension array
232impl<T: ArrayElement, const N: usize> NdArrayView<T> for [T; N] {
233    type Iter<'a>
234        = std::slice::Iter<'a, T>
235    where
236        T: 'a;
237
238    fn ndim(&self) -> usize {
239        1
240    }
241
242    fn dim(&self, idx: usize) -> Result<usize, Error> {
243        if idx == 0 {
244            Ok(N)
245        } else {
246            Err(error::fmt!(
247                ArrayError,
248                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
249                idx,
250                1
251            ))
252        }
253    }
254
255    fn as_slice(&self) -> Option<&[T]> {
256        Some(self)
257    }
258
259    fn iter(&self) -> Self::Iter<'_> {
260        self.as_slice().iter()
261    }
262}
263
264/// impl NdArrayView for one dimension slice
265impl<T: ArrayElement> NdArrayView<T> for &[T] {
266    type Iter<'a>
267        = std::slice::Iter<'a, T>
268    where
269        Self: 'a,
270        T: 'a;
271
272    fn ndim(&self) -> usize {
273        1
274    }
275
276    fn dim(&self, idx: usize) -> Result<usize, Error> {
277        if idx == 0 {
278            Ok(self.len())
279        } else {
280            Err(error::fmt!(
281                ArrayError,
282                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
283                idx,
284                1
285            ))
286        }
287    }
288
289    fn as_slice(&self) -> Option<&[T]> {
290        Some(self)
291    }
292
293    fn iter(&self) -> Self::Iter<'_> {
294        <[T]>::iter(self)
295    }
296}
297
298/// impl NdArrayView for two dimensions vector
299impl<T: ArrayElement> NdArrayView<T> for Vec<Vec<T>> {
300    type Iter<'a>
301        = std::iter::Flatten<std::slice::Iter<'a, Vec<T>>>
302    where
303        T: 'a;
304
305    fn ndim(&self) -> usize {
306        2
307    }
308
309    fn dim(&self, idx: usize) -> Result<usize, Error> {
310        match idx {
311            0 => Ok(self.len()),
312            1 => {
313                let dim1 = self.first().map_or(0, |v| v.len());
314                if self.as_slice().iter().any(|v2| v2.len() != dim1) {
315                    return Err(error::fmt!(ArrayError, "Irregular array shape"));
316                }
317                Ok(dim1)
318            }
319            _ => Err(error::fmt!(
320                ArrayError,
321                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
322                idx,
323                2
324            )),
325        }
326    }
327
328    fn as_slice(&self) -> Option<&[T]> {
329        None
330    }
331
332    fn iter(&self) -> Self::Iter<'_> {
333        self.as_slice().iter().flatten()
334    }
335}
336
337/// impl NdArrayView for two dimensions array
338impl<T: ArrayElement, const M: usize, const N: usize> NdArrayView<T> for [[T; M]; N] {
339    type Iter<'a>
340        = std::iter::Flatten<std::slice::Iter<'a, [T; M]>>
341    where
342        T: 'a;
343
344    fn ndim(&self) -> usize {
345        2
346    }
347
348    fn dim(&self, idx: usize) -> Result<usize, Error> {
349        match idx {
350            0 => Ok(N),
351            1 => Ok(M),
352            _ => Err(error::fmt!(
353                ArrayError,
354                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
355                idx,
356                2
357            )),
358        }
359    }
360
361    fn as_slice(&self) -> Option<&[T]> {
362        Some(unsafe { std::slice::from_raw_parts(self.as_ptr() as *const T, N * M) })
363    }
364
365    fn iter(&self) -> Self::Iter<'_> {
366        self.as_slice().iter().flatten()
367    }
368}
369
370/// impl NdArrayView for two dimensions slices
371impl<T: ArrayElement, const M: usize> NdArrayView<T> for &[[T; M]] {
372    type Iter<'a>
373        = std::iter::Flatten<std::slice::Iter<'a, [T; M]>>
374    where
375        Self: 'a,
376        T: 'a;
377
378    fn ndim(&self) -> usize {
379        2
380    }
381
382    fn dim(&self, idx: usize) -> Result<usize, Error> {
383        match idx {
384            0 => Ok(self.len()),
385            1 => Ok(M),
386            _ => Err(error::fmt!(
387                ArrayError,
388                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
389                idx,
390                2
391            )),
392        }
393    }
394
395    fn as_slice(&self) -> Option<&[T]> {
396        Some(unsafe { std::slice::from_raw_parts(self.as_ptr() as *const T, self.len() * M) })
397    }
398
399    fn iter(&self) -> Self::Iter<'_> {
400        <[[T; M]]>::iter(self).flatten()
401    }
402}
403
404/// impl NdArrayView for three dimensions vector
405impl<T: ArrayElement> NdArrayView<T> for Vec<Vec<Vec<T>>> {
406    type Iter<'a>
407        = std::iter::Flatten<std::iter::Flatten<std::slice::Iter<'a, Vec<Vec<T>>>>>
408    where
409        T: 'a;
410
411    fn ndim(&self) -> usize {
412        3
413    }
414
415    fn dim(&self, idx: usize) -> Result<usize, Error> {
416        match idx {
417            0 => Ok(self.len()),
418            1 => {
419                let dim1 = self.first().map_or(0, |v| v.len());
420                if self.as_slice().iter().any(|v2| v2.len() != dim1) {
421                    return Err(error::fmt!(ArrayError, "Irregular array shape"));
422                }
423                Ok(dim1)
424            }
425            2 => {
426                let dim2 = self
427                    .first()
428                    .and_then(|v2| v2.first())
429                    .map_or(0, |v3| v3.len());
430
431                if self
432                    .as_slice()
433                    .iter()
434                    .flat_map(|v2| v2.as_slice().iter())
435                    .any(|v3| v3.len() != dim2)
436                {
437                    return Err(error::fmt!(ArrayError, "Irregular array shape"));
438                }
439                Ok(dim2)
440            }
441            _ => Err(error::fmt!(
442                ArrayError,
443                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
444                idx,
445                3
446            )),
447        }
448    }
449
450    fn as_slice(&self) -> Option<&[T]> {
451        None
452    }
453
454    fn iter(&self) -> Self::Iter<'_> {
455        self.as_slice().iter().flatten().flatten()
456    }
457}
458
459/// impl NdArrayView for three dimensions array
460impl<T: ArrayElement, const M: usize, const N: usize, const L: usize> NdArrayView<T>
461    for [[[T; M]; N]; L]
462{
463    type Iter<'a>
464        = std::iter::Flatten<std::iter::Flatten<std::slice::Iter<'a, [[T; M]; N]>>>
465    where
466        T: 'a;
467
468    fn ndim(&self) -> usize {
469        3
470    }
471
472    fn dim(&self, idx: usize) -> Result<usize, Error> {
473        match idx {
474            0 => Ok(L),
475            1 => Ok(N),
476            2 => Ok(M),
477            _ => Err(error::fmt!(
478                ArrayError,
479                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
480                idx,
481                3
482            )),
483        }
484    }
485
486    fn as_slice(&self) -> Option<&[T]> {
487        Some(unsafe { std::slice::from_raw_parts(self.as_ptr() as *const T, L * N * M) })
488    }
489
490    fn iter(&self) -> Self::Iter<'_> {
491        self.as_slice().iter().flatten().flatten()
492    }
493}
494
495impl<T: ArrayElement, const M: usize, const N: usize> NdArrayView<T> for &[[[T; M]; N]] {
496    type Iter<'a>
497        = std::iter::Flatten<std::iter::Flatten<std::slice::Iter<'a, [[T; M]; N]>>>
498    where
499        Self: 'a,
500        T: 'a;
501
502    fn ndim(&self) -> usize {
503        3
504    }
505
506    fn dim(&self, idx: usize) -> Result<usize, Error> {
507        match idx {
508            0 => Ok(self.len()),
509            1 => Ok(N),
510            2 => Ok(M),
511            _ => Err(error::fmt!(
512                ArrayError,
513                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
514                idx,
515                3
516            )),
517        }
518    }
519
520    fn as_slice(&self) -> Option<&[T]> {
521        Some(unsafe { std::slice::from_raw_parts(self.as_ptr() as *const T, self.len() * N * M) })
522    }
523
524    fn iter(&self) -> Self::Iter<'_> {
525        <[[[T; M]; N]]>::iter(self).flatten().flatten()
526    }
527}
528
529use crate::{Error, error};
530#[cfg(feature = "ndarray")]
531use ndarray::{ArrayView, Axis, Dimension};
532use std::slice;
533
534use super::{MAX_ARRAY_BUFFER_SIZE, MAX_ARRAY_DIM_LEN};
535
536#[cfg(feature = "ndarray")]
537impl<T, D> NdArrayView<T> for ArrayView<'_, T, D>
538where
539    T: ArrayElement,
540    D: Dimension,
541{
542    type Iter<'a>
543        = ndarray::iter::Iter<'a, T, D>
544    where
545        Self: 'a,
546        T: 'a;
547
548    fn ndim(&self) -> usize {
549        self.ndim()
550    }
551
552    fn dim(&self, index: usize) -> Result<usize, Error> {
553        let len = self.ndim();
554        if index < len {
555            Ok(self.len_of(Axis(index)))
556        } else {
557            Err(error::fmt!(
558                ArrayError,
559                "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
560                index,
561                3
562            ))
563        }
564    }
565
566    fn iter(&self) -> Self::Iter<'_> {
567        self.iter()
568    }
569
570    fn as_slice(&self) -> Option<&[T]> {
571        self.as_slice()
572    }
573}
574
575#[cfg(test)]
576mod tests {
577    use super::*;
578
579    #[test]
580    fn test_f64_element_type() {
581        assert_eq!(<f64 as ArrayElementSealed>::type_tag(), 10);
582    }
583}
questdb/ingress/ndarr.rs

questdb/ingress/
ndarr.rs
