github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/ext/dsort/sort.go (about)

     1  // Package dsort provides APIs for distributed archive file shuffling.
     2  /*
     3   * Copyright (c) 2018-2024, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package dsort
     6  
     7  import (
     8  	"math/rand"
     9  	"sort"
    10  	"strconv"
    11  	"time"
    12  
    13  	"github.com/NVIDIA/aistore/cmn/debug"
    14  	"github.com/NVIDIA/aistore/ext/dsort/shard"
    15  )
    16  
    17  type (
    18  	alphaByKey struct {
    19  		err        error
    20  		records    *shard.Records
    21  		keyType    string
    22  		decreasing bool
    23  	}
    24  )
    25  
    26  // interface guard
    27  var _ sort.Interface = (*alphaByKey)(nil)
    28  
    29  func (s *alphaByKey) Len() int      { return s.records.Len() }
    30  func (s *alphaByKey) Swap(i, j int) { s.records.Swap(i, j) }
    31  
    32  func (s *alphaByKey) Less(i, j int) bool {
    33  	var (
    34  		err  error
    35  		less bool
    36  	)
    37  	if s.decreasing {
    38  		less, err = s.records.Less(j, i, s.keyType)
    39  	} else {
    40  		less, err = s.records.Less(i, j, s.keyType)
    41  	}
    42  	if err != nil {
    43  		s.err = err
    44  	}
    45  	return less
    46  }
    47  
    48  // sorts records by each Record.Key in the order determined by the `alg` algorithm.
    49  func sortRecords(r *shard.Records, alg *Algorithm) (err error) {
    50  	switch alg.Kind {
    51  	case None:
    52  		return nil
    53  	case Shuffle:
    54  		var (
    55  			rnd  *rand.Rand
    56  			seed = time.Now().Unix()
    57  		)
    58  		if alg.Seed != "" {
    59  			seed, err = strconv.ParseInt(alg.Seed, 10, 64)
    60  			debug.AssertNoErr(err)
    61  		}
    62  		rnd = rand.New(rand.NewSource(seed))
    63  		for i := range r.Len() { // https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle
    64  			j := rnd.Intn(i + 1)
    65  			r.Swap(i, j)
    66  		}
    67  	default:
    68  		keys := &alphaByKey{records: r, decreasing: alg.Decreasing, keyType: alg.ContentKeyType}
    69  		sort.Sort(keys)
    70  		err = keys.err
    71  	}
    72  	return
    73  }