globalsort: use partitioned prefix to store intermediate files (#62905) (#62963)

close #62904

Author: ti-chi-bot

br/pkg/storage/memstore.go

@@ -69,9 +69,6 @@ func (s *MemStorage) DeleteFile(ctx context.Context, name string) error {
 	default:
 		// continue on
 	}
-	if !path.IsAbs(name) {
-		return errors.Errorf("file name is not an absolute path: %s", name)
-	}
 	s.rwm.Lock()
 	defer s.rwm.Unlock()
 	if _, ok := s.dataStore[name]; !ok {
@@ -102,9 +99,6 @@ func (s *MemStorage) WriteFile(ctx context.Context, name string, data []byte) er
 	default:
 		// continue on
 	}
-	if !path.IsAbs(name) {
-		return errors.Errorf("file name is not an absolute path: %s", name)
-	}
 	fileData := slices.Clone(data)
 	s.rwm.Lock()
 	defer s.rwm.Unlock()
@@ -128,9 +122,6 @@ func (s *MemStorage) ReadFile(ctx context.Context, name string) ([]byte, error)
 	default:
 		// continue on
 	}
-	if !path.IsAbs(name) {
-		return nil, errors.Errorf("file name is not an absolute path: %s", name)
-	}
 	theFile, ok := s.loadMap(name)
 	if !ok {
 		return nil, errors.Errorf("cannot find the file: %s", name)
@@ -148,9 +139,6 @@ func (s *MemStorage) FileExists(ctx context.Context, name string) (bool, error)
 	default:
 		// continue on
 	}
-	if !path.IsAbs(name) {
-		return false, errors.Errorf("file name is not an absolute path: %s", name)
-	}
 	_, ok := s.loadMap(name)
 	return ok, nil
 }
@@ -161,9 +149,6 @@ func (s *MemStorage) Open(ctx context.Context, filePath string, o *ReaderOption)
 	if err := ctx.Err(); err != nil {
 		return nil, errors.Trace(err)
 	}
-	if !path.IsAbs(filePath) {
-		return nil, errors.Errorf("file name is not an absolute path: %s", filePath)
-	}
 	theFile, ok := s.loadMap(filePath)
 	if !ok {
 		return nil, errors.Errorf("cannot find the file: %s", filePath)
@@ -248,9 +233,6 @@ func (s *MemStorage) Create(ctx context.Context, name string, _ *WriterOption) (
 	default:
 		// continue on
 	}
-	if !path.IsAbs(name) {
-		return nil, errors.Errorf("file name is not an absolute path: %s", name)
-	}
 	s.rwm.Lock()
 	defer s.rwm.Unlock()
 	theFile := new(memFile)
@@ -269,9 +251,6 @@ func (s *MemStorage) Rename(ctx context.Context, oldFileName, newFileName string
 	default:
 		// continue on
 	}
-	if !path.IsAbs(newFileName) {
-		return errors.Errorf("new file name is not an absolute path: %s", newFileName)
-	}
 	s.rwm.Lock()
 	defer s.rwm.Unlock()
 	theFile, ok := s.dataStore[oldFileName]

pkg/disttask/importinto/planner_test.go

@@ -344,9 +344,11 @@ func TestSplitForOneSubtask(t *testing.T) {
 			multiFileStat = s.MultipleFilesStats
 		}).
 		Build(store, "/mock-test", "0")
