Speedup fixed row encoding

arrow-buffer/src/bigint/mod.rs

@@ -21,7 +21,7 @@ use num::cast::AsPrimitive;
 use num::{BigInt, FromPrimitive, ToPrimitive};
 use std::cmp::Ordering;
 use std::num::ParseIntError;
-use std::ops::{BitAnd, BitOr, BitXor, Neg, Shl, Shr};
+use std::ops::{BitAnd, BitOr, BitXor, Neg, Not, Shl, Shr};
 use std::str::FromStr;
 
 mod div;
@@ -126,6 +126,15 @@ impl From<i64> for i256 {
     }
 }
 
+impl Not for i256 {
+    type Output = i256;
+
+    #[inline]
+    fn not(self) -> Self::Output {
+        Self::from_parts(!self.low, !self.high)
+    }
+}
+
 /// Parse `s` with any sign and leading 0s removed
 fn parse_impl(s: &str, negative: bool) -> Result<i256, ParseI256Error> {
     if s.len() <= 38 {

arrow-row/src/fixed.rs

@@ -51,6 +51,13 @@ pub trait FixedLengthEncoding: Copy {
 
     fn encode(self) -> Self::Encoded;
 
+    fn encode_desc(self) -> Self::Encoded {
+        let mut encoded = self.encode();
+        // Flip bits to reverse order
+        encoded.as_mut().iter_mut().for_each(|v| *v = !*v);
+        encoded
+    }
+
     fn decode(encoded: Self::Encoded) -> Self;
 }
 
@@ -61,6 +68,10 @@ impl FixedLengthEncoding for bool {
         [self as u8]
     }
 
+    fn encode_desc(self) -> [u8; 1] {
+        [!self as u8]
+    }
+
     fn decode(encoded: Self::Encoded) -> Self {
         encoded[0] != 0
     }
@@ -71,13 +82,24 @@ macro_rules! encode_signed {
         impl FixedLengthEncoding for $t {
             type Encoded = [u8; $n];
 
+            #[inline]
             fn encode(self) -> [u8; $n] {
                 let mut b = self.to_be_bytes();
                 // Toggle top "sign" bit to ensure consistent sort order
                 b[0] ^= 0x80;
                 b
             }
 
+            #[inline]
+            fn encode_desc(self) -> [u8; $n] {
+                // fast path for descending order
+                let b = !self;
+                let mut b = b.to_be_bytes();
+                // Toggle top "sign" bit to ensure consistent sort order
+                b[0] ^= 0x80;
+                b
+            }
+
             fn decode(mut encoded: Self::Encoded) -> Self {
                 // Toggle top "sign" bit
                 encoded[0] ^= 0x80;
@@ -99,10 +121,16 @@ macro_rules! encode_unsigned {
         impl FixedLengthEncoding for $t {
             type Encoded = [u8; $n];
 
+            // #[inline]
             fn encode(self) -> [u8; $n] {
                 self.to_be_bytes()
             }
 
+            // #[inline]
+            fn encode_desc(self) -> [u8; $n] {
+                (!self).to_be_bytes()
+            }
+
             fn decode(encoded: Self::Encoded) -> Self {
                 Self::from_be_bytes(encoded)
             }
@@ -224,22 +252,38 @@ pub fn encode<T: FixedLengthEncoding>(
     nulls: &NullBuffer,
     opts: SortOptions,
 ) {
-    for (value_idx, is_valid) in nulls.iter().enumerate() {
-        let offset = &mut offsets[value_idx + 1];
-        let end_offset = *offset + T::ENCODED_LEN;
-        if is_valid {
-            let to_write = &mut data[*offset..end_offset];
-            to_write[0] = 1;
-            let mut encoded = values[value_idx].encode();
-            if opts.descending {
-                // Flip bits to reverse order
-                encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
+    #[inline]
+    fn encode<const DESC: bool, T: FixedLengthEncoding>(
+        data: &mut [u8],
+        offsets: &mut [usize],
+        values: &[T],
+        nulls: &NullBuffer,
+        opts: SortOptions,
+    ) {
+        let null_sentinel = null_sentinel(opts);
+        for ((is_valid, offset), val) in
+            nulls.iter().zip(offsets[1..].iter_mut()).zip(values.iter())
+        {
+            let end_offset = *offset + T::ENCODED_LEN;
+            if is_valid {
+                let to_write = &mut data[*offset..end_offset];
+                to_write[0] = 1;
+                let encoded = if DESC {
+                    val.encode_desc()
+                } else {
+                    val.encode()
+                };
+
+                to_write[1..].copy_from_slice(encoded.as_ref())
+            } else {
+                data[*offset] = null_sentinel;
             }
-            to_write[1..].copy_from_slice(encoded.as_ref())
-        } else {
-            data[*offset] = null_sentinel(opts);
+            *offset = end_offset;
         }
-        *offset = end_offset;
+    }
+    match opts.descending {
+        true => encode::<true, T>(data, offsets, values, nulls, opts),
+        false => encode::<false, T>(data, offsets, values, nulls, opts),
     }
 }
 
@@ -251,20 +295,27 @@ pub fn encode_not_null<T: FixedLengthEncoding>(
     values: &[T],
     opts: SortOptions,
 ) {
-    for (value_idx, val) in values.iter().enumerate() {
-        let offset = &mut offsets[value_idx + 1];
-        let end_offset = *offset + T::ENCODED_LEN;
-
-        let to_write = &mut data[*offset..end_offset];
-        to_write[0] = 1;
-        let mut encoded = val.encode();
-        if opts.descending {
-            // Flip bits to reverse order
-            encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
+    #[inline]
+    fn encode<const DESC: bool, T: FixedLengthEncoding>(
+        data: &mut [u8],
+        offsets: &mut [usize],
+        values: &[T],
+    ) {
+        for (val, offset) in values.iter().zip(offsets[1..].iter_mut()) {
+            let to_write = &mut data[*offset..*offset + T::ENCODED_LEN];
+            to_write[0] = 1;
+            let encoded = if DESC {
+                val.encode_desc()
+            } else {
+                val.encode()
+            };
+            to_write[1..].copy_from_slice(encoded.as_ref());
+            *offset += T::ENCODED_LEN;
         }
-        to_write[1..].copy_from_slice(encoded.as_ref());
-
-        *offset = end_offset;
+    }
+    match opts.descending {
+        true => encode::<true, T>(data, offsets, values),
+        false => encode::<false, T>(data, offsets, values),
     }
 }