从0设计一个图数据自动生成工具

1.引言

1.1 背景

项目中使用到了图数据库nebula-graph，初次使用，不太确定在我们的使用场景中性能如何，于是就想针对我们的数据模型和使用场景进行一次压测，需要一个压测和生成规模数据的工具。

1.2 现状

压测工具目前有现成的，如K6,nebula官网提供了一个nebula-k6插件。

但是关于批量生成数据始终没有找到合适的工具。网上搜了好多次，也包括在nebula论坛里咨询，了解下来目前行业现状如下：

1. 目前业界做图DB性能压测大部分基于社交数据，典型的如ldbc的snb社交数据集
2. 生成数据方面目前也只找到ldbc提供的批量生成社交数据的工具，包括nebula提供的压测工具nebula-bench, 也是基于ldbc工具来生成数据；

   

不同业务场景需要的数据模型肯定有差别，如果只是对比各种图数据之间的性能差异，统一使用社交数据倒没什么问题。但是如果要验证具体业务模型的性能 ，可能只有使用尽量贴近业务场景的数据验证出来的结果才有参考意义。

于是就想能不能自己设计一个工具，最好数据模型里的点、边结构可以自由配置，包括生成的数据量也可以配置，这样就有可能解决更多人的问题，于是写了这篇文档。

2.设计思路

要实现图数据的批量生成，要解决以下几个问题：

• 如何Mock真实场景中的数据；
• 如何配置图的点、边结构以及生成规则；
• 如何抽象任务和数据，以灵活适配不同业务场景；
• 如何设计运行流程能既简单又保留可预见的扩展性；

2.1 语言选择

鉴于只是做一个小工具，就不太想用编译型语言，于是选了python作为开发语言。

2.2 数据生成器

2.2.1 Mock通用数据

我们的业务是分析用户和所参与活动的关系图谱，对Mock数据的需求包括：

• 自动生成业务标识ID；
• 生成用户姓名
• 生成活动主题
• 生成活动主办方的企业名称
• 生成邮箱和电话
• 生成指定长度的随机数字
• 生成指定范围内的枚举值
• 一个字段去引用另一个字段的值
• 基于已有字段的算术表达式运算来生成新字段的值
• 对于边的srcVID和dstVID，需要从原点和目标点中按一定规则来选择；
• ……

所以需要一个Mock数据的类库，来生成常用的数据。个人知道的相关类库有两个：

• mock.js
• faker

相比来说，faker在多国语言方面提供的更全面，另外也支持以插件形式扩展自定义的数据生成器，加之它有python版本，于是就选用了faker。

2.2.2 Mock专用数据

对于通用的数据，faker是能直接生成的，例如：手机、邮箱 、姓名之类，但是对于以下几个场景，可能要自定义数据生成器：

• 引用字段值
• 表达式运算
• id标识
• 边的srcVID和dstVID来引用指定TAG下的点ID；

分析需求并调研python的语法后，对应思路如下：

• 引用字段可以使用python字符串中的format函数来实现；
• 表达式运算可以使用python中的eval函数来实现；
• id标识可以用一个有状态的自增序列实现（见下面流程图右半部分）；
• 边的srcVID和dstVID问题比较特殊，需要找到对应点的vid生成规则来间接生成，流程如下：

  

2.3 规则配置

规则配置主要包括几部分：

1. 图的点/边结构及相关属性字段
2. 每个字段的生成规则
3. 每个节点（点或边）的生成数量配置
4. 数据输出配置

2.3.1 图结构配置

引入图数据库时看过nebula的一些介绍，里面有一个数据导入工具nebula-importer，这个工具的配置文件对我们配置图结构有参考意义。

schema:
  type: vertex
  vertex:
    vid:
      type: string
      index: 0
    tags:
      - name: student
        props:
          - name: name
            type: string
            index: 1
          - name: age
            type: int
            index: 2
          - name: gender
            type: string
            index: 3

上面是一个nebula-importer中关于一个点结构的配置示例，配置了一个tag名称为student, 有name、age、gender三个属性的点结构。相应的，还有对边的配置，具体这里就不展开，参考：使用Nebula Importer。

类似的，在我们这个图数据生成工具中，尽量也沿用类似的格式来配置点和边结构。

2.3.2 生成规则配置

参考faker和mock.js的相关函数定义，一个字段值的生成规则大致可以抽象为两部分：

1. 规则名称：选用哪个规则生成数据，例如random_int表示区间范围内的随机整数；
2. 规则参数：还以random_int为例，可以指定一个min表示最小值，max表示最大值，来限定随机取值的区间范围；

