| 212 | |
| 213 | #[inline(never)] |
| 214 | fn take_impl<IndexType: ArrowPrimitiveType>( |
| 215 | values: &dyn Array, |
| 216 | indices: &PrimitiveArray<IndexType>, |
| 217 | ) -> Result<ArrayRef, ArrowError> { |
| 218 | if indices.is_empty() { |
| 219 | return Ok(new_empty_array(values.data_type())); |
| 220 | } |
| 221 | downcast_primitive_array! { |
| 222 | values => Ok(Arc::new(take_primitive(values, indices)?)), |
| 223 | DataType::Boolean => { |
| 224 | let values = values.as_any().downcast_ref::<BooleanArray>().unwrap(); |
| 225 | Ok(Arc::new(take_boolean(values, indices))) |
| 226 | } |
| 227 | DataType::Utf8 => { |
| 228 | Ok(Arc::new(take_bytes(values.as_string::<i32>(), indices)?)) |
| 229 | } |
| 230 | DataType::LargeUtf8 => { |
| 231 | Ok(Arc::new(take_bytes(values.as_string::<i64>(), indices)?)) |
| 232 | } |
| 233 | DataType::Utf8View => { |
| 234 | Ok(Arc::new(take_byte_view(values.as_string_view(), indices)?)) |
| 235 | } |
| 236 | DataType::List(_) => { |
| 237 | Ok(Arc::new(take_list::<_, Int32Type>(values.as_list(), indices)?)) |
| 238 | } |
| 239 | DataType::LargeList(_) => { |
| 240 | Ok(Arc::new(take_list::<_, Int64Type>(values.as_list(), indices)?)) |
| 241 | } |
| 242 | DataType::ListView(_) => { |
| 243 | Ok(Arc::new(take_list_view::<_, Int32Type>(values.as_list_view(), indices)?)) |
| 244 | } |
| 245 | DataType::LargeListView(_) => { |
| 246 | Ok(Arc::new(take_list_view::<_, Int64Type>(values.as_list_view(), indices)?)) |
| 247 | } |
| 248 | DataType::FixedSizeList(_, length) => { |
| 249 | let values = values |
| 250 | .as_any() |
| 251 | .downcast_ref::<FixedSizeListArray>() |
| 252 | .unwrap(); |
| 253 | Ok(Arc::new(take_fixed_size_list( |
| 254 | values, |
| 255 | indices, |
| 256 | *length as u32, |
| 257 | )?)) |
| 258 | } |
| 259 | DataType::Map(_, _) => { |
| 260 | let list_arr = ListArray::from(values.as_map().clone()); |
| 261 | let list_data = take_list::<_, Int32Type>(&list_arr, indices)?; |
| 262 | let builder = list_data.into_data().into_builder().data_type(values.data_type().clone()); |
| 263 | Ok(Arc::new(MapArray::from(unsafe { builder.build_unchecked() }))) |
| 264 | } |
| 265 | DataType::Struct(fields) => { |
| 266 | let array: &StructArray = values.as_struct(); |
| 267 | let arrays = array |
| 268 | .columns() |
| 269 | .iter() |
| 270 | .map(|a| take_impl(a.as_ref(), indices)) |
| 271 | .collect::<Result<Vec<ArrayRef>, _>>()?; |