Skip to content

Commit e850962

Browse files
arhamchopraarhamchopra
committed
Move numpyReshape/numpyShape/npyTypeFromPyType to NumpyConversions
These functions are pure numpy utilities, not CppNode-specific. Moving them to NumpyConversions.h/.cpp keeps ArrowCppNodes.cpp focused on node wiring and makes the utilities reusable. Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>
1 parent 23a9469 commit e850962

3 files changed

Lines changed: 78 additions & 74 deletions

File tree

‎cpp/csp/python/NumpyConversions.cpp‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,5 +202,71 @@ void validateNumpyTypeVsCspType( const CspTypePtr & type, char numpy_type_char )
202202
}
203203
}
204204

205+
DialectGenericType numpyReshape( DialectGenericType flatData, const std::vector<int64_t> & dims )
206+
{
207+
// PyArray_Newshape accepts a PyArray_Dims struct directly, avoiding a numpy array allocation for dims.
208+
// Stack-allocate the dims for typical 2-8 dimensional arrays.
209+
npy_intp dimsBuf[8];
210+
npy_intp ndims = static_cast<npy_intp>( dims.size() );
211+
npy_intp * dimsPtr;
212+
213+
std::vector<npy_intp> heapDims;
214+
if( ndims <= 8 )
215+
{
216+
for( npy_intp i = 0; i < ndims; ++i )
217+
dimsBuf[i] = static_cast<npy_intp>( dims[i] );
218+
dimsPtr = dimsBuf;
219+
}
220+
else
221+
{
222+
heapDims.resize( ndims );
223+
for( npy_intp i = 0; i < ndims; ++i )
224+
heapDims[i] = static_cast<npy_intp>( dims[i] );
225+
dimsPtr = heapDims.data();
226+
}
227+
228+
PyArray_Dims newshape{ dimsPtr, static_cast<int>( ndims ) };
229+
auto * flatPyArr = reinterpret_cast<PyArrayObject *>( toPythonBorrowed( flatData ) );
230+
PyObjectPtr reshaped{ PyObjectPtr::own(
231+
reinterpret_cast<PyObject *>( PyArray_Newshape( flatPyArr, &newshape, NPY_CORDER ) ) ) };
232+
if( !reshaped.get() )
233+
CSP_THROW( PythonPassthrough, "" );
234+
235+
return fromPython<DialectGenericType>( reshaped.get() );
236+
}
237+
238+
std::vector<int64_t> numpyShape( DialectGenericType ndarray )
239+
{
240+
auto * pyArr = reinterpret_cast<PyArrayObject *>( toPythonBorrowed( ndarray ) );
241+
int ndim = PyArray_NDIM( pyArr );
242+
npy_intp * shape = PyArray_SHAPE( pyArr );
243+
std::vector<int64_t> result( ndim );
244+
for( int i = 0; i < ndim; ++i )
245+
result[i] = shape[i];
246+
return result;
247+
}
248+
249+
int npyTypeFromPyType( DialectGenericType pyTypeObj )
250+
{
251+
auto & cspType = pyTypeAsCspType( toPythonBorrowed( pyTypeObj ) );
252+
auto cspTypeEnum = cspType -> type();
253+
254+
// Temporal types don't have NPY_TYPE<> specializations — handle them explicitly
255+
if( cspTypeEnum == CspType::Type::DATETIME || cspTypeEnum == CspType::Type::DATE )
256+
return NPY_DATETIME;
257+
if( cspTypeEnum == CspType::Type::TIMEDELTA )
258+
return NPY_TIMEDELTA;
259+
260+
return csp::PartialSwitchCspType<csp::CspType::Type::DOUBLE, csp::CspType::Type::INT64,
261+
csp::CspType::Type::BOOL, csp::CspType::Type::STRING>::invoke(
262+
cspType.get(),
263+
[]( auto tag ) -> int
264+
{
265+
using CValueType = typename decltype( tag )::type;
266+
return NPY_TYPE<CValueType>::value;
267+
}
268+
);
269+
}
270+
205271

206272
}

‎cpp/csp/python/NumpyConversions.h‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,15 @@ void stringFromNumpyStr( void* data, std::string& out, char numpy_type, int elem
240240

241241
void validateNumpyTypeVsCspType( const CspTypePtr & type, char numpy_type_char );
242242

243+
// Reshape a flat 1D numpy array into an NDArray using the given dimensions.
244+
CSPIMPL_EXPORT DialectGenericType numpyReshape( DialectGenericType flatData, const std::vector<int64_t> & dims );
245+
246+
// Extract shape from a numpy NDArray as a vector of dimension sizes.
247+
CSPIMPL_EXPORT std::vector<int64_t> numpyShape( DialectGenericType ndarray );
248+
249+
// Convert a Python type object (passed as DialectGenericType) to an NPY type constant.
250+
CSPIMPL_EXPORT int npyTypeFromPyType( DialectGenericType pyTypeObj );
251+
243252
}
244253

245254
#endif

‎cpp/csp/python/adapters/ArrowCppNodes.cpp‎

