Path forward on IO operations

[WardBrian]

This issue serves to document plans that myself, @bob-carpenter, and @mitzimorris have all discussed in person.

Issues:

• the existing IO code re-uses very little between the different services.
• the existing IO code is slow.
• the existing IO code depends on things like the order of comments in the files.

Steps:

• Split IO code into two passes. One pass should ignore all comments and use a fast reader like numpy's fromtxt to read in the draws. This can be used for all services. The second pass should deal with the comments. This will be service-specific. This change can be done any time
• Replace the second pass of the above with the newer JSON outputs rather than CSV comments. This requires a minimum supported cmdstan version: Remove support for CmdStan pre-2.34 in a future version #685
• The first step could later be replaced with some sort of binary output. This would require upstream Stan changes

[eroesch]

Thanks for creating the issue!

When I spoke to Mitzi, the other day, I had something like that in mind, but I was not thinking of using numpy, rather navigate through the file with seek() [1]. I cannot imagine this would be slower than numpy.fromtxt().

with open("bernoulli-20250416172900_1.csv", "r") as f:

    # Move to byte 1678 in my test file
    # which is the start of the chain data
    f.seek(1678)  # probably NOT fixed across files though
    beg_pos = f.tell()

    # Read in the length of the file
    f.seek(0, 2)
    end_pos = f.tell()

    # Move back to exclude the tail information about timing
    offset = 129  # probably fixed across files
    remaining = end_pos - 1678 - offset

    f.seek(1678)
    data = f.read(remaining)
    print(data)

This is very fast. The only annoying thing about that solution is that the beg_pos where the chain data starts is possibly not the same, depending on the number of characters in the preamble that's commented. This is not the case of the tail of the csv, because the formatting is, I believe, forcing a dedicated structure likely to always be the same size in bytes.

Arguably, we could undershoot roughly where we think we're close to the data (like in my example), then either filter out the few remaining lines to the data. That's very likely to be faster than the current solution, which is to read each line, check if they start with "#" and filter out, line after line.

That would work, I think, but it would be cleaner to separate the preamble from the data directly in Stan, adding a new flag like save_cmdstan_config to save only the data of the chains.

The other solution I was thinking about was using grep, which is quite optimised already. :) .. but I haven't tested it.

We probably need a way to evaluate the time that these steps in cmdstanpy actually take. What do you use?

Have a great Easter break!

References:

• https://docs.python.org/3/tutorial/inputoutput.html#methods-of-file-objects

[WardBrian]

I'm reasonable certain what is slow is not seeking around the file or skipping the # rows, but actually parsing the numbers into floats. Right now, sampling does this in an all-python loop:

cmdstanpy/cmdstanpy/stanfit/mcmc.py

Lines 465 to 468 in 0c25ec4

	for i in range(sampling_iter_start, num_draws):
	line = fd.readline().strip()
	xs = line.split(',')
	self._draws[i, chain, :] = [float(x) for x in xs]

