<commit_before>package org.apache.spark.sql.sources.druid

import org.apache.spark.sql.catalyst.analysis.HiveTypeCoercion
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.planning.PhysicalOperation
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
import org.apache.spark.sql.sources.LogicalRelation
import org.apache.spark.sql.types.{StringType, DoubleType, IntegerType, LongType}
import org.sparklinedata.druid.metadata.{DruidColumn, DruidDimension, DruidDataType, DruidMetric}
import org.sparklinedata.druid._


abstract class DruidTransforms {
  self: DruidPlanner =>

  type DruidTransform = PartialFunction[(DruidQueryBuilder, LogicalPlan), Option[DruidQueryBuilder]]

  val druidRelationTransform: DruidTransform = {
    case (_, PhysicalOperation(projectList, filters,
    l@LogicalRelation(d@DruidRelation(info, None)))) => {
      var dqb : Option[DruidQueryBuilder] = Some(DruidQueryBuilder(info))
      dqb = projectList.foldLeft(dqb) { (dqB, e) =>
        dqB.flatMap(projectExpression(_, e))
      }

      if ( dqb.isDefined) {
        /*
         * Filter Rewrites:
         * - A conjunct is a predicate on the Time Dimension => rewritten to Interval constraint
         * - A expression containing comparisons on Dim Columns.
         */
        val iCE: IntervalConditionExtractor = new IntervalConditionExtractor(dqb.get)
        filters.foldLeft(dqb) { (dqB, e) =>
          dqB.flatMap{b =>
            intervalFilterExpression(b, iCE, e).orElse(
              dimFilterExpression(b,e).map(p => b.filter(p))
            )
          }
        }
      } else None
    }
  }

  val aggregateTransform: DruidTransform = {
    case (dqb, agg@Aggregate(gEs, aEs, child)) => {
      plan(dqb, child).flatMap { dqb =>

        val timeElemExtractor = new TimeElementExtractor(dqb)

        val dqb1 = gEs.foldLeft(Some(dqb).asInstanceOf[Option[DruidQueryBuilder]]) { (dqb, e) =>
          dqb.flatMap(groupingExpression(_, timeElemExtractor, e))
        }

        val allAggregates =
          aEs.flatMap(_ collect { case a: AggregateExpression => a})
        // Collect all aggregate expressions that can be computed partially.
        val partialAggregates =
          aEs.flatMap(_ collect { case p: PartialAggregate => p})

        // Only do partial aggregation if supported by all aggregate expressions.
        if (allAggregates.size == partialAggregates.size) {
          val dqb2 = partialAggregates.foldLeft(dqb1) {
            (dqb, ne) =>
              dqb.flatMap(aggregateExpression(_, ne))
          }

          dqb2.map(_.aggregateOp(agg))
        } else {
          None
        }

      }
    }
  }


