Go gota库如何为DataFrame添加自增行索引列

问题背景 #

通过如下示例代码，可以实现基于行值派生新增列的功能，代码运行效果正常：

package main

import (
    "fmt"
    "strings"

    "github.com/go-gota/gota/dataframe"
    "github.com/go-gota/gota/series"
)

func main() {
    csvStr := `accountId,deposit,Withdrawals
anil0001,50,10
vikas0002,10,10
ravi0003,20,10
user1111,NaN,20`

    df := dataframe.ReadCSV(strings.NewReader(csvStr))

    // 行内元素通过列索引访问
    indexDeposit := 1
    indexWithdrawals := 2

    // Rapply 按行遍历数据，可通过 s.Elem(index) 或 s.Val(index) 访问行内元素
    // 按列遍历可使用 Capply
    s := df.Rapply(func(s series.Series) series.Series {
        deposit, err := s.Elem(indexDeposit).Int()
        if err != nil {
            return series.Ints("NAN")
        }
        withdrawal, err := s.Elem(indexWithdrawals).Int()
        if err != nil {
            return series.Ints("NAN")
        }
        return series.Ints(deposit - withdrawal)
    })

    // 新生成的序列通过Mutate方法追加到原DataFrame
    df = df.Mutate(s.Col("X0")).
        Rename("deposit_Withdrawals_diff", "X0")

    fmt.Println(df)
}

当前需求为给每一行添加从1开始的行号索引列，后续需要将原DataFrame与其子集做join关联操作，因此需要依赖该索引列做匹配。
期望输出效果如下：

index,accountId,deposit,Withdrawals
1,anil0001,50,10
2,vikas0002,10,10
3,ravi0003,20,10
4,user1111,NaN,20

查阅接口后发现series类型没有提供GetIndex/Index类方法，无法直接在Rapply回调中获取当前行号。

实现方法 #

不需要通过Rapply生成索引列，直接基于DataFrame总行数预先生成从1开始的连续整数序列，构造为series后通过Mutate追加到表中，再调整列顺序即可，实现代码如下：

package main

import (
	"fmt"
	"strings"

	"github.com/go-gota/gota/dataframe"
	"github.com/go-gota/gota/series"
)

func main() {
	csvStr := `accountId,deposit,Withdrawals
anil0001,50,10
vikas0002,10,10
ravi0003,20,10
user1111,NaN,20`

	df := dataframe.ReadCSV(strings.NewReader(csvStr))

	// 生成从1开始的行索引列
	rowNum := df.Nrow()
	indexVals := make([]int, rowNum)
	for i := range indexVals {
		indexVals[i] = i + 1 // 需要从0开始计数直接赋值i即可
	}
	indexCol := series.New(indexVals, series.Int, "index")

	// 追加索引列并调整列顺序，将索引列放在最左侧
	df = df.Mutate(indexCol).Select([]string{"index", "accountId", "deposit", "Withdrawals"})

	fmt.Println(df)
}

运行上述代码即可得到和预期完全一致的输出。

注意事项 #

• 后续如果要做DataFrame子集筛选，要在筛选操作前先加索引列，筛选后的子集行会保留原始行的索引值，直接就能用来和原表做join。
• 这种预生成序列的方式比逐行Rapply回调性能更好，索引值稳定，不会受遍历逻辑影响。

内容的提问来源于stack exchange，提问作者Yuseferi