make catalyst expressions sql compatible (#16)

Author: fqaiser94

README.md

@@ -31,7 +31,7 @@ As such, you can pull in the current stable release by simply adding a library d
 For example, for an SBT project, simply add the following line to your `build.sbt`:
 
 ```
-libraryDependencies += "com.github.fqaiser94" %% "mse" % "0.2.2"
+libraryDependencies += "com.github.fqaiser94" %% "mse" % "0.2.4"
 ```
 
 For other types of projects (e.g. Maven, Gradle), see the installation instructions at this [link](https://search.maven.org/artifact/com.github.fqaiser94/mse_2.11). 
@@ -43,7 +43,7 @@ You will also need to provide your PySpark application/s with the path to the MS
 
 ```bash
 pip install mse
-curl https://repo1.maven.org/maven2/com/github/fqaiser94/mse_2.11/0.2.2/mse_2.11-0.2.2.jar --output mse.jar
+curl https://repo1.maven.org/maven2/com/github/fqaiser94/mse_2.11/0.2.4/mse_2.11-0.2.4.jar --output mse.jar
 pyspark --jars mse.jar
 ```
 
@@ -345,6 +345,12 @@ arrayOfStructs.withColumn("array", transform($"array", elem => elem.dropFields("
 // +----------------+
 ```
 
+# SQL installation and usage
+
+The underlying Catalyst Expressions are SQL compatible. 
+Unfortunately, Spark only added public APIs for plugging in custom Catalyst Expressions into the FunctionRegistry in Spark 3.0.0 
+(which is at the time of writing is still in preview). You can find a project with an example of how to do this [here](https://github.com/fqaiser94/mse-sql-example). 
+
 # Catalyst Optimization Rules
 
 We also provide some Catalyst optimization rules that can be plugged into a Spark session to get even better performance. This is as simple as including the following two lines of code at the start of your Scala Spark program:  
@@ -386,7 +392,7 @@ As you can see, the successive `add_fields` method calls have been collapsed int
 
 Theoretically, this should improve performance but for the most part, you won't notice much difference unless you're doing some particularly intense struct manipulation and/or working with a particularly large dataset.  
 
-Unfortunately, to the best of our knowledge, there is currently no way to plug in custom Catalyst optimization rules using the Python APIs. 
+Unfortunately, to the best of our knowledge, there is currently no way to plug in custom Catalyst optimization rules directly using the Python APIs. 
 
 # Questions/Thoughts/Concerns?
 

python/README.md

@@ -21,7 +21,7 @@ For example:
 
 ```bash
 pip install mse
-curl https://repo1.maven.org/maven2/com/github/fqaiser94/mse_2.11/0.2.2/mse_2.11-0.2.2.jar --output mse.jar
+curl https://repo1.maven.org/maven2/com/github/fqaiser94/mse_2.11/0.2.4/mse_2.11-0.2.4.jar --output mse.jar
 pyspark --jars mse.jar
 ```
 

python/mse/tests/methods_tests.py

@@ -69,13 +69,13 @@ def test_withField(self):
         with self.subTest("throw error if withField is called on a column that is not struct dataType"):
             self.assertRaisesRegex(
                 AnalysisException,
-                "struct should be struct data type. struct is integer",
+                "Only struct is allowed to appear at first position, got: integer",
                 lambda: non_struct_df.withColumn("a", col("a").withField("a", lit(2))).collect())
 
         with self.subTest("throw error if null fieldName supplied"):
             self.assertRaisesRegex(
                 AnalysisException,
-                "fieldNames cannot contain null",
+                "Only non-null foldable string expressions are allowed to appear at even position",
                 lambda: struct_level1_df.withColumn("a", col("a").withField(None, lit(2))).collect())
 
         with self.subTest("add new field to struct"):
@@ -194,19 +194,19 @@ def test_withFieldRenamed(self):
         with self.subTest("throw error if withFieldRenamed is called on a column that is not struct dataType"):
             self.assertRaisesRegex(
                 AnalysisException,
-                "struct should be struct data type. struct is integer",
+                "Only struct is allowed to appear at first position, got: integer",
                 lambda: non_struct_df.withColumn("a", col("a").withFieldRenamed("a", "z")))
 
         with self.subTest("throw error if null existingFieldName supplied"):
             self.assertRaisesRegex(
                 AnalysisException,
-                "existingFieldName cannot be null",
+                "Only non-null foldable string expressions are allowed to appear after first position.",
                 lambda: struct_df.withColumn("a", col("a").withFieldRenamed(None, "z")))
 
         with self.subTest("throw error if null newFieldName supplied"):
             self.assertRaisesRegex(
                 AnalysisException,
-                "newFieldName cannot be null",
+                "Only non-null foldable string expressions are allowed to appear after first position.",
                 lambda: struct_df.withColumn("a", col("a").withFieldRenamed("a", None)))
 
         with self.subTest("rename field in struct"):
@@ -286,13 +286,13 @@ def test_dropFields(self):
         with self.subTest("throw error if withField is called on a column that is not struct dataType"):
             self.assertRaisesRegex(
                 AnalysisException,
-                "struct should be struct data type. struct is integer",
+                "Only struct is allowed to appear at first position, got: integer",
                 lambda: non_struct_df.withColumn("a", col("a").dropFields("a")))
 
         with self.subTest("throw error if null fieldName supplied"):
             self.assertRaisesRegex(
                 AnalysisException,
-                "fieldNames cannot contain null",
+                "Only non-null foldable string expressions are allowed after first position.",
                 lambda: struct_df.withColumn("a", col("a").dropFields(None)))
 
         with self.subTest("drop field in struct"):

python/setup.py

@@ -8,7 +8,7 @@
 setuptools.setup(
     name='mse',
     packages=['mse'],
-    version='0.1.3',
+    version='0.1.4',
     license='Apache license 2.0',
     description='Make Structs Easy (MSE)',
     long_description=long_description,

python/testing/ReusedPySparkTestCase.py

@@ -28,7 +28,7 @@ def conf(cls):
         Override this in subclasses to supply a more specific conf
         """
 
-        return SparkConf().set("spark.driver.extraClassPath", "./mse.jar")
+        return SparkConf().set("spark.jars", "./mse.jar")
 
     @classmethod
     def setUpClass(cls):

src/main/scala/com/github/fqaiser94/mse/methods.scala

@@ -1,9 +1,9 @@
 package com.github.fqaiser94.mse
 
-import org.apache.spark.sql.{Column, ColumnName}
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
-import org.apache.spark.sql.catalyst.expressions.{AddFields, DropFields, Expression, RenameFields}
+import org.apache.spark.sql.catalyst.expressions.{AddFields, DropFields, Expression, Literal, RenameFields}
 import org.apache.spark.sql.functions._
+import org.apache.spark.sql.{Column, ColumnName}
 
 object methods {
 
@@ -21,7 +21,7 @@ object methods {
       * @since 2.4.4
       */
     def withField(fieldName: String, fieldValue: Column): Column = withExpr {
-      AddFields(expr, fieldName, fieldValue.expr)
+      AddFields(expr :: Literal(fieldName) :: fieldValue.expr :: Nil)
     }
 
     /**
@@ -33,7 +33,7 @@ object methods {
       * @since 2.4.4
       */
     def dropFields(fieldNames: String*): Column = withExpr {
-      DropFields(expr, fieldNames: _*)
+      DropFields(expr +: fieldNames.toList.map(Literal(_)))
     }
 
     /**
@@ -45,7 +45,7 @@ object methods {
       * @since 2.4.4
       */
     def withFieldRenamed(existingFieldName: String, newFieldName: String): Column = withExpr {
-      RenameFields(expr, existingFieldName, newFieldName)
+      RenameFields(expr :: Literal(existingFieldName) :: Literal(newFieldName) :: Nil)
     }
 
     private def withExpr(newExpr: Expression): Column = new Column(newExpr)

src/main/scala/org/apache/spark/sql/catalyst/expressions/AddFields.scala

@@ -4,55 +4,67 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodeGenerator, CodegenContext, ExprCode, FalseLiteral}
-import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.types.{StringType, StructField, StructType}
+import org.apache.spark.unsafe.types.UTF8String
 
 /**
-  *
-  * Adds/replaces fields in a struct.
-  * Returns null if struct is null.
-  * If there are multiple existing fields with the one of the fieldNames, they will all be replaced.
-  *
-  * @param struct           : The struct to add fields to.
-  * @param fieldNames       : The names of the fieldExpressions to add to given struct.
-  * @param fieldExpressions : The expressions to assign to each fieldName in fieldNames.
-  */
+ * Adds/replaces fields in a struct.
+ * Returns null if struct is null.
+ * If multiple fields already exist with the one of the given fieldNames, they will all be replaced.
+ */
 // scalastyle:off line.size.limit
 @ExpressionDescription(
-  usage = "_FUNC_(struct, fieldName, field) - Adds/replaces field in given struct.",
-  examples =
-    """
+  usage = "_FUNC_(struct, name1, val1, name2, val2, ...) - Adds/replaces fields in struct by name.",
+  examples = """
     Examples:
-      > SELECT _FUNC_({"a":1}, "b", 2);
-       {"a":1,"b":2}
+      > SELECT _FUNC_({"a":1}, "b", 2, "c", 3);
+       {"a":1,"b":2,"c":3}
   """)
 // scalastyle:on line.size.limit
-case class AddFields(struct: Expression, fieldNames: Seq[String], fieldExpressions: Seq[Expression]) extends Expression {
+case class AddFields(children: Seq[Expression]) extends Expression {
 
-  override def children: Seq[Expression] = struct +: fieldExpressions
+  private lazy val struct: Expression = children.head
+  private lazy val (nameExprs, valExprs) = children.drop(1).grouped(2).map {
+    case Seq(name, value) => (name, value)
+  }.toList.unzip
+  private lazy val fieldNames = nameExprs.map(_.eval().asInstanceOf[UTF8String].toString)
+  private lazy val pairs = fieldNames.zip(valExprs)
+
+  override def nullable: Boolean = struct.nullable
+
+  private lazy val ogStructType: StructType =
+    struct.dataType.asInstanceOf[StructType]
 
   override lazy val dataType: StructType = {
     val existingFields = ogStructType.fields.map { x => (x.name, x) }
-    val addOrReplaceFields = pairs.map { case (fieldName, field) => (fieldName, StructField(fieldName, field.dataType, field.nullable)) }
+    val addOrReplaceFields = pairs.map { case (fieldName, field) =>
+      (fieldName, StructField(fieldName, field.dataType, field.nullable))
+    }
     val newFields = loop(existingFields, addOrReplaceFields).map(_._2)
     StructType(newFields)
   }
 
-  override def nullable: Boolean = struct.nullable
-
   override def checkInputDataTypes(): TypeCheckResult = {
+    if (children.size % 2 == 0) {
+      return TypeCheckResult.TypeCheckFailure(s"$prettyName expects an odd number of arguments.")
+    }
+
     val typeName = struct.dataType.typeName
-    if (typeName != StructType(Nil).typeName)
+    val expectedStructType = StructType(Nil).typeName
+    if (typeName != expectedStructType) {
       return TypeCheckResult.TypeCheckFailure(
-        s"struct should be struct data type. struct is $typeName")
-
-    if (fieldNames.contains(null))
-      return TypeCheckResult.TypeCheckFailure("fieldNames cannot contain null")
+        s"Only $expectedStructType is allowed to appear at first position, got: $typeName.")
+    }
 
-    if (fieldExpressions.contains(null))
-      return TypeCheckResult.TypeCheckFailure("fieldExpressions cannot contain null")
+    if (nameExprs.contains(null) || nameExprs.exists(e => !(e.foldable && e.dataType == StringType))) {
+      return TypeCheckResult.TypeCheckFailure(
+        s"Only non-null foldable ${StringType.catalogString} expressions are allowed to appear at even position.")
+    }
 
-    if (fieldNames.size != fieldExpressions.size)
-      return TypeCheckResult.TypeCheckFailure("fieldNames and fieldExpressions cannot have different lengths")
+    if (valExprs.contains(null)) {
+      return TypeCheckResult.TypeCheckFailure(
+        s"Only non-null expressions are allowed to appear at odd positions after first position.")
+    }
 
     TypeCheckResult.TypeCheckSuccess
   }
@@ -62,32 +74,36 @@ case class AddFields(struct: Expression, fieldNames: Seq[String], fieldExpressio
     if (structValue == null) {
       null
     } else {
-      val existingValues: Seq[(FieldName, Any)] = ogStructType.fieldNames.zip(structValue.asInstanceOf[InternalRow].toSeq(ogStructType))
-      val addOrReplaceValues: Seq[(FieldName, Any)] = pairs.map { case (fieldName, expression) => (fieldName, expression.eval(input)) }
+      val existingValues: Seq[(FieldName, Any)] =
+        ogStructType.fieldNames.zip(structValue.asInstanceOf[InternalRow].toSeq(ogStructType))
+      val addOrReplaceValues: Seq[(FieldName, Any)] =
+        pairs.map { case (fieldName, expression) => (fieldName, expression.eval(input)) }
       val newValues = loop(existingValues, addOrReplaceValues).map(_._2)
       InternalRow.fromSeq(newValues)
     }
   }
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     val structGen = struct.genCode(ctx)
-    val addOrReplaceFieldsGens = fieldExpressions.map(_.genCode(ctx))
+    val addOrReplaceFieldsGens = valExprs.map(_.genCode(ctx))
     val resultCode: String = {
       val structVar = structGen.value
       type NullCheck = String
       type NonNullValue = String
-      val existingFieldsCode: Seq[(FieldName, (NullCheck, NonNullValue))] = ogStructType.fields.zipWithIndex.map {
-        case (structField, i) =>
-          val nullCheck = s"$structVar.isNullAt($i)"
-          val nonNullValue = CodeGenerator.getValue(structVar, structField.dataType, i.toString)
-          (structField.name, (nullCheck, nonNullValue))
-      }
-      val addOrReplaceFieldsCode: Seq[(FieldName, (NullCheck, NonNullValue))] = fieldNames.zip(addOrReplaceFieldsGens).map {
-        case (fieldName, fieldExprCode) =>
-          val nullCheck = fieldExprCode.isNull.code
-          val nonNullValue = fieldExprCode.value.code
-          (fieldName, (nullCheck, nonNullValue))
-      }
+      val existingFieldsCode: Seq[(FieldName, (NullCheck, NonNullValue))] =
+        ogStructType.fields.zipWithIndex.map {
+          case (structField, i) =>
+            val nullCheck = s"$structVar.isNullAt($i)"
+            val nonNullValue = CodeGenerator.getValue(structVar, structField.dataType, i.toString)
+            (structField.name, (nullCheck, nonNullValue))
+        }
+      val addOrReplaceFieldsCode: Seq[(FieldName, (NullCheck, NonNullValue))] =
+        fieldNames.zip(addOrReplaceFieldsGens).map {
+          case (fieldName, fieldExprCode) =>
+            val nullCheck = fieldExprCode.isNull.code
+            val nonNullValue = fieldExprCode.value.code
+            (fieldName, (nullCheck, nonNullValue))
+        }
       val newFieldsCode = loop(existingFieldsCode, addOrReplaceFieldsCode)
       val rowClass = classOf[GenericInternalRow].getName
       val rowValuesVar = ctx.freshName("rowValues")
@@ -138,17 +154,14 @@ case class AddFields(struct: Expression, fieldNames: Seq[String], fieldExpressio
 
   override def prettyName: String = "add_fields"
 
-  private lazy val ogStructType: StructType =
-    struct.dataType.asInstanceOf[StructType]
-
-  private val pairs = fieldNames.zip(fieldExpressions)
-
   private type FieldName = String
 
   /**
-    * Recursively loops through addOrReplaceFields, adding or replacing fields by FieldName.
-    */
-  private def loop[V](existingFields: Seq[(String, V)], addOrReplaceFields: Seq[(String, V)]): Seq[(String, V)] = {
+   * Recursively loop through addOrReplaceFields, adding or replacing fields by FieldName.
+   */
+  @scala.annotation.tailrec
+  private def loop[V](existingFields: Seq[(String, V)],
+                      addOrReplaceFields: Seq[(String, V)]): Seq[(String, V)] = {
     if (addOrReplaceFields.nonEmpty) {
       val existingFieldNames = existingFields.map(_._1)
       val newField@(newFieldName, _) = addOrReplaceFields.head
@@ -172,6 +185,13 @@ case class AddFields(struct: Expression, fieldNames: Seq[String], fieldExpressio
 }
 
 object AddFields {
+  @deprecated("use AddFields(children: Seq[Expression]) constructor.", "0.2.4")
   def apply(struct: Expression, fieldName: String, fieldExpression: Expression): AddFields =
-    AddFields(struct, Seq(fieldName), Seq(fieldExpression))
+    AddFields(struct :: Literal(fieldName) :: fieldExpression :: Nil)
+
+  @deprecated("use AddFields(children: Seq[Expression]) constructor.", "0.2.4")
+  def apply(struct: Expression, fieldNames: Seq[String], fieldExpressions: Seq[Expression]): AddFields = {
+    val exprs = fieldNames.zip(fieldExpressions).flatMap { case (name, expr) => Seq(Literal(name), expr) }
+    AddFields(struct +: exprs)
+  }
 }