Implement a Mars Operand#

Use read_csv as an example to illustrate how to implement a Mars operand.

Define Operand Class#

All Mars operands inherit from the base class Operand, it defines the basic properties of operand, each module has it’s own child class, such as DataFrameOperand, TensorOperand, etc. For tilebale operand, it also needs to inherit from TileableOperandMixin to implement tile and execute functions. So we firstly define operand class and set output types in init method, the types could be DataFrame or Tensor which depends on the type of operand’s output data, __call__ method is also needed for creating a Mars dataframe.

# NOTE: Use relative import if in Mars modules
from mars.dataframe.operands import DataFrameOperand, DataFrameOperandMixin
from mars.core import OutputType
from mars.serialization.serializables import StringField

class SimpleReadCSV(DataFrameOperand, DataFrameOperandMixin):
    path = StringField('path')

    def __init__(self, **kw):

    def __call__(self, index_value=None, columns_value=None,
                 dtypes=None, chunk_bytes=None):
        # np.nan means its size is unknown on axis 0
        shape = (np.nan, len(dtypes))
        return self.new_dataframe(

For the SimpleReadCSV operand, the property path means the path of csv file, we use a StringField to indicate the property’s type which is useful for serialization. If the type is uncertain, AnyField will work.

Implement Tile Method#

Tile method is the next goal, this method will split the computing task into several sub tasks. Ideally, these tasks can be assigned on different executors in parallel. In the specific case of read_csv, each sub task read a block of bytes from the file, so we need calculate the offset and length of each block in the tile function. As we use the same class for both coarse-grained and fine-grained operand, offset, length and other properties are added to record information for fine-grained operand.

import os

import numpy as np

from mars.core import OutputType
from mars.dataframe.operands import DataFrameOperand, DataFrameOperandMixin
from mars.serialization.serializables import AnyField, StringField, Int64Field
from mars.utils import parse_readable_size

class SimpleReadCSV(DataFrameOperand, DataFrameOperandMixin):
    path = StringField("path")
    chunk_bytes = AnyFiled('chunk_bytes')
    offset = Int64Field("offset")
    length = Int64Field("length")

    def __init__(self, **kw):

    def tile(cls, op: "SimpleReadCSV"):
        # out is operand's output in coarse-grained graph
        out = op.outputs[0]

        file_path = op.path
        file_size = os.path.getsize(file_path)

        # split file into chunks
        chunk_bytes = int(parse_readable_size(op.chunk_bytes)[0])
        offset = 0
        index_num = 0
        out_chunks = []
        while offset < file_size:
            chunk_op = op.copy().reset_key()
            chunk_op.path = file_path
            # offset and length for current chunk
            chunk_op.offset = offset
            chunk_op.length = min(chunk_bytes, file_size - offset)
            # calculate chunk's meta, including shape, index_value, columns_value
            # here we use np.nan to represent unknown shape
            shape = (np.nan, len(out.dtypes))
            # use `op.new_chunk` to create a dataframe chunk
            new_chunk = chunk_op.new_chunk(
                index=(index_num, 0),
            offset += chunk_bytes
            index_num += 1

        # create a new tileable which holds `chunks` for generating fine-grained graph
        new_op = op.copy()
        # `nsplits` records the split info for each axis. For read_csv,
        # the output dataframe is split into multiple parts on axis 0 and
        # keep one chunk on axis 1, so the nsplits will be
        # like ((np.nan, np.nan, ...), (out.shape[1],))
        nsplits = ((np.nan,) * len(out_chunks), (out.shape[1],))
        return new_op.new_dataframes(

Implement Execute Method#

When sub task is delivered to executor, Mars will call operand’s execute method to perform calculations. When it comes to read_csv, we need read the block from the file according to the offset and length, however the offset is a rough position as we can’t read a csv file from the middle of a line, using readline to find the starts and ends at delimiter boundaries.

from io import BytesIO

import pandas as pd

from mars.dataframe.utils import build_empty_df

def _find_chunk_start_end(f, offset, size):
    if f.tell() == 0:
        start = 0
        start = f.tell() + size)
    end = f.tell()
    return start, end

class SimpleReadCSV(DataFrameOperand, DataFrameOperandMixin):
    def execute(cls, ctx: Union[dict, Context], op: "SimpleReadCSV"):
        out = op.outputs[0]
        with open(op.path, 'rb') as f:
            start, end = _find_chunk_start_end(f, op.offset, op.length)
            if end == start:
                # the last chunk may be empty
                data = build_empty_df(out.dtypes)
                if start == 0:
                    # The first chunk contains header, skip header rows
                    data = pd.read_csv(BytesIO( - start)),
                    data = pd.read_csv(BytesIO( - start)),

        ctx[out.key] = data

After reading the chunk data by pd.read_csv, we store the results in ctx. SimpleReadCSV only has one output here, for operand like SVD that has multiple outputs, we can store them separately using output’s keys.

Define User Interface#

Finally, we need define function read_csv exposed to users. In this function, besides creating a SimpleReadCSV operand, a sample data is taken to infer some meta information of Mars DataFrame, such as dtypes, columns, index, etc.

from mars.dataframe.utils import parse_index

def read_csv(file_path, chunk_bytes='16M'):
    # use first 10 lines to infer
    with open(file_path, 'rb') as f:
        head_lines = b"".join([f.readline() for _ in range(10)])

    mini_df = pd.read_csv(BytesIO(head_lines))
    index_value = parse_index(mini_df.index)
    columns_value = parse_index(mini_df.columns, store_data=True)
    dtypes = mini_df.dtypes
    op = SimpleReadCSV(path=file_path, chunk_bytes=chunk_bytes)
    return op(

Write Tests#

Mars uses pytest for testing, we can add tests under the tests subdirectory of the specific module and follow the current examples of tests. Define a test function and use the shared fixture setup to run your tests under the test environment.

def test_simple_read_csv_execution(setup):
    with tempfile.TemporaryDirectory() as tempdir:
        file_path = os.path.join(tempdir, "test.csv")
        # write to a csv file
        raw = pd.DataFrame({
            'int': range(10),
            'float': np.random.rand(10),
            'str': [f'value_{i}' for i in range(10)]
        }).to_csv(file_path, index=False)
        mdf = read_csv(file_path).execute().fetch()
        pd.testing.assert_frame_equal(raw, mdf)

When tests pass locally, we can submit a pull requests on GitHub, the test suite will run automatically on GitHub Actions and Azure Pipelines continuous integration services, if all checks have passed, it means the pull request is up to the quality of merging.

Documenting Your Code#

If the changes add APIs to Mars modules, we should document our code in docs directory, it can be done following the regarding documentation.