cleanup

Author: sidosera

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/promql/LabelSet.java

@@ -0,0 +1,254 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.optimizer.rules.logical.promql;
+
+import org.elasticsearch.xpack.esql.core.expression.Alias;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
+import org.elasticsearch.xpack.esql.core.expression.MetadataAttribute;
+import org.elasticsearch.xpack.esql.core.expression.NameId;
+import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.expression.function.grouping.TimeSeriesWithout;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+sealed public interface LabelSet permits LabelSet.FixedLabelSet, LabelSet.ExclusionLabelSet {
+    /**
+     * Create a label set with known declared.
+     */
+    static LabelSet of(List<Attribute> labels) {
+        return new LabelSet.FixedLabelSet(labels);
+    }
+
+    /**
+     * Project to intersection with input's visible declared. Kills uncertainty.
+     * Carries declared for null synthesis (PromQL requires all BY declared in output).
+     */
+    static LabelSet by(LabelSet input, List<Attribute> labels) {
+        return new LabelSet.FixedLabelSet(intersection(input.declaredLabels(), labels), labels);
+    }
+
+    /**
+     * Remove declared. Introduces uncertainty. Accumulates exclusions.
+     */
+    static LabelSet without(LabelSet input, List<Attribute> excluded) {
+        List<Attribute> visible = difference(input.declaredLabels(), excluded);
+        if (input instanceof LabelSet.ExclusionLabelSet w) {
+            var accumulated = new ArrayList<>(w.excludedLabels());
+            accumulated.addAll(excluded);
+            return new LabelSet.ExclusionLabelSet(visible, accumulated);
+        }
+        return new LabelSet.ExclusionLabelSet(visible, excluded);
+    }
+
+    /**
+     * Intersect with available declared. If available is empty, trust the input.
+     */
+    static LabelSet intersectWithLabels(LabelSet input, List<Attribute> available) {
+        if (available.isEmpty()) {
+            return input;
+        }
+        return new LabelSet.FixedLabelSet(intersection(input.declaredLabels(), available));
+    }
+
+    /**
+     * Empty label set.
+     */
+    static LabelSet none() {
+        return new LabelSet.FixedLabelSet(List.of());
+    }
+
+    /**
+     * Set difference by field name.
+     */
+    static List<Attribute> difference(List<Attribute> from, List<Attribute> toRemove) {
+        Set<String> removeNames = new HashSet<>();
+        for (Attribute attr : toRemove) {
+            removeNames.add(fieldName(attr));
+        }
+        return difference(from, removeNames);
+    }
+
+    /**
+     * Set difference by pre-computed field name set.
+     */
+    static List<Attribute> difference(List<Attribute> from, Set<String> toRemove) {
+        List<Attribute> result = new ArrayList<>();
+        for (Attribute attr : from) {
+            if (toRemove.contains(fieldName(attr)) == false) {
+                result.add(attr);
+            }
+        }
+        return result;
+    }
+
+    /**
+     * Set intersection by field name.
+     */
+    static List<Attribute> intersection(List<Attribute> requested, List<Attribute> available) {
+        Set<String> availableNames = new HashSet<>();
+        for (Attribute attr : available) {
+            availableNames.add(fieldName(attr));
+        }
+        List<Attribute> result = new ArrayList<>();
+        for (Attribute attr : requested) {
+            if (availableNames.contains(fieldName(attr))) {
+                result.add(attr);
+            }
+        }
+        return result;
+    }
+
+    /**
+     * Resolves requested labels against a plan output, matching first by identity then by field name.
+     */
+    static List<Attribute> resolveLabels(List<Attribute> requested, List<Attribute> visibleOutput) {
+        List<Attribute> resolved = new ArrayList<>();
+        for (Attribute attribute : requested) {
+            if (visibleOutput.contains(attribute)) {
+                resolved.add(attribute);
+                continue;
+            }
+            Attribute byName = findAttributeByFieldName(visibleOutput, fieldName(attribute));
+            if (byName != null) {
+                resolved.add(byName);
+            }
+        }
+        return resolved;
+    }
+
+    static Attribute findAttributeByFieldName(List<Attribute> attributes, String fieldNameToFind) {
+        for (Attribute attribute : attributes) {
+            if (fieldName(attribute).equals(fieldNameToFind)) {
+                return attribute;
+            }
+        }
+        return null;
+    }
+
+    static Attribute findAttributeById(List<Attribute> attributes, NameId id) {
+        for (Attribute attribute : attributes) {
+            if (attribute.id().equals(id)) {
+                return attribute;
+            }
+        }
+        return null;
+    }
+
+    static String fieldName(Attribute attr) {
+        if (attr instanceof FieldAttribute fieldAttr) {
+            return fieldAttr.fieldName().string();
+        }
+        return attr.name();
+    }
+
+    /**
+     * Labels known to be present.
+     */
+    List<Attribute> declaredLabels();
+
+    /**
+     * Labels known to be excluded. Empty for exact (Fixed) label sets.
+     */
+    List<Attribute> excludedLabels();
+
+    /**
+     * Declared output labels (for BY null synthesis). Defaults to declared().
+     */
+    default List<Attribute> byDeclaredLabels() {
+        return declaredLabels();
+    }
+
+    /**
+     * True if the label set is exact (no unknowns).
+     */
+    boolean stationary();
+
+    /**
+     * Resolve into concrete grouping and aggregate output expressions.
+     * parentDemand carries the parent's label requirements -- if the parent is WITHOUT,
+     * a _timeseries grouping is added even for Fixed aggregates.
+     */
+    LabelSet.GroupingShape toGroupingShape(Source source, LabelSet parentDemand);
+
+    /**
+     * Match visible declared against plan output. Returns resolved and missing declared.
+     */
+    default LabelSet.Matched match(List<Attribute> planOutput) {
+        List<Attribute> resolved = resolveLabels(declaredLabels(), planOutput);
+        List<Attribute> missing = difference(declaredLabels(), resolved);
+        return new LabelSet.Matched(resolved, missing);
+    }
+
+    record GroupingShape(List<? extends NamedExpression> groupings, List<? extends NamedExpression> outputs) {}
+
+    record Matched(List<Attribute> resolved, List<Attribute> missing) {}
+
+    record FixedLabelSet(List<Attribute> declaredLabels, List<Attribute> byDeclaredLabels) implements LabelSet {
+        FixedLabelSet(List<Attribute> declared) {
+            this(declared, declared);
+        }
+
+        @Override
+        public List<Attribute> excludedLabels() {
+            return List.of();
+        }
+
+        @Override
+        public List<Attribute> byDeclaredLabels() {
+            return byDeclaredLabels;
+        }
+
+        @Override
+        public boolean stationary() {
+            return true;
+        }
+
+        @Override
+        public LabelSet.GroupingShape toGroupingShape(Source source, LabelSet parentDemand) {
+            var groupings = new ArrayList<NamedExpression>(declaredLabels);
+            var outputs = new ArrayList<NamedExpression>();
+            for (Attribute a : declaredLabels) {
+                if (MetadataAttribute.isTimeSeriesAttributeName(a.name()) == false) {
+                    outputs.add(a);
+                }
+            }
+            // If parent demands _timeseries (WITHOUT flowing downward), add it
+            if (parentDemand.stationary() == false) {
+                Alias tw = new TimeSeriesWithout(source, List.copyOf(parentDemand.excludedLabels())).toAttribute();
+                groupings.add(tw);
+                outputs.add(tw.toAttribute());
+            }
+            return new LabelSet.GroupingShape(groupings, outputs);
+        }
+    }
+
+    record ExclusionLabelSet(List<Attribute> declaredLabels, List<Attribute> excludedLabels) implements LabelSet {
+        @Override
+        public boolean stationary() {
+            return false;
+        }
+
+        @Override
+        public LabelSet.GroupingShape toGroupingShape(Source source, LabelSet parentDemand) {
+            Alias tw = new TimeSeriesWithout(source, List.copyOf(excludedLabels)).toAttribute();
+            var groupings = new ArrayList<NamedExpression>();
+            groupings.add(tw);
+            groupings.addAll(declaredLabels);
+            var outputs = new ArrayList<NamedExpression>();
+            outputs.add(tw.toAttribute());
+            outputs.addAll(declaredLabels);
+            return new LabelSet.GroupingShape(groupings, outputs);
+        }
+    }
+
+}