GH-32276: [C++][FlightRPC] Add option to align RecordBatch buffers given to IPC reader by EnricoMi · Pull Request #44279 · apache/arrow

EnricoMi · 2024-10-01T08:37:03Z

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 data type-specific alignment, cannot handle such data: #43552.
https://arrow.apache.org/docs/format/Columnar.html#buffer-alignment-and-padding

What changes are included in this PR?

This adds option arrow::ipc::IpcReadOptions.ensure_alignment of type arrow::ipc::Alignment to configure how RecordBatch array buffers decoded by IPC are realigned. It supports no realignment (default), data type-specific alignment and 64-byte alignment.

Implementation mirrors that of align_buffers in arrow-rs (apache/arrow-rs#4681).

Are these changes tested?

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

Are there any user-facing changes?

Adds option IpcReadOptions.ensure_alignment and enum type Alignment.

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

EnricoMi · 2024-10-02T05:59:51Z

@pitrou do you think this fix is viable?

lidavidm

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

lidavidm · 2024-10-06T23:40:57Z

cpp/src/arrow/array/data.h

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

lidavidm · 2024-10-06T23:41:28Z

cpp/src/arrow/array/data.h

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

much better!

lidavidm · 2024-10-06T23:42:12Z

python/pyarrow/tests/test_ipc.py

Where did this come from?

EnricoMi · 2024-10-07T12:32:46Z

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 · 2024-10-14T01:53:25Z

cpp/src/arrow/ipc/reader.cc

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

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 · 2024-10-22T15:21:40Z

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 · 2024-10-26T09:40:49Z

@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)

…torage type

Co-authored-by: David Li <li.davidm96@gmail.com>

This reverts commit 069c30f7fb87f16d2d8d236ca2e3da1de19dec30.

Co-authored-by: Antoine Pitrou <pitrou@free.fr>

docs/source/cpp/flight.rst

pitrou · 2025-05-15T15:04:26Z

Thanks @EnricoMi . This looks good to me except for the nit above.

Co-authored-by: Antoine Pitrou <pitrou@free.fr>

lidavidm · 2025-05-19T23:32:12Z

Thanks @EnricoMi!

conbench-apache-arrow · 2025-05-20T07:22:59Z

After merging your PR, Conbench analyzed the 4 benchmarking runs that have been run so far on merge-commit f45594d.

There were no benchmark performance regressions. 🎉

The full Conbench report has more details. It also includes information about 15 possible false positives for unstable benchmarks that are known to sometimes produce them.

