<commit_before>package cz.seznam.euphoria.hadoop.input;

import cz.seznam.euphoria.guava.shaded.com.google.common.collect.AbstractIterator;
import cz.seznam.euphoria.core.client.io.DataSource;
import cz.seznam.euphoria.core.client.io.Partition;
import cz.seznam.euphoria.core.client.io.Reader;
import cz.seznam.euphoria.hadoop.HadoopUtils;
import cz.seznam.euphoria.hadoop.SerializableWritable;
import cz.seznam.euphoria.core.client.util.Pair;
import lombok.SneakyThrows;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableUtils;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;

class HadoopInputFormatAdapter implements DataSource<Pair<?, ?>> {

  private final Class<? extends InputFormat> hadoopFormatCls;
  private final SerializableWritable<Configuration> conf;

  private transient InputFormat<?, ?> hadoopFormatInstance;

  public HadoopInputFormatAdapter(Class<? extends InputFormat> hadoopFormatCls,
                                  SerializableWritable<Configuration> conf)
  {
    this.hadoopFormatCls = Objects.requireNonNull(hadoopFormatCls);
    this.conf = Objects.requireNonNull(conf);
  }

  @Override
  @SneakyThrows
  public List<Partition<Pair<?, ?>>> getPartitions() {
    Configuration c = conf.getWritable();
    return getHadoopFormatInstance()
            .getSplits(HadoopUtils.createJobContext(c))
            .stream()
            .map(split -> new HadoopPartition(hadoopFormatCls, c, split))
            .collect(Collectors.toList());
  }

  @Override
  public boolean isBounded() {
    return true;
  }

 /**
   * Retrieves the instance or create new if not exists.
   * You must always pass a valid configuration object
   * or {@code NullPointerException} might be thrown.
   */
  private InputFormat<?, ?> getHadoopFormatInstance()
          throws InstantiationException, IllegalAccessException
  {
    if (hadoopFormatInstance == null) {
      hadoopFormatInstance = HadoopUtils.instantiateHadoopFormat(
              hadoopFormatCls,
              InputFormat.class,
              conf.getWritable());
    }

    return hadoopFormatInstance;
  }

  /**
   * Wraps Hadoop {@link RecordReader}
   */
  private static class HadoopReader
          extends AbstractIterator<Pair<?, ?>>
          implements Reader<Pair<?, ?>>
  {
    private final RecordReader<?, ?> hadoopReader;
    private final Configuration conf;

    public HadoopReader(RecordReader<?, ?> hadoopReader, Configuration conf) {
      this.hadoopReader = Objects.requireNonNull(hadoopReader);
      this.conf = Objects.requireNonNull(conf);
    }

    @Override
    @SneakyThrows
    protected Pair<?, ?> computeNext() {
      if (hadoopReader.nextKeyValue()) {

        // ~ cast to Writable (may throw ClassCastException
        // when input format doesn't support Writables)
        Writable key = (Writable) hadoopReader.getCurrentKey();
        Writable value = (Writable) hadoopReader.getCurrentValue();

        // ~ clone Writables since they are reused
        // between calls to RecordReader#nextKeyValue
        return Pair.of(WritableUtils.clone(key, conf), WritableUtils.clone(value, conf));
      } else {
        return endOfData();
      }
    }

    @Override
    public void close() throws IOException {
      hadoopReader.close();
    }
  }

  /**
   * Wraps Hadoop {@link InputSplit}
   */
  private static class HadoopPartition implements Partition<Pair<?, ?>> {

    private final Class<? extends InputFormat> hadoopFormatCls;
    private final Configuration conf;
    private final InputSplit hadoopSplit;

    public HadoopPartition(Class<? extends InputFormat> hadoopFormatCls,
                           Configuration conf,
                           InputSplit hadoopSplit)
    {
      this.hadoopFormatCls = Objects.requireNonNull(hadoopFormatCls);
      this.conf = Objects.requireNonNull(conf);
      this.hadoopSplit = Objects.requireNonNull(hadoopSplit);
    }

    @Override
    @SneakyThrows
    public Set<String> getLocations() {
      return Arrays.stream(hadoopSplit.getLocations())
              .collect(Collectors.toSet());
    }

    @Override
    @SneakyThrows
    public Reader<Pair<?, ?>> openReader() throws IOException {
      TaskAttemptContext ctx = HadoopUtils.createTaskContext(conf, 0);
      RecordReader<?, ?> reader =
              HadoopUtils.instantiateHadoopFormat(
                      hadoopFormatCls,
                      InputFormat.class,
                      conf)
                      .createRecordReader(hadoopSplit, ctx);

      reader.initialize(hadoopSplit, ctx);

      return new HadoopReader(reader, conf);
    }
  }

}<commit_msg>#! Type-safe HadoopInputFormatAdapter.<commit_after>package cz.seznam.euphoria.hadoop.input;

