diff --git a/README.md b/README.md index 98bf9d48..883cbd0a 100644 --- a/README.md +++ b/README.md @@ -380,6 +380,32 @@ df.write .save() ``` + +### Setting a custom column type + +If you need to manually set a column type, you can use the `redshift_type` column metadata. For example, if you desire to override +the `Spark SQL Schema -> Redshift SQL` type matcher to assign a user-defined column type, you can do the following: + +```scala +import org.apache.spark.sql.types.MetadataBuilder + +// Specify the custom width of each column +val columnTypeMap = Map( + "language_code" -> "CHAR(2)", + "country_code" -> "CHAR(2)", + "url" -> "BPCHAR(111)" +) + +var df = ... // the dataframe you'll want to write to Redshift + +// Apply each column metadata customization +columnTypeMap.foreach { case (colName, colType) => + val metadata = new MetadataBuilder().putString("redshift_type", colType).build() + df = df.withColumn(colName, df(colName).as(colName, metadata)) +} +``` + + ## Migration Guide diff --git a/src/main/scala/com/databricks/spark/redshift/RedshiftJDBCWrapper.scala b/src/main/scala/com/databricks/spark/redshift/RedshiftJDBCWrapper.scala index 6b3d2572..171a5bcb 100644 --- a/src/main/scala/com/databricks/spark/redshift/RedshiftJDBCWrapper.scala +++ b/src/main/scala/com/databricks/spark/redshift/RedshiftJDBCWrapper.scala @@ -232,26 +232,32 @@ private[redshift] class JDBCWrapper { val sb = new StringBuilder() schema.fields.foreach { field => { val name = field.name - val typ: String = field.dataType match { - case IntegerType => "INTEGER" - case LongType => "BIGINT" - case DoubleType => "DOUBLE PRECISION" - case FloatType => "REAL" - case ShortType => "INTEGER" - case ByteType => "SMALLINT" // Redshift does not support the BYTE type. - case BooleanType => "BOOLEAN" - case StringType => - if (field.metadata.contains("maxlength")) { - s"VARCHAR(${field.metadata.getLong("maxlength")})" - } else { - "TEXT" - } - case BinaryType => "BLOB" - case TimestampType => "TIMESTAMP" - case DateType => "DATE" - case t: DecimalType => s"DECIMAL(${t.precision},${t.scale})" - case _ => throw new IllegalArgumentException(s"Don't know how to save $field to JDBC") + val typ: String = if (field.metadata.contains("redshift_type")) { + field.metadata.getString("redshift_type") + } + else { + field.dataType match { + case IntegerType => "INTEGER" + case LongType => "BIGINT" + case DoubleType => "DOUBLE PRECISION" + case FloatType => "REAL" + case ShortType => "INTEGER" + case ByteType => "SMALLINT" // Redshift does not support the BYTE type. + case BooleanType => "BOOLEAN" + case StringType => + if (field.metadata.contains("maxlength")) { + s"VARCHAR(${field.metadata.getLong("maxlength")})" + } else { + "TEXT" + } + case BinaryType => "BLOB" + case TimestampType => "TIMESTAMP" + case DateType => "DATE" + case t: DecimalType => s"DECIMAL(${t.precision},${t.scale})" + case _ => throw new IllegalArgumentException(s"Don't know how to save $field to JDBC") + } } + val nullable = if (field.nullable) "" else "NOT NULL" sb.append(s""", "${name.replace("\"", "\\\"")}" $typ $nullable""".trim) }} diff --git a/src/test/scala/com/databricks/spark/redshift/RedshiftSourceSuite.scala b/src/test/scala/com/databricks/spark/redshift/RedshiftSourceSuite.scala index 4d623584..64b872c0 100644 --- a/src/test/scala/com/databricks/spark/redshift/RedshiftSourceSuite.scala +++ b/src/test/scala/com/databricks/spark/redshift/RedshiftSourceSuite.scala @@ -387,6 +387,25 @@ class RedshiftSourceSuite assert(createTableCommand === expectedCreateTableCommand) } + test("configuring redshift_type on columns") { + val bpcharMetadata = new MetadataBuilder().putString("redshift_type", "BPCHAR(2)").build() + val nvarcharMetadata = new MetadataBuilder().putString("redshift_type", "NVARCHAR(123)").build() + + val schema = StructType( + StructField("bpchar_str", StringType, metadata = bpcharMetadata) :: + StructField("bpchar_str", StringType, metadata = nvarcharMetadata) :: + StructField("default_str", StringType) :: + Nil) + + val df = testSqlContext.createDataFrame(sc.emptyRDD[Row], schema) + val createTableCommand = + DefaultRedshiftWriter.createTableSql(df, MergedParameters.apply(defaultParams)).trim + val expectedCreateTableCommand = + """CREATE TABLE IF NOT EXISTS "PUBLIC"."test_table" ("bpchar_str" BPCHAR(2),""" + + """ "bpchar_str" NVARCHAR(123), "default_str" TEXT)""" + assert(createTableCommand === expectedCreateTableCommand) + } + test("Respect SaveMode.ErrorIfExists when table exists") { val mockRedshift = new MockRedshift( defaultParams("url"),