From e167f1c7fa75ac718f70424bd55c3eeb32df3b87 Mon Sep 17 00:00:00 2001 From: Seebs Date: Thu, 29 Jul 2021 20:08:33 -0500 Subject: [PATCH] fancier shard-sorting This is a rework of Nia's radix sort. Still using stdlib sort for the tail ends of things, and should probably replace it at some point because it's still woefully inefficient, but this gets decent performance, and lets us do the fancy thing of doing quick partial sorting by record-key-only to get to shards, then deciding whether to sort by value-then-record (as for a set field) or just by record (as for int fields), which lets us reduce the amount of re-sorting the same data by different criteria we do. We also use a messy code-duplication basically-bubblesort for the inner loops because it's much cheaper for small N. This also lets us use field-aware sorting for shards, sorting them correctly for a corresponding field type, and add corresponding API support and fragment support for an option to tell the fragment code that we already ordered things in the order that's most efficient there, to avoid a second sort that we don't otherwise need. --- api.go | 16 +- fragment.go | 20 ++- ingest/codec.go | 32 ++-- ingest/op.go | 402 +++++++++++++++++++++++++++++++++++++++------- ingest/op_test.go | 53 ++++++ 5 files changed, 450 insertions(+), 73 deletions(-) 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 + } + } + } +}