<commit_before>package com.azavea.rf.batch.export.spark

import com.azavea.rf.batch.export.json.S3ExportStatus
import com.azavea.rf.batch._
import com.azavea.rf.batch.ast._
import com.azavea.rf.batch.dropbox._
import com.azavea.rf.batch.export._
import com.azavea.rf.batch.util._
import com.azavea.rf.batch.util.conf._
import com.azavea.rf.datamodel._
import com.azavea.rf.tool.ast.MapAlgebraAST
import com.azavea.rf.common.S3.putObject

import com.azavea.maml.serve._
import com.azavea.maml.eval._
import com.dropbox.core.v2.DbxClientV2
import com.dropbox.core.v2.files.{CreateFolderErrorException, WriteMode}
import com.typesafe.scalalogging.LazyLogging
import io.circe.parser._
import io.circe.syntax._
import cats.data.Validated._
import cats.implicits._
import com.dropbox.core.v2.DbxClientV2
import com.dropbox.core.v2.files.{CreateFolderErrorException, WriteMode}
import com.typesafe.scalalogging.LazyLogging
import geotrellis.proj4.{CRS, LatLng}
import geotrellis.raster._
import geotrellis.raster.histogram.Histogram
import geotrellis.raster.io._
import geotrellis.raster.io.geotiff.{GeoTiff, MultibandGeoTiff, SinglebandGeoTiff}
import geotrellis.spark._
import geotrellis.spark.io._
import geotrellis.spark.io.file._
import geotrellis.spark.io.hadoop._
import geotrellis.spark.io.s3._
import geotrellis.spark.regrid._
import geotrellis.spark.tiling._
import geotrellis.vector.MultiPolygon
import org.apache.hadoop.fs.Path
import spray.json.DefaultJsonProtocol._
import org.apache.spark.rdd.RDD
import org.apache.spark._

import java.util.UUID


object Export extends SparkJob with Config with LazyLogging {

  val jobName = "Export"

  val defaultRasterSize = 4000

  def astExport(
    ed: ExportDefinition,
    ast: MapAlgebraAST,
    sceneLocs: Map[UUID, String],
    projLocs: Map[UUID, List[(UUID, String)]],
    conf: HadoopConfiguration
  )(implicit sc: SparkContext): Unit = {
    interpretRDD(ast, ed.input.resolution, sceneLocs, projLocs) match {
      case Invalid(errs) => throw InterpreterException(errs)
      case Valid(rdd) => {
        val crs: CRS = rdd.metadata.crs

        if(!ed.output.stitch) {
          val targetRDD: TileLayerRDD[SpatialKey] =
            ed.output.rasterSize match {
              case Some(size) => rdd.regrid(size)
              case None => rdd.regrid(defaultRasterSize)
            }

          val mt: MapKeyTransform = targetRDD.metadata.layout.mapTransform

          /* Create GeoTiffs and output them */
          val singles: RDD[(SpatialKey, SinglebandGeoTiff)] =
            rdd.map({ case (key, tile) => (key, SinglebandGeoTiff(tile, mt(key), crs)) })

          writeGeoTiffs[Tile, SinglebandGeoTiff](singles, ed, conf)
        } else {
          /* Stitch the Layer into a single GeoTiff and output it */
          val single: SinglebandGeoTiff = GeoTiff(rdd.stitch, crs)

          writeGeoTiff[Tile, SinglebandGeoTiff](single, ed, conf, singlePath)
        }
      }
    }
  }

  /** Get a LayerReader and an attribute store for the catalog located at the provided URI
    *
    *  @param ed export job's layer definition
    */
  def getRfLayerManagement(
    ed: ExportLayerDefinition
  )(implicit @transient sc: SparkContext): (FilteringLayerReader[LayerId], AttributeStore) = {
    ed.ingestLocation.getScheme match {
      case "s3" | "s3a" | "s3n" =>
        val (bucket, prefix) = S3.parse(ed.ingestLocation)
        val reader = S3LayerReader(bucket, prefix)
        (reader, reader.attributeStore)
      case "file" =>
        val reader = FileLayerReader(ed.ingestLocation.getPath)
        (reader, reader.attributeStore)
    }
  }

