github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/reshard.go (about)

     1  // Copyright 2019 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package bigslice
     6  
     7  import (
     8  	"fmt"
     9  
    10  	"github.com/grailbio/bigslice/slicefunc"
    11  	"github.com/grailbio/bigslice/sliceio"
    12  	"github.com/grailbio/bigslice/typecheck"
    13  )
    14  
    15  type reshardSlice struct {
    16  	name   Name
    17  	nshard int
    18  	Slice
    19  }
    20  
    21  // Reshard returns a slice that is resharded to the given
    22  // number of shards; this is done by re-shuffling to the
    23  // provided number of shards.
    24  func Reshard(slice Slice, nshard int) Slice {
    25  	if err := canMakeCombiningFrame(slice); err != nil {
    26  		typecheck.Panic(1, err.Error())
    27  	}
    28  	if slice.NumShard() == nshard {
    29  		return slice
    30  	}
    31  	return &reshardSlice{MakeName("reshard"), nshard, slice}
    32  }
    33  
    34  func (r *reshardSlice) Name() Name             { return r.name }
    35  func (*reshardSlice) NumDep() int              { return 1 }
    36  func (r *reshardSlice) NumShard() int          { return r.nshard }
    37  func (r *reshardSlice) Dep(i int) Dep          { return Dep{r.Slice, true, nil, false} }
    38  func (*reshardSlice) Combiner() slicefunc.Func { return slicefunc.Nil }
    39  
    40  func (r *reshardSlice) Reader(shard int, deps []sliceio.Reader) sliceio.Reader {
    41  	if len(deps) != 1 {
    42  		panic(fmt.Errorf("expected one dep, got %d", len(deps)))
    43  	}
    44  	return deps[0]
    45  }