dataframe - 编写 UDF 以在 Java 中的 Map 中查找，给出不支持的文字类型类 java.util.HashMap

Question

我正在使用带有 spark v2.4.1 的 java8

我正在尝试使用 UDF 使用 Map 进行查找，如下所示

数据:

+-----+-----+-----+
|code1|code2|code3|
+-----+-----+-----+
|1    |7    |  5  |
|2    |7    |  4  |
|3    |7    |  3  |
|4    |7    |  2  |
|5    |7    |  1  |
+-----+-----+-----+

预期数据:

+-----+-----+-----+
|code1|code2|code3|
+-----+-----+-----+
|1    |7    |51   |
|2    |7    |41   |
|3    |7    |31   |
|4    |7    |21   |
|5    |7    |11   |
+-----+-----+-----+

Map<Integer,Integer> map= new HashMap<>();
      map.put(1,11);
      map.put(2,21);
      map.put(3,31);
      map.put(4,41);
      map.put(5,51);



public static UDF2 userDefinedFunction= new UDF2<java.util.Map<Integer, Integer> ,Integer, Integer>() 
 {
        private static final long serialVersionUID = 1L;
        
        @Override
        public Integer call(java.util.Map<Integer, Integer> map, Integer score) throws Exception {
            return map.get(score);
        }
    };


  Dataset<Row> resultDs= dataDs.withColumn("code3",
           functions.callUDF("userDefinedFunction",col("code3"),lit(map) ) )

错误:

java.lang.RuntimeException: 不支持的文字类型类 java.util.HashMap

这里有什么问题？如何使用 JavaAPI 在 UDF 中传递/处理 HashMap 参数

数据:

    List<String[]> stringAsList = new ArrayList<>();
    stringAsList.add(new String[] { "1","7","5" });
    stringAsList.add(new String[] { "2","7","4" });
    stringAsList.add(new String[] { "3","7","3" });
    stringAsList.add(new String[] { "4","7","2" });
    stringAsList.add(new String[] { "5","7","1" });
    
    JavaSparkContext sparkContext = new JavaSparkContext(sparkSession.sparkContext());
    JavaRDD<Row> rowRDD = sparkContext.parallelize(stringAsList).map((String[] row) -> RowFactory.create(row));

   
    StructType schema = DataTypes
            .createStructType(new StructField[] {
                    DataTypes.createStructField("code1", DataTypes.StringType, false),
                    DataTypes.createStructField("code2", DataTypes.StringType, false),
                    DataTypes.createStructField("code3", DataTypes.StringType, false)
                  
            });

    Dataset<Row> dataDf= sparkSession.sqlContext().createDataFrame(rowRDD, schema).toDF();

    
    Dataset<Row> dataDs = dataDf
                    .withColumn("code1", col("code1").cast(DataTypes.IntegerType))
                    .withColumn("code2", col("code2").cast(DataTypes.IntegerType))
                    .withColumn("code3", col("code3").cast(DataTypes.IntegerType))
                      ;

Answer 1

您可以使用partial 将查找映射或数组等传递给udf。看看这个例子。

from functools import partial
from pyspark.sql.functions import udf

fruit_dict = {"O": "Orange", "A": "Apple", "G": "Grape"}
df = spark.createDataFrame([("A", 20), ("G", 30), ("O", 10)], ["Label", "Count"])
def decipher_fruit(label, fruit_map):
  label_names = list(fruit_map.keys())
  if label in label_names:
    return fruit_map[label]
  return None

decipher_fruit_udf = udf(partial(decipher_fruit, fruit_map = fruit_dict), StringType())
df2 = df.withColumn("fruit_name", decipher_fruit_udf("label"))
display(df2)