github.com/weaviate/weaviate@v1.24.6/usecases/replica/batch.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package replica
    13  
    14  import (
    15  	"sort"
    16  
    17  	"github.com/go-openapi/strfmt"
    18  	"github.com/weaviate/weaviate/entities/storobj"
    19  	"github.com/weaviate/weaviate/usecases/objects"
    20  )
    21  
    22  // indexedBatch holds an indexed list of objects
    23  type indexedBatch struct {
    24  	Data []*storobj.Object
    25  	// Index is z-index used to maintain object's order
    26  	Index []int
    27  }
    28  
    29  // createBatch creates indexedBatch from xs
    30  func createBatch(xs []*storobj.Object) indexedBatch {
    31  	var bi indexedBatch
    32  	bi.Data = xs
    33  	bi.Index = make([]int, len(xs))
    34  	for i := 0; i < len(xs); i++ {
    35  		bi.Index[i] = i
    36  	}
    37  	return bi
    38  }
    39  
    40  // cluster data object by shard
    41  func cluster(bi indexedBatch) []shardPart {
    42  	index := bi.Index
    43  	data := bi.Data
    44  	sort.Slice(index, func(i, j int) bool {
    45  		return data[index[i]].BelongsToShard < data[index[j]].BelongsToShard
    46  	})
    47  	clusters := make([]shardPart, 0, 16)
    48  	// partition
    49  	cur := data[index[0]]
    50  	j := 0
    51  	for i := 1; i < len(index); i++ {
    52  		if data[index[i]].BelongsToShard == cur.BelongsToShard {
    53  			continue
    54  		}
    55  		clusters = append(clusters, shardPart{
    56  			Shard: cur.BelongsToShard,
    57  			Node:  cur.BelongsToNode, Data: data,
    58  			Index: index[j:i],
    59  		})
    60  		j = i
    61  		cur = data[index[j]]
    62  
    63  	}
    64  	clusters = append(clusters, shardPart{
    65  		Shard: cur.BelongsToShard,
    66  		Node:  cur.BelongsToNode, Data: data,
    67  		Index: index[j:],
    68  	})
    69  	return clusters
    70  }
    71  
    72  // shardPart represents a data partition belonging to a physical shard
    73  type shardPart struct {
    74  	Shard string // one-to-one mapping between Shard and Node
    75  	Node  string
    76  
    77  	Data  []*storobj.Object
    78  	Index []int // index for data
    79  }
    80  
    81  func (b *shardPart) ObjectIDs() []strfmt.UUID {
    82  	xs := make([]strfmt.UUID, len(b.Index))
    83  	for i, idx := range b.Index {
    84  		xs[i] = b.Data[idx].ID()
    85  	}
    86  	return xs
    87  }
    88  
    89  func (b *shardPart) Extract() ([]objects.Replica, []strfmt.UUID) {
    90  	xs := make([]objects.Replica, len(b.Index))
    91  	ys := make([]strfmt.UUID, len(b.Index))
    92  
    93  	for i, idx := range b.Index {
    94  		p := b.Data[idx]
    95  		xs[i] = objects.Replica{ID: p.ID(), Deleted: false, Object: p}
    96  		ys[i] = p.ID()
    97  	}
    98  	return xs, ys
    99  }