Move state to ChangePointStates

Author: jan-elastic

x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/ChangePointLongAggregatorFunction.java

@@ -7,47 +7,28 @@
 
 package org.elasticsearch.compute.aggregation;
 
-import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.util.BigArrays;
-import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.compute.aggregation.ChangePointStates.GroupingState;
+import org.elasticsearch.compute.aggregation.ChangePointStates.SingleState;
 import org.elasticsearch.compute.ann.Aggregator;
 import org.elasticsearch.compute.ann.GroupingAggregator;
 import org.elasticsearch.compute.ann.IntermediateState;
 import org.elasticsearch.compute.data.Block;
-import org.elasticsearch.compute.data.BlockFactory;
-import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.LongBlock;
 import org.elasticsearch.compute.operator.DriverContext;
-import org.elasticsearch.core.Releasable;
-import org.elasticsearch.xcontent.ToXContent;
-import org.elasticsearch.xcontent.XContentBuilder;
-import org.elasticsearch.xcontent.XContentFactory;
-import org.elasticsearch.xpack.core.ml.utils.NamedXContentObjectHelper;
-import org.elasticsearch.xpack.ml.aggs.MlAggsHelper;
-import org.elasticsearch.xpack.ml.aggs.changepoint.ChangePointDetector;
-import org.elasticsearch.xpack.ml.aggs.changepoint.ChangeType;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.function.Function;
 
 /**
  * Change point detection for series of long values.
  */
 // TODO: make .java.st from this to support different types
 @Aggregator(
     includeTimestamps = true,
-    value = { @IntermediateState(name = "timestamps", type = "LONG_BLOCK"), @IntermediateState(name = "values", type = "LONG_BLOCK") }
+    value = { @IntermediateState(name = "timestamps", type = "LONG_BLOCK"), @IntermediateState(name = "values", type = "DOUBLE_BLOCK") }
 )
 @GroupingAggregator(
     includeTimestamps = true,
-    value = { @IntermediateState(name = "timestamps", type = "LONG_BLOCK"), @IntermediateState(name = "values", type = "LONG_BLOCK") }
+    value = { @IntermediateState(name = "timestamps", type = "LONG_BLOCK"), @IntermediateState(name = "values", type = "DOUBLE_BLOCK") }
 )
 class ChangePointLongAggregator {
 
@@ -59,11 +40,11 @@ public static void combine(SingleState state, long timestamp, long value) {
         state.add(timestamp, value);
     }
 
-    public static void combineIntermediate(SingleState state, LongBlock timestamps, LongBlock values) {
+    public static void combineIntermediate(SingleState state, LongBlock timestamps, DoubleBlock values) {
         int start = values.getFirstValueIndex(0);
         int end = start + values.getValueCount(0);
         for (int i = start; i < end; i++) {
-            combine(state, timestamps.getLong(i), values.getLong(i));
+            state.add(timestamps.getLong(i), values.getDouble(i));
         }
     }
 
@@ -79,7 +60,13 @@ public static void combine(GroupingState current, int groupId, long timestamp, l
         current.add(groupId, timestamp, value);
     }
 
-    public static void combineIntermediate(GroupingState current, int groupId, LongBlock timestamps, LongBlock values, int otherPosition) {
+    public static void combineIntermediate(
+        GroupingState current,
+        int groupId,
+        LongBlock timestamps,
+        DoubleBlock values,
+        int otherPosition
+    ) {
         current.combine(groupId, timestamps, values, otherPosition);
     }
 
@@ -90,190 +77,4 @@ public static void combineStates(GroupingState current, int currentGroupId, Grou
     public static Block evaluateFinal(GroupingState state, IntVector selected, DriverContext driverContext) {
         return state.evaluateFinal(selected, driverContext.blockFactory());
     }
-
-    public static class SingleState implements Releasable {
-        private final BigArrays bigArrays;
-        private int count;
-        private LongArray timestamps;
-        private LongArray values;
-
-        private SingleState(BigArrays bigArrays) {
-            this.bigArrays = bigArrays;
-            count = 0;
-            timestamps = bigArrays.newLongArray(0);
-            values = bigArrays.newLongArray(0);
-        }
-
-        void add(long timestamp, long value) {
-            count++;
-            timestamps = bigArrays.grow(timestamps, count);
-            timestamps.set(count - 1, timestamp);
-            values = bigArrays.grow(values, count);
-            values.set(count - 1, value);
-        }
-
-        void toIntermediate(Block[] blocks, int offset, DriverContext driverContext) {
-            blocks[offset] = toBlock(timestamps, driverContext.blockFactory());
-            blocks[offset + 1] = toBlock(values, driverContext.blockFactory());
-        }
-
-        Block toBlock(LongArray arr, BlockFactory blockFactory) {
-            if (arr.size() == 0) {
-                return blockFactory.newConstantNullBlock(1);
-            }
-            if (values.size() == 1) {
-                return blockFactory.newConstantLongBlockWith(arr.get(0), 1);
-            }
-            try (LongBlock.Builder builder = blockFactory.newLongBlockBuilder((int) arr.size())) {
-                builder.beginPositionEntry();
-                for (int id = 0; id < count; id++) {
-                    builder.appendLong(arr.get(id));
-                }
-                builder.endPositionEntry();
-                return builder.build();
-            }
-        }
-
-        BytesRef toBytesRef() {
-            // TODO: this copying doesn't account for memory
-            List<TimeAndValue> list = new ArrayList<>(count);
-            for (int i = 0; i < count; i++) {
-                list.add(new TimeAndValue(timestamps.get(i), values.get(i)));
-            }
-            Collections.sort(list);
-            double[] values = new double[count];
-            for (int i = 0; i < count; i++) {
-                values[i] = list.get(i).value;
-            }
-            MlAggsHelper.DoubleBucketValues bucketValues = new MlAggsHelper.DoubleBucketValues(null, values);
-            ChangeType changeType = ChangePointDetector.getChangeType(bucketValues);
-            try (XContentBuilder xContentBuilder = XContentFactory.jsonBuilder()) {
-                xContentBuilder.startObject();
-                NamedXContentObjectHelper.writeNamedObject(xContentBuilder, ToXContent.EMPTY_PARAMS, "type", changeType);
-                xContentBuilder.endObject();
-                String xContent = Strings.toString(xContentBuilder);
-                return new BytesRef(xContent);
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
-        }
-
-        public void add(LongBlock timestamps, LongBlock values, int otherPosition) {
-            final int valueCount = timestamps.getValueCount(otherPosition);
-            final int firstIndex = timestamps.getFirstValueIndex(otherPosition);
-            for (int i = 0; i < valueCount; i++) {
-                add(timestamps.getLong(firstIndex + i), values.getLong(firstIndex + i));
-            }
-        }
-
-        public Block toBlock(BlockFactory blockFactory) {
-            // TODO: this needs to output multiple columns or a composite object, not a JSON blob.
-            return blockFactory.newConstantBytesRefBlockWith(toBytesRef(), 1);
-        }
-
-        record TimeAndValue(long timestamp, long value) implements Comparable<TimeAndValue> {
-            @Override
-            public int compareTo(TimeAndValue other) {
-                return Long.compare(timestamp, other.timestamp);
-            }
-        }
-
-        @Override
-        public void close() {
-            timestamps.close();
-            values.close();
-        }
-    }
-
-    public static class GroupingState implements Releasable {
-        private final BigArrays bigArrays;
-        private final Map<Integer, SingleState> states;
-
-        private GroupingState(BigArrays bigArrays) {
-            this.bigArrays = bigArrays;
-            states = new HashMap<>();
-        }
-
-        void add(int groupId, long timestamp, long value) {
-            SingleState state = states.computeIfAbsent(groupId, key -> new SingleState(bigArrays));
-            state.add(timestamp, value);
-        }
-
-        void combine(int groupId, LongBlock timestamps, LongBlock values, int otherPosition) {
-            if (timestamps.getValueCount(otherPosition) == 0) {
-                return;
-            }
-            SingleState state = states.computeIfAbsent(groupId, key -> new SingleState(bigArrays));
-            state.add(timestamps, values, otherPosition);
-        }
-
-        void combineState(int groupId, GroupingState otherState, int otherGroupId) {
-            SingleState other = otherState.states.get(otherGroupId);
-            if (other == null) {
-                return;
-            }
-            var state = states.computeIfAbsent(groupId, key -> new SingleState(bigArrays));
-            for (int i = 0; i < other.timestamps.size(); i++) {
-                state.add(state.timestamps.get(i), state.values.get(i));
-            }
-        }
-
-        void toIntermediate(Block[] blocks, int offset, IntVector selected, DriverContext driverContext) {
-            blocks[offset] = toBlock(s -> s.timestamps, driverContext.blockFactory(), selected);
-            blocks[offset + 1] = toBlock(s -> s.values, driverContext.blockFactory(), selected);
-        }
-
-        public Block evaluateFinal(IntVector selected, BlockFactory blockFactory) {
-            // TODO: this needs to output multiple columns or a composite object, not a JSON blob.
-            try (BytesRefBlock.Builder builder = blockFactory.newBytesRefBlockBuilder(selected.getPositionCount())) {
-                for (int s = 0; s < selected.getPositionCount(); s++) {
-                    builder.appendBytesRef(states.get(selected.getInt(s)).toBytesRef());
-                }
-                return builder.build();
-            }
-        }
-
-        Block toBlock(Function<SingleState, LongArray> getArray, BlockFactory blockFactory, IntVector selected) {
-            if (states.isEmpty()) {
-                return blockFactory.newConstantNullBlock(selected.getPositionCount());
-            }
-            try (LongBlock.Builder builder = blockFactory.newLongBlockBuilder(selected.getPositionCount())) {
-                for (int s = 0; s < selected.getPositionCount(); s++) {
-                    int selectedGroup = selected.getInt(s);
-                    SingleState state = states.get(selectedGroup);
-                    LongArray values = getArray.apply(state);
-                    int count = 0;
-                    long first = 0;
-                    for (int i = 0; i < state.count; i++) {
-                        long value = values.get(i);
-                        switch (count) {
-                            case 0 -> first = value;
-                            case 1 -> {
-                                builder.beginPositionEntry();
-                                builder.appendLong(first);
-                                builder.appendLong(value);
-                            }
-                            default -> builder.appendLong(value);
-                        }
-                        count++;
-                    }
-                    switch (count) {
-                        case 0 -> builder.appendNull();
-                        case 1 -> builder.appendLong(first);
-                        default -> builder.endPositionEntry();
-                    }
-                }
-                return builder.build();
-            }
-        }
-
-        void enableGroupIdTracking(SeenGroupIds seenGroupIds) {}
-
-        @Override
-        public void close() {
-            for (SingleState state : states.values()) {
-                state.close();
-            }
-        }
-    }
 }