Ray Data 公共 API 全景基于 autosummary 索引解析 Dataset、DataIterator 與聚合體系【免費下載鏈接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.項目地址: 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 個核心公共 APIDataset、DataIterator、Schema、DatasetSummary、GroupedData、AggregateFn、AggregateFnV2并結(jié)合 python/ray/data 下的源碼實現(xiàn)與同目錄 API 參考文檔講清每個類的職責(zé)邊界、典型用法與底層原理。讀完本文你將能夠看懂 Ray Data API 文檔的組織與生成機(jī)制掌握數(shù)據(jù)集從創(chuàng)建、變換到消費的完整調(diào)用鏈并具備編寫自定義分布式聚合AggregateFnV2的實戰(zhàn)能力。一、_autogen.rst是什么Ray Data API 文檔的生成樞紐_autogen.rst是一份 Sphinx 的autosummary指令文件文件頭注釋明確說明了它的用途僅用于自動生成 API 文檔不應(yīng)被納入 toctree。任何需要出現(xiàn)在 API 文檔中的類都需要先加入這份清單隨后 Sphinx 會為清單中的每個條目生成對應(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幾個值得注意的細(xì)節(jié).. currentmodule:: ray.data將后續(xù)條目的默認(rèn)模塊設(shè)置為ray.data因此清單中的Schema實際指向ray.data.Schema而帶模塊前綴的stats.DatasetSummary、grouped_data.GroupedData、aggregate.AggregateFn、aggregate.AggregateFnV2則分別指向ray.data.stats、ray.data.grouped_data、ray.data.aggregate三個子模塊。:nosignatures:表示生成條目時不展示函數(shù)簽名適合以類為單位的索引頁:template: autosummary/class_v2.rst指定了類模板:toctree:讓生成的 RST 文件落入 toctree形成可導(dǎo)航的子頁面。這 7 個條目含 8 個類恰好構(gòu)成了 Ray Data 用戶最常接觸的門面 API數(shù)據(jù)集本體Dataset、迭代消費接口DataIterator、類型描述Schema、統(tǒng)計匯總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。整個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)鍵信息可以拆解為三點塊Block即并行單元一個 Dataset 在物理上被切分為若干 Arrow 格式的塊每塊對應(yīng)一個 RayObjectRef塊的數(shù)量決定了并行度。這意味著用戶看到的是一個邏輯上的數(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)變換與消費分離、執(zhí)行惰性化Dataset有兩類操作——變換transformation產(chǎn)生新 Dataset如map_batches()、sort()、random_shuffle()、repartition()、join()消費consumption產(chǎn)生具體值如iter_batches()、take_all()、max()。變換是惰性的只有下游消費觸發(fā)時才真正執(zhí)行這為執(zhí)行引擎做算子融合與流式調(diào)度留出了空間 ds ray.data.range(1000) ds.max(id) # 消費類操作立即觸發(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計算策略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等已廢棄接口后者仍保留在文檔中以幫助用戶平滑遷移。三、消費側(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()可獲取一個迭代器 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 中每個 Trainer Actor 應(yīng)當(dāng)通過ray.train.get_dataset_shard(train)獲取屬于自己的迭代器分片實現(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)橋梁。底層實現(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)計Schema與DatasetSummary4.1Schema數(shù)據(jù)集的列類型描述Schema定義在 python/ray/data/dataset.py是對底層 PyArrow Schema或 Pandas Block Schema的一層封裝base_schema屬性保存底層 Arrow / Pandas schemanames屬性返回用戶可見的列名列表實現(xiàn)中會過濾掉__bsp_stub這類讀路徑注入的物理占位列types屬性返回各列的 Arrow 類型對非 Arrow 兼容類型統(tǒng)一返回object實現(xiàn)中針對pd.ArrowDtype、pd.StringDtype、掩碼類型BaseMaskedDtype做了到pyarrow.DataType的轉(zhuǎn)換并將 Ray 的張量擴(kuò)展類型create_arrow_fixed_shape_tensor_type納入類型系統(tǒng)。一個值得注意的實現(xiàn)細(xì)節(jié)Schema在構(gòu)造時會快照當(dāng)前的DataContextcopy.deepcopy(DataContext.get_current())注釋表明數(shù)據(jù)集創(chuàng)建時刻的配置決定了其行為即 Dataset 的配置語義在創(chuàng)建時被凍結(jié)避免運行期上下文漂移。4.2DatasetSummary數(shù)據(jù)集統(tǒng)計匯總DatasetSummary定義在 python/ray/data/stats.py是一個標(biāo)注了stabilityalpha的dataclass用于承載對數(shù)據(jù)集列計算出的統(tǒng)計信息如 min/max/count/missing_pct 等。其實現(xiàn)有一個值得稱道的細(xì)節(jié)由于聚合結(jié)果可能與原列類型不一致例如字符串列的count是 int64DatasetSummary內(nèi)部把統(tǒng)計拆成兩張 PyArrow 表——_stats_matching_column_dtype與原列同類型的統(tǒng)計保留原 dtype和_stats_mismatching_column_dtype如 count、missing_pct 等類型不同的統(tǒng)計并在to_pandas()時合并為單個 DataFrame遇到TensorExtensionType等擴(kuò)展類型轉(zhuǎn)換失敗時會降級為逐列轉(zhuǎn)換、將問題列轉(zhuǎn)為 null 類型見 python/ray/data/stats.py。這保證了統(tǒng)計結(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í)行任何分組邏輯只是返回一個持有dataset、key、num_partitions的輕量句柄真正的分組發(fā)生在調(diào)用aggregate()之時。aggregate(*aggs)將分組鍵與聚合算子打包成一個Aggregate邏輯算子接入 Dataset 的邏輯計劃LogicalPlan并返回新的Dataset。結(jié)果數(shù)據(jù)集有n 1列第一列是分組鍵其后依次是各聚合結(jié)果當(dāng)分組鍵為None時鍵列省略。GroupedData還提供map_groups()方法支持對每個組應(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 個條目基類AggregateFnV2、AggregateFn與 16 個內(nèi)置實現(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)為每個組創(chuàng)建初始累加器如 0、空列表或空字典accumulate_row/accumulate_block逐行或逐塊更新累加器兩者必須且只能提供一個支持向量化批量處理merge(c1, c2)合并不同 worker 產(chǎn)生的兩個累加器finalize(acc)可選將最終累加器轉(zhuǎn)換為輸出不提供則原樣返回。經(jīng)典示例按組計數(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對每個分組或整個數(shù)據(jù)集用zero_factory創(chuàng)建初始累加器塊聚合Block Aggregationaggregate_block獨立作用于每個塊產(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。兩個泛型參數(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)。一個簡單的自定義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): # 對單個塊做向量化求和返回局部累加器 ... 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ò)展點。這意味著用戶既可以開箱即用地落盤也可以基于 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è)計遵循小而穩(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ù)千核集群上實現(xiàn)自定義分布式聚合這套 API 都提供了從入口到擴(kuò)展的完整路徑?！久赓M下載鏈接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.項目地址: https://gitcode.com/gh_mirrors/ra/ray創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考