mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-08-28 10:54:59 +00:00
Fb 1874 dataframe-csv-loader featurebase subcommand (#2341)
Embeded dataframe-csv-loader command as featurebase subservice
This commit is contained in:
tgruben
GitHub
parent
ac3ffac8e6
commit
f2a13c8bde
5 changed files with 503 additions and 10 deletions
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
37
cmd/dataframe-csv-loader.go
Normal file
37
cmd/dataframe-csv-loader.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -105,6 +105,7 @@ at https://docs.molecula.cloud/.
|
|||
rc.AddCommand(newKeygenCommand(logdest))
|
||||
rc.AddCommand(newCLICommand(logdest))
|
||||
rc.AddCommand(newDAXCommand(stderr))
|
||||
rc.AddCommand(newDataframeCsvLoaderCommand(logdest))
|
||||
|
||||
rc.SetOutput(stderr)
|
||||
return rc
|
||||
|
|
|
|||
391
ctl/dataframe-csv-loader.go
Normal file
391
ctl/dataframe-csv-loader.go
Normal file
|
|
@ -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
|
||||
}
|
||||
74
ctl/dataframe-csv-loader_test.go
Normal file
74
ctl/dataframe-csv-loader_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue