miller/pkg/transformers/merge_fields.go

package transformers

import (
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/johnkerl/miller/v6/pkg/cli"
	"github.com/johnkerl/miller/v6/pkg/lib"
	"github.com/johnkerl/miller/v6/pkg/transformers/utils"
	"github.com/johnkerl/miller/v6/pkg/types"
)

const verbNameMergeFields = "merge-fields"

var mergeFieldsOptions = []OptionSpec{
	{Flag: "-a", Arg: "{sum,count,...}", Type: "enum", Desc: "Names of accumulators: one or more of the accumulators listed below.", Values: []string{"count", "null_count", "distinct_count", "mode", "antimode", "sum", "mean", "mad", "var", "stddev", "meaneb", "skewness", "kurtosis", "min", "max", "minlen", "maxlen"}},
	{Flag: "-f", Arg: "{a,b,c}", Type: "csv-list", Desc: "Value-field names on which to compute statistics. Requires -o."},
	{Flag: "-r", Arg: "{a,b,c}", Type: "csv-list", Desc: "Regular expressions for value-field names on which to compute statistics. Requires -o."},
	{Flag: "-c", Arg: "{a,b,c}", Type: "csv-list", Desc: "Substrings for collapse mode: all fields which have the same names after removing substrings will be accumulated together. Please see examples below."},
	{Flag: "-i", Type: "bool", Desc: "Use interpolated percentiles, like R's type=7; default like type=1. Not sensical for string-valued fields."},
	{Flag: "-o", Arg: "{name}", Type: "string", Desc: "Output field basename for -f/-r."},
	{Flag: "-k", Type: "bool", Desc: "Keep the input fields which contributed to the output statistics; the default is to omit them."},
	{Flag: "-S", Type: "bool", Desc: "No-op flag for backward compatibility with Miller 5."},
	{Flag: "-F", Type: "bool", Desc: "No-op flag for backward compatibility with Miller 5."},
}

var MergeFieldsSetup = TransformerSetup{
	Verb:         verbNameMergeFields,
	UsageFunc:    transformerMergeFieldsUsage,
	ParseCLIFunc: transformerMergeFieldsParseCLI,
	IgnoresInput: false,
	Options:      mergeFieldsOptions,
}

type mergeByType int

const (
	e_MERGE_BY_NAME_LIST mergeByType = iota
	e_MERGE_BY_NAME_REGEX
	e_MERGE_BY_COLLAPSING
	e_MERGE_UNSPECIFIED
)

func transformerMergeFieldsUsage(
	o *os.File,
) {
	argv0 := "mlr"
	verb := verbNameMergeFields
	fmt.Fprintf(o, "Usage: %s %s [options]\n", argv0, verb)
	fmt.Fprintf(o, "Computes univariate statistics for each input record, accumulated across\n")
	fmt.Fprintf(o, "specified fields.\n")
	WriteVerbOptions(o, mergeFieldsOptions)
	fmt.Fprintf(o, "\n")
	fmt.Fprintf(o, "Accumulators for -a:\n")
	utils.ListStats1Accumulators(o)
	fmt.Fprintf(o, "\n")
	fmt.Fprintf(o, "String-valued data make sense unless arithmetic on them is required,\n")
	fmt.Fprintf(o, "e.g. for sum, mean, interpolated percentiles, etc. In case of mixed data,\n")
	fmt.Fprintf(o, "numbers are less than strings.\n")
	fmt.Fprintf(o, "\n")
	fmt.Fprintf(o, "Example input data: \"a_in_x=1,a_out_x=2,b_in_y=4,b_out_x=8\".\n")
	fmt.Fprintf(o, "Example: %s %s -a sum,count -f a_in_x,a_out_x -o foo\n", argv0, verb)
	fmt.Fprintf(o, "  produces \"b_in_y=4,b_out_x=8,foo_sum=3,foo_count=2\" since \"a_in_x,a_out_x\" are\n")
	fmt.Fprintf(o, "  summed over.\n")
	fmt.Fprintf(o, "Example: %s %s -a sum,count -r in_,out_ -o bar\n", argv0, verb)
	fmt.Fprintf(o, "  produces \"bar_sum=15,bar_count=4\" since all four fields are summed over.\n")
	fmt.Fprintf(o, "Example: %s %s -a sum,count -c in_,out_\n", argv0, verb)
	fmt.Fprintf(o, "  produces \"a_x_sum=3,a_x_count=2,b_y_sum=4,b_y_count=1,b_x_sum=8,b_x_count=1\"\n")
	fmt.Fprintf(o, "  since \"a_in_x\" and \"a_out_x\" both collapse to \"a_x\", \"b_in_y\" collapses to\n")
	fmt.Fprintf(o, "  \"b_y\", and \"b_out_x\" collapses to \"b_x\".\n")
}

