[client] Replace Netty PooledByteBufAllocator with bump-pointer ChunkedAllocationManager for Arrow memory allocation.

fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScanner.java

@@ -37,6 +37,7 @@
 import org.apache.fluss.rpc.gateway.TabletServerGateway;
 import org.apache.fluss.rpc.messages.LimitScanRequest;
 import org.apache.fluss.rpc.messages.LimitScanResponse;
+import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.ChunkedAllocationManager;
 import org.apache.fluss.types.RowType;
 import org.apache.fluss.utils.CloseableIterator;
 import org.apache.fluss.utils.SchemaUtil;
@@ -66,6 +67,8 @@ public class LimitBatchScanner implements BatchScanner {
     private final SchemaGetter schemaGetter;
     private final KvFormat kvFormat;
     private final int targetSchemaId;
+    /** The chunked allocation manager factory to reuse memory for arrow log write batch. */
+    private final ChunkedAllocationManager.ChunkedFactory chunkedFactory;
 
     /**
      * A cache for schema projection mapping from source schema to target. Use HashMap here, because
@@ -116,6 +119,7 @@ public LimitBatchScanner(
 
         this.kvFormat = tableInfo.getTableConfig().getKvFormat();
         this.endOfInput = false;
+        this.chunkedFactory = new ChunkedAllocationManager.ChunkedFactory();
     }
 
     @Nullable
@@ -164,7 +168,8 @@ private List<InternalRow> parseLimitScanResponse(LimitScanResponse limitScanResp
             }
         } else {
             LogRecordReadContext readContext =
-                    LogRecordReadContext.createReadContext(tableInfo, false, null, schemaGetter);
+                    LogRecordReadContext.createReadContext(
+                            tableInfo, false, null, schemaGetter, chunkedFactory);
             LogRecords records = MemoryLogRecords.pointToByteBuffer(recordsBuffer);
             for (LogRecordBatch logRecordBatch : records.batches()) {
                 // A batch of log record maybe little more than limit, thus we need slice the
@@ -203,5 +208,7 @@ private InternalRow maybeProject(InternalRow originRow) {
     @Override
     public void close() throws IOException {
         scanFuture.cancel(true);
+        // Release off-heap memory held by the chunked allocation manager factory.
+        chunkedFactory.close();
     }
 }

fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogFetcher.java

@@ -53,6 +53,7 @@
 import org.apache.fluss.rpc.protocol.ApiError;
 import org.apache.fluss.rpc.protocol.Errors;
 import org.apache.fluss.rpc.util.PredicateMessageUtils;
+import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.ChunkedAllocationManager;
 import org.apache.fluss.shaded.netty4.io.netty.buffer.ByteBuf;
 import org.apache.fluss.utils.IOUtils;
 import org.apache.fluss.utils.Projection;
@@ -94,6 +95,7 @@ public class LogFetcher implements Closeable {
     //  currently can only do project when generate scanRecord instead of doing project while read
     //  bytes from remote file.
     private final LogRecordReadContext remoteReadContext;
+    private final ChunkedAllocationManager.ChunkedFactory chunkedFactory;
     @Nullable private final Projection projection;
     @Nullable private final org.apache.fluss.rpc.messages.PbPredicate cachedPbPredicate;
     private final int filterSchemaId;
@@ -128,10 +130,13 @@ public LogFetcher(
             SchemaGetter schemaGetter) {
         this.tablePath = tableInfo.getTablePath();
         this.isPartitioned = tableInfo.isPartitioned();
+        this.chunkedFactory = new ChunkedAllocationManager.ChunkedFactory();
         this.readContext =
-                LogRecordReadContext.createReadContext(tableInfo, false, projection, schemaGetter);
+                LogRecordReadContext.createReadContext(
+                        tableInfo, false, projection, schemaGetter, chunkedFactory);
         this.remoteReadContext =
-                LogRecordReadContext.createReadContext(tableInfo, true, projection, schemaGetter);
+                LogRecordReadContext.createReadContext(
+                        tableInfo, true, projection, schemaGetter, chunkedFactory);
         this.projection = projection;
         this.cachedPbPredicate =
                 recordBatchFilter != null
@@ -603,6 +608,7 @@ public synchronized void close() throws IOException {
             IOUtils.closeQuietly(remoteLogDownloader, "remoteLogDownloader");
             readContext.close();
             remoteReadContext.close();
+            chunkedFactory.close();
             isClosed = true;
             LOG.info("Fetcher for {} is closed.", tablePath);
         }

fluss-client/src/main/java/org/apache/fluss/client/write/RecordAccumulator.java

@@ -35,8 +35,8 @@
 import org.apache.fluss.record.LogRecordBatchStatisticsCollector;
 import org.apache.fluss.row.arrow.ArrowWriter;
 import org.apache.fluss.row.arrow.ArrowWriterPool;
-import org.apache.fluss.row.arrow.memory.BufferAllocatorUtil;
 import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator;
+import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.ChunkedAllocationManager;
 import org.apache.fluss.utils.CopyOnWriteMap;
 import org.apache.fluss.utils.MathUtils;
 import org.apache.fluss.utils.clock.Clock;
@@ -63,6 +63,7 @@
 
 import static org.apache.fluss.record.LogRecordBatchFormat.NO_BATCH_SEQUENCE;
 import static org.apache.fluss.record.LogRecordBatchFormat.NO_WRITER_ID;
+import static org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocatorUtil.createBufferAllocator;
 import static org.apache.fluss.utils.Preconditions.checkNotNull;
 
 /* This file is based on source code of Apache Kafka Project (https://kafka.apache.org/), licensed by the Apache
@@ -99,6 +100,9 @@ public final class RecordAccumulator {
     /** The arrow buffer allocator to allocate memory for arrow log write batch. */
     private final BufferAllocator bufferAllocator;
 
+    /** The chunked allocation manager factory, stored for explicit native memory release. */
+    private final ChunkedAllocationManager.ChunkedFactory chunkedFactory;
+
     /** The pool of lazily created arrow {@link ArrowWriter}s for arrow log write batch. */
     private final ArrowWriterPool arrowWriterPool;
 
@@ -134,7 +138,8 @@ public final class RecordAccumulator {
                 Math.max(1, (int) conf.get(ConfigOptions.CLIENT_WRITER_BATCH_SIZE).getBytes());
 
         this.writerBufferPool = LazyMemorySegmentPool.createWriterBufferPool(conf);
-        this.bufferAllocator = BufferAllocatorUtil.createBufferAllocator();
+        this.chunkedFactory = new ChunkedAllocationManager.ChunkedFactory();
+        this.bufferAllocator = createBufferAllocator(chunkedFactory);
         this.arrowWriterPool = new ArrowWriterPool(bufferAllocator);
         this.incomplete = new IncompleteBatches();
         this.nodesDrainIndex = new HashMap<>();
@@ -964,6 +969,8 @@ public void close() {
         // Release all the memory segments.
         bufferAllocator.releaseBytes(bufferAllocator.getAllocatedMemory());
         bufferAllocator.close();
+        // Release native memory held by the chunked allocation manager factory.
+        chunkedFactory.close();
     }
 
     /** Per table bucket and write batches. */

fluss-common/src/main/java/org/apache/fluss/record/LogRecordReadContext.java

@@ -25,8 +25,10 @@
 import org.apache.fluss.row.InternalRow;
 import org.apache.fluss.row.InternalRow.FieldGetter;
 import org.apache.fluss.row.ProjectedRow;
-import org.apache.fluss.row.arrow.memory.BufferAllocatorUtil;
+import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.AllocationManager;
 import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator;
+import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocatorUtil;
+import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.ChunkedAllocationManager;
 import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.fluss.types.DataType;
 import org.apache.fluss.types.RowType;
@@ -68,6 +70,29 @@ public static LogRecordReadContext createReadContext(
             boolean readFromRemote,
             @Nullable Projection projection,
             SchemaGetter schemaGetter) {
+        return createReadContext(
+                tableInfo,
+                readFromRemote,
+                projection,
+                schemaGetter,
+                new ChunkedAllocationManager.ChunkedFactory());
+    }
+
+    /**
+     * Creates a {@link LogRecordReadContext} with a custom {@link AllocationManager.Factory}.
+     *
+     * @param tableInfo the table info of the table to read
+     * @param readFromRemote whether the data is read from remote storage
+     * @param projection the projection to apply, or null for all fields
+     * @param schemaGetter the schema getter to resolve schema by id
+     * @param allocationManagerFactory the factory for creating Arrow memory allocations
+     */
+    public static LogRecordReadContext createReadContext(
+            TableInfo tableInfo,
+            boolean readFromRemote,
+            @Nullable Projection projection,
+            SchemaGetter schemaGetter,
+            AllocationManager.Factory allocationManagerFactory) {
         RowType rowType = tableInfo.getRowType();
         LogFormat logFormat = tableInfo.getTableConfig().getLogFormat();
         // only for arrow log format, the projection can be push downed to the server side
@@ -84,7 +109,12 @@ public static LogRecordReadContext createReadContext(
                 // so set the rowType as is.
                 int[] selectedFields = projection.getProjection();
                 return createArrowReadContext(
-                        rowType, schemaId, selectedFields, false, schemaGetter);
+                        rowType,
+                        schemaId,
+                        selectedFields,
+                        false,
+                        schemaGetter,
+                        allocationManagerFactory);
             } else {
                 // arrow data that returned from server has been projected (in order)
                 RowType projectedRowType = projection.projectInOrder(rowType);
@@ -95,7 +125,8 @@ public static LogRecordReadContext createReadContext(
                         schemaId,
                         selectedFields,
                         projectionPushDowned,
-                        schemaGetter);
+                        schemaGetter,
+                        allocationManagerFactory);
             }
         } else if (logFormat == LogFormat.INDEXED) {
             int[] selectedFields = projection.getProjection();
@@ -113,9 +144,11 @@ private static LogRecordReadContext createArrowReadContext(
             int schemaId,
             int[] selectedFields,
             boolean projectionPushDowned,
-            SchemaGetter schemaGetter) {
+            SchemaGetter schemaGetter,
+            AllocationManager.Factory allocationManagerFactory) {
         // TODO: use a more reasonable memory limit
-        BufferAllocator allocator = BufferAllocatorUtil.createBufferAllocator();
+        BufferAllocator allocator =
+                BufferAllocatorUtil.createBufferAllocator(allocationManagerFactory);
         FieldGetter[] fieldGetters = buildProjectedFieldGetters(dataRowType, selectedFields);
         return new LogRecordReadContext(
                 LogFormat.ARROW,
@@ -139,7 +172,13 @@ private static LogRecordReadContext createArrowReadContext(
     public static LogRecordReadContext createArrowReadContext(
             RowType rowType, int schemaId, SchemaGetter schemaGetter) {
         int[] selectedFields = IntStream.range(0, rowType.getFieldCount()).toArray();
-        return createArrowReadContext(rowType, schemaId, selectedFields, false, schemaGetter);
+        return createArrowReadContext(
+                rowType,
+                schemaId,
+                selectedFields,
+                false,
+                schemaGetter,
+                new ChunkedAllocationManager.ChunkedFactory());
     }
 
     @VisibleForTesting
@@ -150,7 +189,12 @@ public static LogRecordReadContext createArrowReadContext(
             boolean projectionPushDowned) {
         int[] selectedFields = IntStream.range(0, rowType.getFieldCount()).toArray();
         return createArrowReadContext(
-                rowType, schemaId, selectedFields, projectionPushDowned, schemaGetter);
+                rowType,
+                schemaId,
+                selectedFields,
+                projectionPushDowned,
+                schemaGetter,
+                new ChunkedAllocationManager.ChunkedFactory());
     }
 
     /**

fluss-common/src/main/java/org/apache/fluss/row/arrow/memory/ArrowRoundingPolicy.java → fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ArrowRoundingPolicy.java

@@ -17,7 +17,7 @@
  * under the License.
  */
 
-package org.apache.fluss.row.arrow.memory;
+package org.apache.fluss.shaded.arrow.org.apache.arrow.memory;
 
 import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.rounding.RoundingPolicy;
 import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.util.CommonUtil;

fluss-common/src/main/java/org/apache/fluss/row/arrow/memory/BufferAllocatorUtil.java → fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/BufferAllocatorUtil.java

@@ -17,19 +17,26 @@
  * under the License.
  */
 
-package org.apache.fluss.row.arrow.memory;
+package org.apache.fluss.shaded.arrow.org.apache.arrow.memory;
 
-import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.AllocationListener;
-import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator;
-import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.RootAllocator;
+import javax.annotation.Nullable;
 
-import static org.apache.fluss.row.arrow.memory.ArrowRoundingPolicy.ARROW_ROUNDING_POLICY;
+import static org.apache.fluss.shaded.arrow.org.apache.arrow.memory.ArrowRoundingPolicy.ARROW_ROUNDING_POLICY;
 
 /** Utility class for creating Arrow BufferAllocators with the custom ArrowRoundingPolicy. */
 public class BufferAllocatorUtil {
 
     /** Creates a {@link BufferAllocator} configured with the {@link ArrowRoundingPolicy}. */
-    public static BufferAllocator createBufferAllocator() {
-        return new RootAllocator(AllocationListener.NOOP, Long.MAX_VALUE, ARROW_ROUNDING_POLICY);
+    public static BufferAllocator createBufferAllocator(
+            @Nullable AllocationManager.Factory allocationManagerFactory) {
+        ImmutableConfig.Builder builder =
+                ImmutableConfig.builder()
+                        .listener(AllocationListener.NOOP)
+                        .maxAllocation(Long.MAX_VALUE)
+                        .roundingPolicy(ARROW_ROUNDING_POLICY);
+        if (allocationManagerFactory != null) {
+            builder.allocationManagerFactory(allocationManagerFactory);
+        }
+        return new RootAllocator(builder.build());
     }
 }