Adding resolution of complex ArrayTypes

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala

@@ -90,7 +90,7 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
       | failure("illegal character")
     )
 
-    override def identChar = letter | elem('.') | elem('_')
+    override def identChar = letter | elem('.') | elem('_') | elem('[') | elem(']')
 
     override def whitespace: Parser[Any] = rep(
       whitespaceChar
@@ -390,6 +390,7 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
     FALSE ^^^ Literal(false, BooleanType) |
     cast |
     "(" ~> expression <~ ")" |
+    "[" ~> literal <~ "]" |
     function |
     "-" ~> literal ^^ UnaryMinus |
     ident ^^ UnresolvedAttribute |

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala

@@ -20,7 +20,7 @@ package org.apache.spark.sql.catalyst.plans.logical
 import org.apache.spark.sql.catalyst.errors.TreeNodeException
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.QueryPlan
-import org.apache.spark.sql.catalyst.types.StructType
+import org.apache.spark.sql.catalyst.types.{DataType, ArrayType, StructType}
 import org.apache.spark.sql.catalyst.trees
 
 abstract class LogicalPlan extends QueryPlan[LogicalPlan] {
@@ -54,9 +54,41 @@ abstract class LogicalPlan extends QueryPlan[LogicalPlan] {
   /**
    * Optionally resolves the given string to a
    * [[catalyst.expressions.NamedExpression NamedExpression]]. The attribute is expressed as
-   * as string in the following form: `[scope].AttributeName.[nested].[fields]...`.
+   * as string in the following form: `[scope].AttributeName.[nested].[fields]...`. Fields
+   * can contain ordinal expressions, such as `field[i][j][k]...`.
    */
   def resolve(name: String): Option[NamedExpression] = {
+    def expandFunc(expType: (Expression, DataType), field: String): (Expression, DataType) = {
+      val (exp, t) = expType
+      val ordinalRegExp = """(\[(\d+)\])""".r
+      val fieldName = if (field.matches("\\w*(\\[\\d\\])+")) {
+        field.substring(0, field.indexOf("["))
+      } else {
+        field
+      }
+      t match {
+        case ArrayType(elementType) =>
+          val ordinals = ordinalRegExp.findAllIn(field).matchData.map(_.group(2))
+          (ordinals.foldLeft(exp)((v1: Expression, v2: String) => GetItem(v1, Literal(v2.toInt))), elementType)
+        case StructType(fields) =>
+          // Note: this only works if we are not on the top-level!
+          val structField = fields.find(_.name == fieldName)
+          if (!structField.isDefined) {
+            throw new TreeNodeException(
+              this, s"Trying to resolve Attribute but field ${fieldName} is not defined")
+          }
+          structField.get.dataType match {
+            case ArrayType(elementType) =>
+              val ordinals = ordinalRegExp.findAllIn(field).matchData.map(_.group(2))
+              (ordinals.foldLeft(GetField(exp, fieldName).asInstanceOf[Expression])((v1: Expression, v2: String) => GetItem(v1, Literal(v2.toInt))), elementType)
+            case _ =>
+              (GetField(exp, fieldName), structField.get.dataType)
+          }
+        case _ =>
+          expType
+      }
+    }
+
     val parts = name.split("\\.")
     // Collect all attributes that are output by this nodes children where either the first part
     // matches the name or where the first part matches the scope and the second part matches the
@@ -67,16 +99,40 @@ abstract class LogicalPlan extends QueryPlan[LogicalPlan] {
       val remainingParts =
         if (option.qualifiers.contains(parts.head) && parts.size > 1) parts.drop(1) else parts
       if (option.name == remainingParts.head) (option, remainingParts.tail.toList) :: Nil else Nil
+      // TODO from rebase!
+      /*val remainingParts = if (option.qualifiers contains parts.head) parts.drop(1) else parts
+      val relevantRemaining =
+        if (remainingParts.head.matches("\\w*\\[(\\d+)\\]")) { // array field name
+          remainingParts.head.substring(0, remainingParts.head.indexOf("["))
+        } else {
+          remainingParts.head
+        }
+      if (option.name == relevantRemaining) (option, remainingParts.tail.toList) :: Nil else Nil*/
     }
 
     options.distinct match {
-      case (a, Nil) :: Nil => Some(a) // One match, no nested fields, use it.
+      case (a, Nil) :: Nil => {
+        a.dataType match {
+          case ArrayType(elementType) =>
+            val expression = expandFunc((a: Expression, a.dataType), name)._1
+            Some(Alias(expression, name)())
+          case _ => Some(a)
+        }
+      } // One match, no nested fields, use it.
       // One match, but we also need to extract the requested nested field.
       case (a, nestedFields) :: Nil =>
         a.dataType match {
           case StructType(fields) =>
-            Some(Alias(nestedFields.foldLeft(a: Expression)(GetField), nestedFields.last)())
-          case _ => None // Don't know how to resolve these field references
+            // this is compatibility reasons with earlier code! TODO: why only nestedFields and not parts?
+            if ((parts(0) :: nestedFields).forall(!_.matches("\\w*\\[\\d+\\]+"))) { // not nested arrays, only fields
+              Some(Alias(nestedFields.foldLeft(a: Expression)(GetField), nestedFields.last)())
+            } else {
+              val expression = parts.foldLeft((a: Expression, a.dataType))(expandFunc)._1
+              Some(Alias(expression, nestedFields.last)())
+            }
+          case _ =>
+            val expression = parts.foldLeft((a: Expression, a.dataType))(expandFunc)._1
+            Some(Alias(expression, nestedFields.last)())
         }
       case Nil => None         // No matches.
       case ambiguousReferences =>

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/types/dataTypes.scala

@@ -29,11 +29,15 @@ abstract class DataType {
     case e: Expression if e.dataType == this => true
     case _ => false
   }
+
+  def isPrimitive(): Boolean = false
 }
 
 case object NullType extends DataType
 
-trait PrimitiveType
+trait PrimitiveType extends DataType {
+  override def isPrimitive() = true
+}
 
 abstract class NativeType extends DataType {
   type JvmType

sql/core/src/main/scala/org/apache/spark/sql/parquet/ParquetTestData.scala

@@ -139,7 +139,7 @@ private[sql] object ParquetTestData {
         |optional group longs {
           |repeated int64 values;
         |}
-        |required group booleanNumberPairs {
+        |repeated group entries {
           |required double value;
           |optional boolean truth;
         |}
@@ -153,8 +153,23 @@ private[sql] object ParquetTestData {
       |}
     """.stripMargin
 
+  val testNestedSchema3 =
+    """
+      |message TestNested3 {
+      |required int32 x;
+        |repeated group booleanNumberPairs {
+          |required int32 key;
+          |repeated group value {
+            |required double nestedValue;
+            |optional boolean truth;
+          |}
+        |}
+      |}
+    """.stripMargin
+
   val testNestedDir1 = Utils.createTempDir()
   val testNestedDir2 = Utils.createTempDir()
+  val testNestedDir3 = Utils.createTempDir()
 
   lazy val testNestedData1 = new ParquetRelation(testNestedDir1.toURI.toString)
   lazy val testNestedData2 = new ParquetRelation(testNestedDir2.toURI.toString)
@@ -285,6 +300,32 @@ private[sql] object ParquetTestData {
     writer.close()
   }
 
+  def writeNestedFile3() {
+    testNestedDir3.delete()
+    val path: Path = testNestedDir3
+    val schema: MessageType = MessageTypeParser.parseMessageType(testNestedSchema3)
+
+    val r1 = new SimpleGroup(schema)
+    r1.add(0, 1)
+    val g1 = r1.addGroup(1)
+    g1.add(0, 1)
+    val ng1 = g1.addGroup(1)
+    ng1.add(0, 1.5)
+    ng1.add(1, false)
+    val ng2 = g1.addGroup(1)
+    ng2.add(0, 2.5)
+    ng2.add(1, true)
+    val g2 = r1.addGroup(1)
+    g2.add(0, 2)
+    val ng3 = g2.addGroup(1)
+    ng3.add(0, 3.5)
+    ng3.add(1, false)
+
+    val writeSupport = new TestGroupWriteSupport(schema)
+    val writer = new ParquetWriter[Group](path, writeSupport)
+    writer.write(r1)
+    writer.close()
+  }
 
   def readNestedFile(path: File, schemaString: String): Unit = {
     val configuration = new Configuration()

sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetQuerySuite.scala

@@ -34,6 +34,11 @@ import org.apache.spark.sql.SchemaRDD
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.types.IntegerType
 import org.apache.spark.util.Utils
+import org.apache.spark.sql.catalyst.types.{StringType, IntegerType, DataType}
+import org.apache.spark.sql.{parquet, SchemaRDD}
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import scala.Tuple2
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 
 // Implicits
 import org.apache.spark.sql.test.TestSQLContext._
@@ -432,9 +437,9 @@ class ParquetQuerySuite extends QueryTest with FunSuiteLike with BeforeAndAfterA
     assert(result(0)(2)(0) === (1.toLong << 32))
     assert(result(0)(2)(1) === (1.toLong << 33))
     assert(result(0)(2)(2) === (1.toLong << 34))
-    assert(result(0)(3).size === 2)
-    assert(result(0)(3)(0) === 2.5)
-    assert(result(0)(3)(1) === false)
+    assert(result(0)(3)(0).size === 2)
+    assert(result(0)(3)(0)(0) === 2.5)
+    assert(result(0)(3)(0)(1) === false)
     assert(result(0)(4).size === 2)
     assert(result(0)(4)(0).size === 2)
     assert(result(0)(4)(1).size === 1)
@@ -452,23 +457,61 @@ class ParquetQuerySuite extends QueryTest with FunSuiteLike with BeforeAndAfterA
     assert(tmp(0)(0) === "Julien Le Dem")
   }
 
+  test("Projection in addressbook") {
+    implicit def anyToRow(value: Any): Row = value.asInstanceOf[Row]
+    val data = TestSQLContext
+      .parquetFile(ParquetTestData.testNestedDir1.toString)
+      .toSchemaRDD
+    data.registerAsTable("data")
+    val tmp = sql("SELECT owner, contacts[1].name FROM data").collect()
+    assert(tmp.size === 2)
+    assert(tmp(0).size === 2)
+    assert(tmp(0)(0) === "Julien Le Dem")
+    assert(tmp(0)(1) === "Chris Aniszczyk")
+    assert(tmp(1)(0) === "A. Nonymous")
+    assert(tmp(1)(1) === null)
+  }
+
   test("Simple query on nested int data") {
     implicit def anyToRow(value: Any): Row = value.asInstanceOf[Row]
     val data = TestSQLContext
       .parquetFile(ParquetTestData.testNestedDir2.toString)
       .toSchemaRDD
     data.registerAsTable("data")
-    val tmp = sql("SELECT booleanNumberPairs.value, booleanNumberPairs.truth FROM data").collect()
-    assert(tmp(0)(0) === 2.5)
-    assert(tmp(0)(1) === false)
-    val result = sql("SELECT outerouter FROM data").collect()
-    // TODO: why does this not work?
-    //val result = sql("SELECT outerouter.values FROM data").collect()
-    // TODO: .. or this:
-    // val result = sql("SELECT outerouter[0] FROM data").collect()
-    assert(result(0)(0)(0)(0)(0) === 7)
-    assert(result(0)(0)(0)(1)(0) === 8)
-    assert(result(0)(0)(1)(0)(0) === 9)
+    val result1 = sql("SELECT entries[0].value FROM data").collect()
+    assert(result1.size === 1)
+    assert(result1(0).size === 1)
+    assert(result1(0)(0) === 2.5)
+    val result2 = sql("SELECT entries[0] FROM data").collect()
+    assert(result2.size === 1)
+    assert(result2(0)(0).size === 2)
+    assert(result2(0)(0)(0) === 2.5)
+    assert(result2(0)(0)(1) === false)
+    val result3 = sql("SELECT outerouter FROM data").collect()
+    assert(result3(0)(0)(0)(0)(0) === 7)
+    assert(result3(0)(0)(0)(1)(0) === 8)
+    assert(result3(0)(0)(1)(0)(0) === 9)
+  }
+
+  test("nested structs") {
+    implicit def anyToRow(value: Any): Row = value.asInstanceOf[Row]
+    ParquetTestData.writeNestedFile3()
+    val data = TestSQLContext
+      .parquetFile(ParquetTestData.testNestedDir3.toString)
+      .toSchemaRDD
+    data.registerAsTable("data")
+    val result1 = sql("SELECT booleanNumberPairs[0].value[0].truth FROM data").collect()
+    assert(result1.size === 1)
+    assert(result1(0).size === 1)
+    assert(result1(0)(0) === false)
+    val result2 = sql("SELECT booleanNumberPairs[0].value[1].truth FROM data").collect()
+    assert(result2.size === 1)
+    assert(result2(0).size === 1)
+    assert(result2(0)(0) === true)
+    val result3 = sql("SELECT booleanNumberPairs[1].value[0].truth FROM data").collect()
+    assert(result3.size === 1)
+    assert(result3(0).size === 1)
+    assert(result3(0)(0) === false)
   }
 
   /**