diff --git a/apply.go b/apply.go index cc9b769a0..fe57c1a51 100644 --- a/apply.go +++ b/apply.go @@ -7,6 +7,7 @@ import ( "io" "os" "path/filepath" + "sort" "strings" "sync" @@ -260,14 +261,14 @@ func (e *executor) executeApplyShard(ctx context.Context, qcx *Qcx, index string func NewShardFile(ctx context.Context, name string, mem memory.Allocator, e *executor) (*ShardFile, error) { if !e.dataFrameExists(name) { - return &ShardFile{dest: name, executor: e}, nil + return &ShardFile{dest: name, executor: e, strings: make(map[key][]string)}, nil } // else read in existing table, err := e.getDataTable(ctx, name, mem) if err != nil { return nil, err } - return &ShardFile{table: table, schema: table.Schema(), dest: name, executor: e}, nil + return &ShardFile{table: table, schema: table.Schema(), dest: name, executor: e, strings: make(map[key][]string)}, nil } type NameType struct { @@ -291,6 +292,8 @@ func cast(v interface{}) arrow.DataType { return arrow.PrimitiveTypes.Float64 case float64: return arrow.PrimitiveTypes.Float64 + case *arrow.StringType: + return arrow.BinaryTypes.String default: vprint.VV("%T .... %v", v, v) } @@ -305,6 +308,11 @@ func (cr *ChangesetRequest) ArrowSchema() *arrow.Schema { return arrow.NewSchema(fields, nil) } +type key struct { + col int + chunk int +} + type ShardFile struct { table arrow.Table schema *arrow.Schema @@ -313,6 +321,7 @@ type ShardFile struct { columns []interface{} dest string executor *executor + strings map[key][]string } func compareSchema(s1, s2 *arrow.Schema) bool { @@ -365,6 +374,8 @@ func (sf *ShardFile) buildAppenders(maxid int64) { sf.columns[i] = make([]int64, newSize) case arrow.PrimitiveTypes.Float64: sf.columns[i] = make([]float64, newSize) + case arrow.BinaryTypes.String: + sf.columns[i] = make([]string, newSize) } } sf.added = newSize @@ -381,6 +392,11 @@ func (sf *ShardFile) SetFloatValue(col int, row int64, val float64) { v[row-sf.beforeRows] = val } +func (sf *ShardFile) SetStringValue(col int, row int64, val string) { + v := sf.columns[col].([]string) + v[row-sf.beforeRows] = val +} + func (sf *ShardFile) Process(cs *ChangesetRequest) error { err := sf.process(cs) if err != nil { @@ -394,9 +410,34 @@ func (sf *ShardFile) Process(cs *ChangesetRequest) error { return os.Rename(rtemp+sf.executor.TableExtension(), sf.dest+sf.executor.TableExtension()) } +func (sf *ShardFile) LoadBlobs() error { + for col := 0; col < len(sf.schema.Fields()); col++ { + column := sf.table.Column(col) + switch column.DataType() { + case arrow.BinaryTypes.String: + for i, chunk := range column.Data().Chunks() { + stringData := chunk.(*array.String) + k := key{col: col, chunk: i} + for j := 0; j < stringData.Len(); j++ { + v := stringData.Value(j) + sf.strings[k] = append(sf.strings[k], v) + + } + } + } + } + return nil +} + +func (sf *ShardFile) ReplaceString(col, chunk, l int, s string) { + sf.strings[key{col: col, chunk: chunk}][l] = s +} + func (sf *ShardFile) process(cs *ChangesetRequest) error { offset := 0 if sf.table != nil { + // need to load blobs prior + sf.LoadBlobs() column := sf.table.Column(0) resolver := dataframe.NewChunkResolver(column) for i, rowid := range cs.ShardIds { @@ -414,6 +455,10 @@ func (sf *ShardFile) process(cs *ChangesetRequest) error { case arrow.PrimitiveTypes.Float64: v := column.Data().Chunk(chunk).(*array.Float64).Float64Values() v[l] = cs.Columns[col].([]float64)[i] + case arrow.BinaryTypes.String: + // TODO(twg) 2023/01/09 How to update existing? + new := cs.Columns[col].([]string)[i] + sf.ReplaceString(col, chunk, l, new) default: panic(fmt.Sprintf("Unknown Type %v", column.DataType())) } @@ -433,6 +478,8 @@ func (sf *ShardFile) process(cs *ChangesetRequest) error { sf.SetIntValue(col, rowid, cs.Columns[col].([]int64)[i]) case arrow.PrimitiveTypes.Float64: sf.SetFloatValue(col, rowid, cs.Columns[col].([]float64)[i]) + case arrow.BinaryTypes.String: + sf.SetStringValue(col, rowid, cs.Columns[col].([]string)[i]) default: panic(fmt.Sprintf("2 Unknown Type %v", sf.schema.Field(col).Type)) } @@ -443,15 +490,65 @@ func (sf *ShardFile) process(cs *ChangesetRequest) error { return nil } +type twoSlices struct { + id_slice []int + lists_slice [][]string +} + +type SortByOther twoSlices + +func (sbo SortByOther) Len() int { + return len(sbo.id_slice) +} + +func (sbo SortByOther) Swap(i, j int) { + sbo.id_slice[i], sbo.id_slice[j] = sbo.id_slice[j], sbo.id_slice[i] + sbo.lists_slice[i], sbo.lists_slice[j] = sbo.lists_slice[j], sbo.lists_slice[i] +} + +func (sbo SortByOther) Less(i, j int) bool { + return sbo.id_slice[i] < sbo.id_slice[j] +} + +func (sf *ShardFile) buildFromStrings(idx int, mem memory.Allocator) []arrow.Array { + ids := make([]int, 0) + lists := make([][]string, 0) + for k, v := range sf.strings { + if k.col == idx { // ugh not ordered :( + ids = append(ids, k.chunk) + lists = append(lists, v) + } + } + // sort ids/lists + parts := twoSlices{id_slice: ids, lists_slice: lists} + sort.Sort(SortByOther(parts)) + + builder := array.NewStringBuilder(mem) + chunks := make([]arrow.Array, 0) + for _, v := range parts.lists_slice { + builder.AppendValues(v, nil) + newChunk := builder.NewArray() + chunks = append(chunks, newChunk) + } + return chunks +} + func (sf *ShardFile) Save(name string) error { parts := make([]arrow.Array, 0) mem := memory.NewGoAllocator() for col := 0; col < len(sf.schema.Fields()); col++ { chunks := make([]arrow.Array, 0) if sf.table != nil { - // we append if there was existing parquet file + // we append if there was existing file column := sf.table.Column(col) - chunks = append(chunks, column.Data().Chunks()...) + // if primative type + switch column.DataType() { + case arrow.BinaryTypes.String: + chunks = sf.buildFromStrings(col, mem) + default: + chunks = append(chunks, column.Data().Chunks()...) + } + // else binary type } switch sf.schema.Field(col).Type { case arrow.PrimitiveTypes.Int64: @@ -480,6 +577,18 @@ func (sf *ShardFile) Save(name string) error { return err } parts = append(parts, record) + case arrow.BinaryTypes.String: + if sf.added > 0 { + fbuild := array.NewStringBuilder(mem) + fbuild.AppendValues(sf.columns[col].([]string), nil) // TODO(twg) 2022/09/28 need to handle null + newChunk := fbuild.NewArray() + chunks = append(chunks, newChunk) + } + record, err := array.Concatenate(chunks, mem) + if err != nil { + return err + } + parts = append(parts, record) default: vprint.VV("UNKNOWN %T", sf.schema.Field(col).Type) } diff --git a/arrow.go b/arrow.go index 4538d8fc2..6d23f8389 100644 --- a/arrow.go +++ b/arrow.go @@ -194,6 +194,8 @@ func (st *basicTable) Get(column, row int) interface{} { case *arrow.Float64Type: v := chunk.(*array.Float64).Float64Values() return v[i] + case *arrow.StringType: + return chunk.(*array.String).Value(i) } return 0 } @@ -223,8 +225,10 @@ func builderFrom(mem memory.Allocator, dt arrow.DataType, size int64) array.Buil bldr = array.NewFloat32Builder(mem) case *arrow.Float64Type: bldr = array.NewFloat64Builder(mem) + case *arrow.StringType: + bldr = array.NewStringBuilder(mem) default: - panic(fmt.Errorf("npy2root: invalid Arrow type %v", dt)) + panic(fmt.Errorf("builderFrom: invalid Arrow type %v", dt)) } bldr.Reserve(int(size)) return bldr @@ -254,8 +258,10 @@ func appendData(bldr array.Builder, v interface{}) { bldr.Append(v.(float32)) case *array.Float64Builder: bldr.Append(v.(float64)) + case *array.StringBuilder: + bldr.Append(v.(string)) default: - panic(fmt.Errorf("npy2root: invalid Arrow builder type %T", bldr)) + panic(fmt.Errorf("appendData: invalid Arrow builder type %T", bldr)) } } diff --git a/ctl/dataframe-csv-loader.go b/ctl/dataframe-csv-loader.go index 17c252e23..b80e7ce60 100644 --- a/ctl/dataframe-csv-loader.go +++ b/ctl/dataframe-csv-loader.go @@ -31,6 +31,7 @@ var ( func init() { gob.Register(arrow.PrimitiveTypes.Int64) gob.Register(arrow.PrimitiveTypes.Float64) + gob.Register(arrow.BinaryTypes.String) } // TODO(rdp): add refresh token to this as well @@ -154,7 +155,7 @@ func (cmd *DataframeCsvLoaderCommand) Run(ctx context.Context) (err error) { } fields := make([]arrow.Field, 0) fields = append(fields, arrow.Field{Name: "_ID", Type: arrow.PrimitiveTypes.Int64}) - fileScanner := bufio.NewScanner(readFile) + fileScanner := bufio.NewScanner(readFile) // TODO(twg) 2023/01/11 need to convert to the go CSV reader for more robust string support fileScanner.Split(bufio.ScanLines) // need for really long csv lines var buf []byte @@ -176,6 +177,8 @@ func (cmd *DataframeCsvLoaderCommand) Run(ctx context.Context) (err error) { fields = append(fields, arrow.Field{Name: name, Type: arrow.PrimitiveTypes.Int64}) } else if strings.HasSuffix(col, "__F") { fields = append(fields, arrow.Field{Name: name, Type: arrow.PrimitiveTypes.Float64}) + } else if strings.HasSuffix(col, "__S") { + fields = append(fields, arrow.Field{Name: name, Type: arrow.BinaryTypes.String}) } else { return errors.New("invalid format for type") } @@ -268,6 +271,8 @@ func (cmd *DataframeCsvLoaderCommand) Run(ctx context.Context) (err error) { continue } shardFile.SetFloatValue(i, shardRow, val) + case arrow.BinaryTypes.String: + shardFile.SetStringValue(i, shardRow, rec) default: return errors.New("unhandled arrow type type") } @@ -314,6 +319,11 @@ func (s *ShardDiff) SetFloatValue(col int, row int64, val float64) { s.columns[col] = append(slice, val) } +func (s *ShardDiff) SetStringValue(col int, row int64, val string) { + slice := s.columns[col].([]string) + s.columns[col] = append(slice, val) +} + func (s *ShardDiff) SetNulll(col int, row uint64) { s.null[pair{col: col, row: row}] = struct{}{} } @@ -329,6 +339,8 @@ func (s *ShardDiff) Setup(schema *arrow.Schema) { s.columns = append(s.columns, make([]int64, 0)) case arrow.PrimitiveTypes.Float64: s.columns = append(s.columns, make([]float64, 0)) + case arrow.BinaryTypes.String: + s.columns = append(s.columns, make([]string, 0)) } } } diff --git a/ctl/dataframe-csv-loader_test.go b/ctl/dataframe-csv-loader_test.go index b3cac1ac1..b85a982d5 100644 --- a/ctl/dataframe-csv-loader_test.go +++ b/ctl/dataframe-csv-loader_test.go @@ -66,6 +66,33 @@ func TestDataframeCsvLoaderCommand(t *testing.T) { cm.Path = file.Name() cm.Index = index + err = cm.Run(ctx) + if err != nil { + t.Fatalf("DataframeCsvLoader Run doesn't work: %s", err) + } + }) + t.Run("strings", func(t *testing.T) { + cmLog := logger.NewStandardLogger(io.Discard) + cm := NewDataframeCsvLoaderCommand(cmLog) + file, err := testhook.TempFile(t, "import_string.csv") + if err != nil { + t.Fatalf("creating tempfile: %v", err) + } + _, err = file.Write([]byte("id,val__S\nA,ab\nB,cd\nC,ef")) + if err != nil { + t.Fatalf("writing to tempfile: %v", err) + } + ctx := context.Background() + if err != nil { + t.Fatal(err) + } + index := "strings" + cmd.API.CreateIndex(ctx, index, pilosa.IndexOptions{Keys: true}) + + cm.Host = cmd.API.Node().URI.HostPort() + cm.Path = file.Name() + cm.Index = index + err = cm.Run(ctx) if err != nil { t.Fatalf("DataframeCsvLoader Run doesn't work: %s", err) diff --git a/dataframe_test.go b/dataframe_test.go index 1e041356b..f5836c975 100644 --- a/dataframe_test.go +++ b/dataframe_test.go @@ -51,19 +51,20 @@ func TestExecutor_Apply(t *testing.T) { } t.Run("dataframe ingest", func(t *testing.T) { - // func (c *Client) ApplyDataframeChangeset(indexName string, cr *pilosa.ChangesetRequest, shard uint64) (map[string]interface{}, error) { cr := &pilosa.ChangesetRequest{} // for each row a list of columns cr.Columns = []interface{}{ []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, []int64{2, 4, 6, 8, 10, 12, 14, 16, 18, 20}, []float64{1, 1.414, 1.732, 2, 2.236, 2.449, 2.646, 2.828, 3, 3.162}, + []string{"A", "B", "C", "D", "E", "F", "G", "H", "I", "J"}, } cr.ShardIds = []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} cr.SimpleSchema = []pilosa.NameType{ {Name: "_ID", DataType: arrow.PrimitiveTypes.Int64}, {Name: "ival", DataType: arrow.PrimitiveTypes.Int64}, {Name: "fval", DataType: arrow.PrimitiveTypes.Float64}, + {Name: "sval", DataType: arrow.BinaryTypes.String}, } shard := uint64(0) err := api.ApplyDataframeChangeset(ctx, indexName, cr, shard) @@ -72,7 +73,7 @@ func TestExecutor_Apply(t *testing.T) { } }) t.Run("dataframe schema", func(t *testing.T) { - expectedJSON := `[{"Name":"_ID","Type":"int64"},{"Name":"ival","Type":"int64"},{"Name":"fval","Type":"float64"}]` + expectedJSON := `[{"Name":"_ID","Type":"int64"},{"Name":"ival","Type":"int64"},{"Name":"fval","Type":"float64"},{"Name":"sval","Type":"utf8"}]` parts, err := api.GetDataframeSchema(ctx, indexName) if err != nil { t.Fatal(err) @@ -138,7 +139,7 @@ func TestExecutor_Apply(t *testing.T) { if res, err := api.Query(ctx, &pilosa.QueryRequest{Index: indexName, Query: pql}); err != nil { t.Fatal(err) } else { - expectedJSON := `{"Results":[{"_ID":[2,4,6],"fval":[1.414,2,2.449],"ival":[4,8,12]}],"Err":null,"Profile":null}` + expectedJSON := `{"Results":[{"_ID":[2,4,6],"fval":[1.414,2,2.449],"ival":[4,8,12],"sval":["B","D","F"]}],"Err":null,"Profile":null}` w := new(bytes.Buffer) if err := json.NewEncoder(w).Encode(res); err != nil { t.Fatal(err) @@ -165,6 +166,44 @@ func TestExecutor_Apply(t *testing.T) { } } }) + t.Run("dataframe ingest update", func(t *testing.T) { + cr := &pilosa.ChangesetRequest{} + // for each row a list of columns + cr.Columns = []interface{}{ + []int64{1}, + []int64{20}, + []float64{10}, + []string{"A2"}, + } + cr.ShardIds = []int64{1} + cr.SimpleSchema = []pilosa.NameType{ + {Name: "_ID", DataType: arrow.PrimitiveTypes.Int64}, + {Name: "ival", DataType: arrow.PrimitiveTypes.Int64}, + {Name: "fval", DataType: arrow.PrimitiveTypes.Float64}, + {Name: "sval", DataType: arrow.BinaryTypes.String}, + } + shard := uint64(0) + err := api.ApplyDataframeChangeset(ctx, indexName, cr, shard) + if err != nil { + t.Fatal(err) + } + }) + t.Run("dataframe arrow filter with header", func(t *testing.T) { + pql := `Arrow(ConstRow(columns=[1]),header=["ival","fval","sval"])` + if res, err := api.Query(ctx, &pilosa.QueryRequest{Index: indexName, Query: pql}); err != nil { + t.Fatal(err) + } else { + expectedJSON := `{"Results":[{"_ID":[1],"fval":[10],"ival":[20],"sval":["A2"]}],"Err":null,"Profile":null}` + w := new(bytes.Buffer) + if err := json.NewEncoder(w).Encode(res); err != nil { + t.Fatal(err) + } + got := strings.Trim(w.String(), "\t \n") + if strings.Compare(got, expectedJSON) != 0 { + t.Fatalf("expected: %v got: %v", expectedJSON, got) + } + } + }) t.Run("dataframe delete", func(t *testing.T) { err := api.DeleteDataframe(ctx, indexName) if err != nil { diff --git a/go.mod b/go.mod index 4548cccb7..20f12e008 100644 --- a/go.mod +++ b/go.mod @@ -4,7 +4,7 @@ replace github.com/go-avro/avro => github.com/pilosa/avro v0.0.0-20200626214113- replace github.com/gomem/gomem => github.com/tgruben/gomem v0.0.0-20221021111114-79fdc77dcf61 -replace robpike.io/ivy => github.com/tgruben/ivy v0.0.0-20221107170120-634b546dcdac +replace robpike.io/ivy => github.com/tgruben/ivy v0.0.0-20230111144143-b80a659caeaf require ( github.com/CAFxX/gcnotifier v0.0.0-20220409005548-0153238b886a diff --git a/go.sum b/go.sum index 8c291fb52..7657c22ad 100644 --- a/go.sum +++ b/go.sum @@ -1058,8 +1058,8 @@ github.com/subosito/gotenv v1.2.0/go.mod h1:N0PQaV/YGNqwC0u51sEeR/aUtSLEXKX9iv69 github.com/syndtr/goleveldb v1.0.0/go.mod h1:ZVVdQEZoIme9iO1Ch2Jdy24qqXrMMOU6lpPAyBWyWuQ= github.com/tgruben/gomem v0.0.0-20221021111114-79fdc77dcf61 h1:3RJ3IN/m4w4B7PeYF8PPOdGLHqu045LIxtymF4lQO7g= github.com/tgruben/gomem v0.0.0-20221021111114-79fdc77dcf61/go.mod h1:avNE+ynGJYvQNY+/5Gk6aHxtlOtt2aO6V9ShcONBvvI= -github.com/tgruben/ivy v0.0.0-20221107170120-634b546dcdac h1:xVEKycwTG+9q18T1zA+ESrkIMIMQMmlsGOFirh5VxT8= -github.com/tgruben/ivy v0.0.0-20221107170120-634b546dcdac/go.mod h1:/COPfnSdd23BhmpDXFKsYPk1kK2sJpUr21TsZ1nsgcg= +github.com/tgruben/ivy v0.0.0-20230111144143-b80a659caeaf h1:+bnHPov8gMIztSP1zxMbUXei8IiVNXbQre4+u9lpkQ0= +github.com/tgruben/ivy v0.0.0-20230111144143-b80a659caeaf/go.mod h1:/COPfnSdd23BhmpDXFKsYPk1kK2sJpUr21TsZ1nsgcg= github.com/tidwall/btree v0.3.0/go.mod h1:huei1BkDWJ3/sLXmO+bsCNELL+Bp2Kks9OLyQFkzvA8= github.com/tidwall/btree v1.1.0/go.mod h1:TzIRzen6yHbibdSfK6t8QimqbUnoxUSrZfeW7Uob0q4= github.com/tidwall/buntdb v1.2.0/go.mod h1:XLza/dhlwzO6dc5o/KWor4kfZSt3BP8QV+77ZMKfI58= diff --git a/http_handler.go b/http_handler.go index cf779bd39..7a98351da 100644 --- a/http_handler.go +++ b/http_handler.go @@ -4207,6 +4207,7 @@ func (h *Handler) handleGetHealth(w http.ResponseWriter, r *http.Request) { func init() { gob.Register(arrow.PrimitiveTypes.Int64) gob.Register(arrow.PrimitiveTypes.Float64) + gob.Register(arrow.BinaryTypes.String) } // EXPERIMENTAL API MAY CHANGE