Replacing this with np.loadtxt (probably combining with the extra logic in that file that loads the warmup if they're present) should be both straightforward and a good bit faster, because the equivalent of the double loop we're doing there ends up being done in C code

[WardBrian]

I did some rough timings on a fake CSV file with 1000 rows and 45 columns

def read_manual(filename):
    # similar to current cmdstanpy code
    draws = np.empty((1000, 45), dtype=float)
    with open(filename, 'r') as f:
        for i in range(1000):
            line = f.readline().strip()
            xs = line.split(',')
            draws[i, :] = [float(x) for x in xs]
    return draws

15.5 ms ± 526 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

---

def read_np(filename):
    with open(filename, 'r') as f:
        return np.loadtxt(f, delimiter=',', comments='#')

11.9 ms ± 215 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

---

def read_pandas(filename):
    with open(filename, 'r') as f:
        return pd.read_csv(f, engine='c', comment='#').to_numpy()

8.39 ms ± 314 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

---

So actually using pandas seems to be the fastest, even if we actually want the numpy data out. And we can probably get it to ~twice as fast as our current code just by switching.

engine='pyarrow' for pandas is actually the fastest (3.8 ms ± 246 μs), but doesn't allow the comment argument, so we can't use it with the current Stan CSV files. And it would be an extra dependency, at which point we could also consider options like dask etc that also claim to have fast CSV readers.

[eroesch]

Nice. You are right, pandas is probably the good enough route, minimising hassle. Shall I have a go at incorporating it? @mitzimorris, any other idea?
Have a great rest of the week!

[mitzimorris]

thanks @eroesch and looking forward to the PR.

[ahartikainen]

Hey, you are worrying about the wrong thing.

CSV files created with stan mcmc process are wide (100s - 100k columns) and usually not that long (<1M rows). (this is due to nD array transformed to 1D presentation)

It should not matter do you read small file in 250ms or 500ms. What should matter is that can you read 100k columns with 100 rows.

Pandas will choke, np.load_txt is not your friend either.

I think best option is to read file line by line and then profile what takes least amount of time to convert one line to floats.

np.frombuffer/np.fromxyz (? I don't remember what was the best but one of the np.from functions) I think has best performance when file is read with "rb" flag.

Also you can profile how long it takes to iterate through files line by line and check first non-whitespace char.

Edit. When I imply load_text is slow, I mean for a full file

[bob-carpenter]

Here's a simple test of default I/O in pandas for data frames with headers. There aren't any comment lines, so I don't know if that would actually slow this down.

$ python3 ptest2.py
  100 ×   1000 write:   0.1s read:    0.0s
  100 ×  10000 write:   1.2s read:    0.4s
  100 × 100000 write:  25.3s read:   11.2s
 1000 ×   1000 write:   1.1s read:    0.2s
 1000 ×  10000 write:  13.7s read:    2.9s
 1000 × 100000 write: 242.2s read:  102.0s
10000 ×   1000 write:  11.2s read:    1.8s
10000 ×  10000 write: 136.3s read:   27.9s

I think this is a lot faster than whatever we have currently. A C++ program should be able to read data like this very quickly. 10K x 10K is less than 1GB of data at double precision. I ran out of patience before the 10K x 100K timing completed.

Here's the code.

import time
import numpy as np
import pandas as pd

def eval(rows, cols):
    df = pd.DataFrame(np.random.randn(rows, cols))
    fn = f"data_{rows}x{cols}.csv"
    with open(fn, "w", newline="") as f:
        f.write(",".join(f'"{c}"' for c in df.columns) + "\n")
        t0 = time.perf_counter()
        df.to_csv(f, index=False, header=False)
    t1 = time.perf_counter()
    df2 = pd.read_csv(fn, engine="c", comment="#")
    t2 = time.perf_counter()
    print(f"{rows:5d} × {cols:6d} write: {(t1 - t0):5.1f}s read:  {(t2 - t1):5.1f}s")

for rows in [100, 1000, 10_000]:
    for cols in [1000, 10_000, 100_000]:
        eval(rows, cols)

[ahartikainen]

Nice to see that pandas IO has been going forward. I did this test couple of years ago (or more than couple) and I think 100k columns never finished.

[eroesch]

I have done more tests this weekend, and I am getting between 2-4x speedup. I'll get cracking on it. There are quite a few places in the codebase using that original algorithm, and I suspect a few edge cases. I'll proceed with caution!

[amas0]

Sounds like pandas is a reasonable direction here, but I wanted to chime in and say that you can get around the engine='pyarrow' comment limitation by filtering out the comment lines in Python and feeding a file-like object back to these optimized readers.

I created a Stan CSV file, 2000 sampling iterations and 2000 parameters. First, I tested the pandas C engine with conversion to numpy:

import pandas as pd

def read_pandas_c(file):
    return pd.read_csv(file, engine="c", comment="#").to_numpy()

804 ms ± 63.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Compared to manually parsing out the comments and using the pyarrow engine + conversion to numpy:

import io
import pandas as pd

def read_pandas_pyarrow(file):
    with open(file, "rb") as f:
        no_comments = (line for line in f if not line.startswith(b"#"))
        no_comment_filelike = io.BytesIO(b"".join(no_comments))
        return pd.read_csv(no_comment_filelike, engine="pyarrow").to_numpy()

312 ms ± 16.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

So you don't pay much of a real performance cost by filtering out the comments manually and are able to get the pyarrow benefits. The main difference between the C and pyarrow engines seems to be that the latter is multithreaded. You'd still have to pay the cost of adding the dependency, but it's a pretty noticeable performance jump for large files.

[WardBrian]

Nice find @amas0. I suspect you might be able to get event faster if we wrote our own little TextIO wrapper to avoid needing to .join the entire no_comments iterator, but that would depend on how they actually use it under the hood

@eroesch are you still interested in working on this issue?

[eroesch]

I am, yes, but haven't done much yet.