<commit_before>/*
 * Licensed to Crate under one or more contributor license agreements.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership.  Crate licenses this file
 * to you 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.
 *
 * However, if you have executed another commercial license agreement
 * with Crate these terms will supersede the license and you may use the
 * software solely pursuant to the terms of the relevant commercial
 * agreement.
 */

package io.crate.execution;

import io.crate.Streamer;
import io.crate.concurrent.CompletableFutures;
import io.crate.data.BatchIterator;
import io.crate.data.Bucket;
import io.crate.data.CollectingBatchIterator;
import io.crate.data.Row;
import io.crate.data.RowConsumer;
import io.crate.execution.jobs.PageBucketReceiver;
import io.crate.execution.jobs.PageResultListener;

import javax.annotation.Nonnull;
import javax.annotation.concurrent.GuardedBy;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collector;

public class IncrementalPageBucketReceiver<T> implements PageBucketReceiver {

    private final Function<T, Iterable<Row>> finisher;
    private final BiConsumer<T, Row> accumulator;
    private final T state;
    private final AtomicInteger remainingUpstreams;
    private final CompletableFuture<Iterable<Row>> processingFuture = new CompletableFuture<>();
    private final Executor executor;
    private final Streamer<?>[] streamers;

    private volatile Throwable lastThrowable = null;
    private final BatchIterator<Row> lazyBatchIterator;
    @GuardedBy("accumulateRowsFunction")
    private CompletableFuture<?> currentlyAccumulating;
    private final Function<Bucket, Void> accumulateRowsFunction;

    public IncrementalPageBucketReceiver(Collector<Row, T, Iterable<Row>> collector,
                                         RowConsumer rowConsumer,
                                         Executor executor,
                                         Streamer<?>[] streamers,
                                         int upstreamsCount) {
        this.state = collector.supplier().get();
        this.accumulator = collector.accumulator();
        this.finisher = collector.finisher();
        this.executor = executor;
        this.streamers = streamers;
        this.remainingUpstreams = new AtomicInteger(upstreamsCount);
        lazyBatchIterator = CollectingBatchIterator.newInstance(
            () -> {},
            t -> {},
            () -> processingFuture,
            true);
        rowConsumer.accept(lazyBatchIterator, null);
        accumulateRowsFunction = (Bucket rows) -> {
            try {
                for (Row row : rows) {
                    accumulator.accept(state, row);
                }
            } catch (Throwable e) {
                lastThrowable = e;
                throw e;
            }
            return null;
        };
    }

    @Override
    public void setBucket(int bucketIdx, Bucket rows, boolean isLast, PageResultListener pageResultListener) {
        pageResultListener.needMore(!isLast);

        // We make sure only one accumulation operation runs at a time because the state is not thread-safe.
        synchronized (accumulateRowsFunction) {
            if (currentlyAccumulating == null) {
                currentlyAccumulating = CompletableFutures
                    .supplyAsync(() -> accumulateRowsFunction.apply(rows), executor);
            } else {
                currentlyAccumulating = currentlyAccumulating.whenComplete((r, t) -> {
                    if (t == null) {
                        accumulateRowsFunction.apply(rows);
                    } else if (t instanceof RuntimeException) {
                        lastThrowable = t;
                        throw (RuntimeException) t;
                    } else {
                        lastThrowable = t;
                        throw new RuntimeException(t);
                    }
                });
            }
        }

        if (isLast) {
            if (remainingUpstreams.decrementAndGet() == 0) {
                currentlyAccumulating.whenComplete((r, t) -> consumeRows());
            }
        }
    }

    @Override
    public Streamer<?>[] streamers() {
        return streamers;
    }

    @Override
    public CompletableFuture<?> completionFuture() {
        return processingFuture;
    }

    @Override
    public void consumeRows() {
        if (lastThrowable == null) {
            processingFuture.complete(finisher.apply(state));
        } else {
            processingFuture.completeExceptionally(lastThrowable);
        }
    }

