{"id":13400253,"url":"https://github.com/uber/petastorm","last_synced_at":"2025-04-10T23:25:28.225Z","repository":{"id":37923914,"uuid":"137539238","full_name":"uber/petastorm","owner":"uber","description":"Petastorm library enables single machine or distributed training and evaluation of deep learning models from datasets in Apache Parquet format. It supports ML frameworks such as Tensorflow, Pytorch, and PySpark and can be used from pure Python code.","archived":false,"fork":false,"pushed_at":"2023-12-02T05:11:31.000Z","size":2819,"stargazers_count":1828,"open_issues_count":178,"forks_count":281,"subscribers_count":37,"default_branch":"master","last_synced_at":"2025-04-03T13:14:06.461Z","etag":null,"topics":["deep-learning","machine-learning","parquet","parquet-files","pyarrow","pyspark","pytorch","sysml","tensorflow"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/uber.png","metadata":{"files":{"readme":"README.rst","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2018-06-15T23:15:29.000Z","updated_at":"2025-03-28T09:01:44.000Z","dependencies_parsed_at":"2023-02-18T01:45:59.420Z","dependency_job_id":"ba7a6ea3-ae05-40b9-86ea-8e0fc6a7c7ce","html_url":"https://github.com/uber/petastorm","commit_stats":{"total_commits":672,"total_committers":52,"mean_commits":"12.923076923076923","dds":0.5119047619047619,"last_synced_commit":"0b0775af42539189d913702d9695566431dabd8a"},"previous_names":[],"tags_count":114,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/uber%2Fpetastorm","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/uber%2Fpetastorm/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/uber%2Fpetastorm/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/uber%2Fpetastorm/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/uber","download_url":"https://codeload.github.com/uber/petastorm/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248313706,"owners_count":21082898,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","machine-learning","parquet","parquet-files","pyarrow","pyspark","pytorch","sysml","tensorflow"],"created_at":"2024-07-30T19:00:49.962Z","updated_at":"2025-04-10T23:25:28.202Z","avatar_url":"https://github.com/uber.png","language":"Python","readme":"\nPetastorm\n=========\n\n.. image:: https://github.com/uber/petastorm/actions/workflows/unittest.yml/badge.svg?branch=master\n   :target: https://github.com/uber/petastorm/actions/workflows/unittest.yml\n   :alt: Build Status\n\n.. image:: https://codecov.io/gh/uber/petastorm/branch/master/graph/badge.svg\n   :target: https://codecov.io/gh/uber/petastorm/branch/master\n   :alt: Code coverage\n\n.. image:: https://img.shields.io/badge/License-Apache%202.0-blue.svg\n   :target: https://img.shields.io/badge/License-Apache%202.0-blue.svg\n   :alt: License\n\n.. image:: https://badge.fury.io/py/petastorm.svg\n   :target: https://pypi.org/project/petastorm\n   :alt: Latest Version\n\n.. inclusion-marker-start-do-not-remove\n\n.. contents::\n\n\nPetastorm is an open source data access library developed at Uber ATG. This library enables single machine or\ndistributed training and evaluation of deep learning models directly from datasets in Apache Parquet\nformat. Petastorm supports popular Python-based machine learning (ML) frameworks such as\n`Tensorflow \u003chttp://www.tensorflow.org/\u003e`_, `PyTorch \u003chttps://pytorch.org/\u003e`_, and\n`PySpark \u003chttp://spark.apache.org/docs/latest/api/python/pyspark.html\u003e`_. It can also be used from pure Python code.\n\nDocumentation web site: `\u003chttps://petastorm.readthedocs.io\u003e`_\n\n\n\nInstallation\n------------\n\n.. code-block:: bash\n\n    pip install petastorm\n\n\nThere are several extra dependencies that are defined by the ``petastorm`` package that are not installed automatically.\nThe extras are: ``tf``, ``tf_gpu``, ``torch``, ``opencv``, ``docs``, ``test``.\n\nFor example to trigger installation of GPU version of tensorflow and opencv, use the following pip command:\n\n.. code-block:: bash\n\n    pip install petastorm[opencv,tf_gpu]\n\n\n\nGenerating a dataset\n--------------------\n\nA dataset created using Petastorm is stored in `Apache Parquet \u003chttps://parquet.apache.org/\u003e`_ format.\nOn top of a Parquet schema, petastorm also stores higher-level schema information that makes multidimensional arrays into a native part of a petastorm dataset. \n\nPetastorm supports extensible data codecs. These enable a user to use one of the standard data compressions (jpeg, png) or implement her own.\n\nGenerating a dataset is done using PySpark.\nPySpark natively supports Parquet format, making it easy to run on a single machine or on a Spark compute cluster.\nHere is a minimalistic example writing out a table with some random data.\n\n\n.. code-block:: python\n\n   import numpy as np\n   from pyspark.sql import SparkSession\n   from pyspark.sql.types import IntegerType\n\n   from petastorm.codecs import ScalarCodec, CompressedImageCodec, NdarrayCodec\n   from petastorm.etl.dataset_metadata import materialize_dataset\n   from petastorm.unischema import dict_to_spark_row, Unischema, UnischemaField\n\n   # The schema defines how the dataset schema looks like\n   HelloWorldSchema = Unischema('HelloWorldSchema', [\n       UnischemaField('id', np.int32, (), ScalarCodec(IntegerType()), False),\n       UnischemaField('image1', np.uint8, (128, 256, 3), CompressedImageCodec('png'), False),\n       UnischemaField('array_4d', np.uint8, (None, 128, 30, None), NdarrayCodec(), False),\n   ])\n\n\n   def row_generator(x):\n       \"\"\"Returns a single entry in the generated dataset. Return a bunch of random values as an example.\"\"\"\n       return {'id': x,\n               'image1': np.random.randint(0, 255, dtype=np.uint8, size=(128, 256, 3)),\n               'array_4d': np.random.randint(0, 255, dtype=np.uint8, size=(4, 128, 30, 3))}\n\n\n   def generate_petastorm_dataset(output_url='file:///tmp/hello_world_dataset'):\n       rowgroup_size_mb = 256\n\n       spark = SparkSession.builder.config('spark.driver.memory', '2g').master('local[2]').getOrCreate()\n       sc = spark.sparkContext\n\n       # Wrap dataset materialization portion. Will take care of setting up spark environment variables as\n       # well as save petastorm specific metadata\n       rows_count = 10\n       with materialize_dataset(spark, output_url, HelloWorldSchema, rowgroup_size_mb):\n\n           rows_rdd = sc.parallelize(range(rows_count))\\\n               .map(row_generator)\\\n               .map(lambda x: dict_to_spark_row(HelloWorldSchema, x))\n\n           spark.createDataFrame(rows_rdd, HelloWorldSchema.as_spark_schema()) \\\n               .coalesce(10) \\\n               .write \\\n               .mode('overwrite') \\\n               .parquet(output_url)\n\n\n- ``HelloWorldSchema`` is an instance of a ``Unischema`` object.\n  ``Unischema`` is capable of rendering types of its fields into different\n  framework specific formats, such as: Spark ``StructType``, Tensorflow\n  ``tf.DType`` and numpy ``numpy.dtype``.\n- To define a dataset field, you need to specify a ``type``, ``shape``, a\n  ``codec`` instance and whether the field is nullable for each field of the\n  ``Unischema``.\n- We use PySpark for writing output Parquet files. In this example, we launch\n  PySpark on a local box (``.master('local[2]')``). Of course for a larger\n  scale dataset generation we would need a real compute cluster.\n- We wrap spark dataset generation code with the ``materialize_dataset``\n  context manager.  The context manager is responsible for configuring row\n  group size at the beginning and write out petastorm specific metadata at the\n  end.\n- The row generating code is expected to return a Python dictionary indexed by\n  a field name. We use ``row_generator`` function for that. \n- ``dict_to_spark_row`` converts the dictionary into a ``pyspark.Row``\n  object while ensuring schema ``HelloWorldSchema`` compliance (shape,\n  type and is-nullable condition are tested).\n- Once we have a ``pyspark.DataFrame`` we write it out to a parquet\n  storage. The parquet schema is automatically derived from\n  ``HelloWorldSchema``.\n\nPlain Python API\n----------------\nThe ``petastorm.reader.Reader`` class is the main entry point for user\ncode that accesses the data from an ML framework such as Tensorflow or Pytorch.\nThe reader has multiple features such as:\n\n- Selective column readout\n- Multiple parallelism strategies: thread, process, single-threaded (for debug)\n- N-grams readout support\n- Row filtering (row predicates)\n- Shuffling\n- Partitioning for multi-GPU training\n- Local caching\n\nReading a dataset is simple using the ``petastorm.reader.Reader`` class which can be created using the\n``petastorm.make_reader`` factory method:\n\n.. code-block:: python\n\n   from petastorm import make_reader\n\n    with make_reader('hdfs://myhadoop/some_dataset') as reader:\n       for row in reader:\n           print(row)\n\n``hdfs://...`` and ``file://...`` are supported URL protocols.\n\nOnce a ``Reader`` is instantiated, you can use it as an iterator.\n\nTensorflow API\n--------------\n\nTo hookup the reader into a tensorflow graph, you can use the ``tf_tensors``\nfunction:\n\n.. code-block:: python\n\n    from petastorm.tf_utils import tf_tensors\n\n    with make_reader('file:///some/localpath/a_dataset') as reader:\n       row_tensors = tf_tensors(reader)\n       with tf.Session() as session:\n           for _ in range(3):\n               print(session.run(row_tensors))\n\nAlternatively, you can use new ``tf.data.Dataset`` API;\n\n.. code-block:: python\n\n    from petastorm.tf_utils import make_petastorm_dataset\n\n    with make_reader('file:///some/localpath/a_dataset') as reader:\n        dataset = make_petastorm_dataset(reader)\n        iterator = dataset.make_one_shot_iterator()\n        tensor = iterator.get_next()\n        with tf.Session() as sess:\n            sample = sess.run(tensor)\n            print(sample.id)\n\nPytorch API\n-----------\n\nAs illustrated in\n`pytorch_example.py \u003chttps://github.com/uber/petastorm/blob/master/examples/mnist/pytorch_example.py\u003e`_,\nreading a petastorm dataset from pytorch\ncan be done via the adapter class ``petastorm.pytorch.DataLoader``,\nwhich allows custom pytorch collating function and transforms to be supplied.\n\nBe sure you have ``torch`` and ``torchvision`` installed:\n\n.. code-block:: bash\n\n    pip install torchvision\n\nThe minimalist example below assumes the definition of a ``Net`` class and\n``train`` and ``test`` functions, included in ``pytorch_example``:\n\n.. code-block:: python\n\n    import torch\n    from petastorm.pytorch import DataLoader\n\n    torch.manual_seed(1)\n    device = torch.device('cpu')\n    model = Net().to(device)\n    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)\n\n    def _transform_row(mnist_row):\n        transform = transforms.Compose([\n            transforms.ToTensor(),\n            transforms.Normalize((0.1307,), (0.3081,))\n        ])\n        return (transform(mnist_row['image']), mnist_row['digit'])\n\n\n    transform = TransformSpec(_transform_row, removed_fields=['idx'])\n\n    with DataLoader(make_reader('file:///localpath/mnist/train', num_epochs=10,\n                                transform_spec=transform, seed=1, shuffle_rows=True), batch_size=64) as train_loader:\n        train(model, device, train_loader, 10, optimizer, 1)\n    with DataLoader(make_reader('file:///localpath/mnist/test', num_epochs=10,\n                                transform_spec=transform), batch_size=1000) as test_loader:\n        test(model, device, test_loader)\n\nIf you are working with very large batch sizes and do not need support for Decimal/strings we provide a ``petastorm.pytorch.BatchedDataLoader`` that can buffer using Torch tensors (``cpu`` or ``cuda``) with a signficantly higher throughput.\n\nIf the size of your dataset can fit into system memory, you can use an in-memory version dataloader ``petastorm.pytorch.InMemBatchedDataLoader``. This dataloader only reades the dataset once, and caches data in memory to avoid additional I/O for multiple epochs.\n\nSpark Dataset Converter API\n---------------------------\n\nSpark converter API simplifies the data conversion from Spark to TensorFlow or PyTorch.\nThe input Spark DataFrame is first materialized in the parquet format and then loaded as\na ``tf.data.Dataset`` or ``torch.utils.data.DataLoader``.\n\nThe minimalist example below assumes the definition of a compiled ``tf.keras`` model and a\nSpark DataFrame containing a feature column followed by a label column.\n\n.. code-block:: python\n\n    from petastorm.spark import SparkDatasetConverter, make_spark_converter\n    import tensorflow.compat.v1 as tf  # pylint: disable=import-error\n\n    # specify a cache dir first.\n    # the dir is used to save materialized spark dataframe files\n    spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, 'hdfs:/...')\n\n    df = ... # `df` is a spark dataframe\n\n    # create a converter from `df`\n    # it will materialize `df` to cache dir.\n    converter = make_spark_converter(df)\n\n    # make a tensorflow dataset from `converter`\n    with converter.make_tf_dataset() as dataset:\n        # the `dataset` is `tf.data.Dataset` object\n        # dataset transformation can be done if needed\n        dataset = dataset.map(...)\n        # we can train/evaluate model on the `dataset`\n        model.fit(dataset)\n        # when exiting the context, the reader of the dataset will be closed\n\n    # delete the cached files of the dataframe.\n    converter.delete()\n\nThe minimalist example below assumes the definition of a ``Net`` class and\n``train`` and ``test`` functions, included in\n`pytorch_example.py \u003chttps://github.com/uber/petastorm/blob/master/examples/mnist/pytorch_example.py\u003e`_,\nand a Spark DataFrame containing a feature column followed by a label column.\n\n.. code-block:: python\n\n    from petastorm.spark import SparkDatasetConverter, make_spark_converter\n\n    # specify a cache dir first.\n    # the dir is used to save materialized spark dataframe files\n    spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, 'hdfs:/...')\n\n    df_train, df_test = ... # `df_train` and `df_test` are spark dataframes\n    model = Net()\n\n    # create a converter_train from `df_train`\n    # it will materialize `df_train` to cache dir. (the same for df_test)\n    converter_train = make_spark_converter(df_train)\n    converter_test = make_spark_converter(df_test)\n\n    # make a pytorch dataloader from `converter_train`\n    with converter_train.make_torch_dataloader() as dataloader_train:\n        # the `dataloader_train` is `torch.utils.data.DataLoader` object\n        # we can train model using the `dataloader_train`\n        train(model, dataloader_train, ...)\n        # when exiting the context, the reader of the dataset will be closed\n\n    # the same for `converter_test`\n    with converter_test.make_torch_dataloader() as dataloader_test:\n        test(model, dataloader_test, ...)\n\n    # delete the cached files of the dataframes.\n    converter_train.delete()\n    converter_test.delete()\n\n\nAnalyzing petastorm datasets using PySpark and SQL\n--------------------------------------------------\n\nA Petastorm dataset can be read into a Spark DataFrame using PySpark, where you can\nuse a wide range of Spark tools to analyze and manipulate the dataset.\n\n.. code-block:: python\n\n   # Create a dataframe object from a parquet file\n   dataframe = spark.read.parquet(dataset_url)\n\n   # Show a schema\n   dataframe.printSchema()\n\n   # Count all\n   dataframe.count()\n\n   # Show a single column\n   dataframe.select('id').show()\n\nSQL can be used to query a Petastorm dataset:\n\n.. code-block:: python\n\n   spark.sql(\n      'SELECT count(id) '\n      'from parquet.`file:///tmp/hello_world_dataset`').collect()\n\nYou can find a full code sample here: `pyspark_hello_world.py \u003chttps://github.com/uber/petastorm/blob/master/examples/hello_world/petastorm_dataset/pyspark_hello_world.py\u003e`_,\n\nNon Petastorm Parquet Stores\n----------------------------\nPetastorm can also be used to read data directly from Apache Parquet stores. To achieve that, use\n``make_batch_reader`` (and not ``make_reader``). The following table summarizes the differences\n``make_batch_reader`` and ``make_reader`` functions.\n\n\n==================================================================  =====================================================\n``make_reader``                                                     ``make_batch_reader``\n==================================================================  =====================================================\nOnly Petastorm datasets (created using materializes_dataset)        Any Parquet store (some native Parquet column types\n                                                                    are not supported yet.\n------------------------------------------------------------------  -----------------------------------------------------\nThe reader returns one record at a time.                            The reader returns batches of records. The size of the\n                                                                    batch is not fixed and defined by Parquet row-group\n                                                                    size.\n------------------------------------------------------------------  -----------------------------------------------------\nPredicates passed to ``make_reader`` are evaluated per single row.  Predicates passed to ``make_batch_reader`` are evaluated per batch.\n------------------------------------------------------------------  -----------------------------------------------------\nCan filter parquet file based on the ``filters`` argument.          Can filter parquet file based on the ``filters`` argument\n==================================================================  =====================================================\n\n\nTroubleshooting\n---------------\n\nSee the Troubleshooting_ page and please submit a ticket_ if you can't find an\nanswer.\n\n\nSee also\n--------\n\n1. Gruener, R., Cheng, O., and Litvin, Y. (2018) *Introducing Petastorm: Uber ATG's Data Access Library for Deep Learning*. URL: https://eng.uber.com/petastorm/\n2. QCon.ai 2019: `\"Petastorm: A Light-Weight Approach to Building ML Pipelines\" \u003chttps://www.infoq.com/presentations/petastorm-ml-pipelines/\u003e`_.\n\n\n.. _Troubleshooting: docs/troubleshoot.rst\n.. _ticket: https://github.com/uber/petastorm/issues/new\n.. _Development: docs/development.rst\n\nHow to Contribute\n=================\n\nWe prefer to receive contributions in the form of GitHub pull requests. Please send pull requests against the ``github.com/uber/petastorm`` repository.\n\n- If you are looking for some ideas on what to contribute, check out `github issues \u003chttps://github.com/uber/petastorm/issues\u003e`_ and comment on the issue.\n- If you have an idea for an improvement, or you'd like to report a bug but don't have time to fix it please a `create a github issue \u003chttps://github.com/uber/petastorm/issues/new\u003e`_.\n\nTo contribute a patch:\n\n- Break your work into small, single-purpose patches if possible. It's much harder to merge in a large change with a lot of disjoint features.\n- Submit the patch as a GitHub pull request against the master branch. For a tutorial, see the GitHub guides on forking a repo and sending a pull request.\n- Include a detailed describtion of the proposed change in the pull request.\n- Make sure that your code passes the unit tests. You can find instructions how to run the unit tests `here \u003chttps://github.com/uber/petastorm/blob/master/docs/development.rst\u003e`_.\n- Add new unit tests for your code.\n\nThank you in advance for your contributions!\n\n\nSee the Development_ for development related information.\n\n\n.. inclusion-marker-end-do-not-remove\n   Place contents above here if they should also appear in read-the-docs.\n   Contents below are already part of the read-the-docs table of contents.\n\n","funding_links":[],"categories":["分布式机器学习","Deep Learning Framework","Python","AI","tensorflow","Optimization Tools"],"sub_categories":["Deployment \u0026 Distribution"],"project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fuber%2Fpetastorm","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fuber%2Fpetastorm","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fuber%2Fpetastorm/lists"}