[Python][Dataset][Parquet] Enable Pre-Buffering by default for Parquet s3 datasets

[Akshay-A-Kulkarni]

Describe the enhancement requested

Summary

I'm seeing a discrepancy in read times between parquet and dataset modules for reading (more notice-able on fairly large) partitioned parquet tables.

In a nutshell, pyarrow.dataset.to_table() seems to be considerably slower than pyarrow.parquet.read_table() in most cases I have tested so far. The difference in read times for some synthetic data (see Appendix) I tested on was approximately 12x.

Details

import sys
sys.version

>>> '3.8.17 (default, Jul 5 2023, 21:04:15) \n[GCC 11.2.0]'

import pyarrow as pa
pa.__version__

>>> '12.0.1'

The test was done on a machine similiar in specs to an m5.4xlarge (nCPU=16 | Memory=64GiB) EC2 Instance

path = 's3://BUCKET/PREFIX/PARTITIONED_DATA_BASE_DIR'

import pyarrow.parquet as pq
table_pq = pq.read_table(path, partitioning='hive')
print(table_pq.schema)
print(table_pq.shape)

a: string
b: int64
c: double
part_col: dictionary<values=string, indices=int32, ordered=0>
(100000000, 4)

CPU times: user 21.2 s, sys: 5.4 s, total: 26.6 s
Wall time: 10.6 s

path = 's3://BUCKET/PREFIX/PARTITIONED_DATA_BASE_DIR'

import pyarrow.dataset as ds

table_ds = ds.dataset(
 path, 
 partitioning='hive', 
 format='parquet'
).to_table()

print(table_ds.schema)
print(table_ds.shape)

a: string
b: int64
c: double
part_col: dictionary<values=string, indices=int32, ordered=0>
(100000000, 4)

CPU times: user 49.7 s, sys: 6.67 s, total: 56.4 s
Wall time: 1min 57s

From doing some digging, I found that the root cause of the issue seems to be pre_buffer, which is not enabled for pyarrow.dataset by default.

Once I add the pre_buffering via ParquetFileFormat & ParquetFragmentScanOptions the issue gets resolved.

%%time
table_ds_buffer = ds.dataset(
 PATH_TO_S3_PARQUET, 
 partitioning='hive', 
 format=ds.ParquetFileFormat(
 default_fragment_scan_options=ds.ParquetFragmentScanOptions(
 pre_buffer=True
 )
 )
).to_table()



print(table_ds_buffer.schema)
print(table_ds_buffer.shape)

a: string
b: int64
c: double
part_col: string
(100000000, 4)

CPU times: user 21 s, sys: 5.27 s, total: 26.2 s
Wall time: 9.4 s

The issue seems to be exacerbated if the parquet datasets with partitions (maybe more file parts in general ?). Doing the read test on a single parquet was still worse but not as worse as the partition one.

Do we know the reasoning for this? As far as I understand, for S3-based parquet datasets, this could be turned on by default. Otherwise maybe we can add it as an argument for the user to enable in the dataset() call?

Appendix

Code to generate test data

%%time
from string import ascii_letters
import pyarrow.dataset as ds
import pyarrow as pa

num_rows = 100_000_000
num_files = 1 # 32/64
num_partitions = 4

import pandas as pd
table = pa.Table.from_pydict({
 'a': [str(i) for i in range(num_rows)],
 'b': range(num_rows),
 'c': [float(i) for i in range(num_rows)],
 'part_col':np.random.choice(list(ascii_letters[:num_partitions]), size=num_rows,
 # p=[0.8,0.05,0.1,0.05] ## for creating row num skew between partitions
 )
})


