From a7a81b6044fd3a71dc9797a1b09f9a68350630ed Mon Sep 17 00:00:00 2001
From: David Li <li.davidm96@gmail.com>
Date: Tue, 18 Feb 2025 01:39:53 -0500
Subject: [PATCH] Add unit tests

---
 arrow/memory/allocator.go                  |  1 +
 arrow/memory/buffer_test.go                |  3 +-
 arrow/memory/mallocator/mallocator_test.go | 47 ++++++++++++++++++++++
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/arrow/memory/allocator.go b/arrow/memory/allocator.go
index 8c0d06a4..9a3a21ce 100644
--- a/arrow/memory/allocator.go
+++ b/arrow/memory/allocator.go
@@ -34,6 +34,7 @@ type AlignedAllocator interface {
 	ReallocateAligned(size int, buf []byte, alloc []byte) (newBuf []byte, newAlloc []byte)
 }
 
+// MakeAlignedAllocator makes an AlignedAllocator for the given Allocator.  If the Allocator already implements AlignedAllocator, it is returned as-is.  There is usually no need to call this directly as Buffer will call this for you.
 func MakeAlignedAllocator(mem Allocator) AlignedAllocator {
 	if aligned, ok := mem.(AlignedAllocator); ok {
 		return aligned
diff --git a/arrow/memory/buffer_test.go b/arrow/memory/buffer_test.go
index 2a7c251b..f95e74ea 100644
--- a/arrow/memory/buffer_test.go
+++ b/arrow/memory/buffer_test.go
@@ -18,6 +18,7 @@ package memory_test
 
 import (
 	"testing"
+	"unsafe"
 
 	"github.com/apache/arrow-go/v18/arrow/memory"
 	"github.com/stretchr/testify/assert"
@@ -32,10 +33,10 @@ func TestNewResizableBuffer(t *testing.T) {
 
 	exp := 10
 	buf.Resize(exp)
-	// assert.Equal(t, 10, mem.CurrentAlloc())
 	assert.NotNil(t, buf.Bytes())
 	assert.Equal(t, exp, len(buf.Bytes()))
 	assert.Equal(t, exp, buf.Len())
+	assert.Equal(t, uintptr(0), uintptr(unsafe.Pointer(&buf.Buf()[0])) % 64)
 
 	buf.Release() // refCount == 1
 	assert.NotNil(t, buf.Bytes())
diff --git a/arrow/memory/mallocator/mallocator_test.go b/arrow/memory/mallocator/mallocator_test.go
index 6ad47ecd..647b7535 100644
--- a/arrow/memory/mallocator/mallocator_test.go
+++ b/arrow/memory/mallocator/mallocator_test.go
@@ -22,7 +22,9 @@ package mallocator_test
 import (
 	"fmt"
 	"testing"
+	"unsafe"
 
+	"github.com/apache/arrow-go/v18/arrow/memory"
 	"github.com/apache/arrow-go/v18/arrow/memory/mallocator"
 	"github.com/stretchr/testify/assert"
 )
@@ -125,3 +127,48 @@ func TestMallocatorReallocateNegative(t *testing.T) {
 		a.Reallocate(-1, buf)
 	})
 }
+
+func TestMallocatorAligned(t *testing.T) {
+	a := mallocator.NewMallocator()
+
+	aligned := memory.MakeAlignedAllocator(a)
+	assert.NotSame(t, a, aligned)
+
+	assert.Equal(t, int64(0), a.AllocatedBytes())
+
+	buf, alloc := aligned.AllocateAligned(64)
+
+	assert.GreaterOrEqual(t, int64(128), a.AllocatedBytes())
+
+	assert.Equal(t, uintptr(0), uintptr(unsafe.Pointer(&buf[0])) % 64)
+
+	buf, alloc = aligned.ReallocateAligned(96, buf, alloc)
+	assert.Equal(t, uintptr(0), uintptr(unsafe.Pointer(&buf[0])) % 64)
+
+	aligned.Free(alloc)
+	assert.Equal(t, int64(0), a.AllocatedBytes())
+}
+
+func TestMallocatorAlignedBuffer(t *testing.T) {
+	a := mallocator.NewMallocator()
+
+	buf := memory.NewResizableBuffer(a)
+	defer buf.Release()
+	assert.Equal(t, int64(0), a.AllocatedBytes())
+
+	buf.Reserve(20)
+	// 20 -> round up to 64 -> add 64 bytes padding
+	assert.GreaterOrEqual(t, int64(128), a.AllocatedBytes())
+	addr := uintptr(unsafe.Pointer(&buf.Buf()[0]))
+	assert.Equal(t, uintptr(0), addr % 64)
+
+	buf.Reserve(40)
+	// There should have been no reallocation
+	assert.Equal(t, addr, uintptr(unsafe.Pointer(&buf.Buf()[0])))
+
+	buf.Reserve(100)
+	// Now the buffer is reallocated, make sure the new address is also aligned
+	assert.NotEqual(t, addr, uintptr(unsafe.Pointer(&buf.Buf()[0])))
+	addr = uintptr(unsafe.Pointer(&buf.Buf()[0]))
+	assert.Equal(t, uintptr(0), addr % 64)
+}