mirror of
https://github.com/johnkerl/miller.git
synced 2026-07-23 16:08:43 +00:00
d536318ed6
Miller assumes 64-bit integers, but in Go, the int type varies in size depending on the architecture: 32-bit architectures have int equivalent to int32. As a result, the supported range of integer values is greatly reduced on 32-bit architectures compared to what is suggested by the documentation. This patch explicitly uses int64 wherever 64-bit integers are assumed. Test cases affected by the behaviour of the random generator are updated to reflect the new values (the existing seed doesn't produce the same behaviour since the way random values are generated has changed). Signed-off-by: Stephen Kitt <steve@sk2.org>
437 lines
13 KiB
Go
437 lines
13 KiB
Go
package transformers
|
|
|
|
import (
|
|
"bytes"
|
|
"container/list"
|
|
"fmt"
|
|
"net/url"
|
|
"os"
|
|
"strings"
|
|
|
|
"github.com/johnkerl/miller/internal/pkg/cli"
|
|
"github.com/johnkerl/miller/internal/pkg/mlrval"
|
|
"github.com/johnkerl/miller/internal/pkg/output"
|
|
"github.com/johnkerl/miller/internal/pkg/types"
|
|
)
|
|
|
|
// ----------------------------------------------------------------
|
|
const verbNameSplit = "split"
|
|
const splitDefaultOutputFileNamePrefix = "split"
|
|
|
|
var SplitSetup = TransformerSetup{
|
|
Verb: verbNameSplit,
|
|
UsageFunc: transformerSplitUsage,
|
|
ParseCLIFunc: transformerSplitParseCLI,
|
|
IgnoresInput: false,
|
|
}
|
|
|
|
func transformerSplitUsage(
|
|
o *os.File,
|
|
doExit bool,
|
|
exitCode int,
|
|
) {
|
|
fmt.Fprintf(o, "Usage: %s %s [options] {filename}\n", "mlr", verbNameSplit)
|
|
fmt.Fprintf(o,
|
|
`Options:
|
|
-n {n}: Cap file sizes at N records.
|
|
-m {m}: Produce M files, round-robining records among them.
|
|
-g {a,b,c}: Write separate files with records having distinct values for fields named a,b,c.
|
|
Exactly one of -m, -n, or -g must be supplied.
|
|
--prefix {p} Specify filename prefix; default "`+splitDefaultOutputFileNamePrefix+`".
|
|
--suffix {s} Specify filename suffix; default is from mlr output format, e.g. "csv".
|
|
-a Append to existing file(s), if any, rather than overwriting.
|
|
-v Send records along to downstream verbs as well as splitting to files.
|
|
-h|--help Show this message.
|
|
Any of the output-format command-line flags (see mlr -h). For example, using
|
|
mlr --icsv --from myfile.csv split --ojson -n 1000
|
|
the input is CSV, but the output files are JSON.
|
|
|
|
Examples: Suppose myfile.csv has 1,000,000 records.
|
|
|
|
100 output files, 10,000 records each. First 10,000 records in split_1.csv, next in split_2.csv, etc.
|
|
mlr --csv --from myfile.csv split -n 10000
|
|
|
|
10 output files, 100,000 records each. Records 1,11,21,etc in split_1.csv, records 2,12,22, etc in split_2.csv, etc.
|
|
mlr --csv --from myfile.csv split -m 10
|
|
Same, but with JSON output.
|
|
mlr --csv --from myfile.csv split -m 10 -o json
|
|
|
|
Same but instead of split_1.csv, split_2.csv, etc. there are test_1.dat, test_2.dat, etc.
|
|
mlr --csv --from myfile.csv split -m 10 --prefix test --suffix dat
|
|
Same, but written to the /tmp/ directory.
|
|
mlr --csv --from myfile.csv split -m 10 --prefix /tmp/test --suffix dat
|
|
|
|
If the shape field has values triangle and square, then there will be split_triangle.csv and split_square.csv.
|
|
mlr --csv --from myfile.csv split -g shape
|
|
|
|
If the color field has values yellow and green, and the shape field has values triangle and square,
|
|
then there will be split_yellow_triangle.csv, split_yellow_square.csv, etc.
|
|
mlr --csv --from myfile.csv split -g color,shape
|
|
|
|
See also the "tee" DSL function which lets you do more ad-hoc customization.
|
|
`)
|
|
if doExit {
|
|
os.Exit(exitCode)
|
|
}
|
|
}
|
|
|
|
func transformerSplitParseCLI(
|
|
pargi *int,
|
|
argc int,
|
|
args []string,
|
|
mainOptions *cli.TOptions,
|
|
doConstruct bool, // false for first pass of CLI-parse, true for second pass
|
|
) IRecordTransformer {
|
|
|
|
// Skip the verb name from the current spot in the mlr command line
|
|
argi := *pargi
|
|
verb := args[argi]
|
|
argi++
|
|
|
|
var n int64 = 0
|
|
var doMod bool = false
|
|
var doSize bool = false
|
|
var groupByFieldNames []string = nil
|
|
var emitDownstream bool = false
|
|
var doAppend bool = false
|
|
var outputFileNamePrefix string = splitDefaultOutputFileNamePrefix
|
|
var outputFileNameSuffix string = "uninit"
|
|
haveOutputFileNameSuffix := false
|
|
|
|
var localOptions *cli.TOptions = nil
|
|
if mainOptions != nil {
|
|
copyThereof := *mainOptions // struct copy
|
|
localOptions = ©Thereof
|
|
}
|
|
|
|
// Parse local flags.
|
|
for argi < argc /* variable increment: 1 or 2 depending on flag */ {
|
|
opt := args[argi]
|
|
if !strings.HasPrefix(opt, "-") {
|
|
break // No more flag options to process
|
|
}
|
|
if args[argi] == "--" {
|
|
break // All transformers must do this so main-flags can follow verb-flags
|
|
}
|
|
argi++
|
|
|
|
if opt == "-h" || opt == "--help" {
|
|
transformerSplitUsage(os.Stdout, true, 0)
|
|
|
|
} else if opt == "-n" {
|
|
n = cli.VerbGetIntArgOrDie(verb, opt, args, &argi, argc)
|
|
doSize = true
|
|
|
|
} else if opt == "-m" {
|
|
n = cli.VerbGetIntArgOrDie(verb, opt, args, &argi, argc)
|
|
doMod = true
|
|
|
|
} else if opt == "-g" {
|
|
groupByFieldNames = cli.VerbGetStringArrayArgOrDie(verb, opt, args, &argi, argc)
|
|
|
|
} else if opt == "--prefix" {
|
|
outputFileNamePrefix = cli.VerbGetStringArgOrDie(verb, opt, args, &argi, argc)
|
|
|
|
} else if opt == "--suffix" {
|
|
outputFileNameSuffix = cli.VerbGetStringArgOrDie(verb, opt, args, &argi, argc)
|
|
haveOutputFileNameSuffix = true
|
|
|
|
} else if opt == "-a" {
|
|
doAppend = true
|
|
|
|
} else if opt == "-v" {
|
|
emitDownstream = true
|
|
|
|
} else {
|
|
// This is inelegant. For error-proofing we advance argi already in our
|
|
// loop (so individual if-statements don't need to). However,
|
|
// ParseWriterOptions expects it unadvanced.
|
|
largi := argi - 1
|
|
if cli.FLAG_TABLE.Parse(args, argc, &largi, localOptions) {
|
|
// This lets mlr main and mlr split have different output formats.
|
|
// Nothing else to handle here.
|
|
argi = largi
|
|
} else {
|
|
transformerSplitUsage(os.Stderr, true, 1)
|
|
}
|
|
}
|
|
}
|
|
|
|
doGroup := groupByFieldNames != nil
|
|
if !doMod && !doSize && !doGroup {
|
|
fmt.Fprintf(os.Stderr, "mlr %s: At least one of -m, -n, or -g is required.\n", verb)
|
|
os.Exit(1)
|
|
}
|
|
if (doMod && doSize) || (doMod && doGroup) || (doSize && doGroup) {
|
|
fmt.Fprintf(os.Stderr, "mlr %s: Only one of -m, -n, or -g is required.\n", verb)
|
|
os.Exit(1)
|
|
}
|
|
|
|
cli.FinalizeWriterOptions(&localOptions.WriterOptions)
|
|
if !haveOutputFileNameSuffix {
|
|
outputFileNameSuffix = localOptions.WriterOptions.OutputFileFormat
|
|
}
|
|
|
|
*pargi = argi
|
|
if !doConstruct { // All transformers must do this for main command-line parsing
|
|
return nil
|
|
}
|
|
|
|
transformer, err := NewTransformerSplit(
|
|
n,
|
|
doMod,
|
|
doSize,
|
|
groupByFieldNames,
|
|
emitDownstream,
|
|
doAppend,
|
|
outputFileNamePrefix,
|
|
outputFileNameSuffix,
|
|
&localOptions.WriterOptions,
|
|
)
|
|
if err != nil {
|
|
// Error message already printed out
|
|
os.Exit(1)
|
|
}
|
|
|
|
return transformer
|
|
}
|
|
|
|
// ----------------------------------------------------------------
|
|
type TransformerSplit struct {
|
|
n int64
|
|
outputFileNamePrefix string
|
|
outputFileNameSuffix string
|
|
emitDownstream bool
|
|
ungroupedCounter int64
|
|
groupByFieldNames []string
|
|
recordWriterOptions *cli.TWriterOptions
|
|
doAppend bool
|
|
|
|
// For doSize ungrouped: only one file open at a time
|
|
outputHandler output.OutputHandler
|
|
previousQuotient int64
|
|
|
|
// For all other cases: multiple files open at a time
|
|
outputHandlerManager output.OutputHandlerManager
|
|
|
|
recordTransformerFunc RecordTransformerFunc
|
|
}
|
|
|
|
func NewTransformerSplit(
|
|
n int64,
|
|
doMod bool,
|
|
doSize bool,
|
|
groupByFieldNames []string,
|
|
emitDownstream bool,
|
|
doAppend bool,
|
|
outputFileNamePrefix string,
|
|
outputFileNameSuffix string,
|
|
recordWriterOptions *cli.TWriterOptions,
|
|
) (*TransformerSplit, error) {
|
|
|
|
tr := &TransformerSplit{
|
|
n: n,
|
|
outputFileNamePrefix: outputFileNamePrefix,
|
|
outputFileNameSuffix: outputFileNameSuffix,
|
|
emitDownstream: emitDownstream,
|
|
ungroupedCounter: 0,
|
|
groupByFieldNames: groupByFieldNames,
|
|
recordWriterOptions: recordWriterOptions,
|
|
doAppend: doAppend,
|
|
|
|
outputHandler: nil,
|
|
previousQuotient: -1,
|
|
}
|
|
|
|
tr.outputHandlerManager = output.NewFileOutputHandlerManager(recordWriterOptions, doAppend)
|
|
|
|
if groupByFieldNames != nil {
|
|
tr.recordTransformerFunc = tr.splitGrouped
|
|
} else if doMod {
|
|
tr.recordTransformerFunc = tr.splitModUngrouped
|
|
} else {
|
|
tr.recordTransformerFunc = tr.splitSizeUngrouped
|
|
}
|
|
|
|
return tr, nil
|
|
}
|
|
|
|
func (tr *TransformerSplit) Transform(
|
|
inrecAndContext *types.RecordAndContext,
|
|
outputRecordsAndContexts *list.List, // list of *types.RecordAndContext
|
|
inputDownstreamDoneChannel <-chan bool,
|
|
outputDownstreamDoneChannel chan<- bool,
|
|
) {
|
|
HandleDefaultDownstreamDone(inputDownstreamDoneChannel, outputDownstreamDoneChannel)
|
|
tr.recordTransformerFunc(inrecAndContext, outputRecordsAndContexts, inputDownstreamDoneChannel,
|
|
outputDownstreamDoneChannel)
|
|
}
|
|
|
|
func (tr *TransformerSplit) splitModUngrouped(
|
|
inrecAndContext *types.RecordAndContext,
|
|
outputRecordsAndContexts *list.List, // list of *types.RecordAndContext
|
|
inputDownstreamDoneChannel <-chan bool,
|
|
outputDownstreamDoneChannel chan<- bool,
|
|
) {
|
|
if !inrecAndContext.EndOfStream {
|
|
remainder := 1 + (tr.ungroupedCounter % tr.n)
|
|
filename := tr.makeUngroupedOutputFileName(remainder)
|
|
|
|
err := tr.outputHandlerManager.WriteRecordAndContext(inrecAndContext, filename)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "mlr: file-write error: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
if tr.emitDownstream {
|
|
outputRecordsAndContexts.PushBack(inrecAndContext)
|
|
}
|
|
|
|
tr.ungroupedCounter++
|
|
|
|
} else {
|
|
outputRecordsAndContexts.PushBack(inrecAndContext) // end-of-stream marker
|
|
errs := tr.outputHandlerManager.Close()
|
|
if len(errs) > 0 {
|
|
for _, err := range errs {
|
|
fmt.Fprintf(os.Stderr, "mlr: file-close error: %v\n", err)
|
|
}
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (tr *TransformerSplit) splitSizeUngrouped(
|
|
inrecAndContext *types.RecordAndContext,
|
|
outputRecordsAndContexts *list.List, // list of *types.RecordAndContext
|
|
inputDownstreamDoneChannel <-chan bool,
|
|
outputDownstreamDoneChannel chan<- bool,
|
|
) {
|
|
var err error
|
|
if !inrecAndContext.EndOfStream {
|
|
quotient := 1 + (tr.ungroupedCounter / tr.n)
|
|
|
|
if quotient != tr.previousQuotient {
|
|
if tr.outputHandler != nil {
|
|
err = tr.outputHandler.Close()
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "mlr: file-close error: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
filename := tr.makeUngroupedOutputFileName(quotient)
|
|
tr.outputHandler, err = output.NewFileOutputHandler(
|
|
filename,
|
|
tr.recordWriterOptions,
|
|
tr.doAppend,
|
|
)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "mlr: file-open error: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
tr.previousQuotient = quotient
|
|
}
|
|
|
|
err = tr.outputHandler.WriteRecordAndContext(inrecAndContext)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "mlr: file-write error: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
if tr.emitDownstream {
|
|
outputRecordsAndContexts.PushBack(inrecAndContext)
|
|
}
|
|
|
|
tr.ungroupedCounter++
|
|
|
|
} else {
|
|
outputRecordsAndContexts.PushBack(inrecAndContext) // end-of-stream marker
|
|
|
|
if tr.outputHandler != nil {
|
|
err := tr.outputHandler.Close()
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "mlr: file-close error: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (tr *TransformerSplit) splitGrouped(
|
|
inrecAndContext *types.RecordAndContext,
|
|
outputRecordsAndContexts *list.List, // list of *types.RecordAndContext
|
|
inputDownstreamDoneChannel <-chan bool,
|
|
outputDownstreamDoneChannel chan<- bool,
|
|
) {
|
|
if !inrecAndContext.EndOfStream {
|
|
var filename string
|
|
groupByFieldValues, ok := inrecAndContext.Record.GetSelectedValues(tr.groupByFieldNames)
|
|
if !ok {
|
|
filename = fmt.Sprintf("%s_ungrouped.%s", tr.outputFileNamePrefix, tr.outputFileNameSuffix)
|
|
} else {
|
|
filename = tr.makeGroupedOutputFileName(groupByFieldValues)
|
|
}
|
|
err := tr.outputHandlerManager.WriteRecordAndContext(inrecAndContext, filename)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "mlr: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
if tr.emitDownstream {
|
|
outputRecordsAndContexts.PushBack(inrecAndContext)
|
|
}
|
|
|
|
} else {
|
|
outputRecordsAndContexts.PushBack(inrecAndContext) // emit end-of-stream marker
|
|
|
|
errs := tr.outputHandlerManager.Close()
|
|
if len(errs) > 0 {
|
|
for _, err := range errs {
|
|
fmt.Fprintf(os.Stderr, "mlr: file-close error: %v\n", err)
|
|
}
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
}
|
|
|
|
// makeUngroupedOutputFileName example: "split_53.csv"
|
|
func (tr *TransformerSplit) makeUngroupedOutputFileName(k int64) string {
|
|
return fmt.Sprintf("%s_%d.%s", tr.outputFileNamePrefix, k, tr.outputFileNameSuffix)
|
|
}
|
|
|
|
// makeGroupedOutputFileName example: "split_orange.csv"
|
|
func (tr *TransformerSplit) makeGroupedOutputFileName(
|
|
groupByFieldValues []*mlrval.Mlrval,
|
|
) string {
|
|
var buffer bytes.Buffer
|
|
buffer.WriteString(tr.outputFileNamePrefix)
|
|
for _, groupByFieldValue := range groupByFieldValues {
|
|
buffer.WriteString("_")
|
|
buffer.WriteString(url.QueryEscape(groupByFieldValue.String()))
|
|
}
|
|
buffer.WriteString(".")
|
|
buffer.WriteString(tr.outputFileNameSuffix)
|
|
return buffer.String()
|
|
}
|
|
|
|
// makeGroupedIndexedOutputFileName example: "split_yellow_53.csv"
|
|
func (tr *TransformerSplit) makeGroupedIndexedOutputFileName(
|
|
groupByFieldValues []*mlrval.Mlrval,
|
|
index int,
|
|
) string {
|
|
// URL-escape the fields which come from data and which may have '/'
|
|
// etc within. Don't URL-escape the prefix since people may want to
|
|
// use prefixes like '/tmp/split' to write to the /tmp directory, etc.
|
|
var buffer bytes.Buffer
|
|
buffer.WriteString(tr.outputFileNamePrefix)
|
|
for _, groupByFieldValue := range groupByFieldValues {
|
|
buffer.WriteString("_")
|
|
buffer.WriteString(url.QueryEscape(groupByFieldValue.String()))
|
|
}
|
|
buffer.WriteString(fmt.Sprintf("_%d", index))
|
|
buffer.WriteString(".")
|
|
buffer.WriteString(tr.outputFileNameSuffix)
|
|
return buffer.String()
|
|
}
|