exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/xsort/SpilledRuns.java - drill - Git at Google

 /*
  * 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.drill.exec.physical.impl.xsort;

 import java.io.IOException;
 import java.util.LinkedList;
 import java.util.List;

 import org.apache.drill.common.AutoCloseables;
 import org.apache.drill.common.exceptions.UserException;
 import org.apache.drill.exec.ops.OperatorContext;
 import org.apache.drill.exec.physical.impl.spill.SpillSet;
 import org.apache.drill.exec.physical.impl.xsort.SortImpl.SortResults;
 import org.apache.drill.exec.record.BatchSchema;
 import org.apache.drill.exec.record.VectorContainer;
 import org.apache.drill.exec.record.VectorInitializer;

 import com.google.common.collect.Lists;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 /**
  * Represents the set of spilled batches, including methods to spill and/or
  * merge a set of batches to produce a new spill file.
  */
 public class SpilledRuns {
   private static final Logger logger = LoggerFactory.getLogger(SpilledRuns.class);

   /**
    * Manages the set of spill directories and files.
    */
   private final SpillSet spillSet;
   private final LinkedList<SpilledRun> spilledRuns = Lists.newLinkedList();

   /**
    * Manages the copier used to merge a collection of batches into
    * a new set of batches.
    */
   private final PriorityQueueCopierWrapper copierHolder;
   private BatchSchema schema;

   private final OperatorContext context;

   public SpilledRuns(OperatorContext opContext, SpillSet spillSet, PriorityQueueCopierWrapper copier) {
     this.context = opContext;
     this.spillSet = spillSet;
     copierHolder = copier;
   }

   public void setSchema(BatchSchema schema) {
     this.schema = schema;
     for (BatchGroup b : spilledRuns) {
       b.setSchema(schema);
     }
     copierHolder.close();
   }

   public int size() { return spilledRuns.size(); }
   public boolean hasSpilled() { return spillSet.hasSpilled(); }
   public long getWriteBytes() { return spillSet.getWriteBytes(); }

   public static List<BatchGroup> prepareSpillBatches(LinkedList<? extends BatchGroup> source, int spillCount) {
     List<BatchGroup> batchesToSpill = Lists.newArrayList();
     spillCount = Math.min(source.size(), spillCount);
     assert spillCount > 0 : "Spill count to mergeAndSpill must not be zero";
     for (int i = 0; i < spillCount; i++) {
       batchesToSpill.add(source.pollFirst());
     }
     return batchesToSpill;
   }

   public void mergeAndSpill(List<BatchGroup> batchesToSpill, int spillBatchRowCount, VectorInitializer allocHelper) {
     spilledRuns.add(safeMergeAndSpill(batchesToSpill, spillBatchRowCount, allocHelper));
     logger.trace("Completed spill: memory = {}",
         context.getAllocator().getAllocatedMemory());
   }

   public void mergeRuns(int targetCount, long mergeMemoryPool,
                   int spillBatchRowCount, VectorInitializer allocHelper) {

     long allocated = context.getAllocator().getAllocatedMemory();
     mergeMemoryPool -= context.getAllocator().getAllocatedMemory();
     logger.trace("Merging {} on-disk runs, alloc. memory = {}, avail. memory = {}",
         targetCount, allocated, mergeMemoryPool);

     // Determine the number of runs to merge. The count should be the
     // target count. However, to prevent possible memory overrun, we
     // double-check with actual spill batch size and only spill as much
     // as fits in the merge memory pool.

     int mergeCount = 0;
     long mergeSize = 0;
     for (SpilledRun run : spilledRuns) {
       long batchSize = run.getBatchSize();
       if (mergeSize + batchSize > mergeMemoryPool) {
         break;
       }
       mergeSize += batchSize;
       mergeCount++;
       if (mergeCount == targetCount) {
         break;
       }
     }

     // Must always spill at least 2, even if this creates an over-size
     // spill file. But, if this is a final consolidation, we may have only
     // a single batch.

     mergeCount = Math.max(mergeCount, 2);
     mergeCount = Math.min(mergeCount, spilledRuns.size());

     // Do the actual spill.

     List<BatchGroup> batchesToSpill = prepareSpillBatches(spilledRuns, mergeCount);
     mergeAndSpill(batchesToSpill, spillBatchRowCount, allocHelper);
   }

   private SpilledRun safeMergeAndSpill(List<? extends BatchGroup> batchesToSpill, int spillBatchRowCount, VectorInitializer allocHelper) {
     try {
       return doMergeAndSpill(batchesToSpill, spillBatchRowCount, allocHelper);
     }
     // If error is a User Exception, just use as is.

     catch (UserException ue) { throw ue; }
     catch (Throwable ex) {
       throw UserException.resourceError(ex)
             .message("External Sort encountered an error while spilling to disk")
             .build(logger);
     }
   }

   private SpilledRun doMergeAndSpill(List<? extends BatchGroup> batchesToSpill,
                         int spillBatchRowCount, VectorInitializer allocHelper) throws Throwable {

     // Merge the selected set of matches and write them to the
     // spill file. After each write, we release the memory associated
     // with the just-written batch.

     String outputFile = spillSet.getNextSpillFile();
     SpilledRun newGroup = null;
     VectorContainer dest = new VectorContainer();
     try (AutoCloseable ignored = AutoCloseables.all(batchesToSpill);
          PriorityQueueCopierWrapper.BatchMerger merger = copierHolder.startMerge(schema, batchesToSpill,
                                          dest, spillBatchRowCount, allocHelper)) {
       newGroup = new SpilledRun(spillSet, outputFile, context.getAllocator());
       logger.trace("Spilling {} batches, into spill batches of {} rows, to {}",
           batchesToSpill.size(), spillBatchRowCount, outputFile);

       // The copier will merge records from the buffered batches into
       // the outputContainer up to targetRecordCount number of rows.
       // The actual count may be less if fewer records are available.

       while (merger.next()) {

         // Add a new batch of records (given by merger.getOutput()) to the spill
         // file.
         //
         // note that addBatch also clears the merger's output container

         newGroup.spillBatch(dest);
       }
       context.injectChecked(ExternalSortBatch.INTERRUPTION_WHILE_SPILLING, IOException.class);
       newGroup.closeWriter();
       logger.trace("Spilled {} output batches, each of {} bytes, {} records, to {}",
                    merger.getBatchCount(), merger.getEstBatchSize(),
                    spillBatchRowCount, outputFile);
       newGroup.setBatchSize(merger.getEstBatchSize());
       return newGroup;
     } catch (Throwable e) {
       // we only need to clean up newGroup if spill failed
       try {
         if (newGroup != null) {
           AutoCloseables.close(e, newGroup);
         }
       } catch (Throwable t) { /* close() may hit the same IO issue; just ignore */ }

       throw e;
     }
   }

   public SortResults finalMerge(List<? extends BatchGroup> bufferedBatches,
                     VectorContainer container, int mergeRowCount, VectorInitializer allocHelper) {
     List<BatchGroup> allBatches = new LinkedList<>();
     allBatches.addAll(bufferedBatches);
     bufferedBatches.clear();
     allBatches.addAll(spilledRuns);
     spilledRuns.clear();
     logger.debug("Starting merge phase. Runs = {}, Alloc. memory = {}",
         allBatches.size(), context.getAllocator().getAllocatedMemory());
     return copierHolder.startMerge(schema, allBatches, container, mergeRowCount, allocHelper);
   }

   public void close() {
     if (spillSet.getWriteBytes() > 0) {
       logger.debug("End of sort. Total write bytes: {}, Total read bytes: {}",
                    spillSet.getWriteBytes(), spillSet.getWriteBytes());
     }
     AutoCloseables.closeWithUserException(() -> BatchGroup.closeAll(spilledRuns),
         copierHolder::close, spillSet::close);
   }
 }
	/*
	* 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.drill.exec.physical.impl.xsort;

	import java.io.IOException;
	import java.util.LinkedList;
	import java.util.List;

	import org.apache.drill.common.AutoCloseables;
	import org.apache.drill.common.exceptions.UserException;
	import org.apache.drill.exec.ops.OperatorContext;
	import org.apache.drill.exec.physical.impl.spill.SpillSet;
	import org.apache.drill.exec.physical.impl.xsort.SortImpl.SortResults;
	import org.apache.drill.exec.record.BatchSchema;
	import org.apache.drill.exec.record.VectorContainer;
	import org.apache.drill.exec.record.VectorInitializer;

	import com.google.common.collect.Lists;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	/**
	* Represents the set of spilled batches, including methods to spill and/or
	* merge a set of batches to produce a new spill file.
	*/
	public class SpilledRuns {
	private static final Logger logger = LoggerFactory.getLogger(SpilledRuns.class);

	/**
	* Manages the set of spill directories and files.
	*/
	private final SpillSet spillSet;
	private final LinkedList<SpilledRun> spilledRuns = Lists.newLinkedList();

	/**
	* Manages the copier used to merge a collection of batches into
	* a new set of batches.
	*/
	private final PriorityQueueCopierWrapper copierHolder;
	private BatchSchema schema;

	private final OperatorContext context;

	public SpilledRuns(OperatorContext opContext, SpillSet spillSet, PriorityQueueCopierWrapper copier) {
	this.context = opContext;
	this.spillSet = spillSet;
	copierHolder = copier;
	}

	public void setSchema(BatchSchema schema) {
	this.schema = schema;
	for (BatchGroup b : spilledRuns) {
	b.setSchema(schema);
	}
	copierHolder.close();
	}

	public int size() { return spilledRuns.size(); }
	public boolean hasSpilled() { return spillSet.hasSpilled(); }
	public long getWriteBytes() { return spillSet.getWriteBytes(); }

	public static List<BatchGroup> prepareSpillBatches(LinkedList<? extends BatchGroup> source, int spillCount) {
	List<BatchGroup> batchesToSpill = Lists.newArrayList();
	spillCount = Math.min(source.size(), spillCount);
	assert spillCount > 0 : "Spill count to mergeAndSpill must not be zero";
	for (int i = 0; i < spillCount; i++) {
	batchesToSpill.add(source.pollFirst());
	}
	return batchesToSpill;
	}

	public void mergeAndSpill(List<BatchGroup> batchesToSpill, int spillBatchRowCount, VectorInitializer allocHelper) {
	spilledRuns.add(safeMergeAndSpill(batchesToSpill, spillBatchRowCount, allocHelper));
	logger.trace("Completed spill: memory = {}",
	context.getAllocator().getAllocatedMemory());
	}

	public void mergeRuns(int targetCount, long mergeMemoryPool,
	int spillBatchRowCount, VectorInitializer allocHelper) {

	long allocated = context.getAllocator().getAllocatedMemory();
	mergeMemoryPool -= context.getAllocator().getAllocatedMemory();
	logger.trace("Merging {} on-disk runs, alloc. memory = {}, avail. memory = {}",
	targetCount, allocated, mergeMemoryPool);

	// Determine the number of runs to merge. The count should be the
	// target count. However, to prevent possible memory overrun, we
	// double-check with actual spill batch size and only spill as much
	// as fits in the merge memory pool.

	int mergeCount = 0;
	long mergeSize = 0;
	for (SpilledRun run : spilledRuns) {
	long batchSize = run.getBatchSize();
	if (mergeSize + batchSize > mergeMemoryPool) {
	break;
	}
	mergeSize += batchSize;
	mergeCount++;
	if (mergeCount == targetCount) {
	break;
	}
	}

	// Must always spill at least 2, even if this creates an over-size
	// spill file. But, if this is a final consolidation, we may have only
	// a single batch.

	mergeCount = Math.max(mergeCount, 2);
	mergeCount = Math.min(mergeCount, spilledRuns.size());

	// Do the actual spill.

	List<BatchGroup> batchesToSpill = prepareSpillBatches(spilledRuns, mergeCount);
	mergeAndSpill(batchesToSpill, spillBatchRowCount, allocHelper);
	}

	private SpilledRun safeMergeAndSpill(List<? extends BatchGroup> batchesToSpill, int spillBatchRowCount, VectorInitializer allocHelper) {
	try {
	return doMergeAndSpill(batchesToSpill, spillBatchRowCount, allocHelper);
	}
	// If error is a User Exception, just use as is.

	catch (UserException ue) { throw ue; }
	catch (Throwable ex) {
	throw UserException.resourceError(ex)
	.message("External Sort encountered an error while spilling to disk")
	.build(logger);
	}
	}

	private SpilledRun doMergeAndSpill(List<? extends BatchGroup> batchesToSpill,
	int spillBatchRowCount, VectorInitializer allocHelper) throws Throwable {

	// Merge the selected set of matches and write them to the
	// spill file. After each write, we release the memory associated
	// with the just-written batch.

	String outputFile = spillSet.getNextSpillFile();
	SpilledRun newGroup = null;
	VectorContainer dest = new VectorContainer();
	try (AutoCloseable ignored = AutoCloseables.all(batchesToSpill);
	PriorityQueueCopierWrapper.BatchMerger merger = copierHolder.startMerge(schema, batchesToSpill,
	dest, spillBatchRowCount, allocHelper)) {
	newGroup = new SpilledRun(spillSet, outputFile, context.getAllocator());
	logger.trace("Spilling {} batches, into spill batches of {} rows, to {}",
	batchesToSpill.size(), spillBatchRowCount, outputFile);

	// The copier will merge records from the buffered batches into
	// the outputContainer up to targetRecordCount number of rows.
	// The actual count may be less if fewer records are available.

	while (merger.next()) {

	// Add a new batch of records (given by merger.getOutput()) to the spill
	// file.
	//
	// note that addBatch also clears the merger's output container

	newGroup.spillBatch(dest);
	}
	context.injectChecked(ExternalSortBatch.INTERRUPTION_WHILE_SPILLING, IOException.class);
	newGroup.closeWriter();
	logger.trace("Spilled {} output batches, each of {} bytes, {} records, to {}",
	merger.getBatchCount(), merger.getEstBatchSize(),
	spillBatchRowCount, outputFile);
	newGroup.setBatchSize(merger.getEstBatchSize());
	return newGroup;
	} catch (Throwable e) {
	// we only need to clean up newGroup if spill failed
	try {
	if (newGroup != null) {
	AutoCloseables.close(e, newGroup);
	}
	} catch (Throwable t) { /* close() may hit the same IO issue; just ignore */ }

	throw e;
	}
	}

	public SortResults finalMerge(List<? extends BatchGroup> bufferedBatches,
	VectorContainer container, int mergeRowCount, VectorInitializer allocHelper) {
	List<BatchGroup> allBatches = new LinkedList<>();
	allBatches.addAll(bufferedBatches);
	bufferedBatches.clear();
	allBatches.addAll(spilledRuns);
	spilledRuns.clear();
	logger.debug("Starting merge phase. Runs = {}, Alloc. memory = {}",
	allBatches.size(), context.getAllocator().getAllocatedMemory());
	return copierHolder.startMerge(schema, allBatches, container, mergeRowCount, allocHelper);
	}

	public void close() {
	if (spillSet.getWriteBytes() > 0) {
	logger.debug("End of sort. Total write bytes: {}, Total read bytes: {}",
	spillSet.getWriteBytes(), spillSet.getWriteBytes());
	}
	AutoCloseables.closeWithUserException(() -> BatchGroup.closeAll(spilledRuns),
	copierHolder::close, spillSet::close);
	}
	}