1include!(concat!(env!("OUT_DIR"), "/simd_lanes.rs"));
19
20use std::hash::Hash;
21#[cfg(feature = "simd")]
22use std::simd::num::SimdUint;
23use std::simd::{LaneCount, SupportedLaneCount};
24
25#[cfg(feature = "fast_hash")]
26use ahash::AHashMap;
27use minarrow::{
28 BooleanAVT, BooleanArray, CategoricalAVT, CategoricalArray, FloatAVT, FloatArray, Integer,
29 IntegerAVT, IntegerArray, StringAVT, StringArray, Vec64,
30};
31#[cfg(feature = "datetime")]
32use minarrow::{DatetimeAVT, DatetimeArray};
33use num_traits::{Float, Signed};
34#[cfg(not(feature = "fast_hash"))]
35use std::collections::HashMap;
36
37use crate::kernels::logical::not_bool;
38use minarrow::enums::error::KernelError;
39
40#[inline(always)]
43fn prealloc_vec<T: Copy>(len: usize) -> Vec64<T> {
44 let mut v = Vec64::<T>::with_capacity(len);
45 unsafe { v.set_len(len) };
47 v
48}
49
50#[cfg(feature = "simd")]
53mod simd_impl {
54 use core::simd::{LaneCount, Simd, SimdElement, SupportedLaneCount};
55
56 use num_traits::Zero;
57
58 #[inline(always)]
60 pub fn negate_dense<T, const LANES: usize>(a: &[T], out: &mut [T])
61 where
62 T: SimdElement + core::ops::Neg<Output = T> + Copy + Zero,
63 LaneCount<LANES>: SupportedLaneCount,
64 Simd<T, LANES>: core::ops::Sub<Output = Simd<T, LANES>>,
65 {
66 let mut i = 0;
67 while i + LANES <= a.len() {
68 let v = Simd::<T, LANES>::from_slice(&a[i..i + LANES]);
69 let res = Simd::<T, LANES>::splat(T::zero()) - v;
70 res.copy_to_slice(&mut out[i..i + LANES]);
71 i += LANES;
72 }
73 for j in i..a.len() {
75 out[j] = -a[j];
76 }
77 }
78}
79
80macro_rules! impl_unary_neg_int {
84 ($fn_name:ident, $ty:ty, $lanes:expr) => {
85 #[inline(always)]
98 pub fn $fn_name(window: IntegerAVT<$ty>) -> IntegerArray<$ty> {
99 let (arr, offset, len) = window;
100 let src = &arr.data[offset..offset + len];
101 let mut data = prealloc_vec::<$ty>(len);
102
103 #[cfg(feature = "simd")]
104 simd_impl::negate_dense::<$ty, $lanes>(src, &mut data);
105
106 #[cfg(not(feature = "simd"))]
107 for i in 0..len {
108 data[i] = -src[i];
109 }
110
111 IntegerArray {
112 data: data.into(),
113 null_mask: arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len)),
114 }
115 }
116 };
117}
118
119#[cfg(feature = "datetime")]
120macro_rules! impl_unary_neg_datetime {
122 ($fn_name:ident, $ty:ty, $lanes:expr) => {
123 #[inline(always)]
136 pub fn $fn_name(window: DatetimeAVT<$ty>) -> DatetimeArray<$ty> {
137 let (arr, offset, len) = window;
138 let src = &arr.data[offset..offset + len];
139 let mut data = prealloc_vec::<$ty>(len);
140
141 #[cfg(feature = "simd")]
142 simd_impl::negate_dense::<$ty, $lanes>(src, &mut data);
143
144 #[cfg(not(feature = "simd"))]
145 for i in 0..len {
146 data[i] = -src[i];
147 }
148
149 DatetimeArray {
150 data: data.into(),
151 null_mask: arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len)),
152 time_unit: arr.time_unit.clone(),
153 }
154 }
155 };
156}
157
158#[cfg(feature = "datetime")]
159#[inline(always)]
176pub fn unary_negate_int_datetime<T>(window: DatetimeAVT<T>) -> DatetimeArray<T>
177where
178 T: Signed + Copy + 'static,
179{
180 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<i32>() {
181 return unsafe { std::mem::transmute(unary_neg_datetime_i32(std::mem::transmute(window))) };
182 }
183 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<i64>() {
184 return unsafe { std::mem::transmute(unary_neg_datetime_i64(std::mem::transmute(window))) };
185 }
186 unreachable!("unsupported datetime type for negation")
187}
188
189#[cfg(feature = "datetime")]
190impl_unary_neg_datetime!(unary_neg_datetime_i32, i32, W32);
191#[cfg(feature = "datetime")]
192impl_unary_neg_datetime!(unary_neg_datetime_i64, i64, W64);
193#[cfg(feature = "extended_numeric_types")]
194impl_unary_neg_int!(unary_neg_i8, i8, W8);
195#[cfg(feature = "extended_numeric_types")]
196impl_unary_neg_int!(unary_neg_i16, i16, W16);
197impl_unary_neg_int!(unary_neg_i32, i32, W32);
198impl_unary_neg_int!(unary_neg_i64, i64, W64);
199
200#[inline(always)]
219pub fn unary_negate_int<T>(window: IntegerAVT<T>) -> IntegerArray<T>
220where
221 T: Signed + Copy + 'static,
222{
223 macro_rules! dispatch {
224 ($t:ty, $f:ident) => {
225 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<$t>() {
226 return unsafe { std::mem::transmute($f(std::mem::transmute(window))) };
227 }
228 };
229 }
230 #[cfg(feature = "extended_numeric_types")]
231 dispatch!(i8, unary_neg_i8);
232 #[cfg(feature = "extended_numeric_types")]
233 dispatch!(i16, unary_neg_i16);
234 dispatch!(i32, unary_neg_i32);
235 dispatch!(i64, unary_neg_i64);
236
237 unreachable!("unsupported integer type")
238}
239
240pub fn unary_negate_u32_to_i32(window: IntegerAVT<u32>) -> IntegerArray<i32> {
253 let (arr, offset, len) = window;
254 let src = &arr.data[offset..offset + len];
255 let mut data = prealloc_vec::<i32>(len);
256
257 for (dst, &v) in data.iter_mut().zip(src) {
258 *dst = -(v as i32);
259 }
260
261 IntegerArray {
262 data: data.into(),
263 null_mask: arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len)),
264 }
265}
266
267pub fn unary_negate_u64_to_i64(window: IntegerAVT<u64>) -> IntegerArray<i64> {
280 let (arr, offset, len) = window;
281 let src = &arr.data[offset..offset + len];
282 let mut data = prealloc_vec::<i64>(len);
283
284 #[cfg(feature = "simd")]
285 {
286 use core::simd::Simd;
287 const LANES: usize = W64;
288 let mut i = 0;
289 while i + LANES <= len {
290 let v = Simd::<u64, LANES>::from_slice(&src[i..i + LANES]).cast::<i64>();
291 (Simd::<i64, LANES>::splat(0) - v).copy_to_slice(&mut data[i..i + LANES]);
292 i += LANES;
293 }
294 for j in i..len {
295 data[j] = -(src[j] as i64);
296 }
297 }
298 #[cfg(not(feature = "simd"))]
299 for (dst, &v) in data.iter_mut().zip(src) {
300 *dst = -(v as i64);
301 }
302
303 IntegerArray {
304 data: data.into(),
305 null_mask: arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len)),
306 }
307}
308
309macro_rules! impl_unary_neg_float {
313 ($fname:ident, $ty:ty, $lanes:expr) => {
314 #[inline(always)]
327 pub fn $fname(window: FloatAVT<$ty>) -> FloatArray<$ty> {
328 let (arr, offset, len) = window;
329 let src = &arr.data[offset..offset + len];
330 let mut data = prealloc_vec::<$ty>(len);
331
332 #[cfg(feature = "simd")]
333 simd_impl::negate_dense::<$ty, $lanes>(src, &mut data);
334
335 #[cfg(not(feature = "simd"))]
336 for i in 0..len {
337 data[i] = -src[i];
338 }
339
340 FloatArray {
341 data: data.into(),
342 null_mask: arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len)),
343 }
344 }
345 };
346}
347
348impl_unary_neg_float!(unary_neg_f32, f32, W32);
349impl_unary_neg_float!(unary_neg_f64, f64, W64);
350
351#[inline(always)]
368pub fn unary_negate_float<T>(window: FloatAVT<T>) -> FloatArray<T>
369where
370 T: Float + Copy + 'static,
371{
372 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() {
373 return unsafe { std::mem::transmute(unary_neg_f32(std::mem::transmute(window))) };
374 }
375 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f64>() {
376 return unsafe { std::mem::transmute(unary_neg_f64(std::mem::transmute(window))) };
377 }
378 unreachable!("unsupported float type")
379}
380
381pub fn unary_not_bool<const LANES: usize>(
383 arr_window: BooleanAVT<()>,
384) -> Result<BooleanArray<()>, KernelError>
385where
386 LaneCount<LANES>: SupportedLaneCount,
387{
388 not_bool::<LANES>(arr_window)
389}
390
391pub fn unary_reverse_str<T: Integer>(arr: StringAVT<T>) -> StringArray<T> {
408 let (array, offset, len) = arr;
409 let offsets = &array.offsets;
410 let data_buf = &array.data;
411 let mask = array.null_mask.as_ref();
412
413 let total_bytes = if len == 0 {
415 0
416 } else {
417 let start = offsets[offset].to_usize();
418 let end = offsets[offset + len].to_usize();
419 end - start
420 };
421
422 let mut out_offsets = Vec64::<T>::with_capacity(len + 1);
424 let mut out_data = Vec64::<u8>::with_capacity(total_bytes);
425 unsafe {
426 out_offsets.set_len(len + 1);
427 }
428 out_offsets[0] = T::zero();
429
430 for i in 0..len {
431 if mask.map_or(true, |m| m.get(offset + i)) {
433 let start = offsets[offset + i].to_usize();
434 let end = offsets[offset + i + 1].to_usize();
435 let s = unsafe { std::str::from_utf8_unchecked(&data_buf[start..end]) };
436 for ch in s.chars().rev() {
437 let mut buf = [0u8; 4];
438 let encoded = ch.encode_utf8(&mut buf);
439 out_data.extend_from_slice(encoded.as_bytes());
440 }
441 }
442 out_offsets[i + 1] = T::from_usize(out_data.len());
444 }
445
446 let out_null_mask = mask.map(|m| m.slice_clone(offset, len));
447
448 StringArray {
449 offsets: out_offsets.into(),
450 data: out_data.into(),
451 null_mask: out_null_mask,
452 }
453}
454
455pub fn unary_reverse_dict<T: Integer + Hash>(arr: CategoricalAVT<T>) -> CategoricalArray<T> {
472 let (array, offset, len) = arr;
473 let mask = array.null_mask.as_ref();
474
475 let windowed_codes = array.data[offset..offset + len].to_vec();
477
478 #[cfg(feature = "fast_hash")]
480 let mut remap: AHashMap<T, T> = AHashMap::new();
481 #[cfg(not(feature = "fast_hash"))]
482 let mut remap: HashMap<T, T> = HashMap::new();
483 let mut new_uniques = Vec64::<String>::new();
484 let mut new_codes = Vec64::<T>::with_capacity(len);
485
486 for &code in &windowed_codes {
487 if !remap.contains_key(&code) {
488 let reversed = array.unique_values[code.to_usize()]
489 .chars()
490 .rev()
491 .collect::<String>();
492 remap.insert(code, T::from_usize(new_uniques.len()));
493 new_uniques.push(reversed);
494 }
495 new_codes.push(remap[&code]);
496 }
497
498 let out_null_mask = mask.map(|m| m.slice_clone(offset, len));
499
500 CategoricalArray {
501 data: new_codes.into(),
502 unique_values: new_uniques,
503 null_mask: out_null_mask,
504 }
505}
506
507#[cfg(test)]
508mod tests {
509 use minarrow::structs::variants::categorical::CategoricalArray;
510 use minarrow::structs::variants::float::FloatArray;
511 use minarrow::structs::variants::integer::IntegerArray;
512 use minarrow::structs::variants::string::StringArray;
513 use minarrow::{Bitmask, BooleanArray, MaskedArray};
514
515 use super::*;
516
517 fn bm(bits: &[bool]) -> Bitmask {
520 let mut m = Bitmask::new_set_all(bits.len(), false);
521 for (i, &b) in bits.iter().enumerate() {
522 m.set(i, b);
523 }
524 m
525 }
526
527 fn expect_int<T: PartialEq + std::fmt::Debug>(
528 arr: &IntegerArray<T>,
529 values: &[T],
530 valid: &[bool],
531 ) {
532 assert_eq!(arr.data.as_slice(), values);
533 let mask = arr.null_mask.as_ref().expect("mask missing");
534 for (i, &v) in valid.iter().enumerate() {
535 assert_eq!(mask.get(i), v, "mask bit {}", i);
536 }
537 }
538
539 fn expect_float<T: num_traits::Float + std::fmt::Debug>(
540 arr: &FloatArray<T>,
541 values: &[T],
542 valid: &[bool],
543 eps: T,
544 ) {
545 assert_eq!(arr.data.len(), values.len());
546 for (a, b) in arr.data.iter().zip(values.iter()) {
547 assert!((*a - *b).abs() <= eps, "value mismatch {:?} vs {:?}", a, b);
548 }
549 let mask = arr.null_mask.as_ref().expect("mask missing");
550 for (i, &v) in valid.iter().enumerate() {
551 assert_eq!(mask.get(i), v, "mask bit {}", i);
552 }
553 }
554
555 #[cfg(feature = "extended_numeric_types")]
558 #[test]
559 fn neg_i8_dense() {
560 let arr = IntegerArray::<i8>::from_slice(&[1, -2, 127]);
561 let out = unary_neg_i8((&arr, 0, arr.len()));
562 assert_eq!(out.data.as_slice(), &[-1, 2, -127]);
563 assert!(out.null_mask.is_none());
564 }
565
566 #[cfg(feature = "extended_numeric_types")]
567 #[test]
568 fn neg_i16_masked() {
569 let mut arr = IntegerArray::<i16>::from_slice(&[-4, 12, 8, 0]);
570 arr.null_mask = Some(bm(&[true, false, true, true]));
571 let out = unary_neg_i16((&arr, 0, arr.len()));
572 expect_int(&out, &[4, -12, -8, 0], &[true, false, true, true]);
573 }
574
575 #[test]
576 fn neg_i32_empty() {
577 let arr = IntegerArray::<i32>::from_slice(&[]);
578 let out = unary_neg_i32((&arr, 0, arr.len()));
579 assert_eq!(out.data.len(), 0);
580 }
581
582 #[test]
583 fn neg_i64_all_nulls() {
584 let mut arr = IntegerArray::<i64>::from_slice(&[5, 10]);
585 arr.null_mask = Some(bm(&[false, false]));
586 let out = unary_neg_i64((&arr, 0, arr.len()));
587 expect_int(&out, &[-5, -10], &[false, false]);
588 }
589
590 #[cfg(feature = "extended_numeric_types")]
591 #[test]
592 fn neg_dispatch_i16() {
593 let mut arr = IntegerArray::<i16>::from_slice(&[-2, 4]);
594 arr.null_mask = Some(bm(&[true, true]));
595 let out = unary_negate_int((&arr, 0, arr.len()));
596 expect_int(&out, &[2, -4], &[true, true]);
597 }
598
599 #[test]
602 fn neg_u32_to_i32() {
603 let mut arr = IntegerArray::<u32>::from_slice(&[1, 2, 100]);
604 arr.null_mask = Some(bm(&[true, false, true]));
605 let out = unary_negate_u32_to_i32((&arr, 0, arr.len()));
606 expect_int(&out, &[-1, -2, -100], &[true, false, true]);
607 }
608
609 #[test]
610 fn neg_u64_to_i64() {
611 let arr = IntegerArray::<u64>::from_slice(&[3, 4, 0]);
612 let out = unary_negate_u64_to_i64((&arr, 0, arr.len()));
613 assert_eq!(out.data.as_slice(), &[-3, -4, 0]);
614 }
615
616 #[test]
619 fn neg_f32_dense() {
620 let arr = FloatArray::<f32>::from_slice(&[0.5, -1.5, 2.0]);
621 let out = unary_neg_f32((&arr, 0, arr.len()));
622 assert_eq!(out.data.as_slice(), &[-0.5, 1.5, -2.0]);
623 assert!(out.null_mask.is_none());
624 }
625
626 #[test]
627 fn neg_f64_masked() {
628 let mut arr = FloatArray::<f64>::from_slice(&[1.1, -2.2, 3.3]);
629 arr.null_mask = Some(bm(&[true, false, true]));
630 let out = unary_neg_f64((&arr, 0, arr.len()));
631 expect_float(&out, &[-1.1, 2.2, -3.3], &[true, false, true], 1e-12);
632 }
633
634 #[test]
635 fn neg_dispatch_f64() {
636 let arr = FloatArray::<f64>::from_slice(&[2.2, -4.4]);
637 let out = unary_negate_float((&arr, 0, arr.len()));
638 assert_eq!(out.data.as_slice(), &[-2.2, 4.4]);
639 }
640
641 #[test]
644 fn not_bool_basic() {
645 let arr = BooleanArray::from_slice(&[true, false, true, false]);
646 let out = unary_not_bool::<W64>((&arr, 0, arr.len())).unwrap();
647 assert_eq!(out.data.as_slice(), &[0b00001010]);
648 assert!(out.null_mask.is_none());
649 }
650
651 #[test]
652 fn not_bool_masked() {
653 let mut arr = BooleanArray::from_slice(&[false, false, true, true]);
654 arr.null_mask = Some(bm(&[true, false, true, true]));
655 let out = unary_not_bool::<W64>((&arr, 0, arr.len())).unwrap();
656 assert_eq!(out.data.as_slice(), &[0b00001100]);
657 assert_eq!(out.null_mask, arr.null_mask);
658 }
659
660 #[test]
663 fn reverse_str_basic() {
664 let arr = StringArray::<u32>::from_slice(&["ab", "xyz", ""]);
665 let out = unary_reverse_str((&arr, 0, arr.len()));
666 assert_eq!(out.get(0), Some("ba"));
667 assert_eq!(out.get(1), Some("zyx"));
668 assert_eq!(out.get(2), Some(""));
669 }
670
671 #[test]
672 fn reverse_str_basic_chunk() {
673 let arr = StringArray::<u32>::from_slice(&["xxx", "ab", "xyz", ""]);
674 let out = unary_reverse_str((&arr, 1, 3)); assert_eq!(out.get(0), Some("ba"));
676 assert_eq!(out.get(1), Some("zyx"));
677 assert_eq!(out.get(2), Some(""));
678 }
679
680 #[test]
681 fn reverse_str_with_nulls() {
682 let mut arr = StringArray::<u32>::from_slice(&["apple", "banana", "carrot"]);
683 arr.null_mask = Some(bm(&[true, false, true]));
684 let out = unary_reverse_str((&arr, 0, arr.len()));
685 assert_eq!(out.get(0), Some("elppa"));
686 assert_eq!(out.get(1), None);
687 assert_eq!(out.get(2), Some("torrac"));
688 assert_eq!(out.null_mask, arr.null_mask);
689 }
690
691 #[test]
692 fn reverse_str_with_nulls_chunk() {
693 let mut arr = StringArray::<u32>::from_slice(&["zero", "apple", "banana", "carrot"]);
694 arr.null_mask = Some(bm(&[true, true, false, true]));
695 let out = unary_reverse_str((&arr, 1, 3)); assert_eq!(out.get(0), Some("elppa"));
697 assert_eq!(out.get(1), None);
698 assert_eq!(out.get(2), Some("torrac"));
699 assert_eq!(
700 out.null_mask.as_ref().unwrap().as_slice(),
701 bm(&[true, false, true]).as_slice()
702 );
703 }
704
705 #[test]
708 fn reverse_dict_basic() {
709 let arr = CategoricalArray::<u32>::from_values(["cat", "dog", "bird", "cat"]);
710 let out = unary_reverse_dict((&arr, 0, arr.data.len()));
711 let uniq: Vec<String> = out.unique_values.iter().map(|s| s.clone()).collect();
712 assert!(uniq.contains(&"tac".to_string()));
713 assert!(uniq.contains(&"god".to_string()));
714 assert!(uniq.contains(&"drib".to_string()));
715 assert_eq!(out.data, arr.data);
716 }
717
718 #[test]
719 fn reverse_dict_basic_chunk() {
720 let arr = CategoricalArray::<u32>::from_values(["z", "cat", "dog", "bird"]);
721 let out = unary_reverse_dict((&arr, 1, 3));
722 let uniq: Vec<String> = out.unique_values.iter().map(|s| s.clone()).collect();
723 assert!(uniq.contains(&"tac".to_string()));
724 assert!(uniq.contains(&"god".to_string()));
725 assert!(uniq.contains(&"drib".to_string()));
726 assert_eq!(&out.data[..], &[0, 1, 2]);
727 }
728
729 #[test]
730 fn test_unary_reverse_str_empty_and_all_nulls() {
731 let arr0 = StringArray::<u32>::from_slice(&[]);
733 let out0 = unary_reverse_str((&arr0, 0, arr0.len()));
734 assert_eq!(out0.len(), 0);
735
736 let mut arr1 = StringArray::<u32>::from_slice(&["a", "b"]);
738 arr1.null_mask = Some(bm(&[false, false]));
739 let out1 = unary_reverse_str((&arr1, 0, arr1.len()));
740 assert_eq!(out1.get(0), None);
742 assert_eq!(out1.get(1), None);
743 assert_eq!(out1.null_mask, arr1.null_mask);
744 }
745
746 #[test]
747 fn test_unary_reverse_str_empty_and_all_nulls_chunk() {
748 let mut arr = StringArray::<u32>::from_slice(&["x", "a", "b", "y"]);
750 arr.null_mask = Some(bm(&[true, false, false, true]));
751 let out = unary_reverse_str((&arr, 1, 2)); assert_eq!(out.get(0), None);
753 assert_eq!(out.get(1), None);
754 assert_eq!(out.null_mask, Some(bm(&[false, false])));
755 }
756}