CODEWITHSUNDEEP
apache-sparkpysparkapache-spark-sqlrdd
bmc
bmc

Reputation: 857

pyspark - argmax of Row(double<array>) within 1 column

I have the following situation.

+--------------------+
|                   p|
+--------------------+
|[0.99998416412131...|
|[0.99998416412131...|
|[0.99998416412131...|
|[0.99998416412131...|
|[0.99998416412131...|
+--------------------+

This is a list of Row() objects.

[Row(p=[0.9999841641213133, 5.975696995141415e-06, 1.3699249952858219e-06, 1.4817184271708493e-06, 2.9022272149130313e-07, 1.4883436072406822e-06, 2.2234697862933896e-06, 3.006502154124559e-06]),
 Row(p=[0.9999841641213133, 5.975696995141415e-06, 1.3699249952858219e-06, 1.4817184271708493e-06, 2.9022272149130313e-07, 1.4883436072406822e-06, 2.2234697862933896e-06, 3.006502154124559e-06]),
 Row(p=[0.9999841641213133, 5.975696995141415e-06, 1.3699249952858219e-06, 1.4817184271708493e-06, 2.9022272149130313e-07, 1.4883436072406822e-06, 2.2234697862933896e-06, 3.006502154124559e-06]),
 Row(p=[0.9999841641213133, 5.975696995141415e-06, 1.3699249952858219e-06, 1.4817184271708493e-06, 2.9022272149130313e-07, 1.4883436072406822e-06, 2.2234697862933896e-06, 3.006502154124559e-06]),
 Row(p=[0.9999841641213133, 5.975696995141415e-06, 1.3699249952858219e-06, 1.4817184271708493e-06, 2.9022272149130313e-07, 1.4883436072406822e-06, 2.2234697862933896e-06, 3.006502154124559e-06])]

I am trying to filter down this column into a new column named "maxClass" that returned the np.argmax(row)[0] for all rows. Below is my best shot at it, but I simply cannot get the linguistics of using this package.

def f(row):
    return np.argmax(np.array(row.p))[0]
results=probs.rdd.map(lambda x:f(x))
results

Upvotes: 1

Views: 1726

Answers (1)

abiratsis
abiratsis

Reputation: 7336

For the sake of completeness and as pault suggested here is a solution without using UDF and numpy. Instead array_position and array_max is used:

import pyspark.sql.functions as f

df = spark.createDataFrame([
  ([0.9999841641213133, 5.975696995141415e-06, 1.3699249952858219e-06, 1.4817184271708493e-06, 2.9022272149130313e-07, 1.4883436072406822e-06, 2.2234697862933896e-06, 3.006502154124559e-06],),
  ([0.9999841641213134, 0.99999, 1.3699249952858219e-06, 1.4817184271708493e-06, 2.9022272149130313e-07, 1.4883436072406822e-06, 2.2234697862933896e-06, 3.006502154124559e-06],),
  ([0.9999841641213135, 5.975696995141415e-06, 1.3699249952858219e-06, 1.4817184271708493e-06, 2.9022272149130313e-07, 1.4883436072406822e-06, 2.2234697862933896e-06, 3.006502154124559e-06],)]) \
.toDF("p")

df.select(
  f.expr('array_position(cast(p as array<decimal(16, 16)>), cast(array_max(p) as decimal(16, 16))) - 1').alias("max_indx")
).show()

# +--------+
# |max_indx|
# +--------+
# |       0|
# |       1|
# |       0|
# +--------+

Upvotes: 2

Related Questions