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、核心设计模式
| 组件 | 接口 | 职责 |
|---|---|---|
| InMemoryDataSource | TableProvider + DataSourceRegister | 入口类,推断 Schema,创建 Table;SPI 注册短别名 "inmemory" |
| InMemoryTable | Table + SupportsRead + SupportsWrite | 表抽象,声明能力(BATCH_READ/WRITE),创建读写构建器 |
| InMemoryScanBuilder | SupportsPushDownRequiredColumns + SupportsPushDownFilters | 列裁剪 + 谓词下推 优化 |
| InMemoryBatch | Batch | 分区规划,决定并行度 |
| InMemoryPartitionReader | PartitionReader<InternalRow> | 实际数据读取 + 过滤执行 + 列投影 |
| InMemoryBatchWrite | BatchWrite | 两阶段提交协调 |
| InMemoryDataWriter | DataWriter<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 控制并行度
**粗体** _斜体_ [链接](http://example.com) `代码` - 列表 > 引用。你还可以使用@来通知其他用户。