import cz.seznam.euphoria.guava.shaded.com.google.common.collect.AbstractIterator;
import cz.seznam.euphoria.core.client.io.DataSource;
import cz.seznam.euphoria.core.client.io.Partition;
import cz.seznam.euphoria.core.client.io.Reader;
import cz.seznam.euphoria.hadoop.HadoopUtils;
import cz.seznam.euphoria.hadoop.SerializableWritable;
import cz.seznam.euphoria.core.client.util.Pair;
import lombok.SneakyThrows;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableUtils;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;

class HadoopInputFormatAdapter<K extends Writable, V extends Writable>
    implements DataSource<Pair<K, V>> {

  private final Class<? extends InputFormat<K, V>> hadoopFormatCls;
  private final SerializableWritable<Configuration> conf;

  private transient InputFormat<K, V> hadoopFormatInstance;

  public HadoopInputFormatAdapter(Class<? extends InputFormat<K, V>> hadoopFormatCls,
                                  SerializableWritable<Configuration> conf)
  {
    this.hadoopFormatCls = Objects.requireNonNull(hadoopFormatCls);
    this.conf = Objects.requireNonNull(conf);
  }

  @Override
  @SneakyThrows
  public List<Partition<Pair<K, V>>> getPartitions() {
    Configuration c = conf.getWritable();
    return getHadoopFormatInstance()
            .getSplits(HadoopUtils.createJobContext(c))
            .stream()
            .map(split -> new HadoopPartition<>(hadoopFormatCls, c, split))
            .collect(Collectors.toList());
  }

  @Override
  public boolean isBounded() {
    return true;
  }

 /**
   * Retrieves the instance or create new if not exists.
   * You must always pass a valid configuration object
   * or {@code NullPointerException} might be thrown.
   */
  private InputFormat<K, V> getHadoopFormatInstance()
          throws InstantiationException, IllegalAccessException
  {
    if (hadoopFormatInstance == null) {
      hadoopFormatInstance = HadoopUtils.instantiateHadoopFormat(
              hadoopFormatCls,
              InputFormat.class,
              conf.getWritable());
    }

    return hadoopFormatInstance;
  }

  /**
   * Wraps Hadoop {@link RecordReader}
   */
  private static class HadoopReader<K extends Writable, V extends Writable>
          extends AbstractIterator<Pair<K, V>>
          implements Reader<Pair<K, V>>
  {
    private final RecordReader<?, ?> hadoopReader;
    private final Configuration conf;

    public HadoopReader(RecordReader<?, ?> hadoopReader, Configuration conf) {
      this.hadoopReader = Objects.requireNonNull(hadoopReader);
      this.conf = Objects.requireNonNull(conf);
    }

    @Override
    @SneakyThrows
    protected Pair<K, V> computeNext() {
      if (hadoopReader.nextKeyValue()) {

        // ~ cast to Writable (may throw ClassCastException
        // when input format doesn't support Writables)
        K key = (K) hadoopReader.getCurrentKey();
        V value = (V) hadoopReader.getCurrentValue();

        // ~ clone Writables since they are reused
        // between calls to RecordReader#nextKeyValue
        return Pair.of(WritableUtils.clone(key, conf), WritableUtils.clone(value, conf));
      } else {
        return endOfData();
      }
    }

    @Override
    public void close() throws IOException {
      hadoopReader.close();
    }
  }

  /**
   * Wraps Hadoop {@link InputSplit}
   */
  private static class HadoopPartition<K extends Writable, V extends Writable>
      implements Partition<Pair<K, V>> {

    private final Class<? extends InputFormat<K, V>> hadoopFormatCls;
    private final Configuration conf;
    private final InputSplit hadoopSplit;

    public HadoopPartition(Class<? extends InputFormat<K, V>> hadoopFormatCls,
                           Configuration conf,
                           InputSplit hadoopSplit)
    {
      this.hadoopFormatCls = Objects.requireNonNull(hadoopFormatCls);
      this.conf = Objects.requireNonNull(conf);
      this.hadoopSplit = Objects.requireNonNull(hadoopSplit);
    }

    @Override
    @SneakyThrows
    public Set<String> getLocations() {
      return Arrays.stream(hadoopSplit.getLocations())
              .collect(Collectors.toSet());
    }

    @Override
    @SneakyThrows
    public Reader<Pair<K, V>> openReader() throws IOException {
      TaskAttemptContext ctx = HadoopUtils.createTaskContext(conf, 0);
      RecordReader<K, V> reader =
              HadoopUtils.instantiateHadoopFormat(
                      hadoopFormatCls,
                      InputFormat.class,
                      conf)
                      .createRecordReader(hadoopSplit, ctx);

      reader.initialize(hadoopSplit, ctx);

      return new HadoopReader<>(reader, conf);
    }
  }

}
{'fraction_non_alphanumeric': 0.0983525697938246, 'fraction_numerical': 0.00019729703067968826, 'mean_word_length': 3.121138211382114, 'pattern_counts': {'":': 0, '<': 64, '<?xml version=': 0, '>': 66, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3147289', 'n_tokens_mistral': 2978, 'n_tokens_neox': 2867, 'n_words': 609}