Lines changed: 3 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
#include <csp/adapters/arrow/StructToRecordBatch.h>
1111
#include <csp/engine/CppNode.h>
1212
#include <csp/engine/Dictionary.h>
13-
#include <csp/engine/PartialSwitchCspType.h>
1413
#include <csp/python/Conversions.h>
1514
#include <csp/python/Exception.h>
1615
#include <csp/python/InitHelper.h>
@@ -27,76 +26,6 @@ using namespace csp::adapters::arrow;
2726
namespace csp::cppnodes
2827
{
2928

30-
// Reshape callback that uses numpy's PyArray_Newshape (avoids allocating a dims array per call)
31-
static DialectGenericType numpyReshape( DialectGenericType flatData, const std::vector<int64_t> & dims )
32-
{
33-
// PyArray_Newshape accepts a PyArray_Dims struct directly, avoiding a numpy array allocation for dims.
34-
// Stack-allocate the dims for typical 2-8 dimensional arrays.
35-
npy_intp dimsBuf[8];
36-
npy_intp ndims = static_cast<npy_intp>( dims.size() );
37-
npy_intp * dimsPtr;
38-
39-
std::vector<npy_intp> heapDims;
40-
if( ndims <= 8 )
41-
{
42-
for( npy_intp i = 0; i < ndims; ++i )
43-
dimsBuf[i] = static_cast<npy_intp>( dims[i] );
44-
dimsPtr = dimsBuf;
45-
}
46-
else
47-
{
48-
heapDims.resize( ndims );
49-
for( npy_intp i = 0; i < ndims; ++i )
50-
heapDims[i] = static_cast<npy_intp>( dims[i] );
51-
dimsPtr = heapDims.data();
52-
}
53-
54-
PyArray_Dims newshape{ dimsPtr, static_cast<int>( ndims ) };
55-
auto * flatPyArr = reinterpret_cast<PyArrayObject *>( csp::python::toPythonBorrowed( flatData ) );
56-
csp::python::PyObjectPtr reshaped{ csp::python::PyObjectPtr::own(
57-
reinterpret_cast<PyObject *>( PyArray_Newshape( flatPyArr, &newshape, NPY_CORDER ) ) ) };
58-
if( !reshaped.get() )
59-
CSP_THROW( csp::python::PythonPassthrough, "" );
60-
61-
return csp::python::fromPython<DialectGenericType>( reshaped.get() );
62-
}
63-
64-
// Shape callback: extract shape from a numpy NDArray
65-
static std::vector<int64_t> numpyShape( DialectGenericType ndarray )
66-
{
67-
auto * pyArr = reinterpret_cast<PyArrayObject *>( csp::python::toPythonBorrowed( ndarray ) );
68-
int ndim = PyArray_NDIM( pyArr );
69-
npy_intp * shape = PyArray_SHAPE( pyArr );
70-
std::vector<int64_t> result( ndim );
71-
for( int i = 0; i < ndim; ++i )
72-
result[i] = shape[i];
73-
return result;
74-
}
75-
76-
// Convert a Python type object (passed as DialectGenericType) to an NPY type constant
77-
// using the same pyTypeAsCspType + PartialSwitchCspType pattern as the parquet adapter.
78-
static int npyTypeFromPyType( DialectGenericType pyTypeObj )
79-
{
80-
auto & cspType = csp::python::pyTypeAsCspType( csp::python::toPythonBorrowed( pyTypeObj ) );
81-
auto cspTypeEnum = cspType -> type();
82-
83-
// Temporal types don't have NPY_TYPE<> specializations — handle them explicitly
84-
if( cspTypeEnum == CspType::Type::DATETIME || cspTypeEnum == CspType::Type::DATE )
85-
return NPY_DATETIME;
86-
if( cspTypeEnum == CspType::Type::TIMEDELTA )
87-
return NPY_TIMEDELTA;
88-
89-
return csp::PartialSwitchCspType<csp::CspType::Type::DOUBLE, csp::CspType::Type::INT64,
90-
csp::CspType::Type::BOOL, csp::CspType::Type::STRING>::invoke(
91-
cspType.get(),
92-
[]( auto tag ) -> int
93-
{
94-
using CValueType = typename decltype( tag )::type;
95-
return csp::python::NPY_TYPE<CValueType>::value;
96-
}
97-
);
98-
}
99-
10029
// PyCapsule destructors for ArrowSchema/ArrowArray (local copies to avoid including ArrowInputAdapter.h)
10130
static void releaseArrowSchemaCapsule( PyObject * capsule )
10231
{
@@ -172,7 +101,7 @@ DECLARE_CPPNODE( record_batches_to_struct )
172101
{
173102
auto dimsColName = numpyDimensionNames -> get<std::string>( colName );
174103
customReaders.push_back(
175-
csp::python::createNumpyNDArrayReader( s_schema, colName, dimsColName, structField, numpyReshape ) );
104+
csp::python::createNumpyNDArrayReader( s_schema, colName, dimsColName, structField, csp::python::numpyReshape ) );
176105
}
177106
else
178107
{
@@ -285,13 +214,13 @@ DECLARE_CPPNODE( struct_to_record_batches )
285214
"Struct field '" << fieldName << "' not found on struct type '" << structMeta -> name() << "'" );
286215

287216
auto pyTypeObj = numpyElementTypes -> get<DialectGenericType>( fieldName );
288-
int npyType = npyTypeFromPyType( pyTypeObj );
217+
int npyType = csp::python::npyTypeFromPyType( pyTypeObj );
289218

290219
if( numpyDimensionNames && numpyDimensionNames -> exists( colName ) )
291220
{
292221
auto dimsColName = numpyDimensionNames -> get<std::string>( colName );
293222
customWriters.push_back(
294-
csp::python::createNumpyNDArrayWriter( colName, dimsColName, structField, npyType, numpyShape ) );
223+
csp::python::createNumpyNDArrayWriter( colName, dimsColName, structField, npyType, csp::python::numpyShape ) );
295224
}
296225
else
297226
{

0 commit comments

Comments
 (0)