Ray Data 公共 API 全景基于 autosummary 索引解析 Dataset、DataIterator 與聚合體系【免費(fèi)下載鏈接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.項(xiàng)目地址: https://gitcode.com/gh_mirrors/ra/rayRay Data 是 Ray 面向 AI 工作負(fù)載的分布式數(shù)據(jù)處理引擎圍繞 Arrow 塊Block構(gòu)建了一套延遲變換 流式執(zhí)行的數(shù)據(jù)集抽象。本文以倉庫中 doc/source/data/api/_autogen.rst 這份自動生成 API 文檔的索引清單為骨架逐一解析其中列出的 8 個(gè)核心公共 APIDataset、DataIterator、Schema、DatasetSummary、GroupedData、AggregateFn、AggregateFnV2并結(jié)合 python/ray/data 下的源碼實(shí)現(xiàn)與同目錄 API 參考文檔講清每個(gè)類的職責(zé)邊界、典型用法與底層原理。讀完本文你將能夠看懂 Ray Data API 文檔的組織與生成機(jī)制掌握數(shù)據(jù)集從創(chuàng)建、變換到消費(fèi)的完整調(diào)用鏈并具備編寫自定義分布式聚合AggregateFnV2的實(shí)戰(zhàn)能力。一、_autogen.rst是什么Ray Data API 文檔的生成樞紐_autogen.rst是一份 Sphinx 的autosummary指令文件文件頭注釋明確說明了它的用途僅用于自動生成 API 文檔不應(yīng)被納入 toctree。任何需要出現(xiàn)在 API 文檔中的類都需要先加入這份清單隨后 Sphinx 會為清單中的每個(gè)條目生成對應(yīng)的ray.data.class.rst文件再由頂層 RST 文件包含這些生成產(chǎn)物。該文件的完整清單如下.. currentmodule:: ray.data .. autosummary:: :nosignatures: :template: autosummary/class_v2.rst :toctree: DataIterator Dataset Schema stats.DatasetSummary grouped_data.GroupedData aggregate.AggregateFn aggregate.AggregateFnV2幾個(gè)值得注意的細(xì)節(jié).. currentmodule:: ray.data將后續(xù)條目的默認(rèn)模塊設(shè)置為ray.data因此清單中的Schema實(shí)際指向ray.data.Schema而帶模塊前綴的stats.DatasetSummary、grouped_data.GroupedData、aggregate.AggregateFn、aggregate.AggregateFnV2則分別指向ray.data.stats、ray.data.grouped_data、ray.data.aggregate三個(gè)子模塊。:nosignatures:表示生成條目時(shí)不展示函數(shù)簽名適合以類為單位的索引頁:template: autosummary/class_v2.rst指定了類模板:toctree:讓生成的 RST 文件落入 toctree形成可導(dǎo)航的子頁面。這 7 個(gè)條目含 8 個(gè)類恰好構(gòu)成了 Ray Data 用戶最常接觸的門面 API數(shù)據(jù)集本體Dataset、迭代消費(fèi)接口DataIterator、類型描述Schema、統(tǒng)計(jì)匯總DatasetSummary、分組中間態(tài)GroupedData以及兩代自定義聚合基類AggregateFn/AggregateFnV2。從源碼結(jié)構(gòu)看這一清單與 python/ray/data 模塊根目錄下的模塊劃分一一對應(yīng)dataset.py定義Dataset與Schemaiterator.py定義DataIteratorgrouped_data.py定義GroupedDataaggregate.py定義聚合基類stats.py定義DatasetSummary。整個(gè)ray.data包還有read_api.py、block.py、datatype.py、context.py等模塊共同支撐起完整的數(shù)據(jù)集能力。二、核心數(shù)據(jù)抽象Dataset—— 分布式數(shù)據(jù)集合Dataset是 Ray Data 的基石。源碼 python/ray/data/dataset.py 給出了精確定義A Dataset is a distributed data collection for data loading and processing. Datasets are distributed pipelines that produceObjectRef[Block]outputs, where each block holds data in Arrow format, representing a shard of the overall data collection. The block also determines the unit of parallelism.關(guān)鍵信息可以拆解為三點(diǎn)塊Block即并行單元一個(gè) Dataset 在物理上被切分為若干 Arrow 格式的塊每塊對應(yīng)一個(gè) RayObjectRef塊的數(shù)量決定了并行度。這意味著用戶看到的是一個(gè)邏輯上的數(shù)據(jù)集底層卻是分布式對象存儲中的一組可并行處理的分片。三種創(chuàng)建途徑從外部存儲本地磁盤、S3、HDFS 等通過read_*()API 讀入從內(nèi)存數(shù)據(jù)通過from_*()API 構(gòu)造從合成數(shù)據(jù)通過range_*()API 生成。例如import ray # 合成數(shù)據(jù) ds ray.data.range(1000) # 內(nèi)存數(shù)據(jù) ds ray.data.from_items([{col1: i, col2: i * 2} for i in range(1000)]) # 外部存儲對象存儲/文件系統(tǒng) ds ray.data.read_parquet(s3://bucket/path) # 寫回外部存儲 ds.write_csv(s3://bucket/output)變換與消費(fèi)分離、執(zhí)行惰性化Dataset有兩類操作——變換transformation產(chǎn)生新 Dataset如map_batches()、sort()、random_shuffle()、repartition()、join()消費(fèi)consumption產(chǎn)生具體值如iter_batches()、take_all()、max()。變換是惰性的只有下游消費(fèi)觸發(fā)時(shí)才真正執(zhí)行這為執(zhí)行引擎做算子融合與流式調(diào)度留出了空間 ds ray.data.range(1000) ds.max(id) # 消費(fèi)類操作立即觸發(fā)執(zhí)行 999 ds.random_shuffle() # 變換類操作返回新的 Dataset未物化 shape: (1000, 1) ╭───────╮ │ id │ │ --- │ │ int64 │ ╰───────╯ (Dataset isnt materialized)在 doc/source/data/api/dataset.rst 中Dataset的 API 參考還延伸出三組周邊 APICompute Strategy計(jì)算策略ActorPoolStrategy與TaskPoolStrategy用于控制map_batches等變換采用 Actor 池可復(fù)用狀態(tài)、適合 GPU 與重量級初始化還是輕量 Task適合無狀態(tài)短任務(wù)來執(zhí)行。Schema 與 DatasetSummary見本文第四、五節(jié)。Developer API 與 Deprecated APIto_pandas_refs、to_numpy_refs、to_arrow_refs、iter_internal_ref_bundles等內(nèi)部能力以及to_random_access_dataset、iter_tf_batches等已廢棄接口后者仍保留在文檔中以幫助用戶平滑遷移。三、消費(fèi)側(cè)抽象DataIterator—— 單次遍歷讀取全部記錄DataIterator是定義在 python/ray/data/iterator.py 的抽象基類abc.ABC用于從Dataset讀取記錄。它的核心語義源碼 docstring是For Datasets, each iteration call represents a complete read of all items in the Dataset.即每次對DataIterator的迭代都是一次對 Dataset 全量數(shù)據(jù)的完整讀取。Dataset.iterator()可獲取一個(gè)迭代器 ds ray.data.range(5) ds.iterator() DataIterator(shape: (5, 1) ╭───────╮ │ id │ │ --- │ │ int64 │ ╰───────╯ (Dataset isnt materialized))DataIterator是 Ray Data 與下游訓(xùn)練框架銜接的關(guān)鍵接口在 Ray Train 中每個(gè) Trainer Actor 應(yīng)當(dāng)通過ray.train.get_dataset_shard(train)獲取屬于自己的迭代器分片實(shí)現(xiàn)數(shù)據(jù)在分布式訓(xùn)練 worker 間的切分迭代器提供iter_batches()、iter_torch_batches()、iter_tf_batches()等批量讀取方法可直接產(chǎn)出 PyTorch / TensorFlow 所需的張量批次是數(shù)據(jù)處理管線 → 訓(xùn)練循環(huán)之間的標(biāo)準(zhǔn)橋梁。底層實(shí)現(xiàn)類為ray.data._internal.iterator.iterator_impl.DataIteratorImpl它持有_to_ref_bundle_iterator()等內(nèi)部方法將迭代操作映射到底層的 RefBundle 流式執(zhí)行之上見 python/ray/data/_internal/iterator/iterator_impl.py相關(guān)參考見 doc/source/data/api/data_iterator.rst。四、類型與統(tǒng)計(jì)Schema與DatasetSummary4.1Schema數(shù)據(jù)集的列類型描述Schema定義在 python/ray/data/dataset.py是對底層 PyArrow Schema或 Pandas Block Schema的一層封裝base_schema屬性保存底層 Arrow / Pandas schemanames屬性返回用戶可見的列名列表實(shí)現(xiàn)中會過濾掉__bsp_stub這類讀路徑注入的物理占位列types屬性返回各列的 Arrow 類型對非 Arrow 兼容類型統(tǒng)一返回object實(shí)現(xiàn)中針對pd.ArrowDtype、pd.StringDtype、掩碼類型BaseMaskedDtype做了到pyarrow.DataType的轉(zhuǎn)換并將 Ray 的張量擴(kuò)展類型create_arrow_fixed_shape_tensor_type納入類型系統(tǒng)。一個(gè)值得注意的實(shí)現(xiàn)細(xì)節(jié)Schema在構(gòu)造時(shí)會快照當(dāng)前的DataContextcopy.deepcopy(DataContext.get_current())注釋表明數(shù)據(jù)集創(chuàng)建時(shí)刻的配置決定了其行為即 Dataset 的配置語義在創(chuàng)建時(shí)被凍結(jié)避免運(yùn)行期上下文漂移。4.2DatasetSummary數(shù)據(jù)集統(tǒng)計(jì)匯總DatasetSummary定義在 python/ray/data/stats.py是一個(gè)標(biāo)注了stabilityalpha的dataclass用于承載對數(shù)據(jù)集列計(jì)算出的統(tǒng)計(jì)信息如 min/max/count/missing_pct 等。其實(shí)現(xiàn)有一個(gè)值得稱道的細(xì)節(jié)由于聚合結(jié)果可能與原列類型不一致例如字符串列的count是 int64DatasetSummary內(nèi)部把統(tǒng)計(jì)拆成兩張 PyArrow 表——_stats_matching_column_dtype與原列同類型的統(tǒng)計(jì)保留原 dtype和_stats_mismatching_column_dtype如 count、missing_pct 等類型不同的統(tǒng)計(jì)并在to_pandas()時(shí)合并為單個(gè) DataFrame遇到TensorExtensionType等擴(kuò)展類型轉(zhuǎn)換失敗時(shí)會降級為逐列轉(zhuǎn)換、將問題列轉(zhuǎn)為 null 類型見 python/ray/data/stats.py。這保證了統(tǒng)計(jì)結(jié)果既能保持類型嚴(yán)謹(jǐn)又對 pandas 用戶足夠友好。五、分組與聚合GroupedData與兩代AggregateFn5.1GroupedData惰性分組中間態(tài)GroupedData定義在 python/ray/data/grouped_data.pydocstring 明確指出Represents a grouped dataset created by callingDataset.groupby(). The actual groupby is deferred until an aggregation is applied.也就是說ds.groupby(key)并不會立即執(zhí)行任何分組邏輯只是返回一個(gè)持有dataset、key、num_partitions的輕量句柄真正的分組發(fā)生在調(diào)用aggregate()之時(shí)。aggregate(*aggs)將分組鍵與聚合算子打包成一個(gè)Aggregate邏輯算子接入 Dataset 的邏輯計(jì)劃LogicalPlan并返回新的Dataset。結(jié)果數(shù)據(jù)集有n 1列第一列是分組鍵其后依次是各聚合結(jié)果當(dāng)分組鍵為None時(shí)鍵列省略。GroupedData還提供map_groups()方法支持對每個(gè)組應(yīng)用用戶自定義的批處理函數(shù)fn并接受compute、batch_format、num_cpus、num_gpus、concurrency等資源與并行度參數(shù)適合分組后做復(fù)雜變換的場景。5.2 內(nèi)置聚合函數(shù)在 doc/source/data/api/aggregate.rst 中聚合 API 共列出 18 個(gè)條目基類AggregateFnV2、AggregateFn與 16 個(gè)內(nèi)置實(shí)現(xiàn)Count、Sum、Min、Max、Mean、Std、AbsMax、Quantile、Unique、AsList、CountDistinct、ValueCounter、MissingValuePercentage、ZeroPercentage、ApproximateQuantile、ApproximateTopK。這些內(nèi)置聚合統(tǒng)一通過Dataset.aggregate()或Dataset.groupby().aggregate()使用。例如import ray from ray.data.aggregate import Mean, ApproximateTopK ds ray.data.from_items( [{group: A, score: 1.0}, {group: A, score: 3.0}, {group: B, score: 2.0}] ) # 全量聚合 print(ds.aggregate(Mean(onscore)).take_all()) # 分組聚合 print(ds.groupby(group).aggregate(Mean(onscore)).take_all())5.3AggregateFn已廢棄基于 init/merge/accumulate 的自定義聚合AggregateFn定義在 python/ray/data/aggregate.py源碼中已被標(biāo)記Deprecated提示請使用AggregateFnV2代替但它定義了自定義聚合最經(jīng)典的四段式接口理解它有助于理解 V2init(key)為每個(gè)組創(chuàng)建初始累加器如 0、空列表或空字典accumulate_row/accumulate_block逐行或逐塊更新累加器兩者必須且只能提供一個(gè)支持向量化批量處理merge(c1, c2)合并不同 worker 產(chǎn)生的兩個(gè)累加器finalize(acc)可選將最終累加器轉(zhuǎn)換為輸出不提供則原樣返回。經(jīng)典示例按組計(jì)數(shù)import ray from ray.data.aggregate import AggregateFn count_agg AggregateFn( initlambda k: 0, accumulate_rowlambda counter, row: counter 1, mergelambda c1, c2: c1 c2, namecustom_count, ) ds ray.data.from_items([{group: A}, {group: B}, {group: A}]) result ds.groupby(group).aggregate(count_agg).take_all() # result: [{group: A, custom_count: 2}, {group: B, custom_count: 1}]5.4AggregateFnV2當(dāng)前推薦的高效聚合接口AggregateFnV2定義在 python/ray/data/aggregate.py繼承自AggregateFn并參數(shù)化為Generic[AccumulatorType, AggOutputType]是當(dāng)前官方推薦的自定義聚合方式。它的執(zhí)行被明確規(guī)范為四步初始化Initialization對每個(gè)分組或整個(gè)數(shù)據(jù)集用zero_factory創(chuàng)建初始累加器塊聚合Block Aggregationaggregate_block獨(dú)立作用于每個(gè)塊產(chǎn)出該塊的局部聚合結(jié)果可向量化性能關(guān)鍵合并Combinationcombine將各塊的局部結(jié)果合并為統(tǒng)一累加器終結(jié)Finalization可選的finalize將最終累加器轉(zhuǎn)換為輸出格式。構(gòu)造參數(shù)為name輸出列名如sum(my_col)、zero_factory初始零值工廠如求和用lambda: 0、求最小用lambda: float(inf)、on聚合列名None表示整行如Count()與ignore_nulls是否跳過 null。兩個(gè)泛型參數(shù)的含義在源碼中有清晰示例Count(AggregateFnV2[int, int]) # 累加器: int, 輸出: int Sum(AggregateFnV2[Union[int, float], ...]) # 累加器: 數(shù)值, 輸出: 數(shù)值 Mean(AggregateFnV2[List[...], float]) # 累加器: [sum, count], 輸出: float Std(AggregateFnV2[List[...], float]) # 累加器: [M2, mean, count], 輸出: float即AccumulatorType是中間狀態(tài)類型aggregate_block的產(chǎn)出、combine的輸入輸出、finalize的入?yún)ggOutputType是寫回結(jié)果集的最終類型。Mean用[sum, count]作為累加器再算出均值、Std用[M2, mean, count]承載二階矩——這種復(fù)合累加器正是 V2 接口表達(dá)力的體現(xiàn)。一個(gè)簡單的自定義Sum子類骨架如下from typing import Union from ray.data.aggregate import AggregateFnV2 class MySum(AggregateFnV2[Union[int, float], Union[int, float]]): def __init__(self, on: str): super().__init__( namefsum({on}), zero_factorylambda: 0, onon, ignore_nullsTrue, ) def aggregate_block(self, block, column_nameNone, *args, **kwargs): # 對單個(gè)塊做向量化求和返回局部累加器 ... def combine(self, left_acc, right_acc): return left_acc right_acc # finalize 可省略默認(rèn)原樣返回累加器作為輸出六、數(shù)據(jù)加載與保存 API 全景圍繞上述核心類Ray Data 提供了極廣的數(shù)據(jù)進(jìn)出通道分別收錄在 doc/source/data/api/loading_data.rst 與 doc/source/data/api/saving_data.rst 中這也是 API 參考頁的主體。6.1 加載側(cè)Loading Data API公共 API 按數(shù)據(jù)源組織主要包括內(nèi)存/框架互轉(zhuǎn)from_pandas、from_numpy、from_arrow、from_items、from_torch、from_tf、from_dask、from_modin、from_mars、from_daft、from_spark、from_huggingface文件與列式格式read_parquet、read_csv、read_json、read_orc、read_avro、read_numpy、read_text、read_tfrecords、read_mcap、read_webdataset、read_zarrAI 多模態(tài)數(shù)據(jù)read_images、read_videos、read_audio、read_lerobot、read_binary_files湖倉與數(shù)據(jù)庫read_bigquery、read_clickhouse、read_mongo、read_sql、read_snowflake、read_delta、read_delta_sharing_tables、read_hudi、read_iceberg、read_lance、read_kafka、read_databricks_tables、read_unity_catalog含Catalog/DatabricksUnityCatalog/ReaderFormat類合成數(shù)據(jù)range、range_tensor分區(qū)與重排輔助datasource.Partitioning、datasource.PartitionStyle、datasource.PathPartitionFilter、datasource.PathPartitionParser、FileShuffleConfig開發(fā)者 APIfrom_arrow_refs、from_numpy_refs、from_pandas_refs、Datasource、datasource.FileBasedDatasource、read_datasource、ReadTask以及元數(shù)據(jù)提供者datasource.BaseFileMetadataProvider、datasource.DefaultFileMetadataProvider、datasource.FileMetadataProvider。6.2 保存?zhèn)萐aving Data API保存?zhèn)扰c加載側(cè)對稱公共 API 包括Dataset.write_parquet、write_csv、write_json、write_orc、write_numpy、write_tfrecords、write_images、write_lance、write_mongo、write_clickhouse、write_bigquery、write_snowflake、write_sql、write_iceberg以及框架互轉(zhuǎn)to_pandas、to_numpy、to_dask、to_modin、to_mars、to_daft、to_spark開發(fā)者側(cè)還有Datasink、Dataset.write_datasink、datasource.RowBasedFileDatasink、datasource.BlockBasedFileDatasink、datasource.WriteResult、datasource.WriteReturnType、datasource.FilenameProvider等擴(kuò)展點(diǎn)。這意味著用戶既可以開箱即用地落盤也可以基于 Datasink 定制自己的輸出格式。七、擴(kuò)展閱讀與配套 API執(zhí)行配置ExecutionOptions與ExecutionResources見 doc/source/data/api/execution_options.rst用于控制數(shù)據(jù)集執(zhí)行的資源上限與策略數(shù)據(jù)類型DataType與TypeCategory見 doc/source/data/api/datatype.rst描述數(shù)據(jù)集支持的數(shù)據(jù)類型分類Checkpoint / 表達(dá)式 / DataContext分別見 checkpoint.rst、expressions.rst、data_context.rstpreprocessor.rst 與 llm.rst 覆蓋預(yù)處理與 LLM 生態(tài)集成API 總?cè)肟赿oc/source/data/api/api.md 通過 toctree 將上述所有 API 頁面組織為完整的 Ray Data API 章節(jié)_autogen.rst生成的類文檔即嵌入其中。整體上Ray Data 的 API 設(shè)計(jì)遵循小而穩(wěn)的門面 可擴(kuò)展的內(nèi)部接口原則Dataset、DataIterator、Schema等門面類保持穩(wěn)定而AggregateFnV2、Datasource、Datasink、BlockAccessor等接口則向開發(fā)者開放深度定制能力。無論是快速做數(shù)據(jù)預(yù)處理還是在訓(xùn)練管線中切分?jǐn)?shù)據(jù)批次或是在數(shù)千核集群上實(shí)現(xiàn)自定義分布式聚合這套 API 都提供了從入口到擴(kuò)展的完整路徑。【免費(fèi)下載鏈接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.項(xiàng)目地址: https://gitcode.com/gh_mirrors/ra/ray創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考