[CROSSDATA-1896] Pushdown limit when serializing result in executors (#693)

Author: mafernandez-stratio

CHANGELOG.md

@@ -13,6 +13,7 @@ Only listing significant user-visible, not internal code cleanups and minor bug
 * [SS-3782][SS-3857] Serve HDFS conf files properly
 * [POC] Function interface for lineage
 * Configuration: Optional serialization in executors
+* [CROSSDATA-1896] Apply limit optimization before result serialization in executors
 
 ## 2.15.0-24f463e (Built: January 14, 2019 | Released: January 17, 2019)
 

Jenkinsfile

@@ -50,20 +50,6 @@ hose {
                     'POSTGRES_DB=hakama']
             ]
         ],
-        ['GOVERNANCE': [
-            'image': 'stratio/governance-catalog-api:0.3.0',
-            'sleep': 30,
-            'healthcheck': 10000,
-            'env': ['USE_DYNAMIC_AUTHENTICATION=false',
-		            'API_SERVICE_TIMEOUT=5000',
-		            'API_LOG_LEVEL=INFO',
-		            'SERVICE_LOG_LEVEL=INFO',
-		            'POSTGRESQL_HOSTS=%%POSTGRES',
-                    'POSTGRESQL_USER=postgres',
-                    'POSTGRESQL_PASSWORD=postgres',
-                    'POSTGRESQL_DATABASE=hakama']
-            ]
-        ],
         ['ELASTIC': [
             'image': 'stratio/elasticsearch:6.1.1',
             'sleep': 30,

core/src/main/scala/com/stratio/crossdata/metrics/MetricsRegister.scala

@@ -4,6 +4,8 @@ import java.util.concurrent.TimeUnit
 
 import com.codahale.metrics._
 
+import scala.util.Try
+
 object MetricsGlossary {
 
   object Counter {
@@ -192,6 +194,6 @@ class MetricsRegister(metricsRegistry: MetricRegistry) {
   }
 
   private def getFailingQueryCounterKey(message: String, user: String): String =
-    s"[$user][${message.replaceAll(" ", "_")}]"
+    Option(message).map(msg => s"[$user][${msg.replaceAll(" ", "_")}]").getOrElse(s"[$user][null]")
 
 }

core/src/main/scala/org/apache/spark/sql/crossdata/execution/XDQueryExecution.scala

@@ -8,23 +8,22 @@ package org.apache.spark.sql.crossdata.execution
 import com.stratio.common.utils.components.logger.impl.Slf4jLoggerComponent
 import com.stratio.crossdata.common.profiler.PerformanceLogger
 import com.stratio.crossdata.security._
-import org.apache.spark.sql.{SaveMode, SparkSession}
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedAlias, UnresolvedAttribute, UnresolvedRelation}
 import org.apache.spark.sql.catalyst.expressions.{And, Cast, Literal, Sha2}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.crossdata.checkAndApplyRule
-import org.apache.spark.sql.crossdata.XDSession
 import org.apache.spark.sql.crossdata.catalyst.catalog.XDSessionCatalog
 import org.apache.spark.sql.crossdata.catalyst.plans.logical.{XDFilter, XDProject}
 import org.apache.spark.sql.crossdata.execution.auth.AuthDirectivesExtractor
 import org.apache.spark.sql.crossdata.rules.XDCompanionRules
 import org.apache.spark.sql.crossdata.session.XDSharedState
+import org.apache.spark.sql.crossdata.{XDSession, checkAndApplyRule}
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.command.CreateTableLikeCommand
 import org.apache.spark.sql.execution.datasources.CreateTable
 import org.apache.spark.sql.types._
+import org.apache.spark.sql.{SaveMode, SparkSession}
 import org.apache.spark.util.Utils
 
 import scala.util.{Failure, Success, Try}
@@ -462,6 +461,7 @@ class XDQueryExecution(sparkSession: SparkSession, parsedPlan: LogicalPlan, used
   override lazy val sparkPlan: SparkPlan = {
     logPerformance(s"[sparkPlan]") {
       SparkSession.setActiveSession(sparkSession)
+
       // TODO: We use next(), i.e. take the first plan returned by the planner, here for now,
       //       but we will implement to choose the best plan.
       planner.plan(ReturnAnswer(optimizedPlan)).next()

core/src/main/scala/org/apache/spark/sql/crossdata/session/XDSessionStateBuilder.scala

@@ -5,17 +5,19 @@
  */
 package org.apache.spark.sql.crossdata.session
 
-import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.{SparkSession, Strategy}
 import org.apache.spark.sql.catalyst.analysis.Analyzer
 import org.apache.spark.sql.catalyst.catalog.SessionCatalog
 import org.apache.spark.sql.catalyst.parser.ParserInterface
-import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.crossdata.XDSession
 import org.apache.spark.sql.crossdata.catalyst.catalog.{XDCatalogWrapper, XDSessionCatalog}
 import org.apache.spark.sql.crossdata.catalyst.catalog.temporary.implementations.DefaultTemporaryCatalog
 import org.apache.spark.sql.crossdata.execution.{XDQueryExecution, XDSparkSqlParser}
-import org.apache.spark.sql.execution.QueryExecution
-import org.apache.spark.sql.internal.{BaseSessionStateBuilder, SessionState}
+import org.apache.spark.sql.execution.{QueryExecution, SparkPlanner}
+import org.apache.spark.sql.internal.SQLConf.buildConf
+import org.apache.spark.sql.internal.{BaseSessionStateBuilder, SQLConf, SessionState}
 
 
 class XDSessionStateBuilder(sparkSession: SparkSession, parentState: Option[SessionState] = None) extends BaseSessionStateBuilder(sparkSession, parentState){
@@ -28,6 +30,27 @@ class XDSessionStateBuilder(sparkSession: SparkSession, parentState: Option[Sess
 
   override lazy val analyzer: Analyzer = super.analyzer
 
+  object XDLimitRule extends Rule[LogicalPlan] {
+    val applyXDLimitRule = sparkSession.sparkContext.conf.getBoolean("spark.sql.crossdata.limitRule", false)
+    logDebug(s"applyXDLimitRule? $applyXDLimitRule")
+    override def apply(plan: LogicalPlan): LogicalPlan = plan transform {
+      case s@SerializeFromObject(serializer, MapPartitions(func, outputObjAttr, DeserializeToObject(des, attrs, gl@GlobalLimit(_, _)))) =>
+        if(applyXDLimitRule){
+          SerializeFromObject(serializer, MapPartitions(func, outputObjAttr, DeserializeToObject(des, attrs, ReturnAnswer(gl))))
+        } else {
+          s
+        }
+    }
+  }
+
+  override def planner: SparkPlanner = {
+    experimentalMethods.extraOptimizations = experimentalMethods.extraOptimizations :+ XDLimitRule
+    new SparkPlanner(session.sparkContext, conf, experimentalMethods) {
+      override def extraPlanningStrategies: Seq[Strategy] =
+        super.extraPlanningStrategies ++ customPlanningStrategies
+    }
+  }
+
   override lazy val catalog: SessionCatalog = new XDSessionCatalog(
     sparkSession.asInstanceOf[XDSession].config,
     DefaultTemporaryCatalog, //TODO: Next iterations requiring value propagation will need a distributed temp map.