Apache Spark Python Data Source API 完全指南用 PySpark 构建自定义数据源与数据汇【免费下载链接】sparkApache Spark - A unified analytics engine for large-scale data processing项目地址: https://gitcode.com/gh_mirrors/sp/sparkPython Data Source API 是 Apache Spark 4.0 引入的全新特性它允许开发者直接用 Python 实现自定义数据源Source与数据汇Sink读写任意外部系统。本文以当前仓库中的官方教程python/docs/source/tutorial/sql/python_data_source.rst为骨架结合python/pyspark/sql/datasource.py的源码实现与官方测试用例系统讲解如何定义、注册、使用和管理 Python 数据源覆盖批处理读写、流式读写、Limit 下推、准入控制Admission Control与 Arrow 批直传等完整能力。概述为什么需要 Python Data Source API在 Spark 4.0 之前接入自定义数据源必须编写 Scala/Java 代码并实现DataSourceV2系列接口再编译打包、部署依赖门槛较高。Python Data Source API 将这一能力开放给 Python 开发者只需要继承DataSource基类、实现若干方法即可通过spark.read.format(...).load()读取自定义数据源通过df.write.format(...).save()写入自定义数据汇。该 API 的核心入口全部集中在 python/pyspark/sql/datasource.py 中__all__导出的关键类包括类职责DataSource自定义数据源基类提供reader()/writer()/streamReader()/simpleStreamReader()/streamWriter()等工厂方法DataSourceReader批式读取器负责产出分区并逐分区读取数据DataSourceWriter批式写入器执行写入并返回提交消息DataSourceStreamReader流式读取器基于 offset 规划微批SimpleDataSourceStreamReader简化版流式读取器低吞吐场景无需分区规划DataSourceStreamWriter流式写入器按微批提交DataSourceArrowWriter/DataSourceStreamArrowWriter使用 ArrowRecordBatch读写的高性能读写器InputPartition输入分区用于把读取拆分成并行任务WriterCommitMessage写入任务的提交消息回传给commit()/abort()Filter及EqualTo、GreaterThan、IsNull、StringStartsWith等子类谓词下推filter pushdown的过滤条件表示从源码看DataSource是一个ABC抽象类其构造函数__init__接收一个大小写不敏感的 options 字典由CaseInsensitiveDict实现见 datasource.py 第 1361 行附近官方文档明确标注该方法不应被重写。所有未被实现的抽象方法如schema()、reader()默认抛出PySparkNotImplementedError。快速上手最简单的 Batch Reader 示例官方教程给出了一个零依赖的最小示例用于演示如何在没有外部库的情况下快速跑通一个数据源。它生成恰好两行合成数据。Step 1定义数据源from typing import Iterator, Tuple from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition from pyspark.sql.types import IntegerType, StringType, StructField, StructType class SimpleDataSource(DataSource): A simple data source for PySpark that generates exactly two rows of synthetic data. classmethod def name(cls) - str: return simple def schema(self) - StructType: return StructType([ StructField(name, StringType()), StructField(age, IntegerType()) ]) def reader(self, schema: StructType) - DataSourceReader: return SimpleDataSourceReader() class SimpleDataSourceReader(DataSourceReader): def read(self, partition: InputPartition) - Iterator[Tuple]: yield (Alice, 20) yield (Bob, 30)Step 2注册数据源from pyspark.sql import SparkSession spark SparkSession.builder.getOrCreate() spark.dataSource.register(SimpleDataSource)注册时DataSourceRegistration.register()会调用dataSource.name()获取格式名将数据源类通过_wrap_function序列化后交给 JVM 侧的UserDefinedPythonDataSource完成注册见 datasource.py 中DataSourceRegistration.register的实现。Step 3读取数据spark.read.format(simple).load().show() # -------- # | name|age| # -------- # |Alice| 20| # | Bob| 30| # --------这个示例虽然简单却涵盖了自定义数据源的核心要素name()定义格式名、schema()定义读取 schema、reader()返回读取器、读取器的read()方法按分区产出元组迭代器。如果schema()未实现且用户读取时也未显式指定 schemaSpark 会抛出异常——这一点在 datasource.py 中schema()的 docstring 中有明确说明。综合示例同时支持 Batch 与 Streaming 的完整数据源生产场景下一个数据源往往既要支持批式读写也要支持流式读写。要创建一个自定义 Python 数据源你需要继承DataSource基类并按需实现读写方法。能力矩阵按需实现对应方法下表来自官方教程列出了不同能力batch/streaming、source/sink所需实现的方法能力source读sink写batchreader()writer()streamingstreamReader()或simpleStreamReader()streamWriter()也就是说可读的批数据源必须实现reader()可写的批数据源必须实现writer()可读的流式数据源必须实现streamReader()与simpleStreamReader()二者之一simpleStreamReader()只在streamReader()未实现时被调用可写的流式数据源必须实现streamWriter()。定义 Data Source 子类以下示例使用faker库生成合成数据。运行前请确保faker已安装并可在 Python 环境中导入。from typing import Union from pyspark.sql.datasource import ( DataSource, DataSourceReader, DataSourceStreamReader, DataSourceStreamWriter, DataSourceWriter ) from pyspark.sql.types import StructType class FakeDataSource(DataSource): A fake data source for PySpark to generate synthetic data using the faker library. Options: - numRows: specify number of rows to generate. Default value is 3. classmethod def name(cls) - str: return fake def schema(self) - Union[StructType, str]: return name string, date string, zipcode string, state string def reader(self, schema: StructType) - DataSourceReader: return FakeDataSourceReader(schema, self.options) def writer(self, schema: StructType, overwrite: bool) - DataSourceWriter: return FakeDataSourceWriter(self.options) def streamReader(self, schema: StructType) - DataSourceStreamReader: return FakeStreamReader(schema, self.options) def streamWriter(self, schema: StructType, overwrite: bool) - DataSourceStreamWriter: return FakeStreamWriter(self.options)这里schema()返回 DDL 字符串与返回StructType等价源码 docstring 中两种写法都有示例。self.options即用户通过.option(key, value)传入的配置字典注意其中每个值都是字符串类型。实现 Batch ReaderFakeDataSourceReader的read()方法利用faker库按 schema 中的每个字段名动态调用对应的 fake 方法填充数据from typing import Dict class FakeDataSourceReader(DataSourceReader): def __init__(self, schema: StructType, options: Dict[str, str]): self.schema: StructType schema self.options options def read(self, partition): from faker import Faker fake Faker() # Note: every value in this self.options dictionary is a string. num_rows int(self.options.get(numRows, 3)) for _ in range(num_rows): row [] for field in self.schema.fields: value getattr(fake, field.name)() row.append(value) yield tuple(row)注意faker是在read()方法内部导入的这正对应下文的序列化要求方法内引用的库必须在方法内部导入以保证读取器可被 pickle 序列化并分发到执行器。DataSourceReader的基类实现见 datasource.py 第 513 行起还提供了两个可选的性能钩子partitions()默认返回[InputPartition(None)]单个分区覆盖它返回 N 个分区查询规划器就会创建 N 个并行任务。读取大数据集时官方建议覆盖此方法。pushFilters(filters)谓词下推接口传入 Spark 希望下推的过滤条件列表列表内条件为 AND 关系返回仍需 Spark 在扫描后自行求值的过滤条件。默认返回全部过滤条件表示不下推。过滤条件由Filter及EqualTo、GreaterThan、LessThan、In、IsNull、Not、StringStartsWith等子类表示。将 Limit 下推给 Batch Reader当查询只需要前几行时让数据源少拉取数据能显著降低开销例如给 REST 请求加一个页大小参数或给底层 SQL 加上LIMIT子句。实现方式是在读取器中实现pushLimit(limit)方法from typing import Dict from pyspark.sql.datasource import DataSourceReader, InputPartition from pyspark.sql.types import StructType class FakeDataSourceReader(DataSourceReader): def __init__(self, schema: StructType, options: Dict[str, str]): self.schema: StructType schema self.options options self.limit None def pushLimit(self, limit: int) - bool: self.limit limit return True def partitions(self): # A limit makes a single request cheaper than a fan-out, since every # partition opens its own connection to the data source. if self.limit is not None: return [InputPartition(None)] return [InputPartition(i) for i in range(16)] def read(self, partition): num_rows int(self.options.get(numRows, 3)) if self.limit is not None: num_rows min(num_rows, self.limit) ...关于pushLimit源码datasource.py 第 591 行起给出了非常精确的语义约定调用时机在查询规划阶段被调用一次且先于partitions()与read()当查询带有可下推的过滤条件时它位于pushFilters()之后运行。开关配置仅当spark.sql.python.limitPushdown.enabled设为true时才会调用pushLimit默认关闭。该配置在 python/pyspark/sql/worker/utils.py 第 85 行有对应映射实现。仅为提示hintSpark 总是会在扫描之后再次应用 limit因此即使read()返回多于limit的行也不会出错返回True不会导致查询看到少于需求的行数。状态初始化pushFilters不一定被调用无过滤条件的查询不会调用它因此pushLimit依赖的中间状态应放在__init__中初始化同时pushLimit返回False时规划会当作从未调用过它partitions()与read()将在未观察到这些修改的读取器上执行。不适用场景LIMIT 0会被优化为空关系不会到达数据源pushLimit收到的limit恒为正数。仓库的官方测试test_python_datasource.pypython/pyspark/sql/tests/test_python_datasource.py中包含大量围绕pushLimit的用例覆盖了spark.sql.python.limitPushdown.enabled开启/关闭、pushLimit与pushFilters的先后关系、返回False时的回退行为等边界场景可作为实现时的行为参考。此外还有一个细节只有当所有过滤条件都被成功下推时limit 才会被下推——因为 Spark 无法在仍需自行求值的过滤条件之前应用 limit所以要想与过滤条件配合获得下推收益pushFilters应返回空的可迭代对象。实现 Batch Writer批式写入器由三部分组成write()在执行器上逐分区写数据、commit()所有任务成功后在驱动器上统一提交、abort()部分任务失败时执行回滚。下面的示例统计每个分区写入的行数成功后打印总行数失败则打印失败任务数from dataclasses import dataclass from typing import Iterator, List from pyspark.sql.types import Row from pyspark.sql.datasource import DataSource, DataSourceWriter, WriterCommitMessage dataclass class SimpleCommitMessage(WriterCommitMessage): partition_id: int count: int class FakeDataSourceWriter(DataSourceWriter): def write(self, rows: Iterator[Row]) - SimpleCommitMessage: from pyspark import TaskContext context TaskContext.get() partition_id context.partitionId() cnt sum(1 for _ in rows) return SimpleCommitMessage(partition_idpartition_id, countcnt) def commit(self, messages: List[SimpleCommitMessage]) - None: total_count sum(message.count for message in messages) print(fTotal number of rows: {total_count}) def abort(self, messages: List[SimpleCommitMessage]) - None: failed_count sum(message is None for message in messages) print(fNumber of failed tasks: {failed_count})从源码看datasource.py 第 1067 行起DataSourceWriter.write()在每个执行器上被调用一次返回一个可序列化的WriterCommitMessage允许返回None驱动器收集所有任务或任务失败时收集已成功的部分的提交消息后调用commit(messages)或abort(messages)。如果某个写任务失败其提交消息为None。实现 Streaming Reader流式读取器基于 offset 工作。下面的FakeStreamReader在每个微批生成 2 行数据其 offset 在每个微批递增 2class RangePartition(InputPartition): def __init__(self, start: int, end: int): self.start start self.end end class FakeStreamReader(DataSourceStreamReader): def __init__(self, schema, options): self.current 0 def initialOffset(self) - dict: Return the initial start offset of the reader. return {offset: 0} def latestOffset(self) - dict: Return the current latest offset that the next microbatch will read to. self.current 2 return {offset: self.current} def partitions(self, start: dict, end: dict) - list[InputPartition]: Plans the partitioning of the current microbatch defined by start and end offset, it needs to return a sequence of :class:InputPartition object. return [RangePartition(start[offset], end[offset])] def commit(self, end: dict) - None: This is invoked when the query has finished processing data before end offset, this can be used to clean up resource. pass def read(self, partition) - Iterator[Tuple]: Takes a partition as an input and read an iterator of tuples from the data source. start, end partition.start, partition.end for i in range(start, end): yield (i, str(i))源码中DataSourceStreamReaderdatasource.py 第 771 行起对各方法的约定如下initialOffset()返回新查询的起始 offset若查询是从 checkpoint 恢复则从已记录的 offset 重启而非初始 offset。latestOffset(start, limit)返回最新可用 offset。新查询第一个微批的start来自initialOffset()后续微批的start是上一微批的结束 offset无新数据时可原样返回start。Spark 4.2 起新增了start与limit参数旧签名仍受支持但新数据源建议采用新签名。partitions(start, end)根据 start/end offset 规划分区当start end空区间时应返回空序列。read(partition)按分区读取数据返回元组或Row也支持 ArrowRecordBatch。该方法应是静态无状态的不要在多次调用间依赖可变的类成员状态。commit(end)/stop()分别在批处理完成后、查询终止时被回调用于释放资源。替代方案实现 Simple Streaming Reader如果数据源吞吐量低、不需要分区可以实现SimpleDataSourceStreamReader替代DataSourceStreamReader。可读的流式数据源必须实现streamReader()与simpleStreamReader()中的至少一个且simpleStreamReader()仅在streamReader()未实现时被调用from pyspark.sql.datasource import SimpleDataSourceStreamReader class FakeDataSource(DataSource): ... def simpleStreamReader(self, schema: StructType) - SimpleDataSourceStreamReader: return FakeSimpleStreamReader() # omit implementation of streamReader ...下面的FakeSimpleStreamReader与前面流式读取器行为一致每个批次生成 2 行但用更简化的接口实现from typing import Iterator, Tuple from pyspark.sql.datasource import SimpleDataSourceStreamReader class FakeSimpleStreamReader(SimpleDataSourceStreamReader): def initialOffset(self) - dict: Return the initial start offset of the reader. return {offset: 0} def read(self, start: dict) - Tuple[Iterator[Tuple], dict]: Takes start offset as an input, return an iterator of tuples and the end offset (start offset for the next read). The end offset must advance past the start offset when returning data; otherwise Spark raises a validation exception. For example, returning 2 records from start_idx 0 means end should be {offset: 2} (i.e. start 2). When there is no data to read, you may return the same offset as end and start, but you must provide an empty iterator. start_idx start[offset] it iter([(i,) for i in range(start_idx, start_idx 2)]) return (it, {offset: start_idx 2}) def readBetweenOffsets(self, start: dict, end: dict) - Iterator[Tuple]: Takes start and end offset as input and read an iterator of data deterministically. This is called whe query replay batches during restart or after failure. start_idx start[offset] end_idx end[offset] return iter([(i,) for i in range(start_idx, end_idx)]) def commit(self, end: dict) - None: This is invoked when the query has finished processing data before end offset, this can be used to clean up resource. pass关于SimpleDataSourceStreamReaderdatasource.py 第 971 行起源码明确说明它不需要规划数据分区read()允许同时读取数据并规划最新 offset由于它在 Spark 驱动器节点上顺序读取以确定每个批次的结束 offset、不分区并行因此只适合输入速率和批次规模较小的轻量场景高吞吐场景应使用DataSourceStreamReader。另外两点语义约定read()返回的 end offset 必须越过 start offset返回数据时否则 Spark 会抛出校验异常无数据可读时end 与 start 可以相同但必须返回空迭代器。readBetweenOffsets()在查询重启或失败后重放批次时被调用用于确定性重读指定区间数据。流式读取器的准入控制Admission Control为限制每个微批处理的数据量可以实现getDefaultReadLimit()并让latestOffset(start, limit)遵守引擎传入的ReadLimitfrom pyspark.sql.streaming.datasource import ReadAllAvailable, ReadLimit, ReadMaxRows class MyStreamReader(DataSourceStreamReader): def getDefaultReadLimit(self) - ReadLimit: Limit each micro-batch to at most 20 rows. This value is just an example; in practice, configure the limit based on source options (e.g., self.options.get(maxRowsPerBatch)). return ReadMaxRows(20) def latestOffset(self, start: dict, limit: ReadLimit) - dict: Return the latest offset, respecting the provided limit. current start[offset] if isinstance(limit, ReadMaxRows): end min(current limit.max_rows, self.max_available) elif isinstance(limit, ReadAllAvailable): end self.max_available else: raise ValueError(fUnexpected ReadLimit type: {type(limit)}) return {offset: end}当 Spark 使用默认的ReadMaxRows(20)限制时每个微批最多读取 20 行视可用数据量而定若引擎传入ReadAllAvailable读取器则应返回全部剩余行。从源码看ReadLimit家族定义在 python/pyspark/sql/streaming/datasource.py 中Spark 4.2.0 起支持以下内置类型类型字段语义ReadAllAvailable无读取所有可用数据无视源自身的 options 配置ReadMinRowsmin_rows至少读取 N 行不足 N 行时源应推迟产生新 offset等待更多数据到达注意与Trigger.AvailableNow语义不兼容需源自行处理等待问题ReadMaxRowsmax_rows最多读取 N 行ReadMaxFilesmax_files最多读取 N 个文件ReadMaxBytesmax_bytes最多读取 N 字节getDefaultReadLimit()是可选实现默认返回ReadAllAvailable()表示对latestOffset()返回的数据量不加限制。引擎即使源返回了其他 read limit也可能出于触发语义trigger用ReadAllAvailable调用latestOffset()因此源必须始终遵守引擎给定的 readLimit——例如收到ReadAllAvailable时必须忽略通过 options 配置的读取限制。准入控制的价值体现在三个方面控制数据摄入速率避免冲垮下游系统内存管理限制批次大小避免 OOM背压处理以可持续的速率处理数据。仓库提供了完整可运行的示例 examples/src/main/python/sql/streaming/structured_blockchain_admission_control.py它模拟了一条 10000 个区块的区块链BlockchainStreamReader通过getDefaultReadLimit()返回ReadMaxRows(20)使每个微批最多处理 20 个区块块 0-19、块 20-39、块 40-59……最后一个批次可能不足 20 个。示例中每个区块携带区块号、模拟哈希、时间戳与交易数四个字段可用以下命令运行bin/spark-submit examples/src/main/python/sql/streaming/structured_blockchain_admission_control.py另外源码中还有SupportsTriggerAvailableNow混入接口实现它的流式源可通过prepareForTriggerAvailableNow()支持Trigger.AvailableNow语义记录当前最新数据作为查询的目标 offset。实现 Streaming Writer流式写入器的write()在每个执行器上按微批写入数据并返回提交消息commit()/abort()则带batchId参数在驱动器上执行。下面的FakeStreamWriter把每个微批的元数据行数、分区数写入本地路径下的 JSON 文件from typing import Iterator, List, Optional from pyspark.sql import Row from pyspark.sql.datasource import DataSourceStreamWriter, WriterCommitMessage class SimpleCommitMessage(WriterCommitMessage): partition_id: int count: int class FakeStreamWriter(DataSourceStreamWriter): def __init__(self, options): self.options options self.path self.options.get(path) assert self.path is not None def write(self, iterator: Iterator[Row]) - WriterCommitMessage: Write the data and return the commit message of that partition from pyspark import TaskContext context TaskContext.get() partition_id context.partitionId() cnt 0 for row in iterator: cnt 1 return SimpleCommitMessage(partition_idpartition_id, countcnt) def commit(self, messages: List[Optional[SimpleCommitMessage]], batchId: int) - None: Receives a sequence of :class:SimpleCommitMessage when all write tasks succeed and decides what to do with it. In this FakeStreamWriter, we write the metadata of the microbatch (number of rows and partitions) into a json file inside commit(). status dict(num_partitionslen(messages), rowssum(m.count for m in messages)) with open(os.path.join(self.path, f{batchId}.json), a) as file: file.write(json.dumps(status) \n) def abort(self, messages: List[Optional[SimpleCommitMessage]], batchId: int) - None: Receives a sequence of :class:SimpleCommitMessage from successful tasks when some tasks fail and decides what to do with it. In this FakeStreamWriter, we write a failure message into a txt file inside abort(). with open(os.path.join(self.path, f{batchId}.txt), w) as file: file.write(ffailed in batch {batchId})源码中DataSourceStreamWriterdatasource.py 第 1185 行起的约定与批式写入器一致write()返回可序列化提交消息所有任务成功后驱动器调用commit(messages, batchId)任一任务失败则调用abort(messages, batchId)batchId是每个微批唯一递增的整数。若要与原生支持 Arrow 的系统对接可改用DataSourceStreamArrowWriter它的write()接收的是 PyArrowRecordBatch迭代器。序列化要求用户自定义的DataSource、DataSourceReader、DataSourceWriter、DataSourceStreamReader、DataSourceStreamWriter及其方法都必须能够被 pickle 序列化。这是因为读取器/写入器会被序列化后分发到各执行器上运行。对于仅在方法内部使用的库必须在方法内部导入例如def read(self, partition): from pyspark import TaskContext context TaskContext.get()使用 Python 数据源注册 Python 数据源定义完数据源后使用前必须先注册spark.dataSource.register(FakeDataSource)从 Python 数据源读取使用默认 schema 和默认选项读取numRows默认 3 行spark.read.format(fake).load().show() # ----------------------------------- # | name| date|zipcode| state| # ----------------------------------- # |Carlos Cobb|2018-07-15| 73003|Indiana| # | Eric Scott|1991-08-22| 10085| Idaho| # | Amy Martin|1988-10-28| 68076| Oregon| # -----------------------------------使用自定义 schema 读取schema 中字段名会映射到faker的同名生成器spark.read.format(fake).schema(name string, company string).load().show() # ----------------------------------- # |name |company | # ----------------------------------- # |Tanner Brennan |Adams Group | # |Leslie Maxwell |Santiago Group| # |Mrs. Jacqueline Brown|Maynard Inc | # -----------------------------------使用不同行数读取spark.read.format(fake).option(numRows, 5).load().show() # ------------------------------------------- # | name| date|zipcode| state| # ------------------------------------------- # | Pam Mitchell|1988-10-20| 23788| Tennessee| # |Melissa Turner|1996-06-14| 30851| Nevada| # | Brian Ramsey|2021-08-21| 55277| Washington| # | Caitlin Reed|1983-06-22| 89813|Pennsylvania| # | Douglas James|2007-01-18| 46226| Alabama| # -------------------------------------------向 Python 数据源写入写入自定义位置时必须指定mode()子句支持append与overwrite两种模式df spark.range(0, 10, 1, 5) df.write.format(fake).mode(append).save() # You can check the Spark log (standard error) to see the output of the write operation. # Total number of rows: 10上例中spark.range(0, 10, 1, 5)生成 10 行数据并划分为 5 个分区每个分区在write()中计数最终commit()打印总行数。对应mode(overwrite)时writer()工厂方法会收到overwriteTrue标志。在流式查询中使用 Python 数据源注册后Python 数据源既可以作为readStream()的源也可以作为writeStream()的汇通过短名或全名传给format()。从 fake 数据源流式读取并输出到 consolequery spark.readStream.format(fake).load().writeStream.format(console).start() # --- # | id| # --- # | 0| # | 1| # --- # --- # | id| # --- # | 2| # | 3| # ---同一个数据源同时用于流式读取与流式写入query spark.readStream.format(fake).load().writeStream.format(fake).start(/output_path)此时FakeStreamReader每个微批产出 2 行FakeStreamWriter则把每个微批的元数据写入/output_path下的{batchId}.json成功或{batchId}.txt失败。支持直接产出 Arrow Batch 的高性能读取器Python 数据源读取器支持直接产出 Arrow Batchpyarrow.RecordBatch以列式内存格式跳过逐行处理的开销在大数据集场景下可将处理性能提升至一个数量级。启用方式非常简单在DataSourceReader或DataSourceStreamReader的read()方法中产出pyarrow.RecordBatch对象即可。源码中DataSourceReader.read()的返回类型注解为Union[Iterator[Tuple], Iterator[RecordBatch]]明确支持两种产出方式。以下示例实现了一个基于 Arrow Batch 的基础数据源from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition from pyspark.sql import SparkSession import pyarrow as pa # Define the ArrowBatchDataSource class ArrowBatchDataSource(DataSource): A Data Source for testing Arrow Batch Serialization classmethod def name(cls): return arrowbatch def schema(self): return key int, value string def reader(self, schema: str): return ArrowBatchDataSourceReader(schema, self.options) # Define the ArrowBatchDataSourceReader class ArrowBatchDataSourceReader(DataSourceReader): def __init__(self, schema, options): self.schema: str schema self.options options def read(self, partition): # Create Arrow Record Batch keys pa.array([1, 2, 3, 4, 5], typepa.int32()) values pa.array([one, two, three, four, five], typepa.string()) schema pa.schema([(key, pa.int32()), (value, pa.string())]) record_batch pa.RecordBatch.from_arrays([keys, values], schemaschema) yield record_batch def partitions(self): # Define the number of partitions num_part 1 return [InputPartition(i) for i in range(num_part)] # Initialize the Spark Session spark SparkSession.builder.appName(ArrowBatchExample).getOrCreate() # Register the ArrowBatchDataSource spark.dataSource.register(ArrowBatchDataSource) # Load data using the custom data source df spark.read.format(arrowbatch).load() df.show()同样地写入侧也有对应的DataSourceArrowWriterSpark 4.0 引入与DataSourceStreamArrowWriterSpark 4.1 引入它们与普通 writer 的区别仅在于write()接收的是 PyArrowRecordBatch迭代器而非 Row 迭代器适合与原生支持 Arrow 的系统或库对接见 datasource.py 第 1134 行与第 1256 行。使用注意事项官方文档的 Usage Notes 部分给出了三条重要的使用约定名称解析优先级在数据源解析过程中内置的以及 Scala/Java 数据源优先于同名的 Python 数据源。若想显式使用 Python 数据源请确保其名称不与任何非 Python 数据源冲突。同名注册覆盖允许多个 Python 数据源注册同名后注册的会覆盖先注册的。自动注册机制要实现数据源的自动注册请将其导出为名称带pyspark_前缀的顶层模块中的DefaultSource。官方文档以pyspark_huggingface项目为参考示例可自行搜索该开源项目了解具体做法。总结本文从最小可运行的 Batch Reader出发逐步构建了一个同时支持批式读写、流式读写、Limit 下推与 Arrow 直传的完整 Python 数据源。所有接口均可在 python/pyspark/sql/datasource.py 与 python/pyspark/sql/streaming/datasource.py 中查看完整定义与 docstring行为约定可参考官方测试 python/pyspark/sql/tests/test_python_datasource.py完整可运行示例见 examples/src/main/python/sql/streaming/structured_blockchain_admission_control.py。掌握这套 API 后接入 REST API、消息队列、数据库或自定义文件格式等外部系统都只需一份纯 Python 实现无需再触碰 Scala/Java 侧的数据源接口。【免费下载链接】sparkApache Spark - A unified analytics engine for large-scale data processing项目地址: https://gitcode.com/gh_mirrors/sp/spark创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考