miller/pkg/transformers/stats2.go

package transformers

import (
	"fmt"
	"os"
	"strings"

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

const verbNameStats2 = "stats2"

// For joining "x" and "y" into "x...y" for map keys. "," is another natural choice but would break
// if we were ever asked to process field names with commas in them.
const stats2KeySeparator = "\001"

var stats2Options = []OptionSpec{
	{Flag: "-a", Arg: "{linreg-ols,corr,...}", Type: "enum", Desc: "Names of accumulators: one or more of the listed values.", Values: []string{"linreg-ols", "linreg-pca", "r2", "logireg", "corr", "cov"}},
	{Flag: "-f", Arg: "{a,b,c,d}", Type: "csv-list", Desc: "Value-field name-pairs on which to compute statistics. There must be an even number of names."},
	{Flag: "-g", Arg: "{e,f,g}", Type: "csv-list", Desc: "Optional group-by-field names."},
	{Flag: "-v", Type: "bool", Desc: "Print additional output for linreg-pca."},
	{Flag: "-s", Type: "bool", Desc: "Print iterative stats. Useful in tail -f contexts, in which case please avoid pprint-format output since end of input stream will never be seen. Likewise, if input is coming from `tail -f`, be sure to use `--records-per-batch 1`."},
	{Flag: "--fit", Type: "bool", Desc: "Rather than printing regression parameters, applies them to the input data to compute new fit fields. All input records are held in memory until end of input stream. Has effect only for linreg-ols, linreg-pca, and logireg."},
	{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 Stats2Setup = TransformerSetup{
	Verb:         verbNameStats2,
	UsageFunc:    transformerStats2Usage,
	ParseCLIFunc: transformerStats2ParseCLI,
	IgnoresInput: false,
	Options:      stats2Options,
}

func transformerStats2Usage(
	o *os.File,
) {
	argv0 := "mlr"
	verb := verbNameStats2

	fmt.Fprintf(o, "Usage: %s %s [options]\n", argv0, verb)
	fmt.Fprintf(o, "Computes bivariate statistics for one or more given field-name pairs,\n")
	fmt.Fprintf(o, "accumulated across the input record stream.\n")
	WriteVerbOptions(o, stats2Options)
	fmt.Fprintf(o, "Names of accumulators for -a, one or more of:\n")

	utils.ListStats2Accumulators(o)

	fmt.Fprintf(o, "Only one of -s or --fit may be used.\n")
	fmt.Fprintf(o, "Example: %s %s -a linreg-pca -f x,y\n", argv0, verb)
	fmt.Fprintf(o, "Example: %s %s -a linreg-ols,r2 -f x,y -g size,shape\n", argv0, verb)
	fmt.Fprintf(o, "Example: %s %s -a corr -f x,y\n", argv0, verb)
}

func transformerStats2ParseCLI(
	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++

	var accumulatorNameList []string = nil
	var valueFieldNameList []string = nil
	groupByFieldNameList := []string{}
	doVerbose := false
	doIterativeStats := false
	doHoldAndFit := 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":
			transformerStats2Usage(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
			}

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

		case "-v":
			doVerbose = true

		case "-s":
			doIterativeStats = true

		case "--fit":
			doHoldAndFit = true

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

		case "-F":
			// The -F flag isn't used for stats2: all arithmetic here is
			// floating-point. Yet it is supported for step and stats1 for all
			// applicable stats1/step accumulators, so we accept here as well
			// for all applicable stats2 accumulators (i.e. none of them).

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

	if doIterativeStats && doHoldAndFit {
		return nil, cli.VerbErrorf(verb, "cannot combine -I and -H")
	}
	if accumulatorNameList == nil {
		return nil, cli.VerbErrorf(verb, "-a option is required")
	}
	if valueFieldNameList == nil {
		return nil, cli.VerbErrorf(verb, "-f option is required")
	}
	if len(valueFieldNameList)%2 != 0 {
		return nil, cli.VerbErrorf(verb, "argument to -f must have even number of fields")
	}

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

	transformer, err := NewTransformerStats2(
		accumulatorNameList,
		valueFieldNameList,
		groupByFieldNameList,
		doVerbose,
		doIterativeStats,
		doHoldAndFit,
	)
	if err != nil {
		return nil, err
	}

	return transformer, nil
}

type TransformerStats2 struct {
	// Input:
	accumulatorNameList  []string
	valueFieldNameList   []string
	groupByFieldNameList []string

	doVerbose        bool
	doIterativeStats bool
	doHoldAndFit     bool

	// State:
	accumulatorFactory *utils.Stats2AccumulatorFactory

	// Accumulators are indexed by
	//   groupByFieldName . value1FieldName+sep+value2FieldName . accumulatorName . accumulator object
	// This would be
	//   namedAccumulators map[string]map[string]map[string]IStats2Accumulator
	// except we need maps that preserve insertion order.
	namedAccumulators *lib.OrderedMap[*lib.OrderedMap[*lib.OrderedMap[utils.IStats2Accumulator]]]

	groupingKeysToGroupByFieldValues *lib.OrderedMap[[]*mlrval.Mlrval]

	// For hold-and-fit:
	// ordered map from grouping-key to list of RecordAndContext
	recordGroups *lib.OrderedMap[*[]*types.RecordAndContext]
}

func NewTransformerStats2(
	accumulatorNameList []string,
	valueFieldNameList []string,
	groupByFieldNameList []string,
	doVerbose bool,
	doIterativeStats bool,
	doHoldAndFit bool,
) (*TransformerStats2, error) {
	for _, name := range accumulatorNameList {
		if !utils.ValidateStats2AccumulatorName(name) {
			return nil, fmt.Errorf(`mlr stats2: accumulator "%s" not found`, name)
		}
	}

	tr := &TransformerStats2{
		accumulatorNameList:              accumulatorNameList,
		valueFieldNameList:               valueFieldNameList,
		groupByFieldNameList:             groupByFieldNameList,
		doVerbose:                        doVerbose,
		doIterativeStats:                 doIterativeStats,
		doHoldAndFit:                     doHoldAndFit,
		accumulatorFactory:               utils.NewStats2AccumulatorFactory(),
		namedAccumulators:                lib.NewOrderedMap[*lib.OrderedMap[*lib.OrderedMap[utils.IStats2Accumulator]]](),
		groupingKeysToGroupByFieldValues: lib.NewOrderedMap[[]*mlrval.Mlrval](),
		recordGroups:                     lib.NewOrderedMap[*[]*types.RecordAndContext](),
	}
	return tr, nil
}

// Given: accumulate corr,cov on values x,y group by a,b.
// Example input:       Example output:
//   a b x y            a b x_corr x_cov y_corr y_cov
//   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","y"] : {                  <--- value field names
//       "corr" : stats2_corr object,
//       "cov"  : stats2_cov  object
//     }
//   },
//   ["u","v"] : {
//     ["x","y"] : {
//       "corr" : stats2_corr object,
//       "cov"  : stats2_cov  object
//     }
//   },
//   ["u","w"] : {
//     ["x","y"] : {
//       "corr" : stats2_corr object,
//       "cov"  : stats2_cov  object
//     }
//   },
// }
//
// In the iterative case, add to the current record its current group's stats fields.
// In the non-iterative case, produce output only at the end of the input stream.

func (tr *TransformerStats2) Transform(
	inrecAndContext *types.RecordAndContext,
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
	inputDownstreamDoneChannel <-chan bool,
	outputDownstreamDoneChannel chan<- bool,
) error {
	HandleDefaultDownstreamDone(inputDownstreamDoneChannel, outputDownstreamDoneChannel)
	if !inrecAndContext.EndOfStream {

		if err := tr.ingest(inrecAndContext); err != nil {
			return err
		}

		if tr.doIterativeStats {
			// The input record is modified in this case, with new fields appended
			*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext)
		}
		// if tr.doHoldAndFit, the input record is held by the ingestor

	} else { // end of record stream
		if !tr.doIterativeStats { // in the iterative case, already emitted per-record
			if tr.doHoldAndFit {
				tr.fit(outputRecordsAndContexts)
			} else {
				tr.emit(outputRecordsAndContexts, &inrecAndContext.Context)
			}
		}
		*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext) // end-of-stream marker
	}
	return nil
}

func (tr *TransformerStats2) ingest(
	inrecAndContext *types.RecordAndContext,
) error {
	inrec := inrecAndContext.Record

	// E.g. if grouping by "a" and "b", and the current record has a=circle, b=blue,
	// then groupingKey is the string "circle,blue".
	groupingKey, groupByFieldValues, ok := inrec.GetSelectedValuesAndJoined(tr.groupByFieldNameList)
	if !ok {
		return nil
	}

	tr.groupingKeysToGroupByFieldValues.Put(groupingKey, groupByFieldValues)

	groupToValueFields := tr.namedAccumulators.Get(groupingKey)
	if groupToValueFields == nil {
		groupToValueFields = lib.NewOrderedMap[*lib.OrderedMap[utils.IStats2Accumulator]]()
		tr.namedAccumulators.Put(groupingKey, groupToValueFields)
	}

	if tr.doHoldAndFit { // Retain the input record in memory, for fitting and delivery at end of stream
		groupToRecords := tr.recordGroups.Get(groupingKey)
		if groupToRecords == nil {
			records := []*types.RecordAndContext{}
			groupToRecords = &records
			tr.recordGroups.Put(groupingKey, groupToRecords)
		}
		*groupToRecords = append(*groupToRecords, inrecAndContext)
	}

	// for [["x","y"]]
	n := len(tr.valueFieldNameList)
	for i := 0; i < n; i += 2 {
		valueFieldName1 := tr.valueFieldNameList[i]
		valueFieldName2 := tr.valueFieldNameList[i+1]

		key := valueFieldName1 + stats2KeySeparator + valueFieldName2

		valueFieldsToAccumulator := groupToValueFields.Get(key)
		if valueFieldsToAccumulator == nil {
			valueFieldsToAccumulator = lib.NewOrderedMap[utils.IStats2Accumulator]()
			groupToValueFields.Put(key, valueFieldsToAccumulator)
		}

		mval1 := inrec.Get(valueFieldName1)
		mval2 := inrec.Get(valueFieldName2)
		if mval1 == nil || mval2 == nil { // Key absent in current record
			continue
		}
		if mval1.IsVoid() || mval2.IsVoid() { // Key present in current record but with empty value
			continue
		}

		// for ["corr", "cov"]
		for _, accumulatorName := range tr.accumulatorNameList {
			accumulator := valueFieldsToAccumulator.Get(accumulatorName)
			if accumulator == nil {
				accumulator = tr.accumulatorFactory.Make(
					valueFieldName1,
					valueFieldName2,
					accumulatorName,
					tr.doVerbose,
				)
				if accumulator == nil {
					return cli.VerbErrorf(verbNameStats2, "accumulator creation failed")
				}
				valueFieldsToAccumulator.Put(accumulatorName, accumulator)
			}
			accumulator.Ingest(
				mval1.GetNumericToFloatValueOrDie(),
				mval2.GetNumericToFloatValueOrDie(),
			)
		}

		if tr.doIterativeStats {
			tr.populateRecord(
				inrecAndContext.Record,
				valueFieldName1,
				valueFieldName2,
				valueFieldsToAccumulator,
			)
		}
	}
	return nil
}

func (tr *TransformerStats2) emit(
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
	context *types.Context,
) {
	for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
		outrec := mlrval.NewMlrmapAsRecord()

		// Add in a=s,b=t fields:
		groupingKey := pa.Key
		groupByFieldValues := tr.groupingKeysToGroupByFieldValues.Get(groupingKey)
		for i, groupByFieldName := range tr.groupByFieldNameList {
			outrec.PutReference(groupByFieldName, groupByFieldValues[i].Copy())
		}

		// Add in fields such as x_y_corr, etc.
		groupToValueFields := tr.namedAccumulators.Get(groupingKey)

		// For "x","y"
		for pc := groupToValueFields.Head; pc != nil; pc = pc.Next {

			pairs := strings.Split(pc.Key, stats2KeySeparator)
			valueFieldName1 := pairs[0]
			valueFieldName2 := pairs[1]
			valueFieldsToAccumulator := pc.Value

			tr.populateRecord(outrec, valueFieldName1, valueFieldName2, valueFieldsToAccumulator)

			// For "corr", "linreg"
			for pd := valueFieldsToAccumulator.Head; pd != nil; pd = pd.Next {
				accumulator := pd.Value
				accumulator.Populate(valueFieldName1, valueFieldName2, outrec)
			}
		}

		*outputRecordsAndContexts = append(*outputRecordsAndContexts, types.NewRecordAndContext(outrec, context))
	}
}

func (tr *TransformerStats2) populateRecord(
	outrec *mlrval.Mlrmap,
	valueFieldName1 string,
	valueFieldName2 string,
	valueFieldsToAccumulator *lib.OrderedMap[utils.IStats2Accumulator],
) {
	// For "corr", "linreg"
	for pe := valueFieldsToAccumulator.Head; pe != nil; pe = pe.Next {
		accumulator := pe.Value
		accumulator.Populate(valueFieldName1, valueFieldName2, outrec)
	}
}

func (tr *TransformerStats2) fit(
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
) {
	for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
		groupingKey := pa.Key
		groupToValueFields := pa.Value
		recordsAndContexts := tr.recordGroups.Get(groupingKey)
		if recordsAndContexts == nil {
			continue
		}
		for _, recordAndContext := range *recordsAndContexts {
			record := recordAndContext.Record

			// For "x","y"
			for pb := groupToValueFields.Head; pb != nil; pb = pb.Next {
				pairs := strings.Split(pb.Key, stats2KeySeparator)
				valueFieldName1 := pairs[0]
				valueFieldName2 := pairs[1]
				valueFieldsToAccumulator := pb.Value

				// For "linreg-ols", "logireg"
				for pc := valueFieldsToAccumulator.Head; pc != nil; pc = pc.Next {
					accumulator := pc.Value

					// Note R2, cov, corr, etc have no non-trivial fit-function
					mval1 := record.Get(valueFieldName1)
					mval2 := record.Get(valueFieldName2)
					if mval1 != nil && mval2 != nil {
						accumulator.Fit(
							mval1.GetNumericToFloatValueOrDie(),
							mval2.GetNumericToFloatValueOrDie(),
							record,
						)
					}
				}
			}

			*outputRecordsAndContexts = append(*outputRecordsAndContexts, recordAndContext)
		}
		*recordsAndContexts = (*recordsAndContexts)[:0]
	}
}