    @Override
    public void kill(@Nonnull Throwable t) {
        lastThrowable = t;
        lazyBatchIterator.kill(t);
        processingFuture.completeExceptionally(t);
    }
}<commit_msg>Remove separate lastThrowable attribute from IncrementalPageBucketReceiver<commit_after>/*
 * Licensed to Crate under one or more contributor license agreements.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership.  Crate licenses this file
 * to you 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.
 *
 * However, if you have executed another commercial license agreement
 * with Crate these terms will supersede the license and you may use the
 * software solely pursuant to the terms of the relevant commercial
 * agreement.
 */

package io.crate.execution;

import io.crate.Streamer;
import io.crate.data.BatchIterator;
import io.crate.data.Bucket;
import io.crate.data.CollectingBatchIterator;
import io.crate.data.Row;
import io.crate.data.RowConsumer;
import io.crate.execution.jobs.PageBucketReceiver;
import io.crate.execution.jobs.PageResultListener;
import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;

import javax.annotation.Nonnull;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collector;

public class IncrementalPageBucketReceiver<T> implements PageBucketReceiver {

    private final Function<T, Iterable<Row>> finisher;
    private final BiConsumer<T, Row> accumulator;
    private final T state;
    private final AtomicInteger remainingUpstreams;
    private final CompletableFuture<Iterable<Row>> processingFuture = new CompletableFuture<>();
    private final Executor executor;
    private final Streamer<?>[] streamers;

    private final BatchIterator<Row> lazyBatchIterator;
    private CompletableFuture<?> currentlyAccumulating;

    public IncrementalPageBucketReceiver(Collector<Row, T, Iterable<Row>> collector,
                                         RowConsumer rowConsumer,
                                         Executor executor,
                                         Streamer<?>[] streamers,
                                         int upstreamsCount) {
        this.state = collector.supplier().get();
        this.accumulator = collector.accumulator();
        this.finisher = collector.finisher();
        this.executor = executor;
        this.streamers = streamers;
        this.remainingUpstreams = new AtomicInteger(upstreamsCount);
        lazyBatchIterator = CollectingBatchIterator.newInstance(
            () -> {},
            t -> {},
            () -> processingFuture,
            true);
        rowConsumer.accept(lazyBatchIterator, null);
    }

    private void processRows(Bucket rows) {
        for (Row row : rows) {
            accumulator.accept(state, row);
        }
    }

    @Override
    public void setBucket(int bucketIdx, Bucket rows, boolean isLast, PageResultListener pageResultListener) {
        if (processingFuture.isCompletedExceptionally()) {
            pageResultListener.needMore(false);
            return;
        } else {
            pageResultListener.needMore(!isLast);
        }

        // We make sure only one accumulation operation runs at a time because the state is not thread-safe.
        synchronized (state) {
            if (currentlyAccumulating == null) {
                try {
                    currentlyAccumulating = CompletableFuture.runAsync(() -> processRows(rows), executor);
                } catch (EsRejectedExecutionException e) {
                    processingFuture.completeExceptionally(e);
                }
            } else {
                currentlyAccumulating = currentlyAccumulating.whenComplete((r, t) -> {
                    if (t == null) {
                        processRows(rows);
                    } else if (t instanceof RuntimeException) {
                        processingFuture.completeExceptionally(t);
                        throw (RuntimeException) t;
                    } else {
                        processingFuture.completeExceptionally(t);
                        throw new RuntimeException(t);
                    }
                });
            }
        }
        if (isLast) {
            if (remainingUpstreams.decrementAndGet() == 0) {
                currentlyAccumulating.whenComplete((r, t) -> consumeRows());
            }
        }
    }

    @Override
    public Streamer<?>[] streamers() {
        return streamers;
    }

    @Override
    public CompletableFuture<?> completionFuture() {
        return processingFuture;
    }

    @Override
    public void consumeRows() {
        processingFuture.complete(finisher.apply(state));
    }

    @Override
    public void kill(@Nonnull Throwable t) {
        lazyBatchIterator.kill(t);
        processingFuture.completeExceptionally(t);
    }
}
{'fraction_non_alphanumeric': 0.0811110085817108, 'fraction_numerical': 0.0009227646027498385, 'mean_word_length': 2.376323987538941, 'pattern_counts': {'":': 0, '<': 34, '<?xml version=': 0, '>': 47, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4442983', 'n_tokens_mistral': 2668, 'n_tokens_neox': 2598, 'n_words': 757}