DX-64328 Array types for Gandiva (dremio#58)

Add List input and output types for Gandiva functions. Add new reference implementations for array_contains and array_remove, tested via integration with Dremio. int32, int64, double and float list types have been tested. Support List types in function specification and llvm code generation. Pass back function type information through the expression registry. See 1p here: https://docs.google.com/document/d/1exwXdUUnk5FqZLzVZyTdhqgwxTk0u9bL54aLVNM5Tas/edit
lriggs · Sep 6, 2024 · 46b3e77 · 46b3e77
1 parent 03d4693
commit 46b3e77
Show file tree

Hide file tree

Showing 42 changed files with 2,327 additions and 112 deletions.
diff --git a/cpp/src/arrow/buffer.h b/cpp/src/arrow/buffer.h
@@ -521,10 +521,27 @@ class ARROW_EXPORT ResizableBuffer : public MutableBuffer {
     return Reserve(sizeof(T) * new_nb_elements);
   }
 
+ public:
+   uint8_t* offsetBuffer;
+   int64_t offsetCapacity;
+   uint8_t* validityBuffer;
+   uint8_t* outerValidityBuffer;
+
  protected:
-  ResizableBuffer(uint8_t* data, int64_t size) : MutableBuffer(data, size) {}
+  ResizableBuffer(uint8_t* data, int64_t size) : MutableBuffer(data, size) {
+    offsetBuffer = nullptr;
+    offsetCapacity = 0;
+    validityBuffer = nullptr;
+    outerValidityBuffer = nullptr;
+
+  }
   ResizableBuffer(uint8_t* data, int64_t size, std::shared_ptr<MemoryManager> mm)
-      : MutableBuffer(data, size, std::move(mm)) {}
+      : MutableBuffer(data, size, std::move(mm)) {
+        offsetBuffer = nullptr;
+        offsetCapacity = 0;
+        validityBuffer = nullptr;
+        outerValidityBuffer = nullptr;
+      }
 };
 
 /// \defgroup buffer-allocation-functions Functions for allocating buffers

