/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.trino.execution.scheduler;

import io.airlift.units.DataSize;
import io.trino.Session;
import io.trino.spi.ErrorCode;

import static java.util.Objects.requireNonNull;

public interface PartitionMemoryEstimator
{
    MemoryRequirements getInitialMemoryRequirements(Session session, DataSize defaultMemoryLimit);

    MemoryRequirements getNextRetryMemoryRequirements(Session session, MemoryRequirements previousMemoryRequirements, ErrorCode errorCode);

    class MemoryRequirements
    {
        private final DataSize requiredMemory;
        private final boolean limitReached;

        MemoryRequirements(DataSize requiredMemory, boolean limitReached)
        {
            this.requiredMemory = requireNonNull(requiredMemory, "requiredMemory is null");
            this.limitReached = limitReached;
        }

        public DataSize getRequiredMemory()
        {
            return requiredMemory;
        }

        public boolean isLimitReached()
        {
            return limitReached;
        }
    }
}

@@ -685,16 +685,108 @@
 rCode;%0A%0A
+import java.util.Objects;%0A%0Aimport static com.google.common.base.MoreObjects.toStringHelper;%0A
 import s
@@ -1663,16 +1663,832 @@
        %7D
+%0A%0A        @Override%0A        public boolean equals(Object o)%0A        %7B%0A            if (this == o) %7B%0A                return true;%0A            %7D%0A            if (o == null %7C%7C getClass() != o.getClass()) %7B%0A                return false;%0A            %7D%0A            MemoryRequirements that = (MemoryRequirements) o;%0A            return limitReached == that.limitReached && Objects.equals(requiredMemory, that.requiredMemory);%0A        %7D%0A%0A        @Override%0A        public int hashCode()%0A        %7B%0A            return Objects.hash(requiredMemory, limitReached);%0A        %7D%0A%0A        @Override%0A        public String toString()%0A        %7B%0A            return toStringHelper(this)%0A                    .add(%22requiredMemory%22, requiredMemory)%0A                    .add(%22limitReached%22, limitReached)%0A                    .toString();%0A        %7D
 %0A    %7D%0A%7D

Add equals, hashCode and toString