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    let elem_size = size_of::<T>();
89    let mut total_len = 0;
90    for (i, &element) in array.iter().enumerate() {
91        unsafe {
92            std::ptr::copy_nonoverlapping(
93                &element as *const T as *const u8,
94                buf.as_mut_ptr().add(i * elem_size),
95                elem_size,
96            )
97        }
98        total_len += elem_size;
99    }
100    if total_len != expect_size {
101        return Err(error::fmt!(
102            ArrayError,
103            "Array write buffer length mismatch (actual: {}, expected: {})",
104            total_len,
105            expect_size
106        ));
107    }
108    Ok(())
109}
110
111pub(crate) fn check_and_get_array_bytes_size<A: NdArrayView<T>, T>(
112    array: &A,
113) -> Result<usize, Error>
114where
115    T: ArrayElement,
116{
117    let mut size = std::mem::size_of::<T>();
118    for dim_index in 0..array.ndim() {
119        let dim = array.dim(dim_index)?;
120        if dim > MAX_ARRAY_DIM_LEN {
121            return Err(error::fmt!(
122                ArrayError,
123                "dimension length out of range: dim {}, dim length {}, max length {}",
124                dim_index,
125                dim,
126                MAX_ARRAY_DIM_LEN
127            ));
128        }
129        // following dimension's length may be zero, so check the size in out of loop
130        size *= dim;
131    }
132
133    if size > MAX_ARRAY_BUFFER_SIZE {
134        return Err(error::fmt!(
135            ArrayError,
136            "Array buffer size too big: {}, maximum: {}",
137            size,
138            MAX_ARRAY_BUFFER_SIZE
139        ));
140    }
141    Ok(size)
142}
143
144/// Marker trait for valid array element types.
145///
146/// Implemented for primitive types that can be stored in arrays.
147/// Combines type information with data type classification.
148pub trait ArrayElement: Copy + 'static {}
149
150pub(crate) trait ArrayElementSealed {
151    /// Returns the binary format identifier for array element types compatible
152    /// with QuestDB's io.questdb.cairo.ColumnType numeric type constants.
153    fn type_tag() -> u8;
154}
155
156impl ArrayElement for f64 {}
157
158impl ArrayElementSealed for f64 {
159    fn type_tag() -> u8 {
160        10 // Double
161    }
162}
163
164/// impl NdArrayView for one dimension vector
165impl<T: ArrayElement> NdArrayView<T> for Vec<T> {
166    type Iter<'a>
167        = std::slice::Iter<'a, T>
168    where
169        T: 'a;
170
171    fn ndim(&self) -> usize {
172        1
173    }
174
175    fn dim(&self, idx: usize) -> Result<usize, Error> {
176        if idx == 0 {
177            Ok(self.len())
178        } else {
179            Err(error::fmt!(
180                    ArrayError,
181                    "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
182                    idx,
183                    1
184                ))
185        }
186    }
187
188    fn as_slice(&self) -> Option<&[T]> {
189        Some(self.as_slice())
190    }
191
192    fn iter(&self) -> Self::Iter<'_> {
193        self.as_slice().iter()
194    }
195}
196
197/// impl NdArrayView for one dimension array
198impl<T: ArrayElement, const N: usize> NdArrayView<T> for [T; N] {
199    type Iter<'a>
200        = std::slice::Iter<'a, T>
201    where
202        T: 'a;
203
204    fn ndim(&self) -> usize {
205        1
206    }
207
208    fn dim(&self, idx: usize) -> Result<usize, Error> {
209        if idx == 0 {
210            Ok(N)
211        } else {
212            Err(error::fmt!(
213                    ArrayError,
214                    "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
215                    idx,
216                    1
217                ))
218        }
219    }
220
221    fn as_slice(&self) -> Option<&[T]> {
222        Some(self)
223    }
224
225    fn iter(&self) -> Self::Iter<'_> {
226        self.as_slice().iter()
227    }
228}
229
230/// impl NdArrayView for one dimension slice
231impl<T: ArrayElement> NdArrayView<T> for &[T] {
232    type Iter<'a>
233        = std::slice::Iter<'a, T>
234    where
235        Self: 'a,
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(self.len())
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        <[T]>::iter(self)
261    }
262}
263
264/// impl NdArrayView for two dimensions vector
265impl<T: ArrayElement> NdArrayView<T> for Vec<Vec<T>> {
266    type Iter<'a>
267        = std::iter::Flatten<std::slice::Iter<'a, Vec<T>>>
268    where
269        T: 'a;
270
271    fn ndim(&self) -> usize {
272        2
273    }
274
275    fn dim(&self, idx: usize) -> Result<usize, Error> {
276        match idx {
277            0 => Ok(self.len()),
278            1 => {
279                let dim1 = self.first().map_or(0, |v| v.len());
280                if self.as_slice().iter().any(|v2| v2.len() != dim1) {
281                    return Err(error::fmt!(ArrayError, "Irregular array shape"));
282                }
283                Ok(dim1)
284            }
285            _ => Err(error::fmt!(
286                    ArrayError,
287                    "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
288                    idx,
289                    2
290                )),
291        }
292    }
293
294    fn as_slice(&self) -> Option<&[T]> {
295        None
296    }
297
298    fn iter(&self) -> Self::Iter<'_> {
299        self.as_slice().iter().flatten()
300    }
301}
302
303/// impl NdArrayView for two dimensions array
304impl<T: ArrayElement, const M: usize, const N: usize> NdArrayView<T> for [[T; M]; N] {
305    type Iter<'a>
306        = std::iter::Flatten<std::slice::Iter<'a, [T; M]>>
307    where
308        T: 'a;
309
310    fn ndim(&self) -> usize {
311        2
312    }
313
314    fn dim(&self, idx: usize) -> Result<usize, Error> {
315        match idx {
316            0 => Ok(N),
317            1 => Ok(M),
318            _ => Err(error::fmt!(
319                    ArrayError,
320                    "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
321                    idx,
322                    2
323                )),
324        }
325    }
326
327    fn as_slice(&self) -> Option<&[T]> {
328        Some(unsafe { std::slice::from_raw_parts(self.as_ptr() as *const T, N * M) })
329    }
330
331    fn iter(&self) -> Self::Iter<'_> {
332        self.as_slice().iter().flatten()
333    }
334}
335
336/// impl NdArrayView for two dimensions slices
337impl<T: ArrayElement, const M: usize> NdArrayView<T> for &[[T; M]] {
338    type Iter<'a>
339        = std::iter::Flatten<std::slice::Iter<'a, [T; M]>>
340    where
341        Self: 'a,
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(self.len()),
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, self.len() * M) })
363    }
364
365    fn iter(&self) -> Self::Iter<'_> {
366        <[[T; M]]>::iter(self).flatten()
367    }
368}
369
370/// impl NdArrayView for three dimensions vector
371impl<T: ArrayElement> NdArrayView<T> for Vec<Vec<Vec<T>>> {
372    type Iter<'a>
373        = std::iter::Flatten<std::iter::Flatten<std::slice::Iter<'a, Vec<Vec<T>>>>>
374    where
375        T: 'a;
376
377    fn ndim(&self) -> usize {
378        3
379    }
380
381    fn dim(&self, idx: usize) -> Result<usize, Error> {
382        match idx {
383            0 => Ok(self.len()),
384            1 => {
385                let dim1 = self.first().map_or(0, |v| v.len());
386                if self.as_slice().iter().any(|v2| v2.len() != dim1) {
387                    return Err(error::fmt!(ArrayError, "Irregular array shape"));
388                }
389                Ok(dim1)
390            }
391            2 => {
392                let dim2 = self
393                    .first()
394                    .and_then(|v2| v2.first())
395                    .map_or(0, |v3| v3.len());
396
397                if self
398                    .as_slice()
399                    .iter()
400                    .flat_map(|v2| v2.as_slice().iter())
401                    .any(|v3| v3.len() != dim2)
402                {
403                    return Err(error::fmt!(ArrayError, "Irregular array shape"));
404                }
405                Ok(dim2)
406            }
407            _ => Err(error::fmt!(
408                    ArrayError,
409                    "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
410                    idx,
411                    3
412                )),
413        }
414    }
415
416    fn as_slice(&self) -> Option<&[T]> {
417        None
418    }
419
420    fn iter(&self) -> Self::Iter<'_> {
421        self.as_slice().iter().flatten().flatten()
422    }
423}
424
425/// impl NdArrayView for three dimensions array
426impl<T: ArrayElement, const M: usize, const N: usize, const L: usize> NdArrayView<T>
427    for [[[T; M]; N]; L]
428{
429    type Iter<'a>
430        = std::iter::Flatten<std::iter::Flatten<std::slice::Iter<'a, [[T; M]; N]>>>
431    where
432        T: 'a;
433
434    fn ndim(&self) -> usize {
435        3
436    }
437
438    fn dim(&self, idx: usize) -> Result<usize, Error> {
439        match idx {
440            0 => Ok(L),
441            1 => Ok(N),
442            2 => Ok(M),
443            _ => Err(error::fmt!(
444                    ArrayError,
445                    "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
446                    idx,
447                    3
448                )),
449        }
450    }
451
452    fn as_slice(&self) -> Option<&[T]> {
453        Some(unsafe { std::slice::from_raw_parts(self.as_ptr() as *const T, L * N * M) })
454    }
455
456    fn iter(&self) -> Self::Iter<'_> {
457        self.as_slice().iter().flatten().flatten()
458    }
459}
460
461impl<T: ArrayElement, const M: usize, const N: usize> NdArrayView<T> for &[[[T; M]; N]] {
462    type Iter<'a>
463        = std::iter::Flatten<std::iter::Flatten<std::slice::Iter<'a, [[T; M]; N]>>>
464    where
465        Self: 'a,
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(self.len()),
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, self.len() * N * M) })
488    }
489
490    fn iter(&self) -> Self::Iter<'_> {
491        <[[[T; M]; N]]>::iter(self).flatten().flatten()
492    }
493}
494
495use crate::{error, Error};
496#[cfg(feature = "ndarray")]
497use ndarray::{ArrayView, Axis, Dimension};
498use std::slice;
499
500use super::{MAX_ARRAY_BUFFER_SIZE, MAX_ARRAY_DIM_LEN};
501
502#[cfg(feature = "ndarray")]
503impl<T, D> NdArrayView<T> for ArrayView<'_, T, D>
504where
505    T: ArrayElement,
506    D: Dimension,
507{
508    type Iter<'a>
509        = ndarray::iter::Iter<'a, T, D>
510    where
511        Self: 'a,
512        T: 'a;
513
514    fn ndim(&self) -> usize {
515        self.ndim()
516    }
517
518    fn dim(&self, index: usize) -> Result<usize, Error> {
519        let len = self.ndim();
520        if index < len {
521            Ok(self.len_of(Axis(index)))
522        } else {
523            Err(error::fmt!(
524                    ArrayError,
525                    "Dimension index out of bounds. Requested axis {}, but array only has {} dimension(s)",
526                    index,
527                    3
528                ))
529        }
530    }
531
532    fn iter(&self) -> Self::Iter<'_> {
533        self.iter()
534    }
535
536    fn as_slice(&self) -> Option<&[T]> {
537        self.as_slice()
538    }
539}
540
541#[cfg(test)]
542mod tests {
543    use super::*;
544
545    #[test]
546    fn test_f64_element_type() {
547        assert_eq!(<f64 as ArrayElementSealed>::type_tag(), 10);
548    }
549}
questdb/ingress/ndarr.rs

questdb/ingress/
ndarr.rs
