#58 Add support for df.agg()

### What changes were proposed in this pull request?
Add support for `df.Agg() and `df.AggWithMap()`.

### Why are the changes needed?
Compatibility

### Does this PR introduce _any_ user-facing change?
New functions

### How was this patch tested?
Added test

Closes #100 from grundprinzip/df_agg.

Authored-by: Martin Grund <martin.grund@databricks.com>
Signed-off-by: Martin Grund <martin.grund@databricks.com>

Author: grundprinzip

internal/tests/integration/functions_test.go

@@ -41,6 +41,25 @@ func TestIntegration_BuiltinFunctions(t *testing.T) {
 	assert.Equal(t, 10, len(res))
 }
 
+func TestAggregationFunctions_Agg(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select id, 1, 2, 3 from range(100)")
+	assert.NoError(t, err)
+
+	res, err := df.Agg(ctx, functions.Count(functions.Col("id")))
+	assert.NoError(t, err)
+	cnt, err := res.Count(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, int64(1), cnt)
+
+	res, err = df.AggWithMap(ctx, map[string]string{"id": "sum"})
+	assert.NoError(t, err)
+	rows, err := res.Collect(ctx)
+	assert.NoError(t, err)
+	assert.Len(t, rows, 1)
+	assert.Equal(t, int64(4950), rows[0].At(0))
+}
+
 func TestIntegration_ColumnGetItem(t *testing.T) {
 	ctx := context.Background()
 	spark, err := sql.NewSessionBuilder().Remote("sc://localhost").Build(ctx)

spark/sql/dataframe.go

@@ -41,6 +41,8 @@ type ResultCollector interface {
 type DataFrame interface {
 	// PlanId returns the plan id of the data frame.
 	PlanId() int64
+	Agg(ctx context.Context, exprs ...column.Convertible) (DataFrame, error)
+	AggWithMap(ctx context.Context, exprs map[string]string) (DataFrame, error)
 	// Alias creates a new DataFrame with the specified subquery alias
 	Alias(ctx context.Context, alias string) DataFrame
 	// Cache persists the DataFrame with the default storage level.
@@ -1542,6 +1544,22 @@ func (df *dataFrameImpl) FillNaWithValues(ctx context.Context,
 	return makeDataframeWithFillNaRelation(df, valueLiterals, columns), nil
 }
 
+func (df *dataFrameImpl) Agg(ctx context.Context, cols ...column.Convertible) (DataFrame, error) {
+	return df.GroupBy().Agg(ctx, cols...)
+}
+
+func (df *dataFrameImpl) AggWithMap(ctx context.Context, exprs map[string]string) (DataFrame, error) {
+	funs := make([]column.Convertible, 0)
+	for k, v := range exprs {
+		// Convert the column name to a column expression.
+		col := column.OfDF(df, k)
+		// Convert the value string to an unresolved function name.
+		fun := column.NewUnresolvedFunctionWithColumns(v, col)
+		funs = append(funs, fun)
+	}
+	return df.Agg(ctx, funs...)
+}
+
 func (df *dataFrameImpl) DropNa(ctx context.Context, subset ...string) (DataFrame, error) {
 	rel := &proto.Relation{
 		Common: &proto.RelationCommon{

spark/sql/group.go

@@ -37,7 +37,7 @@ type GroupedData struct {
 
 // Agg compute aggregates and returns the result as a DataFrame. The aggegrate expressions
 // are passed as column.Column arguments.
-func (gd *GroupedData) Agg(ctx context.Context, exprs ...column.Column) (DataFrame, error) {
+func (gd *GroupedData) Agg(ctx context.Context, exprs ...column.Convertible) (DataFrame, error) {
 	if len(exprs) == 0 {
 		return nil, sparkerrors.WithString(sparkerrors.InvalidInputError, "exprs should not be empty")
 	}
@@ -144,7 +144,7 @@ func (gd *GroupedData) numericAgg(ctx context.Context, name string, cols ...stri
 		aggCols = numericCols
 	}
 
-	finalColumns := make([]column.Column, len(aggCols))
+	finalColumns := make([]column.Convertible, len(aggCols))
 	for i, col := range aggCols {
 		finalColumns[i] = column.NewColumn(column.NewUnresolvedFunctionWithColumns(name, functions.Col(col)))
 	}