From 5a6a433aec1ce54d7b91177f1415a0e9b3de3de3 Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Thu, 26 Oct 2023 14:45:38 -0400 Subject: [PATCH] ESQL: Enable block tracking in extract tests This enables memory tracking in the tests for `grok` and `dissect`. They are already tracking their memory. --- .../operator/ColumnExtractOperatorTests.java | 5 ++ .../operator/StringExtractOperatorTests.java | 61 +++++++++++-------- 2 files changed, 40 insertions(+), 26 deletions(-) diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java index 7680fd410a709..78d0485593197 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java @@ -91,4 +91,9 @@ protected void assertSimpleOutput(List input, List results) { protected ByteSizeValue smallEnoughToCircuitBreak() { return ByteSizeValue.ofBytes(between(1, 32)); } + + @Override + protected DriverContext driverContext() { + return breakingDriverContext(); + } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java index 5dff00ec930c4..52cd969f7fd8b 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java @@ -99,31 +99,40 @@ public Block.Ref eval(Page page) { public void close() {} }, new FirstWord("test"), driverContext()); - BytesRefBlock.Builder builder = BytesRefBlock.newBlockBuilder(1); - builder.beginPositionEntry(); - builder.appendBytesRef(new BytesRef("foo1 bar1")); - builder.appendBytesRef(new BytesRef("foo2 bar2")); - builder.endPositionEntry(); - builder.beginPositionEntry(); - builder.appendBytesRef(new BytesRef("foo3 bar3")); - builder.appendBytesRef(new BytesRef("foo4 bar4")); - builder.appendBytesRef(new BytesRef("foo5 bar5")); - builder.endPositionEntry(); - Page page = new Page(builder.build()); - - Page result = operator.process(page); - Block resultBlock = result.getBlock(1); - assertThat(resultBlock.getPositionCount(), equalTo(2)); - assertThat(resultBlock.getValueCount(0), equalTo(2)); - assertThat(resultBlock.getValueCount(1), equalTo(3)); - BytesRefBlock brb = (BytesRefBlock) resultBlock; - BytesRef spare = new BytesRef(""); - int idx = brb.getFirstValueIndex(0); - assertThat(brb.getBytesRef(idx, spare).utf8ToString(), equalTo("foo1")); - assertThat(brb.getBytesRef(idx + 1, spare).utf8ToString(), equalTo("foo2")); - idx = brb.getFirstValueIndex(1); - assertThat(brb.getBytesRef(idx, spare).utf8ToString(), equalTo("foo3")); - assertThat(brb.getBytesRef(idx + 1, spare).utf8ToString(), equalTo("foo4")); - assertThat(brb.getBytesRef(idx + 2, spare).utf8ToString(), equalTo("foo5")); + Page result = null; + try (BytesRefBlock.Builder builder = BytesRefBlock.newBlockBuilder(1)) { + builder.beginPositionEntry(); + builder.appendBytesRef(new BytesRef("foo1 bar1")); + builder.appendBytesRef(new BytesRef("foo2 bar2")); + builder.endPositionEntry(); + builder.beginPositionEntry(); + builder.appendBytesRef(new BytesRef("foo3 bar3")); + builder.appendBytesRef(new BytesRef("foo4 bar4")); + builder.appendBytesRef(new BytesRef("foo5 bar5")); + builder.endPositionEntry(); + result = operator.process(new Page(builder.build())); + } + try { + Block resultBlock = result.getBlock(1); + assertThat(resultBlock.getPositionCount(), equalTo(2)); + assertThat(resultBlock.getValueCount(0), equalTo(2)); + assertThat(resultBlock.getValueCount(1), equalTo(3)); + BytesRefBlock brb = (BytesRefBlock) resultBlock; + BytesRef spare = new BytesRef(""); + int idx = brb.getFirstValueIndex(0); + assertThat(brb.getBytesRef(idx, spare).utf8ToString(), equalTo("foo1")); + assertThat(brb.getBytesRef(idx + 1, spare).utf8ToString(), equalTo("foo2")); + idx = brb.getFirstValueIndex(1); + assertThat(brb.getBytesRef(idx, spare).utf8ToString(), equalTo("foo3")); + assertThat(brb.getBytesRef(idx + 1, spare).utf8ToString(), equalTo("foo4")); + assertThat(brb.getBytesRef(idx + 2, spare).utf8ToString(), equalTo("foo5")); + } finally { + result.releaseBlocks(); + } + } + + @Override + protected DriverContext driverContext() { + return breakingDriverContext(); } }