所以我们可以在图结构中针对vid和每个属性扩展一个子节点genrule来表示生成规则，还以上面的为例：

  vertex:
    vid:
      type: string
      genrule: 
        generator: id      # 使用id自增序列作为生成器
        prefix: student_   # vid前缀
        start: 10000       # id自增器起始数字
    tags:
      - name: student
        props:
          - name: name
            type: string
            genrule: 
              generator: name   # 使用人名作为生成器，不需要参数
          - name: age
            type: int
            genrule: 
              generator: random_int   # 使用指定范围的随机数作为生成器，范围15-20
              min: 15
              max: 30
          - name: gender
            type: string
            genrule: 
              generator: random_element
              elements: ('male', 'female')

为了能使用faker提供的丰富的数据生成函数，我们约定：

• faker中的函数名作为规则名称；
• faker中函数的参数作为规则参数，统一使用命名参数；
• genrule中除generator以外的配置项都会自动当作命名参数传给generator指定的数据生成函数；

这样我们瞬间就能支持丰富的数据生成规则，详细的生成规则参考：faker的provider列表

2.3.3 生成数量配置

生成数量不需要具体到字段属性，一个点（或边）配置一个数量即可。

对于点，生成数量配置如下所示：

schema:
  type: vertex
  genNum: 10000       # 当前vertex生成数量2000 
  vertex:
    ……

对于边，表示的是点与点之间的关系，像点一样直接配置总数量虽然从数据规模角度来看也没毛病，但是这样相当于忽视了具体一个点对边数量的要求，粒度有些粗，可能会导致边数量分布不均或者分布过于均匀。

所以个人更倾向于按具体业务来配，对于某种类型的边，从一个起点出发一般会有多少条此类型的边（例如一名学生一个学期一般能订阅5-10门课程），这样更符合业务的直观理解。相应配置示例如下：

schema:
  type: edge
  edge:
    name: follow
    withRanking: false
    genNumPerVID:      # 一个起点生成多少条follow类型的边
      type: srcVID
      genrule: range|5-10
    srcVID:            # 学生作为边的起点
      type: string
      genrule: oftag|student
    dstVID:            # 课程作为边的终点
      type: string
      genrule: oftag|course
    props:
      - name: degree
        type: double
        index: 3

2.3.4 数据输出配置

上文提到，nebula官方已经有了nebula-importer这样的批量数据导入工具，它接受CSV格式的数据输入，我们就尽量把数据输出成它需要的格式，以最大程度适配已有工具。

从nebula-importer的配置（具体参考：Nebula Importer）以及我们的程序运行需要来看，有以下几点要考虑：

• 每个节点都需要独立的CSV文件
• nebula-importer支持有、无Header两种文件；
• 支持配置每次写文件的批量数据条数，用于提高写文件效率，同时控制内存中缓存的数据量；

  output:
  type: csv                     # 输出格式，目前只支持csv
  path: ./target/data/user.csv  # 输出文件位置
  batchSize: 5000               # 每约5000条数据写一次文件
  csv:
    withHeader: true            # 输出文件中是否有header
  

2.4 实体模型

在我们这个数据生成工具中，大概有三类信息需要定义：

• 规则：用于存放配置文件中描述的数据生成规则和输出规则
• 数据：生成具体一个点（或边）数据的存放实体
• 任务：描述一个待执行的任务，不论是生成数据还是存储数据都会被描述成一个任务;

调研过python中一般如何定义实体，有直接用class定义的，也有用命名元组namedtuple定义的，不过这两种都有一定的缺陷，在实际编写时最终选用了python3.7中新支持dataclass注解来定义实体，既简单又支持类型声明。

2.4.1 规则配置定义

上面【规则配置】部分已经示例比较详细，这里只挑两个重要的示例如下：

• PropConfig：对应点或边结构中的一个具体属性，例如：上面学生示例中的age属性，包括像vid、srcVID、dstVID也可以看作是一个预定义的属性；
• Schema: 对应图结构中一个具体的点或边，里面会包含多个PropConfig对象

@dataclass
class PropConfig:        
    name: str                # 属性名, 考虑到多个tag的情况，对于普通属性统一使用[tag].[prop_name]格式表示
    type: str                  # 属性值类型，例如 int, string, bool
    rule: str                   # 生成规则，如 random_int、id、oftag等
    rule_args: dict       # 生成规则需要的参数

@dataclass
class Schema:            
    tags: list = None    # 节点里的TAG列表，边只有一个TAG，点可能有多个TAG，多个TAG合成一个Schema表示
    type: str = ""            # 节点类型，vertex 或 edge
    gen_num: int = 0    # 节点生成数量
    prop_rules: List[PropConfig] = None  # 属性列表，包括预定义属性和普通属性

2.4.2 任务定义

虽然点和边有很多，并且 每种业务场景的图结构都不一样，但我们这个工具预计就只做两件事情：

1. 生成数据
2. 存储数据

这两件事情会有抽象为两种任务来表示：

• GTask: 表示为一个节点生成数据的任务，例如：上面名为student的点配置就会创建一个生成任务；
• STask: 表示一个节点中待存储的一批数据，一个节点有多批数据就会对应生成多个存储任务，例如：总共要生成10000条student数据，如果一批产生500条数据，则大约可能会先后生成20个存储任务；

