diff --git a/api.go b/api.go index b5d2f83c4..56a226698 100644 --- a/api.go +++ b/api.go @@ -1394,6 +1394,7 @@ type ImportOptions struct { Clear bool IgnoreKeyCheck bool Presorted bool + fullySorted bool // format-aware sorting, internal use only please. // test Tx atomicity if > 0 SimPowerLossAfter int @@ -1427,6 +1428,13 @@ func OptImportOptionsPresorted(b bool) ImportOption { } } +func optImportOptionsFullySorted(b bool) ImportOption { + return func(o *ImportOptions) error { + o.fullySorted = b + return nil + } +} + var ErrAborted = fmt.Errorf("error: update was aborted") func (api *API) ImportAtomicRecord(ctx context.Context, qcx *Qcx, req *AtomicRecord, opts ...ImportOption) error { @@ -1905,8 +1913,8 @@ func (api *API) IngestOperations(ctx context.Context, qcx *Qcx, indexName string func (api *API) applyOperations(ctx context.Context, qcx *Qcx, index *Index, shard uint64, fields map[string]*Field, ops []*ingest.Operation) error { // For each operation, we may have a set of records/fields to clear, and then // also a set of fields to set/remove specific bits in. - opts := &ImportOptions{Presorted: true, IgnoreKeyCheck: true} - funcOpts := []ImportOption{OptImportOptionsIgnoreKeyCheck(true), OptImportOptionsPresorted(true), OptImportOptionsClear(true)} + opts := &ImportOptions{Presorted: true, IgnoreKeyCheck: true, fullySorted: true} + funcOpts := []ImportOption{OptImportOptionsIgnoreKeyCheck(true), OptImportOptionsPresorted(true), optImportOptionsFullySorted(true), OptImportOptionsClear(true)} for _, op := range ops { // ClearRecordIDs should exist only for delete, clear, and write. For clear and write, // we'll have a list of fields, for delete, it should be all the fields. @@ -1957,9 +1965,9 @@ func (api *API) applyOperations(ctx context.Context, qcx *Qcx, index *Index, sha opts.Clear = (op.OpType == ingest.OpRemove) // reslice this rather than regenerating it. i'm so efficient. if opts.Clear { - funcOpts = funcOpts[:3] + funcOpts = funcOpts[:4] } else { - funcOpts = funcOpts[:2] + funcOpts = funcOpts[:3] } // for "set" and "write" ops, we'll be setting bits, for // "remove" ops we'll be clearing them, and for "clear" ops diff --git a/fragment.go b/fragment.go index cf4c1c9f1..23f312627 100644 --- a/fragment.go +++ b/fragment.go @@ -2210,7 +2210,7 @@ func (f *fragment) bulkImport(tx Tx, rowIDs, columnIDs []uint64, options *Import } if f.mutexVector != nil && !options.Clear { - return f.bulkImportMutex(tx, rowIDs, columnIDs) + return f.bulkImportMutex(tx, rowIDs, columnIDs, options) } return f.bulkImportStandard(tx, rowIDs, columnIDs, options) } @@ -2253,8 +2253,12 @@ func (f *fragment) bulkImportStandard(tx Tx, rowIDs, columnIDs []uint64, options rowSet := make(map[uint64]struct{}) lastRowID := uint64(1 << 63) + // It's possible for the ingest API to have already sorted things in + // the row-first order we want for this import. + if !options.fullySorted { + sort.Sort(rowColumnSet{r: rowIDs, c: columnIDs}) + } // replace columnIDs with calculated positions to avoid allocation. - sort.Sort(rowColumnSet{r: rowIDs, c: columnIDs}) prevRow, prevCol := ^uint64(0), ^uint64(0) next := 0 for i := 0; i < len(columnIDs); i++ { @@ -2525,14 +2529,20 @@ func sliceDifference(original, remove []uint64) []uint64 { // mutex restrictions. Because the mutex requirements must be checked // against storage, this method must acquire a write lock on the fragment // during the entire process, and it handles every bit independently. -func (f *fragment) bulkImportMutex(tx Tx, rowIDs, columnIDs []uint64) error { +func (f *fragment) bulkImportMutex(tx Tx, rowIDs, columnIDs []uint64, options *ImportOptions) error { f.mu.Lock() defer f.mu.Unlock() + // if ingest promises that this is "fully sorted", then we have been + // promised that (1) there's no duplicate entries that need to be + // pruned, (2) the input is sorted by row IDs and then column IDs, + // meaning that we will generate positions in strictly sequential order. + if !options.fullySorted { p := parallelSlices{cols: columnIDs, rows: rowIDs} p.fullPrune() columnIDs = p.cols rowIDs = p.rows + } // create a mask of columns we care about columns := roaring.NewSliceBitmap(columnIDs...) @@ -2551,6 +2561,10 @@ func (f *fragment) bulkImportMutex(tx Tx, rowIDs, columnIDs []uint64) error { // positions are sorted by columns, but not by absolute // position. we might want them sorted, though. if pos < prev { + if options.fullySorted { + fmt.Printf("HELP! was promised fully sorted input, but previous position was %d, now generated %d\n", + prev, pos) + } unsorted = true } prev = pos diff --git a/ingest/codec.go b/ingest/codec.go index a0ab071cc..108baada1 100644 --- a/ingest/codec.go +++ b/ingest/codec.go @@ -76,6 +76,7 @@ type jsonFieldCodec struct { epoch int64 // used only by Timestamp fields scratch []uint64 // reusable scratch space for sets of values lookup KeyLookupFunc + fieldType FieldType } // JSONCodec is a Codec which accepts a JSON map of record keys/ids to updated value maps. @@ -103,7 +104,7 @@ func (codec *JSONCodec) AddTimeQuantumField(name string, lookup KeyLookupFunc) e if lookup != nil { fieldCodec.translate = (*FieldOperation).TranslateUnsigned } - return codec.addField(name, fieldCodec, lookup) + return codec.addField(name, FieldTypeTimeQuantum, fieldCodec, lookup) } func (codec *JSONCodec) AddSetField(name string, lookup KeyLookupFunc) error { @@ -112,7 +113,7 @@ func (codec *JSONCodec) AddSetField(name string, lookup KeyLookupFunc) error { if lookup != nil { fieldCodec.translate = (*FieldOperation).TranslateUnsigned } - return codec.addField(name, fieldCodec, lookup) + return codec.addField(name, FieldTypeSet, fieldCodec, lookup) } func (codec *JSONCodec) AddIntField(name string, lookup KeyLookupFunc) error { @@ -121,7 +122,7 @@ func (codec *JSONCodec) AddIntField(name string, lookup KeyLookupFunc) error { if lookup != nil { fieldCodec.translate = (*FieldOperation).TranslateSigned } - return codec.addField(name, fieldCodec, lookup) + return codec.addField(name, FieldTypeInt, fieldCodec, lookup) } func (codec *JSONCodec) AddMutexField(name string, lookup KeyLookupFunc) error { @@ -130,13 +131,13 @@ func (codec *JSONCodec) AddMutexField(name string, lookup KeyLookupFunc) error { if lookup != nil { fieldCodec.translate = (*FieldOperation).TranslateUnsigned } - return codec.addField(name, fieldCodec, lookup) + return codec.addField(name, FieldTypeMutex, fieldCodec, lookup) } func (codec *JSONCodec) AddBoolField(name string) error { fieldCodec := &jsonFieldCodec{} fieldCodec.decode = fieldCodec.DecodeBoolValue - return codec.addField(name, fieldCodec, nil) + return codec.addField(name, FieldTypeBool, fieldCodec, nil) } // TimestampField is used to store seconds since unix epoch. The numeric values @@ -146,7 +147,7 @@ func (codec *JSONCodec) AddBoolField(name string) error { func (codec *JSONCodec) AddTimestampField(name string, timeScale time.Duration, epoch int64) error { fieldCodec := &jsonFieldCodec{scaleUnit: int64(timeScale), epoch: epoch} fieldCodec.decode = fieldCodec.DecodeTimeValue - return codec.addField(name, fieldCodec, nil) + return codec.addField(name, FieldTypeTimeStamp, fieldCodec, nil) } // AddDecimalField adds a decimal field, which is stored as integer values @@ -155,10 +156,10 @@ func (codec *JSONCodec) AddTimestampField(name string, timeScale time.Duration, func (codec *JSONCodec) AddDecimalField(name string, decimalScale int64) error { fieldCodec := &jsonFieldCodec{scaleUnit: int64(math.Pow10(int(decimalScale)))} fieldCodec.decode = fieldCodec.DecodeDecimalValue - return codec.addField(name, fieldCodec, nil) + return codec.addField(name, FieldTypeDecimal, fieldCodec, nil) } -func (codec *JSONCodec) addField(name string, fieldCodec *jsonFieldCodec, lookup KeyLookupFunc) error { +func (codec *JSONCodec) addField(name string, fieldType FieldType, fieldCodec *jsonFieldCodec, lookup KeyLookupFunc) error { if _, ok := codec.fields[name]; ok { return fmt.Errorf("duplicate field %q", name) } @@ -166,6 +167,7 @@ func (codec *JSONCodec) addField(name string, fieldCodec *jsonFieldCodec, lookup fieldCodec.valueKeys = NewStringTable() fieldCodec.lookup = lookup } + fieldCodec.fieldType = fieldType codec.fields[name] = fieldCodec return nil } @@ -383,11 +385,16 @@ func (j *jsonFieldCodec) DecodeDecimalValue(recID uint64, dataType jsonparser.Va } func (codec *JSONCodec) ParseKeyedRecords(data []byte) (err error) { + seen := make(map[uint64]struct{}) return jsonparser.ObjectEach(data, func(key []byte, value []byte, dataType jsonparser.ValueType, offset int) error { id, err := codec.recKeys.ID(key) if err != nil { return err } + if _, ok := seen[id]; ok { + return fmt.Errorf("key %q duplicated in input", key) + } + seen[id] = struct{}{} if codec.currentOp.OpType == OpWrite { codec.currentOp.ClearRecordIDs = append(codec.currentOp.ClearRecordIDs, id) } @@ -463,6 +470,8 @@ func (codec *JSONCodec) ParseOperation(data []byte) (op *Operation, err error) { return op, err } +// Parse reads a request, but does not sort the results at all or divide +// them into shards. func (codec *JSONCodec) Parse(r io.Reader) (req *Request, err error) { data, err := ioutil.ReadAll(r) if err != nil { @@ -501,6 +510,7 @@ func (codec *JSONCodec) ParseBytes(data []byte) (req *Request, err error) { return nil, fmt.Errorf("trying to find record key mapping: %w", err) } } + req = &Request{FieldTypes: make(map[string]FieldType, len(codec.fields))} valueMaps := map[string]func(*FieldOperation) error{} for name, fieldCodec := range codec.fields { // make closure survive iteration @@ -514,6 +524,7 @@ func (codec *JSONCodec) ParseBytes(data []byte) (req *Request, err error) { return fieldCodec.translate(fo, fieldMap) } } + req.FieldTypes[name] = fieldCodec.fieldType } for _, op := range ops { // For Clear and Write, we need to translate/sort our record ID @@ -524,7 +535,6 @@ func (codec *JSONCodec) ParseBytes(data []byte) (req *Request, err error) { return nil, fmt.Errorf("mapping record keys for clear op: %w", err) } } - op.Sort() } // For clear/delete, that's all we need to do; there's no meaningful fieldops under them. if op.OpType == OpClear || op.OpType == OpDelete { @@ -542,14 +552,14 @@ func (codec *JSONCodec) ParseBytes(data []byte) (req *Request, err error) { } } // Sort by column keys, for now. - fieldOp.Sort() // Write op will also want to clear every field we saw. if op.OpType == OpWrite { op.ClearFields = append(op.ClearFields, field) } } } - return &Request{Ops: ops}, nil + req.Ops = ops + return req, nil } type errFieldNotFound struct { diff --git a/ingest/op.go b/ingest/op.go index 1006de992..3a1382e0e 100644 --- a/ingest/op.go +++ b/ingest/op.go @@ -16,7 +16,7 @@ package ingest import ( "fmt" - "sort" + "math/bits" "github.com/molecula/featurebase/v2/shardwidth" ) @@ -41,6 +41,28 @@ var opNames = []string{ "delete", } +type FieldType string + +const ( + FieldTypeSet = "set" + FieldTypeInt = "int" + FieldTypeTimeQuantum = "time" + FieldTypeTimeStamp = "timestamp" + FieldTypeDecimal = "decimal" + FieldTypeMutex = "mutex" + FieldTypeBool = "bool" +) + +var fieldTypeSorts = map[FieldType]func(*FieldOperation){ + FieldTypeSet: (*FieldOperation).SortByValues, + FieldTypeInt: (*FieldOperation).SortByRecords, + FieldTypeTimeQuantum: (*FieldOperation).SortByValues, + FieldTypeDecimal: (*FieldOperation).SortByRecords, + FieldTypeMutex: (*FieldOperation).SortByValues, + FieldTypeTimeStamp: (*FieldOperation).SortByRecords, + FieldTypeBool: (*FieldOperation).SortByValues, +} + func (o OpType) String() string { if int(o) < len(opNames) { return opNames[o] @@ -92,69 +114,31 @@ type FieldOperation struct { Signed []int64 } -var _ sort.Interface = &FieldOperation{} - -// We implement sort.Interface here so we don't have to write sort code. -// This should probably be replaced by smarter sorting later if it's a -// performance issue. -func (f *FieldOperation) Len() int { - return len(f.RecordIDs) -} - -func (f *FieldOperation) Less(i, j int) bool { - return f.RecordIDs[i] < f.RecordIDs[j] -} - -func (f *FieldOperation) Swap(i, j int) { - f.RecordIDs[i], f.RecordIDs[j] = f.RecordIDs[j], f.RecordIDs[i] - if f.Values != nil { - f.Values[i], f.Values[j] = f.Values[j], f.Values[i] - } - if f.Signed != nil { - f.Signed[i], f.Signed[j] = f.Signed[j], f.Signed[i] - } -} - -// Sort sorts the values by record ID. It is not a stable sort. -func (f *FieldOperation) Sort() { - sort.Sort(f) -} - -// Similarly, do this for Operation, which is used only in the Clear case. -var _ sort.Interface = &Operation{} - -// We implement sort.Interface here so we don't have to write sort code. -// This should probably be replaced by smarter sorting later if it's a -// performance issue. -func (o *Operation) Len() int { - return len(o.ClearRecordIDs) -} - -func (o *Operation) Less(i, j int) bool { - return o.ClearRecordIDs[i] < o.ClearRecordIDs[j] -} - -func (o *Operation) Swap(i, j int) { - o.ClearRecordIDs[i], o.ClearRecordIDs[j] = o.ClearRecordIDs[j], o.ClearRecordIDs[i] -} - // Sort sorts the values by record ID. It is not a stable sort. func (o *Operation) Sort() { - sort.Sort(o) + // I am aware that this is a crime, but it avoids rewriting + // the code and justifies FieldOperation handling the "only record + // IDs" case. + f := FieldOperation{RecordIDs: o.ClearRecordIDs} + f.SortByRecords() } type ShardedFieldOperation map[uint64]*FieldOperation // Shard() divides the FieldOperation's values up into corresponding chunks -// based on the shards of record IDs. Record IDs should be sorted before this -// happens. +// based on the shards of record IDs. Does not further sort IDs within those +// chunks. func (f *FieldOperation) Shard() ShardedFieldOperation { if len(f.RecordIDs) == 0 { return nil } + return f.SortToShards() +} + +// ShardInto puts the shards it finds into a target map. +func (f *FieldOperation) ShardInto(target ShardedFieldOperation) { shards, ends := shardwidth.FindShards(f.RecordIDs) prev := 0 - op := make(ShardedFieldOperation, len(shards)) for i, shard := range shards { endIndex := ends[i] subOp := &FieldOperation{RecordIDs: f.RecordIDs[prev:endIndex]} @@ -164,10 +148,311 @@ func (f *FieldOperation) Shard() ShardedFieldOperation { if len(f.Signed) > 0 { subOp.Signed = f.Signed[prev:endIndex] } - op[shard] = subOp + target[shard] = subOp prev = endIndex } - return op +} + +func ShardIDs(ids []uint64) (out map[uint64][]uint64) { + shards, ends := shardwidth.FindShards(ids) + prev := 0 + out = make(map[uint64][]uint64, len(shards)) + for i, shard := range shards { + endIndex := ends[i] + out[shard] = ids[prev:endIndex] + prev = endIndex + } + return out +} + +// SortToShards() uses a pseudo-radix-sort to divide inputs into +// shards; the individual shards are not sorted. +func (f *FieldOperation) SortToShards() ShardedFieldOperation { + if len(f.RecordIDs) == 0 { + return nil + } + diffMask := uint64(0) + prev := f.RecordIDs[0] + for _, r := range f.RecordIDs[1:] { + diffMask |= r ^ prev + prev = r + } + bitsRemaining := bits.Len64(diffMask) + if bitsRemaining <= shardwidth.Exponent { + return map[uint64]*FieldOperation{f.RecordIDs[0] >> shardwidth.Exponent: f} + } + output := make(ShardedFieldOperation) + sortToShardsInto(f, bitsRemaining-8, output) + return output +} + +// sortToShardsInto puts the shards it finds into the given map, so that +// as we split off buckets, they can be inserted into the same map. +func sortToShardsInto(f *FieldOperation, shift int, into ShardedFieldOperation) { + if shift < shardwidth.Exponent { + shift = shardwidth.Exponent + } + nextShift := shift - 8 + if nextShift < shardwidth.Exponent { + nextShift = shardwidth.Exponent + } + // count things that belong in each of the 256 buckets + var buckets [256]int + var starts [256]int + + // compute the buckets ourselves + for _, r := range f.RecordIDs { + b := (r >> shift) & 0xFF + buckets[b]++ + } + total := 0 + // compute starting points of each bucket, converting the + // bucket counts into ends + for i := range buckets { + starts[i] = total + total += buckets[i] + buckets[i] = total + } + // starts[n] is the index of the first thing that should + // go in that bucket, buckets[n] is the index of the first + // thing that shouldn't + var bucketOp FieldOperation + for bucket, start := range starts { + end := buckets[bucket] + if end <= start { + continue + } + for j := start; j < end; j++ { + want := int((f.RecordIDs[j] >> shift) & 0xFF) + for want != bucket { + // move this to the beginning of the + // bucket it wants to be in, swapping + // the thing there here + dst := starts[want] + f.RecordIDs[j], f.RecordIDs[dst] = f.RecordIDs[dst], f.RecordIDs[j] + if f.Values != nil { + f.Values[j], f.Values[dst] = f.Values[dst], f.Values[j] + } + if f.Signed != nil { + f.Signed[j], f.Signed[dst] = f.Signed[dst], f.Signed[j] + } + starts[want]++ + want = int((f.RecordIDs[j] >> shift) & 0xFF) + } + } + // If shift == shardwidth.Exponent, then this is a completed + // shard and can go into the sharded output. otherwise, we + // can subdivide it. + bucketOp.RecordIDs = f.RecordIDs[start:end] + if f.Values != nil { + bucketOp.Values = f.Values[start:end] + } + if f.Signed != nil { + bucketOp.Signed = f.Signed[start:end] + } + if shift == shardwidth.Exponent { + x := bucketOp + into[f.RecordIDs[start]>>shardwidth.Exponent] = &x + } else { + sortToShardsInto(&bucketOp, nextShift, into) + } + } +} + +const shardMask = ((uint64(1) << shardwidth.Exponent) - 1) + +// SortByValues sorts the operation by values first, then by record +// ID within each value. This is the best ordering for set/mutex fields, +// where we'll want to generate positions in that order. For these +// purposes, a time quantum or bool counts as a kind of a set. +func (f *FieldOperation) SortByValues() { + keys := make([]uint64, len(f.RecordIDs)) + for i, v := range f.RecordIDs { + keys[i] = (f.Values[i] << shardwidth.Exponent) | (v & shardMask) + } + f.SortByKeys(keys) +} + +// SortByRecords sorts the operation by record ID, and not by value at +// all. This makes the most sense for int fields and the like. +func (f *FieldOperation) SortByRecords() { + f.SortByKeys(f.RecordIDs) +} + +// SortByKeys reorganizes the record IDs and values of f according to the +// corresponding members of keys. +func (f *FieldOperation) SortByKeys(keys []uint64) { + if len(f.RecordIDs) < 2 { + return + } + diffMask := uint64(0) + prev := keys[0] + for _, r := range keys[1:] { + diffMask |= r ^ prev + prev = r + } + bitsRemaining := bits.Len64(diffMask) + sortPartialByKeys(f, keys, bitsRemaining-8) +} + +// simpleSort sorts a FieldOperation by external keys, or record IDs. It's a +// horribly naive bubble sort because N is small and a more complex algorithm +// doesn't help as much as you'd hope. This beats using stdlib sort by about +// a factor of two for those small N, for larger N we're using the radix sort +// that calls this. +func simpleSort(f *FieldOperation, keys []uint64) { + if keys != nil { + // sorting by record IDs + if f.Values != nil && f.Signed != nil { + for i := 1; i < len(keys); i++ { + for j := i; j > 0 && keys[j-1] > keys[j]; j-- { + keys[j-1], keys[j] = keys[j], keys[j-1] + f.RecordIDs[j-1], f.RecordIDs[j] = f.RecordIDs[j], f.RecordIDs[j-1] + f.Values[j-1], f.Values[j] = f.Values[j], f.Values[j-1] + f.Signed[j-1], f.Signed[j] = f.Signed[j], f.Signed[j-1] + + } + } + } else if f.Values != nil { + for i := 1; i < len(keys); i++ { + for j := i; j > 0 && keys[j-1] > keys[j]; j-- { + keys[j-1], keys[j] = keys[j], keys[j-1] + f.RecordIDs[j-1], f.RecordIDs[j] = f.RecordIDs[j], f.RecordIDs[j-1] + + f.Values[j-1], f.Values[j] = f.Values[j], f.Values[j-1] + } + } + } else if f.Signed != nil { + for i := 1; i < len(keys); i++ { + for j := i; j > 0 && keys[j-1] > keys[j]; j-- { + keys[j-1], keys[j] = keys[j], keys[j-1] + f.RecordIDs[j-1], f.RecordIDs[j] = f.RecordIDs[j], f.RecordIDs[j-1] + f.Signed[j-1], f.Signed[j] = f.Signed[j], f.Signed[j-1] + } + } + } else { + for i := 1; i < len(keys); i++ { + for j := i; j > 0 && keys[j-1] > keys[j]; j-- { + keys[j-1], keys[j] = keys[j], keys[j-1] + f.RecordIDs[j-1], f.RecordIDs[j] = f.RecordIDs[j], f.RecordIDs[j-1] + } + } + } + } else { + if f.Values != nil && f.Signed != nil { + for i := 1; i < len(f.RecordIDs); i++ { + for j := i; j > 0 && f.RecordIDs[j-1] > f.RecordIDs[j]; j-- { + f.RecordIDs[j-1], f.RecordIDs[j] = f.RecordIDs[j], f.RecordIDs[j-1] + f.Values[j-1], f.Values[j] = f.Values[j], f.Values[j-1] + f.Signed[j-1], f.Signed[j] = f.Signed[j], f.Signed[j-1] + } + } + } else if f.Values != nil { + for i := 1; i < len(f.RecordIDs); i++ { + for j := i; j > 0 && f.RecordIDs[j-1] > f.RecordIDs[j]; j-- { + f.RecordIDs[j-1], f.RecordIDs[j] = f.RecordIDs[j], f.RecordIDs[j-1] + f.Values[j-1], f.Values[j] = f.Values[j], f.Values[j-1] + } + } + } else if f.Signed != nil { + for i := 1; i < len(f.RecordIDs); i++ { + for j := i; j > 0 && f.RecordIDs[j-1] > f.RecordIDs[j]; j-- { + f.RecordIDs[j-1], f.RecordIDs[j] = f.RecordIDs[j], f.RecordIDs[j-1] + f.Signed[j-1], f.Signed[j] = f.Signed[j], f.Signed[j-1] + } + } + } else { + // why do we only have record IDs? I don't know + for i := 1; i < len(f.RecordIDs); i++ { + for j := i; j > 0 && f.RecordIDs[j-1] > f.RecordIDs[j]; j-- { + f.RecordIDs[j-1], f.RecordIDs[j] = f.RecordIDs[j], f.RecordIDs[j-1] + } + } + } + } +} + +func sortPartialByKeys(f *FieldOperation, keys []uint64, shift int) { + if shift < 0 { + shift = 0 + } + externalKeys := &f.RecordIDs[0] != &keys[0] + nextShift := shift - 8 + if nextShift < 0 { + nextShift = 0 + } + // count things that belong in each of the 256 buckets + var buckets [256]int + var starts [256]int + // compute the buckets ourselves + for _, r := range keys { + b := (r >> shift) & 0xFF + buckets[b]++ + } + total := 0 + // compute starting points of each bucket, converting the + // bucket counts into ends + for i := range buckets { + starts[i] = total + total += buckets[i] + buckets[i] = total + } + // starts[n] is the index of the first thing that should + // go in that bucket, buckets[n] is the index of the first + // thing that shouldn't + // var newbuckets [256]int + var bucketOp FieldOperation + for bucket, start := range starts { + end := buckets[bucket] + if end <= start { + continue + } + for j := start; j < end; j++ { + want := int((keys[j] >> shift) & 0xFF) + for want != bucket { + // move this to the beginning of the + // bucket it wants to be in, swapping + // the thing there here + dst := starts[want] + keys[j], keys[dst] = keys[dst], keys[j] + // we do this to allow you to just pass in the records as keys + if externalKeys { + f.RecordIDs[j], f.RecordIDs[dst] = f.RecordIDs[dst], f.RecordIDs[j] + } + if f.Values != nil { + f.Values[j], f.Values[dst] = f.Values[dst], f.Values[j] + } + if f.Signed != nil { + f.Signed[j], f.Signed[dst] = f.Signed[dst], f.Signed[j] + } + starts[want]++ + want = int((keys[j] >> shift) & 0xFF) + } + } + // If shift == shardwidth.Exponent, then this is a completed + // shard and can go into the sharded output. otherwise, we + // can subdivide it. + if shift > 0 { + bucketOp.RecordIDs = f.RecordIDs[start:end] + if f.Values != nil { + bucketOp.Values = f.Values[start:end] + } + if f.Signed != nil { + bucketOp.Signed = f.Signed[start:end] + } + // if there's not very many, sort naively instead + if end-start > 32 { + sortPartialByKeys(&bucketOp, keys[start:end], nextShift) + } else { + // naive stdlib sort + if externalKeys { + simpleSort(&bucketOp, keys[start:end]) + } else { + simpleSort(&bucketOp, nil) + } + } + } + } } // AddPair adds a record ID/value pair where the value is unsigned, as @@ -248,13 +533,15 @@ type ShardOperations struct { // Request is a complete ingest request, which may be any combination // of operations, which may apply to multiple shards. type Request struct { - Ops []*Operation + FieldTypes map[string]FieldType + Ops []*Operation } // ShardedRequest is an ingest request, split up into individual per-shard // operations. type ShardedRequest struct { - Ops map[uint64][]*Operation + FieldTypes map[string]FieldType + Ops map[uint64][]*Operation } // Shard converts a request into the same request, only sharded. @@ -280,7 +567,12 @@ func (r *Request) Shard() (*ShardedRequest, error) { } for field, fieldOp := range op.FieldOps { sharded := fieldOp.Shard() + sorter := fieldTypeSorts[r.FieldTypes[field]] + if sorter == nil { + sorter = (*FieldOperation).SortByRecords + } for shard, data := range sharded { + sorter(data) shardOp, ok := shards[shard] if !ok { if op.OpType == OpWrite { diff --git a/ingest/op_test.go b/ingest/op_test.go index db857ac5e..1988201f4 100644 --- a/ingest/op_test.go +++ b/ingest/op_test.go @@ -15,6 +15,7 @@ package ingest_test import ( + "math/rand" "reflect" "testing" @@ -132,3 +133,55 @@ func TestOpSharding(t *testing.T) { } } } + +func TestFancySharding(t *testing.T) { + const shardLimit = 700 + const recordCount = 5000 + grr := rand.New(rand.NewSource(0)) + for i := 0; i < 100; i++ { + f := &ingest.FieldOperation{RecordIDs: make([]uint64, recordCount), Values: make([]uint64, recordCount)} + shards := make([]int, shardLimit) + for j := range f.RecordIDs { + v := uint64(grr.Int63n(shardLimit << shardwidth.Exponent)) + f.RecordIDs[j] = v + f.Values[j] = uint64(grr.Int63n(8)) + shards[v>>shardwidth.Exponent]++ + } + + sharded := f.SortToShards() + for shard, data := range sharded { + if len(data.RecordIDs) != shards[shard] { + t.Errorf("shard %d: expected %d items, got %d", shard, shards[shard], len(data.RecordIDs)) + } + for _, v := range data.RecordIDs { + if (v >> shardwidth.Exponent) != shard { + t.Errorf("shard %d: got %x, which should be in %d", shard, v, v>>shardwidth.Exponent) + } + } + // expect sorted-ness + data.SortByRecords() + prev := data.RecordIDs[0] + for i, next := range data.RecordIDs[1:] { + if next < prev { + t.Errorf("index %d: prev %d, next %d", i+1, prev, next) + } + prev = next + } + data.SortByValues() + prevV, prevRec := data.Values[0], data.RecordIDs[0] + for i, nextRec := range data.RecordIDs[1:] { + nextV := data.Values[i+1] + if nextV < prevV { + t.Errorf("index %d: prev value %d, next value %d", i+1, prevV, nextV) + } + if nextV == prevV { + if nextRec < prevRec { + t.Errorf("index %d, value %d: prev rec %d, next rec %d", i+1, nextV, prevRec, nextRec) + } + } + prevV = nextV + prevRec = nextRec + } + } + } +}