Merge pull request #914 from FlorentinD/listComprehension

Implement list-comprehension

Author: DarthMax

morpheus-spark-cypher/src/main/scala/org/opencypher/morpheus/impl/SparkSQLExprMapper.scala

@@ -26,7 +26,7 @@
  */
 package org.opencypher.morpheus.impl
 
-import org.apache.spark.sql.catalyst.expressions.CaseWhen
+import org.apache.spark.sql.catalyst.expressions.{ArrayFilter, ArrayTransform, CaseWhen, ExprId, LambdaFunction, NamedLambdaVariable}
 import org.apache.spark.sql.functions.{array_contains => _, translate => _, _}
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.{Column, DataFrame}
@@ -88,6 +88,9 @@ object SparkSQLExprMapper {
       */
     def asSparkSQLExpr(implicit header: RecordHeader, df: DataFrame, parameters: CypherMap): Column = {
       val outCol = expr match {
+        case v: LambdaVar =>
+          val sparkType = v.cypherType.toSparkType.getOrElse(throw IllegalStateException(s"No valid dataType for LambdaVar $v"))
+          new Column(NamedLambdaVariable(v.name, sparkType, nullable = v.cypherType.isNullable, ExprId(v.hashCode.toLong)))
         // Evaluate based on already present data; no recursion
         case _: Var | _: HasLabel | _: HasType | _: StartNode | _: EndNode => column_for(expr)
         // Evaluate bottom-up
@@ -329,6 +332,25 @@ object SparkSQLExprMapper {
         case _: ListSliceFrom => list_slice(child0, Some(child1), None)
         case _: ListSliceTo => list_slice(child0, None, Some(child1))
 
+        case ListComprehension(variable, innerPredicate, extractExpression, listExpr) =>
+          val lambdaVar = variable.asSparkSQLExpr.expr match {
+            case v: NamedLambdaVariable => v
+            case err => throw IllegalStateException(s"$variable should be converted into a NamedLambdaVariable instead of $err")
+          }
+          val filteredExpr = innerPredicate match {
+            case Some(filterExpr) =>
+              val filterFunc = LambdaFunction(filterExpr.asSparkSQLExpr.expr, Seq(lambdaVar))
+              ArrayFilter(listExpr.asSparkSQLExpr.expr, filterFunc)
+            case None => listExpr.asSparkSQLExpr.expr
+          }
+          val result = extractExpression match{
+            case Some(extractExpr) =>
+              val extractFunc = LambdaFunction(extractExpr.asSparkSQLExpr.expr, Seq(lambdaVar))
+              ArrayTransform(filteredExpr, extractFunc)
+            case None => filteredExpr
+          }
+          new Column(result)
+
         case MapExpression(items) => expr.cypherType.material match {
           case CTMap(_) =>
             val innerColumns = items.map {

morpheus-tck/src/test/resources/failing_blacklist

@@ -20,12 +20,9 @@ Feature "UnwindAcceptance": Scenario "Unwinding a collected expression"
 Feature "UnwindAcceptance": Scenario "Unwind does not remove variables from scope"
 Feature "TypeConversionFunctions": Scenario "`toInteger()` handling mixed number types"
 Feature "TypeConversionFunctions": Scenario "`toInteger()` handling Any type"
-Feature "TypeConversionFunctions": Scenario "`toInteger()` on a list of strings"
 Feature "TypeConversionFunctions": Scenario "`toFloat()` on mixed number types"
 Feature "TypeConversionFunctions": Scenario "`toFloat()` handling Any type"
-Feature "TypeConversionFunctions": Scenario "`toFloat()` on a list of strings"
 Feature "TypeConversionFunctions": Scenario "`toString()` should work on Any type"
-Feature "TypeConversionFunctions": Scenario "`toString()` on a list of integers"
 Feature "TypeConversionFunctions": Scenario "`toBoolean()` on invalid types #1"
 Feature "TypeConversionFunctions": Scenario "`toBoolean()` on invalid types #2"
 Feature "TypeConversionFunctions": Scenario "`toBoolean()` on invalid types #3"
@@ -192,7 +189,6 @@ Feature "Aggregation": Scenario "`min()` over mixed values"
 Feature "Aggregation": Scenario "`max()` over mixed values"
 Feature "Aggregation": Scenario "`max()` over mixed numeric values"
 Feature "ListOperations": Scenario "IN should return true if correct list found despite other lists having nulls"
-Feature "ListOperations": Scenario "Size of list comprehension"
 Feature "ListOperations": Scenario "IN should return false when matching a number with a string - list version"
 Feature "ListOperations": Scenario "IN should return false when matching a number with a string"
 Feature "ListOperations": Scenario "IN should return true when LHS and RHS contain a nested list - singleton version"

morpheus-testing/src/test/scala/org/opencypher/morpheus/impl/acceptance/ExpressionTests.scala

@@ -1468,4 +1468,78 @@ class ExpressionTests extends MorpheusTestSuite with ScanGraphInit with Checkers
       }
     }
   }
+
+  describe("list comprehension") {
+    it("supports list comprehension with static mapping") {
+      val result = morpheus.cypher(
+        """
+          |WITH [1, 2, 3] AS things
+          |RETURN [n IN things | 1] AS value
+        """.stripMargin)
+
+      result.records.toMaps shouldEqual Bag(
+        CypherMap("value" -> List(1, 1, 1))
+      )
+    }
+
+    it("supports list comprehension with simple mapping") {
+      val result = morpheus.cypher(
+        """
+          |WITH [1, 2, 3] AS things
+          |RETURN [n IN things | n*3] AS value
+        """.stripMargin)
+
+      result.records.toMaps shouldEqual Bag(
+        CypherMap("value" -> List(3, 6, 9))
+      )
+    }
+
+    it("supports list comprehension with more complex mapping") {
+      val result = morpheus.cypher(
+        """
+          |WITH ['1', '2', '3'] AS things
+          |RETURN [n IN things | toInteger(n)*3 + toInteger(n)] AS value
+        """.stripMargin)
+
+      result.records.toMaps shouldEqual Bag(
+        CypherMap("value" -> List(4, 8, 12))
+      )
+    }
+
+    it("supports list comprehension with inner predicate") {
+      val result = morpheus.cypher(
+        """
+          |WITH [1, 2, 3] AS things
+          |RETURN [n IN things WHERE n > 2] AS value
+        """.stripMargin)
+
+      result.records.toMaps shouldEqual Bag(
+        CypherMap("value" -> List(3))
+      )
+    }
+
+    it("supports list comprehension with inner predicate and more complex mapping") {
+      val result = morpheus.cypher(
+        """
+          |WITH ['1', '2', '3'] AS things
+          |RETURN [n IN things WHERE toInteger(n) > 2 | toInteger(n)*3 + toInteger(n)] AS value
+        """.stripMargin)
+
+      result.records.toMaps shouldEqual Bag(
+        CypherMap("value" -> List(12))
+      )
+    }
+
+    it("supports nested list comprehensions") {
+      val result = morpheus.cypher(
+        """
+          |WITH [[1,2,3], [2,2,3], [3,4]] AS things
+          |RETURN [n IN things | [n IN n WHERE n < 2]] AS value
+        """.stripMargin)
+
+      result.records.toMaps shouldEqual Bag(
+        CypherMap("value" -> List(List(1), List(), List()))
+      )
+    }
+  }
 }

okapi-ir/src/main/scala/org/opencypher/okapi/ir/api/expr/Expr.scala

@@ -196,6 +196,8 @@ final case class SimpleVar(name: String)(val cypherType: CypherType) extends Ret
   override def withOwner(expr: Var): SimpleVar = SimpleVar(expr.name)(expr.cypherType)
 }
 
+final case class LambdaVar(name: String)(val cypherType: CypherType) extends Var
+
 final case class StartNode(rel: Expr)(val cypherType: CypherType) extends Expr {
   type This = StartNode
 
@@ -1099,6 +1101,18 @@ final case class ListSliceFrom(list: Expr, from: Expr) extends ListSlice(Some(fr
 
 final case class ListSliceTo(list: Expr, to: Expr) extends ListSlice(None, Some(to))
 
+final case class ListComprehension(variable: Expr, innerPredicate: Option[Expr], extractExpression: Option[Expr], expr : Expr) extends Expr {
+  override def withoutType: String = {
+    val p = innerPredicate.map(" WHERE " + _.withoutType).getOrElse("")
+    val e = extractExpression.map(" | " + _.withoutType).getOrElse("")
+    s"[${variable.withoutType} IN ${expr.withoutType}$p$e]"
+  }
+  override def cypherType: CypherType = extractExpression match {
+    case Some(x) => CTList(x.cypherType)
+    case None => expr.cypherType
+  }
+}
+
 final case class ContainerIndex(container: Expr, index: Expr)(val cypherType: CypherType) extends Expr {
 
   override def withoutType: String = s"${container.withoutType}[${index.withoutType}]"

okapi-ir/src/main/scala/org/opencypher/okapi/ir/impl/ExpressionConverter.scala

@@ -34,7 +34,7 @@ import org.opencypher.okapi.ir.api.expr._
 import org.opencypher.okapi.ir.impl.parse.{functions => f}
 import org.opencypher.okapi.ir.impl.typer.SignatureConverter.Signature
 import org.opencypher.okapi.ir.impl.typer.{InvalidArgument, InvalidContainerAccess, MissingParameter, UnTypedExpr, WrongNumberOfArguments}
-import org.opencypher.v9_0.expressions.{RegexMatch, functions}
+import org.opencypher.v9_0.expressions.{ExtractScope, LogicalVariable, RegexMatch, functions}
 import org.opencypher.v9_0.{expressions => ast}
 import SignatureTyping._
 
@@ -62,18 +62,21 @@ final class ExpressionConverter(context: IRBuilderContext) {
     }
   }
 
-  def convert(e: ast.Expression): Expr = {
+  def convert(e: ast.Expression) (implicit lambdaVars: Map[String, CypherType]): Expr = {
 
     lazy val child0: Expr = convert(e.arguments.head)
 
     lazy val child1: Expr = convert(e.arguments(1))
 
     lazy val child2: Expr = convert(e.arguments(2))
 
-    lazy val convertedChildren: List[Expr] = e.arguments.toList.map(convert)
+    lazy val convertedChildren: List[Expr] = e.arguments.toList.map(convert(_))
 
     e match {
-      case ast.Variable(name) => Var(name)(context.knownTypes.getOrElse(e, throw UnTypedExpr(e)))
+      case ast.Variable(name) => lambdaVars.get(name) match {
+          case Some(varType) => LambdaVar(name)(varType)
+          case None => Var(name)(context.knownTypes.getOrElse(e, throw UnTypedExpr(e)))
+        }
       case p@ast.Parameter(name, _) => Param(name)(parameterType(p))
 
       // Literals
@@ -276,6 +279,16 @@ final class ExpressionConverter(context: IRBuilderContext) {
       case ast.ListSlice(list, None, Some(to)) => ListSliceTo(convert(list), convert(to))
       case ast.ListSlice(list, Some(from), None) => ListSliceFrom(convert(list), convert(from))
 
+      case ast.ListComprehension(ExtractScope(variable, innerPredicate, extractExpression), expr) =>
+        val listExpr = convert(expr)(lambdaVars)
+        val listInnerType = listExpr.cypherType match {
+          case CTList(inner) => inner
+          case err => throw IllegalArgumentException("a list to step over", err, "Wrong list comprehension type")
+        }
+        val updatedLambdaVars: Map[String, CypherType] = lambdaVars + (variable.name -> listInnerType)
+        ListComprehension(convert(variable)(updatedLambdaVars), innerPredicate.map(convert(_)(updatedLambdaVars)),
+          extractExpression.map(convert(_)(updatedLambdaVars)), listExpr)
+
       case ast.ContainerIndex(container, index) =>
         val convertedContainer = convert(container)
         val elementType = convertedContainer.cypherType.material match {

okapi-ir/src/main/scala/org/opencypher/okapi/ir/impl/IRBuilderContext.scala

@@ -61,7 +61,7 @@ final case class IRBuilderContext(
     patternConverter.convert(p, knownTypes, qgn.getOrElse(workingGraph.qualifiedGraphName))
   }
 
-  def convertExpression(e: ast.Expression): Expr = exprConverter.convert(e)
+  def convertExpression(e: ast.Expression): Expr = exprConverter.convert(e)(lambdaVars = Map())
 
   def schemaFor(qgn: QualifiedGraphName): PropertyGraphSchema = queryLocalCatalog.schema(qgn)
 

okapi-ir/src/test/scala/org/opencypher/okapi/ir/impl/ExpressionConverterTest.scala

@@ -439,10 +439,29 @@ class ExpressionConverterTest extends BaseTestSuite with Neo4jAstTestSupport {
     )
   }
 
+  describe("list comprehension") {
+    val intVar = LambdaVar("x")(CTInteger)
+    it("can convert list comprehension with static mapping") {
+      convert("[x IN [1,2] | 1]") shouldEqual ListComprehension(intVar, None, Some(IntegerLit(1)), ListLit(List(IntegerLit(1), IntegerLit(2))))
+    }
+
+    it("can convert list comprehension with unary mapping") {
+      convert("[x IN [1,2] | toString(x)]")
+    }
+
+    it("can convert list comprehension with 2 var-calls") {
+      convert("[x IN [1,2] | x + x * 2]") shouldEqual ListComprehension(intVar, None, Some(Add(intVar, Multiply(intVar, IntegerLit(2)))), ListLit(List(IntegerLit(1), IntegerLit(2))))
+    }
+
+    it("can convert list comprehension with inner predicate") {
+      convert("[x IN [1,2] WHERE x < 1 | 1]") shouldEqual ListComprehension(intVar, Some(LessThan(intVar, IntegerLit(1))), Some(IntegerLit(1)), ListLit(List(IntegerLit(1), IntegerLit(2))))
+    }
+  }
+
   implicit def toVar(s: Symbol): Var = all.find(_.name == s.name).get
 
   private def convert(e: ast.Expression): Expr =
-    new ExpressionConverter(testContext).convert(e)
+    new ExpressionConverter(testContext).convert(e)(Map())
 
   implicit class TestExpr(expr: Expr) {
     def shouldEqual(other: Expr): Assertion = {