Add new tail -n +N, head -n -N options (#2071)

Inspired by GNU head & tail, they match their behavior while supporting
the usual grouping operations.

Co-authored-by: John Kerl <kerl.john.r@gmail.com>

pkg/transformers/head.go

@@ -28,6 +28,8 @@ func transformerHeadUsage(
 	fmt.Fprintf(o, "Options:\n")
 	fmt.Fprintf(o, "-g {a,b,c} Optional group-by-field names for head counts, e.g. a,b,c.\n")
 	fmt.Fprintf(o, "-n {n} Head-count to print. Default 10.\n")
+	fmt.Fprintf(o, "           A negative count, e.g. -n -2, passes through all but the last n records,\n")
+	fmt.Fprintf(o, "           optionally by category.\n")
 	fmt.Fprintf(o, "-h|--help Show this message.\n")
 }
 
@@ -68,16 +70,6 @@ func transformerHeadParseCLI(
 			}
 			headCount = n
 
-			// This is a bit of a hack. In our Getoptify routine we preprocess
-			// the command line sending '-xyz' to '-x -y -z', but leaving
-			// '--xyz' as-is. Also, Unix-like tools often support 'head -n4'
-			// and 'tail -n4' in addition to 'head -n 4' and 'tail -n 4'.  Our
-			// getoptify paradigm, combined with syntax familiar to users,
-			// means we get '-n -4' here. So, take the absolute value to handle this.
-			if headCount < 0 {
-				headCount = -headCount
-			}
-
 		} else if opt == "-g" {
 			names, err := cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
 			if err != nil {
@@ -116,6 +108,7 @@ type TransformerHead struct {
 	recordTransformerFunc RecordTransformerFunc
 	unkeyedRecordCount    int64
 	keyedRecordCounts     map[string]int64
+	recordListsByGroup    map[string]*[]*types.RecordAndContext
 
 	// See ChainTransformer
 	wroteDownstreamDone bool
@@ -126,15 +119,23 @@ func NewTransformerHead(
 	groupByFieldNames []string,
 ) (*TransformerHead, error) {
 
+	allButLast := headCount < 0
+	if allButLast {
+		headCount = -headCount
+	}
+
 	tr := &TransformerHead{
 		headCount:           headCount,
 		groupByFieldNames:   groupByFieldNames,
 		unkeyedRecordCount:  0,
 		keyedRecordCounts:   make(map[string]int64),
+		recordListsByGroup:  make(map[string]*[]*types.RecordAndContext),
 		wroteDownstreamDone: false,
 	}
 
-	if groupByFieldNames == nil {
+	if allButLast {
+		tr.recordTransformerFunc = tr.transformAllButLast
+	} else if groupByFieldNames == nil {
 		tr.recordTransformerFunc = tr.transformUnkeyed
 	} else {
 		tr.recordTransformerFunc = tr.transformKeyed
@@ -207,3 +208,36 @@ func (tr *TransformerHead) transformKeyed(
 		*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext)
 	}
 }
+
+func (tr *TransformerHead) transformAllButLast(
+	inrecAndContext *types.RecordAndContext,
+	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
+	inputDownstreamDoneChannel <-chan bool,
+	outputDownstreamDoneChannel chan<- bool,
+) {
+	if !inrecAndContext.EndOfStream {
+		inrec := inrecAndContext.Record
+
+		groupingKey, ok := inrec.GetSelectedValuesJoined(tr.groupByFieldNames)
+		if !ok {
+			return
+		}
+
+		recordListForGroup := tr.recordListsByGroup[groupingKey]
+		if recordListForGroup == nil { // first time
+			recordListForGroup = &[]*types.RecordAndContext{}
+			tr.recordListsByGroup[groupingKey] = recordListForGroup
+		}
+
+		*recordListForGroup = append(*recordListForGroup, inrecAndContext)
+		for int64(len(*recordListForGroup)) > tr.headCount {
+			// Emit records that have fallen out of the window.
+			*outputRecordsAndContexts = append(*outputRecordsAndContexts, (*recordListForGroup)[0])
+			(*recordListForGroup)[0] = nil // release the backing-array slot's reference
+			*recordListForGroup = (*recordListForGroup)[1:]
+		}
+
+	} else {
+		*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext)
+	}
+}

pkg/transformers/tail.go

@@ -26,8 +26,11 @@ func transformerTailUsage(
 	fmt.Fprintln(o, "Passes through the last n records, optionally by category.")
 
 	fmt.Fprintf(o, "Options:\n")
-	fmt.Fprintf(o, "-g {a,b,c} Optional group-by-field names for head counts, e.g. a,b,c.\n")
-	fmt.Fprintf(o, "-n {n} Head-count to print. Default 10.\n")
+	fmt.Fprintf(o, "-g {a,b,c} Optional group-by-field names for tail counts, e.g. a,b,c.\n")
+	fmt.Fprintf(o, "-n {n} Tail-count to print. Default 10.\n")
+	fmt.Fprintf(o, "           A leading '+' means start at the nth record rather than print\n")
+	fmt.Fprintf(o, "           the last n: e.g. -n +3 passes through all but the first 2\n")
+	fmt.Fprintf(o, "           records, optionally by category.\n")
 	fmt.Fprintf(o, "-h|--help Show this message.\n")
 }
 
@@ -45,6 +48,7 @@ func transformerTailParseCLI(
 	argi++
 
 	tailCount := int64(10)
+	fromStart := false
 	var groupByFieldNames []string = nil
 
 	for argi < argc /* variable increment: 1 or 2 depending on flag */ {
@@ -62,22 +66,15 @@ func transformerTailParseCLI(
 			return nil, cli.ErrHelpRequested
 
 		} else if opt == "-n" {
+			if argi < argc && strings.HasPrefix(args[argi], "+") {
+				fromStart = true
+			}
 			n, err := cli.VerbGetIntArg(verb, opt, args, &argi, argc)
 			if err != nil {
 				return nil, err
 			}
 			tailCount = n
 
-			// This is a bit of a hack. In our Getoptify routine we preprocess
-			// the command line sending '-xyz' to '-x -y -z', but leaving
-			// '--xyz' as-is. Also, Unix-like tools often support 'head -n4'
-			// and 'tail -n4' in addition to 'head -n 4' and 'tail -n 4'.  Our
-			// getoptify paradigm, combined with syntax familiar to users,
-			// means we get '-n -4' here. So, take the absolute value to handle this.
-			if tailCount < 0 {
-				tailCount = -tailCount
-			}
-
 		} else if opt == "-g" {
 			names, err := cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
 			if err != nil {
@@ -98,6 +95,7 @@ func transformerTailParseCLI(
 
 	transformer, err := NewTransformerTail(
 		tailCount,
+		fromStart,
 		groupByFieldNames,
 	)
 	if err != nil {
@@ -109,24 +107,43 @@ func transformerTailParseCLI(
 
 type TransformerTail struct {
 	// input
-	tailCount         int64
+	count             int64 // last-n count, or skip-count in from-start mode
 	groupByFieldNames []string
 
 	// state
+	recordTransformerFunc RecordTransformerFunc
 	// map from string to record slices
 	recordListsByGroup *lib.OrderedMap[*[]*types.RecordAndContext]
+	// for the from-start mode: per-group counts of records seen so far
+	countsByGroup map[string]int64
 }
 
 func NewTransformerTail(
 	tailCount int64,
+	fromStart bool,
 	groupByFieldNames []string,
 ) (*TransformerTail, error) {
 
 	tr := &TransformerTail{
-		tailCount:         tailCount,
 		groupByFieldNames: groupByFieldNames,
 
 		recordListsByGroup: lib.NewOrderedMap[*[]*types.RecordAndContext](),
+		countsByGroup:      make(map[string]int64),
+	}
+
+	if fromStart {
+		// '-n +N' is a 1-based record index, i.e. skip the first n-1.
+		tr.count = tailCount - 1
+		if tr.count < 0 {
+			tr.count = 0
+		}
+		tr.recordTransformerFunc = tr.transformFromStart
+	} else {
+		if tailCount < 0 {
+			tailCount = -tailCount
+		}
+		tr.count = tailCount
+		tr.recordTransformerFunc = tr.transformLastN
 	}
 
 	return tr, nil
@@ -139,6 +156,15 @@ func (tr *TransformerTail) Transform(
 	outputDownstreamDoneChannel chan<- bool,
 ) {
 	HandleDefaultDownstreamDone(inputDownstreamDoneChannel, outputDownstreamDoneChannel)
+	tr.recordTransformerFunc(inrecAndContext, outputRecordsAndContexts, inputDownstreamDoneChannel, outputDownstreamDoneChannel)
+}
+
+func (tr *TransformerTail) transformLastN(
+	inrecAndContext *types.RecordAndContext,
+	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
+	inputDownstreamDoneChannel <-chan bool,
+	outputDownstreamDoneChannel chan<- bool,
+) {
 	if !inrecAndContext.EndOfStream {
 		inrec := inrecAndContext.Record
 
@@ -155,7 +181,8 @@ func (tr *TransformerTail) Transform(
 		}
 
 		*recordListForGroup = append(*recordListForGroup, inrecAndContext)
-		for int64(len(*recordListForGroup)) > tr.tailCount {
+		for int64(len(*recordListForGroup)) > tr.count {
+			(*recordListForGroup)[0] = nil // release the backing-array slot's reference
 			*recordListForGroup = (*recordListForGroup)[1:]
 		}
 
@@ -166,3 +193,28 @@ func (tr *TransformerTail) Transform(
 		*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext) // end-of-stream marker
 	}
 }
+
+func (tr *TransformerTail) transformFromStart(
+	inrecAndContext *types.RecordAndContext,
+	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
+	inputDownstreamDoneChannel <-chan bool,
+	outputDownstreamDoneChannel chan<- bool,
+) {
+	if !inrecAndContext.EndOfStream {
+		inrec := inrecAndContext.Record
+
+		groupingKey, ok := inrec.GetSelectedValuesJoined(tr.groupByFieldNames)
+		if !ok {
+			return
+		}
+
+		tr.countsByGroup[groupingKey]++
+		if tr.countsByGroup[groupingKey] > tr.count {
+			// Emit records now that we skipped the requested number of them.
+			*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext)
+		}
+
+	} else {
+		*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext) // end-of-stream marker
+	}
+}