I asked a related question about a more idiomatic way to select rows from a PyArrow table based on contents of a column. The answer from @joris looks great. But it looks like selecting rows purely in PyArrow with a row mask has performance issues with sparse selections.
Is there a more efficient way to do this, but stay purely in PyArrow without going back and forth between PyArrow and numpy?
Test case:
import pyarrow as pa
import pyarrow.parquet as pq
import pandas as pd
import numpy as np
# Example table for data schema:
# Alternating rows with index 0 and 1
irow = np.arange(2**20)
dt = 17
df0 = pd.DataFrame({'timestamp': np.array((irow//2)*dt, dtype=np.int64),
'index': np.array(irow%2, dtype=np.int16),
'value': np.array(irow*0, dtype=np.int32)},
columns=['timestamp','index','value'])
ii = df0['index'] == 0
df0.loc[ii,'value'] = irow[ii]//2
ii = df0['index'] == 1
df0.loc[ii,'value'] = (np.sin(df0.loc[ii,'timestamp']*0.01)*10000).astype(np.int32)
# Insert rows with index 2 every 16 timestamps
irow = np.arange(10000)
subsample = 16
df1 = pd.DataFrame({'timestamp': np.array(irow*dt*subsample, dtype=np.int64),
'index': np.full_like(irow, 2, dtype=np.int16),
'value': np.array(irow*irow, dtype=np.int32)},
columns=['timestamp','index','value'],
index=irow*subsample*2+1.5)
df2=pd.concat([df0,df1]).sort_index()
df2.index = pd.RangeIndex(len(df2))
print(df2)
table2 = pa.Table.from_pandas(df2)
# which prints:
timestamp index value
0 0 0 0
1 0 1 0
2 0 2 0
3 17 0 1
4 17 1 1691
... ... ... ...
1058571 8912845 1 9945
1058572 8912862 0 524286
1058573 8912862 1 9978
1058574 8912879 0 524287
1058575 8912879 1 9723
[1058576 rows x 3 columns]
to verify there is sparse content with index=2:
print(df2[df2['index']==2])
# which prints
timestamp index value
2 0 2 0
35 272 2 1
68 544 2 4
101 816 2 9
134 1088 2 16
... ... ... ...
329837 2718640 2 99900025
329870 2718912 2 99920016
329903 2719184 2 99940009
329936 2719456 2 99960004
329969 2719728 2 99980001
[10000 rows x 3 columns]
and benchmarking:
import time
# My method, which sloshes back and forth between PyArrow and numpy
def select_by_index_np(table, ival):
value_index = table.column('index').to_numpy()
row_indices = np.nonzero(value_index==ival)[0]
return table.take(pa.array(row_indices))
# Stay in PyArrow: see https://stackoverflow.com/a/64579502/44330
def select_by_index(table, ival):
value_index = table.column('index')
index_type = value_index.type.to_pandas_dtype()
mask = pc.equal(value_index, index_type(ival))
return table.filter(mask)
def run_timing_test(table, ival, select_algorithm, nrep=100):
t1 = time.time_ns()
for _ in range(nrep):
tsel = select_algorithm(table, ival)
t2 = time.time_ns()
print('%.0fus %20s(%s) -> %s' %
((t2-t1)/1000/nrep,
select_algorithm.__name__,
ival,
tsel.column('value').to_numpy()))
run_timing_test(table2, 0, select_by_index)
run_timing_test(table2, 0, select_by_index_np)
run_timing_test(table2, 1, select_by_index)
run_timing_test(table2, 1, select_by_index_np)
run_timing_test(table2, 2, select_by_index)
run_timing_test(table2, 2, select_by_index_np)
# which prints
7639us select_by_index(0) -> [ 0 1 2 ... 524285 524286 524287]
7780us select_by_index_np(0) -> [ 0 1 2 ... 524285 524286 524287]
7789us select_by_index(1) -> [ 0 1691 3334 ... 9945 9978 9723]
8204us select_by_index_np(1) -> [ 0 1691 3334 ... 9945 9978 9723]
3840us select_by_index(2) -> [ 0 1 4 ... 99940009 99960004 99980001]
1611us select_by_index_np(2) -> [ 0 1 4 ... 99940009 99960004 99980001]
The two methods are comparable when selected rows are a substantial portion of the table, but when they are very small, select_by_index_np which marshals to numpy, determines indices where the rows of the mask are True, and marshals back to PyArrow, is faster!
Is there an efficient way to do this but stay in PyArrow? (I don't see any pyarrow.compute equivalent to numpy.nonzero)