Optimize aggregation push down column logic (pingcap#143)

wanfangdata · Dec 27, 2017 · 8bbec4f · 8bbec4f
1 parent fe033c1
commit 8bbec4f
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/src/main/scala/org/apache/spark/sql/TiStrategy.scala b/src/main/scala/org/apache/spark/sql/TiStrategy.scala
@@ -15,6 +15,7 @@
 
 package org.apache.spark.sql
 
+import com.pingcap.tikv.expression.scalar.TiScalarFunction
 import java.time.ZonedDateTime
 
 import com.pingcap.tikv.expression.{ExpressionBlacklist, TiByItem, TiColumnRef, TiExpr}
@@ -35,6 +36,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
 import scala.collection.JavaConversions._
+import scala.collection.mutable.ArrayBuffer
 import scala.collection.{JavaConversions, mutable}
 
 // TODO: Too many hacks here since we hijack the planning
@@ -143,6 +145,15 @@ class TiStrategy(context: SQLContext) extends Strategy with Logging {
     dagRequest
   }
 
+  def extractColumnFromFilter(tiFilter: TiExpr, result: ArrayBuffer[TiColumnRef]): Unit =
+    tiFilter match {
+      case fun: TiScalarFunction =>
+        fun.getArgs.foreach(extractColumnFromFilter(_, result))
+      case col: TiColumnRef =>
+        result.add(col)
+      case _ =>
+    }
+
   def filterToDAGRequest(
     filters: Seq[Expression],
     source: TiDBRelation,
@@ -266,6 +277,7 @@ class TiStrategy(context: SQLContext) extends Strategy with Logging {
   }
 
   def groupAggregateProjection(
+    filters: Seq[Expression],
     groupingExpressions: Seq[NamedExpression],
     aggregateExpressions: Seq[AggregateExpression],
     resultExpressions: Seq[NamedExpression],
@@ -375,6 +387,10 @@ class TiStrategy(context: SQLContext) extends Strategy with Logging {
 
     val projectSeq: Seq[Attribute] = projects.asInstanceOf[Seq[Attribute]]
     projectSeq.foreach(attr => dagReq.addRequiredColumn(TiColumnRef.create(attr.name)))
+    val pushDownCols = ArrayBuffer[TiColumnRef]()
+    val tiFilters: Seq[TiExpr] = filters.collect { case BasicExpression(expr) => expr }
+    tiFilters.foreach(extractColumnFromFilter(_, pushDownCols))
+    pushDownCols.foreach(dagReq.addRequiredColumn)
 
     aggregate.AggUtils.planAggregateWithoutDistinct(
       groupingExpressions,
@@ -446,6 +462,7 @@ class TiStrategy(context: SQLContext) extends Strategy with Logging {
           ) if isValidAggregates(groupingExpressions, aggregateExpressions, filters, source) =>
         val dagReq: TiDAGRequest = filterToDAGRequest(filters, source)
         groupAggregateProjection(
+          filters,
           groupingExpressions,
           aggregateExpressions,
           resultExpressions,

diff --git a/tikv-client-lib-java b/tikv-client-lib-java
+1 −1		kvproto
+17 −7		src/main/java/com/pingcap/tikv/Main.java
+14 −1		src/main/java/com/pingcap/tikv/expression/TiColumnRef.java
+46 −7		src/main/java/com/pingcap/tikv/meta/TiDAGRequest.java