# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations from typing import ( Any, Iterable, Optional, Sequence, ) import pyarrow as pa from pyarrow.interchange.column import _PyArrowColumn class _PyArrowDataFrame: """ A data frame class, with only the methods required by the interchange protocol defined. A "data frame" represents an ordered collection of named columns. A column's "name" must be a unique string. Columns may be accessed by name or by position. This could be a public data frame class, or an object with the methods and attributes defined on this DataFrame class could be returned from the ``__dataframe__`` method of a public data frame class in a library adhering to the dataframe interchange protocol specification. """ def __init__( self, df: pa.Table | pa.RecordBatch, nan_as_null: bool = False, allow_copy: bool = True ) -> None: """ Constructor - an instance of this (private) class is returned from `pa.Table.__dataframe__` or `pa.RecordBatch.__dataframe__`. """ self._df = df # ``nan_as_null`` is a keyword intended for the consumer to tell the # producer to overwrite null values in the data with ``NaN`` (or # ``NaT``). if nan_as_null is True: raise RuntimeError( "nan_as_null=True currently has no effect, " "use the default nan_as_null=False" ) self._nan_as_null = nan_as_null self._allow_copy = allow_copy def __dataframe__( self, nan_as_null: bool = False, allow_copy: bool = True ) -> _PyArrowDataFrame: """ Construct a new exchange object, potentially changing the parameters. ``nan_as_null`` is a keyword intended for the consumer to tell the producer to overwrite null values in the data with ``NaN``. It is intended for cases where the consumer does not support the bit mask or byte mask that is the producer's native representation. ``allow_copy`` is a keyword that defines whether or not the library is allowed to make a copy of the data. For example, copying data would be necessary if a library supports strided buffers, given that this protocol specifies contiguous buffers. """ return _PyArrowDataFrame(self._df, nan_as_null, allow_copy) @property def metadata(self) -> dict[str, Any]: """ The metadata for the data frame, as a dictionary with string keys. The contents of `metadata` may be anything, they are meant for a library to store information that it needs to, e.g., roundtrip losslessly or for two implementations to share data that is not (yet) part of the interchange protocol specification. For avoiding collisions with other entries, please add name the keys with the name of the library followed by a period and the desired name, e.g, ``pandas.indexcol``. """ # The metadata for the data frame, as a dictionary with string keys. # Add schema metadata here (pandas metadata or custom metadata) if self._df.schema.metadata: schema_metadata = {"pyarrow." + k.decode('utf8'): v.decode('utf8') for k, v in self._df.schema.metadata.items()} return schema_metadata else: return {} def num_columns(self) -> int: """ Return the number of columns in the DataFrame. """ return self._df.num_columns def num_rows(self) -> int: """ Return the number of rows in the DataFrame, if available. """ return self._df.num_rows def num_chunks(self) -> int: """ Return the number of chunks the DataFrame consists of. """ if isinstance(self._df, pa.RecordBatch): return 1 else: # pyarrow.Table can have columns with different number # of chunks so we take the number of chunks that # .to_batches() returns as it takes the min chunk size # of all the columns (to_batches is a zero copy method) batches = self._df.to_batches() return len(batches) def column_names(self) -> Iterable[str]: """ Return an iterator yielding the column names. """ return self._df.schema.names def get_column(self, i: int) -> _PyArrowColumn: """ Return the column at the indicated position. """ return _PyArrowColumn(self._df.column(i), allow_copy=self._allow_copy) def get_column_by_name(self, name: str) -> _PyArrowColumn: """ Return the column whose name is the indicated name. """ return _PyArrowColumn(self._df.column(name), allow_copy=self._allow_copy) def get_columns(self) -> Iterable[_PyArrowColumn]: """ Return an iterator yielding the columns. """ return [ _PyArrowColumn(col, allow_copy=self._allow_copy) for col in self._df.columns ] def select_columns(self, indices: Sequence[int]) -> _PyArrowDataFrame: """ Create a new DataFrame by selecting a subset of columns by index. """ return _PyArrowDataFrame( self._df.select(list(indices)), self._nan_as_null, self._allow_copy ) def select_columns_by_name( self, names: Sequence[str] ) -> _PyArrowDataFrame: """ Create a new DataFrame by selecting a subset of columns by name. """ return _PyArrowDataFrame( self._df.select(list(names)), self._nan_as_null, self._allow_copy ) 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