From 09bc4d48a0d796c1e3ec2daf9cbea5babed2f5ca Mon Sep 17 00:00:00 2001 From: Nhat Nguyen Date: Thu, 28 Aug 2025 17:08:35 -0700 Subject: [PATCH] Avoid excessive memory requests in singleton builders --- .../lucene/read/SingletonDoubleBuilder.java | 15 +++++++++++---- .../compute/lucene/read/SingletonLongBuilder.java | 15 +++++++++++---- 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/SingletonDoubleBuilder.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/SingletonDoubleBuilder.java index 204805ae87e51..ffaa2e5676f46 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/SingletonDoubleBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/SingletonDoubleBuilder.java @@ -7,6 +7,7 @@ package org.elasticsearch.compute.lucene.read; +import org.apache.lucene.util.RamUsageEstimator; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BlockFactory; import org.elasticsearch.core.Releasable; @@ -23,10 +24,13 @@ public final class SingletonDoubleBuilder implements BlockLoader.SingletonDouble private final BlockFactory blockFactory; private int count; + private long memoryUsed; public SingletonDoubleBuilder(int expectedCount, BlockFactory blockFactory) { this.blockFactory = blockFactory; - blockFactory.adjustBreaker(valuesSize(expectedCount)); + long memory = valuesSize(expectedCount); + blockFactory.adjustBreaker(memory); + memoryUsed += memory; this.values = new double[expectedCount]; } @@ -69,7 +73,9 @@ public Block build() { if (values.length != count) { throw new IllegalStateException("expected " + values.length + " values but got " + count); } - return blockFactory.newDoubleArrayVector(values, count).asBlock(); + var result = blockFactory.newDoubleArrayVector(values, count, memoryUsed).asBlock(); + memoryUsed = 0; + return result; } @Override @@ -83,10 +89,11 @@ public BlockLoader.SingletonDoubleBuilder appendLongs(BlockDocValuesReader.ToDou @Override public void close() { - blockFactory.adjustBreaker(-valuesSize(values.length)); + blockFactory.adjustBreaker(-memoryUsed); } static long valuesSize(int count) { - return (long) count * Double.BYTES; + return RamUsageEstimator.alignObjectSize((long) RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + (long) Double.BYTES * count) + + Block.PAGE_MEM_OVERHEAD_PER_BLOCK; } } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/SingletonLongBuilder.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/SingletonLongBuilder.java index de076f8b5d269..5ffa3e845a1b5 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/SingletonLongBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/SingletonLongBuilder.java @@ -7,6 +7,7 @@ package org.elasticsearch.compute.lucene.read; +import org.apache.lucene.util.RamUsageEstimator; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BlockFactory; import org.elasticsearch.core.Releasable; @@ -20,12 +21,15 @@ public final class SingletonLongBuilder implements BlockLoader.SingletonLongBuil private final long[] values; private final BlockFactory blockFactory; + private long memoryUsed; private int count; public SingletonLongBuilder(int expectedCount, BlockFactory blockFactory) { this.blockFactory = blockFactory; - blockFactory.adjustBreaker(valuesSize(expectedCount)); + final long memory = valuesSize(expectedCount); + blockFactory.adjustBreaker(memory); + memoryUsed += memory; this.values = new long[expectedCount]; } @@ -68,7 +72,9 @@ public Block build() { if (values.length != count) { throw new IllegalStateException("expected [" + values.length + "] values but got [" + count + "]"); } - return blockFactory.newLongArrayVector(values, count).asBlock(); + var result = blockFactory.newLongArrayVector(values, count, memoryUsed).asBlock(); + memoryUsed = 0; + return result; } @Override @@ -86,10 +92,11 @@ public BlockLoader.SingletonLongBuilder appendLongs(long[] values, int from, int @Override public void close() { - blockFactory.adjustBreaker(-valuesSize(values.length)); + blockFactory.adjustBreaker(-memoryUsed); } static long valuesSize(int count) { - return (long) count * Long.BYTES; + return RamUsageEstimator.alignObjectSize((long) RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + (long) Long.BYTES * count) + + Block.PAGE_MEM_OVERHEAD_PER_BLOCK; } }