forked from rocketlaunchr/dataframe-go
/
csv.go
113 lines (86 loc) · 2.24 KB
/
csv.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
// Copyright 2018-20 PJ Engineering and Business Solutions Pty. Ltd. All rights reserved.
package exports
import (
"context"
"encoding/csv"
"io"
dataframe "github.com/rocketlaunchr/dataframe-go"
)
// CSVExportOptions contains options for ExportToCSV function.
type CSVExportOptions struct {
// NullString is used to set what nil values should be encoded to.
// Common options are NULL, \N, NaN, NA.
NullString *string
// Range is used to export a subset of rows from the dataframe.
Range dataframe.Range
// Separator is the field delimiter. A common option is ',', which is
// the default if CSVExportOptions is not provided.
Separator rune
// UseCRLF determines the line terminator.
// When true, it is set to \r\n.
UseCRLF bool
}
// ExportToCSV exports a Dataframe to a CSV file.
func ExportToCSV(ctx context.Context, w io.Writer, df *dataframe.DataFrame, options ...CSVExportOptions) error {
df.Lock()
defer df.Unlock()
header := []string{}
var r dataframe.Range
nullString := "NaN" // Default will be "NaN"
cw := csv.NewWriter(w)
if len(options) > 0 {
cw.Comma = options[0].Separator
cw.UseCRLF = options[0].UseCRLF
r = options[0].Range
if options[0].NullString != nil {
nullString = *options[0].NullString
}
}
for _, aSeries := range df.Series {
header = append(header, aSeries.Name())
}
if err := cw.Write(header); err != nil {
return err
}
nRows := df.NRows(dataframe.DontLock)
if nRows > 0 {
s, e, err := r.Limits(nRows)
if err != nil {
return err
}
flushCount := 0
for row := s; row <= e; row++ {
if err := ctx.Err(); err != nil {
return err
}
flushCount++
// flush after every 100 writes
if flushCount > 100 { // flush in the 101th count
cw.Flush()
if err := cw.Error(); err != nil {
return err
}
flushCount = 1
}
sVals := []string{}
for _, aSeries := range df.Series {
val := aSeries.Value(row)
if val == nil {
sVals = append(sVals, nullString)
} else {
sVals = append(sVals, aSeries.ValueString(row, dataframe.DontLock))
}
}
// Write every row
if err := cw.Write(sVals); err != nil {
return err
}
}
}
// flush before exit
cw.Flush()
if err := cw.Error(); err != nil {
return err
}
return nil
}