[SPARK-33934][SQL] Support user defined script command wrapper like hive

AngersZhuuuu · AngersZhuuuu · commit c39e3cfb8ee2 · 2020-12-30T17:45:44.000+08:00
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -2918,6 +2918,14 @@ object SQLConf {
       .checkValues(LegacyBehaviorPolicy.values.map(_.toString))
       .createWithDefault(LegacyBehaviorPolicy.EXCEPTION.toString)
 
+  val SCRIPT_TRANSFORMATION_COMMAND_WRAPPER =
+    buildConf("spark.sql.scriptTransformation.commandWrapper")
+      .internal()
+      .doc("Command wrapper for executor to execute transformation script.")
+      .version("3.2.0")
+      .stringConf
+      .createWithDefault("/bin/bash -c")
+
   val SCRIPT_TRANSFORMATION_EXIT_TIMEOUT =
     buildConf("spark.sql.scriptTransformation.exitTimeoutInSeconds")
       .internal()
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala
@@ -17,16 +17,17 @@
 
 package org.apache.spark.sql.execution
 
-import java.io.{BufferedReader, InputStream, InputStreamReader, OutputStream}
+import java.io._
 import java.nio.charset.StandardCharsets
+import java.util.ArrayList
 import java.util.concurrent.TimeUnit
 
 import scala.collection.JavaConverters._
 import scala.util.control.NonFatal
 
 import org.apache.hadoop.conf.Configuration
 
-import org.apache.spark.{SparkException, TaskContext}
+import org.apache.spark.{SparkException, SparkFiles, TaskContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
@@ -46,6 +47,8 @@ trait BaseScriptTransformationExec extends UnaryExecNode {
   def child: SparkPlan
   def ioschema: ScriptTransformationIOSchema
 
+  type ProcParameters = (OutputStream, Process, InputStream, CircularBuffer)
+
   protected lazy val inputExpressionsWithoutSerde: Seq[Expression] = {
     input.map(Cast(_, StringType).withTimeZone(conf.sessionLocalTimeZone))
   }
@@ -55,8 +58,11 @@ trait BaseScriptTransformationExec extends UnaryExecNode {
   override def outputPartitioning: Partitioning = child.outputPartitioning
 
   override def doExecute(): RDD[InternalRow] = {
+    val hadoopConf = sqlContext.sessionState.newHadoopConf()
+    hadoopConf.set(SQLConf.SCRIPT_TRANSFORMATION_COMMAND_WRAPPER.key,
+      conf.getConf(SQLConf.SCRIPT_TRANSFORMATION_COMMAND_WRAPPER))
     val broadcastedHadoopConf =
-      new SerializableConfiguration(sqlContext.sessionState.newHadoopConf())
+      new SerializableConfiguration(hadoopConf)
 
     child.execute().mapPartitions { iter =>
       if (iter.hasNext) {
@@ -69,9 +75,11 @@ trait BaseScriptTransformationExec extends UnaryExecNode {
     }
   }
 
-  protected def initProc: (OutputStream, Process, InputStream, CircularBuffer) = {
-    val cmd = List("/bin/bash", "-c", script)
+  protected def initProc(hadoopConf: Configuration): ProcParameters = {
+    val wrapper = splitArgs(hadoopConf.get(SQLConf.SCRIPT_TRANSFORMATION_COMMAND_WRAPPER.key))
+    val cmd = wrapper.toList ++ List(script)
     val builder = new ProcessBuilder(cmd.asJava)
+      .directory(new File(SparkFiles.getRootDirectory()))
 
     val proc = builder.start()
     val inputStream = proc.getInputStream
@@ -181,6 +189,55 @@ trait BaseScriptTransformationExec extends UnaryExecNode {
     }
   }
 
+  def splitArgs(args: String): Array[String] = {
+    val OUTSIDE = 1
+    val SINGLEQ = 2
+    val DOUBLEQ = 3
+    val argList = new ArrayList[String]
+    val ch = args.toCharArray
+    val clen = ch.length
+    var state = OUTSIDE
+    var argstart = 0
+    var c = 0
+    while (c <= clen) {
+      val last = c == clen
+      var lastState = state
+      var endToken = false
+      if (!last) {
+        if (ch(c) == '\'') {
+          if (state == OUTSIDE) {
+            state = SINGLEQ
+          } else if (state == SINGLEQ) {
+            state = OUTSIDE
+          }
+          endToken = state != lastState
+        } else if (ch(c) == '"') {
+          if (state == OUTSIDE) {
+            state = DOUBLEQ
+          } else if (state == DOUBLEQ) {
+            state = OUTSIDE
+          }
+          endToken = state != lastState
+        } else if (ch(c) == ' ') {
+          if (state == OUTSIDE) {
+            endToken = true
+          }
+        }
+      }
+      if (last || endToken) {
+        if (c == argstart) {
+          // unquoted space
+        } else {
+          argList.add(args.substring(argstart, c))
+        }
+        argstart = c + 1
+        lastState = state
+      }
+      c += 1
+    }
+    argList.toArray(new Array[String](0))
+  }
+
   private lazy val outputFieldWriters: Seq[String => Any] = output.map { attr =>
     val converter = CatalystTypeConverters.createToCatalystConverter(attr.dataType)
     attr.dataType match {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkScriptTransformationExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkScriptTransformationExec.scala
@@ -48,7 +48,7 @@ case class SparkScriptTransformationExec(
       inputIterator: Iterator[InternalRow],
       hadoopConf: Configuration): Iterator[InternalRow] = {
 
-    val (outputStream, proc, inputStream, stderrBuffer) = initProc
+    val (outputStream, proc, inputStream, stderrBuffer) = initProc(hadoopConf)
 
     val outputProjection = new InterpretedProjection(inputExpressionsWithoutSerde, child.output)
 
diff --git a/sql/core/src/test/resources/test_script.py b/sql/core/src/test/resources/test_script.py
@@ -1,3 +1,5 @@
+#!/usr/bin/python
+
 # Licensed to the Apache Software Foundation (ASF) under one or more
 # contributor license agreements.  See the NOTICE file distributed with
 # this work for additional information regarding copyright ownership.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/BaseScriptTransformationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/BaseScriptTransformationSuite.scala
@@ -470,6 +470,180 @@ abstract class BaseScriptTransformationSuite extends SparkPlanTest with SQLTestU
           Row("3\u00014\u00015") :: Nil)
     }
   }
+
+  test("SPARK-33934: Check default execute command wrapper") {
+    assume(TestUtils.testCommandAvailable("python"))
+    val scriptFilePath = copyAndGetResourceFile("test_script.py", ".py").getAbsoluteFile
+    withTempView("v") {
+      val df = Seq(
+        (1, "1", 1.0, BigDecimal(1.0), new Timestamp(1)),
+        (2, "2", 2.0, BigDecimal(2.0), new Timestamp(2)),
+        (3, "3", 3.0, BigDecimal(3.0), new Timestamp(3))
+      ).toDF("a", "b", "c", "d", "e") // Note column d's data type is Decimal(38, 18)
+      df.createTempView("v")
+
+      // test '/bin/bash -c python /path/to/script.py'
+      checkAnswer(
+        sql(
+          s"""
+             |SELECT
+             |TRANSFORM(a, b, c, d, e)
+             |  ROW FORMAT DELIMITED
+             |  FIELDS TERMINATED BY '\t'
+             |  USING 'python $scriptFilePath' AS (a, b, c, d, e)
+             |  ROW FORMAT DELIMITED
+             |  FIELDS TERMINATED BY '\t'
+             |FROM v
+        """.stripMargin), identity, df.select(
+          'a.cast("string"),
+          'b.cast("string"),
+          'c.cast("string"),
+          'd.cast("string"),
+          'e.cast("string")).collect())
+
+      // test '/bin/bash -c /path/to/script.py' with script not executable
+      val e1 = intercept[TestFailedException] {
+        checkAnswer(
+          sql(
+            s"""
+               |SELECT
+               |TRANSFORM(a, b, c, d, e)
+               |  ROW FORMAT DELIMITED
+               |  FIELDS TERMINATED BY '\t'
+               |  USING '$scriptFilePath' AS (a, b, c, d, e)
+               |  ROW FORMAT DELIMITED
+               |  FIELDS TERMINATED BY '\t'
+               |FROM v
+        """.stripMargin), identity, df.select(
+            'a.cast("string"),
+            'b.cast("string"),
+            'c.cast("string"),
+            'd.cast("string"),
+            'e.cast("string")).collect())
+      }.getMessage
+      assert(e1.contains("Permission denied"))
+
+      // test '/bin/bash -c /path/to/script.py' with script executable
+      scriptFilePath.setExecutable(true)
+      checkAnswer(
+        sql(
+          s"""
+             |SELECT
+             |TRANSFORM(a, b, c, d, e)
+             |  ROW FORMAT DELIMITED
+             |  FIELDS TERMINATED BY '\t'
+             |  USING '$scriptFilePath' AS (a, b, c, d, e)
+             |  ROW FORMAT DELIMITED
+             |  FIELDS TERMINATED BY '\t'
+             |FROM v
+        """.stripMargin), identity, df.select(
+          'a.cast("string"),
+          'b.cast("string"),
+          'c.cast("string"),
+          'd.cast("string"),
+          'e.cast("string")).collect())
+
+      scriptFilePath.setExecutable(false)
+      sql(s"ADD FILE ${scriptFilePath.getAbsolutePath}")
+
+      // test '/bin/bash -c script.py'
+      val e2 = intercept[TestFailedException] {
+        checkAnswer(
+          sql(
+            s"""
+               |SELECT TRANSFORM(a, b, c, d, e)
+               |  ROW FORMAT DELIMITED
+               |  FIELDS TERMINATED BY '\t'
+               |  USING '${scriptFilePath.getName}' AS (a, b, c, d, e)
+               |  ROW FORMAT DELIMITED
+               |  FIELDS TERMINATED BY '\t'
+               |FROM v
+        """.stripMargin), identity, df.select(
+            'a.cast("string"),
+            'b.cast("string"),
+            'c.cast("string"),
+            'd.cast("string"),
+            'e.cast("string")).collect())
+      }.getMessage()
+      assert(e2.contains("command not found"))
+    }
+  }
+
+  test("SPARK-33934: Check execute command wrapper is empty") {
+    assume(TestUtils.testCommandAvailable("python"))
+    val scriptFilePath = copyAndGetResourceFile(
+      "test_script.py", "_test_empty.py").getAbsoluteFile
+    withTempView("v") {
+      withSQLConf(SQLConf.SCRIPT_TRANSFORMATION_COMMAND_WRAPPER.key -> "") {
+        val df = Seq(
+          (1, "1", 1.0, BigDecimal(1.0), new Timestamp(1)),
+          (2, "2", 2.0, BigDecimal(2.0), new Timestamp(2)),
+          (3, "3", 3.0, BigDecimal(3.0), new Timestamp(3))
+        ).toDF("a", "b", "c", "d", "e") // Note column d's data type is Decimal(38, 18)
+        df.createTempView("v")
+
+        scriptFilePath.setExecutable(true)
+        sql(s"ADD FILE ${scriptFilePath.getAbsolutePath}")
+        sql(
+          s"""
+             |SELECT TRANSFORM(a)
+             |  ROW FORMAT DELIMITED
+             |  FIELDS TERMINATED BY '\t'
+             |  USING 'pwd' AS (a)
+             |  ROW FORMAT DELIMITED
+             |  FIELDS TERMINATED BY '&'
+             |FROM (SELECT 1 AS a) TEMP
+        """.stripMargin).show(false)
+
+        sql(
+          s"""
+             |SELECT TRANSFORM(a)
+             |  ROW FORMAT DELIMITED
+             |  FIELDS TERMINATED BY '\t'
+             |  USING 'ls' AS (a)
+             |  ROW FORMAT DELIMITED
+             |  FIELDS TERMINATED BY '&'
+             |FROM (SELECT 1 AS a) TEMP
+        """.stripMargin).show(false)
+
+        // test 'python script.py'
+        checkAnswer(
+          sql(
+            s"""
+               |SELECT TRANSFORM(a, b, c, d, e)
+               |  ROW FORMAT DELIMITED
+               |  FIELDS TERMINATED BY '\t'
+               |  USING 'python ${scriptFilePath.getName}' AS (a, b, c, d, e)
+               |  ROW FORMAT DELIMITED
+               |  FIELDS TERMINATED BY '\t'
+               |FROM v
+        """.stripMargin), identity, df.select(
+            'a.cast("string"),
+            'b.cast("string"),
+            'c.cast("string"),
+            'd.cast("string"),
+            'e.cast("string")).collect())
+
+        // test 'script.py'
+        checkAnswer(
+          sql(
+            s"""
+               |SELECT TRANSFORM(a, b, c, d, e)
+               |  ROW FORMAT DELIMITED
+               |  FIELDS TERMINATED BY '\t'
+               |  USING '${scriptFilePath.getName}' AS (a, b, c, d, e)
+               |  ROW FORMAT DELIMITED
+               |  FIELDS TERMINATED BY '\t'
+               |FROM v
+        """.stripMargin), identity, df.select(
+            'a.cast("string"),
+            'b.cast("string"),
+            'c.cast("string"),
+            'd.cast("string"),
+            'e.cast("string")).collect())
+      }
+    }
+  }
 }
 
 case class ExceptionInjectingOperator(child: SparkPlan) extends UnaryExecNode {
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/execution/HiveScriptTransformationExec.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/execution/HiveScriptTransformationExec.scala
@@ -140,7 +140,7 @@ case class HiveScriptTransformationExec(
       inputIterator: Iterator[InternalRow],
       hadoopConf: Configuration): Iterator[InternalRow] = {
 
-    val (outputStream, proc, inputStream, stderrBuffer) = initProc
+    val (outputStream, proc, inputStream, stderrBuffer) = initProc(hadoopConf)
 
     val (inputSerde, inputSoi) = initInputSerDe(ioschema, input).getOrElse((null, null))
 

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+#!/usr/bin/python`
	`2`	`+`
`1`	`3`	`# Licensed to the Apache Software Foundation (ASF) under one or more`
`2`	`4`	`# contributor license agreements. See the NOTICE file distributed with`
`3`	`5`	`# this work for additional information regarding copyright ownership.`