  def multibandExport(
    ed: ExportDefinition,
    layers: Array[ExportLayerDefinition],
    mask: Option[MultiPolygon],
    conf: HadoopConfiguration
  )(implicit @transient sc: SparkContext): Unit = {
    val rdds = layers.map { ld =>
      val (reader, store) = getRfLayerManagement(ld)
      val layerId = LayerId(ld.layerId.toString, ed.input.resolution)
      val md = store.readMetadata[TileLayerMetadata[SpatialKey]](layerId)
      val hist = ld.colorCorrections.map { _ => store.read[Array[Histogram[Double]]](LayerId(layerId.name, 0), "histogram") }

      val q = reader.query[SpatialKey, MultibandTile, TileLayerMetadata[SpatialKey]](layerId)

      val query: ContextRDD[SpatialKey, MultibandTile, TileLayerMetadata[SpatialKey]] =
        mask
          .fold(q)(mp => q.where(Intersects(mp)))
          .result
          .withContext({ rdd => rdd.mapValues({ tile =>
            val ctile = (ld.colorCorrections |@| hist) map { _.colorCorrect(tile, _) } getOrElse tile

            ed.output.render.flatMap(_.bands.map(_.toSeq)) match {
              case Some(seq) if seq.nonEmpty => ctile.subsetBands(seq)
              case _ => ctile
            }
          })})

      (ed.output.rasterSize, ed.output.crs) match {
        case (Some(rs), Some(crs)) =>
          query.reproject(ZoomedLayoutScheme(crs, rs))._2
        case (None, Some(crs)) =>
          query.reproject(ZoomedLayoutScheme(crs, defaultRasterSize))._2
        case (Some(rs), None) =>
          query.regrid(rs)
        case (None, None) =>
          query.regrid(defaultRasterSize)
      }
    }

    /** Tile merge with respect to layer initial ordering */
    if(!ed.output.stitch) {
      val result: RDD[(SpatialKey, MultibandGeoTiff)] =
        rdds
          .zipWithIndex
          .map { case (rdd, i) =>
            val md = rdd.metadata<commit_msg>multibandExport can now perfrom overzooming<commit_after>package com.azavea.rf.batch.export.spark

import com.azavea.rf.batch.export.json.S3ExportStatus
import com.azavea.rf.batch._
import com.azavea.rf.batch.ast._
import com.azavea.rf.batch.dropbox._
import com.azavea.rf.batch.export._
import com.azavea.rf.batch.util._
import com.azavea.rf.batch.util.conf._
import com.azavea.rf.datamodel._
import com.azavea.rf.tool.ast.MapAlgebraAST
import com.azavea.rf.common.S3.putObject

import com.azavea.maml.serve._
import com.azavea.maml.eval._
import com.dropbox.core.v2.DbxClientV2
import com.dropbox.core.v2.files.{CreateFolderErrorException, WriteMode}
import com.typesafe.scalalogging.LazyLogging
import io.circe.parser._
import io.circe.syntax._
import cats.data.Validated._
import cats.implicits._
import com.dropbox.core.v2.DbxClientV2
import com.dropbox.core.v2.files.{CreateFolderErrorException, WriteMode}
import com.typesafe.scalalogging.LazyLogging
import geotrellis.proj4.{CRS, LatLng}
import geotrellis.raster._
import geotrellis.raster.histogram.Histogram
import geotrellis.raster.io._
import geotrellis.raster.io.geotiff.{GeoTiff, MultibandGeoTiff, SinglebandGeoTiff}
import geotrellis.spark._
import geotrellis.spark.io._
import geotrellis.spark.io.file._
import geotrellis.spark.io.hadoop._
import geotrellis.spark.io.s3._
import geotrellis.spark.regrid._
import geotrellis.spark.resample._
import geotrellis.spark.tiling._
import geotrellis.vector.MultiPolygon
import org.apache.hadoop.fs.Path
import spray.json.DefaultJsonProtocol._
import org.apache.spark.rdd.RDD
import org.apache.spark._

import java.util.UUID


object Export extends SparkJob with Config with LazyLogging {

  val jobName = "Export"

  val defaultRasterSize = 4000

  def astExport(
    ed: ExportDefinition,
    ast: MapAlgebraAST,
    sceneLocs: Map[UUID, String],
    projLocs: Map[UUID, List[(UUID, String)]],
    conf: HadoopConfiguration
  )(implicit sc: SparkContext): Unit = {
    interpretRDD(ast, ed.input.resolution, sceneLocs, projLocs) match {
      case Invalid(errs) => throw InterpreterException(errs)
      case Valid(rdd) => {
        val crs: CRS = rdd.metadata.crs

        if(!ed.output.stitch) {
          val targetRDD: TileLayerRDD[SpatialKey] =
            ed.output.rasterSize match {
              case Some(size) => rdd.regrid(size)
              case None => rdd.regrid(defaultRasterSize)
            }

          val mt: MapKeyTransform = targetRDD.metadata.layout.mapTransform

          /* Create GeoTiffs and output them */
          val singles: RDD[(SpatialKey, SinglebandGeoTiff)] =
            rdd.map({ case (key, tile) => (key, SinglebandGeoTiff(tile, mt(key), crs)) })

          writeGeoTiffs[Tile, SinglebandGeoTiff](singles, ed, conf)
        } else {
          /* Stitch the Layer into a single GeoTiff and output it */
          val single: SinglebandGeoTiff = GeoTiff(rdd.stitch, crs)

          writeGeoTiff[Tile, SinglebandGeoTiff](single, ed, conf, singlePath)
        }
      }
    }
  }

  /** Get a LayerReader and an attribute store for the catalog located at the provided URI
    *
    *  @param ed export job's layer definition
    */
  def getRfLayerManagement(
    ed: ExportLayerDefinition
  )(implicit @transient sc: SparkContext): (FilteringLayerReader[LayerId], AttributeStore) = {
    ed.ingestLocation.getScheme match {
      case "s3" | "s3a" | "s3n" =>
        val (bucket, prefix) = S3.parse(ed.ingestLocation)
        val reader = S3LayerReader(bucket, prefix)
        (reader, reader.attributeStore)
      case "file" =>
        val reader = FileLayerReader(ed.ingestLocation.getPath)
        (reader, reader.attributeStore)
    }
  }

  def multibandExport(
    ed: ExportDefinition,
    layers: Array[ExportLayerDefinition],
    mask: Option[MultiPolygon],
    conf: HadoopConfiguration
  )(implicit @transient sc: SparkContext): Unit = {
    val rdds = layers.map { ld =>
      val (reader, store) = getRfLayerManagement(ld)
      val requestedLayerId = LayerId(ld.layerId.toString, ed.input.resolution)

      val maxAvailableZoom =
        store
          .layerIds
          .filter { case LayerId(name, _) => name == requestedLayerId.name }
          .map { _.zoom }
          .max

      val maxLayerId = requestedLayerId.copy(zoom = maxAvailableZoom)
      val maxMetadata = store.readMetadata[TileLayerMetadata[SpatialKey]](maxLayerId)
      val maxMapTransform = maxMetadata.mapTransform

      val layoutScheme = ZoomedLayoutScheme(maxMetadata.crs, math.min(maxMetadata.tileCols, maxMetadata.tileRows))

      val requestedMapTransform = layoutScheme.levelForZoom(requestedLayerId.zoom).layout.mapTransform

      lazy val requestedQuery = reader.query[SpatialKey, MultibandTile, TileLayerMetadata[SpatialKey]](requestedLayerId)
      lazy val maxQuery = reader.query[SpatialKey, MultibandTile, TileLayerMetadata[SpatialKey]](maxLayerId)

      val hist =
        ld
          .colorCorrections
          .map { _ => store.read[Array[Histogram[Double]]](requestedLayerId.copy(zoom = 0), "histogram") }

      val queryLayer = (q: BoundLayerQuery[SpatialKey, TileLayerMetadata[SpatialKey], RDD[(SpatialKey, MultibandTile)] with Metadata[TileLayerMetadata[SpatialKey]]]) =>
        mask
          .fold(q)(mp => q.where(Intersects(mp)))
          .result
          .withContext({ rdd => rdd.mapValues({ tile =>
            val ctile = (ld.colorCorrections |@| hist) map { _.colorCorrect(tile, _) } getOrElse tile

            ed.output.render.flatMap(_.bands.map(_.toSeq)) match {
              case Some(seq) if seq.nonEmpty => ctile.subsetBands(seq)
              case _ => ctile
            }
          })})

      val query: MultibandTileLayerRDD[SpatialKey] =
        if (requestedLayerId.zoom <= maxAvailableZoom)
          queryLayer(requestedQuery)
        else
          queryLayer(maxQuery).resample(maxAvailableZoom, requestedLayerId.zoom)

      val md = query.metadata

      (ed.output.rasterSize, ed.output.crs) match {
        case (Some(rs), Some(crs)) =>
          query.reproject(ZoomedLayoutScheme(crs, rs))._2
        case (None, Some(crs)) =>
          query.reproject(ZoomedLayoutScheme(crs, defaultRasterSize))._2
        case (Some(rs), None) =>
          query.regrid(rs)
        case (None, None) =>
          query.regrid(defaultRasterSize)
      }
    }

    /** Tile merge with respect to layer initial ordering */
    if(!ed.output.stitch) {
      val result: RDD[(SpatialKey, MultibandGeoTiff)] =
        rdds
          .zipWithIndex
          .map { case (rdd, i) =>
            val md = rdd.metadata
{'fraction_non_alphanumeric': 0.1103648424543947, 'fraction_numerical': 0.003648424543946932, 'mean_word_length': 3.1276522929500343, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 43, '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': '5398830', 'n_tokens_mistral': 4082, 'n_tokens_neox': 3947, 'n_words': 780}