Impl scatter fix width

zanmato1984 · Jun 12, 2024 · e748eeb · e748eeb
1 parent 3f639d9
commit e748eeb
Showing 1 changed file with 25 additions and 0 deletions.
diff --git a/cpp/src/arrow/array/scatter.cc b/cpp/src/arrow/array/scatter.cc
@@ -42,6 +42,24 @@ Status ScatterBitmap(const uint8_t* in_bitmap, const uint8_t* mask_bitmap,
   return Status::OK();
 }
 
+Result<std::shared_ptr<Buffer>> ScatterBuffer(const Buffer& src, const uint8_t* mask,
+                                              int64_t byte_width, int64_t length,
+                                              MemoryPool* pool) {
+  ARROW_ASSIGN_OR_RAISE(auto out, AllocateBuffer(length * byte_width, pool));
+  auto* in_data = src.data();
+  auto* out_data = out->mutable_data();
+  int64_t i_in = 0, i_out = 0;
+  VisitNullBitmapInline(
+      mask, /*valid_bits_offset=*/0, length, kUnknownNullCount,
+      [&] {
+        std::memcpy(out_data + i_out++ * byte_width, in_data + i_in++ * byte_width,
+                    byte_width);
+      },
+      [&] { ++i_out; });
+
+  return Status::OK();
+}
+
 struct ScatterImpl {
   explicit ScatterImpl(const ArrayData& in, const BooleanArray& mask, MemoryPool* pool)
       : in_(in), mask_(mask), pool_(pool), out_(std::shared_ptr<ArrayData>()) {
@@ -98,6 +116,13 @@ struct ScatterImpl {
     return Status::OK();
   }
 
+  Status Visit(const FixedWidthType& type) {
+    DCHECK_EQ(type.bit_width() % 8, 0);
+    return ScatterBuffer(*in_.buffers[1], out_->buffers[0]->data(), type.bit_width() / 8,
+                         out_->length, pool_)
+        .Value(&out_->buffers[1]);
+  }
+
   Status Visit(const DataType&) {
     return Status::NotImplemented("Scatter not implemented for type ", *out_->type);
   }