diff --git a/cpp/src/gandiva/CMakeLists.txt b/cpp/src/gandiva/CMakeLists.txt
@@ -45,6 +45,7 @@ set_source_files_properties(${GANDIVA_PRECOMPILED_CC_PATH} PROPERTIES GENERATED
 
 set(SRC_FILES
     annotator.cc
+    array_ops.cc
     bitmap_accumulator.cc
     cache.cc
     cast_time.cc
@@ -68,6 +69,7 @@ set(SRC_FILES
     function_ir_builder.cc
     function_registry.cc
     function_registry_arithmetic.cc
+    function_registry_array.cc
     function_registry_datetime.cc
     function_registry_hash.cc
     function_registry_math_ops.cc
@@ -233,6 +235,7 @@ endfunction()
 
 add_gandiva_test(internals-test
                  SOURCES
+                 array_ops_test.cc
                  bitmap_accumulator_test.cc
                  engine_llvm_test.cc
                  function_registry_test.cc

diff --git a/cpp/src/gandiva/annotator.cc b/cpp/src/gandiva/annotator.cc
@@ -46,15 +46,27 @@ FieldDescriptorPtr Annotator::MakeDesc(FieldPtr field, bool is_output) {
   int data_idx = buffer_count_++;
   int validity_idx = buffer_count_++;
   int offsets_idx = FieldDescriptor::kInvalidIdx;
+  int child_offsets_idx = FieldDescriptor::kInvalidIdx;
   if (arrow::is_binary_like(field->type()->id())) {
     offsets_idx = buffer_count_++;
   }
+
+  if (field->type()->id() == arrow::Type::LIST) {
+    offsets_idx = buffer_count_++;
+    if (arrow::is_binary_like(field->type()->field(0)->type()->id())) {
+      child_offsets_idx = buffer_count_++;
+    }
+  }
   int data_buffer_ptr_idx = FieldDescriptor::kInvalidIdx;
   if (is_output) {
     data_buffer_ptr_idx = buffer_count_++;
   }
+  int child_valid_buffer_ptr_idx = FieldDescriptor::kInvalidIdx;
+  if (field->type()->id() == arrow::Type::LIST) {
+    child_valid_buffer_ptr_idx = buffer_count_++;
+  }
   return std::make_shared<FieldDescriptor>(field, data_idx, validity_idx, offsets_idx,
-                                           data_buffer_ptr_idx);
+                                           data_buffer_ptr_idx, child_offsets_idx, child_valid_buffer_ptr_idx);
 }
 
 int Annotator::AddHolderPointer(void* holder) {
@@ -80,17 +92,76 @@ void Annotator::PrepareBuffersForField(const FieldDescriptor& desc,
   if (desc.HasOffsetsIdx()) {
     uint8_t* offsets_buf = const_cast<uint8_t*>(array_data.buffers[buffer_idx]->data());
     eval_batch->SetBuffer(desc.offsets_idx(), offsets_buf, array_data.offset);
-    ++buffer_idx;
+
+    if (desc.HasChildOffsetsIdx()) {
+      if (is_output) {
+        // if list field is output field, we should put buffer pointer into eval batch
+        // for resizing
+        uint8_t* child_offsets_buf = reinterpret_cast<uint8_t*>(
+            array_data.child_data.at(0)->buffers[buffer_idx].get());
+        eval_batch->SetBuffer(desc.child_data_offsets_idx(), child_offsets_buf,
+                              array_data.child_data.at(0)->offset);
+
+        uint8_t* child_valid_buf = reinterpret_cast<uint8_t*>(
+            array_data.child_data.at(0)->buffers[0].get());
+        eval_batch->SetBuffer(desc.child_data_validity_idx(), child_valid_buf,
+                              array_data.child_data.at(0)->offset);
+
+      } else {
+        // if list field is input field, just put buffer data into eval batch
+        uint8_t* child_offsets_buf = const_cast<uint8_t*>(
+            array_data.child_data.at(0)->buffers[buffer_idx]->data());
+        eval_batch->SetBuffer(desc.child_data_offsets_idx(), child_offsets_buf,
+                              array_data.child_data.at(0)->offset);
+
+        uint8_t* child_valid_buf = const_cast<uint8_t*>(
+            array_data.child_data.at(0)->buffers[0]->data());
+        eval_batch->SetBuffer(desc.child_data_offsets_idx(), child_valid_buf,
+                              array_data.child_data.at(0)->offset);
+      }
+    }
+    if (array_data.type->id() != arrow::Type::LIST ||
+        arrow::is_binary_like(array_data.type->field(0)->type()->id())) {
+        // primitive type list data buffer index is 1
+        // binary like type list data buffer index is 2
+        ++buffer_idx;
+        }
+  }
+
+  int const childDataIndex = 0;
+  if (array_data.type->id() != arrow::Type::LIST) {
+    uint8_t* data_buf = const_cast<uint8_t*>(array_data.buffers[buffer_idx]->data());
+    eval_batch->SetBuffer(desc.data_idx(), data_buf, array_data.offset);
+  } else {
+    uint8_t* data_buf =
+        const_cast<uint8_t*>(array_data.child_data.at(childDataIndex)->buffers[buffer_idx]->data());
+    eval_batch->SetBuffer(desc.data_idx(), data_buf, array_data.child_data.at(0)->offset);
+
+    int const childDataBufferIndex = 0;
+    if (array_data.child_data.at(childDataIndex)->buffers[childDataBufferIndex] ) {
+    uint8_t* child_valid_buf = const_cast<uint8_t*>(
+            array_data.child_data.at(childDataIndex)->buffers[childDataBufferIndex]->data());
+        eval_batch->SetBuffer(desc.child_data_validity_idx(), child_valid_buf, 0);
+    }
+
   }
 
-  uint8_t* data_buf = const_cast<uint8_t*>(array_data.buffers[buffer_idx]->data());
-  eval_batch->SetBuffer(desc.data_idx(), data_buf, array_data.offset);
   if (is_output) {
     // pass in the Buffer object for output data buffers. Can be used for resizing.
-    uint8_t* data_buf_ptr =
-        reinterpret_cast<uint8_t*>(array_data.buffers[buffer_idx].get());
-    eval_batch->SetBuffer(desc.data_buffer_ptr_idx(), data_buf_ptr, array_data.offset);
+
+    if (array_data.type->id() != arrow::Type::LIST) {
+      uint8_t* data_buf_ptr =
+          reinterpret_cast<uint8_t*>(array_data.buffers[buffer_idx].get());
+      eval_batch->SetBuffer(desc.data_buffer_ptr_idx(), data_buf_ptr, array_data.offset);
+    } else {
+      // list data buffer is in child data buffer
+      uint8_t* data_buf_ptr = reinterpret_cast<uint8_t*>(
+          array_data.child_data.at(0)->buffers[buffer_idx].get());
+      eval_batch->SetBuffer(desc.data_buffer_ptr_idx(), data_buf_ptr,
+                            array_data.child_data.at(0)->offset);
+    }
   }
+
 }
 
 EvalBatchPtr Annotator::PrepareEvalBatch(const arrow::RecordBatch& record_batch,
@@ -106,7 +177,6 @@ EvalBatchPtr Annotator::PrepareEvalBatch(const arrow::RecordBatch& record_batch,
       // skip columns not involved in the expression.
       continue;
     }
-
     PrepareBuffersForField(*(found->second), *(record_batch.column_data(i)),
                            eval_batch.get(), false /*is_output*/);
   }