From a2e14df15ce00715bca0d5fea6b93c1d99cb00af Mon Sep 17 00:00:00 2001 From: tgruben Date: Mon, 12 Dec 2022 15:13:58 -0600 Subject: [PATCH] Fb 1874 dataframe-csv-loader featurebase subcommand (#2341) Embeded dataframe-csv-loader command as featurebase subservice (cherry picked from commit f2a13c8bde2ccc848b5f91bcd19907382049006b) --- client/client.go | 10 - cmd/dataframe-csv-loader.go | 37 +++ cmd/root.go | 1 + ctl/dataframe-csv-loader.go | 391 +++++++++++++++++++++++++++++++ ctl/dataframe-csv-loader_test.go | 74 ++++++ 5 files changed, 503 insertions(+), 10 deletions(-) create mode 100644 cmd/dataframe-csv-loader.go create mode 100644 ctl/dataframe-csv-loader.go create mode 100644 ctl/dataframe-csv-loader_test.go diff --git a/client/client.go b/client/client.go index b41f32616..0b313b971 100644 --- a/client/client.go +++ b/client/client.go @@ -1842,15 +1842,5 @@ func (c *Client) ApplyDataframeChangeset(indexName string, cr *pilosa.ChangesetR }) } err = eg.Wait() - - // status, body, err := c.HTTPRequest(http.MethodPost, path, buffer.Bytes(), headers) - /* - var result map[string]interface{} - err = json.Unmarshal(body, &result) - if err != nil { - return nil, errors.Wrap(err, "unmarshalling response") - } - */ - return nil, err } diff --git a/cmd/dataframe-csv-loader.go b/cmd/dataframe-csv-loader.go new file mode 100644 index 000000000..592ecaa3e --- /dev/null +++ b/cmd/dataframe-csv-loader.go @@ -0,0 +1,37 @@ +// Copyright 2021 Molecula Corp. All rights reserved. +package cmd + +import ( + "github.com/molecula/featurebase/v3/ctl" + "github.com/molecula/featurebase/v3/logger" + "github.com/spf13/cobra" +) + +// newImportCommand runs the FeatureBase import subcommand for ingesting bulk data. +func newDataframeCsvLoaderCommand(logdest logger.Logger) *cobra.Command { + cmd := ctl.NewDataframeCsvLoaderCommand(logdest) + loaderCmd := &cobra.Command{ + Use: "dataframe-csv-loader", + Short: "load dataframe integer and floating point values into featurebase", + Long: ` +`, + RunE: usageErrorWrapper(cmd), + } + flags := loaderCmd.Flags() + flags.StringVar(&cmd.Path, "csv", "", "path to csv input file") + flags.StringVar(&cmd.Host, "host", "localhost:10101", "host:port of FeatureBase.") + flags.StringVar(&cmd.Pprof, "pprof", cmd.Pprof, "host:port to listen for profiling requests at /debug/pprof and /debug/fgprof.") + flags.StringVar(&cmd.AuthToken, "auth-token", "", "Authentication token") + flags.StringVar(&cmd.Index, "index", "", "Destination Index. ") + flags.IntVar(&cmd.MaxCapacity, "buffer", 0, "Maximum size of of the line buffer defaults to go bufio default ") + ctl.SetTLSConfig( + flags, "", + &cmd.TLS.CertificatePath, + &cmd.TLS.CertificateKeyPath, + &cmd.TLS.CACertPath, + &cmd.TLS.SkipVerify, + &cmd.TLS.EnableClientVerification, + ) + + return loaderCmd +} diff --git a/cmd/root.go b/cmd/root.go index cbda03892..7aa5d47e3 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -106,6 +106,7 @@ at https://docs.featurebase.com/. rc.AddCommand(newKeygenCommand(logdest)) rc.AddCommand(newCLICommand(logdest)) rc.AddCommand(newDAXCommand(stderr)) + rc.AddCommand(newDataframeCsvLoaderCommand(logdest)) rc.SetOutput(stderr) return rc diff --git a/ctl/dataframe-csv-loader.go b/ctl/dataframe-csv-loader.go new file mode 100644 index 000000000..17c252e23 --- /dev/null +++ b/ctl/dataframe-csv-loader.go @@ -0,0 +1,391 @@ +// Copyright 2021 Molecula Corp. All rights reserved. +package ctl + +import ( + "bufio" + "context" + "crypto/tls" + "encoding/gob" + "fmt" + "io" + "math" + "net/http" + "os" + "runtime" + "strconv" + "strings" + + "github.com/apache/arrow/go/v10/arrow" + pilosa "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/client" + "github.com/molecula/featurebase/v3/idk" + "github.com/molecula/featurebase/v3/logger" + "github.com/pkg/errors" +) + +var ( + mask = uint64(pilosa.ShardWidth - 1) + Sentinal = uint64(math.MaxUint64) +) + +func init() { + gob.Register(arrow.PrimitiveTypes.Int64) + gob.Register(arrow.PrimitiveTypes.Float64) +} + +// TODO(rdp): add refresh token to this as well + +// NewDataframeCsvLoaderCommand +type DataframeCsvLoaderCommand struct { + tlsConfig *tls.Config + + Host string + + Index string + + // Filepath to the csv file + Path string + + // max line length of csv file + MaxCapacity int + + // Host:port on which to listen for pprof. + Pprof string `json:"pprof"` + + TLS idk.TLSConfig + AuthToken string `flag:"auth-token" help:"Authentication Token for FeatureBase"` + SchemaManager idk.SchemaManager `flag:"-"` + + // Reusable client. + client *client.Client + index *client.Index + needTranslation bool + + // Standard input/output + log logger.Logger +} + +// Logger returns the command's associated Logger to maintain CommandWithTLSSupport interface compatibility +func (cmd *DataframeCsvLoaderCommand) Logger() logger.Logger { + return cmd.log +} + +// NewDataframeCsvLoaderCommand returns a new instance of DataframeCsvLoaderCommand. +func NewDataframeCsvLoaderCommand(logdest logger.Logger) *DataframeCsvLoaderCommand { + return &DataframeCsvLoaderCommand{ + log: logdest, + Pprof: "localhost:0", + } +} + +func (cmd *DataframeCsvLoaderCommand) setupClient() (*tls.Config, error) { + var tlsConfig *tls.Config + var err error + opts := []client.ClientOption{} + if cmd.TLS.CertificatePath != "" { + tlsConfig, err = idk.GetTLSConfig(&cmd.TLS, cmd.log) + if err != nil { + return nil, errors.Wrap(err, "getting TLS config") + } + opts = append(opts, client.OptClientTLSConfig(tlsConfig)) + } else { + opts = append(opts, + client.OptClientRetries(2), + client.OptClientTotalPoolSize(1000), + client.OptClientPoolSizePerRoute(400), + ) + } + cmd.client, err = client.NewClient([]string{cmd.Host}, opts...) + if err != nil { + return nil, err + } + cmd.client.AuthToken = cmd.AuthToken + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + return tlsConfig, nil +} + +func (cmd *DataframeCsvLoaderCommand) Setup() (err error) { + // setup logging + cmd.log = logger.NewStandardLogger(os.Stderr) + if cmd.Pprof != "" { + go func() { + runtime.SetBlockProfileRate(10000000) // 1 sample per 10 ms + runtime.SetMutexProfileFraction(100) // 1% sampling + cmd.log.Printf("Listening for /debug/pprof/ and /debug/fgprof on '%s'", cmd.Pprof) + cmd.log.Printf("%v", http.ListenAndServe(cmd.Pprof, nil)) + }() + } + // set up Pilosa client + _, err = cmd.setupClient() + if err != nil { + return errors.Wrap(err, "setting up client") + } + + if cmd.AuthToken != "" { + cmd.AuthToken = "Bearer " + cmd.AuthToken // Gets added to context + } + + return nil +} + +// Run executes the dataload. +func (cmd *DataframeCsvLoaderCommand) Run(ctx context.Context) (err error) { + err = cmd.Setup() + if err != nil { + return err + } + logger := cmd.Logger() + close, err := startProfilingServer(cmd.Pprof, logger) + if err != nil { + return errors.Wrap(err, "starting profiling server") + } + defer close() + + // Validate arguments. + if cmd.Path == "" { + return fmt.Errorf("%w: --csv flag required", UsageError) + } + + readFile, err := os.Open(cmd.Path) + if err != nil { + return err + } + fields := make([]arrow.Field, 0) + fields = append(fields, arrow.Field{Name: "_ID", Type: arrow.PrimitiveTypes.Int64}) + fileScanner := bufio.NewScanner(readFile) + fileScanner.Split(bufio.ScanLines) + // need for really long csv lines + var buf []byte + if cmd.MaxCapacity > 0 { + buf = make([]byte, cmd.MaxCapacity) + fileScanner.Buffer(buf, cmd.MaxCapacity) + } + total := 0 + if fileScanner.Scan() { + total++ + t := fileScanner.Text() + p := strings.Split(t, ",") + for _, col := range p[1:] { + col = strings.TrimSpace(col) + cmd.Logger().Infof("checking %v", col) + name := col[:strings.LastIndex(col, "__")] + cmd.Logger().Infof("name:%v", name) + if strings.HasSuffix(col, "__I") { + 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 { + return errors.New("invalid format for type") + } + } + + } else { + return errors.Wrap(fileScanner.Err(), "No header") + } + schema, err := cmd.client.Schema() + if err != nil { + return err + } + idx := schema.Index(cmd.Index) + if idx.Opts().Keys() { + cmd.needTranslation = true + cmd.index = idx + } + + arrowSchema := arrow.NewSchema(fields, nil) + keys := make([]string, 0) + lookup := make(map[string]uint64) + if cmd.needTranslation { + + for fileScanner.Scan() { + t := fileScanner.Text() + r := t[:strings.Index(t, ",")] + _, ok := lookup[r] + if !ok { + keys = append(keys, r) + lookup[r] = Sentinal + } + + } + cmd.Logger().Infof("Translate Keys %d", total) + ids, err := cmd.client.CreateIndexKeys(cmd.index, keys...) + if err != nil { + return err + } + lookup = ids + } + sharder := &Sharder{ + shards: make(map[uint64]*ShardDiff), + schema: arrowSchema, + index: cmd.Index, + log: cmd.log, + } + readFile.Seek(0, io.SeekStart) + fileScanner = bufio.NewScanner(readFile) + if cmd.MaxCapacity > 0 { + fileScanner.Buffer(buf, cmd.MaxCapacity) + } + fileScanner.Split(bufio.ScanLines) + fileScanner.Scan() // skip the header + cmd.Logger().Infof("Build the dataframe input package in memory") + id := uint64(0) + for fileScanner.Scan() { + records := strings.Split(fileScanner.Text(), ",") + if cmd.needTranslation { + id = lookup[records[0]] + } else { + id, err = strconv.ParseUint(records[0], 10, 64) + if err != nil { + return err + } + } + shard := id / pilosa.ShardWidth + shardFile, err := sharder.GetShard(shard) + if err != nil { + return err + } + shardRow := int64(id & mask) + shardFile.SetRow(shardRow) + for i, rec := range records { + if i == 0 { + shardFile.SetIntValue(i, shardRow, int64(id)) + } else { + rec = strings.TrimSpace(rec) + switch arrowSchema.Field(i).Type { + case arrow.PrimitiveTypes.Int64: + val, err := strconv.ParseInt(rec, 10, 64) + if err != nil { + shardFile.SetIntValue(i, shardRow, 0) + continue + } + shardFile.SetIntValue(i, shardRow, val) + case arrow.PrimitiveTypes.Float64: + val, err := strconv.ParseFloat(rec, 64) + if err != nil { + shardFile.SetFloatValue(i, shardRow, 0) + continue + } + shardFile.SetFloatValue(i, shardRow, val) + default: + return errors.New("unhandled arrow type type") + } + } + } + } + cmd.Logger().Infof("sending package to featurebase") + err = sharder.Store(arrowSchema, cmd.client) + if err != nil { + return err + } + return err +} + +type pair struct { + col int + row uint64 +} + +type ShardDiff struct { + columns []interface{} + rows []int64 + null map[pair]struct{} + shard uint64 + // Standard input/output + log logger.Logger +} + +func NewShardDiff(shard uint64, log logger.Logger) (*ShardDiff, error) { + return &ShardDiff{shard: shard, log: log}, nil +} + +type Number interface { + int64 | float64 +} + +func (s *ShardDiff) SetIntValue(col int, row int64, val int64) { + slice := s.columns[col].([]int64) + s.columns[col] = append(slice, val) +} + +func (s *ShardDiff) SetFloatValue(col int, row int64, val float64) { + slice := s.columns[col].([]float64) + s.columns[col] = append(slice, val) +} + +func (s *ShardDiff) SetNulll(col int, row uint64) { + s.null[pair{col: col, row: row}] = struct{}{} +} + +func (s *ShardDiff) SetRow(row int64) { + s.rows = append(s.rows, row) +} + +func (s *ShardDiff) Setup(schema *arrow.Schema) { + for _, f := range schema.Fields() { + switch f.Type { + case arrow.PrimitiveTypes.Int64: + s.columns = append(s.columns, make([]int64, 0)) + case arrow.PrimitiveTypes.Float64: + s.columns = append(s.columns, make([]float64, 0)) + } + } +} + +func makeSimpleSchema(a *arrow.Schema) []pilosa.NameType { + nt := make([]pilosa.NameType, len(a.Fields())) + + for i := 0; i < len(a.Fields()); i++ { + f := a.Field(i) + nt[i] = pilosa.NameType{Name: f.Name, DataType: f.Type} + } + return nt +} + +func (s *ShardDiff) Store(index string, schema *arrow.Schema, fb *client.Client) error { + s.log.Infof("dataframe for shard %v:%v:", index, s.shard) + request := &pilosa.ChangesetRequest{} + request.Columns = s.columns + request.ShardIds = s.rows + request.SimpleSchema = makeSimpleSchema(schema) + _, err := fb.ApplyDataframeChangeset(index, request, s.shard) + return err +} + +func (s *ShardDiff) IsValid(col int, row int) bool { + _, ok := s.null[pair{col: col, row: uint64(row)}] + return !ok +} + +type Sharder struct { + shards map[uint64]*ShardDiff + index string + schema *arrow.Schema + // Standard input/output + log logger.Logger +} + +func (s *Sharder) GetShard(shard uint64) (*ShardDiff, error) { + f, ok := s.shards[shard] + if ok { + return f, nil + } + f, err := NewShardDiff(shard, s.log) + f.Setup(s.schema) + if err != nil { + return nil, err + } + s.shards[shard] = f + return f, nil +} + +func (s *Sharder) Store(schema *arrow.Schema, client *client.Client) error { + for _, f := range s.shards { + err := f.Store(s.index, schema, client) + if err != nil { + return err + } + } + return nil +} diff --git a/ctl/dataframe-csv-loader_test.go b/ctl/dataframe-csv-loader_test.go new file mode 100644 index 000000000..b3cac1ac1 --- /dev/null +++ b/ctl/dataframe-csv-loader_test.go @@ -0,0 +1,74 @@ +// Copyright 2021 Molecula Corp. All rights reserved. +package ctl + +import ( + "context" + "io" + "testing" + + pilosa "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/logger" + "github.com/molecula/featurebase/v3/server" + "github.com/molecula/featurebase/v3/test" + "github.com/molecula/featurebase/v3/testhook" +) + +func TestDataframeCsvLoaderCommand(t *testing.T) { + cluster := test.MustRunCluster(t, 1, []server.CommandOption{server.OptCommandServerOptions(pilosa.OptServerIsDataframeEnabled(true))}) + defer cluster.Close() + cmd := cluster.GetNode(0) + t.Run("basic", func(t *testing.T) { + cmLog := logger.NewStandardLogger(io.Discard) + cm := NewDataframeCsvLoaderCommand(cmLog) + file, err := testhook.TempFile(t, "import.csv") + if err != nil { + t.Fatalf("creating tempfile: %v", err) + } + _, err = file.Write([]byte("id,val__I\n1,2\n3,4\n5,6")) + if err != nil { + t.Fatalf("writing to tempfile: %v", err) + } + ctx := context.Background() + if err != nil { + t.Fatal(err) + } + index := "non-keyed" + cmd.API.CreateIndex(ctx, index, pilosa.IndexOptions{Keys: false}) + + 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) + } + }) + t.Run("keyed", func(t *testing.T) { + cmLog := logger.NewStandardLogger(io.Discard) + cm := NewDataframeCsvLoaderCommand(cmLog) + file, err := testhook.TempFile(t, "import_key.csv") + if err != nil { + t.Fatalf("creating tempfile: %v", err) + } + _, err = file.Write([]byte("id,val__I\nA,2\nB,4\nC,6")) + if err != nil { + t.Fatalf("writing to tempfile: %v", err) + } + ctx := context.Background() + if err != nil { + t.Fatal(err) + } + index := "keyed" + 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) + } + }) +}