<commit_before>package edu.umass.cs.ciir.waltz.coders.sorter;

import ciir.jfoley.chai.collections.util.Comparing;
import ciir.jfoley.chai.collections.util.MapFns;
import ciir.jfoley.chai.collections.util.QuickSort;
import ciir.jfoley.chai.fn.SinkFn;
import ciir.jfoley.chai.io.IO;
import edu.umass.cs.ciir.waltz.coders.Coder;
import edu.umass.cs.ciir.waltz.coders.kinds.FixedSize;

import java.io.*;
import java.util.*;

/**
 * Break External Sorting into two steps:
 *
 * The writing writes sorted runs to files within a directory.
 * The reading stage combines these while reading in order.
 *
 * <ol>
 *   <li>{link: ExternalSortingWriter} Writing (this class)</li>
 *   <li>{link: ExternalSortingReader} Reading</li>
 * </ol>
 * @author jfoley.
 */
public class ExternalSortingWriter<T> implements Flushable, Closeable, SinkFn<T> {
  private final File dir;
  final Coder<Long> countCoder;
  final Coder<T> objCoder;
  final Comparator<T> cmp;
  private final int maxItemsInMemory;
  private final int mergeFactor;
  private final ArrayList<T> buffer;
  private int nextId;
  Map<Integer, List<Integer>> runsByLevel = new HashMap<>();
  private int maxLevelRuns;

  public ExternalSortingWriter(File dir, Coder<T> coder) {
    this(dir, coder, Comparing.defaultComparator());
  }
  public ExternalSortingWriter(File dir, Coder<T> coder, Comparator<T> comparator) {
    // Leverage some defaults from Galago, because why not?
    this(dir, coder, comparator, 50 * 1024 * 1024, 10);
  }
  public ExternalSortingWriter(File dir, Coder<T> coder, Comparator<T> comparator, int maxItemsInMemory, int mergeFactor) {
    assert(dir.isDirectory());
    this.dir = dir;
    this.countCoder = FixedSize.longs;
    this.objCoder = coder;
    this.cmp = comparator;
    this.maxItemsInMemory = maxItemsInMemory;
    this.mergeFactor = mergeFactor;
    this.buffer = new ArrayList<>(maxItemsInMemory);
    this.nextId = 0;
    this.maxLevelRuns = 0;
  }

  @Override
  public void close() throws IOException {
    // push all runs to topmost level:
    flush();
  }

  /**
   * Push items out of memory buffer and onto disk.
   * @throws IOException
   */
  @Override
  public void flush() throws IOException {
    int currentId = nextId++;
    try (OutputStream output = IO.openOutputStream(nameForId(currentId).getAbsolutePath())) {
      QuickSort.sort(cmp, buffer);
      // write run to file.
      countCoder.write(output, (long) buffer.size());
      for (T t : buffer) {
        objCoder.write(output, t);
      }
      // ditch memory as soon as possible.
      buffer.clear();
      // add to lowest rung of runs collection.
      MapFns.extendListInMap(runsByLevel, 0, currentId);
    }
    // check and see if we need to mergeRuns()
    mergeRuns();
  }


  public static class RunReader<T> implements Closeable, Comparable<RunReader<T>> {
    private final InputStream source;
    private final long count;
    private final Coder<T> itemCoder;
    private final Comparator<T> cmp;
    private T nextItem;
    private int index;

    public RunReader(Comparator<T> cmp, Coder<Long> countCoder, Coder<T> itemCoder, InputStream source) {
      this.cmp = cmp;
      this.itemCoder = itemCoder;
      this.source = source;
      this.count = countCoder.read(source);
      index = 0;
      next();
    }

    public T peek() {
      return nextItem;
    }

    public boolean hasNext() {
      return nextItem != null;
    }