ds.write_dataset(
 table,
 path,
 format='parquet',
 partitioning_flavor='hive',
 partitioning=['part_col'],
 max_rows_per_file = (num_rows//(num_files*num_partitions)), 
 existing_data_behavior='overwrite_or_ignore'
)

Component(s)

Parquet, Python

[mapleFU]

Nice catch! I guess to_table should prebuffer the data, let me take a look

[mapleFU]

to_table:

 return self.scanner(
 schema=schema,
 columns=columns,
 filter=filter,
 batch_size=batch_size,
 batch_readahead=batch_readahead,
 fragment_readahead=fragment_readahead,
 fragment_scan_options=fragment_scan_options,
 use_threads=use_threads,
 memory_pool=memory_pool
 ).to_table()

It call DatasetScanner::ToTable. By default, it has a prefetch depth(fragment_readahead and batch_readahead), but I think that is not fully used in scanner. And fragment_scan_options would be used in C++ reader's code

@westonpace would you mind take a look or give some advices?

[westonpace]

fragment_readahead and batch_readahead control how many files/row-groups to read at a time. pre_buffer controls how an individual row group is read. So these are separate properties. pre_buffer is probably always a good thing when reading from S3. However, when reading from local disk I think pre_buffer can sometimes lead to greater memory consumption. Is pre_buffer=True the default for read_table?

[mapleFU]

I guess

table_ds = ds.dataset(
 path, 
 partitioning='hive', 
 format='parquet'
).to_table()

print(table_ds.schema)
print(table_ds.shape)

Will not buffer when read from single Parquet dataset, so might cause performance lost.

def read_table(source, *, columns=None, use_threads=True, metadata=None,
 schema=None, use_pandas_metadata=False, read_dictionary=None,
 memory_map=False, buffer_size=0, partitioning="hive",
 filesystem=None, filters=None, use_legacy_dataset=False,
 ignore_prefixes=None, pre_buffer=True,
 coerce_int96_timestamp_unit=None,
 decryption_properties=None, thrift_string_size_limit=None,
 thrift_container_size_limit=None):

And this will default call pre_buffer

[mapleFU]

You can require pre_buffer first, or use larger buffer_size to do buffer read in Parquet.

• pre_buffer: might issue io request before do cpu read
• buffer_size: Apart from Local SSD, default S3 may prefer about 1MB read. I guess it could be better.

[westonpace]

I did some recent testing a few weeks ago and I also found that buffer_size was better when it came to memory usage. pre_buffer will fetch the entire column for the entire row group (or row groups) that you are reading.

However, if pre_buffer is the default for read_table then I think it is ok to make it the default for datasets.

[mapleFU]

@westonpace Since Acero Scan node would default prefetch some fragments, would it cause memory usage worse than previous read_table?

[westonpace]

Possibly. I think there are two concerns users generally have. Either they want "max speed" (use as little memory as possible but, if there is a speed/memory tradeoff, prefer speed) or they want "least memory" (if there is a speed/memory tradeoff, prefer memory)

I tried to simplify things in #35889 (this is why I was running experiments a few weeks ago) and came up with:

parquet::ReaderProperties

 case ParquetScanStrategy::kLeastMemory:
 properties.enable_buffered_stream();
 properties.set_buffer_size(8 * 1024 * 1024);
 case ParquetScanStrategy::kMaxSpeed:
 properties.disable_buffered_stream();

parquet::ArrowReaderProperties

 case ParquetScanStrategy::kLeastMemory:
 properties.set_batch_size(acero::ExecPlan::kMaxBatchSize);
 properties.set_pre_buffer(false);
 case ParquetScanStrategy::kMaxSpeed:
 properties.set_batch_size(64 * 1024 * 1024);
 properties.set_pre_buffer(true);
 properties.set_cache_options(io::CacheOptions::LazyDefaults());

I'm pretty sure that the kLeastMemory options use very low memory (even when scanning large files). I wasn't convinced that kMaxSpeed was much faster (but I only tested local disks and not S3).

I am very surprised to see a 10x difference due to pre-buffering. I don't think any of our experiments (S3 or not) ever showed a difference that was that drastic.

[westonpace]

Since Acero Scan node would default prefetch some fragments, would it cause memory usage worse than previous read_table?

Ah, I misunderstood this question. Yes, scan node / datasets will always use more memory than read_table because it is reading multiple files at once. Hopefully we can still do some buffering though.

[Akshay-A-Kulkarni]

fragment_readahead and batch_readahead control how many files/row-groups to read at a time. pre_buffer controls how an individual row group is read. So these are separate properties. pre_buffer is probably always a good thing when reading from S3. However, when reading from local disk I think pre_buffer can sometimes lead to greater memory consumption. Is pre_buffer=True the default for read_table?

@westonpace With the limited understanding that I have, if memory consumption on local fs is an issue, could we check the filesystem on dataset()call and if its S3 enable pre_buffering for parquet datasets then?