1、整体架构

spark-connector-test/
├── pom.xml                          # Maven 配置 (Spark 3.5.1, Scala 2.12, Java 8)
└── src/main/java/.../connector/
    ├── InMemoryDataSource.java      # 入口类 (TableProvider + DataSourceRegister)
    ├── InMemoryTable.java           # 表抽象 (SupportsRead + SupportsWrite)
    ├── read/                        # 读取链路
    │   ├── InMemoryScanBuilder.java     # 列裁剪 + 谓词下推
    │   ├── InMemoryScan.java            # 扫描描述符
    │   ├── InMemoryBatch.java           # 分区规划
    │   ├── InMemoryInputPartition.java  # 分区元数据
    │   ├── InMemoryPartitionReaderFactory.java
    │   └── InMemoryPartitionReader.java # 实际数据读取
    ├── write/                       # 写入链路
    │   ├── InMemoryWriteBuilder.java    # 写入构建器
    │   ├── InMemoryBatchWrite.java      # 两阶段提交协调
    │   ├── InMemoryDataWriterFactory.java
    │   ├── InMemoryDataWriter.java      # 实际数据写入
    │   └── InMemoryWriterCommitMessage.java
    └── store/
        └── InMemoryStore.java       # 内存存储(模拟外部数据源)

2、读取流程

┌─────────────────────────────────────────────────────────────────────────────┐
│                              Driver 端                                       │
├─────────────────────────────────────────────────────────────────────────────┤
│  spark.read().format("inmemory").option("tableName","t1").load()            │
│       ↓                                                                      │
│  InMemoryDataSource.inferSchema()         → 推断 Schema                     │
│       ↓                                                                      │
│  InMemoryDataSource.getTable()            → 创建 InMemoryTable               │
│       ↓                                                                      │
│  InMemoryTable.newScanBuilder()           → 创建 InMemoryScanBuilder         │
│       ↓                                                                      │
│  ScanBuilder.pruneColumns()               → 列裁剪优化                       │
│  ScanBuilder.pushFilters()                → 谓词下推优化                     │
│       ↓                                                                      │
│  ScanBuilder.build()                      → 创建 InMemoryScan                │
│       ↓                                                                      │
│  InMemoryScan.toBatch()                   → 创建 InMemoryBatch               │
│       ↓                                                                      │
│  InMemoryBatch.planInputPartitions()      → 规划分区 [P0, P1, ...]           │
│  InMemoryBatch.createReaderFactory()      → 创建 ReaderFactory               │
└─────────────────────────────────────────────────────────────────────────────┘
                              ↓ 序列化发送到 Executor
┌─────────────────────────────────────────────────────────────────────────────┐
│                              Executor 端                                     │
├─────────────────────────────────────────────────────────────────────────────┤
│  ReaderFactory.createReader(partition)    → 创建 PartitionReader             │
│       ↓                                                                      │
│  while (reader.next()) {                  → 迭代读取                         │
│      InternalRow row = reader.get();                                         │
│  }                                                                           │
│  reader.close()                                                               │
└─────────────────────────────────────────────────────────────────────────────┘

3、写入流程

┌─────────────────────────────────────────────────────────────────────────────┐
│                              Driver 端                                       │
├─────────────────────────────────────────────────────────────────────────────┤
│  df.write().format("inmemory").mode("append").save()                        │
│       ↓                                                                      │
│  InMemoryTable.newWriteBuilder()          → 创建 InMemoryWriteBuilder        │
│       ↓                                                                      │
│  WriteBuilder.truncate() (Overwrite模式)  → 设置清空标记                     │
│       ↓                                                                      │
│  WriteBuilder.buildForBatch()             → 创建 InMemoryBatchWrite          │
│       ↓                                                                      │
│  BatchWrite.createBatchWriterFactory()    → 创建 DataWriterFactory           │
│       ↓                                                                      │
│  [等待所有 Task 完成]                                                        │
│       ↓                                                                      │
│  全部成功 → BatchWrite.commit(messages[])  → 提交                            │
│  任一失败 → BatchWrite.abort(messages[])   → 回滚                            │
└─────────────────────────────────────────────────────────────────────────────┘
                              ↓ 序列化发送到 Executor
┌─────────────────────────────────────────────────────────────────────────────┐
│                              Executor 端                                     │
├─────────────────────────────────────────────────────────────────────────────┤
│  DataWriterFactory.createWriter(partitionId, taskId) → DataWriter           │
│       ↓                                                                      │
│  DataWriter.write(row) × N               → 缓冲写入                         │
│       ↓                                                                      │
│  成功 → DataWriter.commit()               → 返回 WriterCommitMessage         │
│  失败 → DataWriter.abort()                → 清理缓冲区                       │
└─────────────────────────────────────────────────────────────────────────────┘

4、核心设计模式

组件接口职责
InMemoryDataSourceTableProvider + DataSourceRegister入口类,推断 Schema,创建 Table;SPI 注册短别名 "inmemory"
InMemoryTableTable + SupportsRead + SupportsWrite表抽象,声明能力(BATCH_READ/WRITE),创建读写构建器
InMemoryScanBuilderSupportsPushDownRequiredColumns + SupportsPushDownFilters列裁剪 + 谓词下推 优化
InMemoryBatchBatch分区规划,决定并行度
InMemoryPartitionReaderPartitionReader<InternalRow>实际数据读取 + 过滤执行 + 列投影
InMemoryBatchWriteBatchWrite两阶段提交协调
InMemoryDataWriterDataWriter<InternalRow>实际数据写入,缓冲批量提交
InMemoryStore单例内存存储,线程安全

5、核心特点

1、SPI 注册机制:通过 META-INF/services/org.apache.spark.sql.sources.DataSourceRegister 注册,支持 format("inmemory") 短别名

2、列裁剪:只读取需要的列,减少 I/O

SELECT name, age FROM table → 只返回 2 列
3、谓词下推:WHERE 条件下推到数据源执行

支持 GreaterThan、LessThan、IsNotNull、EqualTo
不支持的过滤条件由 Spark 内存处理
4、两阶段提交 (2PC):保证写入原子性

  • Task 级别:数据先缓冲,commit() 时写入
  • Job 级别:所有 Task 成功后才真正提交
    5、并行读写:通过 numPartitions 控制并行度

journey
37 声望25 粉丝