<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.sql.sources.druid

import org.apache.spark.Logging
import org.apache.spark.sql.Strategy
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
import org.apache.spark.sql.execution.{PhysicalRDD, Project, SparkPlan, Union}
import org.apache.spark.sql.types.DoubleType
import org.sparklinedata.druid._
import org.sparklinedata.druid.query.QuerySpecTransforms

private[druid] class DruidStrategy(val planner: DruidPlanner) extends Strategy
with PredicateHelper with DruidPlannerHelper with Logging {

  override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
    case l => {

      val p: Seq[SparkPlan] = for (dqb <- planner.plan(null, l);
                                   a <- dqb.aggregateOper
      ) yield {

        val druidOpSchema = new DruidOperatorSchema(dqb)
        val pAgg = new PostAggregate(druidOpSchema)

        /*
         * 2. Interval is either in the SQL, or use the entire datasource interval
         */
        val intervals = dqb.queryIntervals.get

        /*
         * 3. Setup GroupByQuerySpec
         */
        var qs: QuerySpec = new GroupByQuerySpec(dqb.drInfo.druidDS.name,
          dqb.dimensions,
          dqb.limitSpec,
          dqb.havingSpec,
          dqb.granularitySpec,
          dqb.filterSpec,
          dqb.aggregations,
          dqb.postAggregations,
          intervals.map(_.toString)
        )

        /*
         * 4. apply QuerySpec transforms
         */
        qs = QuerySpecTransforms.transform(dqb.drInfo, qs)

        /*
         * 5. Setup DruidRelation
         */
        val dq = DruidQuery(qs,
          dqb.drInfo.options.queryHistoricalServers && pAgg.canBeExecutedInHistorical,
          intervals,
          Some(druidOpSchema.operatorDruidAttributes))

        planner.debugTranslation(
          s"""
             | DruidQuery:
             |   ${Utils.queryToString(dq)}
                """.stripMargin

        )

        val dR: DruidRelation = DruidRelation(dqb.drInfo, Some(dq))(planner.sqlContext)

        var druidPhysicalOp : SparkPlan = PhysicalRDD.createFromDataSource(
          druidOpSchema.operatorSchema,
          dR.buildInternalScan,
          dR)

        if (dq.queryHistoricalServer) {
          // add an Agg on top of druidQuery
          druidPhysicalOp = pAgg.aggOp(druidPhysicalOp).head
        }

        if ( druidPhysicalOp != null ) {
          val projections = buildProjectionList(dqb.aggregateOper.get,
            dqb.aggExprToLiteralExpr, druidOpSchema)
          Project(projections, druidPhysicalOp)
        } else null
      }

      val pL = p.toList

      if (pL.size < 2) pL else Seq(Union(pL))

    }
  }

  def buildProjectionList(aggOp: Aggregate,
                          grpExprToFillInLiteralExpr: Map[Expression, Expression],
                          druidOpSchema : DruidOperatorSchema):
  Seq[NamedExpression] = {

    /*
     * Replace aggregationExprs with fillIn expressions setup for this GroupingSet.
     * These are for Grouping__Id and for grouping expressions that are missing(null) for
     * this Grouping Set.
     */
    val aEs = aggOp.aggregateExpressions.map { aE => grpExprToFillInLiteralExpr.getOrElse(aE, aE) }

    val druidPushDownExprMap = druidOpSchema.pushedDownExprToDruidAttr
    val avgExpressions = druidOpSchema.avgExpressions

    aEs.map(_.transformUp {
      case e: Expression if avgExpressions.contains(e) => {
        val (s,c) = avgExpressions(e)
        val (sDAttr, cDAttr) = (druidOpSchema.druidAttrMap(s), druidOpSchema.druidAttrMap(c))
        Cast(
        Divide(
          Cast(AttributeReference(sDAttr.name, sDAttr.dataType)(sDAttr.exprId), DoubleType),
          Cast(AttributeReference(cDAttr.name, sDAttr.dataType)(cDAttr.exprId), DoubleType)
        ),
          e.dataType)
      }
      case ne: AttributeReference if druidPushDownExprMap.contains(ne) &&
        druidPushDownExprMap(ne).dataType != ne.dataType => {
        val dA = druidPushDownExprMap(ne)
        Alias(Cast(
          AttributeReference(dA.name, dA.dataType)(dA.exprId), ne.dataType), dA.name)(dA.exprId)
      }
      case ne: AttributeReference if druidPushDownExprMap.contains(ne) &&
        druidPushDownExprMap(ne).name != ne.name => {
        val dA = druidPushDownExprMap(ne)
        Alias(AttributeReference(dA.name, dA.dataType)(dA.exprId), dA.name)(dA.exprId)
      }
      case ne: AttributeReference if druidPushDownExprMap.contains(ne) => {
        ne<commit_msg>fix bug in datatype for generated count expression for Avg. expressions<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.sql.sources.druid

import org.apache.spark.Logging
import org.apache.spark.sql.Strategy
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
import org.apache.spark.sql.execution.{PhysicalRDD, Project, SparkPlan, Union}
import org.apache.spark.sql.types.DoubleType
import org.sparklinedata.druid._
import org.sparklinedata.druid.query.QuerySpecTransforms

private[druid] class DruidStrategy(val planner: DruidPlanner) extends Strategy
with PredicateHelper with DruidPlannerHelper with Logging {

  override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
    case l => {

      val p: Seq[SparkPlan] = for (dqb <- planner.plan(null, l);
                                   a <- dqb.aggregateOper
      ) yield {

        val druidOpSchema = new DruidOperatorSchema(dqb)
        val pAgg = new PostAggregate(druidOpSchema)

        /*
         * 2. Interval is either in the SQL, or use the entire datasource interval
         */
        val intervals = dqb.queryIntervals.get

        /*
         * 3. Setup GroupByQuerySpec
         */
        var qs: QuerySpec = new GroupByQuerySpec(dqb.drInfo.druidDS.name,
          dqb.dimensions,
          dqb.limitSpec,
          dqb.havingSpec,
          dqb.granularitySpec,
          dqb.filterSpec,
          dqb.aggregations,
          dqb.postAggregations,
          intervals.map(_.toString)
        )

        /*
         * 4. apply QuerySpec transforms
         */
        qs = QuerySpecTransforms.transform(dqb.drInfo, qs)

        /*
         * 5. Setup DruidRelation
         */
        val dq = DruidQuery(qs,
          dqb.drInfo.options.queryHistoricalServers && pAgg.canBeExecutedInHistorical,
          intervals,
          Some(druidOpSchema.operatorDruidAttributes))

        planner.debugTranslation(
          s"""
             | DruidQuery:
             |   ${Utils.queryToString(dq)}
                """.stripMargin

        )

        val dR: DruidRelation = DruidRelation(dqb.drInfo, Some(dq))(planner.sqlContext)

        var druidPhysicalOp : SparkPlan = PhysicalRDD.createFromDataSource(
          druidOpSchema.operatorSchema,
          dR.buildInternalScan,
          dR)

        if (dq.queryHistoricalServer) {
          // add an Agg on top of druidQuery
          druidPhysicalOp = pAgg.aggOp(druidPhysicalOp).head
        }

        if ( druidPhysicalOp != null ) {
          val projections = buildProjectionList(dqb.aggregateOper.get,
            dqb.aggExprToLiteralExpr, druidOpSchema)
          Project(projections, druidPhysicalOp)
        } else null
      }

      val pL = p.toList

      if (pL.size < 2) pL else Seq(Union(pL))

    }
  }

  def buildProjectionList(aggOp: Aggregate,
                          grpExprToFillInLiteralExpr: Map[Expression, Expression],
                          druidOpSchema : DruidOperatorSchema):
  Seq[NamedExpression] = {

    /*
     * Replace aggregationExprs with fillIn expressions setup for this GroupingSet.
     * These are for Grouping__Id and for grouping expressions that are missing(null) for
     * this Grouping Set.
     */
    val aEs = aggOp.aggregateExpressions.map { aE => grpExprToFillInLiteralExpr.getOrElse(aE, aE) }

    val druidPushDownExprMap = druidOpSchema.pushedDownExprToDruidAttr
    val avgExpressions = druidOpSchema.avgExpressions

    aEs.map(_.transformUp {
      case e: Expression if avgExpressions.contains(e) => {
        val (s,c) = avgExpressions(e)
        val (sDAttr, cDAttr) = (druidOpSchema.druidAttrMap(s), druidOpSchema.druidAttrMap(c))
        Cast(
        Divide(
          Cast(AttributeReference(sDAttr.name, sDAttr.dataType)(sDAttr.exprId), DoubleType),
          Cast(AttributeReference(cDAttr.name, cDAttr.dataType)(cDAttr.exprId), DoubleType)
        ),
          e.dataType)
      }
      case ne: AttributeReference if druidPushDownExprMap.contains(ne) &&
        druidPushDownExprMap(ne).dataType != ne.dataType => {
        val dA = druidPushDownExprMap(ne)
        Alias(Cast(
          AttributeReference(dA.name, dA.dataType)(dA.exprId), ne.dataType), dA.name)(dA.exprId)
      }
      case ne: AttributeReference if druidPushDownExprMap.contains(ne) &&
        druidPushDownExprMap(ne).name != ne.name => {
        val dA = druidPushDownExprMap(ne)
        Alias(AttributeReference(dA.name, dA.dataType)(dA.exprId), dA.name)(dA.exprId)
      }
      case ne: AttributeReference if druidPushDownExprMap.contains(ne) => {
        ne
{'fraction_non_alphanumeric': 0.0882844411437059, 'fraction_numerical': 0.0016709988860007428, 'mean_word_length': 2.5981963927855714, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2654461', 'n_tokens_mistral': 3356, 'n_tokens_neox': 3209, 'n_words': 759}