import pyarrow.dataset as ds

dataset = ds.dataset('./data', format='parquet', partitioning=['dt'])
print('n_rows:', dataset.count_rows())

for batch in dataset.to_batches():
    # Pandas
    df = batch.to_pandas()
    
    # Pandas는 모든 데이터를 불러오는데, batch 에서 특정 column 만 불러올 수 있습니다.
    for i in range(batch.num_rows):
        col0 = batch.column('col_0')[0].as_py()

        

from pyarrow.parquet import ParquetDataset

dataset = ParquetDataset('./data',
                         memory_map=True,                         
                         use_legacy_dataset=False)
df = pd.read_parquet('./data')

file_rows = [frag.count_rows() for frag in dataset.fragments]

print('Pandas shape :', df.shape)
print('Pandas  size :', sys.getsizeof(df))
print('Pyarrow size :', sys.getsizeof(dataset))
print('files        :', dataset.files)
print('fragments    :', dataset.fragments)
print('files rows   :', file_rows)
print('column size  :', len(dataset.schema))

import pyarrow.parquet as pq

dataset = pq.ParquetDataset('./data', use_legacy_dataset=False)
total_rows = 0
for frag in dataset.fragments:
    for i, row_group in enumerate(frag.row_groups):
        total_rows += row_group.num_rows

print(total_rows)  # 전체 rows 갯수.

import pyarrow.parquet as pq

dataset = pq.ParquetDataset('./data', use_legacy_dataset=False)

for frag in dataset.fragments:
    for batch in frag.to_batches():
        df = batch.to_pandas()
        row = batch.take(pa.array([0]))  # 특정 row 를 꺼낼 수 있음
        
        print('num rows    :', batch.num_rows)
        print('Pandas shape:', df.shape)
        print(row['gender'])

import pyarrow.parquet as pq

dataset = pq.ParquetDataset('./data', use_legacy_dataset=False)

for frag in dataset.fragments:
    for batch in frag.to_batches():
        for i in range(batch.num_rows):
            user_id = batch.column('user_id')[i].as_py()
            name = batch.column('name')[i].as_py()
            age = batch.column('age')[i].as_py()
            print(user_id, name, age)

import random
from bisect import bisect_right
from typing import List, Iterator, Optional, Tuple

import pandas as pd
from pyarrow.dataset import ParquetFileFragment
from pyarrow.lib import RecordBatch
from pyarrow.parquet import ParquetDataset
from torch.utils.data import Dataset


class PyArrowDataset(Dataset):
    """
    Restriction
     - Don't shuffle in Dataloader. this is for efficiency to precess large dataset.
       If you need to shuffle, do it before this custom dataset. (like in SparkSQL)
       But the algorithm supports random access.
    """

    def __init__(self, source: str, seed: int = 123):
        self.source = source
        self.seed = seed

        # Pyarrow
        self.dataset = ParquetDataset(source, use_legacy_dataset=False)
        self.fragments: List[ParquetFileFragment] = self.dataset.fragments
        self._batches: Optional[Iterator[RecordBatch]] = None
        self._batch: Optional[RecordBatch] = None

        # Indexing meta information to make search faster
        self._cumulative_n_rows: List[int] = []
        self._batch_idx: int = 0

        # Index
        self._fragment_idx = 0

        # Initialization
        self._init()

    def _init(self):
        random.seed(self.seed)

        self._cumulative_n_rows = [frag.count_rows() for frag in self.fragments]
        for i in range(1, len(self._cumulative_n_rows)):
            self._cumulative_n_rows[i] += self._cumulative_n_rows[i - 1]

    def _get_next(self, idx: int) -> Tuple[RecordBatch, int, int]:
        def get_prev_cum_frag_size(_fragment_idx):
            if _fragment_idx >= 1:
                return self._cumulative_n_rows[_fragment_idx - 1]
            return 0

        # Calculate fragment idx
        fragment_idx = self._fragment_idx
        fragment_changed = False
        _prev_size = get_prev_cum_frag_size(fragment_idx)
        _cur_size = self._cumulative_n_rows[self._fragment_idx]
        if (idx < _prev_size) or (idx >= _cur_size):
            fragment_idx = bisect_right(self._cumulative_n_rows, idx)
            assert fragment_idx < len(self.fragments)
            # fragment_idx %= len(self.fragments)
            fragment_changed = self._fragment_idx != fragment_idx
            self._fragment_idx = fragment_idx

            self._batch_idx = 0
            del self._batches
            del self._batch
            self._batches = None
            self._batch = None

        # Calculate batch idx
        _prev_size = get_prev_cum_frag_size(fragment_idx)
        batch_idx = idx - _prev_size
        batch_changed = batch_idx < self._batch_idx

        # Calculate batches of the fragment
        if self._batches is None or fragment_changed or batch_changed:
            if self._batches:
                self._batches.clear()

            self.batches = self.fragments[fragment_idx].to_batches()
            self._batch = None

        if self._batch is None:
            self._batch = next(self.batches)
            self._batch_idx = 0

        while True:
            if self._batch_idx <= batch_idx < self._batch_idx + self._batch.num_rows:
                break

            self._batch_idx += self._batch.num_rows
            self._batch = next(self.batches)

        return self._batch, fragment_idx, batch_idx - self._batch_idx

    def __del__(self):
        del self.dataset
        del self.fragments
        del self._batches
        del self._batch

    def __len__(self):
        return self._cumulative_n_rows[-1]


    def __getitem__(self, idx):
        batch, fragment_idx, batch_idx = self._get_next(idx)
        return batch.column('idx')[batch_idx].as_py()

dataset = PyArrowDataset("./data")
print(dataset[50000])
print(dataset[5000])
print(dataset[1])
print(dataset[500000])
print(dataset[151547])
print(dataset[127])
print(dataset[878])