@dataclass
class GTask:             # 数据生成任务定义
    nkey: NKey           # 任务所属节点的唯一标识符
    schema: Schema       # 节点结构及生成规则配置
    batch_size: int      # 与存储有关，表示每生成 多少条数落一次盘


@dataclass
class STask:             # 存储任务定义
    nkey: NKey           # 任务所属节点的唯一标识符
    output: OutConfig    # 任务所属节点输出配置
    schema: Schema       # 任务所属节点结构配置
    data_index: int      # 当前数据批次的起始编号
    datalist: List       # 当前批次的数据列表

2.4.3 数据实体定义

使用数据生成器Mock的数据需要有一个实体对象作为载体来临时存放，就是图DB中具体的点或边对象，定义如下：

@dataclass
class Vertex:                     # 点对象定义
    vid: str = ""                     # vid
    props: Dict[str, any] = None  # 点的属性集

@dataclass
class Edge:                       # 边对象定义
    srcVID: str = ""              # 源点vid
    dstVID: str = ""              # 目标点vid
    rank: int = 0               
    props: Dict[str, any] = None  # 边的属性集

2.5 运行设计

整体项目运行流程如下：

下面会就图中同个关键部分稍微作些说明。

2.5.1 配置管理

负责配置文件解析和配置项的管理，对整个项目的运行起支撑作用，包括：

• 图结构信息的解析
• 数据生成规则的解析和校验
• 节点生成数量的解析
• 节点配置的管理
• 通用配置项的管理

2.5.2 任务队列

负责待运行任务的管理，在生成、存储两个业务之间搭起一座桥梁，对项目的模块划分起着关键的连接和解耦作用。此模块实际包含两种队列：

• 生成队列：管理所有待运行的生成数据任务，在队列初始化时自动从节点配置一次性构造所有的生成任务；
• 存储队列：管理待执行的存储任务，是一个带缓冲的阻塞队列，当队列满时会阻塞写，当队列空时会阻塞读，在一定程度上能动态调节生成任务和存储任务之间的速度差异；

2.5.3 生成线程

只负责生成数据，工作流程为：

1. 从生成队列读取待运行的生成任务，运行任务就是循环生成指定数量的点或边；
2. 具体到一个点或一条边内部则是逐属性进行，按属性配置的生成规则来调指定函数生成数据；
3. 生成的数据不会立即执行存储，而是会先缓到一个内存buffer；
4. 当到达批量写的阀值时，buffer中的数据会被打包成一个存储任务，压入存储队列并清空buffer，继续回到第2步执行；
5. 当任务退出则表示此任务的所有数据已经生成完，回到第1步继续取下一个任务运行，
6. 当没有任务可以运行时，线程退出，并同时写一个None标志到存储列表，告诉存储线程后续不再有新的数据；

2.5.4 存储线程

只负责存储数据，工作流程为：

1. 从存储队列读取待执行的存储任务；
2. 如果写的是第一条数据，则会判断是否需要写Header，需要则先写一条Header头；
3. 将内存对象转换为csv需要的存储格式，写盘；

2.5.5 运行效率

对于大规模数据下，如何提高数据的生成速度，有多线程和多进程两种可能方式。

多线程：多个线程同时生成数据和存储数据，但有几个限制：

• 一种类型的数据只能生成一个文件，多线程有可能出现文件并发写错乱。为此便设计了一种存储多队列方案，存储线程和生成线程配对出现，都绑定到一个独立的存储队列；
• python中有全局解释器锁（GIL）的限制，目前多线程并不能真正利用多核心；

多进程：脚本内部支持多进程运行比较复杂，而且多进程之间内存不共享，会涉及到规模数据拷贝问题，提升效果并不一定明显。更建议手动起多个进程的方式，大概思路如下：

• 通过配置文件切分来实现任务切分，可以按节点切分，也可以把单个生成数量大的节点切分为多个生成数量较小的节点，只需要ID序列的起始值错开就行；
• 对每份配置文件都独立起一个进程运行；
• 为避免相同节点因文件同名而产生数据覆盖问题，可以给每个配置文件配一个独立的目录；

3.后续

这个项目做的过程中，WarpAI帮了很大的忙，帮助解答了很多python小白类的问题，在解决技术问题方面，ChatGPT确实比搜索好用很多。

项目已经实现，代码地址：GitHub - golfxiao/graph-datagen

设计过程中想到几个可能对使用者有价值的需求点，先列在这里，具体视需求反馈情况和个人时间而定：

1. 支持对接图DB，读取指定图空间中的点和边结构来自动生成图的schema配置；
2. 自动拆分大任务的配置为多个子任务的配置，以支持大规模生成数据Case下的多进程同时运行；
3. 支持兼容更多图DB类型的数据输出格式；