fix issue of column name contains space

bomeng · bomeng · commit 88935c5b0eb3 · 2016-04-07T20:25:33.000-07:00
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
@@ -97,7 +97,7 @@ object JdbcUtils extends Logging {
    * Returns a PreparedStatement that inserts a row into table via conn.
    */
   def insertStatement(conn: Connection, table: String, rddSchema: StructType): PreparedStatement = {
-    val columns = rddSchema.fields.map(_.name).mkString(",")
+    val columns = rddSchema.fields.map(f => escapeColumnName(f.name)).mkString(",")
     val placeholders = rddSchema.fields.map(_ => "?").mkString(",")
     val sql = s"INSERT INTO $table ($columns) VALUES ($placeholders)"
     conn.prepareStatement(sql)
@@ -245,14 +245,23 @@ object JdbcUtils extends Logging {
     Array[Byte]().iterator
   }
 
+  /**
+   * The utility to add backtick if the column name has space
+   * @param columnName the input column name
+   * @return the escaped column name, add backtick if name contains space
+   */
+  private def escapeColumnName(columnName: String): String = {
+    if (columnName.contains(" ")) s"`$columnName`" else columnName
+  }
+
   /**
    * Compute the schema string for this RDD.
    */
   def schemaString(df: DataFrame, url: String): String = {
     val sb = new StringBuilder()
     val dialect = JdbcDialects.get(url)
     df.schema.fields foreach { field => {
-      val name = field.name
+      val name = escapeColumnName(field.name)
       val typ: String = getJdbcType(field.dataType, dialect).databaseTypeDefinition
       val nullable = if (field.nullable) "" else "NOT NULL"
       sb.append(s", $name $typ $nullable")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCWriteSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCWriteSuite.scala
@@ -52,7 +52,8 @@ class JDBCWriteSuite extends SharedSQLContext with BeforeAndAfter {
     conn1.prepareStatement("insert into test.people values ('mary', 2)").executeUpdate()
     conn1.prepareStatement("drop table if exists test.people1").executeUpdate()
     conn1.prepareStatement(
-      "create table test.people1 (name TEXT(32) NOT NULL, theid INTEGER NOT NULL)").executeUpdate()
+      "create table test.people1 (name TEXT(32) NOT NULL, `the id` INTEGER NOT NULL)")
+      .executeUpdate()
     conn1.commit()
 
     sql(
@@ -151,4 +152,10 @@ class JDBCWriteSuite extends SharedSQLContext with BeforeAndAfter {
     assert(2 === sqlContext.read.jdbc(url1, "TEST.PEOPLE1", properties).count)
     assert(2 === sqlContext.read.jdbc(url1, "TEST.PEOPLE1", properties).collect()(0).length)
   }
+
+  test("SPARK-14460: Insert into table with column containing space") {
+    val df = sqlContext.createDataFrame(sparkContext.parallelize(arr2x2), schema2)
+    df.write.insertInto("PEOPLE1")
+    assert(2 === sqlContext.read.jdbc(url1, "TEST.PEOPLE1", properties).count)
+  }
 }
