feat(dataframe): SeriesGroupBy has apply method, returns new Series

Author: dustmop

dataframe/index.go

@@ -78,6 +78,10 @@ func (i *Index) Type() string {
 func (i *Index) Attr(name string) (starlark.Value, error) {
 	switch name {
 	case "name":
+		if i == nil {
+			// TODO(dustmop): Add a test that covers this
+			return starlark.None, nil
+		}
 		return starlark.String(i.name), nil
 	case "str":
 		return &stringMethods{subject: i}, nil

dataframe/series.go

@@ -450,6 +450,16 @@ func builtinAttrNames(methods map[string]*starlark.Builtin) []string {
 	return names
 }
 
+// name returns of the name of the series
+func seriesAttrName(self *Series) (starlark.Value, error) {
+	return starlark.String(self.name), nil
+}
+
+// size returns the number of elements in the series
+func seriesAttrSize(self *Series) (starlark.Value, error) {
+	return starlark.MakeInt(self.Len()), nil
+}
+
 func seriesGet(_ *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var key starlark.Value
 	if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &key); err != nil {
@@ -746,6 +756,14 @@ func newSeriesFromFloats(vals []float64, index *Index, name string) *Series {
 	}
 }
 
+func newSeriesFromStrings(texts []string, index *Index, name string) *Series {
+	results := make([]interface{}, len(texts))
+	for i, txt := range texts {
+		results[i] = txt
+	}
+	return newSeriesFromObjects(results, index, name)
+}
+
 func newSeriesFromObjects(vals []interface{}, index *Index, name string) *Series {
 	return &Series{
 		dtype:   "object",

dataframe/series_all_methods.go

@@ -24,11 +24,11 @@ var seriesAttributes = map[string]seriesAttrImpl{
 	"is_monotonic_increasing": attrNoImplSeries("is_monotonic_increasing"),
 	"is_unique":               attrNoImplSeries("is_unique"),
 	"loc":                     attrNoImplSeries("loc"),
-	"name":                    attrNoImplSeries("name"),
+	"name":                    seriesAttrName,
 	"nbytes":                  attrNoImplSeries("nbytes"),
 	"ndim":                    attrNoImplSeries("ndim"),
 	"shape":                   attrNoImplSeries("shape"),
-	"size":                    attrNoImplSeries("size"),
+	"size":                    seriesAttrSize,
 	"values":                  attrNoImplSeries("values"),
 }
 

dataframe/series_groupby_result.go

@@ -12,6 +12,7 @@ import (
 type SeriesGroupByResult struct {
 	lhsLabel string
 	rhsLabel string
+	// TODO(dustmop): convert to map[string]Series
 	grouping map[string][]string
 }
 
@@ -24,6 +25,7 @@ var (
 var seriesGroupByResultMethods = map[string]*starlark.Builtin{
 	"count": starlark.NewBuiltin("count", seriesGroupByResultCount),
 	"sum":   starlark.NewBuiltin("sum", seriesGroupByResultSum),
+	"apply": starlark.NewBuiltin("apply", seriesGroupByResultApply),
 }
 
 // Freeze has no effect on the immutable SeriesGroupByResult
@@ -113,6 +115,55 @@ func seriesGroupByResultCount(_ *starlark.Thread, b *starlark.Builtin, args star
 	return newSeriesFromInts(vals, index, self.rhsLabel), nil
 }
 
+// apply method returns a Series that is built by calling the given
+// function, and passing each grouped series as an argument to it
+func seriesGroupByResultApply(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var (
+		funcVal starlark.Value
+		self    = b.Receiver().(*SeriesGroupByResult)
+	)
+
+	if err := starlark.UnpackArgs("apply", args, kwargs,
+		"function", &funcVal,
+	); err != nil {
+		return nil, err
+	}
+
+	funcObj, ok := funcVal.(*starlark.Function)
+	if !ok {
+		return nil, fmt.Errorf("first argument must be a function")
+	}
+
+	sortedKeys := getSortedKeys(self.grouping)
+	builder := newTypedSliceBuilder(len(sortedKeys))
+	indexNames := make([]string, len(sortedKeys))
+
+	for i, groupName := range sortedKeys {
+		values := self.grouping[groupName]
+		// TODO(dustmop): Pass actual index here
+		index := NewIndex(nil, groupName)
+		series := newSeriesFromStrings(values, index, groupName)
+		arguments := starlark.Tuple{series}
+		// Call function, passing the series to it
+		res, err := starlark.Call(thread, funcObj, arguments, nil)
+		if err != nil {
+			return nil, err
+		}
+		obj, ok := toScalarMaybe(res)
+		if !ok {
+			return nil, fmt.Errorf("could not convert: %v", res)
+		}
+		// Accumulate the new series, and build the new index
+		builder.push(obj)
+		indexNames[i] = groupName
+	}
+	if err := builder.error(); err != nil {
+		return nil, err
+	}
+	s := builder.toSeries(NewIndex(indexNames, self.lhsLabel), self.rhsLabel)
+	return &s, nil
+}
+
 func getSortedKeys(m map[string][]string) []string {
 	keys := make([]string, 0, len(m))
 	for k := range m {

dataframe/testdata/dataframe_groupby.expect.txt

@@ -11,3 +11,12 @@ cat    3
 dog    2
 Name: breed, dtype: int64
 Index(['cat', 'dog'], dtype='object', name='species')
+
+species
+cat    tabby cat, black cat, calico cat
+dog               doberman dog, pug dog
+Name: breed, dtype: object
+
+dataframe.Series
+breed
+species

dataframe/testdata/dataframe_groupby.star

@@ -1,6 +1,11 @@
 load("dataframe.star", "dataframe")
 
 
+def example_animals(series):
+  examples = ['{} {}'.format(series[i], series.name) for i in range(series.size)]
+  return ', '.join(examples)
+
+
 def f():
   df = dataframe.DataFrame({"IDs": ["cat", "dog", "eel", "dog", "cat", "frog", "cat", "eel"],
                             "count": [1, 2, 3, 4, 5, 6, 7, 8]})
@@ -19,6 +24,14 @@ def f():
   num_breeds = df.groupby(['species'])['breed'].count()
   print(num_breeds)
   print(num_breeds.index)
+  print('')
+
+  list_of_examples = df.groupby(['species'])['breed'].apply(example_animals)
+  print(list_of_examples)
+  print('')
+  print(type(list_of_examples))
+  print(list_of_examples.name)
+  print(list_of_examples.index.name)
 
 
 f()