<commit_before>package com.azavea.rf.batch

import java.util.UUID

import cats.data.NonEmptyList
import cats.data.Validated._
import cats.implicits._
import com.azavea.rf.tool.ast._
import com.azavea.rf.tool.ast.MapAlgebraAST._
import com.azavea.rf.tool.eval._
import com.azavea.rf.tool.eval.Interpreter.Interpreted
import geotrellis.raster.{MultibandTile, Tile}
import geotrellis.spark._
import geotrellis.spark.io._
import geotrellis.spark.io.s3._
import geotrellis.spark.render._
import org.apache.spark.SparkContext
import org.apache.spark.rdd.RDD

// --- //

package object ast {

  /** Perform a binary operation on RDDs, while preserving any metadata they had. */
  private def binary(
    f: (RDD[(SpatialKey, Tile)], RDD[(SpatialKey, Tile)]) => RDD[(SpatialKey, Tile)],
    tlr0: TileLayerRDD[SpatialKey],
    tlr1: TileLayerRDD[SpatialKey]
  ): TileLayerRDD[SpatialKey] = {
    TileLayerRDD(f(tlr0, tlr1), tlr0.metadata.combine(tlr1.metadata))
  }

  /** Evaluate an AST of RDD Sources. Assumes that the AST's
    * [[NodeMetadata]] has already been replaced, if applicable.
    */
  def interpretRDD(
    ast: MapAlgebraAST,
    sourceMapping: Map[UUID, RFMLRaster],
    zoom: Int,
    sceneLocs: Map[UUID, String],
    projLocs: Map[UUID, List[(UUID, String)]]
  )(implicit sc: SparkContext): Interpreted[TileLayerRDD[SpatialKey]] = {

    /* Perform Map Algebra over a validated RDD-filled AST */
    def eval(
      ast: MapAlgebraAST,
      rdds: Map[UUID, TileLayerRDD[SpatialKey]]
    ): TileLayerRDD[SpatialKey] = ast match {
      case Source(id, _) => rdds(id)
      case Addition(args, _, _) =>
        args.map(eval(_, rdds)).reduce((acc,r) => binary({_ + _}, acc, r))
      case Subtraction(args, _, _) =>
        args.map(eval(_, rdds)).reduce((acc,r) => binary({_ - _}, acc, r))
      case Multiplication(args, _, _) =>
        args.map(eval(_, rdds)).reduce((acc,r) => binary({_ * _}, acc, r))
      case Division(args, _, _) =>
        args.map(eval(_, rdds)).reduce((acc,r) => binary({_ / _}, acc, r))
      case Classification(arg :: _, _, _, classMap) =>
        eval(arg, rdds).withContext(_.color(classMap.toColorMap))
      case op => throw new Exception(s"Unimplemented Operation given! ${op}")
    }

    /* Guarantee correctness before performing Map Algebra */
    val pure = Interpreter.interpretPure[Unit](ast, sourceMapping)
    val rdds = sourceMapping.mapValues(r => fetch(r, zoom, sceneLocs, projLocs)).sequence

    (pure |@| rdds).map({ case (_, rs) => eval(ast, rs) })
  }

  /** This requires that for each [[RFMLRaster]] that a band number be specified. */
  private def fetch(
    raster: RFMLRaster,
    zoom: Int,
    sceneLocs: Map[UUID, String],
    projLocs: Map[UUID, List[(UUID, String)]]
  )(implicit sc: SparkContext): Interpreted[TileLayerRDD[SpatialKey]] = raster match {

    case ProjectRaster(id, None) => Invalid(NonEmptyList.of(NoBandGiven(id)))
    case ProjectRaster(id, Some(band)) => getStores(id, projLocs) match {
      case None => Invalid(NonEmptyList.of(AttributeStoreFetchError(id)))
      case Some(stores) => {
        val rdds: List[TileLayerRDD[SpatialKey]] =
          stores.map({ case (sceneId, store) =>
            S3LayerReader(store)
              .read[SpatialKey, MultibandTile, TileLayerMetadata[SpatialKey]](LayerId(sceneId.toString, zoom))
              .withContext(rdd => rdd.mapValues(_.band(band)))
          })
<commit_msg>Add Min/Max support in RDD AST exports<commit_after>package com.azavea.rf.batch

import java.util.UUID

import cats.data.NonEmptyList
import cats.data.Validated._
import cats.implicits._
import com.azavea.rf.tool.ast._
import com.azavea.rf.tool.ast.MapAlgebraAST._
import com.azavea.rf.tool.eval._
import com.azavea.rf.tool.eval.Interpreter.Interpreted
import geotrellis.raster.{MultibandTile, Tile}
import geotrellis.spark._
import geotrellis.spark.io._
import geotrellis.spark.io.s3._
import geotrellis.spark.render._
import org.apache.spark.SparkContext
import org.apache.spark.rdd.RDD

// --- //

package object ast {

  /** Perform a binary operation on RDDs, while preserving any metadata they had. */
  private def binary(
    f: (RDD[(SpatialKey, Tile)], RDD[(SpatialKey, Tile)]) => RDD[(SpatialKey, Tile)],
    tlr0: TileLayerRDD[SpatialKey],
    tlr1: TileLayerRDD[SpatialKey]
  ): TileLayerRDD[SpatialKey] = {
    TileLayerRDD(f(tlr0, tlr1), tlr0.metadata.combine(tlr1.metadata))
  }

  /** Evaluate an AST of RDD Sources. Assumes that the AST's
    * [[NodeMetadata]] has already been replaced, if applicable.
    */
  def interpretRDD(
    ast: MapAlgebraAST,
    sourceMapping: Map[UUID, RFMLRaster],
    zoom: Int,
    sceneLocs: Map[UUID, String],
    projLocs: Map[UUID, List[(UUID, String)]]
  )(implicit sc: SparkContext): Interpreted[TileLayerRDD[SpatialKey]] = {

    /* Perform Map Algebra over a validated RDD-filled AST */
    @SuppressWarnings(Array("TraversableHead"))
    def eval(
      ast: MapAlgebraAST,
      rdds: Map[UUID, TileLayerRDD[SpatialKey]]
    ): TileLayerRDD[SpatialKey] = ast match {
      case Source(id, _) => rdds(id)
      case Addition(args, _, _) =>
        args.map(eval(_, rdds)).reduce((acc,r) => binary({_ + _}, acc, r))
      case Subtraction(args, _, _) =>
        args.map(eval(_, rdds)).reduce((acc,r) => binary({_ - _}, acc, r))
      case Multiplication(args, _, _) =>
        args.map(eval(_, rdds)).reduce((acc,r) => binary({_ * _}, acc, r))
      case Division(args, _, _) =>
        args.map(eval(_, rdds)).reduce((acc,r) => binary({_ / _}, acc, r))
      case Classification(arg :: _, _, _, classMap) =>
        eval(arg, rdds).withContext(_.color(classMap.toColorMap))
      case Max(args, _, _) => {
        val kids: List[TileLayerRDD[SpatialKey]] = args.map(eval(_, rdds))

        /* The head call will never fail */
        ContextRDD(kids.head.localMax(kids.tail), kids.map(_.metadata).reduce(_ combine _))
      }
      case Min(args, _, _) => {
        val kids: List[TileLayerRDD[SpatialKey]] = args.map(eval(_, rdds))

        /* The head call will never fail */
        ContextRDD(kids.head.localMin(kids.tail), kids.map(_.metadata).reduce(_ combine _))
      }
      case op => throw new Exception(s"Unimplemented Operation given! ${op}")
    }

    /* Guarantee correctness before performing Map Algebra */
    val pure = Interpreter.interpretPure[Unit](ast, sourceMapping)
    val rdds = sourceMapping.mapValues(r => fetch(r, zoom, sceneLocs, projLocs)).sequence

    (pure |@| rdds).map({ case (_, rs) => eval(ast, rs) })
  }

  /** This requires that for each [[RFMLRaster]] that a band number be specified. */
  private def fetch(
    raster: RFMLRaster,
    zoom: Int,
    sceneLocs: Map[UUID, String],
    projLocs: Map[UUID, List[(UUID, String)]]
  )(implicit sc: SparkContext): Interpreted[TileLayerRDD[SpatialKey]] = raster match {

    case ProjectRaster(id, None) => Invalid(NonEmptyList.of(NoBandGiven(id)))
    case ProjectRaster(id, Some(band)) => getStores(id, projLocs) match {
      case None => Invalid(NonEmptyList.of(AttributeStoreFetchError(id)))
      case Some(stores) => {
        val rdds: List[TileLayerRDD[SpatialKey]] =
          stores.map({ case (sceneId, store) =>
            S3LayerReader(store)
              .read[SpatialKey, MultibandTile, TileLayerMetadata[SpatialKey]](LayerId(sceneId.toString, zoom))
              .withContext(rdd => rdd.mapValues(_.band(band)))
          })

{'fraction_non_alphanumeric': 0.1461901550910317, 'fraction_numerical': 0.002157788267026298, 'mean_word_length': 3.541334966319657, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 45, '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': '5452791', 'n_tokens_mistral': 2648, 'n_tokens_neox': 2566, 'n_words': 537}