tustvold commented on code in PR #4705:
URL: https://github.com/apache/arrow-rs/pull/4705#discussion_r1296001041
##########
arrow-select/src/take.rs:
##########
@@ -676,6 +692,65 @@ where
Ok(PrimitiveArray::<UInt32Type>::from(values))
}
+/// To avoid generating take implementations for every index type, instead we
+/// only generate for UInt32 and UInt64 and coerce inputs to these types
+trait ToIndices {
+ type T: ArrowPrimitiveType;
+
+ fn to_indices(&self) -> PrimitiveArray<Self::T>;
+}
+
+macro_rules! to_indices_reinterpret {
+ ($t:ty, $o:ty) => {
+ impl ToIndices for PrimitiveArray<$t> {
+ type T = $o;
+
+ fn to_indices(&self) -> PrimitiveArray<$o> {
+ let cast =
+ ScalarBuffer::new(self.values().inner().clone(), 0,
self.len());
+ PrimitiveArray::new(cast, self.nulls().cloned())
+ }
+ }
+ };
+}
+
+macro_rules! to_indices_identity {
+ ($t:ty) => {
+ impl ToIndices for PrimitiveArray<$t> {
+ type T = $t;
+
+ fn to_indices(&self) -> PrimitiveArray<$t> {
+ self.clone()
+ }
+ }
+ };
+}
+
+macro_rules! to_indices_widening {
+ ($t:ty, $o:ty) => {
+ impl ToIndices for PrimitiveArray<$t> {
+ type T = UInt32Type;
+
+ fn to_indices(&self) -> PrimitiveArray<$o> {
+ let cast = self.values().iter().copied().map(|x| x as
_).collect();
+ PrimitiveArray::new(cast, self.nulls().cloned())
+ }
+ }
+ };
+}
+
+to_indices_widening!(UInt8Type, UInt32Type);
+to_indices_widening!(Int8Type, UInt32Type);
+
+to_indices_widening!(UInt16Type, UInt32Type);
+to_indices_widening!(Int16Type, UInt32Type);
Review Comment:
8-bit and 16-bit indices are very rare in my experience, if this becomes a
performance issue for people we can easily change the specializations based on
that feedback
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
