tustvold commented on code in PR #4705:
URL: https://github.com/apache/arrow-rs/pull/4705#discussion_r1295999657
##########
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);
+
+to_indices_identity!(UInt32Type);
+to_indices_reinterpret!(Int32Type, UInt32Type);
+
+to_indices_identity!(UInt64Type);
+to_indices_reinterpret!(Int64Type, UInt64Type);
Review Comment:
32-bit and 64-bit indices are both relatively common and I think justify
having separate specializations
--
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]
