GH-32276: [C++][FlightRPC] Align RecordBatch buffers given to IPC

[EnricoMi]

Rationale for this change

Data retrieved via IPC is expected to provide memory-aligned arrays, but data retrieved via C++ Flight client is mis-aligned. Datafusion (Rust), which requires proper alignment, cannot handle such data: #43552.

What changes are included in this PR?

This aligns RecordBatch array buffers decoded by IPC if mis-aligned according to the data type byte width.
Implementation mirrors that of align_buffers in arrow-rs (apache/arrow-rs#4681).

Are these changes tested?

Configuration flag tested in unit test.
Manually end-to-end tested that memory alignment fixes issue with reproduction code provided in #43552.

Are there any user-facing changes?

Memory alignment is checked and fixed by default. This is configurable via IpcReadOptions.ensure_memory_alignment.

• GitHub Issue: [C++][FlightRPC] Flight generates misaligned buffers #32276

[EnricoMi]

@pitrou do you think this fix is viable?

[lidavidm]

Thanks for picking this up. This seems reasonable to me at a quick glance.

[lidavidm]

nit: put this include with the rest of the Arrow includes (and use quotes to be consistent)

[EnricoMi]

done

[lidavidm]

nit: you could iterate with for (auto& child : child_data) and avoid the explicit index?

[EnricoMi]

much better!

[lidavidm]

Where did this come from?

[EnricoMi]

Fixed.

[EnricoMi]

While attempting to write some unit tests I found there is util::EnsureAlignment:

arrow/cpp/src/arrow/util/align_util.cc

Lines 169 to 205 in e62fbaa

	Result<std::shared_ptr<ArrayData>> EnsureAlignment(std::shared_ptr<ArrayData> array_data,
	int64_t alignment,
	MemoryPool* memory_pool) {
	if (!CheckAlignment(*array_data, alignment)) {
	std::vector<std::shared_ptr<Buffer>> buffers = array_data->buffers;
	Type::type type_id = GetTypeForBuffers(*array_data);
	for (size_t i = 0; i < buffers.size(); ++i) {
	if (buffers[i]) {
	int64_t expected_alignment = alignment;
	if (alignment == kValueAlignment) {
	expected_alignment =
	RequiredValueAlignmentForBuffer(type_id, static_cast<int>(i));
	}
	ARROW_ASSIGN_OR_RAISE(
	buffers[i],
	EnsureAlignment(std::move(buffers[i]), expected_alignment, memory_pool));
	}
	}
	for (auto& it : array_data->child_data) {
	ARROW_ASSIGN_OR_RAISE(it, EnsureAlignment(std::move(it), alignment, memory_pool));
	}
	if (array_data->type->id() == Type::DICTIONARY) {
	ARROW_ASSIGN_OR_RAISE(
	array_data->dictionary,
	EnsureAlignment(std::move(array_data->dictionary), alignment, memory_pool));
	}
	auto new_array_data = ArrayData::Make(
	array_data->type, array_data->length, std::move(buffers), array_data->child_data,
	array_data->dictionary, array_data->GetNullCount(), array_data->offset);
	return new_array_data;
	} else {
	return array_data;
	}
	}

I will try to reuse that method rather than re-implementing it. There is also test infrastructure for misaligned array data.

[lidavidm]

Can you rebase? Tests appear to be failing

[lidavidm]

Can we use the memory pool in context.options.memory_pool?

[EnricoMi]

Sure! Ideally we should use the buffer's memory manager rather than the default CPU manager:

arrow/cpp/src/arrow/memory_pool.cc

Lines 907 to 916 in 5ad0b3e

	static std::unique_ptr<PoolBuffer> MakeUnique(MemoryPool* pool, int64_t alignment) {
	std::shared_ptr<MemoryManager> mm;
	if (pool == nullptr) {
	pool = default_memory_pool();
	mm = default_cpu_memory_manager();
	} else {
	mm = CPUDevice::memory_manager(pool);
	}
	return std::make_unique<PoolBuffer>(std::move(mm), pool, alignment);
	}

[EnricoMi]

Test arrow-ipc-read-write-test fails with a SIGSEGV in line 45 of

arrow/cpp/src/arrow/util/align_util.cc

Lines 44 to 52 in bcb4653

	Type::type GetTypeForBuffers(const ArrayData& array) {
	Type::type type_id = array.type->storage_id();
	if (type_id == Type::DICTIONARY) {
	return ::arrow::internal::checked_pointer_cast<DictionaryType>(array.type)
	->index_type()
	->id();
	}
	return type_id;
	}

https://github.com/apache/arrow/actions/runs/11462607112/job/31894398411?pr=44279#step:13:1548

That test complains a lot about /home/enrico/Work/git/arrow/cpp/src/arrow/status.cc:152: Invalid: RequiredValueAlignmentForBuffer called with invalid type id 29 due to RequiredValueAlignmentForBuffer being called for DICTIONARY data type in line 62:

arrow/cpp/src/arrow/util/align_util.cc

Lines 56 to 76 in bcb4653

	bool CheckSelfAlignment(const ArrayData& array, int64_t alignment) {
	if (alignment == kValueAlignment) {
	Type::type type_id = GetTypeForBuffers(array);
	for (std::size_t i = 0; i < array.buffers.size(); i++) {
	if (array.buffers[i]) {
	int expected_alignment =
	RequiredValueAlignmentForBuffer(type_id, static_cast<int>(i));
	if (!CheckAlignment(*array.buffers[i], expected_alignment)) {
	return false;
	}
	}
	}
	} else {
	for (const auto& buffer : array.buffers) {
	if (buffer) {
	if (!CheckAlignment(*buffer, alignment)) return false;
	}
	}
	}
	return true;
	}

Looks like util::EnsureAlignment and util::CheckAlignment is not quite ready for this use case.

[EnricoMi]

@westonpace @sanjibansg
Any pointers in which situations array.type can be null (causing GetDtypeForBuffers to segfault)?
Any rational on why RequiredValueAlignmentForBuffer does not support DICTIONARY data type?
#44279 (comment)

[lidavidm]

I think this is reasonable and we just need to update EnsureAlignment to cover DICTIONARY

[EnricoMi]

Problem was that tests define an ExtensionType with storage type DICTIONARY, which caused an invalid cast, returning an empty pointer that was dereferenced. That has been fixes in a1ad4da.

This means that user code that defines ExtensionTypes with other storage types than DICTIONARY may see unexpected behaviour. What would be a good default behaviour in GetTypeForBuffer? Return Type::UINT8 to not realign those unknown data types? Or how can we tell from an unknown ExtensionTypes what to return by GetTypeForBuffers?

[lidavidm]

I'm not sure I follow. An extension type should be treated the same as its storage type. I think the problem with the current code is that it blindly casts the array type to the storage type, not accounting for the fact that the point of the storage type is to allow the extension type.

[EnricoMi]

Any non-dictionary type (including extension types) therefore should be covered with array.type->storage_id().
Any dictionary extension types are expected to provide a DictionaryType via array.type->storage_type(), so they are covered with array.type->storage_type()->index_type()->id().

Than that fix should be sound and safe.

arrow/cpp/src/arrow/util/align_util.cc

Lines 47 to 59 in 7a01029

	Type::type GetTypeForBuffers(const ArrayData& array) {
	Type::type type_id = array.type->storage_id();
	if (type_id == Type::DICTIONARY) {
	// return index type id, provided by the DictionaryType array.type or
	// array.type->storage_type() if array.type is an ExtensionType
	std::shared_ptr<DataType> dict_type = array.type;
	if (array.type->id() == Type::EXTENSION) {
	dict_type = checked_pointer_cast<ExtensionType>(array.type)->storage_type();
	}
	return checked_pointer_cast<DictionaryType>(dict_type)->index_type()->id();
	}
	return type_id;
	}