performing a number of operations with a single function

Question

I have a dataframe and I need to agregate it. I do it this way.

df
      .groupBy($"key1", $"key2")
      .agg(
        sum($"total_in").as("total_in")
        , sum($"val1_in").as("val1_in")
        , sum($"val2_in").as("val2_in")
        , sum($"val3_in").as("val3_in")
        , sum($"val4_in").as("val4_in")
        , sum($"total_out").as("total_out")
        , sum($"val1_out").as("val1_out")
        , sum($"val2_out").as("val2_out")
        , sum($"val3_out").as("val3_out")
        , sum($"val4_out").as("val4_out"))

It doesn't look all that bad, but if the number of operation in agg becomes larger or a case when or some other logic gets added to every sum, this code will look really cumbersome.

I know it can be done. Maybe passing the columns to a function and iterating other them. Something like

.agg(
  map(x => func())
)

but this is where I'm stuck. How do I pass only the columns I need? Is it possible to alias them in the same place?

score 0 · Answer 1 · answered Sep 11 '19 at 17:54

Use vararg syntax:

In python:

cols = ['total_in', 'val1_in', 'val2_in', ...]
.agg(*(sum(col).alias(col) for col in cols))

In Scala:

cols = Seq("total_in", "val1_in", "val2_in", ...)
.agg(
  sum(col(cols.head)).as(cols.head), 
  cols.tail.map(x => sum(col(x)).as(x)): _*
)

score 0 · Accepted Answer · answered Sep 11 '19 at 18:27

0

I would do it like this :

val aggExpresion =
    df.columns
      .filterNot(Seq("key1","key2").contains(_))
      .map(s => sum(col(s)).as(s))

df.agg(aggExpresion.head,aggExpresion.tail:_*)

answered Sep 11 '19 at 18:27

Raphael Roth

23,651
9
69
117

performing a number of operations with a single function

2 Answers2