diff --git a/docs/changelog/127201.yaml b/docs/changelog/127201.yaml new file mode 100644 index 0000000000000..5a205c65785b0 --- /dev/null +++ b/docs/changelog/127201.yaml @@ -0,0 +1,5 @@ +pr: 127201 +summary: Emit ordinal output block for values aggregate +area: ES|QL +type: enhancement +issues: [] diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesBytesRefAggregator.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesBytesRefAggregator.java index f326492664fb8..9018a1b7b73fb 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesBytesRefAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesBytesRefAggregator.java @@ -19,7 +19,10 @@ import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BlockFactory; import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.IntBlock; import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.OrdinalBytesRefBlock; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.core.Releasables; @@ -128,7 +131,7 @@ public void close() { */ public static class GroupingState implements GroupingAggregatorState { private final LongLongHash values; - private final BytesRefHash bytes; + private BytesRefHash bytes; private GroupingState(BigArrays bigArrays) { LongLongHash _values = null; @@ -237,34 +240,78 @@ Block toBlock(BlockFactory blockFactory, IntVector selected) { ids[selectedCounts[group]++] = id; } } + if (OrdinalBytesRefBlock.isDense(selected.getPositionCount(), Math.toIntExact(values.size()))) { + return buildOrdinalOutputBlock(blockFactory, selected, selectedCounts, ids); + } else { + return buildOutputBlock(blockFactory, selected, selectedCounts, ids); + } + } finally { + blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + } + } - /* - * Insert the ids in order. - */ - BytesRef scratch = new BytesRef(); - try (BytesRefBlock.Builder builder = blockFactory.newBytesRefBlockBuilder(selected.getPositionCount())) { - int start = 0; - for (int s = 0; s < selected.getPositionCount(); s++) { - int group = selected.getInt(s); - int end = selectedCounts[group]; - int count = end - start; - switch (count) { - case 0 -> builder.appendNull(); - case 1 -> append(builder, ids[start], scratch); - default -> { - builder.beginPositionEntry(); - for (int i = start; i < end; i++) { - append(builder, ids[i], scratch); - } - builder.endPositionEntry(); + Block buildOutputBlock(BlockFactory blockFactory, IntVector selected, int[] selectedCounts, int[] ids) { + /* + * Insert the ids in order. + */ + BytesRef scratch = new BytesRef(); + try (BytesRefBlock.Builder builder = blockFactory.newBytesRefBlockBuilder(selected.getPositionCount())) { + int start = 0; + for (int s = 0; s < selected.getPositionCount(); s++) { + int group = selected.getInt(s); + int end = selectedCounts[group]; + int count = end - start; + switch (count) { + case 0 -> builder.appendNull(); + case 1 -> append(builder, ids[start], scratch); + default -> { + builder.beginPositionEntry(); + for (int i = start; i < end; i++) { + append(builder, ids[i], scratch); + } + builder.endPositionEntry(); + } + } + start = end; + } + return builder.build(); + } + } + + Block buildOrdinalOutputBlock(BlockFactory blockFactory, IntVector selected, int[] selectedCounts, int[] ids) { + BytesRefVector dict = null; + IntBlock ordinals = null; + BytesRefBlock result = null; + var dictArray = bytes.takeBytesRefsOwnership(); + bytes = null; // transfer ownership to dictArray + try (var builder = blockFactory.newIntBlockBuilder(selected.getPositionCount())) { + int start = 0; + for (int s = 0; s < selected.getPositionCount(); s++) { + int group = selected.getInt(s); + int end = selectedCounts[group]; + int count = end - start; + switch (count) { + case 0 -> builder.appendNull(); + case 1 -> builder.appendInt(Math.toIntExact(values.getKey2(ids[start]))); + default -> { + builder.beginPositionEntry(); + for (int i = start; i < end; i++) { + builder.appendInt(Math.toIntExact(values.getKey2(ids[i]))); } + builder.endPositionEntry(); } - start = end; } - return builder.build(); + start = end; } + ordinals = builder.build(); + dict = blockFactory.newBytesRefArrayVector(dictArray, Math.toIntExact(dictArray.size())); + dictArray = null; // transfer ownership to dict + result = new OrdinalBytesRefBlock(ordinals, dict); + return result; } finally { - blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + if (result == null) { + Releasables.close(dictArray, dict, ordinals); + } } } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesDoubleAggregator.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesDoubleAggregator.java index 752cd53a140f7..f5b0d519dd890 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesDoubleAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesDoubleAggregator.java @@ -216,33 +216,36 @@ Block toBlock(BlockFactory blockFactory, IntVector selected) { ids[selectedCounts[group]++] = id; } } + return buildOutputBlock(blockFactory, selected, selectedCounts, ids); + } finally { + blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + } + } - /* - * Insert the ids in order. - */ - try (DoubleBlock.Builder builder = blockFactory.newDoubleBlockBuilder(selected.getPositionCount())) { - int start = 0; - for (int s = 0; s < selected.getPositionCount(); s++) { - int group = selected.getInt(s); - int end = selectedCounts[group]; - int count = end - start; - switch (count) { - case 0 -> builder.appendNull(); - case 1 -> append(builder, ids[start]); - default -> { - builder.beginPositionEntry(); - for (int i = start; i < end; i++) { - append(builder, ids[i]); - } - builder.endPositionEntry(); + Block buildOutputBlock(BlockFactory blockFactory, IntVector selected, int[] selectedCounts, int[] ids) { + /* + * Insert the ids in order. + */ + try (DoubleBlock.Builder builder = blockFactory.newDoubleBlockBuilder(selected.getPositionCount())) { + int start = 0; + for (int s = 0; s < selected.getPositionCount(); s++) { + int group = selected.getInt(s); + int end = selectedCounts[group]; + int count = end - start; + switch (count) { + case 0 -> builder.appendNull(); + case 1 -> append(builder, ids[start]); + default -> { + builder.beginPositionEntry(); + for (int i = start; i < end; i++) { + append(builder, ids[i]); } + builder.endPositionEntry(); } - start = end; } - return builder.build(); + start = end; } - } finally { - blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + return builder.build(); } } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesFloatAggregator.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesFloatAggregator.java index 91f1730ab3111..4cfbf329a895d 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesFloatAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesFloatAggregator.java @@ -223,33 +223,36 @@ Block toBlock(BlockFactory blockFactory, IntVector selected) { ids[selectedCounts[group]++] = id; } } + return buildOutputBlock(blockFactory, selected, selectedCounts, ids); + } finally { + blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + } + } - /* - * Insert the ids in order. - */ - try (FloatBlock.Builder builder = blockFactory.newFloatBlockBuilder(selected.getPositionCount())) { - int start = 0; - for (int s = 0; s < selected.getPositionCount(); s++) { - int group = selected.getInt(s); - int end = selectedCounts[group]; - int count = end - start; - switch (count) { - case 0 -> builder.appendNull(); - case 1 -> append(builder, ids[start]); - default -> { - builder.beginPositionEntry(); - for (int i = start; i < end; i++) { - append(builder, ids[i]); - } - builder.endPositionEntry(); + Block buildOutputBlock(BlockFactory blockFactory, IntVector selected, int[] selectedCounts, int[] ids) { + /* + * Insert the ids in order. + */ + try (FloatBlock.Builder builder = blockFactory.newFloatBlockBuilder(selected.getPositionCount())) { + int start = 0; + for (int s = 0; s < selected.getPositionCount(); s++) { + int group = selected.getInt(s); + int end = selectedCounts[group]; + int count = end - start; + switch (count) { + case 0 -> builder.appendNull(); + case 1 -> append(builder, ids[start]); + default -> { + builder.beginPositionEntry(); + for (int i = start; i < end; i++) { + append(builder, ids[i]); } + builder.endPositionEntry(); } - start = end; } - return builder.build(); + start = end; } - } finally { - blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + return builder.build(); } } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesIntAggregator.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesIntAggregator.java index c4f595d938aa9..38e5ad99cf581 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesIntAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesIntAggregator.java @@ -223,33 +223,36 @@ Block toBlock(BlockFactory blockFactory, IntVector selected) { ids[selectedCounts[group]++] = id; } } + return buildOutputBlock(blockFactory, selected, selectedCounts, ids); + } finally { + blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + } + } - /* - * Insert the ids in order. - */ - try (IntBlock.Builder builder = blockFactory.newIntBlockBuilder(selected.getPositionCount())) { - int start = 0; - for (int s = 0; s < selected.getPositionCount(); s++) { - int group = selected.getInt(s); - int end = selectedCounts[group]; - int count = end - start; - switch (count) { - case 0 -> builder.appendNull(); - case 1 -> append(builder, ids[start]); - default -> { - builder.beginPositionEntry(); - for (int i = start; i < end; i++) { - append(builder, ids[i]); - } - builder.endPositionEntry(); + Block buildOutputBlock(BlockFactory blockFactory, IntVector selected, int[] selectedCounts, int[] ids) { + /* + * Insert the ids in order. + */ + try (IntBlock.Builder builder = blockFactory.newIntBlockBuilder(selected.getPositionCount())) { + int start = 0; + for (int s = 0; s < selected.getPositionCount(); s++) { + int group = selected.getInt(s); + int end = selectedCounts[group]; + int count = end - start; + switch (count) { + case 0 -> builder.appendNull(); + case 1 -> append(builder, ids[start]); + default -> { + builder.beginPositionEntry(); + for (int i = start; i < end; i++) { + append(builder, ids[i]); } + builder.endPositionEntry(); } - start = end; } - return builder.build(); + start = end; } - } finally { - blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + return builder.build(); } } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesLongAggregator.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesLongAggregator.java index 8ae5da509151e..4bfc230d7e1f7 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesLongAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/ValuesLongAggregator.java @@ -216,33 +216,36 @@ Block toBlock(BlockFactory blockFactory, IntVector selected) { ids[selectedCounts[group]++] = id; } } + return buildOutputBlock(blockFactory, selected, selectedCounts, ids); + } finally { + blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + } + } - /* - * Insert the ids in order. - */ - try (LongBlock.Builder builder = blockFactory.newLongBlockBuilder(selected.getPositionCount())) { - int start = 0; - for (int s = 0; s < selected.getPositionCount(); s++) { - int group = selected.getInt(s); - int end = selectedCounts[group]; - int count = end - start; - switch (count) { - case 0 -> builder.appendNull(); - case 1 -> append(builder, ids[start]); - default -> { - builder.beginPositionEntry(); - for (int i = start; i < end; i++) { - append(builder, ids[i]); - } - builder.endPositionEntry(); + Block buildOutputBlock(BlockFactory blockFactory, IntVector selected, int[] selectedCounts, int[] ids) { + /* + * Insert the ids in order. + */ + try (LongBlock.Builder builder = blockFactory.newLongBlockBuilder(selected.getPositionCount())) { + int start = 0; + for (int s = 0; s < selected.getPositionCount(); s++) { + int group = selected.getInt(s); + int end = selectedCounts[group]; + int count = end - start; + switch (count) { + case 0 -> builder.appendNull(); + case 1 -> append(builder, ids[start]); + default -> { + builder.beginPositionEntry(); + for (int i = start; i < end; i++) { + append(builder, ids[i]); } + builder.endPositionEntry(); } - start = end; } - return builder.build(); + start = end; } - } finally { - blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + return builder.build(); } } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/X-ValuesAggregator.java.st b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/X-ValuesAggregator.java.st index 68c6a8640cbd0..f0397de497426 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/X-ValuesAggregator.java.st +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/X-ValuesAggregator.java.st @@ -28,13 +28,20 @@ import org.elasticsearch.compute.ann.GroupingAggregator; import org.elasticsearch.compute.ann.IntermediateState; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BlockFactory; -$if(int||double||float||BytesRef)$ +$if(int||double||float)$ import org.elasticsearch.compute.data.$Type$Block; +$elseif(BytesRef)$ +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.IntBlock; $endif$ import org.elasticsearch.compute.data.IntVector; $if(long)$ import org.elasticsearch.compute.data.LongBlock; $endif$ +$if(BytesRef)$ +import org.elasticsearch.compute.data.OrdinalBytesRefBlock; +$endif$ import org.elasticsearch.compute.operator.DriverContext; $if(BytesRef)$ import org.elasticsearch.core.Releasables; @@ -235,7 +242,7 @@ $if(long||double)$ $elseif(BytesRef)$ private final LongLongHash values; - private final BytesRefHash bytes; + private BytesRefHash bytes; $elseif(int||float)$ private final LongHash values; @@ -364,40 +371,88 @@ $endif$ ids[selectedCounts[group]++] = id; } } +$if(BytesRef)$ + if (OrdinalBytesRefBlock.isDense(selected.getPositionCount(), Math.toIntExact(values.size()))) { + return buildOrdinalOutputBlock(blockFactory, selected, selectedCounts, ids); + } else { + return buildOutputBlock(blockFactory, selected, selectedCounts, ids); + } +$else$ + return buildOutputBlock(blockFactory, selected, selectedCounts, ids); +$endif$ + } finally { + blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + } + } - /* - * Insert the ids in order. - */ + Block buildOutputBlock(BlockFactory blockFactory, IntVector selected, int[] selectedCounts, int[] ids) { + /* + * Insert the ids in order. + */ $if(BytesRef)$ - BytesRef scratch = new BytesRef(); + BytesRef scratch = new BytesRef(); $endif$ - try ($Type$Block.Builder builder = blockFactory.new$Type$BlockBuilder(selected.getPositionCount())) { - int start = 0; - for (int s = 0; s < selected.getPositionCount(); s++) { - int group = selected.getInt(s); - int end = selectedCounts[group]; - int count = end - start; - switch (count) { - case 0 -> builder.appendNull(); - case 1 -> append(builder, ids[start]$if(BytesRef)$, scratch$endif$); - default -> { - builder.beginPositionEntry(); - for (int i = start; i < end; i++) { - append(builder, ids[i]$if(BytesRef)$, scratch$endif$); - } - builder.endPositionEntry(); + try ($Type$Block.Builder builder = blockFactory.new$Type$BlockBuilder(selected.getPositionCount())) { + int start = 0; + for (int s = 0; s < selected.getPositionCount(); s++) { + int group = selected.getInt(s); + int end = selectedCounts[group]; + int count = end - start; + switch (count) { + case 0 -> builder.appendNull(); + case 1 -> append(builder, ids[start]$if(BytesRef)$, scratch$endif$); + default -> { + builder.beginPositionEntry(); + for (int i = start; i < end; i++) { + append(builder, ids[i]$if(BytesRef)$, scratch$endif$); } + builder.endPositionEntry(); } - start = end; } - return builder.build(); + start = end; } - } finally { - blockFactory.adjustBreaker(-selectedCountsSize - idsSize); + return builder.build(); } } $if(BytesRef)$ + Block buildOrdinalOutputBlock(BlockFactory blockFactory, IntVector selected, int[] selectedCounts, int[] ids) { + BytesRefVector dict = null; + IntBlock ordinals = null; + BytesRefBlock result = null; + var dictArray = bytes.takeBytesRefsOwnership(); + bytes = null; // transfer ownership to dictArray + try (var builder = blockFactory.newIntBlockBuilder(selected.getPositionCount())) { + int start = 0; + for (int s = 0; s < selected.getPositionCount(); s++) { + int group = selected.getInt(s); + int end = selectedCounts[group]; + int count = end - start; + switch (count) { + case 0 -> builder.appendNull(); + case 1 -> builder.appendInt(Math.toIntExact(values.getKey2(ids[start]))); + default -> { + builder.beginPositionEntry(); + for (int i = start; i < end; i++) { + builder.appendInt(Math.toIntExact(values.getKey2(ids[i]))); + } + builder.endPositionEntry(); + } + } + start = end; + } + ordinals = builder.build(); + dict = blockFactory.newBytesRefArrayVector(dictArray, Math.toIntExact(dictArray.size())); + dictArray = null; // transfer ownership to dict + result = new OrdinalBytesRefBlock(ordinals, dict); + return result; + } finally { + if (result == null) { + Releasables.close(dictArray, dict, ordinals); + } + } + } + private void append($Type$Block.Builder builder, int id, BytesRef scratch) { BytesRef value = bytes.get(values.getKey2(id), scratch); builder.appendBytesRef(value);