diff --git a/client/batch.go b/client/batch.go index f2039a7df..8e5943189 100644 --- a/client/batch.go +++ b/client/batch.go @@ -86,6 +86,7 @@ type agedTranslation struct { // | uint64 | set | any | // | int64 | int | any | // | float64| decimal | scale | +// | bool | bool | any | // | nil | any | | // // nil values are ignored. @@ -123,6 +124,15 @@ type Batch struct { // values holds the values for each record of an int field values map[string][]int64 + // boolValues is a map[fieldName][idsIndex]bool, which holds the values for + // each record of a bool field. It is a map of maps in order to accomodate + // nil values (they just aren't recorded in the map[int]). + boolValues map[string]map[int]bool + + // boolNulls holds a slice of indices into b.ids for each bool field which + // has nil values. + boolNulls map[string][]uint64 + // times holds a time for each record. (if any of the fields are time fields) times []QuantizedTime @@ -232,6 +242,8 @@ func NewBatch(client *Client, size int, index *Index, fields []*Field, opts ...B headerMap := make(map[string]*Field, len(fields)) rowIDs := make(map[int][]uint64, len(fields)) values := make(map[string][]int64) + boolValues := make(map[string]map[int]bool) + boolNulls := make(map[string][]uint64) tt := make(map[int]map[string][]int, len(fields)) ttSets := make(map[string]map[string][]int) hasTime := false @@ -256,6 +268,8 @@ func NewBatch(client *Client, size int, index *Index, fields []*Field, opts ...B tt[i] = make(map[string][]int) } rowIDs[i] = make([]uint64, 0, size) + case FieldTypeBool: + boolValues[field.Name()] = make(map[int]bool) default: return nil, errors.Errorf("field type '%s' is not currently supported through Batch", typ) } @@ -272,6 +286,8 @@ func NewBatch(client *Client, size int, index *Index, fields []*Field, opts ...B clearRowIDs: make(map[int]map[int]uint64), rowIDSets: make(map[string][][]uint64), values: values, + boolValues: boolValues, + boolNulls: boolNulls, nullIndices: make(map[string][]uint64), toTranslate: tt, toTranslateClear: make(map[int]map[string][]int), @@ -479,28 +495,8 @@ func (b *Batch) Add(rec Row) error { field := b.header[i] switch val := rec.Values[i].(type) { case string: - if field.Opts().Type() != FieldTypeInt { - // nil-extend - for len(b.rowIDs[i]) < curPos { - b.rowIDs[i] = append(b.rowIDs[i], nilSentinel) - } - rowIDs := b.rowIDs[i] - // empty string is not a valid value at this point (Pilosa refuses to translate it) - if val == "" { // - b.rowIDs[i] = append(rowIDs, nilSentinel) - - } else if rowID, ok := b.getRowTranslation(field.Name(), val); ok { - b.rowIDs[i] = append(rowIDs, rowID) - } else { - ints, ok := b.toTranslate[i][val] - if !ok { - ints = make([]int, 0) - } - ints = append(ints, curPos) - b.toTranslate[i][val] = ints - b.rowIDs[i] = append(rowIDs, 0) - } - } else if field.Opts().Type() == FieldTypeInt { + switch field.Opts().Type() { + case FieldTypeInt: if val == "" { // copied from the `case nil:` section for ints and decimals b.values[field.Name()] = append(b.values[field.Name()], 0) @@ -521,6 +517,30 @@ func (b *Batch) Add(rec Row) error { b.toTranslate[i][val] = ints b.values[field.Name()] = append(b.values[field.Name()], 0) } + case FieldTypeBool: + // If we want to support bools as string values, we would do + // that here. + default: + // nil-extend + for len(b.rowIDs[i]) < curPos { + b.rowIDs[i] = append(b.rowIDs[i], nilSentinel) + } + rowIDs := b.rowIDs[i] + // empty string is not a valid value at this point (Pilosa refuses to translate it) + if val == "" { // + b.rowIDs[i] = append(rowIDs, nilSentinel) + + } else if rowID, ok := b.getRowTranslation(field.Name(), val); ok { + b.rowIDs[i] = append(rowIDs, rowID) + } else { + ints, ok := b.toTranslate[i][val] + if !ok { + ints = make([]int, 0) + } + ints = append(ints, curPos) + b.toTranslate[i][val] = ints + b.rowIDs[i] = append(rowIDs, 0) + } } case uint64: // nil-extend @@ -578,8 +598,8 @@ func (b *Batch) Add(rec Row) error { } b.rowIDSets[field.Name()] = append(rowIDSets, val) case nil: - t := field.Opts().Type() - if t == FieldTypeInt || t == FieldTypeDecimal || t == FieldTypeTimestamp { + switch field.Opts().Type() { + case FieldTypeInt, FieldTypeDecimal, FieldTypeTimestamp: b.values[field.Name()] = append(b.values[field.Name()], 0) nullIndices, ok := b.nullIndices[field.Name()] if !ok { @@ -588,7 +608,15 @@ func (b *Batch) Add(rec Row) error { nullIndices = append(nullIndices, uint64(curPos)) b.nullIndices[field.Name()] = nullIndices - } else { + case FieldTypeBool: + boolNulls, ok := b.boolNulls[field.Name()] + if !ok { + boolNulls = make([]uint64, 0) + } + boolNulls = append(boolNulls, uint64(curPos)) + b.boolNulls[field.Name()] = boolNulls + + default: // only append nil to rowIDs if this field already has // rowIDs. Otherwise, this could be a []string or // []uint64 field where we've only seen nil values so @@ -599,6 +627,10 @@ func (b *Batch) Add(rec Row) error { b.rowIDs[i] = append(rowIDs, nilSentinel) } } + + case bool: + b.boolValues[field.Name()][curPos] = val + default: return errors.Errorf("Val %v Type %[1]T is not currently supported. Use string, uint64 (row id), or int64 (integer value)", val) } @@ -711,7 +743,6 @@ func (b *Batch) Import() error { return errors.Wrap(err, "making fragments (flush)") } if b.useShardTransactionalEndpoint { - // TODO handle bool? frags, clearFrags, err = b.makeSingleValFragments(frags, clearFrags) if err != nil { return errors.Wrap(err, "making single val fragments") @@ -1477,6 +1508,60 @@ func (b *Batch) makeSingleValFragments(frags, clearFrags fragments) (fragments, } } } + // ------------------------- + // Boolean fields + // ------------------------- + falseRowOffset := 0 * shardWidth // fragment row 0 + trueRowOffset := 1 * shardWidth // fragment row 1 + + // For bools which have been set to null, clear both the true and false + // values for the record. Because this ends up going through the + // API.ImportRoaringShard() method (which handles `bool` fields the same as + // `mutex` fields), we don't actually set the true and false rows of the + // boolean fragment; rather, we just set the first row to indicate which + // records (for all rows) to clear. + for fieldname, boolNulls := range b.boolNulls { + field := b.headerMap[fieldname] + if field.Opts().Type() != featurebase.FieldTypeBool { + continue + } + for _, pos := range boolNulls { + recID := b.ids[pos] + shard := recID / shardWidth + clearBM := clearFrags.GetOrCreate(shard, field.Name(), "standard") + + fragmentColumn := recID % shardWidth + clearBM.Add(fragmentColumn) + } + } + + // For bools which have been set to a non-nil value, set the appropriate + // value for the record, and unset the opposing values. For example, if the + // bool is set to `false`, then set the bit in the "false" row, and clear + // the bit in the "true" row. + for fieldname, boolMap := range b.boolValues { + field := b.headerMap[fieldname] + if field.Opts().Type() != featurebase.FieldTypeBool { + continue + } + + for pos, boolVal := range boolMap { + recID := b.ids[pos] + + shard := recID / shardWidth + bitmap := frags.GetOrCreate(shard, field.Name(), "standard") + clearBM := clearFrags.GetOrCreate(shard, field.Name(), "standard") + + fragmentColumn := recID % shardWidth + clearBM.Add(fragmentColumn) + + if boolVal { + bitmap.Add(trueRowOffset + fragmentColumn) + } else { + bitmap.Add(falseRowOffset + fragmentColumn) + } + } + } return frags, clearFrags, nil } @@ -1699,6 +1784,14 @@ func (b *Batch) reset() { delete(clearMap, k) } } + for _, boolsMap := range b.boolValues { + for k := range boolsMap { + delete(boolsMap, k) + } + } + for k := range b.boolNulls { + delete(b.boolNulls, k) // TODO pool these slices + } for i := range b.toTranslateID { b.toTranslateID[i] = "" } diff --git a/client/batch_test.go b/client/batch_test.go index 3dd382e0b..caed90cfa 100644 --- a/client/batch_test.go +++ b/client/batch_test.go @@ -31,6 +31,7 @@ func TestAgainstCluster(t *testing.T) { client := NewTestClient(t, c) t.Run("string-slice-combos", func(t *testing.T) { testStringSliceCombos(t, c, client) }) t.Run("import-batch-ints", func(t *testing.T) { testImportBatchInts(t, c, client) }) + t.Run("import-batch-bools", func(t *testing.T) { testImportBatchBools(t, c, client) }) t.Run("import-batch-sorting", func(t *testing.T) { testImportBatchSorting(t, c, client) }) t.Run("test-trim-null", func(t *testing.T) { testTrimNull(t, c, client) }) t.Run("test-string-slice-empty-and-nil", func(t *testing.T) { testStringSliceEmptyAndNil(t, c, client) }) @@ -1818,7 +1819,7 @@ func mutexNilClearKey(t *testing.T, c *test.Cluster, client *Client) { if err != nil { t.Fatalf("importing: %v", err) } - resp, err := client.Query(idx.RawQuery(`Row(mut="a")`)) + resp, _ := client.Query(idx.RawQuery(`Row(mut="a")`)) errorIfNotEqual(t, resp.Result().Row().Keys, []string{"0", "2"}) r.ID = "2" @@ -1838,3 +1839,45 @@ func mutexNilClearKey(t *testing.T, c *test.Cluster, client *Client) { } errorIfNotEqual(t, resp.Result().Row().Keys, []string{"0"}) } + +func testImportBatchBools(t *testing.T, c *test.Cluster, client *Client) { + schema := NewSchema() + idx := schema.Index("test-import-batch-bools") + field := idx.Field("boolcol", OptFieldTypeBool()) + err := client.SyncSchema(schema) + if err != nil { + t.Fatalf("syncing schema: %v", err) + } + + b, err := NewBatch(client, 3, idx, []*Field{field}, OptUseShardTransactionalEndpoint(true)) + if err != nil { + t.Fatalf("getting batch: %v", err) + } + r := Row{Values: make([]interface{}, 1)} + + r.ID = uint64(0) + r.Values[0] = bool(false) + err = b.Add(r) + if err != nil { + t.Fatalf("adding after import: %v", err) + } + r.ID = uint64(1) + r.Values[0] = bool(true) + err = b.Add(r) + if err != nil { + t.Fatalf("adding second after import: %v", err) + } + + err = b.Import() + if err != nil { + t.Fatalf("second import: %v", err) + } + + resp, err := client.Query(idx.RawQuery("Count(All())")) + if err != nil { + t.Fatalf("querying: %v", err) + } + if res := resp.Results()[0]; res.Count() != 2 { + t.Fatalf("unexpected result: %+v", res) + } +} diff --git a/idk/ingest.go b/idk/ingest.go index 094b9e2ce..58ad6600b 100644 --- a/idk/ingest.go +++ b/idk/ingest.go @@ -1530,9 +1530,10 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor case DeleteSentinel: if hasMutex { //need to clear the mutex rec.Clears[valIdx] = nil //? maybe - } else { //TODO(twg) set fields not supported - } + // else { //TODO(twg) set fields not supported + + // } default: rec.Values[valIdx], err = idkField.PilosafyVal(rawRec[i]) } @@ -1604,10 +1605,15 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor fields = append(fields, m.index.Field(fld.DestName(), pilosaclient.OptFieldTypeBool())) valIdx := len(fields) - 1 recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { - rec.Values[valIdx], err = idkField.PilosafyVal(rawRec[i]) + val, err := idkField.PilosafyVal(rawRec[i]) + if err != nil { + return errors.Wrapf(err, "booling '%v' of %[1]T", val) + } + if b, ok := val.(bool); ok { + rec.Values[valIdx] = b + } return errors.Wrapf(err, "converting field %d:%+v, val:%+v", i, idkField, rawRec[i]) }) - // TODO, unpacked bools aren't actually supported by importbatch } else { fields = append(fields, boolField, boolFieldExists) fieldIdx := len(fields) - 2 diff --git a/idk/ingest_test.go b/idk/ingest_test.go index 5b1d8e3b2..942d9499f 100644 --- a/idk/ingest_test.go +++ b/idk/ingest_test.go @@ -20,6 +20,7 @@ import ( "github.com/molecula/featurebase/v3/idk/idktest" "github.com/molecula/featurebase/v3/logger" "github.com/pkg/errors" + "github.com/stretchr/testify/assert" ) func configureTestFlags(main *Main) { @@ -269,6 +270,7 @@ func TestSingleBoolClear(t *testing.T) { ingester.Index = fmt.Sprintf("single_bool_clear%d", rand.Intn(100000)) ingester.BatchSize = 1 ingester.IDField = "id" + ingester.PackBools = "bools" if err := ingester.Run(); err != nil { t.Fatalf("%s: %v", idktest.ErrRunningIngest, err) @@ -300,6 +302,7 @@ func TestSingleBoolClear(t *testing.T) { ingester2.NewSource = func() (Source, error) { return ts2, nil } ingester2.Index = ingester.Index ingester2.IDField = "id" + ingester2.PackBools = "bools" if err := ingester2.Run(); err != nil { t.Fatalf("running ingester2: %v", err) @@ -566,6 +569,7 @@ func TestDelete(t *testing.T) { deleter := NewMain() configureTestFlags(deleter) deleter.Delete = true + deleter.PackBools = "bools" deleter.NewSource = func() (Source, error) { return tsDelete, nil } deleter.Index = indexName deleter.BatchSize = 5 @@ -579,6 +583,7 @@ func TestDelete(t *testing.T) { ingester := NewMain() configureTestFlags(ingester) + ingester.PackBools = "bools" ingester.NewSource = func() (Source, error) { return tsWrite, nil } ingester.PrimaryKeyFields = primaryKeyFields ingester.Index = indexName @@ -874,7 +879,6 @@ func TestBatchFromSchema(t *testing.T) { rawRec: []interface{}{true, uint64(7), false}, rowID: uint64(7), rowVals: []interface{}{true, false}, - err: "field type 'bool' is not currently supported through Batch", }, { name: "mutex field", @@ -1429,6 +1433,7 @@ func TestNilIngest(t *testing.T) { err string batchErr string rdzErrs []string + packBools string } runTest := func(t *testing.T, test testcase, removeIndex bool, server serverInfo, rawRec []interface{}, clearmap map[int]interface{}, values []interface{}) { m := NewMain() @@ -1437,6 +1442,7 @@ func TestNilIngest(t *testing.T) { m.PrimaryKeyFields = test.pkFields m.BatchSize = 2 m.Pprof = "" + m.PackBools = test.packBools m.NewSource = func() (Source, error) { return nil, nil } if server.AuthToken != "" { m.AuthToken = server.AuthToken @@ -1522,12 +1528,13 @@ func TestNilIngest(t *testing.T) { Vals2: []interface{}{nil, nil}, Vals3: []interface{}{nil, nil}, }, { - name: "bools null", - pkFields: []string{"user_id"}, - rawRec1: []interface{}{"1a", true}, // bool and bool-exists - rawRec2: []interface{}{"1a", DELETE_SENTINEL}, - rawRec3: []interface{}{"1a", nil}, - rowID: "1a", + name: "bools null", + packBools: "bools", + pkFields: []string{"user_id"}, + rawRec1: []interface{}{"1a", true}, // bool and bool-exists + rawRec2: []interface{}{"1a", DELETE_SENTINEL}, + rawRec3: []interface{}{"1a", nil}, + rowID: "1a", schema: []Field{ StringField{NameVal: "user_id"}, BoolField{NameVal: "bool_val_1"}, @@ -1587,3 +1594,114 @@ func TestNilIngest(t *testing.T) { } } + +func TestBoolIngest(t *testing.T) { + rand.Seed(time.Now().UTC().UnixNano()) + indexName := fmt.Sprintf("boolingest%d", rand.Intn(100000)) + primaryKeyFields := []string{"user_id"} + + tests := []struct { + src *testSource + expTrue []string + expFalse []string + expNull []string + }{ + { + src: newTestSource( + []Field{ + StringField{NameVal: "user_id"}, + BoolField{NameVal: "bool_val"}, + }, + [][]interface{}{ + {"a1", true}, + }, + ), + expTrue: []string{"a1"}, + expFalse: nil, + expNull: nil, + }, + { + src: newTestSource( + []Field{ + StringField{NameVal: "user_id"}, + BoolField{NameVal: "bool_val"}, + }, + [][]interface{}{ + {"a1", false}, + }, + ), + expTrue: nil, + expFalse: []string{"a1"}, + expNull: nil, + }, + { + src: newTestSource( + []Field{ + StringField{NameVal: "user_id"}, + BoolField{NameVal: "bool_val"}, + }, + [][]interface{}{ + {"a1", nil}, + }, + ), + expTrue: nil, + expFalse: nil, + expNull: []string{"a1"}, + }, + } + + var ing *Main + defer func() { + ing := ing + if err := ing.PilosaClient().DeleteIndexByName(ing.Index); err != nil { + t.Logf("%s for index %s: %v", idktest.ErrDeletingIndex, ing.Index, err) + } + }() + + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + ingester := NewMain() + configureTestFlags(ingester) + ingester.PackBools = "" + ingester.NewSource = func() (Source, error) { return test.src, nil } + ingester.PrimaryKeyFields = primaryKeyFields + ingester.Index = indexName + ingester.BatchSize = 1 + ingester.UseShardTransactionalEndpoint = true + + // Set ing so the defer can do cleanup. + if i == 0 { + ing = ingester + } + + if err := ingester.Run(); err != nil { + t.Fatalf("%s: %v", idktest.ErrRunningIngest, err) + } + + client := ingester.PilosaClient() + idx := ingester.index + fld := ingester.index.Field("bool_val") + + // Check true. + { + resp, err := client.Query(fld.Row(true)) + assert.NoError(t, err) + assert.Equal(t, test.expTrue, resp.Result().Row().Keys) + } + + // Check false. + { + resp, err := client.Query(fld.Row(false)) + assert.NoError(t, err) + assert.Equal(t, test.expFalse, resp.Result().Row().Keys) + } + + // Check null. + { + resp, err := client.Query(idx.Difference(idx.All(), idx.Union(fld.Row(true), fld.Row(false)))) + assert.NoError(t, err) + assert.Equal(t, test.expNull, resp.Result().Row().Keys) + } + }) + } +} diff --git a/idk/interfaces.go b/idk/interfaces.go index 935760da4..469d7db64 100644 --- a/idk/interfaces.go +++ b/idk/interfaces.go @@ -1131,11 +1131,16 @@ func toBool(val interface{}) (bool, error) { } return vt != 0, nil case string: - switch strings.ToLower(vt) { + vt = strings.ToLower(vt) + vt = strings.TrimSpace(vt) + switch vt { case "", "0", "f", "false": return false, nil + case "1", "t", "true": + return true, nil } - return true, nil + return false, errors.Errorf("couldn't convert %v of %[1]T to bool", vt) + default: if vint, err := toInt64(val); err == nil { return vint != 0, nil diff --git a/idk/interfaces_test.go b/idk/interfaces_test.go index 069cc54ae..b7ab73362 100644 --- a/idk/interfaces_test.go +++ b/idk/interfaces_test.go @@ -15,6 +15,29 @@ func TestDecimalFieldPilosafy(t *testing.T) { } +func TestBoolFieldPilosafy(t *testing.T) { + f := BoolField{NameVal: "boolcol"} + validTrue := []interface{}{true, 1, "t", "true", " T ", " True"} + validFalse := []interface{}{false, 0, "f", "false", " F ", " False"} + invalid := []interface{}{"boat", " test "} + + for i, v := range validTrue { + if ret, err := f.PilosafyVal(v); err != nil || ret != true { + t.Errorf("test: %d, got: %v of %[1]T, err: %v", i, ret, err) + } + } + for i, v := range validFalse { + if ret, err := f.PilosafyVal(v); err != nil || ret != false { + t.Errorf("test: %d, got: %v of %[1]T, err: %v", i, ret, err) + } + } + for i, v := range invalid { + if ret, err := f.PilosafyVal(v); err == nil { + t.Errorf("test: %d, got: %v of %[1]T, err: %v", i, ret, err) + } + } +} + func TestPilosafyVal(t *testing.T) { tests := []struct { field Field diff --git a/idk/kafka/cmd_delete_test.go b/idk/kafka/cmd_delete_test.go index 919b03932..bceb56a2c 100644 --- a/idk/kafka/cmd_delete_test.go +++ b/idk/kafka/cmd_delete_test.go @@ -56,6 +56,7 @@ func TestDeleteConsumerCompoundStringKey(t *testing.T) { m.Index = fmt.Sprintf("cmd_del_comp_indexij%s", topic) m.PrimaryKeyFields = []string{"abc", "db", "user_id"} m.Topics = []string{topic} + m.PackBools = "bools" defer func() { // TODO: for some reason (which I didn't dig into),