File size: 8,405 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 |
# 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
|