string support for dataframe (#2405)

* String support for arrow
2026-08-28 02:44:59 +00:00 · 2023-01-12 12:58:19 -06:00 · 2023-01-12 12:58:19 -06:00 · 4cd3cb02a2
commit 4cd3cb02a2
parent 7f6ea0e6e5
8 changed files with 207 additions and 13 deletions
--- 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)
 		}
--- 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))
 	}
 }

--- 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))
 		}
 	}
 }
--- 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)
--- 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 {
--- 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
--- 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=
--- 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