如何从字符串值中获取数字代码？

Question

我使用pyspark.sql将csv文件导入spark并通过以下方式将其注册为临时表：

import pyspark
from pyspark.sql import SQLContext
sc = pyspark.SparkContext()
from pyspark.sql import HiveContext

sqlCtx= HiveContext(sc)
spark_df = sqlCtx.read.format('com.databricks.spark.csv').options(header='true', inferschema='true').load("./data/geo_file.csv")
spark_df.registerTempTable("geo_table")

在'geo_table'表中，有一个名为'geo_location'的列，其值如下：

  

US＆GT; TX＆GT; 618

     

US＆GT; NJ＆GT; 241

     

US＆GT; NJ

我的问题是，如何将这些文本值转换为数值？在sql或pyspark.sql中？

在熊猫，我会这样做

df["geo_location_categories"] = df["geo_location"].astype('category')
df["geo_location_codes"] = df["geo_location_categories"].cat.codes 

Answer 1

从我的角度来看，有几种方法可以解决您的问题。如果您只需要将“geo_location”列转换为数字列，则可以使用UDF。您可以通过以下方式定义UDF（抱歉，Scala中的代码）：

val toInt = udf[Int, String](str => {
   // convert your geo location string into integer using existing business logic
})

之后，您可以通过以下方式使用此UDF：

var df = spark_df.withColumn("geo_location_codes", toInt(spark_df("geo_location")))

另外我认为应该可以使用UserDefinedType作为列类型。但是，它取决于您使用的Spark版本。请查看此问题以获取更多信息：How to define schema for custom type in Spark SQL?

我希望你会发现这些信息很有用。