Feature: cleanup the memory aspects

Mojo has started enforcing its memory model. Use this as an opportunity
to cleanup the design of the various types and their relationships.

Author: winding-lines

firebolt/MEMORY.md

@@ -0,0 +1,32 @@
+Starting with the Sept 2025 version of Mojo the compiler is starting to enforce
+[lifetimes](https://docs.modular.com/mojo/manual/values/lifetimes/). This document proposes an approach at using memory.
+
+1. Low level memory format
+
+Apache Arrow defines a columnar memory [format](https://arrow.apache.org/docs/format/Columnar.html) that can
+be accessed in many languages, including Python. Firebolt defines
+an API to access this format in Mojo. One of the goals is to allow high performance
+integration between Python and Mojo when it comes to process vast amounts of data.
+
+1. ArrayData owns the data
+
+In Firebolt ArrayData is the low level API that will access the Arrow memory block.
+
+As such it should own the data type, bitmap, buffers and children.
+
+2. Typed arrays own ArrayData
+
+The next level in the API are the typed arrays: PrimitiveArray, ListArray, StructArray and so on.
+
+When constructing a typed array from an ArrowData the typed array will own the ArrayData. 
+
+The typed arrays provide convenient accessors into the ArrayData. For example PrimitiveArray
+provides a `bitmap` and a `buffer`. Since Mojo doesn't currently provide properties these 
+helper accessors will be implemented as functions.
+
+3. The Array trait
+
+All of the typed Arrays are expected to implement the Array trait by providing 2 methods:
+
+- `fn take_data(deinit self) -> ArrayData` creates a standalone ArrayData by destroying the self.
+- `fn data(self) -> ref [self] ArrayData` access a read only copy of the ArrayGata in the typed array.

firebolt/arrays/base.mojo

@@ -4,8 +4,19 @@ from sys.info import sizeof
 
 
 trait Array(Movable, Representable, Sized, Stringable, Writable):
-    fn as_data(self) -> ArrayData:
-        ...
+    fn take_data(deinit self) -> ArrayData:
+        """Construct an ArrayData by consuming self."""
+        pass
+
+    fn as_data[
+        self_origin: ImmutableOrigin
+    ](ref [self_origin]self) -> UnsafePointer[ArrayData, mut=False]:
+        """Return a read only reference to the ArrayData wrapped by self.
+
+        Note that ideally the output type would be `ref [self_origin] ArrayData` but this is not supported yet.
+        https://forum.modular.com/t/how-to-mark-a-trait-as-applying-to-not-register-passable/2265/6?u=mseritan
+        """
+        pass
 
 
 @fieldwise_init
@@ -22,47 +33,55 @@ struct ArrayData(Copyable, Movable, Representable, Stringable, Writable):
     var children: List[ArcPointer[ArrayData]]
     var offset: Int
 
+    fn __copyinit__(out self, existing: Self):
+        self.dtype = existing.dtype.copy()
+        self.length = existing.length
+        self.bitmap = existing.bitmap
+        self.buffers = existing.buffers.copy()
+        self.children = existing.children.copy()
+        self.offset = existing.offset
+
     fn is_valid(self, index: Int) -> Bool:
         return self.bitmap[].unsafe_get(index + self.offset)
 
-    fn as_primitive[T: DataType](self) raises -> PrimitiveArray[T]:
-        return PrimitiveArray[T](self)
+    fn as_primitive[T: DataType](var self) raises -> PrimitiveArray[T]:
+        return PrimitiveArray[T](self^)
 
-    fn as_int8(self) raises -> Int8Array:
-        return Int8Array(self)
+    fn as_int8(var self) raises -> Int8Array:
+        return Int8Array(self^)
 
-    fn as_int16(self) raises -> Int16Array:
-        return Int16Array(self)
+    fn as_int16(var self) raises -> Int16Array:
+        return Int16Array(self^)
 
-    fn as_int32(self) raises -> Int32Array:
-        return Int32Array(self)
+    fn as_int32(var self) raises -> Int32Array:
+        return Int32Array(self^)
 
-    fn as_int64(self) raises -> Int64Array:
-        return Int64Array(self)
+    fn as_int64(var self) raises -> Int64Array:
+        return Int64Array(self^)
 
-    fn as_uint8(self) raises -> UInt8Array:
-        return UInt8Array(self)
+    fn as_uint8(var self) raises -> UInt8Array:
+        return UInt8Array(self^)
 
-    fn as_uint16(self) raises -> UInt16Array:
-        return UInt16Array(self)
+    fn as_uint16(var self) raises -> UInt16Array:
+        return UInt16Array(self^)
 
-    fn as_uint32(self) raises -> UInt32Array:
-        return UInt32Array(self)
+    fn as_uint32(var self) raises -> UInt32Array:
+        return UInt32Array(self^)
 
-    fn as_uint64(self) raises -> UInt64Array:
-        return UInt64Array(self)
+    fn as_uint64(var self) raises -> UInt64Array:
+        return UInt64Array(self^)
 
-    fn as_float32(self) raises -> Float32Array:
-        return Float32Array(self)
+    fn as_float32(var self) raises -> Float32Array:
+        return Float32Array(self^)
 
-    fn as_float64(self) raises -> Float64Array:
-        return Float64Array(self)
+    fn as_float64(var self) raises -> Float64Array:
+        return Float64Array(self^)
 
-    fn as_string(self) raises -> StringArray:
-        return StringArray(self)
+    fn as_string(var self) raises -> StringArray:
+        return StringArray(self^)
 
-    fn as_list(self) raises -> ListArray:
-        return ListArray(self)
+    fn as_list(var self) raises -> ListArray:
+        return ListArray(self^)
 
     fn _dynamic_write[W: Writer](self, index: Int, mut writer: W):
         """Write to the given stream dispatching on the dtype."""
@@ -117,3 +136,20 @@ struct ArrayData(Copyable, Movable, Representable, Stringable, Writable):
 
     fn __repr__(self) -> String:
         return String.write(self)
+
+    fn append_to_array(
+        deinit self: ArrayData, mut combined: ArrayData, start: Int
+    ) -> Int:
+        """Append the content self to the combined array, consumes self.
+
+        Args:
+            combined: Array to append to.
+            start: Position where to append.
+
+        Returns:
+            The new start position.
+        """
+        combined.bitmap[].extend(self.bitmap[], start, self.length)
+        combined.buffers.extend(self.buffers^)
+        combined.children.extend(self.children^)
+        return start + self.length

firebolt/arrays/binary.mojo

@@ -6,24 +6,27 @@ from ..dtypes import *
 
 struct StringArray(Array):
     var data: ArrayData
-    var bitmap: ArcPointer[Bitmap]
-    var offsets: ArcPointer[Buffer]
-    var values: ArcPointer[Buffer]
     var capacity: Int
 
-    fn __init__(out self, data: ArrayData) raises:
+    fn __init__(out self, var data: ArrayData) raises:
         if data.dtype != string:
             raise Error(
                 "Unexpected dtype '{}' instead of 'string'.".format(data.dtype)
             )
         elif len(data.buffers) != 2:
             raise Error("StringArray requires exactly two buffers")
 
-        self.data = data
-        self.bitmap = data.bitmap
-        self.offsets = data.buffers[0]
-        self.values = data.buffers[1]
         self.capacity = data.length
+        self.data = data^
+
+    fn bitmap(self) -> ref [self.data.bitmap] ArcPointer[Bitmap]:
+        return self.data.bitmap
+
+    fn offsets(self) -> ref [self.data.buffers] ArcPointer[Buffer]:
+        return self.data.buffers[0]
+
+    fn values(self) -> ref [self.data.buffers] ArcPointer[Buffer]:
+        return self.data.buffers[1]
 
     fn __init__(out self, capacity: Int = 0):
         var bitmap = Bitmap.alloc(capacity)
@@ -33,68 +36,67 @@ struct StringArray(Array):
         offsets.unsafe_set[DType.uint32](0, 0)
 
         self.capacity = capacity
-        self.bitmap = ArcPointer(bitmap^)
-        self.offsets = ArcPointer(offsets^)
-        self.values = ArcPointer(values^)
         self.data = ArrayData(
             dtype=string,
             length=0,
-            bitmap=self.bitmap,
-            buffers=List(self.offsets, self.values),
+            bitmap=ArcPointer(bitmap^),
+            buffers=List(ArcPointer(offsets^), ArcPointer(values^)),
             children=List[ArcPointer[ArrayData]](),
             offset=0,
         )
 
     fn __moveinit__(out self, deinit existing: Self):
         self.data = existing.data^
-        self.bitmap = existing.bitmap^
-        self.offsets = existing.offsets^
-        self.values = existing.values^
         self.capacity = existing.capacity
 
     fn __len__(self) -> Int:
         return self.data.length
 
-    fn as_data(self) -> ArrayData:
-        return self.data
+    fn as_data[
+        self_origin: ImmutableOrigin
+    ](ref [self_origin]self) -> UnsafePointer[ArrayData, mut=False]:
+        return UnsafePointer(to=self.data)
+
+    fn take_data(deinit self) -> ArrayData:
+        return self.data^
 
     fn grow(mut self, capacity: Int):
-        self.bitmap[].grow(capacity)
-        self.offsets[].grow[DType.uint32](capacity + 1)
+        self.bitmap()[].grow(capacity)
+        self.offsets()[].grow[DType.uint32](capacity + 1)
         self.capacity = capacity
 
     # fn shrink_to_fit(out self):
 
     fn is_valid(self, index: Int) -> Bool:
-        return self.bitmap[].unsafe_get(index)
+        return self.bitmap()[].unsafe_get(index)
 
     fn unsafe_append(mut self, value: String):
         # todo(kszucs): use unsafe set
         var index = self.data.length
-        var last_offset = self.offsets[].unsafe_get[DType.uint32](index)
+        var last_offset = self.offsets()[].unsafe_get[DType.uint32](index)
         var next_offset = last_offset + len(value)
         self.data.length += 1
-        self.bitmap[].unsafe_set(index, True)
-        self.offsets[].unsafe_set[DType.uint32](index + 1, next_offset)
-        self.values[].grow[DType.uint8](next_offset)
-        var dst_address = self.values[].get_ptr_at(Int(last_offset))
+        self.bitmap()[].unsafe_set(index, True)
+        self.offsets()[].unsafe_set[DType.uint32](index + 1, next_offset)
+        self.values()[].grow[DType.uint8](next_offset)
+        var dst_address = self.values()[].get_ptr_at(Int(last_offset))
         var src_address = value.unsafe_ptr()
         memcpy(dst_address, src_address, len(value))
 
     fn unsafe_get(self, index: UInt) -> StringSlice[__origin_of(self)]:
-        var start_offset = self.offsets[].unsafe_get[DType.uint32](
+        var start_offset = self.offsets()[].unsafe_get[DType.uint32](
             index + self.data.offset
         )
-        var end_offset = self.offsets[].unsafe_get[DType.uint32](
+        var end_offset = self.offsets()[].unsafe_get[DType.uint32](
             index + 1 + self.data.offset
         )
-        var address = self.values[].get_ptr_at(Int(start_offset))
+        var address = self.values()[].get_ptr_at(Int(start_offset))
         var length = UInt(Int(end_offset - start_offset))
         return StringSlice[__origin_of(self)](ptr=address, length=length)
 
     fn unsafe_set(mut self, index: Int, value: String) raises:
-        var start_offset = self.offsets[].unsafe_get[DType.int32](index)
-        var end_offset = self.offsets[].unsafe_get[DType.int32](index + 1)
+        var start_offset = self.offsets()[].unsafe_get[DType.int32](index)
+        var end_offset = self.offsets()[].unsafe_get[DType.int32](index + 1)
         var length = Int(end_offset - start_offset)
 
         if length != len(value):
@@ -103,7 +105,7 @@ struct StringArray(Array):
                 " length"
             )
 
-        var dst_address = self.values[].get_ptr_at(Int(start_offset))
+        var dst_address = self.values()[].get_ptr_at(Int(start_offset))
         var src_address = value.unsafe_ptr()
         memcpy(dst_address, src_address, length)
 

firebolt/arrays/chunked_array.mojo

@@ -21,9 +21,9 @@ struct ChunkedArray:
             total_length += chunk.length
         self.length = total_length
 
-    fn __init__(out self, dtype: DataType, chunks: List[ArrayData]):
-        self.dtype = dtype
-        self.chunks = chunks
+    fn __init__(out self, var dtype: DataType, var chunks: List[ArrayData]):
+        self.dtype = dtype^
+        self.chunks = chunks^
         self.length = 0
         self._compute_length()
 
@@ -38,21 +38,19 @@ struct ChunkedArray:
         """
         return self.chunks[index]
 
-    fn combine_chunks(self) -> ArrayData:
+    fn combine_chunks(var self, out combined: ArrayData):
         """Combines all chunks into a single array."""
         var bitmap = ArcPointer(Bitmap.alloc(self.length))
-        var combined = ArrayData(
-            dtype=self.dtype,
+        combined = ArrayData(
+            dtype=self.dtype.copy(),
             length=self.length,
             bitmap=bitmap,
             buffers=List[ArcPointer[Buffer]](),
             children=List[ArcPointer[ArrayData]](),
             offset=0,
         )
         var start = 0
-        for chunk in self.chunks:
-            combined.bitmap[].extend(chunk.bitmap[], start, chunk.length)
-            start += chunk.length
-            combined.buffers.extend(chunk.buffers)
-            combined.children.extend(chunk.children)
-        return combined
+        while self.chunks:
+            var chunk = self.chunks.pop(0)
+            start += chunk^.append_to_array(combined, start)
+        return combined^