Skip to content

Wire Arrow consumers to unified native allocator - #225

Closed
gaurav-amz wants to merge 3 commits into
Bukhtawar:unified-arrow-allocatorfrom
gaurav-amz:wire-unified-allocator
Closed

Wire Arrow consumers to unified native allocator#225
gaurav-amz wants to merge 3 commits into
Bukhtawar:unified-arrow-allocatorfrom
gaurav-amz:wire-unified-allocator

Conversation

@gaurav-amz

Copy link
Copy Markdown

Summary

Stacked on top of opensearch-project#21703. Wires every existing Arrow consumer in OpenSearch to obtain memory through the unified ArrowNativeAllocator introduced there. Today three independent RootAllocator instances exist that don't see each other; this PR collapses them into one node-wide hierarchy with elastic budgeting.

NativeAllocator root  (node-wide ceiling)
  ├── flight pool      Arrow Flight RPC transport
  ├── ingest pool      Parquet writers (already wired in #21703)
  ├── query pool       analytics-engine per-query and fragment allocators
  └── datafusion pool  DataFusion native MemoryPool (mirrored via FFM)

What changes

Per consumer:

  • arrow-flight-rpc — Replace static RootAllocator(Long.MAX_VALUE) with a thin shim over the flight pool. All public-API callers unchanged.
  • analytics-engineQueryContext (per-query, 256 MB cap) and AnalyticsSearchService (service lifetime) carve children from the query pool.
  • analytics-backend-datafusion — Delete dead DataFusionService.rootAllocator field and newChildAllocator() method. Register a NativeAllocatorListener that mirrors datafusion-pool resize to the Rust MemoryPool via df_set_memory_pool_limit (R1).

Cross-cutting:

  • NativeAllocatorListener SPI in arrow-spi for pool-resize callbacks (allocator-agnostic).
  • POOL_QUERY and POOL_DATAFUSION constants + min/max settings.
  • native.allocator.root.limit defaults to node.native_memory.limit × 0.8 when unset, leaving 20% headroom for non-Arrow native (jemalloc, Netty direct, JNI). Reconciles with PR native memory based admission control opensearch-project/OpenSearch#21191's admission control.
  • CircuitBreakerPlugin integration — native_arrow breaker, Durability.PERMANENT. Allocator periodically syncs root usage into the breaker via the rebalance hook, surfacing native usage in _nodes/stats?breaker and rolling up into the parent breaker.
  • datafusion.spill_memory_limit_bytes becomes Dynamic only when df_set_spill_limit is exported by the loaded datafusion library (probed at static init). Otherwise stays NodeScope-only and OpenSearch rejects runtime PUTs cleanly. Lights up automatically when the upstream datafusion crate carrying the symbol is picked up here.
  • extendedPlugins = ['native-allocator-arrow'] declared on all migrated plugins.
  • Defensive rebalance() after pool creation in createComponents so pools reach max capacity even when the rebalancer is disabled by default. Without this, pools sit at min=0 and silently fail allocations — applies to parquet's ingest pool too. Worth folding into Unified native allocator framework to track and monitor Arrow allocations opensearch-project/OpenSearch#21703 itself.

Test plan

  • :libs:opensearch-arrow-spi:test — passes
  • :plugins:native-allocator-arrow:test — passes (15 tests including 5 new for listener fan-out, breaker sync, AC-derived default)
  • :plugins:arrow-flight-rpc:test — passes (12 tests in FlightClientChannelTests, all green)
  • :sandbox:plugins:analytics-engine:test — passes
  • :sandbox:plugins:analytics-backend-datafusion:test — passes (sandbox tests requiring native lib excluded; native lib not built in this checkout)

Notes for reviewers

This PR targets Bukhtawar:unified-arrow-allocator directly. After opensearch-project#21703 merges to main, this rebases onto main and the diff stays the same.

Overlap with opensearch-project#21465 (bowen's arrow-base): there is meaningful overlap if arrow-base lands first — ArrowAllocatorProvider would be deleted there, my shim becomes redundant, and ArrowAllocatorService becomes the canonical injection target. Coordinating with bowen separately on whether ArrowAllocatorService should expose pool semantics so this work slots in cleanly.

Out of scope:

  • Deprecation aliases for parquet.max_native_allocation and datafusion.memory_pool_limit_bytes (settings still work; aliases are a follow-up).
  • Per-pool stats endpoint (_nodes/stats?native_allocator) — today the native_arrow breaker gives aggregate visibility.
  • Rust-side df_set_spill_limit and RwLock<Arc<RuntimeEnv>> work — separate PR on the datafusion side.

Bukhtawar and others added 3 commits May 18, 2026 11:30
Introduces a new plugin (native-allocator-arrow) that owns a single
Arrow RootAllocator for the node, with named pool-level children for
each subsystem (flight, ingest). Each pool has a min (guaranteed floor)
and max (burst ceiling). A background rebalancer (disabled by default)
redistributes unused capacity across pools so active pools can grow
beyond their min up to their max.

The SPI (libs/arrow-spi) defines the Arrow-agnostic interface, pool
config constants, and Writeable stats shape. The plugin implements it
with Arrow's BufferAllocator and registers configurable limits as
dynamic cluster settings.

Also wires the Parquet ingest path to use the unified ingest pool
instead of a standalone RootAllocator.

Signed-off-by: Bukhtawar Khan <bukhtawa@amazon.com>
Signed-off-by: Bukhtawar Khan <bukhtawa@amazon.com>
Builds on opensearch-project#21703 with the framework-side changes that were not part of
its initial scope. Consumer wiring for arrow-flight-rpc and
analytics-engine is handled by opensearch-project#21465 (arrow-base, already merged) and
not duplicated here.

Pools added:
  POOL_QUERY       analytics-engine query execution
  POOL_DATAFUSION  DataFusion native MemoryPool mirror

Both registered with min/max settings dynamically updatable. The flight
and ingest pools registered by opensearch-project#21703 are unchanged.

Cross-cutting changes:

* NativeAllocatorListener SPI in arrow-spi for pool-resize callbacks
  (allocator-agnostic, no Arrow types in the signature). DataFusion
  uses this to mirror datafusion-pool resize to the Rust MemoryPool
  via df_set_memory_pool_limit (R1).

* Derive native.allocator.root.limit from node.native_memory.limit
  * 0.8 when unset, leaving 20% headroom for non-Arrow native usage
  that admission control still needs to throttle on (R5). Reconciles
  with opensearch-project#21191's admission control.

* Implement CircuitBreakerPlugin: register a native_arrow breaker
  with Durability.PERMANENT so root usage is reflected in
  _nodes/stats?breaker and the parent breaker rolls up off-heap
  pressure. Allocator periodically syncs root.getAllocatedMemory
  into the breaker counter via the rebalance hook.

* Probe df_set_spill_limit at NativeBridge static init. When the
  symbol is present, datafusion.spill_memory_limit_bytes becomes
  Dynamic and a listener calls the FFM symbol on resize. When
  absent, the setting stays NodeScope-only and OpenSearch rejects
  runtime PUTs cleanly. Lights up automatically once the upstream
  datafusion crate carrying df_set_spill_limit is picked up here.

* Delete dead DataFusionService.rootAllocator field and
  newChildAllocator() method (no production callers).

* Call rebalance() once after pool creation in createComponents so
  pools reach their max capacity even when the rebalancer is
  disabled by default. Without this, pools sit at min=0 and
  silently fail allocations.

* Add ensureForTesting helper for unit tests that bring up
  consumers without the plugin lifecycle.

Tests cover listener fan-out, multi-listener invocation, exception
isolation, breaker sync via rebalance, AC-derived default resolution.

Signed-off-by: Gaurav Singh <gauravsg@amazon.com>
@gaurav-amz
gaurav-amz force-pushed the wire-unified-allocator branch from 4cdd366 to 71b9e9f Compare May 18, 2026 08:43
@Bukhtawar
Bukhtawar force-pushed the unified-arrow-allocator branch 6 times, most recently from 9762df2 to a40258a Compare May 18, 2026 15:38
@gaurav-amz gaurav-amz closed this May 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants