miller/pkg/transformers/top.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 verbNameTop = "top"
const verbTopDefaultOutputFieldName = "top_idx"

var topOptions = []OptionSpec{
	{Flag: "-f", Arg: "{a,b,c}", Type: "csv-list", Desc: "Value-field names for top counts."},
	{Flag: "-g", Arg: "{d,e,f}", Type: "csv-list", Desc: "Optional group-by-field names for top counts."},
	{Flag: "-n", Arg: "{count}", Type: "int", Desc: "How many records to print per category; default 1."},
	{Flag: "-a", Type: "bool", Desc: "Print all fields for top-value records; default is to print only value and group-by fields. Requires a single value-field name only."},
	{Flag: "--max", Type: "bool", Desc: "Print top largest values. This is the default."},
	{Flag: "--min", Type: "bool", Desc: "Print top smallest values; default is top largest values."},
	{Flag: "-F", Type: "bool", Desc: "Keep top values as floats even if they look like integers (ignored in Miller 6, kept for backward compatibility)."},
	{Flag: "-o", Arg: "{name}", Type: "string", Desc: "Field name for output indices. Default \"top_idx\". Ignored if -a is used."},
}

var TopSetup = TransformerSetup{
	Verb:         verbNameTop,
	UsageFunc:    transformerTopUsage,
	ParseCLIFunc: transformerTopParseCLI,
	IgnoresInput: false,
	Options:      topOptions,
}

func transformerTopUsage(
	o *os.File,
) {
	argv0 := "mlr"
	verb := verbNameTop
	fmt.Fprintf(o, "Usage: %s %s [options]\n", argv0, verb)
	WriteVerbOptions(o, topOptions)
	fmt.Fprintf(o, "Prints the n records with smallest/largest values at specified fields,\n")
	fmt.Fprintf(o, "optionally by category. If -a is given, then the top records are emitted\n")
	fmt.Fprintf(o, "with the same fields as they appeared in the input. Without -a, only fields\n")
	fmt.Fprintf(o, "from -f, fields from -g, and the top-index field are emitted. For more information\n")
	fmt.Fprintf(o, "please see https://miller.readthedocs.io/en/latest/reference-verbs#top\n")
}

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

	// Parse local flags
	topCount := int64(1)
	var valueFieldNames []string = nil
	var groupByFieldNames []string = nil
	showFullRecords := false
	doMax := true
	outputFieldName := verbTopDefaultOutputFieldName

	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":
			transformerTopUsage(os.Stdout)
			return nil, cli.ErrHelpRequested

		case "-n":
			topCount, err = cli.VerbGetIntArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}
		case "-f":
			valueFieldNames, err = cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}
		case "-g":
			groupByFieldNames, err = cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}
		case "-a":
			showFullRecords = true
		case "--max":
			doMax = true
		case "--min":
			doMax = false
		case "-F":
			// Ignored in Miller 6; allowed for command-line backward compatibility
		case "-o":
			outputFieldName, err = cli.VerbGetStringArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}

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

	if valueFieldNames == nil {
		return nil, cli.VerbErrorf(verb, "-f field names required")
	}
	if len(valueFieldNames) > 1 && showFullRecords {
		return nil, cli.VerbErrorf(verb, "-a (full records) incompatible with multiple -f fields")
	}

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

	transformer, _ := NewTransformerTop(
		topCount,
		valueFieldNames,
		groupByFieldNames,
		showFullRecords,
		doMax,
		outputFieldName,
	)

	return transformer, nil
}

type TransformerTop struct {
	topCount          int64
	valueFieldNames   []string
	groupByFieldNames []string
	showFullRecords   bool
	doMax             bool
	outputFieldName   string

	// Two-level map from grouping key (string of joined-together group-by field values),
	// to string value-field name, to *utils.TopKeeper
	groups                           *lib.OrderedMap[*lib.OrderedMap[*utils.TopKeeper]]
	groupingKeysToGroupByFieldValues map[string][]*mlrval.Mlrval
}

func NewTransformerTop(
	topCount int64,
	valueFieldNames []string,
	groupByFieldNames []string,
	showFullRecords bool,
	doMax bool,
	outputFieldName string,
) (*TransformerTop, error) {

	tr := &TransformerTop{
		topCount:          topCount,
		valueFieldNames:   valueFieldNames,
		groupByFieldNames: groupByFieldNames,
		showFullRecords:   showFullRecords,
		doMax:             doMax,
		outputFieldName:   outputFieldName,

		groups:                           lib.NewOrderedMap[*lib.OrderedMap[*utils.TopKeeper]](),
		groupingKeysToGroupByFieldValues: make(map[string][]*mlrval.Mlrval),
	}

	return tr, nil
}

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

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

	// ["s", "t"]
	valueFieldValues, fok := inrec.ReferenceSelectedValues(tr.valueFieldNames)
	groupingKey, groupByFieldValues, gok := inrec.GetSelectedValuesAndJoined(tr.groupByFieldNames)

	// Heterogeneous-data case -- not all sought fields were present in record
	if !fok || !gok {
		return
	}
	iSecondLevel := tr.groups.Get(groupingKey)
	var secondLevel *lib.OrderedMap[*utils.TopKeeper] = nil
	if iSecondLevel == nil {
		secondLevel = lib.NewOrderedMap[*utils.TopKeeper]()
		tr.groups.Put(groupingKey, secondLevel)
		tr.groupingKeysToGroupByFieldValues[groupingKey] = groupByFieldValues
	} else {
		secondLevel = iSecondLevel
	}

	// for "x", "y" and "1", "2"
	for i := range tr.valueFieldNames {
		valueFieldName := tr.valueFieldNames[i]
		valueFieldValue := valueFieldValues[i]

		iTopKeeper := secondLevel.Get(valueFieldName)
		var topKeeper *utils.TopKeeper
		if iTopKeeper == nil {
			topKeeper = utils.NewTopKeeper(tr.topCount, tr.doMax)
			secondLevel.Put(valueFieldName, topKeeper)
		} else {
			topKeeper = iTopKeeper
		}

		var maybeRecordAndContext *types.RecordAndContext = nil
		if tr.showFullRecords {
			maybeRecordAndContext = inrecAndContext
		}
		topKeeper.Add(
			valueFieldValue,
			maybeRecordAndContext,
		)
	}

}

func (tr *TransformerTop) emit(
	inrecAndContext *types.RecordAndContext,
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
) {
	for pa := tr.groups.Head; pa != nil; pa = pa.Next {
		groupingKey := pa.Key
		secondLevel := pa.Value
		groupByFieldValues := tr.groupingKeysToGroupByFieldValues[groupingKey]

		// Above we required that there be only one value field in the
		// show-full-records case. That's because here, we print each record at most
		// once, which would need a change in the format presented as output.
		if tr.showFullRecords {
			for pb := secondLevel.Head; pb != nil; pb = pb.Next {
				topKeeper := pb.Value
				for i := int64(0); i < topKeeper.GetSize(); i++ {
					*outputRecordsAndContexts = append(*outputRecordsAndContexts, topKeeper.TopRecordsAndContexts[i].Copy())
				}
			}

		} else {

			for i := int64(0); i < tr.topCount; i++ {
				newrec := mlrval.NewMlrmapAsRecord()

				// Add in a=s,b=t fields:
				for j := range tr.groupByFieldNames {
					newrec.PutCopy(tr.groupByFieldNames[j], groupByFieldValues[j])
				}

				// Add in fields such as x_top_1=#
				// for "x", "y"
				for pb := secondLevel.Head; pb != nil; pb = pb.Next {
					valueFieldName := pb.Key
					topKeeper := pb.Value
					key := valueFieldName + "_top"
					if i < topKeeper.GetSize() {
						newrec.PutReference(tr.outputFieldName, mlrval.FromInt(i+1))
						newrec.PutReference(key, topKeeper.TopValues[i].Copy())
					} else {
						newrec.PutReference(tr.outputFieldName, mlrval.FromInt(i+1))
						newrec.PutCopy(key, mlrval.VOID)
					}
				}

				*outputRecordsAndContexts = append(*outputRecordsAndContexts, types.NewRecordAndContext(newrec, &inrecAndContext.Context))
			}
		}
	}

	*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext) // emit end-of-stream marker
}