[spark] branch master updated: [SPARK-27278][SQL] Optimize GetMapValue when the map is a foldable and the key is not

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 5888b15  [SPARK-27278][SQL] Optimize GetMapValue when the map is a foldable and the key is not
5888b15 is described below

commit 5888b15d9cdf8272012018f39bf58c8faf68a5e1
Author: Marco Gaido <ma...@gmail.com>
AuthorDate: Mon Apr 1 09:09:06 2019 -0700

    [SPARK-27278][SQL] Optimize GetMapValue when the map is a foldable and the key is not
    
    ## What changes were proposed in this pull request?
    
    When `GetMapValue` contains a foldable Map and a non-foldable key, `SimplifyExtractValueOps` fails to optimize it transforming it into case when statements.
    The PR adds a case for covering this situation too.
    
    ## How was this patch tested?
    
    added UT
    
    Closes #24223 from mgaido91/SPARK-27278.
    
    Authored-by: Marco Gaido <ma...@gmail.com>
    Signed-off-by: Dongjoon Hyun <dh...@apple.com>
---
 .../apache/spark/sql/catalyst/optimizer/ComplexTypes.scala | 14 ++++++++++++++
 .../spark/sql/catalyst/optimizer/complexTypesSuite.scala   | 11 +++++++++++
 2 files changed, 25 insertions(+)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
index db7d6d3..0255128 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
@@ -17,9 +17,13 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
+import scala.collection.mutable
+
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.util.MapData
+import org.apache.spark.sql.types.MapType
 
 /**
  * Simplify redundant [[CreateNamedStructLike]], [[CreateArray]] and [[CreateMap]] expressions.
@@ -59,6 +63,16 @@ object SimplifyExtractValueOps extends Rule[LogicalPlan] {
           Literal(null, ga.dataType)
         }
       case GetMapValue(CreateMap(elems), key) => CaseKeyWhen(key, elems)
+      // The case below happens when the map is foldable, but the key is not, so ConstantFolding
+      // converts the map in a Literal, but the GetMapValue is still there since the key is not
+      // foldable. It cannot happen in any other case.
+      case GetMapValue(Literal(map: MapData, MapType(kt, vt, _)), key) if !key.foldable =>
+        val elems = new mutable.ListBuffer[Literal]
+        map.foreach(kt, vt, (key, value) => {
+          elems.append(Literal(key, kt))
+          elems.append(Literal(value, vt))
+        })
+        CaseKeyWhen(key, elems.result())
     }
   }
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/complexTypesSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/complexTypesSuite.scala
index 5452e72..84c74a6 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/complexTypesSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/complexTypesSuite.scala
@@ -452,4 +452,15 @@ class ComplexTypesSuite extends PlanTest with ExpressionEvalHelper {
     checkEvaluation(GetMapValue(mb0, Literal(Array[Byte](2, 1), BinaryType)), "2")
     checkEvaluation(GetMapValue(mb0, Literal(Array[Byte](3, 4))), null)
   }
+
+  test("SPARK-27278: simplify map access with non-foldable key and foldable map") {
+    val query = relation.select(GetMapValue(CreateMap(Seq(
+      1L, "a",
+      2L, "b")), 'id) as "a")
+    val expected = relation.select(
+        CaseWhen(Seq(
+          (EqualTo('id, 1L), "a"),
+          (EqualTo('id, 2L), "b"))) as "a")
+    checkRule(query, expected)
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org