import cudf
import rmm
import nvtx

df = cudf.DataFrame({
    'a': ['aaaa'] * 1_000_000,
    'b': ['bbbb'] * 1_000_000,
    'c': ['cccc'] * 1_000_000,
})
df = cudf.concat([df] * 20, ignore_index=True)
file_path = '/raid/file.jsonl'
df.to_json(file_path, lines=True, orient='records', engine='cudf')


mr = rmm.mr.CudaMemoryResource()
mr = rmm.mr.StatisticsResourceAdaptor(mr)
rmm.mr.set_current_device_resource(mr)

with nvtx.annotate('single read'):
    df = cudf.read_json(file_path, lines=True)

df_size_mb = df.memory_usage(deep=True).sum() / 1e6
print(f"Dataframe size: {df_size_mb:.2f} MB")
peak_mb = mr.allocation_counts.peak_bytes / 1e6
print(f"Peak memory: {peak_mb:.2f} MB")

mr = rmm.mr.CudaMemoryResource()
mr = rmm.mr.StatisticsResourceAdaptor(mr)
rmm.mr.set_current_device_resource(mr)

with nvtx.annotate('byte_range read'):
    chunk_size = 100_000_000
    data = []
    x =0
    while True:
        try:
            d = cudf.read_json(
                file_path,        
                lines=True, 
                byte_range=(chunk_size * x, chunk_size)
            )
        except:
            # RuntimeError: CUDF failure at: /opt/conda/conda-bld/work/cpp/src/io/utilities/datasource.cpp:219: Offset is past end of file
            break
        data.append(d)    
        x+=1
    df = cudf.concat(data, ignore_index=True)

df_size_mb = df.memory_usage(deep=True).sum() / 1e6
print(f"Dataframe size: {df_size_mb:.2f} MB")
peak_mb = mr.allocation_counts.peak_bytes / 1e6
print(f"Peak memory: {peak_mb:.2f} MB")

Dataframe size: 480.00 MB
Peak memory: 7580.00 MB
Dataframe size: 480.00 MB
Peak memory: 1494.29 MB