Spark 4.0 新特性Python Data Source API 快速上手

1. 什么是 Python Data Source APIPython Data Source API 是 Spark 4.0 引入的新能力它允许开发者在 Python 中直接实现自定义数据源和数据写出逻辑。换句话说你可以像实现一个插件一样为 Spark 增加新的读取来源和写出目标而不必先写 Scala/Java 版本。它支持的能力包括批处理读取批处理写入流式读取流式写入官方给出的能力映射关系很清晰批处理读实现reader()批处理写实现writer()流式读实现streamReader()或simpleStreamReader()流式写实现streamWriter()2. 一个最简单的自定义数据源先看一个最小示例定义一个只返回两行数据的数据源。fromtypingimportIterator,Tuplefrompyspark.sql.datasourceimportDataSource,DataSourceReader,InputPartitionfrompyspark.sql.typesimportIntegerType,StringType,StructField,StructTypeclassSimpleDataSource(DataSource):classmethoddefname(cls)-str:returnsimpledefschema(self)-StructType:returnStructType([StructField(name,StringType()),StructField(age,IntegerType())])defreader(self,schema:StructType)-DataSourceReader:returnSimpleDataSourceReader()classSimpleDataSourceReader(DataSourceReader):defread(self,partition:InputPartition)-Iterator[Tuple]:yield(Alice,20)yield(Bob,30)注册后即可直接读取frompyspark.sqlimportSparkSession sparkSparkSession.builder.getOrCreate()spark.dataSource.register(SimpleDataSource)spark.read.format(simple).load().show()输出结果-------- | name|age| -------- |Alice| 20| | Bob| 30| --------这个例子说明了 Python Data Source API 的核心思路定义名字、定义 schema、实现 reader然后注册给 Spark 使用。3. DataSource 的基本结构一个完整的 Python Data Source通常要继承DataSource然后根据你的能力实现对应方法。官方示例中给出了一份更完整的定义fromtypingimportUnionfrompyspark.sql.datasourceimport(DataSource,DataSourceReader,DataSourceStreamReader,DataSourceStreamWriter,DataSourceWriter)frompyspark.sql.typesimportStructTypeclassFakeDataSource(DataSource):classmethoddefname(cls)-str:returnfakedefschema(self)-Union[StructType,str]:returnname string, date string, zipcode string, state stringdefreader(self,schema:StructType)-DataSourceReader:returnFakeDataSourceReader(schema,self.options)defwriter(self,schema:StructType,overwrite:bool)-DataSourceWriter:returnFakeDataSourceWriter(self.options)defstreamReader(self,schema:StructType)-DataSourceStreamReader:returnFakeStreamReader(schema,self.options)defstreamWriter(self,schema:StructType,overwrite:bool)-DataSourceStreamWriter:returnFakeStreamWriter(self.options)这段代码已经覆盖了批量读批量写流式读流式写如果你只需要其中一种能力只实现对应方法即可。4. 批处理 Reader 怎么写官方示例里的FakeDataSourceReader使用faker库按 schema 生成假数据。fromtypingimportDictfrompyspark.sql.datasourceimportDataSourceReaderfrompyspark.sql.typesimportStructTypeclassFakeDataSourceReader(DataSourceReader):def__init__(self,schema:StructType,options:Dict[str,str]):self.schemaschema self.optionsoptionsdefread(self,partition):fromfakerimportFaker fakeFaker()num_rowsint(self.options.get(numRows,3))for_inrange(num_rows):row[]forfieldinself.schema.fields:valuegetattr(fake,field.name)()row.append(value)yieldtuple(row)这里有两个关键点options里的值全部都是字符串所以要自己做类型转换依赖库faker是在方法内部导入的这和后面要说的序列化要求有关5. 批处理 Writer 怎么写批处理写出要实现DataSourceWriter。官方示例里Writer 会统计每个分区写出的记录数并在commit()中输出总行数。fromdataclassesimportdataclassfromtypingimportIterator,Listfrompyspark.sqlimportRowfrompyspark.sql.datasourceimportDataSourceWriter,WriterCommitMessagedataclassclassSimpleCommitMessage(WriterCommitMessage):partition_id:intcount:intclassFakeDataSourceWriter(DataSourceWriter):defwrite(self,rows:Iterator[Row])-SimpleCommitMessage:frompysparkimportTaskContext contextTaskContext.get()partition_idcontext.partitionId()cntsum(1for_inrows)returnSimpleCommitMessage(partition_idpartition_id,countcnt)defcommit(self,messages:List[SimpleCommitMessage])-None:total_countsum(message.countformessageinmessages)print(fTotal number of rows:{total_count})defabort(self,messages:List[SimpleCommitMessage])-None:failed_countsum(messageisNoneformessageinmessages)print(fNumber of failed tasks:{failed_count})这类写法很适合统计落盘条数汇总分区写出结果自定义成功/失败处理逻辑6. 流式 Reader 怎么写6.1 标准流式 Reader如果你的流式数据源需要自己管理 offset、分区规划、提交状态可以实现DataSourceStreamReader。官方示例中每个 micro-batch 生成 2 行数据。fromtypingimportIterator,Tuplefrompyspark.sql.datasourceimportDataSourceStreamReader,InputPartitionclassRangePartition(InputPartition):def__init__(self,start:int,end:int):self.startstart self.endendclassFakeStreamReader(DataSourceStreamReader):def__init__(self,schema,options):self.current0definitialOffset(self)-dict:return{offset:0}deflatestOffset(self)-dict:self.current2return{offset:self.current}defpartitions(self,start:dict,end:dict)-list[InputPartition]:return[RangePartition(start[offset],end[offset])]defcommit(self,end:dict)-None:passdefread(self,partition)-Iterator[Tuple]:foriinrange(partition.start,partition.end):yield(i,str(i))6.2 简化版流式 Reader如果数据源吞吐量不高也不需要复杂分区可以实现SimpleDataSourceStreamReader。官方说明里明确说streamReader()和simpleStreamReader()二选一即可前者优先。fromtypingimportIterator,Tuplefrompyspark.sql.datasourceimportSimpleDataSourceStreamReaderclassFakeSimpleStreamReader(SimpleDataSourceStreamReader):definitialOffset(self)-dict:return{offset:0}defread(self,start:dict)-Tuple[Iterator[Tuple],dict]:start_idxstart[offset]ititer([(i,)foriinrange(start_idx,start_idx2)])return(it,{offset:start_idx2})defreadBetweenOffsets(self,start:dict,end:dict)-Iterator[Tuple]:returniter([(i,)foriinrange(start[offset],end[offset])])defcommit(self,end:dict)-None:pass7. 流式 Writer 怎么写流式写出对应的是DataSourceStreamWriter。官方示例里它把每个 micro-batch 的元信息写到本地路径中。fromtypingimportIterator,List,Optionalfrompyspark.sqlimportRowfrompyspark.sql.datasourceimportDataSourceStreamWriter,WriterCommitMessageclassSimpleCommitMessage(WriterCommitMessage):partition_id:intcount:intclassFakeStreamWriter(DataSourceStreamWriter):def__init__(self,options):self.optionsoptions self.pathself.options.get(path)assertself.pathisnotNonedefwrite(self,iterator:Iterator[Row])-WriterCommitMessage:frompysparkimportTaskContext contextTaskContext.get()partition_idcontext.partitionId()cnt0for_initerator:cnt1returnSimpleCommitMessage(partition_idpartition_id,countcnt)defcommit(self,messages:List[Optional[SimpleCommitMessage]],batchId:int)-None:statusdict(num_partitionslen(messages),rowssum(m.countforminmessages))withopen(os.path.join(self.path,f{batchId}.json),a)asfile:file.write(json.dumps(status)\n)defabort(self,messages:List[Optional[SimpleCommitMessage]],batchId:int)-None:withopen(os.path.join(self.path,f{batchId}.txt),w)asfile:file.write(ffailed in batch{batchId})这个模式的核心是write()负责分区级写出commit()负责 batch 成功后的汇总处理abort()负责 batch 失败后的兜底逻辑8. 序列化要求一定要注意官方专门强调DataSource、DataSourceReader、DataSourceWriter、DataSourceStreamReader、DataSourceStreamWriter以及它们的方法都必须能够被pickle序列化。因此有一个很重要的实践方法内部用到的库尽量在方法内部导入。比如官方示例中TaskContext不是在文件开头导入而是在read()或write()方法内部导入defread(self,partition):frompysparkimportTaskContext contextTaskContext.get()这是为了避免序列化问题。9. 如何注册和使用 Python Data Source9.1 注册spark.dataSource.register(FakeDataSource)9.2 读取使用默认 schemaspark.read.format(fake).load().show()使用自定义 schemaspark.read.format(fake).schema(name string, company string).load().show()传入 optionspark.read.format(fake).option(numRows,5).load().show()这些都是官方给出的直接用法。9.3 写出写出时要指定mode()支持append和overwrite。dfspark.range(0,10,1,5)df.write.format(fake).mode(append).save()10. 在 Structured Streaming 里使用注册之后同一个 Python Data Source 也可以用于readStream()和writeStream()。10.1 作为流式 sourcequeryspark.readStream.format(fake).load()\.writeStream.format(console).start()10.2 同时作为流式 source 和 sinkqueryspark.readStream.format(fake).load()\.writeStream.format(fake).start(/output_path)这意味着你可以完全用 Python 打通一个自定义流式读写链路。11. Arrow Batch 支持性能提升的重点官方文档最后一部分特别强调了一点Python Data Source Reader 支持直接产出 Arrow Batch这能显著提升性能甚至在大数据场景下带来一个数量级的提升。实现方式很简单在read()方法里直接yield pyarrow.RecordBatch。frompyspark.sql.datasourceimportDataSource,DataSourceReader,InputPartitionfrompyspark.sqlimportSparkSessionimportpyarrowaspaclassArrowBatchDataSource(DataSource):classmethoddefname(cls):returnarrowbatchdefschema(self):returnkey int, value stringdefreader(self,schema:str):returnArrowBatchDataSourceReader(schema,self.options)classArrowBatchDataSourceReader(DataSourceReader):def__init__(self,schema,options):self.schemaschema self.optionsoptionsdefread(self,partition):keyspa.array([1,2,3,4,5],typepa.int32())valuespa.array([one,two,three,four,five],typepa.string())schemapa.schema([(key,pa.int32()),(value,pa.string())])record_batchpa.RecordBatch.from_arrays([keys,values],schemaschema)yieldrecord_batchdefpartitions(self):return[InputPartition(i)foriinrange(1)]sparkSparkSession.builder.appName(ArrowBatchExample).getOrCreate()spark.dataSource.register(ArrowBatchDataSource)dfspark.read.format(arrowbatch).load()df.show()这部分非常值得关注因为它说明 Python Data Source API 不只是“能用”而且在配合 Arrow 后可以更快。12. 使用时的几个注意点官方文档还给了几条很实用的说明12.1 名称冲突时内置和 Scala/Java Data Source 优先如果 Python Data Source 和内置数据源或 Scala/Java 数据源同名默认优先解析后者。因此定义名字时尽量避免冲突。12.2 可以重复注册后注册会覆盖前注册这对调试很方便但也意味着线上环境要注意命名和注册顺序。12.3 支持自动注册官方提到如果你在顶层模块里把数据源导出为DefaultSource并且模块名前缀是pyspark_就可以自动注册。13. 总结Python Data Source API 的价值很直接它让 Spark 4.0 开始真正具备了“用 Python 写自定义数据源”的能力。以前这类扩展大多要落到 JVM 侧现在很多批处理和流处理场景都可以直接用 Python 完成。它最重要的几个亮点是批处理读写可扩展流处理读写可扩展与spark.read / write / readStream / writeStream体系自然融合支持 Arrow Batch 提升性能如果你在做自定义 Connector、测试数据源、内部平台适配层或者想快速接入非标准数据来源这个 API 会非常值得深入。