org.elasticsearch.xpack.esql.plan.logical.Eval Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of x-pack-esql Show documentation
The plugin that powers ESQL for Elasticsearch
The newest version!
/*
 * 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.plan.logical;

import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
import org.elasticsearch.xpack.esql.core.expression.Alias;
import org.elasticsearch.xpack.esql.core.expression.Attribute;
import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
import org.elasticsearch.xpack.esql.core.expression.NameId;
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
import org.elasticsearch.xpack.esql.core.plan.logical.LogicalPlan;
import org.elasticsearch.xpack.esql.core.plan.logical.UnaryPlan;
import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
import org.elasticsearch.xpack.esql.core.tree.Source;
import org.elasticsearch.xpack.esql.plan.GeneratingPlan;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;

public class Eval extends UnaryPlan implements GeneratingPlan {

    private final List fields;
    private List lazyOutput;

    public Eval(Source source, LogicalPlan child, List fields) {
        super(source, child);
        this.fields = fields;
    }

    public List fields() {
        return fields;
    }

    @Override
    public List output() {
        if (lazyOutput == null) {
            lazyOutput = mergeOutputAttributes(fields, child().output());
        }

        return lazyOutput;
    }

    @Override
    public List generatedAttributes() {
        return asAttributes(fields);
    }

    @Override
    public Eval withGeneratedNames(List newNames) {
        checkNumberOfNewNames(newNames);

        return new Eval(source(), child(), renameAliases(fields, newNames));
    }

    private List renameAliases(List originalAttributes, List newNames) {
        AttributeMap.Builder aliasReplacedByBuilder = AttributeMap.builder();
        List newFields = new ArrayList<>(originalAttributes.size());
        for (int i = 0; i < originalAttributes.size(); i++) {
            Alias field = originalAttributes.get(i);
            String newName = newNames.get(i);
            if (field.name().equals(newName)) {
                newFields.add(field);
            } else {
                Alias newField = new Alias(field.source(), newName, field.qualifier(), field.child(), new NameId(), field.synthetic());
                newFields.add(newField);
                aliasReplacedByBuilder.put(field.toAttribute(), newField.toAttribute());
            }
        }
        AttributeMap aliasReplacedBy = aliasReplacedByBuilder.build();

        // We need to also update any references to the old attributes in the new attributes; e.g.
        // EVAL x = 1, y = x + 1
        // renaming x, y to x1, y1
        // so far became
        // EVAL x1 = 1, y1 = x + 1
        // - but x doesn't exist anymore, so replace it by x1 to obtain
        // EVAL x1 = 1, y1 = x1 + 1

        List newFieldsWithUpdatedRefs = new ArrayList<>(originalAttributes.size());
        for (Alias newField : newFields) {
            newFieldsWithUpdatedRefs.add((Alias) newField.transformUp(ReferenceAttribute.class, r -> aliasReplacedBy.resolve(r, r)));
        }

        return newFieldsWithUpdatedRefs;
    }

    @Override
    public boolean expressionsResolved() {
        return Resolvables.resolved(fields);
    }

    @Override
    public UnaryPlan replaceChild(LogicalPlan newChild) {
        return new Eval(source(), newChild, fields);
    }

    @Override
    protected NodeInfo info() {
        return NodeInfo.create(this, Eval::new, child(), fields);
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }
        Eval eval = (Eval) o;
        return child().equals(eval.child()) && Objects.equals(fields, eval.fields);
    }

    @Override
    public int hashCode() {
        return Objects.hash(super.hashCode(), fields);
    }
}