<commit_before>package edu.berkeley.cs.succinct.impl

import edu.berkeley.cs.succinct.{SuccinctIndexedFile, SuccinctRDD}
import org.apache.spark.OneToOneDependency
import org.apache.spark.rdd.RDD
import org.apache.spark.storage.StorageLevel

/**
 * Implementation of SuccinctRDD.
 *
 * @constructor Create a new SuccinctRDDImpl from an RDD of partitions (SuccinctIndexedBuffer) and the target
 *              storage level.
 * @param partitionsRDD The input RDD of partitions.
 * @param targetStorageLevel The storage level for the RDD.
 */
class SuccinctRDDImpl private[succinct](
    val partitionsRDD: RDD[SuccinctIndexedFile],
    val targetStorageLevel: StorageLevel = StorageLevel.MEMORY_ONLY)
  extends SuccinctRDD(partitionsRDD.context, List(new OneToOneDependency(partitionsRDD))) {

  val partitionRanges = partitionsRDD.map(_.getFileRange).collect.sorted

  /** Set the name for the RDD; By default set to "SuccinctRDD" */
  override def setName(_name: String): this.type = {
    if (partitionsRDD.name != null) {
      partitionsRDD.setName(partitionsRDD.name + ", " + _name)
    } else {
      partitionsRDD.setName(_name)
    }
    this
  }

  setName("SuccinctRDD")

  /**
   * Persists the Succinct partitions at the specified storage level, ignoring any existing target
   * storage level.
   */
  override def persist(newLevel: StorageLevel): this.type = {
    partitionsRDD.persist(newLevel)
    this
  }

  /** Un-persists the Succinct partitions using the specified blocking mode. */
  override def unpersist(blocking: Boolean = true): this.type = {
    partitionsRDD.unpersist(blocking)
    this
  }

  /** Persists the Succinct partitions at `targetStorageLevel`, which defaults to MEMORY_ONLY. */
  override def cache(): this.type = {
    partitionsRDD.persist(targetStorageLevel)
    this
  }

  /** Extract data from an RDD **/
  override def extract(offset: Long, length: Int): Array[Byte] = {
    val startPartitionRange = partitionRanges.filter(_.contains(offset))(0)
    val endPartitionRange = partitionRanges.filter(_.contains(offset + length))(0)
    if (startPartitionRange == endPartitionRange) {<commit_msg>Error handling in extract<commit_after>package edu.berkeley.cs.succinct.impl

import edu.berkeley.cs.succinct.{SuccinctIndexedFile, SuccinctRDD}
import org.apache.spark.OneToOneDependency
import org.apache.spark.rdd.RDD
import org.apache.spark.storage.StorageLevel

/**
 * Implementation of SuccinctRDD.
 *
 * @constructor Create a new SuccinctRDDImpl from an RDD of partitions (SuccinctIndexedBuffer) and the target
 *              storage level.
 * @param partitionsRDD The input RDD of partitions.
 * @param targetStorageLevel The storage level for the RDD.
 */
class SuccinctRDDImpl private[succinct](
    val partitionsRDD: RDD[SuccinctIndexedFile],
    val targetStorageLevel: StorageLevel = StorageLevel.MEMORY_ONLY)
  extends SuccinctRDD(partitionsRDD.context, List(new OneToOneDependency(partitionsRDD))) {

  val partitionRanges = partitionsRDD.map(_.getFileRange).collect.sorted

  /** Set the name for the RDD; By default set to "SuccinctRDD" */
  override def setName(_name: String): this.type = {
    if (partitionsRDD.name != null) {
      partitionsRDD.setName(partitionsRDD.name + ", " + _name)
    } else {
      partitionsRDD.setName(_name)
    }
    this
  }

  setName("SuccinctRDD")

  /**
   * Persists the Succinct partitions at the specified storage level, ignoring any existing target
   * storage level.
   */
  override def persist(newLevel: StorageLevel): this.type = {
    partitionsRDD.persist(newLevel)
    this
  }

  /** Un-persists the Succinct partitions using the specified blocking mode. */
  override def unpersist(blocking: Boolean = true): this.type = {
    partitionsRDD.unpersist(blocking)
    this
  }

  /** Persists the Succinct partitions at `targetStorageLevel`, which defaults to MEMORY_ONLY. */
  override def cache(): this.type = {
    partitionsRDD.persist(targetStorageLevel)
    this
  }

  /** Extract data from an RDD **/
  override def extract(offset: Long, length: Int): Array[Byte] = {
    val startPartitionRanges = partitionRanges.filter(_.contains(offset))
    val endPartitionRanges = partitionRanges.filter(_.contains(offset + length))

    if (startPartitionRanges.size != 1) {
      throw new ArrayIndexOutOfBoundsException("Invalid offset = " + offset)
    }

    if (endPartitionRanges.size != 1) {
      throw new ArrayIndexOutOfBoundsException("Invalid length = " + length)
    }

    val startPartitionRange = startPartitionRanges(0)
    val endPartitionRange = endPartitionRanges(0)

    // TODO: Handle case where extracted data spans more than 2 partitions

    if (startPartitionRange == endPartitionRange) {
{'fraction_non_alphanumeric': 0.09585272958104105, 'fraction_numerical': 0.0014811680067710537, 'mean_word_length': 4.034078807241746, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '3438095', 'n_tokens_mistral': 1477, 'n_tokens_neox': 1414, 'n_words': 364}