  def aggregateExpression(dqb: DruidQueryBuilder, pa: PartialAggregate):
  Option[DruidQueryBuilder] = pa match {
    case Count(Literal(1, IntegerType)) => {
      val a = dqb.nextAlias
      Some(dqb.aggregate(FunctionAggregationSpec("count", a, "count")).
        outputAttribute(a, pa, pa.dataType, LongType))
    }
    case c => {
      val a = dqb.nextAlias
      (dqb, c) match {
        case CountDistinctAggregate(dN) =>
          Some(dqb.aggregate(new CardinalityAggregationSpec(a, List(dN))).
            outputAttribute(a, pa, pa.dataType, LongType))
        case SumMinMaxAvgAggregate(t) => t._1 match {
          case "avg" => {
            val dC: DruidColumn = t._2
            val aggFunc = dC.dataType match {
              case DruidDataType.Long => "longSum"
              case _ => "doubleSum"
            }
            val sumAlias = dqb.nextAlias
            val countAlias = dqb.nextAlias

            Some(<commit_msg>countDistinct returns double from Druid<commit_after>package org.apache.spark.sql.sources.druid

import org.apache.spark.sql.catalyst.analysis.HiveTypeCoercion
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.planning.PhysicalOperation
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
import org.apache.spark.sql.sources.LogicalRelation
import org.apache.spark.sql.types.{StringType, DoubleType, IntegerType, LongType}
import org.sparklinedata.druid.metadata.{DruidColumn, DruidDimension, DruidDataType, DruidMetric}
import org.sparklinedata.druid._


abstract class DruidTransforms {
  self: DruidPlanner =>

  type DruidTransform = PartialFunction[(DruidQueryBuilder, LogicalPlan), Option[DruidQueryBuilder]]

  val druidRelationTransform: DruidTransform = {
    case (_, PhysicalOperation(projectList, filters,
    l@LogicalRelation(d@DruidRelation(info, None)))) => {
      var dqb : Option[DruidQueryBuilder] = Some(DruidQueryBuilder(info))
      dqb = projectList.foldLeft(dqb) { (dqB, e) =>
        dqB.flatMap(projectExpression(_, e))
      }

      if ( dqb.isDefined) {
        /*
         * Filter Rewrites:
         * - A conjunct is a predicate on the Time Dimension => rewritten to Interval constraint
         * - A expression containing comparisons on Dim Columns.
         */
        val iCE: IntervalConditionExtractor = new IntervalConditionExtractor(dqb.get)
        filters.foldLeft(dqb) { (dqB, e) =>
          dqB.flatMap{b =>
            intervalFilterExpression(b, iCE, e).orElse(
              dimFilterExpression(b,e).map(p => b.filter(p))
            )
          }
        }
      } else None
    }
  }

  val aggregateTransform: DruidTransform = {
    case (dqb, agg@Aggregate(gEs, aEs, child)) => {
      plan(dqb, child).flatMap { dqb =>

        val timeElemExtractor = new TimeElementExtractor(dqb)

        val dqb1 = gEs.foldLeft(Some(dqb).asInstanceOf[Option[DruidQueryBuilder]]) { (dqb, e) =>
          dqb.flatMap(groupingExpression(_, timeElemExtractor, e))
        }

        val allAggregates =
          aEs.flatMap(_ collect { case a: AggregateExpression => a})
        // Collect all aggregate expressions that can be computed partially.
        val partialAggregates =
          aEs.flatMap(_ collect { case p: PartialAggregate => p})

        // Only do partial aggregation if supported by all aggregate expressions.
        if (allAggregates.size == partialAggregates.size) {
          val dqb2 = partialAggregates.foldLeft(dqb1) {
            (dqb, ne) =>
              dqb.flatMap(aggregateExpression(_, ne))
          }

          dqb2.map(_.aggregateOp(agg))
        } else {
          None
        }

      }
    }
  }


  def aggregateExpression(dqb: DruidQueryBuilder, pa: PartialAggregate):
  Option[DruidQueryBuilder] = pa match {
    case Count(Literal(1, IntegerType)) => {
      val a = dqb.nextAlias
      Some(dqb.aggregate(FunctionAggregationSpec("count", a, "count")).
        outputAttribute(a, pa, pa.dataType, LongType))
    }
    case c => {
      val a = dqb.nextAlias
      (dqb, c) match {
        case CountDistinctAggregate(dN) =>
          Some(dqb.aggregate(new CardinalityAggregationSpec(a, List(dN))).
            outputAttribute(a, pa, pa.dataType, DoubleType))
        case SumMinMaxAvgAggregate(t) => t._1 match {
          case "avg" => {
            val dC: DruidColumn = t._2
            val aggFunc = dC.dataType match {
              case DruidDataType.Long => "longSum"
              case _ => "doubleSum"
            }
            val sumAlias = dqb.nextAlias
            val countAlias = dqb.nextAlias

            Some(
{'fraction_non_alphanumeric': 0.10685979775895053, 'fraction_numerical': 0.0019130910084722602, 'mean_word_length': 2.628656420426376, 'pattern_counts': {'":': 0, '<': 3, '<?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': '18119', 'n_tokens_mistral': 2362, 'n_tokens_neox': 2264, 'n_words': 459}