feat: Use chunks while preparing batch (#11658)

candiduslynx · web-flow · commit 8294e7d28afe · 2023-06-19T10:53:15.000Z
diff --git a/plugins/destination/clickhouse/client/write.go b/plugins/destination/clickhouse/client/write.go
@@ -15,7 +15,7 @@ func (c *Client) WriteTableBatch(ctx context.Context, table *schema.Table, recor
 		return err
 	}
 
-	if err := values.BatchAddRecords(ctx, batch, table, records); err != nil {
+	if err := values.BatchAddRecords(ctx, batch, table.ToArrowSchema(), records); err != nil {
 		_ = batch.Abort()
 		return err
 	}
diff --git a/plugins/destination/clickhouse/typeconv/ch/values/batch.go b/plugins/destination/clickhouse/typeconv/ch/values/batch.go
@@ -5,18 +5,18 @@ import (
 
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 	"github.com/apache/arrow/go/v13/arrow"
-	"github.com/cloudquery/plugin-sdk/v3/schema"
+	"github.com/apache/arrow/go/v13/arrow/array"
 	"golang.org/x/sync/errgroup"
 )
 
-func BatchAddRecords(ctx context.Context, batch driver.Batch, table *schema.Table, records []arrow.Record) error {
+func BatchAddRecords(ctx context.Context, batch driver.Batch, sc *arrow.Schema, records []arrow.Record) error {
+	table := array.NewTableFromRecords(sc, records)
 	eg, _ := errgroup.WithContext(ctx)
-	for n := range table.Columns {
-		n := n
+	for n := 0; n < int(table.NumCols()); n++ {
+		column, chunks := batch.Column(n), table.Column(n).Data().Chunks()
 		eg.Go(func() error {
-			column := batch.Column(n)
-			for i := range records {
-				data, err := FromArray(records[i].Column(n))
+			for _, chunk := range chunks {
+				data, err := FromArray(chunk)
 				if err != nil {
 					return err
 				}

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ func (c Client) WriteTableBatch(ctx context.Context, table schema.Table, recor`
`15`	`15`	`return err`
`16`	`16`	`}`
`17`	`17`
`18`		`- if err := values.BatchAddRecords(ctx, batch, table, records); err != nil {`
	`18`	`+ if err := values.BatchAddRecords(ctx, batch, table.ToArrowSchema(), records); err != nil {`
`19`	`19`	`_ = batch.Abort()`
`20`	`20`	`return err`
`21`	`21`	`}`