Method get_chunks

python/pyarrow/interchange/dataframe.py:176–217

Return an iterator yielding the chunks. By default (None), yields the chunks that the data is stored as by the producer. If given, ``n_chunks`` must be a multiple of ``self.num_chunks()``, meaning the producer must subdivide each chunk before yielding it.

(
        self, n_chunks: Optional[int] = None
    )

        )

    def get_chunks(
        self, n_chunks: Optional[int] = None
    ) -> Iterable[_PyArrowDataFrame]:
        """
        Return an iterator yielding the chunks.

        By default (None), yields the chunks that the data is stored as by the
        producer. If given, ``n_chunks`` must be a multiple of
        ``self.num_chunks()``, meaning the producer must subdivide each chunk
        before yielding it.

        Note that the producer must ensure that all columns are chunked the
        same way.
        """
        # Subdivide chunks
        if n_chunks and n_chunks > 1:
            chunk_size = self.num_rows() // n_chunks
            if self.num_rows() % n_chunks != 0:
                chunk_size += 1
            if isinstance(self._df, pa.Table):
                batches = self._df.to_batches(max_chunksize=chunk_size)
            else:
                batches = []
                for start in range(0, chunk_size * n_chunks, chunk_size):
                    batches.append(self._df.slice(start, chunk_size))
            # In case when the size of the chunk is such that the resulting
            # list is one less chunk then n_chunks -> append an empty chunk
            if len(batches) == n_chunks - 1:
                batches.append(pa.record_batch([[]], schema=self._df.schema))
        # yields the chunks that the data is stored as
        else:
            if isinstance(self._df, pa.Table):
                batches = self._df.to_batches()
            else:
                batches = [self._df]

        # Create an iterator of RecordBatches
        iterator = [_PyArrowDataFrame(batch,
                                      self._nan_as_null,
                                      self._allow_copy)
                    for batch in batches]
        return iterator