Repository: pingcap/tidb-lightning
Pull request #520 (2020-12-10 08:29:03)

Sort index rather than insert it into skiplist
<!--
Thank you for contributing to TiDB-Lightning! Please read the [CONTRIBUTING](https://github.com/pingcap/tidb-lightning/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->
Improve the speed of writing index to pebble.


### What is changed and how it works?
This PR allows IndexEngine store more data in memory, and then flush it when it exceed the size limit of memtable.
Pebble will sort the WriteBatch and flush it directly if the size of WriteBatch exceed the half of memtable.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test
 - Manual test (add detailed scripts or steps below)


Related changes

 - Need to cherry-pick to the release branch


CLAassistant: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/pingcap/tidb-lightning?pullRequest=520) <br/>Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our [Contributor License Agreement](https://cla-assistant.io/pingcap/tidb-lightning?pullRequest=520) before we can accept your contribution.<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/pingcap/tidb-lightning?pullRequest=520) it.</sub>
glorv: @kennytm PTAL
glorv: LGTM
kennytm: LGTM

Patch:
diff --git a/lightning/backend/backend.go b/lightning/backend/backend.go
index d5ab891..965fa5b 100644
--- a/lightning/backend/backend.go
+++ b/lightning/backend/backend.go
@@ -95,10 +95,6 @@ type AbstractBackend interface {
 	// RetryImportDelay returns the duration to sleep when retrying an import
 	RetryImportDelay() time.Duration
 
-	// MaxChunkSize returns the maximum size acceptable by the backend. The
-	// value will be used in `Rows.SplitIntoChunks`.
-	MaxChunkSize() int
-
 	// ShouldPostProcess returns whether KV-specific post-processing should be
 	// performed for this backend. Post-processing includes checksum and analyze.
 	ShouldPostProcess() bool
@@ -108,15 +104,6 @@ type AbstractBackend interface {
 
 	OpenEngine(ctx context.Context, engineUUID uuid.UUID) error
 
-	WriteRows(
-		ctx context.Context,
-		engineUUID uuid.UUID,
-		tableName string,
-		columnNames []string,
-		commitTS uint64,
-		rows Rows,
-	) error
-
 	CloseEngine(ctx context.Context, engineUUID uuid.UUID) error
 
 	ImportEngine(ctx context.Context, engineUUID uuid.UUID) error
@@ -140,6 +127,8 @@ type AbstractBackend interface {
 	//     * Offset (must be 0, 1, 2, ...)
 	//  - PKIsHandle (true = do not generate _tidb_rowid)
 	FetchRemoteTableModels(ctx context.Context, schemaName string) ([]*model.TableInfo, error)
+
+	LocalWriter(ctx context.Context, engineUUID uuid.UUID) (EngineWriter, error)
 }
 
 func fetchRemoteTableModelsFromTLS(ctx context.Context, tls *common.TLS, schema string) ([]*model.TableInfo, error) {
@@ -184,6 +173,12 @@ type ClosedEngine struct {
 	engine
 }
 
+type LocalEngineWriter struct {
+	writer    EngineWriter
+	tableName string
+	ts        uint64
+}
+
 func MakeBackend(ab AbstractBackend) Backend {
 	return Backend{abstract: ab}
 }
@@ -265,25 +260,32 @@ func (engine *OpenedEngine) Flush() error {
 
 // WriteRows writes a collection of encoded rows into the engine.
 func (engine *OpenedEngine) WriteRows(ctx context.Context, columnNames []string, rows Rows) error {
-	var err error
+	writer, err := engine.backend.LocalWriter(ctx, engine.uuid)
+	if err != nil {
+		return err
+	}
+	if err = writer.AppendRows(ctx, engine.tableName, columnNames, engine.ts, rows); err != nil {
+		writer.Close()
+		return err
+	}
+	return writer.Close()
+}
 
-outside:
-	for _, r := range rows.SplitIntoChunks(engine.backend.MaxChunkSize()) {
-		for i := 0; i < maxRetryTimes; i++ {
-			err = engine.backend.WriteRows(ctx, engine.uuid, engine.tableName, columnNames, engine.ts, r)
-			switch {
-			case err == nil:
-				continue outside
-			case common.IsRetryableError(err):
-				// retry next loop
-			default:
-				return err
-			}
-		}
-		return errors.Annotatef(err, "[%s] write rows reach max retry %d and still failed", engine.tableName, maxRetryTimes)
+func (engine *OpenedEngine) LocalWriter(ctx context.Context) (*LocalEngineWriter, error) {
+	w, err := engine.backend.LocalWriter(ctx, engine.uuid)
+	if err != nil {
+		return nil, err
 	}
+	return &LocalEngineWriter{writer: w, ts: engine.ts, tableName: engine.tableName}, nil
+}
 
-	return nil
+// WriteRows writes a collection of encoded rows into the engine.
+func (w *LocalEngineWriter) WriteRows(ctx context.Context, columnNames []string, rows Rows) error {
+	return w.writer.AppendRows(ctx, w.tableName, columnNames, w.ts, rows)
+}
+
+func (w *LocalEngineWriter) Close() error {
+	return w.writer.Close()
 }
 
 // UnsafeCloseEngine closes the engine without first opening it.
@@ -388,3 +390,14 @@ type Rows interface {
 	// capacity with the current instance. The typical usage is `x = x.Clear()`.
 	Clear() Rows
 }
+
+type EngineWriter interface {
+	AppendRows(
+		ctx context.Context,
+		tableName string,
+		columnNames []string,
+		commitTS uint64,
+		rows Rows,
+	) error
+	Close() error
+}
diff --git a/lightning/backend/importer.go b/lightning/backend/importer.go
index a344e71..231b64c 100644
--- a/lightning/backend/importer.go
+++ b/lightning/backend/importer.go
@@ -142,6 +142,10 @@ func (importer *importer) CloseEngine(ctx context.Context, engineUUID uuid.UUID)
 	return nil
 }
 
+func (importer *importer) Flush(_ context.Context, _ uuid.UUID) error {
+	return nil
+}
+
 func (importer *importer) ImportEngine(ctx context.Context, engineUUID uuid.UUID) error {
 	req := &kv.ImportEngineRequest{
 		Uuid:   engineUUID[:],
@@ -168,6 +172,31 @@ func (importer *importer) WriteRows(
 	columnNames []string,
 	ts uint64,
 	rows Rows,
+) (finalErr error) {
+	var err error
+outside:
+	for _, r := range rows.SplitIntoChunks(importer.MaxChunkSize()) {
+		for i := 0; i < maxRetryTimes; i++ {
+			err = importer.WriteRowsToImporter(ctx, engineUUID, ts, r)
+			switch {
+			case err == nil:
+				continue outside
+			case common.IsRetryableError(err):
+				// retry next loop
+			default:
+				return err
+			}
+		}
+		return errors.Annotatef(err, "[%s] write rows reach max retry %d and still failed", tableName, maxRetryTimes)
+	}
+	return nil
+}
+
+func (importer *importer) WriteRowsToImporter(
+	ctx context.Context,
+	engineUUID uuid.UUID,
+	ts uint64,
+	rows Rows,
 ) (finalErr error) {
 	kvs := rows.(kvPairs)
 	if len(kvs) == 0 {
@@ -325,3 +354,20 @@ func checkVersion(component string, expected, actual semver.Version) error {
 func (importer *importer) FetchRemoteTableModels(ctx context.Context, schema string) ([]*model.TableInfo, error) {
 	return fetchRemoteTableModelsFromTLS(ctx, importer.tls, schema)
 }
+
+func (importer *importer) LocalWriter(ctx context.Context, engineUUID uuid.UUID) (EngineWriter, error) {
+	return &ImporterWriter{importer: importer, engineUUID: engineUUID}, nil
+}
+
+type ImporterWriter struct {
+	importer   *importer
+	engineUUID uuid.UUID
+}
+
+func (w *ImporterWriter) Close() error {
+	return nil
+}
+
+func (w *ImporterWriter) AppendRows(ctx context.Context, tableName string, columnNames []string, ts uint64, rows Rows) error {
+	return w.importer.WriteRows(ctx, w.engineUUID, tableName, columnNames, ts, rows)
+}
diff --git a/lightning/backend/local.go b/lightning/backend/local.go
index 132d690..7351473 100644
--- a/lightning/backend/local.go
+++ b/lightning/backend/local.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"encoding/binary"
 	"encoding/json"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"os"
@@ -29,6 +30,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/pebble"
+	"github.com/cockroachdb/pebble/sstable"
 	"github.com/coreos/go-semver/semver"
 	"github.com/google/btree"
 	"github.com/google/uuid"
@@ -68,7 +70,7 @@ const (
 	gRPCKeepAliveTimeout = 3 * time.Second
 	gRPCBackOffMaxDelay  = 3 * time.Second
 
-	LocalMemoryTableSize = 512 << 20
+	LocalMemoryTableSize = 128 << 20
 
 	// See: https://github.com/tikv/tikv/blob/e030a0aae9622f3774df89c62f21b2171a72a69e/etc/config-template.toml#L360
 	regionMaxKeyCount = 1_440_000
@@ -1196,50 +1198,29 @@ func (local *local) FetchRemoteTableModels(ctx context.Context, schemaName strin
 	return fetchRemoteTableModelsFromTLS(ctx, local.tls, schemaName)
 }
 
-func (local *local) WriteRows(
-	ctx context.Context,
-	engineUUID uuid.UUID,
-	tableName string,
-	columnNames []string,
-	ts uint64,
-	rows Rows,
-) (finalErr error) {
-	kvs := rows.(kvPairs)
-	if len(kvs) == 0 {
-		return nil
-	}
+func (local *local) MakeEmptyRows() Rows {
+	return kvPairs(nil)
+}
+
+func (local *local) NewEncoder(tbl table.Table, options *SessionOptions) (Encoder, error) {
+	return NewTableKVEncoder(tbl, options)
+}
 
+func (local *local) LocalWriter(ctx context.Context, engineUUID uuid.UUID) (EngineWriter, error) {
 	e, ok := local.engines.Load(engineUUID)
 	if !ok {
-		return errors.Errorf("could not find engine for %s", engineUUID.String())
+		return nil, errors.Errorf("could not find engine for %s", engineUUID.String())
 	}
 	engineFile := e.(*LocalFile)
-
-	// write to pebble to make them sorted
-	wb := engineFile.db.NewBatch()
-	defer wb.Close()
-	wo := &pebble.WriteOptions{Sync: false}
-
-	size := int64(0)
-	for _, pair := range kvs {
-		wb.Set(pair.Key, pair.Val, wo)
-		size += int64(len(pair.Key) + len(pair.Val))
-	}
-	if err := wb.Commit(wo); err != nil {
-		return err
-	}
-	atomic.AddInt64(&engineFile.Length, int64(len(kvs)))
-	atomic.AddInt64(&engineFile.TotalSize, size)
-	engineFile.Ts = ts
-	return
-}
-
-func (local *local) MakeEmptyRows() Rows {
-	return kvPairs(nil)
+	return openLocalWriter(engineFile, local.localStoreDir, LocalMemoryTableSize), nil
 }
 
-func (local *local) NewEncoder(tbl table.Table, options *SessionOptions) (Encoder, error) {
-	return NewTableKVEncoder(tbl, options)
+func openLocalWriter(f *LocalFile, sstDir string, memtableSizeLimit int64) *LocalWriter {
+	kvsChan := make(chan []common.KvPair, 1024)
+	w := &LocalWriter{sstDir: sstDir, kvsChan: kvsChan, local: f, memtableSizeLimit: memtableSizeLimit}
+	w.consumeWg.Add(1)
+	go w.writeRowsLoop()
+	return w
 }
 
 func (local *local) isIngestRetryable(
@@ -1513,3 +1494,180 @@ func (s *sizeProperties) iter(f func(p *rangeProperty) bool) {
 		return f(prop)
 	})
 }
+
+type LocalWriter struct {
+	writeErr           common.OnceError
+	local              *LocalFile
+	lastKey            []byte
+	consumeWg          sync.WaitGroup
+	kvsChan            chan []common.KvPair
+	sstDir             string
+	memtableSizeLimit  int64
+	writeBatch         []common.KvPair
+	isWriteBatchSorted bool
+}
+
+// If this method return false, it would not change `w.lastKey`
+func (w *LocalWriter) isSorted(kvs []common.KvPair) bool {
+	lastKey := w.lastKey
+	for _, pair := range kvs {
+		if len(lastKey) > 0 && bytes.Compare(lastKey, pair.Key) >= 0 {
+			return false
+		}
+		lastKey = pair.Key
+	}
+	w.lastKey = append(w.lastKey[:0], lastKey...)
+	return true
+}
+
+func (w *LocalWriter) AppendRows(ctx context.Context, tableName string, columnNames []string, ts uint64, rows Rows) error {
+	kvs := rows.(kvPairs)
+	if len(kvs) == 0 {
+		return nil
+	}
+	w.kvsChan <- kvs
+	w.local.Ts = ts
+	return nil
+}
+
+func (w *LocalWriter) Close() error {
+	close(w.kvsChan)
+	w.consumeWg.Wait()
+	return w.writeErr.Get()
+}
+
+func (w *LocalWriter) writeRowsLoop() {
+	batchSize := int64(0)
+	totalSize := int64(0)
+	totalCount := int64(0)
+	w.isWriteBatchSorted = true
+	var writer *sstable.Writer = nil
+	//var wb *pebble.Batch = nil
+	var filePath string
+	defer w.consumeWg.Done()
+	var err error
+	for kvs := range w.kvsChan {
+		hasSort := w.isSorted(kvs)
+		if totalCount > 1000 && hasSort {
+			for _, pair := range kvs {
+				totalSize += int64(len(pair.Key) + len(pair.Val))
+			}
+			if writer == nil {
+				writer, filePath, err = w.createWriter()
+				if err != nil {
+					w.writeErr.Set(err)
+					return
+				}
+				if len(w.writeBatch) > 0 && w.isWriteBatchSorted {
+					if err := writeKVs(writer, w.writeBatch); err != nil {
+						w.writeErr.Set(err)
+						return
+					}
+					w.isWriteBatchSorted = true
+					w.writeBatch = w.writeBatch[:0]
+					totalSize += batchSize
+					batchSize = 0
+				}
+			}
+			if err := writeKVs(writer, kvs); err != nil {
+				w.writeErr.Set(err)
+				return
+			}
+		} else {
+			for _, pair := range kvs {
+				batchSize += int64(len(pair.Key) + len(pair.Val))
+			}
+			if !hasSort {
+				w.isWriteBatchSorted = false
+			}
+			w.writeBatch = append(w.writeBatch, kvs...)
+
+			if batchSize > w.memtableSizeLimit {
+				if err := w.flushKVs(); err != nil {
+					w.writeErr.Set(err)
+					return
+				}
+				if writer == nil {
+					w.lastKey = w.lastKey[:0]
+				}
+				totalSize += batchSize
+				batchSize = 0
+				w.isWriteBatchSorted = true
+			}
+		}
+		totalCount += int64(len(kvs))
+	}
+
+	atomic.AddInt64(&w.local.Length, totalCount)
+	if batchSize > 0 {
+		if err := w.flushKVs(); err != nil {
+			w.writeErr.Set(err)
+			return
+		}
+		totalSize += batchSize
+		log.L().Info("write data by sort index", zap.Int64("bytes", totalSize))
+	}
+	if writer != nil {
+		if err := writer.Close(); err != nil {
+			w.writeErr.Set(err)
+			return
+		}
+		if err := w.local.db.Ingest([]string{filePath}); err != nil {
+			w.writeErr.Set(err)
+			return
+		}
+		log.L().Info("write data by sst writer", zap.Int64("bytes", totalSize))
+	}
+	atomic.AddInt64(&w.local.TotalSize, totalSize)
+}
+
+func (w *LocalWriter) flushKVs() error {
+	writer, filePath, err := w.createWriter()
+	if err != nil {
+		return err
+	}
+	if !w.isWriteBatchSorted {
+		sort.Slice(w.writeBatch, func(i, j int) bool {
+			return bytes.Compare(w.writeBatch[i].Key, w.writeBatch[j].Key) < 0
+		})
+	}
+	if err := writeKVs(writer, w.writeBatch); err != nil {
+		return err
+	}
+	if err := writer.Close(); err != nil {
+		return err
+	}
+	w.writeBatch = w.writeBatch[:0]
+	return w.local.db.Ingest([]string{filePath})
+}
+
+func (w *LocalWriter) createWriter() (*sstable.Writer, string, error) {
+	filePath := filepath.Join(w.sstDir, fmt.Sprintf("%s.sst", uuid.New()))
+	f, err := os.Create(filePath)
+	if err != nil {
+		return nil, filePath, err
+	}
+	writer := sstable.NewWriter(f, sstable.WriterOptions{
+		TablePropertyCollectors: []func() pebble.TablePropertyCollector{
+			func() pebble.TablePropertyCollector {
+				return newRangePropertiesCollector()
+			},
+		},
+		BlockSize: 16 * 1024,
+	})
+	return writer, filePath, nil
+}
+
+func writeKVs(writer *sstable.Writer, kvs []common.KvPair) error {
+	internalKey := sstable.InternalKey{
+		UserKey: []byte{},
+		Trailer: uint64(sstable.InternalKeyKindSet),
+	}
+	for _, p := range kvs {
+		internalKey.UserKey = p.Key
+		if err := writer.Add(internalKey, p.Val); err != nil {
+			return err
+		}
+	}
+	return nil
+}
diff --git a/lightning/backend/tidb.go b/lightning/backend/tidb.go
index 32c2d9a..bf81e0f 100644
--- a/lightning/backend/tidb.go
+++ b/lightning/backend/tidb.go
@@ -316,7 +316,27 @@ func (be *tidbBackend) ImportEngine(context.Context, uuid.UUID) error {
 	return nil
 }
 
-func (be *tidbBackend) WriteRows(ctx context.Context, _ uuid.UUID, tableName string, columnNames []string, _ uint64, r Rows) error {
+func (be *tidbBackend) WriteRows(ctx context.Context, _ uuid.UUID, tableName string, columnNames []string, _ uint64, rows Rows) error {
+	var err error
+outside:
+	for _, r := range rows.SplitIntoChunks(be.MaxChunkSize()) {
+		for i := 0; i < maxRetryTimes; i++ {
+			err = be.WriteRowsToDB(ctx, tableName, columnNames, r)
+			switch {
+			case err == nil:
+				continue outside
+			case common.IsRetryableError(err):
+				// retry next loop
+			default:
+				return err
+			}
+		}
+		return errors.Annotatef(err, "[%s] write rows reach max retry %d and still failed", tableName, maxRetryTimes)
+	}
+	return nil
+}
+
+func (be *tidbBackend) WriteRowsToDB(ctx context.Context, tableName string, columnNames []string, r Rows) error {
 	rows := r.(tidbRows)
 	if len(rows) == 0 {
 		return nil
@@ -498,3 +518,20 @@ func (be *tidbBackend) FetchRemoteTableModels(ctx context.Context, schemaName st
 	})
 	return
 }
+
+func (be *tidbBackend) LocalWriter(ctx context.Context, engineUUID uuid.UUID) (EngineWriter, error) {
+	return &TiDBWriter{be: be, engineUUID: engineUUID}, nil
+}
+
+type TiDBWriter struct {
+	be         *tidbBackend
+	engineUUID uuid.UUID
+}
+
+func (w *TiDBWriter) Close() error {
+	return nil
+}
+
+func (w *TiDBWriter) AppendRows(ctx context.Context, tableName string, columnNames []string, arg1 uint64, rows Rows) error {
+	return w.be.WriteRows(ctx, w.engineUUID, tableName, columnNames, arg1, rows)
+}
diff --git a/lightning/config/config.go b/lightning/config/config.go
index 4865a01..d1afdfd 100644
--- a/lightning/config/config.go
+++ b/lightning/config/config.go
@@ -361,7 +361,7 @@ func NewConfig() *Config {
 		TikvImporter: TikvImporter{
 			Backend:         BackendImporter,
 			OnDuplicate:     ReplaceOnDup,
-			MaxKVPairs:      32,
+			MaxKVPairs:      4096,
 			SendKVPairs:     32768,
 			RegionSplitSize: SplitRegionSize,
 		},
diff --git a/lightning/restore/restore.go b/lightning/restore/restore.go
index 2364723..101cbcd 100644
--- a/lightning/restore/restore.go
+++ b/lightning/restore/restore.go
@@ -1063,12 +1063,16 @@ func (t *TableRestore) restoreEngine(
 		closedEngine, err := rc.backend.UnsafeCloseEngine(ctx, t.tableName, engineID)
 		// If any error occurred, recycle worker immediately
 		if err != nil {
-			rc.closedEngineLimit.Recycle(w)
 			return closedEngine, nil, errors.Trace(err)
 		}
 		return closedEngine, w, nil
 	}
 
+	indexWriter, err := indexEngine.LocalWriter(ctx)
+	if err != nil {
+		return nil, nil, errors.Trace(err)
+	}
+
 	logTask := t.logger.With(zap.Int32("engineNumber", engineID)).Begin(zap.InfoLevel, "encode kv data and write")
 
 	dataEngine, err := rc.backend.OpenEngine(ctx, t.tableName, engineID)
@@ -1108,6 +1112,10 @@ func (t *TableRestore) restoreEngine(
 
 		restoreWorker := rc.regionWorkers.Apply()
 		wg.Add(1)
+		dataWriter, err := dataEngine.LocalWriter(ctx)
+		if err != nil {
+			return nil, nil, errors.Trace(err)
+		}
 		go func(w *worker.Worker, cr *chunkRestore) {
 			// Restore a chunk.
 			defer func() {
@@ -1116,17 +1124,25 @@ func (t *TableRestore) restoreEngine(
 				rc.regionWorkers.Recycle(w)
 			}()
 			metric.ChunkCounter.WithLabelValues(metric.ChunkStateRunning).Inc()
-			err := cr.restore(ctx, t, engineID, dataEngine, indexEngine, rc)
+			if err == nil {
+				err = cr.restore(ctx, t, engineID, dataWriter, indexWriter, rc)
+			}
+			if err == nil {
+				err = dataWriter.Close()
+			}
 			if err == nil {
 				metric.ChunkCounter.WithLabelValues(metric.ChunkStateFinished).Inc()
-				return
+			} else {
+				metric.ChunkCounter.WithLabelValues(metric.ChunkStateFailed).Inc()
+				chunkErr.Set(err)
 			}
-			metric.ChunkCounter.WithLabelValues(metric.ChunkStateFailed).Inc()
-			chunkErr.Set(err)
 		}(restoreWorker, cr)
 	}
 
 	wg.Wait()
+	if err := indexWriter.Close(); err != nil {
+		return nil, nil, errors.Trace(err)
+	}
 
 	// Report some statistics into the log for debugging.
 	totalKVSize := uint64(0)
@@ -1162,6 +1178,7 @@ func (t *TableRestore) restoreEngine(
 			rc.closedEngineLimit.Recycle(dataWorker)
 			return nil, nil, errors.Trace(err)
 		}
+
 		// Currently we write all the checkpoints after data&index engine are flushed.
 		for _, chunk := range cp.Chunks {
 			saveCheckpoint(rc, t, engineID, chunk)
@@ -1754,12 +1771,10 @@ func (cr *chunkRestore) deliverLoop(
 	kvsCh <-chan []deliveredKVs,
 	t *TableRestore,
 	engineID int32,
-	dataEngine, indexEngine *kv.OpenedEngine,
+	dataEngine, indexEngine *kv.LocalEngineWriter,
 	rc *RestoreController,
 ) (deliverTotalDur time.Duration, err error) {
 	var channelClosed bool
-	dataKVs := rc.backend.MakeEmptyRows()
-	indexKVs := rc.backend.MakeEmptyRows()
 
 	deliverLogger := t.logger.With(
 		zap.Int32("engineNumber", engineID),
@@ -1777,8 +1792,11 @@ func (cr *chunkRestore) deliverLoop(
 		offset := cr.chunk.Chunk.Offset
 		rowID := cr.chunk.Chunk.PrevRowIDMax
 		// Fetch enough KV pairs from the source.
+		dataKVs := rc.backend.MakeEmptyRows()
+		indexKVs := rc.backend.MakeEmptyRows()
+
 	populate:
-		for dataChecksum.SumSize()+indexChecksum.SumSize() < minDeliverBytes {
+		for dataChecksum.SumSize() < minDeliverBytes {
 			select {
 			case kvPacket = <-kvsCh:
 				if len(kvPacket) == 0 {
@@ -1817,9 +1835,6 @@ func (cr *chunkRestore) deliverLoop(
 		metric.BlockDeliverKVPairsHistogram.WithLabelValues(metric.BlockDeliverKindData).Observe(float64(dataChecksum.SumKVS()))
 		metric.BlockDeliverKVPairsHistogram.WithLabelValues(metric.BlockDeliverKindIndex).Observe(float64(indexChecksum.SumKVS()))
 
-		dataKVs = dataKVs.Clear()
-		indexKVs = indexKVs.Clear()
-
 		// Update the table, and save a checkpoint.
 		// (the write to the importer is effective immediately, thus update these here)
 		// No need to apply a lock since this is the only thread updating these variables.
@@ -1978,7 +1993,7 @@ func (cr *chunkRestore) restore(
 	ctx context.Context,
 	t *TableRestore,
 	engineID int32,
-	dataEngine, indexEngine *kv.OpenedEngine,
+	dataEngine, indexEngine *kv.LocalEngineWriter,
 	rc *RestoreController,
 ) error {
 	// Create the encoder.
diff --git a/mock/backend.go b/mock/backend.go
index 1f0c240..7e73a98 100644
--- a/mock/backend.go
+++ b/mock/backend.go
@@ -1,7 +1,7 @@
 // Code generated by MockGen. DO NOT EDIT.
-// Source: github.com/pingcap/tidb-lightning/lightning/backend (interfaces: AbstractBackend,Encoder,Rows,Row)
+// Source: github.com/pingcap/tidb-lightning/lightning/backend (interfaces: AbstractBackend,Encoder,Rows,Row,EngineWriter)
 
-// $ mockgen -package mock -mock_names 'AbstractBackend=MockBackend' github.com/pingcap/tidb-lightning/lightning/backend AbstractBackend,Encoder,Rows,Row
+// $ mockgen -package mock -mock_names 'AbstractBackend=MockBackend' github.com/pingcap/tidb-lightning/lightning/backend AbstractBackend,Encoder,Rows,Row,EngineWriter
 
 // Package mock is a generated GoMock package.
 package mock
@@ -126,6 +126,21 @@ func (mr *MockBackendMockRecorder) ImportEngine(arg0, arg1 interface{}) *gomock.
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ImportEngine", reflect.TypeOf((*MockBackend)(nil).ImportEngine), arg0, arg1)
 }
 
+// LocalWriter mocks base method
+func (m *MockBackend) LocalWriter(arg0 context.Context, arg1 uuid.UUID) (backend.EngineWriter, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "LocalWriter", arg0, arg1)
+	ret0, _ := ret[0].(backend.EngineWriter)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// LocalWriter indicates an expected call of LocalWriter
+func (mr *MockBackendMockRecorder) LocalWriter(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LocalWriter", reflect.TypeOf((*MockBackend)(nil).LocalWriter), arg0, arg1)
+}
+
 // MakeEmptyRows mocks base method
 func (m *MockBackend) MakeEmptyRows() backend.Rows {
 	m.ctrl.T.Helper()
@@ -140,20 +155,6 @@ func (mr *MockBackendMockRecorder) MakeEmptyRows() *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MakeEmptyRows", reflect.TypeOf((*MockBackend)(nil).MakeEmptyRows))
 }
 
-// MaxChunkSize mocks base method
-func (m *MockBackend) MaxChunkSize() int {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "MaxChunkSize")
-	ret0, _ := ret[0].(int)
-	return ret0
-}
-
-// MaxChunkSize indicates an expected call of MaxChunkSize
-func (mr *MockBackendMockRecorder) MaxChunkSize() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MaxChunkSize", reflect.TypeOf((*MockBackend)(nil).MaxChunkSize))
-}
-
 // NewEncoder mocks base method
 func (m *MockBackend) NewEncoder(arg0 table.Table, arg1 *backend.SessionOptions) (backend.Encoder, error) {
 	m.ctrl.T.Helper()
@@ -211,20 +212,6 @@ func (mr *MockBackendMockRecorder) ShouldPostProcess() *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ShouldPostProcess", reflect.TypeOf((*MockBackend)(nil).ShouldPostProcess))
 }
 
-// WriteRows mocks base method
-func (m *MockBackend) WriteRows(arg0 context.Context, arg1 uuid.UUID, arg2 string, arg3 []string, arg4 uint64, arg5 backend.Rows) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "WriteRows", arg0, arg1, arg2, arg3, arg4, arg5)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// WriteRows indicates an expected call of WriteRows
-func (mr *MockBackendMockRecorder) WriteRows(arg0, arg1, arg2, arg3, arg4, arg5 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WriteRows", reflect.TypeOf((*MockBackend)(nil).WriteRows), arg0, arg1, arg2, arg3, arg4, arg5)
-}
-
 // MockEncoder is a mock of Encoder interface
 type MockEncoder struct {
 	ctrl     *gomock.Controller
@@ -360,3 +347,54 @@ func (mr *MockRowMockRecorder) ClassifyAndAppend(arg0, arg1, arg2, arg3 interfac
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ClassifyAndAppend", reflect.TypeOf((*MockRow)(nil).ClassifyAndAppend), arg0, arg1, arg2, arg3)
 }
+
+// MockEngineWriter is a mock of EngineWriter interface
+type MockEngineWriter struct {
+	ctrl     *gomock.Controller
+	recorder *MockEngineWriterMockRecorder
+}
+
+// MockEngineWriterMockRecorder is the mock recorder for MockEngineWriter
+type MockEngineWriterMockRecorder struct {
+	mock *MockEngineWriter
+}
+
+// NewMockEngineWriter creates a new mock instance
+func NewMockEngineWriter(ctrl *gomock.Controller) *MockEngineWriter {
+	mock := &MockEngineWriter{ctrl: ctrl}
+	mock.recorder = &MockEngineWriterMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockEngineWriter) EXPECT() *MockEngineWriterMockRecorder {
+	return m.recorder
+}
+
+// AppendRows mocks base method
+func (m *MockEngineWriter) AppendRows(arg0 context.Context, arg1 string, arg2 []string, arg3 uint64, arg4 backend.Rows) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "AppendRows", arg0, arg1, arg2, arg3, arg4)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// AppendRows indicates an expected call of AppendRows
+func (mr *MockEngineWriterMockRecorder) AppendRows(arg0, arg1, arg2, arg3, arg4 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AppendRows", reflect.TypeOf((*MockEngineWriter)(nil).AppendRows), arg0, arg1, arg2, arg3, arg4)
+}
+
+// Close mocks base method
+func (m *MockEngineWriter) Close() error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Close")
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Close indicates an expected call of Close
+func (mr *MockEngineWriterMockRecorder) Close() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockEngineWriter)(nil).Close))
+}


Test patch:
diff --git a/lightning/backend/backend_test.go b/lightning/backend/backend_test.go
index 46f1c36..d241739 100644
--- a/lightning/backend/backend_test.go
+++ b/lightning/backend/backend_test.go
@@ -117,29 +117,28 @@ func (s *backendSuite) TestWriteEngine(c *C) {
 	ctx := context.Background()
 	engineUUID := uuid.MustParse("902efee3-a3f9-53d4-8c82-f12fb1900cd1")
 
-	rows0 := mock.NewMockRows(s.controller)
 	rows1 := mock.NewMockRows(s.controller)
 	rows2 := mock.NewMockRows(s.controller)
 
 	s.mockBackend.EXPECT().
 		OpenEngine(ctx, engineUUID).
 		Return(nil)
-	s.mockBackend.EXPECT().
-		MaxChunkSize().
-		Return(654321)
-	rows0.EXPECT().
-		SplitIntoChunks(654321).
-		Return([]kv.Rows{rows1, rows2})
-	s.mockBackend.EXPECT().
-		WriteRows(ctx, engineUUID, "`db`.`table`", []string{"c1", "c2"}, gomock.Any(), rows1).
+
+	mockWriter := mock.NewMockEngineWriter(s.controller)
+	s.mockBackend.EXPECT().LocalWriter(ctx, gomock.Any()).Return(mockWriter, nil).AnyTimes()
+	mockWriter.EXPECT().
+		AppendRows(ctx, "`db`.`table`", []string{"c1", "c2"}, gomock.Any(), rows1).
 		Return(nil)
-	s.mockBackend.EXPECT().
-		WriteRows(ctx, engineUUID, "`db`.`table`", []string{"c1", "c2"}, gomock.Any(), rows2).
+	mockWriter.EXPECT().Close().Return(nil).AnyTimes()
+	mockWriter.EXPECT().
+		AppendRows(ctx, "`db`.`table`", []string{"c1", "c2"}, gomock.Any(), rows2).
 		Return(nil)
 
 	engine, err := s.backend.OpenEngine(ctx, "`db`.`table`", 1)
 	c.Assert(err, IsNil)
-	err = engine.WriteRows(ctx, []string{"c1", "c2"}, rows0)
+	err = engine.WriteRows(ctx, []string{"c1", "c2"}, rows1)
+	c.Assert(err, IsNil)
+	err = engine.WriteRows(ctx, []string{"c1", "c2"}, rows2)
 	c.Assert(err, IsNil)
 }
 
@@ -149,10 +148,12 @@ func (s *backendSuite) TestWriteToEngineWithNothing(c *C) {
 
 	ctx := context.Background()
 	emptyRows := mock.NewMockRows(s.controller)
+	writer := mock.NewMockEngineWriter(s.controller)
 
 	s.mockBackend.EXPECT().OpenEngine(ctx, gomock.Any()).Return(nil)
-	s.mockBackend.EXPECT().MaxChunkSize().Return(654321)
-	emptyRows.EXPECT().SplitIntoChunks(654321).Return([]kv.Rows{})
+	writer.EXPECT().AppendRows(ctx, gomock.Any(), gomock.Any(), gomock.Any(), emptyRows).Return(nil)
+	writer.EXPECT().Close().Return(nil)
+	s.mockBackend.EXPECT().LocalWriter(ctx, gomock.Any()).Return(writer, nil)
 
 	engine, err := s.backend.OpenEngine(ctx, "`db`.`table`", 1)
 	c.Assert(err, IsNil)
@@ -181,12 +182,12 @@ func (s *backendSuite) TestWriteEngineFailed(c *C) {
 	rows := mock.NewMockRows(s.controller)
 
 	s.mockBackend.EXPECT().OpenEngine(ctx, gomock.Any()).Return(nil)
-	s.mockBackend.EXPECT().MaxChunkSize().Return(123)
-	rows.EXPECT().SplitIntoChunks(123).Return([]kv.Rows{rows})
-
-	s.mockBackend.EXPECT().
-		WriteRows(ctx, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), rows).
+	mockWriter := mock.NewMockEngineWriter(s.controller)
+	s.mockBackend.EXPECT().LocalWriter(ctx, gomock.Any()).Return(mockWriter, nil).AnyTimes()
+	mockWriter.EXPECT().
+		AppendRows(ctx, gomock.Any(), gomock.Any(), gomock.Any(), rows).
 		Return(errors.Annotate(context.Canceled, "fake unrecoverable write error"))
+	mockWriter.EXPECT().Close().Return(nil)
 
 	engine, err := s.backend.OpenEngine(ctx, "`db`.`table`", 1)
 	c.Assert(err, IsNil)
@@ -202,13 +203,12 @@ func (s *backendSuite) TestWriteBatchSendFailedWithRetry(c *C) {
 	rows := mock.NewMockRows(s.controller)
 
 	s.mockBackend.EXPECT().OpenEngine(ctx, gomock.Any()).Return(nil)
-	s.mockBackend.EXPECT().MaxChunkSize().Return(123)
-	rows.EXPECT().SplitIntoChunks(123).Return([]kv.Rows{rows})
-
-	s.mockBackend.EXPECT().
-		WriteRows(ctx, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), rows).
+	mockWriter := mock.NewMockEngineWriter(s.controller)
+	s.mockBackend.EXPECT().LocalWriter(ctx, gomock.Any()).Return(mockWriter, nil).AnyTimes()
+	mockWriter.EXPECT().AppendRows(ctx, gomock.Any(), gomock.Any(), gomock.Any(), rows).
 		Return(errors.New("fake recoverable write batch error")).
-		MinTimes(2)
+		MinTimes(1)
+	mockWriter.EXPECT().Close().Return(nil).MinTimes(1)
 
 	engine, err := s.backend.OpenEngine(ctx, "`db`.`table`", 1)
 	c.Assert(err, IsNil)
@@ -216,30 +216,6 @@ func (s *backendSuite) TestWriteBatchSendFailedWithRetry(c *C) {
 	c.Assert(err, ErrorMatches, ".*fake recoverable write batch error")
 }
 
-func (s *backendSuite) TestWriteBatchSendRecovered(c *C) {
-	s.setUpTest(c)
-	defer s.tearDownTest()
-
-	ctx := context.Background()
-	rows := mock.NewMockRows(s.controller)
-
-	s.mockBackend.EXPECT().OpenEngine(ctx, gomock.Any()).Return(nil)
-	s.mockBackend.EXPECT().MaxChunkSize().Return(123)
-	rows.EXPECT().SplitIntoChunks(123).Return([]kv.Rows{rows})
-
-	s.mockBackend.EXPECT().
-		WriteRows(ctx, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), rows).
-		Return(errors.New("fake recoverable write batch error"))
-	s.mockBackend.EXPECT().
-		WriteRows(ctx, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), rows).
-		Return(nil)
-
-	engine, err := s.backend.OpenEngine(ctx, "`db`.`table`", 1)
-	c.Assert(err, IsNil)
-	err = engine.WriteRows(ctx, nil, rows)
-	c.Assert(err, IsNil)
-}
-
 func (s *backendSuite) TestImportFailedNoRetry(c *C) {
 	s.setUpTest(c)
 	defer s.tearDownTest()
diff --git a/lightning/backend/local_test.go b/lightning/backend/local_test.go
index ab55cf1..ea3cb82 100644
--- a/lightning/backend/local_test.go
+++ b/lightning/backend/local_test.go
@@ -15,10 +15,15 @@ package backend
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
+	"fmt"
+	"github.com/pingcap/tidb-lightning/lightning/common"
 	"math"
 	"math/rand"
+	"os"
 	"path/filepath"
+	"sort"
 
 	"github.com/cockroachdb/pebble"
 	. "github.com/pingcap/check"
@@ -285,3 +290,104 @@ func (s *localSuite) TestRangePropertiesWithPebble(c *C) {
 
 	c.Assert(sstMetas[0][0].Properties.UserProperties, DeepEquals, props)
 }
+
+func testLocalWriter(c *C, needSort bool, partitialSort bool) {
+	dir := c.MkDir()
+	opt := &pebble.Options{
+		MemTableSize:             1024 * 1024,
+		MaxConcurrentCompactions: 16,
+		L0CompactionThreshold:    math.MaxInt32, // set to max try to disable compaction
+		L0StopWritesThreshold:    math.MaxInt32, // set to max try to disable compaction
+		DisableWAL:               true,
+		ReadOnly:                 false,
+	}
+	db, err := pebble.Open(filepath.Join(dir, "test"), opt)
+	c.Assert(err, IsNil)
+	tmpPath := filepath.Join(dir, "tmp")
+	err = os.Mkdir(tmpPath, 0755)
+	c.Assert(err, IsNil)
+	meta := localFileMeta{}
+	_, engineUUID := MakeUUID("ww", 0)
+	f := LocalFile{localFileMeta: meta, db: db, Uuid: engineUUID}
+	w := openLocalWriter(&f, tmpPath, 1024*1024)
+
+	ctx := context.Background()
+	//kvs := make(kvPairs, 1000)
+	var kvs kvPairs
+	value := make([]byte, 128)
+	for i := 0; i < 16; i++ {
+		binary.BigEndian.PutUint64(value[i*8:], uint64(i))
+	}
+	var keys [][]byte
+	for i := 1; i <= 20000; i++ {
+		var kv common.KvPair
+		kv.Key = make([]byte, 16)
+		kv.Val = make([]byte, 128)
+		copy(kv.Val, value)
+		key := rand.Intn(1000)
+		binary.BigEndian.PutUint64(kv.Key, uint64(key))
+		binary.BigEndian.PutUint64(kv.Key[8:], uint64(i))
+		kvs = append(kvs, kv)
+		keys = append(keys, kv.Key)
+	}
+	var rows1 kvPairs
+	var rows2 kvPairs
+	var rows3 kvPairs
+	rows4 := kvs[:12000]
+	if partitialSort {
+		sort.Slice(rows4, func(i, j int) bool {
+			return bytes.Compare(rows4[i].Key, rows4[j].Key) < 0
+		})
+		rows1 = rows4[:6000]
+		rows3 = rows4[6000:]
+		rows2 = kvs[12000:]
+	} else {
+		if needSort {
+			sort.Slice(kvs, func(i, j int) bool {
+				return bytes.Compare(kvs[i].Key, kvs[j].Key) < 0
+			})
+		}
+		rows1 = kvs[:6000]
+		rows2 = kvs[6000:12000]
+		rows3 = kvs[12000:]
+	}
+	err = w.AppendRows(ctx, "", []string{}, 1, rows1)
+	c.Assert(err, IsNil)
+	err = w.AppendRows(ctx, "", []string{}, 1, rows2)
+	c.Assert(err, IsNil)
+	err = w.AppendRows(ctx, "", []string{}, 1, rows3)
+	c.Assert(err, IsNil)
+	err = w.Close()
+	c.Assert(err, IsNil)
+	err = db.Flush()
+	c.Assert(err, IsNil)
+	o := &pebble.IterOptions{}
+	it := db.NewIter(o)
+
+	sort.Slice(keys, func(i, j int) bool {
+		return bytes.Compare(keys[i], keys[j]) < 0
+	})
+	c.Assert(int(f.Length), Equals, 20000)
+	c.Assert(int(f.TotalSize), Equals, 144*20000)
+	valid := it.SeekGE(keys[0])
+	c.Assert(valid, IsTrue)
+	for _, k := range keys {
+		c.Assert(it.Key(), DeepEquals, k)
+		it.Next()
+	}
+}
+
+func (s *localSuite) TestLocalWriterWithSort(c *C) {
+	fmt.Println("TestLocalWriterWithSort")
+	testLocalWriter(c, false, false)
+}
+
+func (s *localSuite) TestLocalWriterWithIngest(c *C) {
+	fmt.Println("TestLocalWriterWithIngest")
+	testLocalWriter(c, true, false)
+}
+
+func (s *localSuite) TestLocalWriterWithIngestUnsort(c *C) {
+	fmt.Println("TestLocalWriterWithIngestUnsort")
+	testLocalWriter(c, true, true)
+}
diff --git a/lightning/lightning_test.go b/lightning/lightning_test.go
index fe670a6..ec514ee 100644
--- a/lightning/lightning_test.go
+++ b/lightning/lightning_test.go
@@ -467,8 +467,8 @@ func (s *lightningServerSuite) TestCheckSystemRequirement(c *C) {
 		},
 	}
 
-	// with max open files 2048, the max table size will be: 524288MB
-	err := failpoint.Enable("github.com/pingcap/tidb-lightning/lightning/backend/GetRlimitValue", "return(2049)")
+	// with max open files 1024, the max table size will be: 65536MB
+	err := failpoint.Enable("github.com/pingcap/tidb-lightning/lightning/backend/GetRlimitValue", "return(8199)")
 	c.Assert(err, IsNil)
 	err = failpoint.Enable("github.com/pingcap/tidb-lightning/lightning/backend/SetRlimitError", "return(true)")
 	c.Assert(err, IsNil)
@@ -487,7 +487,7 @@ func (s *lightningServerSuite) TestCheckSystemRequirement(c *C) {
 	c.Assert(err, IsNil)
 
 	// the min rlimit should be bigger than the default min value (16384)
-	err = failpoint.Enable("github.com/pingcap/tidb-lightning/lightning/backend/GetRlimitValue", "return(2050)")
+	err = failpoint.Enable("github.com/pingcap/tidb-lightning/lightning/backend/GetRlimitValue", "return(8200)")
 	defer failpoint.Disable("github.com/pingcap/tidb-lightning/lightning/backend/GetRlimitValue")
 	c.Assert(err, IsNil)
 	err = checkSystemRequirement(cfg, dbMetas)
diff --git a/lightning/restore/restore_test.go b/lightning/restore/restore_test.go
index 92dcb55..9185a60 100644
--- a/lightning/restore/restore_test.go
+++ b/lightning/restore/restore_test.go
@@ -816,12 +816,20 @@ func (s *chunkRestoreSuite) TestDeliverLoopEmptyData(c *C) {
 
 	mockBackend.EXPECT().OpenEngine(ctx, gomock.Any()).Return(nil).Times(2)
 	mockBackend.EXPECT().MakeEmptyRows().Return(kv.MakeRowsFromKvPairs(nil)).AnyTimes()
-	mockBackend.EXPECT().MaxChunkSize().Return(10000).AnyTimes()
+	mockWriter := mock.NewMockEngineWriter(controller)
+	mockBackend.EXPECT().LocalWriter(ctx, gomock.Any()).Return(mockWriter, nil).AnyTimes()
+	mockWriter.EXPECT().
+		AppendRows(ctx, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(nil).AnyTimes()
 
 	dataEngine, err := importer.OpenEngine(ctx, s.tr.tableName, 0)
 	c.Assert(err, IsNil)
+	dataWriter, err := dataEngine.LocalWriter(ctx)
+	c.Assert(err, IsNil)
 	indexEngine, err := importer.OpenEngine(ctx, s.tr.tableName, -1)
 	c.Assert(err, IsNil)
+	indexWriter, err := indexEngine.LocalWriter(ctx)
+	c.Assert(err, IsNil)
 
 	// Deliver nothing.
 
@@ -830,7 +838,7 @@ func (s *chunkRestoreSuite) TestDeliverLoopEmptyData(c *C) {
 
 	kvsCh := make(chan []deliveredKVs, 1)
 	kvsCh <- []deliveredKVs{}
-	_, err = s.cr.deliverLoop(ctx, kvsCh, s.tr, 0, dataEngine, indexEngine, rc)
+	_, err = s.cr.deliverLoop(ctx, kvsCh, s.tr, 0, dataWriter, indexWriter, rc)
 	c.Assert(err, IsNil)
 }
 
@@ -848,17 +856,23 @@ func (s *chunkRestoreSuite) TestDeliverLoop(c *C) {
 
 	mockBackend.EXPECT().OpenEngine(ctx, gomock.Any()).Return(nil).Times(2)
 	mockBackend.EXPECT().MakeEmptyRows().Return(kv.MakeRowsFromKvPairs(nil)).AnyTimes()
-	mockBackend.EXPECT().MaxChunkSize().Return(10000).AnyTimes()
+	mockWriter := mock.NewMockEngineWriter(controller)
+	mockBackend.EXPECT().LocalWriter(ctx, gomock.Any()).Return(mockWriter, nil).AnyTimes()
 
 	dataEngine, err := importer.OpenEngine(ctx, s.tr.tableName, 0)
 	c.Assert(err, IsNil)
 	indexEngine, err := importer.OpenEngine(ctx, s.tr.tableName, -1)
 	c.Assert(err, IsNil)
 
+	dataWriter, err := dataEngine.LocalWriter(ctx)
+	c.Assert(err, IsNil)
+	indexWriter, err := indexEngine.LocalWriter(ctx)
+	c.Assert(err, IsNil)
+
 	// Set up the expected API calls to the data engine...
 
-	mockBackend.EXPECT().
-		WriteRows(ctx, gomock.Any(), s.tr.tableName, mockCols, gomock.Any(), kv.MakeRowsFromKvPairs([]common.KvPair{
+	mockWriter.EXPECT().
+		AppendRows(ctx, s.tr.tableName, mockCols, gomock.Any(), kv.MakeRowsFromKvPairs([]common.KvPair{
 			{
 				Key: []byte("txxxxxxxx_ryyyyyyyy"),
 				Val: []byte("value1"),
@@ -874,8 +888,8 @@ func (s *chunkRestoreSuite) TestDeliverLoop(c *C) {
 	//
 	// Note: This test assumes data engine is written before the index engine.
 
-	mockBackend.EXPECT().
-		WriteRows(ctx, gomock.Any(), s.tr.tableName, mockCols, gomock.Any(), kv.MakeRowsFromKvPairs([]common.KvPair{
+	mockWriter.EXPECT().
+		AppendRows(ctx, s.tr.tableName, mockCols, gomock.Any(), kv.MakeRowsFromKvPairs([]common.KvPair{
 			{
 				Key: []byte("txxxxxxxx_izzzzzzzz"),
 				Val: []byte("index1"),
@@ -914,7 +928,7 @@ func (s *chunkRestoreSuite) TestDeliverLoop(c *C) {
 	cfg := &config.Config{}
 	rc := &RestoreController{cfg: cfg, saveCpCh: saveCpCh, backend: importer}
 
-	_, err = s.cr.deliverLoop(ctx, kvsCh, s.tr, 0, dataEngine, indexEngine, rc)
+	_, err = s.cr.deliverLoop(ctx, kvsCh, s.tr, 0, dataWriter, indexWriter, rc)
 	c.Assert(err, IsNil)
 	c.Assert(saveCpCh, HasLen, 2)
 	c.Assert(s.cr.chunk.Chunk.Offset, Equals, int64(12))
@@ -1026,6 +1040,10 @@ func (s *chunkRestoreSuite) TestRestore(c *C) {
 	c.Assert(err, IsNil)
 	indexEngine, err := importer.OpenEngine(ctx, s.tr.tableName, -1)
 	c.Assert(err, IsNil)
+	dataWriter, err := dataEngine.LocalWriter(ctx)
+	c.Assert(err, IsNil)
+	indexWriter, err := indexEngine.LocalWriter(ctx)
+	c.Assert(err, IsNil)
 
 	// Expected API sequence
 	// (we don't care about the actual content, this would be checked in the integrated tests)
@@ -1049,7 +1067,7 @@ func (s *chunkRestoreSuite) TestRestore(c *C) {
 	// Now actually start the restore loop.
 
 	saveCpCh := make(chan saveCp, 2)
-	err = s.cr.restore(ctx, s.tr, 0, dataEngine, indexEngine, &RestoreController{
+	err = s.cr.restore(ctx, s.tr, 0, dataWriter, indexWriter, &RestoreController{
 		cfg:      s.cfg,
 		saveCpCh: saveCpCh,
 		backend:  importer,