Fix SPARK-12139: REGEX Column Specification for Hive Queries

Author: janewangfb

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala

@@ -83,6 +83,33 @@ case class UnresolvedTableValuedFunction(
   override lazy val resolved = false
 }
 
+/**
+ * Represents all of the input attributes to a given relational operator, for example in
+ * "SELECT `(id)?+.+` FROM ...".
+ *
+ * @param table an optional table that should be the target of the expansion.  If omitted all
+ *              tables' columns are produced.
+ */
+case class UnresolvedRegex(expr: String, table: Option[String]) extends Star with Unevaluable {
+  override def expand(input: LogicalPlan, resolver: Resolver): Seq[NamedExpression] = {
+    val expandedAttributes: Seq[Attribute] = table match {
+      // If there is no table specified, use all input attributes that match expr
+      case None => input.output.filter(_.name.matches(expr))
+      // If there is a table, pick out attributes that are part of this table that match expr
+      case Some(t) => input.output.filter(_.qualifier.filter(resolver(_, t)).nonEmpty)
+        .filter(_.name.matches(expr))
+    }
+
+    expandedAttributes.zip(input.output).map {
+      case (n: NamedExpression, _) => n
+      case (e, originalAttribute) =>
+        Alias(e, originalAttribute.name)()
+    }
+  }
+
+  override def toString: String = table.map(_ + ".").getOrElse("") + expr
+}
+
 /**
  * Holds the name of an attribute that has yet to be resolved.
  */

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -23,10 +23,9 @@ import javax.xml.bind.DatatypeConverter
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
-
 import org.antlr.v4.runtime.{ParserRuleContext, Token}
 import org.antlr.v4.runtime.tree.{ParseTree, RuleNode, TerminalNode}
-
+import org.apache.spark.SparkEnv
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier}
@@ -1229,25 +1228,56 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
     CaseWhen(branches, Option(ctx.elseExpression).map(expression))
   }
 
+  def enableHiveSupportQuotedIdentifiers() : Boolean = {
+    SparkEnv.get != null &&
+      SparkEnv.get.conf != null &&
+      SparkEnv.get.conf.getBoolean("hive.support.quoted.identifiers", false)
+  }
+
   /**
-   * Create a dereference expression. The return type depends on the type of the parent, this can
-   * either be a [[UnresolvedAttribute]] (if the parent is an [[UnresolvedAttribute]]), or an
-   * [[UnresolvedExtractValue]] if the parent is some expression.
+   * Create a dereference expression. The return type depends on the type of the parent.
+   * If the parent is an [[UnresolvedAttribute]], it can be a [[UnresolvedAttribute]] or
+   * a [[UnresolvedRegex]] for regex quoted in ``; if the parent is some other expression,
+   * it can be [[UnresolvedExtractValue]].
    */
   override def visitDereference(ctx: DereferenceContext): Expression = withOrigin(ctx) {
     val attr = ctx.fieldName.getText
     expression(ctx.base) match {
-      case UnresolvedAttribute(nameParts) =>
+      case unresolved_attr @ UnresolvedAttribute(nameParts) =>
+        if (enableHiveSupportQuotedIdentifiers) {
+          val escapedIdentifier = "`(.+)`".r
+          val ret = Option(ctx.fieldName.getStart).map(_.getText match {
+            case r@escapedIdentifier(i) =>
+              UnresolvedRegex(i, Some(unresolved_attr.name))
+            case _ =>
+              UnresolvedAttribute(nameParts :+ attr)
+          })
+          return ret.get
+        }
+
         UnresolvedAttribute(nameParts :+ attr)
       case e =>
         UnresolvedExtractValue(e, Literal(attr))
     }
   }
 
   /**
-   * Create an [[UnresolvedAttribute]] expression.
+   * Create an [[UnresolvedAttribute]] expression or a [[UnresolvedRegex]] if it is a regex
+   * quoted in ``
    */
   override def visitColumnReference(ctx: ColumnReferenceContext): Expression = withOrigin(ctx) {
+    if (enableHiveSupportQuotedIdentifiers) {
+      val escapedIdentifier = "`(.+)`".r
+      val ret = Option(ctx.getStart).map(_.getText match {
+        case r @ escapedIdentifier(i) =>
+          UnresolvedRegex(i, None)
+        case _ =>
+          UnresolvedAttribute.quoted(ctx.getText)
+      })
+
+      return ret.get
+    }
+
     UnresolvedAttribute.quoted(ctx.getText)
   }
 

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -2624,4 +2624,92 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
     val e = intercept[AnalysisException](sql("SELECT nvl(1, 2, 3)"))
     assert(e.message.contains("Invalid number of arguments"))
   }
+
+  test("SPARK-12139: REGEX Column Specification for Hive Queries") {
+    // hive.support.quoted.identifiers is turned off by default
+    checkAnswer(
+      sql(
+        """
+          |SELECT b
+          |FROM testData2
+          |WHERE a = 1
+        """.stripMargin),
+      Row(1) :: Row(2) :: Nil)
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT t.b
+          |FROM testData2 t
+          |WHERE a = 1
+        """.stripMargin),
+      Row(1) :: Row(2) :: Nil)
+
+    intercept[AnalysisException] {
+      sql(
+        """
+          |SELECT `(a)?+.+`
+          |FROM testData2
+          |WHERE a = 1
+        """.stripMargin)
+    }
+
+    intercept[AnalysisException] {
+      sql(
+        """
+          |SELECT t.`(a)?+.+`
+          |FROM testData2 t
+          |WHERE a = 1
+        """.stripMargin)
+    }
+
+    // now, turn on hive.support.quoted.identifiers
+    sparkContext.conf.set("hive.support.quoted.identifiers", "true")
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT b
+          |FROM testData2
+          |WHERE a = 1
+        """.stripMargin),
+      Row(1) :: Row(2) :: Nil)
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT t.b
+          |FROM testData2 t
+          |WHERE a = 1
+        """.stripMargin),
+      Row(1) :: Row(2) :: Nil)
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT `(a)?+.+`
+          |FROM testData2
+          |WHERE a = 1
+        """.stripMargin),
+      Row(1) :: Row(2) :: Nil)
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT t.`(a)?+.+`
+          |FROM testData2 t
+          |WHERE a = 1
+        """.stripMargin),
+      Row(1) :: Row(2) :: Nil)
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT p.`(key)?+.+`, b, testdata2.`(b)?+.+`
+          |FROM testData p join testData2
+          |ON p.key = testData2.a
+          |WHERE key < 3
+        """.stripMargin),
+      Row("1", 1, 1) :: Row("1", 2, 1) :: Row("2", 1, 2) :: Row("2", 2, 2) ::Nil)
+  }
 }