@@ -8,6 +8,7 @@ package algo
88import (
99 "container/heap"
1010 "sort"
11+ "sync"
1112
1213 "github.com/hypermodeinc/dgraph/v25/codec"
1314 "github.com/hypermodeinc/dgraph/v25/protos/pb"
@@ -360,41 +361,61 @@ func Difference(u, v *pb.List) *pb.List {
360361 return & pb.List {Uids : out }
361362}
362363
363- // MergeSorted merges sorted lists.
364- func MergeSorted (lists []* pb.List ) * pb.List {
364+ func MergeSortedMoreMem (lists []* pb.List ) * pb.List {
365+ numThreads := 10
366+ if len (lists ) > numThreads * numThreads {
367+ k := numThreads
368+ res := []* pb.List {}
369+ var wg sync.WaitGroup
370+ var mutex sync.Mutex
371+ wg .Add (k )
372+ for i := 0 ; i < k ; i ++ {
373+ go func () {
374+ defer wg .Done ()
375+ end := (i + 1 ) * len (lists ) / k
376+ if end > len (lists ) {
377+ end = len (lists )
378+ }
379+ result := MergeSortedMoreMem (lists [i * len (lists )/ k : end ])
380+ mutex .Lock ()
381+ res = append (res , result )
382+ mutex .Unlock ()
383+ }()
384+ }
385+ wg .Wait ()
386+ return MergeSortedMoreMem (res )
387+ } else {
388+ return internalMergeSort (lists )
389+ }
390+ }
391+
392+ func internalMergeSortWithBuffer (lists []* pb.List , buffer []uint64 ) * pb.List {
365393 if len (lists ) == 0 {
366- return new ( pb.List )
394+ return & pb.List { Uids : buffer [: 0 ]}
367395 }
368396
369397 h := & uint64Heap {}
370398 heap .Init (h )
371- maxSz := 0
372399
373400 for i , l := range lists {
374- if l == nil {
401+ if l == nil || len ( l . Uids ) == 0 {
375402 continue
376403 }
377- lenList := len (l .Uids )
378- if lenList > 0 {
379- heap .Push (h , elem {
380- val : l .Uids [0 ],
381- listIdx : i ,
382- })
383- if lenList > maxSz {
384- maxSz = lenList
385- }
386- }
404+ heap .Push (h , elem {
405+ val : l .Uids [0 ],
406+ listIdx : i ,
407+ })
387408 }
388409
389- // Our final output. Give it an approximate capacity as copies are expensive.
390- output := make ([]uint64 , 0 , maxSz )
391- // idx[i] is the element we are looking at for lists[i].
410+ // Use the provided buffer
411+ output := buffer [:0 ]
392412 idx := make ([]int , len (lists ))
393- var last uint64 // Last element added to sorted / final output.
394- for h .Len () > 0 { // While heap is not empty.
395- me := (* h )[0 ] // Peek at the top element in heap.
413+ var last uint64
414+
415+ for h .Len () > 0 {
416+ me := (* h )[0 ]
396417 if len (output ) == 0 || me .val != last {
397- output = append (output , me .val ) // Add if unique.
418+ output = append (output , me .val )
398419 last = me .val
399420 }
400421 l := lists [me .listIdx ]
@@ -404,12 +425,122 @@ func MergeSorted(lists []*pb.List) *pb.List {
404425 idx [me .listIdx ]++
405426 val := l.Uids [idx [me.listIdx ]]
406427 (* h )[0 ].val = val
407- heap .Fix (h , 0 ) // Faster than Pop() followed by Push().
428+ heap .Fix (h , 0 )
408429 }
409430 }
431+
410432 return & pb.List {Uids : output }
411433}
412434
435+ // MergeSorted merges sorted lists.
436+ func internalMergeSort (lists []* pb.List ) * pb.List {
437+ sz := 0
438+ for _ , l := range lists {
439+ if l == nil || len (l .Uids ) == 0 {
440+ continue
441+ }
442+ sz += len (l .Uids )
443+ }
444+ buffer := make ([]uint64 , 0 , sz )
445+ return internalMergeSortWithBuffer (lists , buffer )
446+ }
447+
448+ func MergeSorted (lists []* pb.List ) * pb.List {
449+ // Calculate total capacity needed
450+ totalCap := 0
451+ for _ , l := range lists {
452+ if l != nil {
453+ totalCap += len (l .Uids )
454+ }
455+ }
456+
457+ // Pre-allocate one big buffer
458+ bigBuffer := make ([]uint64 , totalCap )
459+ return mergeSortedWithBuffer (lists , bigBuffer )
460+ }
461+
462+ const numThreads = 10
463+ const numListPerThread = 10
464+
465+ func mergeSortedWithBuffer (lists []* pb.List , buffer []uint64 ) * pb.List {
466+ if len (lists ) < numThreads * numListPerThread {
467+ return internalMergeSort (lists )
468+ }
469+
470+ // Calculate how much buffer each goroutine needs
471+ chunkSizes := make ([]int , numThreads )
472+ totalNeeded := 0
473+
474+ chunkSize := (len (lists ) + numThreads - 1 ) / numThreads
475+
476+ for i := 0 ; i < numThreads ; i ++ {
477+ start := i * chunkSize
478+ end := (i + 1 ) * chunkSize
479+ if end > len (lists ) {
480+ end = len (lists )
481+ }
482+
483+ if start > len (lists ) {
484+ continue
485+ }
486+
487+ chunkCap := 0
488+ for j := start ; j < end ; j ++ {
489+ if lists [j ] != nil {
490+ chunkCap += len (lists [j ].Uids )
491+ }
492+ }
493+ chunkSizes [i ] = chunkCap
494+ totalNeeded += chunkCap
495+ }
496+
497+ // Calculate buffer offsets for each goroutine
498+ bufferOffsets := make ([]int , numThreads )
499+ bufferOffsets [0 ] = 0
500+ for i := 1 ; i < numThreads ; i ++ {
501+ bufferOffsets [i ] = bufferOffsets [i - 1 ] + chunkSizes [i - 1 ]
502+ }
503+
504+ // Distribute buffer slices to each goroutine
505+ intermediateResults := make ([]* pb.List , numThreads )
506+ var wg sync.WaitGroup
507+
508+ wg .Add (numThreads )
509+ for i := 0 ; i < numThreads ; i ++ {
510+ go func (idx int ) {
511+ defer wg .Done ()
512+ start := idx * chunkSize
513+ end := (idx + 1 ) * chunkSize
514+ if end > len (lists ) {
515+ end = len (lists )
516+ }
517+ if start > len (lists ) {
518+ return
519+ }
520+
521+ // Give this goroutine its slice of the big buffer
522+ bufferStart := bufferOffsets [idx ]
523+ bufferEnd := bufferStart + chunkSizes [idx ]
524+ goroutineBuffer := buffer [bufferStart :bufferEnd :bufferEnd ][:0 ]
525+ result := internalMergeSortWithBuffer (lists [start :end ], goroutineBuffer )
526+ intermediateResults [idx ] = result
527+ }(i )
528+ }
529+ wg .Wait ()
530+
531+ // Filter out nil results
532+ validResults := make ([]* pb.List , 0 , numThreads )
533+ for _ , result := range intermediateResults {
534+ if result != nil && len (result .Uids ) > 0 {
535+ validResults = append (validResults , result )
536+ }
537+ }
538+
539+ // Use the remaining part of buffer for final merge
540+ finalBuffer := make ([]uint64 , 0 , totalNeeded )
541+ return internalMergeSortWithBuffer (validResults , finalBuffer )
542+ }
543+
413544// IndexOf performs a binary search on the uids slice and returns the index at
414545// which it finds the uid, else returns -1
415546func IndexOf (u * pb.List , uid uint64 ) int {
0 commit comments