Make sure we have properly aligned arrow buffers after deserialization (#9179)

Author: nepal

ydb/library/yql/minikql/computation/mkql_block_transport.cpp

@@ -3,6 +3,7 @@
 
 #include <ydb/library/yql/minikql/mkql_type_builder.h>
 #include <ydb/library/yql/public/udf/arrow/block_reader.h>
+#include <ydb/library/yql/public/udf/arrow/memory_pool.h>
 #include <ydb/library/yql/utils/rope_over_buffer.h>
 #include <ydb/library/yql/utils/yql_panic.h>
 
@@ -30,22 +31,35 @@ std::shared_ptr<arrow::Buffer> MakeEmptyBuffer() {
     return std::make_shared<arrow::Buffer>(nullptr, 0);
 }
 
+bool HasArrrowAlignment(const void* buf) {
+    return AlignUp(buf, NYql::NUdf::ArrowMemoryAlignment) == buf;
+}
+
 std::shared_ptr<arrow::Buffer> MakeZeroBuffer(size_t byteLen) {
+    using namespace NYql::NUdf;
+    if (!byteLen) {
+        return MakeEmptyBuffer();
+    }
+
     constexpr size_t NullWordCount = (MaxBlockSizeInBytes + sizeof(ui64) - 1) / sizeof(ui64);
-    static ui64 nulls[NullWordCount] = { 0 };
-    if (byteLen <= sizeof(nulls)) {
-        return std::make_shared<arrow::Buffer>(reinterpret_cast<const ui8*>(nulls), byteLen);
+    constexpr size_t ExtraAlignWords = (ArrowMemoryAlignment > sizeof(ui64)) ? (ArrowMemoryAlignment / sizeof(ui64) - 1) : 0;
+    static const ui64 nulls[NullWordCount + ExtraAlignWords] = { 0 };
+
+    // round all buffer length to 64 bytes
+    size_t capacity = AlignUp(byteLen, size_t(64));
+    if (capacity <= NullWordCount * sizeof(ui64)) {
+        return std::make_shared<arrow::Buffer>(AlignUp(reinterpret_cast<const ui8*>(nulls), ArrowMemoryAlignment), byteLen);
     }
 
-    size_t wordCount = (byteLen + sizeof(ui64) - 1) / sizeof(ui64);
-    std::shared_ptr<ui64[]> buf(new ui64[wordCount]);
-    std::fill(buf.get(), buf.get() + wordCount, 0);
-    return std::make_shared<TOwnedArrowBuffer>(TContiguousSpan{ reinterpret_cast<const char*>(buf.get()), byteLen }, buf);
+    auto result = AllocateResizableBuffer(byteLen, GetYqlMemoryPool());
+    ARROW_OK(result->Resize(byteLen));
+    std::memset(result->mutable_data(), 0, byteLen);
+    return result;
 }
 
 std::shared_ptr<arrow::Buffer> MakeZeroBitmap(size_t bitCount) {
     // align up 8 byte boundary
-    size_t byteCount = ((bitCount + 63u) & ~size_t(63u)) >> 3;
+    size_t byteCount = AlignUp(bitCount, size_t(64)) >> 3;
     return MakeZeroBuffer(byteCount);
 }
 
@@ -62,7 +76,7 @@ void StoreNullsSizes(const arrow::ArrayData& data, const IBlockSerializer::TMeta
     }
 
     const ui64 desiredOffset = data.offset % 8;
-    size_t nullBytes = (((size_t)data.length + desiredOffset + 7) & ~7ull) >> 3;
+    size_t nullBytes = AlignUp((size_t)data.length + desiredOffset, size_t(8)) >> 3;
     metaSink(nullBytes);
 }
 
@@ -77,7 +91,7 @@ void StoreNulls(const arrow::ArrayData& data, TRope& dst) {
         return;
     }
     const ui64 desiredOffset = data.offset % 8;
-    size_t nullBytes = (((size_t)data.length + desiredOffset + 7) & ~7ull) >> 3;
+    size_t nullBytes = AlignUp((size_t)data.length + desiredOffset, size_t(8)) >> 3;
     YQL_ENSURE(desiredOffset <= (size_t)data.offset);
     YQL_ENSURE((data.offset - desiredOffset) % 8 == 0);
     const char* nulls = data.GetValues<char>(0, 0) + (data.offset - desiredOffset) / 8;
@@ -90,17 +104,26 @@ void LoadBufferSize(const IBlockDeserializer::TMetadataSource& metaSource, TMayb
 }
 
 std::shared_ptr<arrow::Buffer> LoadBuffer(TRope& source, TMaybe<ui64> size) {
+    using namespace NYql::NUdf;
     YQL_ENSURE(size.Defined(), "Buffer size is not loaded");
     if (!*size) {
-        return std::make_shared<arrow::Buffer>(nullptr, 0);
+        return MakeEmptyBuffer();
     }
 
     YQL_ENSURE(source.size() >= *size, "Premature end of data");
     auto owner = std::make_shared<TRope>(source.Begin(), source.Begin() + *size);
     source.EraseFront(*size);
 
     owner->Compact();
-    return std::make_shared<TOwnedArrowBuffer>(owner->GetContiguousSpan(), owner);
+    auto span = owner->GetContiguousSpan();
+    if (HasArrrowAlignment(span.Data())) {
+        return std::make_shared<TOwnedArrowBuffer>(span, owner);
+    }
+
+    auto result = AllocateResizableBuffer(span.Size(), NYql::NUdf::GetYqlMemoryPool());
+    ARROW_OK(result->Resize((int64_t)span.Size()));
+    std::memcpy(result->mutable_data(), span.Data(), span.Size());
+    return result;
 }
 
 std::shared_ptr<arrow::Buffer> LoadNullsBitmap(TRope& source, TMaybe<ui64> nullCount, TMaybe<ui64> bitmapSize) {

ydb/library/yql/public/udf/arrow/memory_pool.h

@@ -6,6 +6,9 @@
 namespace NYql {
 namespace NUdf {
 
+constexpr size_t ArrowMemoryAlignment = 64;
+static_assert((ArrowMemoryAlignment & (ArrowMemoryAlignment - 1)) == 0, "ArrowMemoryAlignment should be power of 2");
+
 #if UDF_ABI_COMPATIBILITY_VERSION_CURRENT >= UDF_ABI_COMPATIBILITY_VERSION(2, 37)
 arrow::MemoryPool* GetYqlMemoryPool();
 #else