mirrors/arrow
0
0 Fork 0
mirror of https://github.com/apache/arrow.git synced 2026-10-29 09:05:47 +08:00
Code Activity

ARROW-2568: [Python] WIP

Browse Source
This commit is contained in:
Antoine Pitrou committed 2018-05-23 18:05:43 +02:00
1 parent fc7a382da4
commit 3417325cc2
11 files changed
+98 -62
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cpp/src/arrow/python/arrow_to_pandas.cc
+21 -18
View File
@@ -40,7 +40,7 @@
#include "arrow/util/decimal.h"
#include "arrow/util/logging.h"
#include "arrow/util/macros.h"
#include "arrow/util/parallel.h"
#include "arrow/util/thread-pool.h"
#include "arrow/visitor_inline.h"
#include "arrow/compute/api.h"
@@ -1370,14 +1370,14 @@ class DataFrameBlockCreator {
const std::shared_ptr<Table>& table, MemoryPool* pool)
: table_(table), options_(options), pool_(pool) {}
Status Convert(int nthreads, PyObject** output) {
Status Convert(bool use_threads, PyObject** output) {
column_types_.resize(table_->num_columns());
column_block_placement_.resize(table_->num_columns());
type_counts_.clear();
blocks_.clear();
RETURN_NOT_OK(CreateBlocks());
RETURN_NOT_OK(WriteTableToBlocks(nthreads));
RETURN_NOT_OK(WriteTableToBlocks(use_threads));
return GetResultList(output);
}
@@ -1450,23 +1450,26 @@ class DataFrameBlockCreator {
return Status::OK();
}
Status WriteTableToBlocks(int nthreads) {
Status WriteTableToBlocks(bool use_threads) {
auto WriteColumn = [this](int i) {
std::shared_ptr<PandasBlock> block;
RETURN_NOT_OK(this->GetBlock(i, &block));
return block->Write(this->table_->column(i), i, this->column_block_placement_[i]);
};
int num_tasks = table_->num_columns();
nthreads = std::min<int>(nthreads, num_tasks);
if (nthreads == 1) {
for (int i = 0; i < num_tasks; ++i) {
RETURN_NOT_OK(WriteColumn(i));
}
} else {
RETURN_NOT_OK(ParallelFor(nthreads, num_tasks, WriteColumn));
arrow::internal::ThreadPool* pool = arrow::internal::CPUThreadPool();
std::vector<std::future<Status>> futures;
for (int i = 0; i < table_->num_columns(); ++i) {
futures.push_back(pool->Submit(WriteColumn, i));
}
return Status::OK();
auto final_status = Status::OK();
for (auto& fut : futures) {
Status st = fut.get();
if (!st.ok()) {
final_status = st;
}
}
return final_status;
}
Status AppendBlocks(const BlockMap& blocks, PyObject* list) {
@@ -1866,14 +1869,14 @@ Status ConvertColumnToPandas(PandasOptions options, const std::shared_ptr<Column
}
Status ConvertTableToPandas(PandasOptions options, const std::shared_ptr<Table>& table,
int nthreads, MemoryPool* pool, PyObject** out) {
return ConvertTableToPandas(options, std::unordered_set<std::string>(), table, nthreads,
pool, out);
bool use_threads, MemoryPool* pool, PyObject** out) {
return ConvertTableToPandas(options, std::unordered_set<std::string>(), table,
use_threads, pool, out);
}
Status ConvertTableToPandas(PandasOptions options,
const std::unordered_set<std::string>& categorical_columns,
const std::shared_ptr<Table>& table, int nthreads,
const std::shared_ptr<Table>& table, bool use_threads,
MemoryPool* pool, PyObject** out) {
std::shared_ptr<Table> current_table = table;
if (!categorical_columns.empty()) {
@@ -1894,7 +1897,7 @@ Status ConvertTableToPandas(PandasOptions options,
}
DataFrameBlockCreator helper(options, current_table, pool);
return helper.Convert(nthreads, out);
return helper.Convert(use_threads, out);
}
} // namespace py
cpp/src/arrow/python/arrow_to_pandas.h
+2 -2
View File
@@ -68,7 +68,7 @@ Status ConvertColumnToPandas(PandasOptions options, const std::shared_ptr<Column
// tuple item: (indices: ndarray[int32], block: ndarray[TYPE, ndim=2])
ARROW_EXPORT
Status ConvertTableToPandas(PandasOptions options, const std::shared_ptr<Table>& table,
int nthreads, MemoryPool* pool, PyObject** out);
bool use_threads, MemoryPool* pool, PyObject** out);
/// Convert a whole table as efficiently as possible to a pandas.DataFrame.
///
@@ -77,7 +77,7 @@ Status ConvertTableToPandas(PandasOptions options, const std::shared_ptr<Table>&
ARROW_EXPORT
Status ConvertTableToPandas(PandasOptions options,
const std::unordered_set<std::string>& categorical_columns,
const std::shared_ptr<Table>& table, int nthreads,
const std::shared_ptr<Table>& table, bool use_threads,
MemoryPool* pool, PyObject** out);
} // namespace py
cpp/src/arrow/util/parallel.h
+6
View File
@@ -27,6 +27,12 @@
namespace arrow {
// A parallelizer that takes a `Status(int)` function and calls it with
// arguments between 0 and `num_tasks - 1`, using the given number of
// concurrent threads.
// In most cases it's more appropriate to use the global CPU thread pool
// (see arrow/util/thread-pool.h)
template <class FUNCTION>
Status ParallelFor(int nthreads, int num_tasks, FUNCTION&& func) {
std::vector<std::thread> thread_pool;
python/pyarrow/feather.py
+22 -8
View File
@@ -66,8 +66,13 @@ class FeatherReader(ext.FeatherReader):
table = Table.from_arrays(columns, names=names)
return table
def read_pandas(self, columns=None, nthreads=1):
return self.read_table(columns=columns).to_pandas(nthreads=nthreads)
def read_pandas(self, columns=None, nthreads=None, use_threads=False):
if nthreads is not None:
warnings.warn("`nthreads` argument is ignored, "
"pass `use_threads` instead", DeprecationWarning,
stacklevel=2)
return self.read_table(columns=columns).to_pandas(
use_threads=use_threads)
class FeatherWriter(object):
@@ -141,7 +146,7 @@ class FeatherDataset(object):
.format(piece, self.schema,
table.schema))
def read_pandas(self, columns=None, nthreads=1):
def read_pandas(self, columns=None, nthreads=None, use_threads=False):
"""
Read multiple Parquet files as a single pandas DataFrame
@@ -157,7 +162,12 @@ class FeatherDataset(object):
pandas.DataFrame
Content of the file as a pandas DataFrame (of columns)
"""
return self.read_table(columns=columns).to_pandas(nthreads=nthreads)
if nthreads is not None:
warnings.warn("`nthreads` argument is ignored, "
"pass `use_threads` instead", DeprecationWarning,
stacklevel=2)
return self.read_table(columns=columns).to_pandas(
use_threads=use_threads)
def write_feather(df, dest):
@@ -186,7 +196,7 @@ def write_feather(df, dest):
raise
def read_feather(source, columns=None, nthreads=1):
def read_feather(source, columns=None, nthreads=None, use_threads=False):
"""
Read a pandas.DataFrame from Feather format
@@ -196,15 +206,19 @@ def read_feather(source, columns=None, nthreads=1):
columns : sequence, optional
Only read a specific set of columns. If not provided, all columns are
read
nthreads : int, default 1
Number of CPU threads to use when reading to pandas.DataFrame
use_threads: bool, default False
Whether to parallelize reading using multiple threads
Returns
-------
df : pandas.DataFrame
"""
if nthreads is not None:
warnings.warn("`nthreads` argument is ignored, "
"pass `use_threads` instead", DeprecationWarning,
stacklevel=2)
reader = FeatherReader(source)
return reader.read_pandas(columns=columns, nthreads=nthreads)
return reader.read_pandas(columns=columns, use_threads=use_threads)
def read_table(source, columns=None):
python/pyarrow/includes/libarrow.pxd
+1 -1
View File
@@ -959,7 +959,7 @@ cdef extern from "arrow/python/api.h" namespace "arrow::py" nogil:
PandasOptions options,
const unordered_set[c_string]& categorical_columns,
const shared_ptr[CTable]& table,
int nthreads, CMemoryPool* pool,
c_bool use_threads, CMemoryPool* pool,
PyObject** out)
void c_set_default_memory_pool \
python/pyarrow/ipc.py
+11 -5
View File
@@ -17,6 +17,8 @@
# Arrow file and stream reader/writer classes, and other messaging tools
import warnings
import pyarrow as pa
from pyarrow.lib import (Message, MessageReader, # noqa
@@ -168,22 +170,26 @@ def serialize_pandas(df, nthreads=None, preserve_index=True):
return sink.get_result()
def deserialize_pandas(buf, nthreads=None):
def deserialize_pandas(buf, nthreads=None, use_threads=False):
"""Deserialize a buffer protocol compatible object into a pandas DataFrame.
Parameters
----------
buf : buffer
An object compatible with the buffer protocol
nthreads : int, defualt None
The number of threads to use to convert the buffer to a DataFrame,
default all CPUs
use_threads: boolean, default False
Whether to parallelize the conversion using multiple threads
Returns
-------
df : pandas.DataFrame
"""
if nthreads is not None:
warnings.warn("`nthreads` argument is ignored, "
"pass `use_threads` instead", DeprecationWarning,
stacklevel=2)
buffer_reader = pa.BufferReader(buf)
reader = pa.RecordBatchStreamReader(buffer_reader)
table = reader.read_all()
return table.to_pandas(nthreads=nthreads)
return table.to_pandas(use_threads=use_threads)
python/pyarrow/pandas_compat.py
+5 -4
View File
@@ -493,7 +493,7 @@ def _make_datetimetz(tz):
# Converting pyarrow.Table efficiently to pandas.DataFrame
def table_to_blockmanager(options, table, memory_pool, nthreads=1,
def table_to_blockmanager(options, table, memory_pool, use_threads,
categories=None):
from pyarrow.compat import DatetimeTZDtype
@@ -568,7 +568,7 @@ def table_to_blockmanager(options, table, memory_pool, nthreads=1,
block_table.schema.get_field_index(raw_name)
)
blocks = _table_to_blocks(options, block_table, nthreads, memory_pool,
blocks = _table_to_blocks(options, block_table, use_threads, memory_pool,
categories)
# Construct the row index
@@ -728,11 +728,12 @@ def _reconstruct_columns_from_metadata(columns, column_indexes):
return pd.MultiIndex(levels=new_levels, labels=labels, names=columns.names)
def _table_to_blocks(options, block_table, nthreads, memory_pool, categories):
def _table_to_blocks(options, block_table, use_threads, memory_pool,
categories):
# Part of table_to_blockmanager
# Convert an arrow table to Block from the internal pandas API
result = pa.lib.table_to_blocks(options, block_table, nthreads,
result = pa.lib.table_to_blocks(options, block_table, use_threads,
memory_pool, categories)
# Defined above
python/pyarrow/table.pxi
+18 -12
View File
@@ -16,6 +16,7 @@
# under the License.
import json
import warnings
from collections import OrderedDict
@@ -746,7 +747,7 @@ cdef class RecordBatch:
entries.append((name, column))
return OrderedDict(entries)
def to_pandas(self, nthreads=None):
def to_pandas(self, nthreads=None, use_threads=False):
"""
Convert the arrow::RecordBatch to a pandas DataFrame
@@ -754,7 +755,11 @@ cdef class RecordBatch:
-------
pandas.DataFrame
"""
return Table.from_batches([self]).to_pandas(nthreads=nthreads)
if nthreads is not None:
warnings.warn("`nthreads` argument is ignored, "
"pass `use_threads` instead", DeprecationWarning,
stacklevel=2)
return Table.from_batches([self]).to_pandas(use_threads=use_threads)
@classmethod
def from_pandas(cls, df, Schema schema=None, bint preserve_index=True,
@@ -825,7 +830,7 @@ cdef class RecordBatch:
CRecordBatch.Make(schema, num_rows, c_arrays))
def table_to_blocks(PandasOptions options, Table table, int nthreads,
def table_to_blocks(PandasOptions options, Table table, c_bool use_threads,
MemoryPool memory_pool, categories):
cdef:
PyObject* result_obj
@@ -840,7 +845,7 @@ def table_to_blocks(PandasOptions options, Table table, int nthreads,
with nogil:
check_status(
libarrow.ConvertTableToPandas(
options, categorical_columns, c_table, nthreads, pool,
options, categorical_columns, c_table, use_threads, pool,
&result_obj
)
)
@@ -1148,16 +1153,12 @@ cdef class Table:
def to_pandas(self, nthreads=None, strings_to_categorical=False,
memory_pool=None, zero_copy_only=False, categories=None,
integer_object_nulls=False):
integer_object_nulls=False, use_threads=False):
"""
Convert the arrow::Table to a pandas DataFrame
Parameters
----------
nthreads : int, default max(1, multiprocessing.cpu_count() / 2)
For the default, we divide the CPU count by 2 because most modern
computers have hyperthreading turned on, so doubling the CPU count
beyond the number of physical cores does not help
strings_to_categorical : boolean, default False
Encode string (UTF8) and binary types to pandas.Categorical
memory_pool: MemoryPool, optional
@@ -1169,6 +1170,8 @@ cdef class Table:
List of columns that should be returned as pandas.Categorical
integer_object_nulls : boolean, default False
Cast integers with nulls to objects
use_threads: boolean, default False
Whether to parallelize the conversion using multiple threads
Returns
-------
@@ -1177,15 +1180,18 @@ cdef class Table:
cdef:
PandasOptions options
if nthreads is not None:
warnings.warn("`nthreads` argument is ignored, "
"pass `use_threads` instead", DeprecationWarning,
stacklevel=2)
options = PandasOptions(
strings_to_categorical=strings_to_categorical,
zero_copy_only=zero_copy_only,
integer_object_nulls=integer_object_nulls)
self._check_nullptr()
if nthreads is None:
nthreads = cpu_count()
mgr = pdcompat.table_to_blockmanager(options, self, memory_pool,
nthreads, categories)
use_threads, categories)
return pd.DataFrame(mgr)
def to_pydict(self):
python/pyarrow/tests/test_convert_pandas.py
+5 -5
View File
@@ -60,7 +60,7 @@ def _alltypes_example(size=100):
})
def _check_pandas_roundtrip(df, expected=None, nthreads=1,
def _check_pandas_roundtrip(df, expected=None, use_threads=False,
expected_schema=None,
check_dtype=True, schema=None,
preserve_index=False,
@@ -68,9 +68,9 @@ def _check_pandas_roundtrip(df, expected=None, nthreads=1,
klass = pa.RecordBatch if as_batch else pa.Table
table = klass.from_pandas(df, schema=schema,
preserve_index=preserve_index,
nthreads=nthreads)
nthreads=2 if use_threads else 1)
result = table.to_pandas(nthreads=nthreads)
result = table.to_pandas(use_threads=use_threads)
if expected_schema:
assert table.schema.equals(expected_schema)
if expected is None:
@@ -1836,8 +1836,8 @@ class TestConvertMisc(object):
def test_threaded_conversion(self):
df = _alltypes_example()
_check_pandas_roundtrip(df, nthreads=2)
_check_pandas_roundtrip(df, nthreads=2, as_batch=True)
_check_pandas_roundtrip(df, use_threads=True)
_check_pandas_roundtrip(df, use_threads=True, as_batch=True)
def test_category(self):
repeats = 5
python/pyarrow/tests/test_feather.py
+3 -3
View File
@@ -65,7 +65,7 @@ class TestFeatherReader(unittest.TestCase):
def _check_pandas_roundtrip(self, df, expected=None, path=None,
columns=None, null_counts=None,
nthreads=1):
use_threads=False):
if path is None:
path = random_path()
@@ -74,7 +74,7 @@ class TestFeatherReader(unittest.TestCase):
if not os.path.exists(path):
raise Exception('file not written')
result = read_feather(path, columns, nthreads=nthreads)
result = read_feather(path, columns, use_threads=use_threads)
if expected is None:
expected = df
@@ -384,7 +384,7 @@ class TestFeatherReader(unittest.TestCase):
data = {'c{0}'.format(i): [''] * 10
for i in range(100)}
df = pd.DataFrame(data)
self._check_pandas_roundtrip(df, nthreads=4)
self._check_pandas_roundtrip(df, use_threads=True)
def test_nan_as_null(self):
# Create a nan that is not numpy.nan
python/pyarrow/tests/test_ipc.py
+4 -4
View File
@@ -406,9 +406,9 @@ def test_get_record_batch_size():
assert pa.get_record_batch_size(batch) > (N * itemsize)
def _check_serialize_pandas_round_trip(df, nthreads=1):
buf = pa.serialize_pandas(df, nthreads=nthreads)
result = pa.deserialize_pandas(buf, nthreads=nthreads)
def _check_serialize_pandas_round_trip(df, use_threads=False):
buf = pa.serialize_pandas(df, nthreads=2 if use_threads else 1)
result = pa.deserialize_pandas(buf, use_threads=use_threads)
assert_frame_equal(result, df)
@@ -429,7 +429,7 @@ def test_pandas_serialize_round_trip_nthreads():
{'foo': [1.5, 1.6, 1.7], 'bar': list('abc')},
index=index, columns=columns
)
_check_serialize_pandas_round_trip(df, nthreads=2)
_check_serialize_pandas_round_trip(df, use_threads=True)
def test_pandas_serialize_round_trip_multi_index():