func transformerMergeFieldsParseCLI(
	pargi *int,
	argc int,
	args []string,
	_ *cli.TOptions,
	doConstruct bool, // false for first pass of CLI-parse, true for second pass
) (RecordTransformer, error) {

	// Skip the verb name from the current spot in the mlr command line
	argi := *pargi
	verb := args[argi]
	argi++

	accumulatorNameList := []string{}
	valueFieldNameList := []string{}
	outputFieldBasename := ""
	doWhich := e_MERGE_UNSPECIFIED
	keepInputFields := false
	doInterpolatedPercentiles := false

	var err error
	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++

		switch opt {
		case "-h", "--help":
			transformerMergeFieldsUsage(os.Stdout)
			return nil, cli.ErrHelpRequested

		case "-a":
			accumulatorNameList, err = cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}

		case "-f":
			valueFieldNameList, err = cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}
			doWhich = e_MERGE_BY_NAME_LIST

		case "-r":
			valueFieldNameList, err = cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}
			doWhich = e_MERGE_BY_NAME_REGEX

		case "-c":
			valueFieldNameList, err = cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}
			doWhich = e_MERGE_BY_COLLAPSING

		case "-o":
			outputFieldBasename, err = cli.VerbGetStringArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}

		case "-k":
			keepInputFields = true

		case "-i":
			doInterpolatedPercentiles = true

		case "-S", "-F":
			// No-op pass-through for backward compatibility with Miller 5

		default:
			return nil, cli.VerbErrorf(verb, "option \"%s\" not recognized", opt)
		}
	}

	// TODO: libify for use across verbs.
	if len(accumulatorNameList) == 0 {
		return nil, cli.VerbErrorf(verbNameMergeFields, "-a option is required")
	}
	if len(valueFieldNameList) == 0 {
		return nil, cli.VerbErrorf(verbNameMergeFields, "-f option is required")
	}
	if outputFieldBasename == "" {
		if doWhich == e_MERGE_BY_NAME_LIST || doWhich == e_MERGE_BY_NAME_REGEX {
			return nil, cli.VerbErrorf(verbNameMergeFields, "-o output field basename required")
		}
	}

	*pargi = argi
	if !doConstruct { // All transformers must do this for main command-line parsing
		return nil, nil
	}

	transformer, err := NewTransformerMergeFields(
		accumulatorNameList,
		valueFieldNameList,
		outputFieldBasename,
		doWhich,
		doInterpolatedPercentiles,
		keepInputFields,
	)
	if err != nil {
		return nil, err
	}

	return transformer, nil
}

// Given: accumulate count,sum on values x,y group by a,b.
//
// Example input:       Example output:
//   a b x y            a b x_count x_sum y_count y_sum
//   s t 1 2            s t 2       6     2       8
//   u v 3 4            u v 1       3     1       4
//   s t 5 6            u w 1       7     1       9
//   u w 7 9
//
// Multilevel hashmap structure:
// {
//   "s,t" : {                <--- group-by field names
//     "x" : {                  <--- value field name
//       "count" : Stats1CountAccumulator object,
//       "sum"   : Stats1SumAccumulator  object
//     },
//     "y" : {
//       "count" : Stats1CountAccumulator object,
//       "sum"   : Stats1SumAccumulator  object
//     },
//   },
//   "u,v" : {
//     "x" : {
//       "count" : Stats1CountAccumulator object,
//       "sum"   : Stats1SumAccumulator  object
//     },
//     "y" : {
//       "count" : Stats1CountAccumulator object,
//       "sum"   : Stats1SumAccumulator  object
//     },
//   },
//   "u,w" : {
//     "x" : {
//       "count" : Stats1CountAccumulator object,
//       "sum"   : Stats1SumAccumulator  object
//     },
//     "y" : {
//       "count" : Stats1CountAccumulator object,
//       "sum"   : Stats1SumAccumulator  object
//     },
//   },
// }

type TransformerMergeFields struct {
	// Input:
	accumulatorNameList       []string
	valueFieldNameList        []string
	outputFieldBasename       string
	doInterpolatedPercentiles bool
	keepInputFields           bool

	// State:
	accumulatorFactory    *utils.Stats1AccumulatorFactory
	valueFieldNameRegexes []*regexp.Regexp

	// Ordered map from accumulator name to accumulator
	namedAccumulators *lib.OrderedMap[*utils.Stats1NamedAccumulator]

	recordTransformerFunc RecordTransformerFunc
}

