colocated aware random picker

migtests/tests/import-file/validate

@@ -132,7 +132,9 @@ TABLE_NAME_DATA_MAP = {
 def file_import_done_checks(tgt):
 	fetched_row_cnts = tgt.row_count_of_all_tables()
 	print(f"Row counts after import data file: {fetched_row_cnts}")
-	assert fetched_row_cnts == EXPECTED
+	for table_name, rc in EXPECTED.items():
+		assert fetched_row_cnts[table_name] == rc, f"Row count mismatch for table {table_name}, expected: {rc}, got: {fetched_row_cnts[table_name]}"
+
 
 	fetched_row_cnts_non_public = tgt.row_count_of_all_tables("non_public")
 	print(f"Row counts after import data file in non_public schema: {fetched_row_cnts_non_public}")

yb-voyager/cmd/importData.go

@@ -155,6 +155,7 @@ func importDataCommandFn(cmd *cobra.Command, args []string) {
 	// TODO: handle case-sensitive in table names with oracle ff-db
 	// quoteTableNameIfRequired()
 	importFileTasks := discoverFilesToImport()
+	log.Debugf("Discovered import file tasks: %v", importFileTasks)
 	if importerRole == TARGET_DB_IMPORTER_ROLE {
 
 		importType = record.ExportType
@@ -316,6 +317,10 @@ type ImportFileTask struct {
 	FileSize     int64
 }
 
+func (task *ImportFileTask) String() string {
+	return fmt.Sprintf("{ID: %d, FilePath: %s, TableName: %s, RowCount: %d, FileSize: %d}", task.ID, task.FilePath, task.TableNameTup.ForOutput(), task.RowCount, task.FileSize)
+}
+
 // func quoteTableNameIfRequired() {
 // 	if tconf.TargetDBType != ORACLE {
 // 		return
@@ -543,6 +548,8 @@ func importData(importFileTasks []*ImportFileTask) {
 			utils.ErrExit("Failed to classify tasks: %s", err)
 		}
 	}
+	log.Debugf("pending tasks: %v", pendingTasks)
+	log.Debugf("completed tasks: %v", completedTasks)
 
 	//TODO: BUG: we are applying table-list filter on importFileTasks, but here we are considering all tables as per
 	// export-data table-list. Should be fine because we are only disabling and re-enabling, but this is still not ideal.
@@ -569,6 +576,7 @@ func importData(importFileTasks []*ImportFileTask) {
 			utils.PrintAndLog("Tables to import: %v", importFileTasksToTableNames(pendingTasks))
 			prepareTableToColumns(pendingTasks) //prepare the tableToColumns map
 			poolSize := tconf.Parallelism * 2
+			maxTasksInProgress := tconf.Parallelism
 			if tconf.EnableYBAdaptiveParallelism {
 				// in case of adaptive parallelism, we need to use maxParalllelism * 2
 				yb, ok := tdb.(*tgtdb.TargetYugabyteDB)
@@ -584,28 +592,36 @@ func importData(importFileTasks []*ImportFileTask) {
 				controlPlane.UpdateImportedRowCount(importDataAllTableMetrics)
 			}
 
-			for _, task := range pendingTasks {
-				// The code can produce `poolSize` number of batches at a time. But, it can consume only
-				// `parallelism` number of batches at a time.
-				batchImportPool = pool.New().WithMaxGoroutines(poolSize)
-				log.Infof("created batch import pool of size: %d", poolSize)
-
-				taskImporter, err := NewFileTaskImporter(task, state, batchImportPool, progressReporter)
+			useTaskPicker := utils.GetEnvAsBool("USE_TASK_PICKER_FOR_IMPORT", true)
+			if useTaskPicker {
+				err := importTasksViaTaskPicker(pendingTasks, state, progressReporter, poolSize, maxTasksInProgress)
 				if err != nil {
-					utils.ErrExit("Failed to create file task importer: %s", err)
+					utils.ErrExit("Failed to import tasks via task picker: %s", err)
 				}
+			} else {
+				for _, task := range pendingTasks {
+					// The code can produce `poolSize` number of batches at a time. But, it can consume only
+					// `parallelism` number of batches at a time.
+					batchImportPool = pool.New().WithMaxGoroutines(poolSize)
+					log.Infof("created batch import pool of size: %d", poolSize)
 
-				for !taskImporter.AllBatchesSubmitted() {
-					err := taskImporter.SubmitNextBatch()
+					taskImporter, err := NewFileTaskImporter(task, state, batchImportPool, progressReporter)
 					if err != nil {
-						utils.ErrExit("Failed to submit next batch: task:%v err: %s", task, err)
+						utils.ErrExit("Failed to create file task importer: %s", err)
 					}
-				}
 
-				batchImportPool.Wait() // Wait for the file import to finish.
-				taskImporter.PostProcess()
+					for !taskImporter.AllBatchesSubmitted() {
+						err := taskImporter.SubmitNextBatch()
+						if err != nil {
+							utils.ErrExit("Failed to submit next batch: task:%v err: %s", task, err)
+						}
+					}
+
+					batchImportPool.Wait() // wait for file import to finish
+					taskImporter.PostProcess()
+				}
+				time.Sleep(time.Second * 2)
 			}
-			time.Sleep(time.Second * 2)
 		}
 		utils.PrintAndLog("snapshot data import complete\n\n")
 	}
@@ -684,6 +700,74 @@ func importData(importFileTasks []*ImportFileTask) {
 
 }
 
+func importTasksViaTaskPicker(pendingTasks []*ImportFileTask, state *ImportDataState, progressReporter *ImportDataProgressReporter, poolSize int, maxTasksInProgress int) error {
+	// The code can produce `poolSize` number of batches at a time. But, it can consume only
+	// `parallelism` number of batches at a time.
+	batchImportPool = pool.New().WithMaxGoroutines(poolSize)
+	log.Infof("created batch import pool of size: %d", poolSize)
+
+	var taskPicker FileTaskPicker
+	var err error
+	if importerRole == TARGET_DB_IMPORTER_ROLE || importerRole == IMPORT_FILE_ROLE {
+		yb, ok := tdb.(*tgtdb.TargetYugabyteDB)
+		if !ok {
+			return fmt.Errorf("expected tdb to be of type TargetYugabyteDB, got: %T", tdb)
+		}
+		taskPicker, err = NewColocatedAwareRandomTaskPicker(maxTasksInProgress, pendingTasks, state, yb)
+		if err != nil {
+			return fmt.Errorf("create colocated aware randmo task picker: %w", err)
+		}
+	} else {
+		taskPicker, err = NewSequentialTaskPicker(pendingTasks, state)
+		if err != nil {
+			return fmt.Errorf("create sequential task picker: %w", err)
+		}
+	}
+
+	taskImporters := map[int]*FileTaskImporter{}
+
+	for taskPicker.HasMoreTasks() {
+		task, err := taskPicker.NextTask()
+		if err != nil {
+			return fmt.Errorf("get next task: %w", err)
+		}
+		var taskImporter *FileTaskImporter
+		var ok bool
+		taskImporter, ok = taskImporters[task.ID]
+		if !ok {
+			taskImporter, err = NewFileTaskImporter(task, state, batchImportPool, progressReporter)
+			if err != nil {
+				return fmt.Errorf("create file task importer: %s", err)
+			}
+			log.Infof("created file task importer for table: %s, task: %v", task.TableNameTup.ForOutput(), task)
+			taskImporters[task.ID] = taskImporter
+		}
+
+		if taskImporter.AllBatchesSubmitted() {
+			// All batches for this task have been submitted.
+			// task could have been completed (all batches imported) OR still in progress
+			// in case task is done, we should inform task picker so that we stop picking that task.
+			taskDone, err := taskImporter.AllBatchesImported()
+			if err != nil {
+				return fmt.Errorf("check if all batches are imported: task: %v err :%w", task, err)
+			}
+			if taskDone {
+				taskImporter.PostProcess()
+				err = taskPicker.MarkTaskAsDone(task)
+				if err != nil {
+					return fmt.Errorf("mark task as done: task: %v, err: %w", task, err)
+				}
+			}
+			continue
+		}
+		err = taskImporter.SubmitNextBatch()
+		if err != nil {
+			return fmt.Errorf("submit next batch: task:%v err: %s", task, err)
+		}
+	}
+	return nil
+}
+
 func startAdaptiveParallelism() (bool, error) {
 	yb, ok := tdb.(*tgtdb.TargetYugabyteDB)
 	if !ok {

yb-voyager/cmd/importDataFileBatchProducer_test.go

@@ -39,7 +39,7 @@ func TestBasicFileBatchProducer(t *testing.T) {
 
 	fileContents := `id,val
 1, "hello"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -71,7 +71,7 @@ func TestFileBatchProducerBasedOnRowsThreshold(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -122,7 +122,7 @@ func TestFileBatchProducerBasedOnSizeThreshold(t *testing.T) {
 2, "ghijk"
 3, "mnopq"
 4, "stuvw"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -182,7 +182,7 @@ func TestFileBatchProducerThrowsErrorWhenSingleRowGreaterThanMaxBatchSize(t *tes
 2, "ghijk"
 3, "mnopq1234567899876543"
 4, "stuvw"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -218,7 +218,7 @@ func TestFileBatchProducerResumable(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)
@@ -274,7 +274,7 @@ func TestFileBatchProducerResumeAfterAllBatchesProduced(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table", 1)
 	assert.NoError(t, err)
 
 	batchproducer, err := NewFileBatchProducer(task, state)

yb-voyager/cmd/importDataFileCommand.go

@@ -168,7 +168,8 @@ func prepareImportFileTasks() []*ImportFileTask {
 		return result
 	}
 	kvs := strings.Split(fileTableMapping, ",")
-	for i, kv := range kvs {
+	idIndex := 0
+	for _, kv := range kvs {
 		globPattern, table := strings.Split(kv, ":")[0], strings.Split(kv, ":")[1]
 		filePaths, err := dataStore.Glob(globPattern)
 		if err != nil {
@@ -187,12 +188,13 @@ func prepareImportFileTasks() []*ImportFileTask {
 				utils.ErrExit("calculating file size in bytes: %q: %v", filePath, err)
 			}
 			task := &ImportFileTask{
-				ID:           i,
+				ID:           idIndex,
 				FilePath:     filePath,
 				TableNameTup: tableNameTuple,
 				FileSize:     fileSize,
 			}
 			result = append(result, task)
+			idIndex++
 		}
 	}
 	return result

yb-voyager/cmd/importDataFileTaskImporter.go

@@ -75,9 +75,12 @@ func (fti *FileTaskImporter) AllBatchesSubmitted() bool {
 	return fti.batchProducer.Done()
 }
 
-func (fti *FileTaskImporter) AllBatchesImported() error {
-	// TODO: check importDataState for status.
-	panic("not implemented")
+func (fti *FileTaskImporter) AllBatchesImported() (bool, error) {
+	taskStatus, err := fti.state.GetFileImportState(fti.task.FilePath, fti.task.TableNameTup)
+	if err != nil {
+		return false, fmt.Errorf("getting file import state: %s", err)
+	}
+	return taskStatus == FILE_IMPORT_COMPLETED, nil
 }
 
 func (fti *FileTaskImporter) SubmitNextBatch() error {
@@ -234,6 +237,7 @@ func getImportBatchArgsProto(tableNameTup sqlname.NameTuple, filePath string) *t
 	}
 	// If `columns` is unset at this point, no attribute list is passed in the COPY command.
 	fileFormat := dataFileDescriptor.FileFormat
+
 	// from export data with ora2pg, it comes as an SQL file, with COPY command having data.
 	// Import-data also reads it appropriately with the help of sqlDataFile.
 	// But while running COPY for a batch, we need to set the format as TEXT (SQL does not make sense)

yb-voyager/cmd/importDataFileTaskImporter_test.go

@@ -47,7 +47,7 @@ func TestBasicTaskImport(t *testing.T) {
 	fileContents := `id,val
 1, "hello"
 2, "world"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_basic")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_basic", 1)
 	testutils.FatalIfError(t, err)
 
 	progressReporter := NewImportDataProgressReporter(true)
@@ -88,7 +88,7 @@ func TestImportAllBatchesAndResume(t *testing.T) {
 	fileContents := `id,val
 1, "hello"
 2, "world"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_all")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_all", 1)
 	testutils.FatalIfError(t, err)
 
 	progressReporter := NewImportDataProgressReporter(true)
@@ -139,7 +139,7 @@ func TestTaskImportResumable(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_resume")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_resume", 1)
 	testutils.FatalIfError(t, err)
 
 	progressReporter := NewImportDataProgressReporter(true)
@@ -198,7 +198,7 @@ func TestTaskImportResumableNoPK(t *testing.T) {
 2, "world"
 3, "foo"
 4, "bar"`
-	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_resume_no_pk")
+	_, task, err := createFileAndTask(lexportDir, fileContents, ldataDir, "test_table_resume_no_pk", 1)
 	testutils.FatalIfError(t, err)
 
 	progressReporter := NewImportDataProgressReporter(true)