…ers given to IPC reader (apache#44279) ### 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 data type-specific alignment, cannot handle such data: apache#43552. https://arrow.apache.org/docs/format/Columnar.html#buffer-alignment-and-padding ### What changes are included in this PR? This adds option `arrow::ipc::IpcReadOptions.ensure_alignment` of type `arrow::ipc::Alignment` to configure how RecordBatch array buffers decoded by IPC are realigned. It supports no realignment (default), data type-specific alignment and 64-byte alignment. Implementation mirrors that of [`align_buffers` in arrow-rs](https://github.com/apache/arrow-rs/blob/3293a8c2f9062fca93bee2210d540a1d25155bf5/arrow-data/src/data.rs#L698-L711) (apache/arrow-rs#4681). ### Are these changes tested? Configuration flag tested in unit test. Integration test with Flight server. Manually end-to-end tested that memory alignment fixes issue with reproduction code provided in apache#43552. ### Are there any user-facing changes? Adds option `IpcReadOptions.ensure_alignment` and enum type `Alignment`. * GitHub Issue: apache#32276 Lead-authored-by: Enrico Minack <github@enrico.minack.dev> Co-authored-by: David Li <li.davidm96@gmail.com> Co-authored-by: Antoine Pitrou <pitrou@free.fr> Signed-off-by: David Li <li.davidm96@gmail.com>

…ers given to IPC reader (apache#44279) (#6) ### 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 data type-specific alignment, cannot handle such data: apache#43552. https://arrow.apache.org/docs/format/Columnar.html#buffer-alignment-and-padding ### What changes are included in this PR? This adds option `arrow::ipc::IpcReadOptions.ensure_alignment` of type `arrow::ipc::Alignment` to configure how RecordBatch array buffers decoded by IPC are realigned. It supports no realignment (default), data type-specific alignment and 64-byte alignment. Implementation mirrors that of [`align_buffers` in arrow-rs](https://github.com/apache/arrow-rs/blob/3293a8c2f9062fca93bee2210d540a1d25155bf5/arrow-data/src/data.rs#L698-L711) (apache/arrow-rs#4681). ### Are these changes tested? Configuration flag tested in unit test. Integration test with Flight server. Manually end-to-end tested that memory alignment fixes issue with reproduction code provided in apache#43552. ### Are there any user-facing changes? Adds option `IpcReadOptions.ensure_alignment` and enum type `Alignment`. * GitHub Issue: apache#32276 Lead-authored-by: Enrico Minack <github@enrico.minack.dev> Signed-off-by: David Li <li.davidm96@gmail.com> Co-authored-by: David Li <li.davidm96@gmail.com> Co-authored-by: Antoine Pitrou <pitrou@free.fr>

EnricoMi changed the title ~~GH-32276: [C++][Flight] Align RecordBatch buffers retrieved via IPC~~ GH-32276: [C++][FlightRPC] Align RecordBatch buffers retrieved via IPC Oct 1, 2024

github-actions bot added the awaiting review Awaiting review label Oct 1, 2024

EnricoMi changed the title ~~GH-32276: [C++][FlightRPC] Align RecordBatch buffers retrieved via IPC~~ GH-32276: [C++][FlightRPC] Align RecordBatch buffers given to IPC Oct 1, 2024

EnricoMi mentioned this pull request Oct 1, 2024

Align buffers from Python (FFI) apache/arrow-rs#6472

Merged

rok requested a review from lidavidm October 2, 2024 10:23

lidavidm reviewed Oct 6, 2024

View reviewed changes

github-actions bot added awaiting changes Awaiting changes and removed awaiting review Awaiting review labels Oct 6, 2024

EnricoMi force-pushed the flight-client-align-buffers branch from 77cc70a to a5d9e2d Compare October 7, 2024 10:57

github-actions bot added Component: C++ Component: Python awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Oct 7, 2024

github-actions bot added the Component: FlightRPC label Oct 7, 2024

lidavidm reviewed Oct 14, 2024

View reviewed changes

github-actions bot added awaiting changes Awaiting changes awaiting review Awaiting review awaiting committer review Awaiting committer review and removed awaiting change review Awaiting change review awaiting review Awaiting review awaiting changes Awaiting changes labels Oct 14, 2024

EnricoMi force-pushed the flight-client-align-buffers branch 2 times, most recently from 960cb21 to 9909f13 Compare October 22, 2024 08:55

EnricoMi force-pushed the flight-client-align-buffers branch from f2dae5b to d1219d2 Compare October 22, 2024 15:50

rok requested review from felipecrv and lidavidm November 7, 2024 11:53

EnricoMi and others added 18 commits May 8, 2025 10:45

Use existing util::EnsureAlignment

58afa9d

Add C++ Flight data alignment integration test

49c3a3a

Use memory pool from options

bba6320

Fix integration test and lint

a8d7f77

Make GetTypeForBuffers work for ExtensionType with Type::DICTIONARY s…

4180cba

…torage type

Refactor structure of GetTypeForBuffers for Type::DICTIONARY

63a6f0a

Avoid extra shared pointer copy

3b0ca1a

Co-authored-by: David Li <li.davidm96@gmail.com>

Revert "Avoid extra shared pointer copy"

5bafb40

This reverts commit 069c30f7fb87f16d2d8d236ca2e3da1de19dec30.

Default ensure_memory_alignment to false, mention in docs

608978d

Replace RecordBatcBuilder with RecordBatchFromJSON

345e4a2

Avoid shared_ptr ref counter increment

3b247cc

Use meaningful ticket name in flight test integration

f6255e4

Be clear about the 'alignment' term

b65396a

Change ensure_memory_alignment default to False in Python as well

bf43da4

Turn ensure_alignment into enum

7361225

Add 64Byte Alignment

9fe0679

Make Alignment extend int64_t

4dbe7e2

Co-authored-by: Antoine Pitrou <pitrou@free.fr>

Apply code review comments

bf1a7e9

EnricoMi force-pushed the flight-client-align-buffers branch from 63d7251 to bf1a7e9 Compare May 8, 2025 08:45

Fix linting / import

7dbc5da

pitrou reviewed May 15, 2025

View reviewed changes

docs/source/cpp/flight.rst Outdated Show resolved Hide resolved

Fix nasty typos in docs/source/cpp/flight.rst

caf111f

Co-authored-by: Antoine Pitrou <pitrou@free.fr>

lidavidm merged commit f45594d into apache:main May 19, 2025
35 of 36 checks passed

lidavidm removed the awaiting change review Awaiting change review label May 19, 2025

EnricoMi mentioned this pull request May 23, 2025

GH-32276: [C++][FlightRPC] Add option to align RecordBatch buffers given to IPC reader (#44279) G-Research/arrow#6

Merged

	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);
	}

Conversation

EnricoMi commented Oct 1, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

EnricoMi commented Oct 2, 2024

Uh oh!

lidavidm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

EnricoMi commented Oct 7, 2024

Uh oh!

lidavidm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

EnricoMi commented Oct 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

EnricoMi commented Oct 26, 2024

Uh oh!

Uh oh!

pitrou commented May 15, 2025

Uh oh!

Uh oh!

lidavidm commented May 19, 2025

Uh oh!

conbench-apache-arrow bot commented May 20, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

EnricoMi commented Oct 1, 2024 •

edited

Loading

EnricoMi commented Oct 22, 2024 •

edited

Loading