Skip to content

Commit 0f27b36

Browse files
arhamchopraarhamchopra
authored
Add record batch adapter to zero-copy stream record batches into csp (#541)
* Add arrow adapter to stream record batches zero-copy into csp Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com> * Address comments, add tests, cleanup Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com> * Address more comments Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com> * Address more comments Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com> * Small comment cleanup Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com> --------- Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>
1 parent c86b842 commit 0f27b36

9 files changed

Lines changed: 743 additions & 2 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ option(CSP_USE_LD_CLASSIC_MAC "On macOS, link with ld_classic" OFF)
7070

7171
# Extension options
7272
option(CSP_BUILD_KAFKA_ADAPTER "Build kafka adapter" ON)
73+
option(CSP_BUILD_ARROW_ADAPTER "Build arrow adapter" ON)
7374
option(CSP_BUILD_PARQUET_ADAPTER "Build parquet adapter" ON)
7475
option(CSP_BUILD_WS_CLIENT_ADAPTER "Build ws client adapter" ON)
7576

@@ -238,6 +239,10 @@ find_package(csp_autogen REQUIRED)
238239
find_package(DepsBase REQUIRED)
239240

240241
# Adapter dependencies
242+
if(CSP_BUILD_ARROW_ADAPTER)
243+
find_package(DepsArrowAdapter REQUIRED)
244+
endif()
245+
241246
if(CSP_BUILD_KAFKA_ADAPTER)
242247
find_package(DepsKafkaAdapter REQUIRED)
243248
endif()

‎cpp/cmake/modules/FindDepsArrowAdapter.cmake‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
cmake_minimum_required(VERSION 3.7.2)
2+
3+
# ARROW
4+
find_package(Arrow REQUIRED)
5+
include_directories(${ARROW_INCLUDE_DIR})

‎cpp/csp/python/CMakeLists.txt‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ target_compile_definitions(csptypesimpl PUBLIC RAPIDJSON_HAS_STDSTRING=1)
2424
target_link_libraries(csptypesimpl csp_core csp_types)
2525
target_compile_definitions(csptypesimpl PRIVATE CSPTYPESIMPL_EXPORTS=1)
2626

27-
2827
set(CSPIMPL_PUBLIC_HEADERS
2928
Common.h
3029
Conversions.h
@@ -89,7 +88,6 @@ target_link_libraries(cspimpl csptypesimpl csp_core csp_engine )
8988
target_compile_definitions(cspimpl PUBLIC NPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION)
9089
target_compile_definitions(cspimpl PRIVATE CSPIMPL_EXPORTS=1)
9190

92-
9391
## Baselib c++ module
9492
add_library(cspbaselibimpl SHARED cspbaselibimpl.cpp)
9593
target_link_libraries(cspbaselibimpl cspimpl baselibimpl)

‎cpp/csp/python/adapters/ArrowInputAdapter.h‎

Lines changed: 287 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,287 @@
1+
#ifndef _IN_CSP_ENGINE_ARROWINPUTADAPTER_H
2+
#define _IN_CSP_ENGINE_ARROWINPUTADAPTER_H
3+
4+
#include <csp/engine/PullInputAdapter.h>
5+
#include <csp/python/PyObjectPtr.h>
6+
#include <csp/python/Conversions.h>
7+
#include <csp/core/Time.h>
8+
#include <Python.h>
9+
10+
#include <arrow/array.h>
11+
#include <arrow/c/abi.h>
12+
#include <arrow/c/bridge.h>
13+
#include <arrow/type.h>
14+
#include <arrow/table.h>
15+
16+
#include <algorithm>
17+
#include <memory>
18+
#include <string>
19+
20+
namespace csp::python::arrow
21+
{
22+
23+
class RecordBatchIterator
24+
{
25+
public:
26+
RecordBatchIterator() {}
27+
RecordBatchIterator( PyObjectPtr iter, std::shared_ptr<::arrow::Schema> schema ): m_iter( std::move( iter ) ), m_schema( schema )
28+
{
29+
}
30+
31+
std::shared_ptr<::arrow::RecordBatch> next()
32+
{
33+
auto py_tuple = csp::python::PyObjectPtr::own( PyIter_Next( m_iter.get() ) );
34+
if( PyErr_Occurred() )
35+
CSP_THROW( csp::python::PythonPassthrough, "" );
36+
37+
if( py_tuple.get() == NULL )
38+
{
39+
// No more data in the input steam
40+
return nullptr;
41+
}
42+
43+
if( !PyTuple_Check( py_tuple.get() ) )
44+
CSP_THROW( csp::TypeError, "Invalid arrow data, expected tuple (using the PyCapsule C interface) got " << Py_TYPE( py_tuple.get() ) -> tp_name );
45+
46+
auto num_elems = PyTuple_Size( py_tuple.get() );
47+
if( num_elems != 2 )
48+
CSP_THROW( csp::TypeError, "Invalid arrow data, expected tuple (using the PyCapsule C interface) with 2 elements got " << num_elems );
49+
50+
// Extract the record batch
51+
PyObject * py_array = PyTuple_GetItem( py_tuple.get(), 1 );
52+
if( !PyCapsule_IsValid( py_array, "arrow_array" ) )
53+
CSP_THROW( csp::TypeError, "Invalid arrow data, expected tuple from the PyCapsule C interface " );
54+
55+
ArrowArray * c_array = reinterpret_cast<ArrowArray*>( PyCapsule_GetPointer( py_array, "arrow_array" ) );
56+
auto result = ::arrow::ImportRecordBatch( c_array, m_schema );
57+
if( !result.ok() )
58+
CSP_THROW( ValueError, "Failed to load record batches through PyCapsule C Data interface: " << result.status().ToString() );
59+
60+
return result.ValueUnsafe();
61+
}
62+
63+
private:
64+
PyObjectPtr m_iter;
65+
std::shared_ptr<::arrow::Schema> m_schema;
66+
};
67+
68+
void ReleaseArrowSchemaPyCapsule( PyObject * capsule )
69+
{
70+
ArrowSchema * schema = reinterpret_cast<ArrowSchema*>( PyCapsule_GetPointer( capsule, "arrow_schema" ) );
71+
if ( schema -> release != NULL )
72+
schema -> release( schema );
73+
free( schema );
74+
}
75+
76+
void ReleaseArrowArrayPyCapsule( PyObject * capsule )
77+
{
78+
ArrowArray * array = reinterpret_cast<ArrowArray*>( PyCapsule_GetPointer( capsule, "arrow_array" ) );
79+
if ( array -> release != NULL )
80+
array -> release( array );
81+
free( array );
82+
}
83+
84+
class RecordBatchInputAdapter: public PullInputAdapter<std::vector<DialectGenericType>>
85+
{
86+
public:
87+
RecordBatchInputAdapter( Engine * engine, CspTypePtr & type, PyObjectPtr pySchema, std::string tsColName, PyObjectPtr source, bool expectSmallBatches )
88+
: PullInputAdapter<std::vector<DialectGenericType>>( engine, type, PushMode::LAST_VALUE ),
89+
m_tsColName( tsColName ),
90+
m_expectSmallBatches( expectSmallBatches ),
91+
m_finished( false )
92+
{
93+
// Extract the arrow schema
94+
ArrowSchema * c_schema = reinterpret_cast<ArrowSchema*>( PyCapsule_GetPointer( pySchema.get(), "arrow_schema" ) );
95+
auto result = ::arrow::ImportSchema( c_schema );
96+
if( !result.ok() )
97+
CSP_THROW( ValueError, "Failed to load schema for record batches through the PyCapsule C Data interface: " << result.status().ToString() );
98+
m_schema = std::move( result.ValueUnsafe() );
99+
100+
auto tsField = m_schema -> GetFieldByName( m_tsColName );
101+
auto timestampType = std::static_pointer_cast<::arrow::TimestampType>( tsField -> type() );
102+
switch( timestampType -> unit() )
103+
{
104+
case ::arrow::TimeUnit::SECOND:
105+
m_multiplier = csp::NANOS_PER_SECOND;
106+
break;
107+
case ::arrow::TimeUnit::MILLI:
108+
m_multiplier = csp::NANOS_PER_MILLISECOND;
109+
break;
110+
case ::arrow::TimeUnit::MICRO:
111+
m_multiplier = csp::NANOS_PER_MICROSECOND;
112+
break;
113+
case ::arrow::TimeUnit::NANO:
114+
m_multiplier = 1;
115+
break;
116+
default:
117+
CSP_THROW( ValueError, "Unsupported unit type for arrow timestamp column" );
118+
}
119+
120+
m_source = RecordBatchIterator( source, m_schema );
121+
}
122+
123+
int64_t findFirstMatchingIndex()
124+
{
125+
// Find the first index with time equal or greater than `time`
126+
auto first_time = m_tsArray -> Value( 0 );
127+
if( first_time >= m_startTime )
128+
{
129+
// Early break
130+
return 0;
131+
}
132+
133+
auto last_time = m_tsArray -> Value( m_numRows - 1 );
134+
if( last_time < m_startTime )
135+
{
136+
// Early break
137+
return m_numRows;
138+
}
139+
::arrow::TimestampArray::IteratorType it;
140+
auto begin = ++( m_tsArray -> begin() );
141+
if( m_expectSmallBatches )
142+
{
143+
auto predicate = [this]( std::optional<int64_t> new_time ) -> bool { return new_time.value() >= this -> m_startTime; };
144+
it = std::find_if( begin, m_endIt, predicate );
145+
}
146+
else
147+
it = std::lower_bound( begin, m_endIt, m_startTime );
148+
return it.index();
149+
}
150+
151+
int64_t findNextLargerTimestampIndex( int64_t start_idx )
152+
{
153+
// Find the first index with time just greater than the time at start_idx
154+
auto cur_time = m_tsArray -> Value( start_idx );
155+
auto begin = ::arrow::TimestampArray::IteratorType( *m_tsArray, start_idx );
156+
::arrow::TimestampArray::IteratorType it;
157+
if( m_expectSmallBatches )
158+
{
159+
auto predicate = [cur_time]( std::optional<int64_t> new_time ) -> bool { return cur_time != new_time.value(); };
160+
it = std::find_if( begin, m_endIt, predicate );
161+
}
162+
else
163+
{
164+
if( m_arrayLastTime == cur_time )
165+
return m_numRows;
166+
it = std::upper_bound( begin, m_endIt, cur_time );
167+
}
168+
return it.index();
169+
}
170+
171+
std::shared_ptr<::arrow::RecordBatch> updateStateFromNextRecordBatch()
172+
{
173+
std::shared_ptr<::arrow::RecordBatch> rb;
174+
while( ( rb = m_source.next() ) && ( rb -> num_rows() == 0 ) ) {}
175+
if( rb )
176+
{
177+
auto array = rb -> GetColumnByName( m_tsColName );
178+
if( !array )
179+
CSP_THROW( ValueError, "Failed to get timestamp column " << m_tsColName << " from record batch " << rb -> ToString() );
180+
181+
m_tsArray = std::static_pointer_cast<::arrow::TimestampArray>( array );
182+
m_numRows = m_tsArray -> length();
183+
m_endIt = m_tsArray -> end();
184+
m_arrayLastTime = m_tsArray -> Value( m_numRows - 1 );
185+
}
186+
m_curRecordBatch = rb;
187+
return rb;
188+
}
189+
190+
void start( DateTime start, DateTime end ) override
191+
{
192+
// start and end as multiples of the unit in timestamp column
193+
auto start_nanos = start.asNanoseconds();
194+
m_startTime = ( start_nanos % m_multiplier == 0 ) ? start_nanos / m_multiplier : start_nanos / m_multiplier + 1;
195+
m_endTime = end.asNanoseconds() / m_multiplier;
196+
197+
// Find the starting index where time >= start
198+
while( !m_finished )
199+
{
200+
updateStateFromNextRecordBatch();
201+
if( !m_curRecordBatch )
202+
{
203+
m_finished = true;
204+
break;
205+
}
206+
m_curBatchIdx = findFirstMatchingIndex();
207+
if( m_curBatchIdx < m_numRows )
208+
break;
209+
}
210+
PullInputAdapter<std::vector<DialectGenericType>>::start( start, end );
211+
}
212+
213+
DialectGenericType convertRecordBatchToPython( std::shared_ptr<::arrow::RecordBatch> rb )
214+
{
215+
ArrowSchema* rb_schema = ( ArrowSchema* )malloc( sizeof( ArrowSchema ) );
216+
ArrowArray* rb_array = ( ArrowArray* )malloc( sizeof( ArrowArray ) );
217+
::arrow::Status st = ::arrow::ExportRecordBatch( *rb, rb_array, rb_schema );
218+
auto py_schema = csp::python::PyObjectPtr::own( PyCapsule_New( rb_schema, "arrow_schema", ReleaseArrowSchemaPyCapsule ) );
219+
auto py_array = csp::python::PyObjectPtr::own( PyCapsule_New( rb_array, "arrow_array", ReleaseArrowArrayPyCapsule ) );
220+
auto py_tuple = csp::python::PyObjectPtr::own( PyTuple_Pack( 2, py_schema.get(), py_array.get() ) );
221+
return csp::python::fromPython<DialectGenericType>( py_tuple.get() );
222+
}
223+
224+
bool next( DateTime & t, std::vector<DialectGenericType> & value ) override
225+
{
226+
std::vector<DialectGenericType> cur_result;
227+
int64_t cur_ts = 0;
228+
while( !m_finished )
229+
{
230+
// Slice current record batch
231+
auto new_ts = m_tsArray -> Value( m_curBatchIdx );
232+
if( new_ts > m_endTime )
233+
{
234+
// Past the end time
235+
m_finished = true;
236+
break;
237+
}
238+
if( !cur_result.empty() && new_ts != cur_ts )
239+
{
240+
// Next timestamp encountered, return the current list of record batches
241+
break;
242+
}
243+
cur_ts = new_ts;
244+
auto next_idx = findNextLargerTimestampIndex( m_curBatchIdx );
245+
auto slice = m_curRecordBatch -> Slice( m_curBatchIdx, next_idx - m_curBatchIdx );
246+
cur_result.emplace_back( convertRecordBatchToPython( slice ) );
247+
m_curBatchIdx = next_idx;
248+
if( m_curBatchIdx != m_numRows )
249+
{
250+
// All rows for current timestamp have been found
251+
break;
252+
}
253+
// Get the next record batch
254+
updateStateFromNextRecordBatch();
255+
if( !m_curRecordBatch )
256+
{
257+
m_finished = true;
258+
break;
259+
}
260+
m_curBatchIdx = 0;
261+
}
262+
if( !cur_result.empty() )
263+
{
264+
value = std::move( cur_result );
265+
t = csp::DateTime::fromNanoseconds( cur_ts * m_multiplier );
266+
return true;
267+
}
268+
return false;
269+
}
270+
271+
private:
272+
std::string m_tsColName;
273+
RecordBatchIterator m_source;
274+
275+
int m_expectSmallBatches;
276+
bool m_finished;
277+
std::shared_ptr<::arrow::Schema> m_schema;
278+
std::shared_ptr<::arrow::RecordBatch> m_curRecordBatch;
279+
std::shared_ptr<::arrow::TimestampArray> m_tsArray;
280+
::arrow::TimestampArray::IteratorType m_endIt;
281+
int64_t m_arrayLastTime;
282+
int64_t m_multiplier, m_numRows, m_startTime, m_endTime, m_curBatchIdx;
283+
};
284+
285+
};
286+
287+
#endif

‎cpp/csp/python/adapters/CMakeLists.txt‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,32 @@
1+
if(CSP_BUILD_ARROW_ADAPTER)
2+
add_library(arrowadapterimpl SHARED PyArrowInputAdapter.cpp ArrowInputAdapter.h)
3+
4+
if(WIN32)
5+
if(CSP_USE_VCPKG)
6+
set(ARROW_PACKAGES_TO_LINK Arrow::arrow_static)
7+
target_compile_definitions(arrowadapterimpl PUBLIC ARROW_STATIC)
8+
else()
9+
# use dynamic variants
10+
# Until we manage to get the fix for ws3_32.dll in arrow-16 into conda, manually fix the error here
11+
get_target_property(LINK_LIBS Arrow::arrow_shared INTERFACE_LINK_LIBRARIES)
12+
string(REPLACE "ws2_32.dll" "ws2_32" FIXED_LINK_LIBS "${LINK_LIBS}")
13+
set_target_properties(Arrow::arrow_shared PROPERTIES INTERFACE_LINK_LIBRARIES "${FIXED_LINK_LIBS}")
14+
set(ARROW_PACKAGES_TO_LINK arrow_shared)
15+
endif()
16+
else()
17+
if(CSP_USE_VCPKG)
18+
# use static variants
19+
set(ARROW_PACKAGES_TO_LINK arrow_static)
20+
else()
21+
# use dynamic variants
22+
set(ARROW_PACKAGES_TO_LINK arrow)
23+
endif()
24+
endif()
25+
26+
target_link_libraries(arrowadapterimpl csp_core csp_engine cspimpl ${ARROW_PACKAGES_TO_LINK})
27+
target_include_directories(arrowadapterimpl PUBLIC ${ARROW_INCLUDE_DIR})
28+
install(TARGETS arrowadapterimpl RUNTIME DESTINATION ${CSP_RUNTIME_INSTALL_SUBDIR} )
29+
endif()
130

231
if(CSP_BUILD_KAFKA_ADAPTER)
332
add_library(kafkaadapterimpl SHARED kafkaadapterimpl.cpp)

0 commit comments

Comments
 (0)