Define a Batch Size For ParquetDataFrameLoader

lidorb: <p>Hey,</p>
<p>I’m trying to create a DataPipe of .parquet files using ParquetDataFrameLoader.<br>
My train data consists from 8 parquet files with 16 columns and more than 100k rows each.</p>
<p>The code I wrote based on the documentation:</p>
<pre><code class="lang-auto">source_dp = FileLister("/path/to/parquet/files/", masks="*.parquet")
parquet_df_dp = source_dp.load_parquet_as_df()
</code></pre>
<p>My problem is when I try to iterate through the parquet_df_dp.<br>
When I tried to create a batch of 4 rows I got a list of 4 elements where each element is a chunk of 10,000 rows.<br>
<code>batch = next(iter(parquet_df_dp.batch(4)))</code></p>
<p>Unbatching also didn’t work - <code>next(iter(parquet_df_dp.flatten().unbatch()))</code> leads to one value of each columns instead of the entire row.</p>
<p>How can I define my desired batch size using the datapipe?<br>
Thanks in advance.</p>

nivek: <p>What is the output of <code>parquet_df_dp.flatten()</code>? Is it 40,000 rows? If so, you can do <code>parquet_df_dp.flatten().batch(4)</code>.</p>

Wei-Cheng_Chang: <p><a class="mention" href="/u/lidorb">@lidorb</a> , I encounter similar issue. Do you find a way to do batching for <code>load_parquet_as_df()</code> ?</p>

nivek: <aside class="quote no-group" data-username="nivek" data-post="2" data-topic="177432">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nivek/48/39672_2.png" class="avatar"> nivek:</div>
<blockquote>
<p><code>parquet_df_dp.flatten().batch(4)</code></p>
</blockquote>
</aside>
<p><a class="mention" href="/u/wei-cheng_chang">@Wei-Cheng_Chang</a> Have you tried this?</p>

Wei-Cheng_Chang: <p>for my case, <code>load_parquet_as_df()</code> will output a torcharrow object of shape N rows and 3 columns.</p>
<p>After <code>load_parquet_as_df.flatten()</code> will output a list of Nx3 elements (which is not desirable).</p>
<p>For me, I realize the solution is</p>
<p><code>load_parquet_as_df.flatmap().batch()</code>, which output a Bx3 object</p>

nivek: <p>Yes, calling <code>.batch()</code> is necessary if you want batches.</p>

lidorb: <p>Using <code>load_parquet_as_df.flatmap().batch(batch_size)</code> worked great for me.<br>
Thanks <a class="mention" href="/u/wei-cheng_chang">@Wei-Cheng_Chang</a>!</p>

Wei-Cheng_Chang: <p><a class="mention" href="/u/lidorb">@lidorb</a> ,</p>
<p>I realized that using <code>load_parquet_as_df.flatmap().batch(batch_size)</code> is not a very efficient implementation for loading large amount of parquet files. I wrote a customized <code>IterableDataPipe</code> to achieve faster data loading time. An minimal example will be like</p>
<pre><code class="lang-auto">import pyarrow.parquet as pq

@functional_datapipe("load_parquet_batch")
class MyParquetLoaderIterDataPipe(IterDataPipe):
    def __init__(self, source_dp, columns, batch_size):
        self.source_dp = source_dp
        self.columns = columns
        self.batch_size = batch_size
    def __iter__(self):
        for path in self.source_dp:
            parquet_file = pq.ParquetFile(path)
            for batch in parquet_file.iter_batches(batch_size=self.batch_size, columns=self.columns):
                yield batch.to_pydict()
</code></pre>
<p>The I created the following datapipe/dataloader2 for pytorch DDP :</p>
<pre><code class="lang-auto">dp = IterableWrapper(FILE_LIST).shuffle().sharding_filter()
dp = dp.load_parquet_batch(columns=COLUMN_LIST, batch_size=BSZ_PER_GPU)
dp = dp.map(transform_fn)
dist_rs = DistributedReadingService()
dl = DataLoader2(dp, reading_service=dist_rs)
</code></pre>
<p>Hope that helps in your case.</p>

PopHSKE: <aside class="quote no-group" data-username="Wei-Cheng_Chang" data-post="8" data-topic="177432">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/wei-cheng_chang/48/16253_2.png" class="avatar"> Wei-Cheng_Chang:</div>
<blockquote>
<pre><code class="lang-auto">dp = IterableWrapper(FILE_LIST).shuffle().sharding_filter()
dp = dp.load_parquet_batch(columns=COLUMN_LIST, batch_size=BSZ_PER_GPU)
dp = dp.map(transform_fn)
dist_rs = DistributedReadingService()
dl = DataLoader2(dp, reading_service=dist_rs)
</code></pre>
</blockquote>
</aside>
<p><a class="mention" href="/u/wei-cheng_chang">@Wei-Cheng_Chang</a> I get<br>
AttributeError: ‘ShardingFilterIterDataPipe’ object has no attribute 'load_parquet_batch</p>