/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF 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.
 */

package org.apache.carbondata.processing.loading.iterator;

import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.TimeUnit;

import org.apache.carbondata.common.CarbonIterator;
import org.apache.carbondata.common.logging.LogServiceFactory;
import org.apache.carbondata.core.util.CarbonProperties;

import org.apache.log4j.Logger;

/**
 * It is wrapper class to hold the rows in batches when record writer writes the data and allows
 * to iterate on it during data load. It uses blocking queue to coordinate between read and write.
 */
public class CarbonOutputIteratorWrapper extends CarbonIterator<Object[]> {

  private static final Logger LOG =
      LogServiceFactory.getLogService(CarbonOutputIteratorWrapper.class.getName());

  private boolean close;

  /**
   * Number of rows kept in memory at most will be batchSize * queue size
   */
  private int batchSize = CarbonProperties.getInstance().getBatchSize();

  private RowBatch loadBatch = new RowBatch(batchSize);

  private RowBatch readBatch;

  private ArrayBlockingQueue<RowBatch> queue = new ArrayBlockingQueue<>(10);

  public void write(Object[] row) throws InterruptedException {
    if (close) {
      // already might be closed forcefully
      return;
    }
    if (!loadBatch.addRow(row)) {
      loadBatch.readyRead();
      queue.put(loadBatch);
      loadBatch = new RowBatch(batchSize);
    }
  }

  @Override
  public boolean hasNext() {
    if (readBatch == null || !readBatch.hasNext()) {
      try {
        if (!close) {
          readBatch = queue.poll(5, TimeUnit.MINUTES);
          if (readBatch == null) {
            LOG.warn("This scenario should not happen");
            return false;
          }
        } else {
          readBatch = queue.poll();
          if (readBatch == null) {
            return false;
          }
        }
      } catch (InterruptedException e) {
        throw new RuntimeException(e);
      }
    }
    return readBatch.hasNext();
  }

  @Override
  public Object[] next() {
    return readBatch.next();
  }

  public void closeWriter(boolean isForceClose) {
    if (close) {
      // already might be closed forcefully
      return;
    }
    try {
      if (isForceClose) {
        // first make close is set to true, when force close happens because of dead consumer.
        // so that, write() method will stop taking input rows.
        close = true;
        // once write() method stops taking input rows, clear the queue.
        // If queue is cleared before close is set to true, then queue will be again filled
        // by .write() and it can go to blocking put() forever as consumer is dead.
        queue.clear();
        return;
      }
      // below code will ensure that the last RowBatch is consumed properly
      loadBatch.readyRead();
      if (loadBatch.size > 0) {
        queue.put(loadBatch);
      }
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    }
    close = true;
    // It is required if the thread waits for take.
    if (queue.isEmpty()) {
      if (!queue.offer(new RowBatch(0))) {
        LOG.warn("The default last element is not added to queue");
      }
    }
  }

  private static class RowBatch extends CarbonIterator<Object[]> {

    private int counter;

    private Object[][] batch;

    private int size;

    private RowBatch(int size) {
      batch = new Object[size][];
      this.size = size;
    }

    /**
     * Add row to the batch, it can hold rows till the batch size.
     * @param row
     * @return false if the row cannot be added as batch is full.
     */
    public boolean addRow(Object[] row) {
      batch[counter++] = row;
      return counter < size;
    }

    public void readyRead() {
      size = counter;
      counter = 0;
    }

    @Override
    public boolean hasNext() {
      return counter < size;
    }

    @Override
    public Object[] next() {
      assert (counter < size);
      return batch[counter++];
    }
  }

}

@@ -2296,276 +2296,1192 @@
     
-try %7B%0A        if (!close) %7B%0A          readBatch = queue.poll(5, TimeUnit.MINUTES);%0A          if (readBatch == null) %7B%0A            LOG.warn(%22This scenario should not happen%22);%0A            return false;%0A          %7D%0A        %7D else %7B%0A          readBatch = queue.poll();%0A  
+// if readBatch don't have next row, set it to null%0A      readBatch = null;%0A      try %7B%0A        // if the output is not closed, poll a batch from the queue in a loop.%0A        // if the queue is always empty, it will wait for the last default element of the output,%0A        // after that, the output will be closed and the loop will be finished.%0A        while (!close) %7B%0A          readBatch = queue.poll(5, TimeUnit.MINUTES);%0A          if (readBatch == null) %7B%0A            LOG.warn(%22try to poll a row batch again.%22);%0A          %7D else %7B%0A            // if readBatch is not null, break this loop to continue%0A            break;%0A          %7D%0A        %7D%0A%0A        // when the output is closed and readBatch is null, should poll a batch immediately.%0A        // it is a double-checking also of the last poll operation. in some cases, the output is%0A        // closed and the readBatch is null, but the queue is not empty, contain the last load%0A        // batch or the last default batch.%0A        if (close && readBatch == null) %7B%0A          LOG.warn(%22try to poll a row batch one more time.%22);%0A          readBatch = queue.poll();%0A        %7D%0A        // checking again whether readBatch is null or not%0A
     
@@ -3515,26 +3515,24 @@
 %7B%0A          
-  
 return false
@@ -3533,28 +3533,16 @@
  false;%0A
-          %7D%0A
         
@@ -4648,34 +4648,16 @@
 ;%0A    %7D%0A
-    close = true;%0A
     // I
@@ -4848,24 +4848,110 @@
     %7D%0A    %7D%0A
+    // after try to add the default last element, close the output.%0A    close = true;%0A
   %7D%0A%0A  priva

Avoid blocking the loading task when the output queue poll timeout