diff --git a/internal/datanode/data_node.go b/internal/datanode/data_node.go
index d2d8ffdc29..89976b689d 100644
--- a/internal/datanode/data_node.go
+++ b/internal/datanode/data_node.go
@@ -850,7 +850,8 @@ func (node *DataNode) Import(ctx context.Context, req *datapb.ImportTaskRequest)
 	_ = idAllocator.Start()
 	defer idAllocator.Close()
 
-	importWrapper := importutil.NewImportWrapper(ctx, schema, 2, Params.DataNodeCfg.FlushInsertBufferSize, idAllocator, node.chunkManager,
+	segmentSize := int64(Params.DataCoordCfg.SegmentMaxSize) * 1024 * 1024
+	importWrapper := importutil.NewImportWrapper(ctx, schema, 2, segmentSize, idAllocator, node.chunkManager,
 		importFlushReqFunc(node, req, importResult, schema, ts), importResult, reportFunc)
 	err = importWrapper.Import(req.GetImportTask().GetFiles(), req.GetImportTask().GetRowBased(), false)
 	if err != nil {