Feature: create buffers from variadic arguments (#40)

The nicer API simplifies some of the tests.

Author: winding-lines

firebolt/arrays/base.mojo

@@ -33,6 +33,23 @@ struct ArrayData(Copyable, Movable, Representable, Stringable, Writable):
     var children: List[ArcPointer[ArrayData]]
     var offset: Int
 
+    @staticmethod
+    fn from_buffer[
+        dtype: DataType
+    ](var buffer: Buffer, length: Int) -> ArrayData:
+        """Build an ArrayData from a buffer where all the values are not null.
+        """
+        var bitmap = Bitmap.alloc(length)
+        bitmap.unsafe_range_set(0, length, True)
+        return ArrayData(
+            dtype=materialize[dtype](),
+            length=length,
+            bitmap=ArcPointer(bitmap^),
+            buffers=List(ArcPointer(buffer^)),
+            children=List[ArcPointer[ArrayData]](),
+            offset=0,
+        )
+
     fn __copyinit__(out self, existing: Self):
         self.dtype = existing.dtype.copy()
         self.length = existing.length

firebolt/arrays/tests/test_nested.mojo

@@ -29,14 +29,9 @@ fn build_list_of_list[data_type: DataType]() raises -> ListArray:
     )
 
     # Define the PrimitiveArrays.
-    var value_offset = ArcPointer(Buffer.alloc(7))
-    value_offset[].unsafe_set[DType.int32](0, 0)
-    value_offset[].unsafe_set[DType.int32](1, 2)
-    value_offset[].unsafe_set[DType.int32](2, 4)
-    value_offset[].unsafe_set[DType.int32](3, 7)
-    value_offset[].unsafe_set[DType.int32](4, 7)
-    value_offset[].unsafe_set[DType.int32](5, 8)
-    value_offset[].unsafe_set[DType.int32](6, 10)
+    var value_offset = ArcPointer(
+        Buffer.from_values[DType.int32](0, 2, 4, 7, 7, 8, 10)
+    )
 
     var list_bitmap = ArcPointer(Bitmap.alloc(6))
     list_bitmap[].unsafe_range_set(0, 6, True)
@@ -51,11 +46,7 @@ fn build_list_of_list[data_type: DataType]() raises -> ListArray:
     )
 
     # Now define the master array data.
-    var top_offsets = Buffer.alloc(4)
-    top_offsets.unsafe_set[DType.int32](0, 0)
-    top_offsets.unsafe_set[DType.int32](1, 2)
-    top_offsets.unsafe_set[DType.int32](2, 5)
-    top_offsets.unsafe_set[DType.int32](3, 6)
+    var top_offsets = Buffer.from_values[DType.int32](0, 2, 5, 6)
     var top_bitmap = ArcPointer(Bitmap.alloc(4))
     top_bitmap[].unsafe_range_set(0, 4, True)
     return ListArray(
@@ -71,10 +62,11 @@ fn build_list_of_list[data_type: DataType]() raises -> ListArray:
 
 
 def test_list_int_array():
-    var ints = Int64Array(capacity=3)
-    ints.append(1)
-    ints.append(2)
-    ints.append(3)
+    var ints = Int64Array(
+        ArrayData.from_buffer[int64](
+            Buffer.from_values[DType.int64](1, 2, 3), 3
+        )
+    )
     var lists = ListArray(ints^)
     assert_equal(lists.data.dtype, list_(materialize[int64]()))
 

firebolt/buffers.mojo

@@ -66,6 +66,16 @@ struct Buffer(Movable):
         memset_zero(ptr.bitcast[UInt8](), size)
         return Buffer(ptr, size)
 
+    @staticmethod
+    fn from_values[dtype: DType](*values: Scalar[dtype]) -> Buffer:
+        """Build a buffer from a list of values."""
+        var buffer = Self.alloc[dtype](len(values))
+
+        for i in range(len(values)):
+            buffer.unsafe_set[dtype](i, values[i])
+
+        return buffer^
+
     @staticmethod
     fn view[
         I: Intable, //

firebolt/tests/test_buffers.mojo

@@ -55,6 +55,14 @@ def test_buffer_set_get():
     assert_equal(buf.unsafe_get[DType.uint16](1), 44)
 
 
+def test_buffer_from_values():
+    var buf = Buffer.from_values[DType.int64](-3, 9, 81)
+
+    assert_equal(buf.unsafe_get[DType.int64](0), -3)
+    assert_equal(buf.unsafe_get[DType.int64](1), 9)
+    assert_equal(buf.unsafe_get[DType.int64](2), 81)
+
+
 def test_buffer_swap():
     var one = Buffer.alloc(10)
     one.unsafe_set(0, 111)