<commit_before>package nsmc.sql

import com.mongodb.DBObject
import nsmc._
import nsmc.conversion.types.{StructureType}
import nsmc.conversion.{RecordConverter, SchemaAccumulator}

import org.apache.spark.rdd.RDD
import org.apache.spark.sql.SQLContext
import org.apache.spark.sql.catalyst.expressions.Row
import org.apache.spark.sql.catalyst.types.StructType
import org.apache.spark.sql.sources.{TableScan, Filter, PrunedFilteredScan, RelationProvider}

import scala.collection.Iterator

object PartitionRecordConverter {
  // each partition gets its own converter as there's no reason for the record
  // converters to communicate with each other
  def convert(internalSchema: StructureType)(in: Iterator[DBObject]) : Iterator[Row] = {
    val rc = new RecordConverter(internalSchema)
    in.map(mo => rc.getSchemaRecord(mo))
  }
}

case class MongoTableScan(database: String, collection: String)
                        (@transient val sqlContext: SQLContext)
  extends TableScan with Logging {

  // TODO: think the RDD lifecycle through carefully here

  logInfo(s"Registering MongoDB collection '$collection' in database '$database'")

  private val data = sqlContext.sparkContext.mongoCollection[DBObject](database, collection)

  private val partitionSchemas = data.mapPartitions(inferType, preservesPartitioning = true)

  val partialSchemas = partitionSchemas.collect()

  private val accum = new SchemaAccumulator()
  accum.accumulate(partialSchemas.iterator)

  private val inferredSchema = accum.getSchema

  logDebug(s"Computed schema for collection '$collection' in database '$database'")

  private val internalSchema = accum.getInternal

  val schema: StructType = StructType(inferredSchema)

  def inferType(in: Iterator[DBObject]) : Iterator[StructureType] = {
    val accum = new SchemaAccumulator()
    in.foreach(mo => accum.considerDatum(mo))
    Seq(accum.getInternal.asInstanceOf[StructureType]).iterator
  }

  def buildScan: RDD[Row] = {
    val schema = internalSchema
    val converter = PartitionRecordConverter.convert(schema.asInstanceOf[StructureType]) _
    data.mapPartitions(converter, preservesPartitioning = true)
  }
}

class MongoRelationProvider extends RelationProvider {
  def createRelation(sqlContext: SQLContext, parameters: Map[String, String]) = {
    MongoTableScan(parameters("db"), parameters("collection"))(sqlContext)
  }
}
<commit_msg>Minimal hack to implement PrunedFilteredScan for SQL integration.<commit_after>package nsmc.sql

import com.mongodb.DBObject
import nsmc._
import nsmc.conversion.types.{StructureType}
import nsmc.conversion.{RecordConverter, SchemaAccumulator}

import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{StructField, SQLContext}
import org.apache.spark.sql.catalyst.expressions.Row
import org.apache.spark.sql.catalyst.types.StructType
import org.apache.spark.sql.sources.{TableScan, Filter, PrunedFilteredScan, RelationProvider}

import scala.collection.Iterator
import scala.collection.immutable.HashMap

object PartitionRecordConverter {
  // each partition gets its own converter as there's no reason for the record
  // converters to communicate with each other
  def convert(internalSchema: StructureType)(in: Iterator[DBObject]) : Iterator[Row] = {
    val rc = new RecordConverter(internalSchema)
    in.map(mo => rc.getSchemaRecord(mo))
  }
}

object RowProjector {
  def projectRow(wholeRow: Row,
                         positionalMap: Map[String, Int],
                         requiredColumns: Array[String]): Row = {
    val projected = requiredColumns.map(colName => {
      val pos = positionalMap(colName)
      wholeRow(pos)
    })
    Row(projected:_*)
  }
}

case class MongoTableScan(database: String, collection: String)
                        (@transient val sqlContext: SQLContext)
  extends PrunedFilteredScan with Logging {

  // TODO: think the RDD lifecycle through carefully here

  logInfo(s"Registering MongoDB collection '$collection' in database '$database'")

  private val data = sqlContext.sparkContext.mongoCollection[DBObject](database, collection)

  private val partitionSchemas = data.mapPartitions(inferType, preservesPartitioning = true)

  val partialSchemas = partitionSchemas.collect()

  private val accum = new SchemaAccumulator()
  accum.accumulate(partialSchemas.iterator)

  private val inferredSchema = accum.getSchema

  logDebug(s"Computed schema for collection '$collection' in database '$database'")

  private val internalSchema = accum.getInternal

  val schema: StructType = StructType(inferredSchema)

  def inferType(in: Iterator[DBObject]) : Iterator[StructureType] = {
    val accum = new SchemaAccumulator()
    in.foreach(mo => accum.considerDatum(mo))
    Seq(accum.getInternal.asInstanceOf[StructureType]).iterator
  }

  private def makePositionalMap(fields: Seq[StructField]) : Map[String, Int] = {
    HashMap[String, Int](fields.map(f => f.name).zipWithIndex:_*)
  }

  def buildScan(requiredColumns: Array[String], filters: Array[Filter]): RDD[Row] = {
    val schema = internalSchema
    val converter = PartitionRecordConverter.convert(schema.asInstanceOf[StructureType]) _
    val allData = data.mapPartitions(converter, preservesPartitioning = true)
    val positionalMap = makePositionalMap(inferredSchema)
    val projected = allData.map(r => RowProjector.projectRow(r, positionalMap, requiredColumns))
    projected
  }
}

class MongoRelationProvider extends RelationProvider {
  def createRelation(sqlContext: SQLContext, parameters: Map[String, String]) = {
    MongoTableScan(parameters("db"), parameters("collection"))(sqlContext)
  }
}

{'fraction_non_alphanumeric': 0.09552664409196222, 'fraction_numerical': 0.0, 'mean_word_length': 4.913593256059009, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 10, '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': '5707901', 'n_tokens_mistral': 1649, 'n_tokens_neox': 1562, 'n_words': 377}