1package fileio23import (4 "bufio"5 "container/heap"6 "errors"7 "io"8 "os"9 "slices"10)1112// An external merge sort of records of any kind in bounded memory: runs13// of at most sortRunBytes sorted in memory and spilled to files, then14// merged, read back one record at a time. A Parquet view's sort orders15// its rows with it (sourcesort.go), and the order statistics of a16// source's column its numbers (numsort.go).1718// sortRunBytes bounds the records a sort holds in memory at once.19var sortRunBytes = 64 << 202021// recCodec is how a sort compares, sizes and stores its records.22type recCodec[T any] struct {23 cmp func(a, b T) int24 size func(*T) int // about what a record holds in memory25 // fixed is every record's size when they're all one size, so the26 // sort can grow its run to what one holds and no further.27 fixed int28 put func(*bufio.Writer, *T)29 // get reads a record into the one given, io.EOF at a run's end.30 get func(*bufio.Reader, *T) error31}3233// extSort sorts records in bounded memory.34type extSort[T any] struct {35 dir string36 codec recCodec[T]37 recs []T38 bytes int39 runs []string40 limit int41 n int6442}4344func newExtSort[T any](dir string, codec recCodec[T]) *extSort[T] {45 return &extSort[T]{dir: dir, codec: codec, limit: sortRunBytes}46}4748// add takes a record, spilling a run when memory is full.49func (s *extSort[T]) add(r T) error {50 if f := s.codec.fixed; f > 0 && len(s.recs) == cap(s.recs) {51 // Grow no further than a run holds: doubling past it would52 // hold twice the run while copying.53 most := max(s.limit/f, 1)54 grown := make([]T, len(s.recs), len(s.recs)+max(min(max(cap(s.recs), 1024), most-len(s.recs)), 1))55 copy(grown, s.recs)56 s.recs = grown57 }58 s.recs = append(s.recs, r)59 s.bytes += s.codec.size(&s.recs[len(s.recs)-1])60 s.n++61 if s.bytes >= s.limit {62 return s.spill()63 }64 return nil65}6667func (s *extSort[T]) sort() {68 slices.SortFunc(s.recs, s.codec.cmp)69}7071// spill sorts the records in memory and writes them as a run.72func (s *extSort[T]) spill() error {73 s.sort()74 f, err := os.CreateTemp(s.dir, "012-sort-*.run")75 if err != nil {76 return err77 }78 s.runs = append(s.runs, f.Name())79 w := bufio.NewWriterSize(f, 1<<20)80 for i := range s.recs {81 s.codec.put(w, &s.recs[i])82 }83 err = errors.Join(w.Flush(), f.Close())84 clear(s.recs)85 s.recs, s.bytes = s.recs[:0], 086 return err87}8889// sorted is the records in order, read once: those held when none were90// spilled, else the runs merged. Closing it removes the runs.91func (s *extSort[T]) sorted() (*sortedRecs[T], error) {92 out := &sortedRecs[T]{s: s}93 if len(s.runs) == 0 {94 s.sort()95 return out, nil96 }97 if len(s.recs) > 0 {98 if err := s.spill(); err != nil {99 s.remove()100 return nil, err101 }102 }103 out.h = &runHeap[T]{cmp: s.codec.cmp}104 for _, name := range s.runs {105 f, err := os.Open(name)106 if err != nil {107 out.close()108 return nil, err109 }110 r := &runReader[T]{f: f, r: bufio.NewReaderSize(f, 256<<10), get: s.codec.get}111 if r.next() {112 out.h.runs = append(out.h.runs, r)113 } else {114 f.Close()115 out.h.done = append(out.h.done, r)116 }117 }118 heap.Init(out.h)119 return out, nil120}121122// remove deletes the runs.123func (s *extSort[T]) remove() {124 for _, name := range s.runs {125 os.Remove(name)126 }127 s.runs = nil128}129130// sortedRecs reads a sort's records in order.131type sortedRecs[T any] struct {132 s *extSort[T]133 h *runHeap[T] // nil when the records are held134 i int135 last *runReader[T] // the run the record handed out last came from136}137138// next is the next record, valid until the next call; nil at the end139// or on an error, which err says.140func (o *sortedRecs[T]) next() *T {141 if o.h == nil {142 if o.i >= len(o.s.recs) {143 return nil144 }145 o.i++146 return &o.s.recs[o.i-1]147 }148 if o.last != nil { // move past the record handed out last149 if o.last.next() {150 heap.Fix(o.h, 0)151 } else {152 o.last.f.Close()153 heap.Pop(o.h)154 }155 o.last = nil156 }157 if o.h.Len() == 0 {158 return nil159 }160 o.last = o.h.runs[0]161 return &o.last.rec162}163164// err is the first error a run met.165func (o *sortedRecs[T]) err() error {166 if o.h == nil {167 return nil168 }169 return o.h.err()170}171172// close lets go of the runs.173func (o *sortedRecs[T]) close() {174 if o.h != nil {175 o.h.close()176 }177 o.s.remove()178}179180// runReader reads a run's records in order.181type runReader[T any] struct {182 f *os.File183 r *bufio.Reader184 get func(*bufio.Reader, *T) error185 rec T186 fail error187}188189// next reads the next record, false at the end or on an error.190func (r *runReader[T]) next() bool {191 if err := r.get(r.r, &r.rec); err != nil {192 if !errors.Is(err, io.EOF) {193 r.fail = err194 }195 return false196 }197 return true198}199200// runHeap merges runs by their next record.201type runHeap[T any] struct {202 runs []*runReader[T]203 cmp func(a, b T) int204 done []*runReader[T]205}206207func (h *runHeap[T]) Len() int { return len(h.runs) }208func (h *runHeap[T]) Less(i, j int) bool { return h.cmp(h.runs[i].rec, h.runs[j].rec) < 0 }209func (h *runHeap[T]) Swap(i, j int) { h.runs[i], h.runs[j] = h.runs[j], h.runs[i] }210func (h *runHeap[T]) Push(x any) { h.runs = append(h.runs, x.(*runReader[T])) }211func (h *runHeap[T]) Pop() any {212 r := h.runs[len(h.runs)-1]213 h.runs = h.runs[:len(h.runs)-1]214 h.done = append(h.done, r)215 return r216}217218// err is the first error a run met.219func (h *runHeap[T]) err() error {220 for _, r := range append(h.done, h.runs...) {221 if r.fail != nil {222 return r.fail223 }224 }225 return nil226}227228func (h *runHeap[T]) close() {229 for _, r := range h.runs {230 r.f.Close()231 }232 h.runs = nil233}