    public T next() {
      T lastItem = nextItem;
      if(index < count) {
        this.nextItem = itemCoder.read(source);<commit_msg>Implement a bunch of readers<commit_after>package edu.umass.cs.ciir.waltz.coders.sorter;

import ciir.jfoley.chai.collections.iters.ClosingIterator;
import ciir.jfoley.chai.collections.util.Comparing;
import ciir.jfoley.chai.collections.util.MapFns;
import ciir.jfoley.chai.collections.util.QuickSort;
import ciir.jfoley.chai.fn.SinkFn;
import ciir.jfoley.chai.io.FS;
import ciir.jfoley.chai.io.IO;
import edu.umass.cs.ciir.waltz.coders.Coder;
import edu.umass.cs.ciir.waltz.coders.kinds.FixedSize;

import java.io.*;
import java.util.*;

/**
 * Break External Sorting into two steps:
 *
 * The writing writes sorted runs to files within a directory.
 * The reading stage combines these while reading in order.
 *
 * <ol>
 *   <li>{link: ExternalSortingWriter} Writing (this class)</li>
 *   <li>{link: ExternalSortingReader} Reading</li>
 * </ol>
 * @author jfoley.
 */
public class ExternalSortingWriter<T> implements Flushable, Closeable, SinkFn<T> {
  private final File dir;
  final Coder<Long> countCoder;
  final Coder<T> objCoder;
  final Comparator<T> cmp;
  private final int maxItemsInMemory;
  private final int mergeFactor;
  private final ArrayList<T> buffer;
  private int nextId;
  Map<Integer, List<Integer>> runsByLevel = new HashMap<>();
  private int maxLevelRuns;

  public ExternalSortingWriter(File dir, Coder<T> coder) {
    this(dir, coder, Comparing.defaultComparator());
  }
  public ExternalSortingWriter(File dir, Coder<T> coder, Comparator<T> comparator) {
    // Leverage some defaults from Galago, because why not?
    this(dir, coder, comparator, 50 * 1024 * 1024, 10);
  }
  public ExternalSortingWriter(File dir, Coder<T> coder, Comparator<T> comparator, int maxItemsInMemory, int mergeFactor) {
    assert(dir.isDirectory());
    this.dir = dir;
    this.countCoder = FixedSize.longs;
    this.objCoder = coder;
    this.cmp = comparator;
    this.maxItemsInMemory = maxItemsInMemory;
    this.mergeFactor = mergeFactor;
    this.buffer = new ArrayList<>(maxItemsInMemory);
    this.nextId = 0;
    this.maxLevelRuns = 0;
  }

  @Override
  public void close() throws IOException {
    // push all runs to topmost level:
    flush();
  }

  public SortDirectory<T> getOutput() throws IOException {
    return new SortDirectory<>(dir, cmp, countCoder, objCoder);
  }

  /**
   * Push items out of memory buffer and onto disk.
   * @throws IOException
   */
  @Override
  public void flush() throws IOException {
    int currentId = nextId++;
    try (OutputStream output = IO.openOutputStream(nameForId(currentId).getAbsolutePath())) {
      QuickSort.sort(cmp, buffer);
      // write run to file.
      countCoder.write(output, (long) buffer.size());
      for (T t : buffer) {
        objCoder.write(output, t);
      }
      // ditch memory as soon as possible.
      buffer.clear();
      // add to lowest rung of runs collection.
      MapFns.extendListInMap(runsByLevel, 0, currentId);
    }
    // check and see if we need to mergeRuns()
    mergeRuns();
  }

  public static class SortDirectory<T> implements Iterable<T> {
    private final File dir;
    private final Comparator<T> cmp;
    private final Coder<Long> countCoder;
    private final Coder<T> itemCoder;

    public SortDirectory(File dir, Comparator<T> cmp, Coder<Long> countCoder, Coder<T> itemCoder) throws IOException {
      this.dir = dir;
      this.cmp = cmp;
      this.countCoder = countCoder;
      this.itemCoder = itemCoder;
    }

    @Override
    public MergingRunReader<T> iterator() {
      try {
        return MergingRunReader.openDirectory(dir, cmp, countCoder, itemCoder);
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
    }
  }

  public static class MergingRunReader<T> implements Closeable, ClosingIterator<T> {
    private final PriorityQueue<RunReader<T>> queue;

    public MergingRunReader(Collection<RunReader<T>> readers) {
      this.queue = new PriorityQueue<>();
      for (RunReader<T> reader : readers) {
        queue.offer(reader);
      }
    }

    @Override
    public boolean hasNext() {
      return queue.peek().hasNext();
    }

    @Override
    public T next() {
      RunReader<T> minimum = queue.poll();
      T obj = minimum.next();
      if(minimum.hasNext()) {
        queue.offer(minimum);
      } else {
        IO.close(minimum);
      }
      return obj;
    }

    @Override
    public void close() throws IOException {
      for (RunReader<T> tRunReader : queue) {
        tRunReader.close();
      }
      queue.clear();
    }

    public static <T> MergingRunReader<T> openDirectory(File dir, Comparator<T> cmp, Coder<Long> countCoder, Coder<T> itemCoder) throws IOException {
      List<RunReader<T>> readers = new ArrayList<>();
      for (File file : FS.listDirectory(dir)) {
        readers.add(new RunReader<T>(cmp, countCoder, itemCoder, IO.openInputStream(file)));
      }
      return new MergingRunReader<>(readers);
    }
  }

  public static class RunReader<T> implements Closeable, Comparable<RunReader<T>> {
    private final InputStream source;
    private final long count;
    private final Coder<T> itemCoder;
    private final Comparator<T> cmp;
    private T nextItem;
    private int index;

    public RunReader(Comparator<T> cmp, Coder<Long> countCoder, Coder<T> itemCoder, InputStream source) {
      this.cmp = cmp;
      this.itemCoder = itemCoder;
      this.source = source;
      this.count = countCoder.read(source);
      index = 0;
      next();
    }

    public T peek() {
      return nextItem;
    }

    public boolean hasNext() {
      return nextItem != null;
    }

    public T next() {
      T lastItem = nextItem;
      if(index < count) {
        this.nextItem = itemCoder.read(source);
{'fraction_non_alphanumeric': 0.11234262348003873, 'fraction_numerical': 0.0034434520606908426, 'mean_word_length': 3.2264665757162345, 'pattern_counts': {'":': 0, '<': 94, '<?xml version=': 0, '>': 92, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1093304', 'n_tokens_mistral': 3050, 'n_tokens_neox': 2974, 'n_words': 779}