-	_, _, err = external.MockExternalEngineWithWriter(
-		store, writer, "/mock-test", keys, values,
-	)
+	for i := range keys {
+		err := writer.WriteRow(ctx, keys[i], values[i], nil)
+		require.NoError(t, err)
+	}
+	require.NoError(t, writer.Close(ctx))
 	require.NoError(t, err)
 	kvMeta := &external.SortedKVMeta{
 		StartKey:           keys[0],

pkg/executor/test/indexmergereadtest/BUILD.bazel

@@ -2,7 +2,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_test")
 
 go_test(
     name = "indexmergereadtest_test",
-    timeout = "short",
+    timeout = "moderate",
     srcs = [
         "index_merge_reader_test.go",
         "main_test.go",

pkg/lightning/backend/external/bench_test.go

@@ -15,12 +15,14 @@
 package external
 
 import (
+	"bytes"
 	"context"
 	"encoding/hex"
 	goerrors "errors"
 	"flag"
 	"fmt"
 	"io"
+	"strings"
 	"testing"
 	"time"
 
@@ -88,11 +90,9 @@ func writePlainFile(s *writeTestSuite) {
 }
 
 func cleanOldFiles(ctx context.Context, store storage.ExternalStorage, subDir string) {
-	dataFiles, statFiles, err := GetAllFileNames(ctx, store, subDir)
+	filenames, err := GetAllFileNames(ctx, store, subDir)
 	intest.AssertNoError(err)
-	err = store.DeleteFiles(ctx, dataFiles)
-	intest.AssertNoError(err)
-	err = store.DeleteFiles(ctx, statFiles)
+	err = store.DeleteFiles(ctx, filenames)
 	intest.AssertNoError(err)
 }
 
@@ -249,7 +249,7 @@ type readTestSuite struct {
 
 func readFileSequential(t *testing.T, s *readTestSuite) {
 	ctx := context.Background()
-	files, _, err := GetAllFileNames(ctx, s.store, "/"+s.subDir)
+	files, _, err := getKVAndStatFilesByScan(ctx, s.store, "/"+s.subDir)
 	intest.AssertNoError(err)
 
 	buf := make([]byte, s.memoryLimit)
@@ -285,9 +285,32 @@ func readFileSequential(t *testing.T, s *readTestSuite) {
 	)
 }
 
+func getKVAndStatFilesByScan(ctx context.Context,
+	store storage.ExternalStorage,
+	nonPartitionedDir string,
+) ([]string, []string, error) {
+	names, err := GetAllFileNames(ctx, store, nonPartitionedDir)
+	if err != nil {
+		return nil, nil, err
+	}
+	var data, stats []string
+	for _, path := range names {
+		bs := []byte(path)
+		lastIdx := bytes.LastIndexByte(bs, '/')
+		secondLastIdx := bytes.LastIndexByte(bs[:lastIdx], '/')
+		parentDir := path[secondLastIdx+1 : lastIdx]
+		if strings.HasSuffix(parentDir, statSuffix) {
+			stats = append(stats, path)
+		} else {
+			data = append(data, path)
+		}
+	}
+	return data, stats, nil
+}
+
 func readFileConcurrently(t *testing.T, s *readTestSuite) {
 	ctx := context.Background()
-	files, _, err := GetAllFileNames(ctx, s.store, "/"+s.subDir)
+	files, _, err := getKVAndStatFilesByScan(ctx, s.store, "/"+s.subDir)
 	intest.AssertNoError(err)
 
 	conc := min(s.concurrency, len(files))
@@ -336,7 +359,7 @@ func readFileConcurrently(t *testing.T, s *readTestSuite) {
 
 func readMergeIter(t *testing.T, s *readTestSuite) {
 	ctx := context.Background()
-	files, _, err := GetAllFileNames(ctx, s.store, "/"+s.subDir)
+	files, _, err := getKVAndStatFilesByScan(ctx, s.store, "/"+s.subDir)
 	intest.AssertNoError(err)
 
 	if s.beforeCreateReader != nil {
@@ -484,7 +507,7 @@ type mergeTestSuite struct {
 
 func mergeStep(t *testing.T, s *mergeTestSuite) {
 	ctx := context.Background()
-	datas, _, err := GetAllFileNames(ctx, s.store, "/"+s.subDir)
+	datas, _, err := getKVAndStatFilesByScan(ctx, s.store, "/"+s.subDir)
 	intest.AssertNoError(err)
 
 	mergeOutput := "merge_output"
@@ -527,7 +550,7 @@ func mergeStep(t *testing.T, s *mergeTestSuite) {
 
 func newMergeStep(t *testing.T, s *mergeTestSuite) {
 	ctx := context.Background()
-	datas, stats, err := GetAllFileNames(ctx, s.store, "/"+s.subDir)
+	datas, stats, err := getKVAndStatFilesByScan(ctx, s.store, "/"+s.subDir)
 	intest.AssertNoError(err)
 
 	mergeOutput := "merge_output"
@@ -647,7 +670,7 @@ func TestReadAllDataLargeFiles(t *testing.T) {
 	writeExternalOneFile(suite2)
 	t.Logf("minKey: %s, maxKey: %s", minKey, maxKey)
 
-	dataFiles, statFiles, err := GetAllFileNames(ctx, store, "")
+	dataFiles, statFiles, err := getKVAndStatFilesByScan(ctx, store, "")
 	intest.AssertNoError(err)
 	intest.Assert(len(dataFiles) == 2)
 
@@ -800,7 +823,7 @@ func TestReadAllData(t *testing.T) {
 
 finishCreateFiles:
 
-	dataFiles, statFiles, err := GetAllFileNames(ctx, store, "/")
+	dataFiles, statFiles, err := getKVAndStatFilesByScan(ctx, store, "/")
 	require.NoError(t, err)
 	require.Equal(t, 2091, len(dataFiles))
 

pkg/lightning/backend/external/byte_reader_test.go

@@ -214,10 +214,12 @@ func TestSwitchMode(t *testing.T) {
 	t.Logf("seed: %d", seed)
 	st := storage.NewMemStorage()
 	// Prepare
+	var kvAndStat [2]string
 	ctx := context.Background()
 	writer := NewWriterBuilder().
 		SetPropSizeDistance(100).
 		SetPropKeysDistance(2).
+		SetOnCloseFunc(func(summary *WriterSummary) { kvAndStat = summary.MultipleFilesStats[0].Filenames[0] }).
 		BuildOneFile(st, "/test", "0")
 
 	writer.InitPartSizeAndLogger(ctx, 5*1024*1024)
@@ -244,9 +246,9 @@ func TestSwitchMode(t *testing.T) {
 	require.NoError(t, err)
 	pool := membuf.NewPool()
 	ConcurrentReaderBufferSizePerConc = rand.Intn(100) + 1
-	kvReader, err := NewKVReader(context.Background(), "/test/0/one-file", st, 0, 64*1024)
+	kvReader, err := NewKVReader(context.Background(), kvAndStat[0], st, 0, 64*1024)
 	require.NoError(t, err)
-	kvReader.byteReader.enableConcurrentRead(st, "/test/0/one-file", 100, ConcurrentReaderBufferSizePerConc, pool.NewBuffer())
+	kvReader.byteReader.enableConcurrentRead(st, kvAndStat[0], 100, ConcurrentReaderBufferSizePerConc, pool.NewBuffer())
 	modeUseCon := false
 	i := 0
 	for {

pkg/lightning/backend/external/iter_test.go

@@ -305,7 +305,7 @@ func testMergeIterSwitchMode(t *testing.T, f func([]byte, int) []byte) {
 	err := writer.Close(context.Background())
 	require.NoError(t, err)
 
-	dataNames, _, err := GetAllFileNames(context.Background(), st, "")
+	dataNames, _, err := getKVAndStatFilesByScan(context.Background(), st, "testprefix")
 	require.NoError(t, err)
 
 	offsets := make([]uint64, len(dataNames))

pkg/lightning/backend/external/onefile_writer.go

@@ -17,6 +17,7 @@ package external
 import (
 	"context"
 	"encoding/binary"
+	"math/rand"
 	"path/filepath"
 	"slices"
 	"time"
@@ -62,6 +63,7 @@ type OneFileWriter struct {
 	// file information.
 	writerID       string
 	filenamePrefix string
+	rnd            *rand.Rand
 	dataFile       string
 	statFile       string
 	dataWriter     storage.ExternalFileWriter
@@ -97,14 +99,14 @@ func (w *OneFileWriter) lazyInitWriter(ctx context.Context) (err error) {
 		return nil
 	}
 
-	dataFile := filepath.Join(w.filenamePrefix, "one-file")
+	dataFile := filepath.Join(w.getPartitionedPrefix(), "one-file")
 	dataWriter, err := w.store.Create(ctx, dataFile, &storage.WriterOption{
 		Concurrency: maxUploadWorkersPerThread,
 		PartSize:    w.partSize})
 	if err != nil {
 		return err
 	}
-	statFile := filepath.Join(w.filenamePrefix+statSuffix, "one-file")
+	statFile := filepath.Join(w.getPartitionedPrefix()+statSuffix, "one-file")
 	statWriter, err := w.store.Create(ctx, statFile, &storage.WriterOption{
 		Concurrency: maxUploadWorkersPerThread,
 		PartSize:    MinUploadPartSize})
@@ -127,7 +129,7 @@ func (w *OneFileWriter) lazyInitDupFile(ctx context.Context) error {
 		return nil
 	}
 
-	dupFile := filepath.Join(w.filenamePrefix+dupSuffix, "one-file")
+	dupFile := filepath.Join(w.getPartitionedPrefix()+dupSuffix, "one-file")
 	dupWriter, err := w.store.Create(ctx, dupFile, &storage.WriterOption{
 		// too many duplicates will cause duplicate resolution part very slow,
 		// we temporarily use 1 as we don't expect too many duplicates, if there
@@ -347,6 +349,10 @@ func (w *OneFileWriter) closeImpl(ctx context.Context) (err error) {
 	return nil
 }
 
+func (w *OneFileWriter) getPartitionedPrefix() string {
+	return randPartitionedPrefix(w.filenamePrefix, w.rnd)
+}
+
 // caller should make sure the buf is large enough to hold the encoded data.
 func encodeToBuf(buf, key, value []byte) {
 	intest.Assert(len(buf) == lengthBytes*2+len(key)+len(value))