func NewTransformerMergeFields(
	accumulatorNameList []string,
	valueFieldNameList []string,
	outputFieldBasename string,
	doWhich mergeByType,
	doInterpolatedPercentiles bool,
	keepInputFields bool,
) (*TransformerMergeFields, error) {

	for _, accumulatorName := range accumulatorNameList {
		if !utils.ValidateStats1AccumulatorName(accumulatorName) {
			return nil, fmt.Errorf(
				`mlr %s: accumulator "%s" not found`,
				verbNameMergeFields, accumulatorName,
			)
		}
	}

	tr := &TransformerMergeFields{
		accumulatorNameList:       accumulatorNameList,
		valueFieldNameList:        valueFieldNameList,
		outputFieldBasename:       outputFieldBasename,
		doInterpolatedPercentiles: doInterpolatedPercentiles,
		keepInputFields:           keepInputFields,
		accumulatorFactory:        utils.NewStats1AccumulatorFactory(),
		namedAccumulators:         lib.NewOrderedMap[*utils.Stats1NamedAccumulator](),
	}

	tr.valueFieldNameRegexes = make([]*regexp.Regexp, len(valueFieldNameList))
	for i, regexString := range valueFieldNameList {
		// Handles "a.*b"i Miller case-insensitive-regex specification
		regex, err := lib.CompileMillerRegex(regexString)
		if err != nil {
			fmt.Fprintf(
				os.Stderr,
				"%s %s: cannot compile regex [%s]\n",
				"mlr", verbNameCut, regexString,
			)
			os.Exit(1)
		}
		tr.valueFieldNameRegexes[i] = regex
	}

	for _, accumulatorName := range accumulatorNameList {
		accumulator := tr.accumulatorFactory.MakeNamedAccumulator(
			accumulatorName,
			"", // grouping-key used for stats1, not here
			outputFieldBasename,
			doInterpolatedPercentiles,
		)
		tr.namedAccumulators.Put(accumulatorName, accumulator)
	}

	switch doWhich {
	case e_MERGE_BY_NAME_LIST:
		tr.recordTransformerFunc = tr.transformByNameList
	case e_MERGE_BY_NAME_REGEX:
		tr.recordTransformerFunc = tr.transformByNameRegex
	case e_MERGE_BY_COLLAPSING:
		tr.recordTransformerFunc = tr.transformByCollapsing
	default:
		lib.InternalCodingErrorIf(true)
	}

	return tr, nil
}

func (tr *TransformerMergeFields) Transform(
	inrecAndContext *types.RecordAndContext,
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
	inputDownstreamDoneChannel <-chan bool,
	outputDownstreamDoneChannel chan<- bool,
) {
	HandleDefaultDownstreamDone(inputDownstreamDoneChannel, outputDownstreamDoneChannel)
	tr.recordTransformerFunc(inrecAndContext, outputRecordsAndContexts, inputDownstreamDoneChannel, outputDownstreamDoneChannel)
}

func (tr *TransformerMergeFields) transformByNameList(
	inrecAndContext *types.RecordAndContext,
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
	inputDownstreamDoneChannel <-chan bool,
	outputDownstreamDoneChannel chan<- bool,
) {
	if inrecAndContext.EndOfStream {
		*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext) // end-of-stream marker
		return
	}

	inrec := inrecAndContext.Record

	for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
		accumulator := pa.Value
		accumulator.Reset() // re-use from one record to the next
	}

	for _, valueFieldName := range tr.valueFieldNameList {
		mvalue := inrec.Get(valueFieldName)
		if mvalue == nil { // key not present
			continue
		}

		if mvalue.IsVoid() { // key present with empty value
			if !tr.keepInputFields {
				inrec.Remove(valueFieldName)
			}
			continue
		}

		for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
			accumulator := pa.Value
			accumulator.Ingest(mvalue)
		}

		if !tr.keepInputFields {
			inrec.Remove(valueFieldName)
		}
	}

	for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
		accumulator := pa.Value
		key, value := accumulator.Emit()
		inrec.PutReference(key, value)
	}

	*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext)
}

