apache · paleolimbot · Nov 29, 2023 · Nov 27, 2023 · Nov 27, 2023 · Nov 27, 2023
diff --git a/extensions/nanoarrow_device/src/nanoarrow/nanoarrow_device.c b/extensions/nanoarrow_device/src/nanoarrow/nanoarrow_device.c
@@ -408,7 +408,7 @@ static ArrowErrorCode ArrowDeviceArrayViewCopyInternal(struct ArrowDevice* devic
   dst->offset = src->offset;
   dst->null_count = src->null_count;
 
-  for (int i = 0; i < 3; i++) {
+  for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     if (src->layout.buffer_type[i] == NANOARROW_BUFFER_TYPE_NONE) {
       break;
     }

diff --git a/extensions/nanoarrow_ipc/src/nanoarrow/nanoarrow_ipc_decoder.c b/extensions/nanoarrow_ipc/src/nanoarrow/nanoarrow_ipc_decoder.c
@@ -1155,7 +1155,7 @@ static void ArrowIpcDecoderInitFields(struct ArrowIpcField* fields,
   field->array = array;
   field->buffer_offset = *n_buffers;
 
-  for (int i = 0; i < 3; i++) {
+  for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     *n_buffers += array_view->layout.buffer_type[i] != NANOARROW_BUFFER_TYPE_NONE;
   }
 
@@ -1524,7 +1524,7 @@ static int ArrowIpcDecoderWalkSetArrayView(struct ArrowIpcArraySetter* setter,
   array_view->null_count = ns(FieldNode_null_count(field));
   setter->field_i += 1;
 
-  for (int64_t i = 0; i < 3; i++) {
+  for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     if (array_view->layout.buffer_type[i] == NANOARROW_BUFFER_TYPE_NONE) {
       break;
     }

diff --git a/python/.gitignore b/python/.gitignore
@@ -19,6 +19,7 @@
 src/nanoarrow/nanoarrow.c
 src/nanoarrow/nanoarrow.h
 src/nanoarrow/nanoarrow_device.h
+src/nanoarrow/nanoarrow_testing.hpp
 src/nanoarrow/nanoarrow_c.pxd
 src/nanoarrow/*.c
 

diff --git a/python/bootstrap.py b/python/bootstrap.py
@@ -36,6 +36,9 @@ def generate_nanoarrow_pxd(self, file_in, file_out):
         # Strip comments
         content = self.re_comment.sub("", content)
 
+        # Replace NANOARROW_MAX_FIXED_BUFFERS with its value
+        content = self.re_max_buffers.sub("3", content)
+
         # Find types and function definitions
         types = self._find_types(content)
         func_defs = self._find_func_defs(content)
@@ -59,6 +62,7 @@ def generate_nanoarrow_pxd(self, file_in, file_out):
             output.write(b"\n")
             output.write(b"    ctypedef int ArrowErrorCode\n")
             output.write(b"    cdef int NANOARROW_OK\n")
+            output.write(b"    cdef int NANOARROW_MAX_FIXED_BUFFERS\n")
             output.write(b"\n")
 
             for type in types_cython:
@@ -71,6 +75,7 @@ def generate_nanoarrow_pxd(self, file_in, file_out):
 
     def _define_regexes(self):
         self.re_comment = re.compile(r"\s*//[^\n]*")
+        self.re_max_buffers = re.compile(r"NANOARROW_MAX_FIXED_BUFFERS")
         self.re_type = re.compile(
             r"(?P<type>struct|union|enum) (?P<name>Arrow[^ ]+) {(?P<body>[^}]*)}"
         )

diff --git a/src/nanoarrow/array.c b/src/nanoarrow/array.c
@@ -437,7 +437,7 @@ static void ArrowArrayFlushInternalPointers(struct ArrowArray* array) {
   struct ArrowArrayPrivateData* private_data =
       (struct ArrowArrayPrivateData*)array->private_data;
 
-  for (int64_t i = 0; i < 3; i++) {
+  for (int64_t i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     private_data->buffer_data[i] = ArrowArrayBuffer(array, i)->data;
   }
 
@@ -621,7 +621,7 @@ void ArrowArrayViewReset(struct ArrowArrayView* array_view) {
 }
 
 void ArrowArrayViewSetLength(struct ArrowArrayView* array_view, int64_t length) {
-  for (int i = 0; i < 3; i++) {
+  for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     int64_t element_size_bytes = array_view->layout.element_size_bits[i] / 8;
 
     switch (array_view->layout.buffer_type[i]) {
@@ -671,26 +671,13 @@ void ArrowArrayViewSetLength(struct ArrowArrayView* array_view, int64_t length)
 static int ArrowArrayViewSetArrayInternal(struct ArrowArrayView* array_view,
                                           struct ArrowArray* array,
                                           struct ArrowError* error) {
-  // Check length and offset
-  if (array->offset < 0) {
-    ArrowErrorSet(error, "Expected array offset >= 0 but found array offset of %ld",
-                  (long)array->offset);
-    return EINVAL;
-  }
-
-  if (array->length < 0) {
-    ArrowErrorSet(error, "Expected array length >= 0 but found array length of %ld",
-                  (long)array->length);
-    return EINVAL;
-  }
-
   array_view->array = array;
   array_view->offset = array->offset;
   array_view->length = array->length;
   array_view->null_count = array->null_count;
 
   int64_t buffers_required = 0;
-  for (int i = 0; i < 3; i++) {
+  for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     if (array_view->layout.buffer_type[i] == NANOARROW_BUFFER_TYPE_NONE) {
       break;
     }
@@ -749,6 +736,18 @@ static int ArrowArrayViewSetArrayInternal(struct ArrowArrayView* array_view,
 
 static int ArrowArrayViewValidateMinimal(struct ArrowArrayView* array_view,
                                          struct ArrowError* error) {
+  if (array_view->length < 0) {
+    ArrowErrorSet(error, "Expected length >= 0 but found length %ld",
+                  (long)array_view->length);
+    return EINVAL;
+  }
+
+  if (array_view->offset < 0) {
+    ArrowErrorSet(error, "Expected offset >= 0 but found offset %ld",
+                  (long)array_view->offset);
+    return EINVAL;
+  }
+
   // Calculate buffer sizes that do not require buffer access. If marked as
   // unknown, assign the buffer size; otherwise, validate it.
   int64_t offset_plus_length = array_view->offset + array_view->length;
@@ -1103,7 +1102,7 @@ static int ArrowAssertInt8In(struct ArrowBufferView view, const int8_t* values,
 
 static int ArrowArrayViewValidateFull(struct ArrowArrayView* array_view,
                                       struct ArrowError* error) {
-  for (int i = 0; i < 3; i++) {
+  for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     switch (array_view->layout.buffer_type[i]) {
       case NANOARROW_BUFFER_TYPE_DATA_OFFSET:
         if (array_view->layout.element_size_bits[i] == 32) {

diff --git a/src/nanoarrow/array_inline.h b/src/nanoarrow/array_inline.h
@@ -140,7 +140,7 @@ static inline ArrowErrorCode ArrowArrayStartAppending(struct ArrowArray* array)
   }
 
   // Initialize any data offset buffer with a single zero
-  for (int i = 0; i < 3; i++) {
+  for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     if (private_data->layout.buffer_type[i] == NANOARROW_BUFFER_TYPE_DATA_OFFSET &&
         private_data->layout.element_size_bits[i] == 64) {
       NANOARROW_RETURN_NOT_OK(ArrowBufferAppendInt64(ArrowArrayBuffer(array, i), 0));
@@ -163,7 +163,7 @@ static inline ArrowErrorCode ArrowArrayStartAppending(struct ArrowArray* array)
 }
 
 static inline ArrowErrorCode ArrowArrayShrinkToFit(struct ArrowArray* array) {
-  for (int64_t i = 0; i < 3; i++) {
+  for (int64_t i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     struct ArrowBuffer* buffer = ArrowArrayBuffer(array, i);
     NANOARROW_RETURN_NOT_OK(ArrowBufferResize(buffer, buffer->size_bytes, 1));
   }
@@ -278,7 +278,7 @@ static inline ArrowErrorCode _ArrowArrayAppendEmptyInternal(struct ArrowArray* a
   struct ArrowBuffer* buffer;
   int64_t size_bytes;
 
-  for (int i = 0; i < 3; i++) {
+  for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) {
     buffer = ArrowArrayBuffer(array, i);
     size_bytes = private_data->layout.element_size_bits[i] / 8;
 

diff --git a/src/nanoarrow/array_test.cc b/src/nanoarrow/array_test.cc
@@ -1589,12 +1589,12 @@ TEST(ArrayTest, ArrayViewTestBasic) {
   // Expect error for bad offset + length
   array.length = -1;
   EXPECT_EQ(ArrowArrayViewSetArray(&array_view, &array, &error), EINVAL);
-  EXPECT_STREQ(error.message, "Expected array length >= 0 but found array length of -1");
+  EXPECT_STREQ(error.message, "Expected length >= 0 but found length -1");
   array.length = 3;
 
   array.offset = -1;
   EXPECT_EQ(ArrowArrayViewSetArray(&array_view, &array, &error), EINVAL);
-  EXPECT_STREQ(error.message, "Expected array offset >= 0 but found array offset of -1");
+  EXPECT_STREQ(error.message, "Expected offset >= 0 but found offset -1");
   array.offset = 0;
 
   // Expect error for the wrong number of buffers