Skip to content

ESQL: Add a stream view to BreakingBytesRefBuilder #133148

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nik9000 merged 4 commits into from
Aug 22, 2025
Merged

ESQL: Add a stream view to BreakingBytesRefBuilder #133148

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
53bc41c
ESQL: Add a stream view to BreakingBytesRefBuilder
nik9000 Aug 19, 2025
e538aa5
Merge branch 'main' into esql_breaking_stream_view
nik9000 Aug 19, 2025
761e672
Merge branch 'main' into esql_breaking_stream_view
nik9000 Aug 22, 2025
75dd0d5
Merge remote-tracking branch 'nik9000/esql_breaking_stream_view' into…
nik9000 Aug 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions ...sql/compute/src/main/java/org/elasticsearch/compute/operator/BreakingBytesRefBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.RamUsageEstimator;
import org.elasticsearch.common.breaker.CircuitBreaker;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.core.Releasable;

/**
Expand Down Expand Up @@ -162,6 +163,13 @@ public long ramBytesUsed() {
return SHALLOW_SIZE + bytesArrayRamBytesUsed(bytes.bytes.length);
}

/**
* Builds a {@link StreamOutput} view into the {@link BreakingBytesRefBuilder}.
*/
public StreamOutput stream() {
return new Stream();
}

private static long bytesArrayRamBytesUsed(long capacity) {
return RamUsageEstimator.alignObjectSize(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + capacity);
}
Expand All @@ -170,4 +178,32 @@ private static long bytesArrayRamBytesUsed(long capacity) {
public void close() {
breaker.addWithoutBreaking(-ramBytesUsed());
}

private class Stream extends StreamOutput {
@Override
public long position() {
return length();
}

@Override
public void writeByte(byte b) {
append(b);
}

@Override
public void writeBytes(byte[] b, int offset, int length) {
append(b, offset, length);
}

@Override
public void flush() {}

/**
* Closes this stream to further operations. NOOP because we don't want to
* close the builder when we close.
*/
@Override
public void close() {}
}

}
113 changes: 94 additions & 19 deletions ...ompute/src/test/java/org/elasticsearch/compute/operator/BreakingBytesRefBuilderTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,16 @@
import org.apache.lucene.util.RamUsageEstimator;
import org.elasticsearch.common.breaker.CircuitBreaker;
import org.elasticsearch.common.breaker.CircuitBreakingException;
import org.elasticsearch.common.io.stream.BytesRefStreamOutput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.util.MockBigArrays;
import org.elasticsearch.test.ESTestCase;
import org.elasticsearch.xcontent.XContentBuilder;
import org.elasticsearch.xcontent.json.JsonXContent;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.function.Supplier;

import static org.hamcrest.Matchers.equalTo;
Expand All @@ -34,11 +40,6 @@ public void testAddByte() {
testAgainstOracle(() -> new TestIteration() {
final byte b = randomByte();

@Override
public int size() {
return 1;
}

@Override
public void applyToBuilder(BreakingBytesRefBuilder builder) {
builder.append(b);
Expand All @@ -55,11 +56,6 @@ public void testAddBytesRef() {
testAgainstOracle(() -> new TestIteration() {
final BytesRef ref = new BytesRef(randomAlphaOfLengthBetween(1, 100));

@Override
public int size() {
return ref.length;
}

@Override
public void applyToBuilder(BreakingBytesRefBuilder builder) {
builder.append(ref);
Expand Down Expand Up @@ -90,11 +86,6 @@ public void testGrow() {
final int length = between(1, 100);
final byte b = randomByte();

@Override
public int size() {
return length;
}

@Override
public void applyToBuilder(BreakingBytesRefBuilder builder) {
builder.grow(builder.length() + length);
Expand All @@ -111,9 +102,90 @@ public void applyToOracle(BytesRefBuilder oracle) {
});
}

interface TestIteration {
int size();
public void testStream() {
testAgainstOracle(() -> switch (between(0, 3)) {
case 0 -> new XContentTestIteration() {
@Override
protected void apply(XContentBuilder builder) throws IOException {
// Noop
}

@Override
public String toString() {
return "noop";
}
};
case 1 -> new XContentTestIteration() {
private final String value = randomAlphanumericOfLength(10);

@Override
protected void apply(XContentBuilder builder) throws IOException {
builder.value(value);
}

@Override
public String toString() {
return '"' + value + '"';
}
};
case 2 -> new XContentTestIteration() {
private final long value = randomLong();

@Override
protected void apply(XContentBuilder builder) throws IOException {
builder.value(value);
}

@Override
public String toString() {
return Long.toString(value);
}
};
case 3 -> new XContentTestIteration() {
private final String name = randomAlphanumericOfLength(5);
private final String value = randomAlphanumericOfLength(5);

@Override
protected void apply(XContentBuilder builder) throws IOException {
builder.startObject().field(name, value).endObject();
}

@Override
public String toString() {
return name + ": " + value;
}
};
default -> throw new UnsupportedOperationException();
});
}

private abstract static class XContentTestIteration implements TestIteration {
protected abstract void apply(XContentBuilder builder) throws IOException;

@Override
public void applyToBuilder(BreakingBytesRefBuilder builder) {
applyToStream(builder.stream());
}

@Override
public void applyToOracle(BytesRefBuilder oracle) {
BytesRefStreamOutput out = new BytesRefStreamOutput();
applyToStream(out);
oracle.append(out.get());
}

private void applyToStream(StreamOutput out) {
try {
try (XContentBuilder builder = new XContentBuilder(JsonXContent.jsonXContent, out)) {
apply(builder);
}
} catch (IOException e) {
throw new UncheckedIOException(e);
}
}
}

interface TestIteration {
void applyToBuilder(BreakingBytesRefBuilder builder);

void applyToOracle(BytesRefBuilder oracle);
Expand All @@ -131,7 +203,11 @@ private void testAgainstOracle(Supplier<TestIteration> iterations) {
assertThat(builder.bytesRefView(), equalTo(oracle.get()));
while (true) {
TestIteration iteration = iterations.get();
int targetSize = builder.length() + iteration.size();

int prevOracle = oracle.length();
iteration.applyToOracle(oracle);
int size = oracle.length() - prevOracle;
int targetSize = builder.length() + size;
boolean willResize = targetSize >= builder.bytes().length;
if (willResize) {
long resizeMemoryUsage = BreakingBytesRefBuilder.SHALLOW_SIZE + ramForArray(builder.bytes().length);
Expand All @@ -143,7 +219,6 @@ private void testAgainstOracle(Supplier<TestIteration> iterations) {
}
}
iteration.applyToBuilder(builder);
iteration.applyToOracle(oracle);
assertThat(builder.bytesRefView(), equalTo(oracle.get()));
assertThat(
builder.ramBytesUsed(),
Expand Down