Add schema_of_xml_df for Python future use (#438)

See #435 .
databricks · Feb 14, 2020 · 27dce95 · 27dce95
1 parent a71f735
commit 27dce95
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/src/main/scala/com/databricks/spark/xml/package.scala b/src/main/scala/com/databricks/spark/xml/package.scala
@@ -125,6 +125,17 @@ package object xml {
   def schema_of_xml(ds: Dataset[String], options: Map[String, String] = Map.empty): StructType =
     InferSchema.infer(ds.rdd, XmlOptions(options))
 
+  /**
+   * Infers the schema of XML documents as strings.
+   *
+   * @param df one-column DataFrame of XML strings
+   * @param options additional XML parsing options
+   * @return inferred schema for XML
+   */
+  @Experimental
+  def schema_of_xml_df(df: DataFrame, options: Map[String, String] = Map.empty): StructType =
+    schema_of_xml(df.as[String](Encoders.STRING), options);
+
   /**
    * Infers the schema of XML documents when inputs are arrays of strings, each an XML doc.
    *

diff --git a/src/test/scala/com/databricks/spark/xml/XmlSuite.scala b/src/test/scala/com/databricks/spark/xml/XmlSuite.scala
@@ -1146,7 +1146,7 @@ final class XmlSuite extends FunSuite with BeforeAndAfterAll {
        """.stripMargin
     import spark.implicits._
     val df = spark.createDataFrame(Seq((8, xmlData))).toDF("number", "payload")
-    val xmlSchema = schema_of_xml(df.select("payload").as[String])
+    val xmlSchema = schema_of_xml_df(df.select("payload"))
     val expectedSchema = df.schema.add("decoded", xmlSchema)
     val result = df.withColumn("decoded", from_xml(df.col("payload"), xmlSchema))
 
@@ -1179,7 +1179,7 @@ final class XmlSuite extends FunSuite with BeforeAndAfterAll {
        """.stripMargin
     import spark.implicits._
     val df = spark.createDataFrame(Seq((8, xmlData))).toDF("number", "payload")
-    val xmlSchema = schema_of_xml(df.select("payload").as[String])
+    val xmlSchema = schema_of_xml_df(df.select("payload"))
     val result = df.withColumn("decoded", from_xml(df.col("payload"), xmlSchema))
     assert(result.select("decoded._corrupt_record").head().getString(0).nonEmpty)
   }
@@ -1192,7 +1192,7 @@ final class XmlSuite extends FunSuite with BeforeAndAfterAll {
        """.stripMargin
     import spark.implicits._
     val df = spark.createDataFrame(Seq((8, xmlData))).toDF("number", "payload")
-    val xmlSchema = schema_of_xml(df.select("payload").as[String])
+    val xmlSchema = schema_of_xml_df(df.select("payload"))
     val result = from_xml_string(xmlData, xmlSchema)
 
     assert(result.getString(0) === "bar")
@@ -1214,7 +1214,7 @@ final class XmlSuite extends FunSuite with BeforeAndAfterAll {
        """.stripMargin
     import spark.implicits._
     val dfNoError = spark.createDataFrame(Seq((8, xmlDataNoError))).toDF("number", "payload")
-    val xmlSchema = schema_of_xml(dfNoError.select("payload").as[String])
+    val xmlSchema = schema_of_xml_df(dfNoError.select("payload"))
     val df = spark.createDataFrame(Seq((8, xmlData))).toDF("number", "payload")
     val result = df.withColumn("decoded", from_xml(df.col("payload"), xmlSchema))
     assert(result.select("decoded").head().get(0) === null)