Apache Spark — Assign the result of UDF to multiple dataframe columns

by

It is not possible to create multiple top level columns from a single UDF call but you can create a new struct. It requires an UDF with specified returnType:

from pyspark.sql.functions import udf
from pyspark.sql.types import StructType, StructField, FloatType

schema = StructType([
    StructField("foo", FloatType(), False),
    StructField("bar", FloatType(), False)
])

def udf_test(n):
    return (n / 2, n % 2) if n and n != 0.0 else (float('nan'), float('nan'))

test_udf = udf(udf_test, schema)
df = sc.parallelize([(1, 2.0), (2, 3.0)]).toDF(["x", "y"])

foobars = df.select(test_udf("y").alias("foobar"))
foobars.printSchema()
## root
##  |-- foobar: struct (nullable = true)
##  |    |-- foo: float (nullable = false)
##  |    |-- bar: float (nullable = false)

You further flatten the schema with simple select:

foobars.select("foobar.foo", "foobar.bar").show()
## +---+---+
## |foo|bar|
## +---+---+
## |1.0|0.0|
## |1.5|1.0|
## +---+---+

See also Derive multiple columns from a single column in a Spark DataFrame

More Related Contents:

  1. Applying UDFs on GroupedData in PySpark (with functioning python example)
  2. Passing a data frame column and external list to udf under withColumn
  3. Using UDF ignores condition in when
  4. How to pass a constant value to Python UDF?
  5. Spark Equivalent of IF Then ELSE
  6. How to add a constant column in a Spark DataFrame?
  7. Load CSV file with Spark
  8. How do I add a new column to a Spark DataFrame (using PySpark)?
  9. How to perform union on two DataFrames with different amounts of columns in spark?
  10. Does spark predicate pushdown work with JDBC?
  11. ‘PipelinedRDD’ object has no attribute ‘toDF’ in PySpark
  12. Count number of non-NaN entries in each column of Spark dataframe with Pyspark
  13. Spark DataFrame: Computing row-wise mean (or any aggregate operation)
  14. Updating a dataframe column in spark
  15. PySpark converting a column of type ‘map’ to multiple columns in a dataframe
  16. Best way to get the max value in a Spark dataframe column
  17. Pyspark: Replacing value in a column by searching a dictionary
  18. Apache Spark Python Cosine Similarity over DataFrames
  19. Pivot String column on Pyspark Dataframe
  20. GroupBy column and filter rows with maximum value in Pyspark
  21. Reshaping/Pivoting data in Spark RDD and/or Spark DataFrames
  22. Cannot find col function in pyspark
  23. How to zip two array columns in Spark SQL
  24. When to cache a DataFrame?
  25. Multiple condition filter on dataframe
  26. How to return a “Tuple type” in a UDF in PySpark?
  27. How to add suffix and prefix to all columns in python/pyspark dataframe
  28. How to connect HBase and Spark using Python?
  29. Pyspark changing type of column from date to string
  30. PySpark: compute row maximum of the subset of columns and add to an exisiting dataframe

Leave a Comment