<jupyter_start><jupyter_text>Define formats<jupyter_code>def hdf_writer(format='fixed', compression=None):
    def writer(df, filename):
        pd.DataFrame.to_hdf(
            df, filename, 
            key='key',
            format=format,
            complib=compression
    )
    return writer


def parquet_writer(compression=None):
    def writer(df, filename):
        pd.DataFrame.to_parquet(
            df, filename, 
            compression=compression,
    )
    return writer


formats = [
    dict(name='csv', reader=pd.read_csv, writer=pd.DataFrame.to_csv),
    dict(name='hdf-fixed', reader=pd.read_hdf, writer=hdf_writer('fixed')),
    dict(name='hdf-table', reader=pd.read_hdf, writer=hdf_writer('table')),
    dict(name='hdf-fixed-zlib', reader=pd.read_hdf, writer=hdf_writer('fixed', 'zlib')),
    dict(name='hdf-table-zlib', reader=pd.read_hdf, writer=hdf_writer('table', 'zlib')),
    dict(name='pq', reader=pd.read_parquet, writer=parquet_writer()),
    dict(name='pq-snappy', reader=pd.read_parquet, writer=parquet_writer('snappy')),
    dict(name='pq-gzip', reader=pd.read_parquet, writer=parquet_writer('gzip')),
]
def test_format(format, df=None):
    name, read, write = format['name'], format['reader'], format['writer']
    filename = f'temp.{name}'
    try:
        os.remove(filename)
    except FileNotFoundError:
        pass
    w_start = time.time()
    write(df, filename)
    w_end = r_start = time.time()
    df2 = read(filename)
    r_end = time.time()
    size = os.stat(filename).st_size
    result = {
        'format': name, 
        'rows': len(df), 
        'read_s': r_end - r_start,
        'write_s': w_end - w_start,
        'size_B': float(size),
        'size_kB': float(size / 2**10),
        'size_MB': float(size / 2**20)
    }
    return result<jupyter_output><empty_output><jupyter_text>Test on a small dataframe, incl. CSV<jupyter_code>test_df = mock_data(1000000)
results = pd.DataFrame.from_records(
    test_format(fmt, df=test_df) for fmt in formats
)
results = results.sort_values(by='read_s', ascending=False)
results
data = [
    go.Bar(
        x=results['format'],
        y=results['write_s'],
        name='write time'
    ),
    go.Bar(
        x=results['format'],
        y=results['read_s'],
        name='read time'   
    ),
]

layout = go.Layout(
    title=f'Read/Write times for {len(test_df)} rows',
    yaxis = dict(
        title='time / s',
        hoverformat='.1f'
    ),
    template='plotly_dark'
)

fig = go.Figure(data=data, layout=layout)

py.iplot(fig)
plotly.io.write_json(fig, 'read_write_bar.json')
data = [
    go.Scatter(
        x=[result['read_s']],
        y=[result['size_MB']],
        text=[result['format']],
        name=result['format'],
        mode='markers',
        hoverinfo='text+x+y'
    ) for _, result in results.iterrows()
]

layout = go.Layout(
    title=f'Read times and file sizes for {len(test_df)} rows',
    xaxis = dict(
        title='read time / s',
        hoverformat='.1f'
    ),
    yaxis = dict(
        title='size / MB',
        hoverformat='.1f'
    ),
    hovermode='closest',
    template='plotly_dark'
)

fig = go.Figure(data=data, layout=layout)

py.iplot(fig)
plotly.io.write_json(fig, 'read_size.json')<jupyter_output><empty_output><jupyter_text>Test on a larger dataframe, excl. CSV<jupyter_code>test_df = mock_data(int(10e6))
results = pd.DataFrame.from_records(
    test_format(fmt, df=test_df) for fmt in formats[1:]
)
results = results.sort_values(by='read_s', ascending=False)
results
data = [
    go.Bar(
        x=results['format'],
        y=results['write_s'],
        name='write time'
    ),
    go.Bar(
        x=results['format'],
        y=results['read_s'],
        name='read time'   
    ),
]

layout = go.Layout(
    title=f'Read/Write times for {len(test_df)} rows',
    yaxis = dict(
        title='time / s',
        hoverformat='.1f'
    ),
    template='plotly_dark'
)

fig = go.Figure(data=data, layout=layout)

py.iplot(fig)
plotly.io.write_json(fig, 'large_read_write_bar.json')
data = [
    go.Scatter(
        x=[result['read_s']],
        y=[result['size_MB']],
        text=[result['format']],
        name=result['format'],
        mode='markers',
        hoverinfo='text+x+y'
    ) for _, result in results.iterrows()
]

layout = go.Layout(
    title=f'Read times and file sizes for {len(test_df)} rows',
    xaxis = dict(
        title='read time / s',
        hoverformat='.1f',
        range=[0, 1.5]
    ),
    yaxis = dict(
        title='size / MB',
        hoverformat='.1f'
    ),
    hovermode='closest',
    template='plotly_dark'
)

fig = go.Figure(data=data, layout=layout)

py.iplot(fig)
plotly.io.write_json(fig, 'large_read_size.json')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16, 'fraction_numerical': 0.00568421052631579, 'mean_word_length': 3.009282700421941, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '386251', 'license': 'MIT', 'n_tokens_mistral': 1723, 'n_tokens_neox': 1697, 'n_words': 285, 'path': 'notebooks/data-file-format-comparison.ipynb', 'repo_name': 'ig248/tabelio', 'size': '4750'}