kshedden / dstream Goto Github PK

func TestConvert2(t *testing.T) {

data1 := `id,v1,v2,v3
1,2,3,4
1,3,4,5
2,4,5,6
3,5,6,7
3,99,99,99
3,100,101,102
4,200,201,202
`
        times100 := func(v map[string]interface{}, z interface{}) {
                 id := v["id"].([]float64)
                 y := z.([]float64)
                 for i := range id {
                      y[i] = id[i] * 100
                 }
        }
        b1 := bytes.NewReader([]byte(data1))
        d1 := FromCSV(b1).SetFloatVars([]string{"id","v1","v2","v3"}).HasHeader().Done()
        d1 = Apply(d1, "id100", times100, "float64")
        d1 = Convert(d1, "id100", "uint64")
        d1 = DropCols(d1, []string{"id"})
        d1n := d1.Names()
        vmap := make(map[string]int)
        for ix, na := range d1n {
            vmap[na] = ix
        }
        var id100pos int
        id100pos = vmap["id100"]
        fmt.Printf("variable position map: %v\n", vmap)
        cnum := 0
        for d1.Next(){
            cnum++
            fmt.Printf("id100 for chunk %d, using GetPos(), assert []uint64:\n%v\n", cnum, d1.GetPos\
(id100pos).([]uint64))
            fmt.Printf("id100 for chunk %d, using Get(), assert []uint64: %v\n", cnum,d1.Get("id100"\
).([]uint64))
        }
}

func TestToCSV(t * testing.T){
        data1 := `id,v1,v2,v3
1,2,3,4
1,3,4,5
2,4,5,6
3,5,6,7
3,99,99,99
3,100,101,102
4,200,201,202
`
        b1 := bytes.NewReader([]byte(data1))
        d1 := FromCSV(b1).SetFloatVars([]string{"id", "v1", "v2", "v3"}).HasHeader().Done()
        d1 = Segment(d1, []string{"id"})
        d1.Next()
        wtr, err := os.Create("tocsv_test.txt")
        if err != nil {
           t.Fail()
        }
        err = ToCSV(d1, wtr)
        if err != nil {
           fmt.Printf("%v\n",err)
           t.Fail()
        }

        wtr.Close()

        rdr, err := os.Open("tocsv_test.txt")
        if err != nil {
           t.Fail()
        }

        d2 := FromCSV(rdr).SetFloatVars([]string{"id","v1","v2","v3"}).SetChunkSize(2).HasHeader().Done()
        for d2.Next(){
            fmt.Printf("v1 = %v\n", d2.Get("v1"))
        }
        fmt.Printf("d2.NumObs() = %d\n", d2.NumObs())
        if d2.NumObs() != 7 {
           t.Fail()
        }

}

func ToCSV(data Dstream, wtr io.Writer) error {

        data.Reset()

        csw := csv.NewWriter(wtr)
        err := csw.Write(data.Names())
        if err != nil {
                return err
        }

        nvar := data.NumVar()
        rec := make([]string, nvar)
        for data.Next() {
                n := ilen(data.GetPos(0))

                for i := 0; i < n; i++ {
                        for j := 0; j < nvar; j++ {
                                switch x := data.GetPos(j).(type) {
                                case []float64:
                                        rec[j] = fmt.Sprintf("%.8f", x[i])
                                case []uint64:
                                        rec[j] = fmt.Sprintf("%d", x[i])
                                case []string:
                                        rec[j] = x[i]
                                }
                        }
                        fmt.Printf("writing %v\n", rec)
                        err = csw.Write(rec)
                        if err != nil {
                                panic(err)
                        }
                }
        }
        csw.Flush()
        err = csw.Error()
        if err != nil{
           panic(err)
        }

        return nil
}