Skip to content

Commit 23a9469

Browse files
committed
Fix ND string array write OOB and temporal NaT round-trip
- StringListWriter: use flat pointer arithmetic on C-contiguous data instead of PyArray_GETPTR1 which uses strides[0] only, causing OOB access for ndim > 1 arrays - makeStringListItemsWriter: same fix for the factory path - TemporalListWriter: check for NPY_DATETIME_NAT and emit Arrow null instead of writing raw INT64_MIN as a value (which overflows on scaling and loses NaT semantics) - DateListWriter: same NaT check before date conversion Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>
1 parent 6a47af0 commit 23a9469

1 file changed

Lines changed: 38 additions & 15 deletions

File tree

cpp/csp/python/adapters/ArrowNumpyListWriter.h

Lines changed: 38 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -153,9 +153,21 @@ class StringListWriter final : public csp::adapters::arrow::FieldWriter
153153
auto elementSize = PyDataType_ELSIZE( PyArray_DESCR( pyArr ) );
154154
auto charCount = elementSize / sizeof( char32_t );
155155

156+
// Ensure C-contiguous layout for safe flat iteration over ND arrays.
157+
// PyArray_GETPTR1 uses strides[0] only, which is incorrect for ndim > 1.
158+
PyArrayObject * contiguousArr = pyArr;
159+
PyObjectPtr contiguousOwner;
160+
if( !PyArray_IS_C_CONTIGUOUS( pyArr ) || PyArray_NDIM( pyArr ) > 1 )
161+
{
162+
contiguousOwner = PyObjectPtr::own(
163+
reinterpret_cast<PyObject *>( PyArray_GETCONTIGUOUS( pyArr ) ) );
164+
contiguousArr = reinterpret_cast<PyArrayObject *>( contiguousOwner.get() );
165+
}
166+
167+
auto * base = reinterpret_cast<char *>( PyArray_DATA( contiguousArr ) );
156168
for( npy_intp i = 0; i < len; ++i )
157169
{
158-
auto * ptr = reinterpret_cast<char32_t *>( PyArray_GETPTR1( pyArr, i ) );
170+
auto * ptr = reinterpret_cast<char32_t *>( base + i * elementSize );
159171
// Find actual string length (exclude trailing nulls)
160172
size_t actualLen = 0;
161173
for( size_t c = 0; c < charCount; ++c )
@@ -233,19 +245,14 @@ class TemporalListWriter final : public csp::adapters::arrow::FieldWriter
233245
auto unit = datetimeUnitFromDescr( PyArray_DESCR( pyArr ) );
234246
int64_t scaling = scalingFromNumpyDtUnit( unit );
235247

236-
if( scaling == 1 )
237-
{
238-
// Already nanoseconds — bulk append
239-
ARROW_OK_OR_THROW_WRITER(
240-
m_valueBuilder -> AppendValues( data, static_cast<int64_t>( len ) ),
241-
"Failed to append temporal list elements" );
242-
}
243-
else
248+
ARROW_OK_OR_THROW_WRITER(
249+
m_valueBuilder -> Reserve( static_cast<int64_t>( len ) ),
250+
"Failed to reserve temporal value builder" );
251+
for( npy_intp i = 0; i < len; ++i )
244252
{
245-
ARROW_OK_OR_THROW_WRITER(
246-
m_valueBuilder -> Reserve( static_cast<int64_t>( len ) ),
247-
"Failed to reserve temporal value builder" );
248-
for( npy_intp i = 0; i < len; ++i )
253+
if( data[i] == NPY_DATETIME_NAT )
254+
ARROW_OK_OR_THROW_WRITER( m_valueBuilder -> AppendNull(), "Failed to append null temporal element" );
255+
else
249256
m_valueBuilder -> UnsafeAppend( data[i] * scaling );
250257
}
251258
}
@@ -311,7 +318,12 @@ class DateListWriter final : public csp::adapters::arrow::FieldWriter
311318
m_valueBuilder -> Reserve( static_cast<int64_t>( len ) ),
312319
"Failed to reserve date value builder" );
313320
for( npy_intp i = 0; i < len; ++i )
314-
m_valueBuilder -> UnsafeAppend( static_cast<int32_t>( data[i] * scaling / csp::NANOS_PER_DAY ) );
321+
{
322+
if( data[i] == NPY_DATETIME_NAT )
323+
ARROW_OK_OR_THROW_WRITER( m_valueBuilder -> AppendNull(), "Failed to append null date element" );
324+
else
325+
m_valueBuilder -> UnsafeAppend( static_cast<int32_t>( data[i] * scaling / csp::NANOS_PER_DAY ) );
326+
}
315327
}
316328

317329
private:
@@ -484,10 +496,21 @@ inline csp::adapters::arrow::ListItemsWriter makeStringListItemsWriter(
484496
auto elementSize = PyDataType_ELSIZE( PyArray_DESCR( pyArr ) );
485497
auto charCount = elementSize / sizeof( char32_t );
486498

499+
// Ensure C-contiguous layout for safe flat iteration over ND arrays.
500+
PyArrayObject * contiguousArr = pyArr;
501+
PyObjectPtr contiguousOwner;
502+
if( !PyArray_IS_C_CONTIGUOUS( pyArr ) || PyArray_NDIM( pyArr ) > 1 )
503+
{
504+
contiguousOwner = PyObjectPtr::own(
505+
reinterpret_cast<PyObject *>( PyArray_GETCONTIGUOUS( pyArr ) ) );
506+
contiguousArr = reinterpret_cast<PyArrayObject *>( contiguousOwner.get() );
507+
}
508+
509+
auto * base = reinterpret_cast<char *>( PyArray_DATA( contiguousArr ) );
487510
std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
488511
for( npy_intp i = 0; i < len; ++i )
489512
{
490-
auto * ptr = reinterpret_cast<char32_t *>( PyArray_GETPTR1( pyArr, i ) );
513+
auto * ptr = reinterpret_cast<char32_t *>( base + i * elementSize );
491514
size_t actualLen = 0;
492515
for( size_t c = 0; c < charCount; ++c )
493516
{

0 commit comments

Comments
 (0)