Skip to content

Commit c017a63

Browse files
pitrouxhochy
authored andcommitted
ARROW-1035: [Python] Add streaming dataframe reconstruction benchmark
Author: Antoine Pitrou <antoine@python.org> Closes apache#1665 from pitrou/ARROW-1035-streaming-benchmark and squashes the following commits: 32b1956 <Antoine Pitrou> ARROW-1035: Add streaming dataframe reconstruction benchmark
1 parent e0328b0 commit c017a63

2 files changed

Lines changed: 90 additions & 8 deletions

File tree

‎python/benchmarks/common.py‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,21 @@
2323
import numpy as np
2424

2525

26+
KILOBYTE = 1 << 10
27+
MEGABYTE = KILOBYTE * KILOBYTE
28+
29+
2630
def _multiplicate_sequence(base, target_size):
2731
q, r = divmod(target_size, len(base))
2832
return [base] * q + [base[:r]]
2933

3034

31-
def get_random_bytes(n):
32-
rnd = np.random.RandomState(42)
35+
def get_random_bytes(n, *, seed=42):
36+
"""
37+
Generate a random bytes object of size *n*.
38+
Note the result might be compressible.
39+
"""
40+
rnd = np.random.RandomState(seed)
3341
# Computing a huge random bytestring can be costly, so we get at most
3442
# 100KB and duplicate the result as needed
3543
base_size = 100003
@@ -43,22 +51,25 @@ def get_random_bytes(n):
4351
return result
4452

4553

46-
def get_random_ascii(n):
47-
arr = np.frombuffer(get_random_bytes(n), dtype=np.int8) & 0x7f
54+
def get_random_ascii(n, *, seed=42):
55+
"""
56+
Get a random ASCII-only unicode string of size *n*.
57+
"""
58+
arr = np.frombuffer(get_random_bytes(n, seed=seed), dtype=np.int8) & 0x7f
4859
result, _ = codecs.ascii_decode(arr)
4960
assert isinstance(result, str)
5061
assert len(result) == n
5162
return result
5263

5364

54-
def _random_unicode_letters(n):
65+
def _random_unicode_letters(n, *, seed=42):
5566
"""
5667
Generate a string of random unicode letters (slow).
5768
"""
5869
def _get_more_candidates():
5970
return rnd.randint(0, sys.maxunicode, size=n).tolist()
6071

61-
rnd = np.random.RandomState(42)
72+
rnd = np.random.RandomState(seed)
6273
out = []
6374
candidates = []
6475

@@ -75,8 +86,12 @@ def _get_more_candidates():
7586
_1024_random_unicode_letters = _random_unicode_letters(1024)
7687

7788

78-
def get_random_unicode(n):
79-
indices = np.frombuffer(get_random_bytes(n * 2), dtype=np.int16) & 1023
89+
def get_random_unicode(n, *, seed=42):
90+
"""
91+
Get a random non-ASCII unicode string of size *n*.
92+
"""
93+
indices = np.frombuffer(get_random_bytes(n * 2, seed=seed),
94+
dtype=np.int16) & 1023
8095
unicode_arr = np.array(_1024_random_unicode_letters)[indices]
8196

8297
result = ''.join(unicode_arr.tolist())

‎python/benchmarks/streaming.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
import numpy as np
19+
import pandas as pd
20+
import pyarrow as pa
21+
22+
from . import common
23+
from .common import KILOBYTE, MEGABYTE
24+
25+
26+
def generate_chunks(total_size, nchunks, ncols, dtype=np.dtype('int64')):
27+
rowsize = total_size // nchunks // ncols
28+
assert rowsize % dtype.itemsize == 0
29+
return [pd.DataFrame({
30+
'c' + str(col): np.frombuffer(
31+
common.get_random_bytes(rowsize, seed=col + 997 * chunk)).view(dtype)
32+
for col in range(ncols)
33+
})
34+
for chunk in range(nchunks)]
35+
36+
37+
class StreamReader(object):
38+
"""
39+
Benchmark in-memory streaming to a Pandas dataframe.
40+
"""
41+
total_size = 64 * MEGABYTE
42+
ncols = 8
43+
chunk_sizes = [16 * KILOBYTE, 256 * KILOBYTE, 8 * MEGABYTE]
44+
45+
param_names = ['chunk_size']
46+
params = [chunk_sizes]
47+
48+
def setup(self, chunk_size):
49+
# Note we're careful to stream different chunks instead of
50+
# streaming N times the same chunk, so that we avoid operating
51+
# entirely out of L1/L2.
52+
chunks = generate_chunks(self.total_size,
53+
nchunks=self.total_size // chunk_size,
54+
ncols=self.ncols)
55+
batches = [pa.RecordBatch.from_pandas(df)
56+
for df in chunks]
57+
schema = batches[0].schema
58+
sink = pa.BufferOutputStream()
59+
stream_writer = pa.RecordBatchStreamWriter(sink, schema)
60+
for batch in batches:
61+
stream_writer.write_batch(batch)
62+
self.source = sink.get_result()
63+
64+
def time_read_to_dataframe(self, *args):
65+
reader = pa.RecordBatchStreamReader(self.source)
66+
table = reader.read_all()
67+
df = table.to_pandas()

0 commit comments

Comments
 (0)