org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec 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
There is a newer version: 8.16.1
/*
 * 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.physical;

import org.elasticsearch.xpack.esql.core.expression.Attribute;
import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
import org.elasticsearch.xpack.esql.core.tree.NodeUtils;
import org.elasticsearch.xpack.esql.core.tree.Source;

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

import static java.util.Collections.emptySet;

public class FieldExtractExec extends UnaryExec implements EstimatesRowSize {
    private final List attributesToExtract;
    private final Attribute sourceAttribute;
    // attributes to extract as doc values
    private final Set docValuesAttributes;

    private List lazyOutput;

    public FieldExtractExec(Source source, PhysicalPlan child, List attributesToExtract) {
        this(source, child, attributesToExtract, emptySet());
    }

    public FieldExtractExec(Source source, PhysicalPlan child, List attributesToExtract, Set docValuesAttributes) {
        super(source, child);
        this.attributesToExtract = attributesToExtract;
        this.sourceAttribute = extractSourceAttributesFrom(child);
        this.docValuesAttributes = docValuesAttributes;
    }

    public static Attribute extractSourceAttributesFrom(PhysicalPlan plan) {
        for (Attribute attribute : plan.outputSet()) {
            if (EsQueryExec.isSourceAttribute(attribute)) {
                return attribute;
            }
        }
        return null;
    }

    @Override
    protected NodeInfo info() {
        return NodeInfo.create(this, FieldExtractExec::new, child(), attributesToExtract, docValuesAttributes);
    }

    @Override
    public UnaryExec replaceChild(PhysicalPlan newChild) {
        return new FieldExtractExec(source(), newChild, attributesToExtract, docValuesAttributes);
    }

    public List attributesToExtract() {
        return attributesToExtract;
    }

    public Attribute sourceAttribute() {
        return sourceAttribute;
    }

    public Collection docValuesAttributes() {
        return docValuesAttributes;
    }

    public boolean hasDocValuesAttribute(Attribute attr) {
        return docValuesAttributes.contains(attr);
    }

    @Override
    public List output() {
        if (lazyOutput == null) {
            List childOutput = child().output();
            lazyOutput = new ArrayList<>(childOutput.size() + attributesToExtract.size());
            lazyOutput.addAll(childOutput);
            lazyOutput.addAll(attributesToExtract);
        }

        return lazyOutput;
    }

    @Override
    public PhysicalPlan estimateRowSize(State state) {
        state.add(true, attributesToExtract);
        return this;
    }

    @Override
    public int hashCode() {
        return Objects.hash(attributesToExtract, docValuesAttributes, child());
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj) {
            return true;
        }

        if (obj == null || getClass() != obj.getClass()) {
            return false;
        }

        FieldExtractExec other = (FieldExtractExec) obj;
        return Objects.equals(attributesToExtract, other.attributesToExtract)
            && Objects.equals(docValuesAttributes, other.docValuesAttributes)
            && Objects.equals(child(), other.child());
    }

    @Override
    public String nodeString() {
        return nodeName() + NodeUtils.limitedToString(attributesToExtract) + "<" + NodeUtils.limitedToString(docValuesAttributes) + ">";
    }

}