func (tr *TransformerMergeFields) transformByNameRegex(
	inrecAndContext *types.RecordAndContext,
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
	inputDownstreamDoneChannel <-chan bool,
	outputDownstreamDoneChannel chan<- bool,
) {
	if inrecAndContext.EndOfStream {
		*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext) // end-of-stream marker
		return
	}

	inrec := inrecAndContext.Record

	for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
		accumulator := pa.Value
		accumulator.Reset() // re-use from one record to the next
	}

	for pe := inrec.Head; pe != nil; /* increment inside loop*/ {
		valueFieldName := pe.Key

		matched := false
		for _, valueFieldNameRegex := range tr.valueFieldNameRegexes {
			if valueFieldNameRegex.MatchString(pe.Key) {
				matched = true
				break
			}
		}
		if !matched {
			pe = pe.Next
			continue
		}

		mvalue := inrec.Get(valueFieldName)

		if mvalue == nil { // Key not present
			pe = pe.Next
			continue
		}

		if mvalue.IsVoid() { // Key present with empty value
			if !tr.keepInputFields { // We are modifying the record while iterating over it.
				next := pe.Next
				inrec.Unlink(pe)
				pe = next
			} else {
				pe = pe.Next
			}
			continue
		}

		for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
			accumulator := pa.Value
			accumulator.Ingest(mvalue)
		}

		if !tr.keepInputFields { // We are modifying the record while iterating over it.
			next := pe.Next
			inrec.Unlink(pe)
			pe = next
		} else {
			pe = pe.Next
		}
	}

	for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
		accumulator := pa.Value
		key, value := accumulator.Emit()
		inrec.PutReference(key, value)
	}

	*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext)
}

// mlr merge-fields -c in_,out_ -a sum
// a_in_x  1     a_sum_x 3
// a_out_x 2     b_sum_y 4
// b_in_y  4     b_sum_x 8
// b_out_x 8

func (tr *TransformerMergeFields) transformByCollapsing(
	inrecAndContext *types.RecordAndContext,
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
	inputDownstreamDoneChannel <-chan bool,
	outputDownstreamDoneChannel chan<- bool,
) {
	if inrecAndContext.EndOfStream {
		*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext) // end-of-stream marker
		return
	}

	inrec := inrecAndContext.Record
	tr.accumulatorFactory.Reset() // discard cached percentile-keepers

	// Ordered map from short name to accumulator name to accumulator
	collapseAccumulators := lib.NewOrderedMap[*lib.OrderedMap[*utils.Stats1NamedAccumulator]]()

	for pe := inrec.Head; pe != nil; /* increment inside loop */ {
		valueFieldName := pe.Key

		matched := false
		shortName := ""
		for _, valueFieldNameRegex := range tr.valueFieldNameRegexes {
			matched = valueFieldNameRegex.MatchString(pe.Key)
			if matched {
				// TODO: comment re matrix
				shortName = lib.RegexCompiledSub(valueFieldName, valueFieldNameRegex, "", nil)
				break
			}
		}
		if !matched {
			pe = pe.Next
			continue
		}

		mvalue := inrec.Get(valueFieldName)
		if mvalue == nil { // Key aesent
			pe = pe.Next
			continue
		}

		var namedAccumulators *lib.OrderedMap[*utils.Stats1NamedAccumulator]
		iNamedAccumulators := collapseAccumulators.Get(shortName)
		if iNamedAccumulators == nil {
			namedAccumulators = lib.NewOrderedMap[*utils.Stats1NamedAccumulator]()
			for _, accumulatorName := range tr.accumulatorNameList {
				accumulator := tr.accumulatorFactory.MakeNamedAccumulator(
					accumulatorName,
					"", // grouping-key used for stats1, not here
					shortName,
					tr.doInterpolatedPercentiles,
				)
				namedAccumulators.Put(accumulatorName, accumulator)
			}
			collapseAccumulators.Put(shortName, namedAccumulators)
		} else {
			namedAccumulators = iNamedAccumulators
		}

		// The accumulator has been initialized with default values; continue
		// here. (If we were to continue before namedAccumulators.Put(...) we
		// would be failing to construct the accumulator.)
		if mvalue.IsVoid() { // key present with empty value
			if !tr.keepInputFields { // We are modifying the record while iterating over it.
				next := pe.Next
				inrec.Unlink(pe)
				pe = next
			} else {
				pe = pe.Next
			}
			continue
		}

		for pa := namedAccumulators.Head; pa != nil; pa = pa.Next {
			accumulator := pa.Value
			accumulator.Ingest(mvalue)
		}

		if !tr.keepInputFields { // We are modifying the record while iterating over it.
			next := pe.Next
			inrec.Unlink(pe)
			pe = next
		} else {
			pe = pe.Next
		}
	}

	for ps := collapseAccumulators.Head; ps != nil; ps = ps.Next {
		namedAccumulators := ps.Value
		for pa := namedAccumulators.Head; pa != nil; pa = pa.Next {
			accumulator := pa.Value
			key, value := accumulator.Emit()
			inrec.PutReference(key, value)
		}
	}

	*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext)
}