.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto/integrations/flytekit_plugins/dolt/dolt_quickstart_example.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        Click :ref:`here <sphx_glr_download_auto_integrations_flytekit_plugins_dolt_dolt_quickstart_example.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_integrations_flytekit_plugins_dolt_dolt_quickstart_example.py:


Quickstart
----------

In this demo and following example, learn how to use ``DoltTable`` to annotate DataFrame inputs and outputs in the Flyte tasks.

.. youtube:: TfBIuHYkyvU

.. GENERATED FROM PYTHON SOURCE LINES 11-12

First, let's import the libraries.

.. GENERATED FROM PYTHON SOURCE LINES 12-19

.. code-block:: default

    import os
    import sys

    import pandas as pd
    from flytekit import task, workflow
    from flytekitplugins.dolt.schema import DoltConfig, DoltTable


.. GENERATED FROM PYTHON SOURCE LINES 20-21

Next, we initialize Dolt's config.

.. GENERATED FROM PYTHON SOURCE LINES 21-28

.. code-block:: default

    doltdb_path = os.path.join(os.path.dirname(__file__), "foo")

    rabbits_conf = DoltConfig(
        db_path=doltdb_path,
        tablename="rabbits",
    )


.. GENERATED FROM PYTHON SOURCE LINES 29-30

We define a task to create a DataFrame and store the table in Dolt.

.. GENERATED FROM PYTHON SOURCE LINES 30-39

.. code-block:: default



    @task
    def populate_rabbits(a: int) -> DoltTable:
        rabbits = [("George", a), ("Alice", a * 2), ("Sugar Maple", a * 3)]
        df = pd.DataFrame(rabbits, columns=["name", "count"])
        return DoltTable(data=df, config=rabbits_conf)



.. GENERATED FROM PYTHON SOURCE LINES 40-41

``unwrap_rabbits`` task does the exact opposite -- reading the table from Dolt and returning a DataFrame.

.. GENERATED FROM PYTHON SOURCE LINES 41-46

.. code-block:: default

    @task
    def unwrap_rabbits(table: DoltTable) -> pd.DataFrame:
        return table.data



.. GENERATED FROM PYTHON SOURCE LINES 47-48

Our workflow combines the above two tasks:

.. GENERATED FROM PYTHON SOURCE LINES 48-63

.. code-block:: default

    @workflow
    def wf(a: int) -> pd.DataFrame:
        rabbits = populate_rabbits(a=a)
        df = unwrap_rabbits(table=rabbits)
        return df


    if __name__ == "__main__":
        print(f"Running {__file__} main...")
        if len(sys.argv) != 2:
            raise ValueError("Expected 1 argument: a (int)")
        a = int(sys.argv[1])
        result = wf(a=a)
        print(f"Running wf(), returns dataframe\n{result}\n{result.dtypes}")


.. GENERATED FROM PYTHON SOURCE LINES 64-69

Run this task by issuing the following command:

.. prompt:: $

  python quickstart_example.py 1


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.000 seconds)


.. _sphx_glr_download_auto_integrations_flytekit_plugins_dolt_dolt_quickstart_example.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example


    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: dolt_quickstart_example.py <dolt_quickstart_example.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: dolt_quickstart_example.ipynb <dolt_quickstart_example.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_