Add datetime, timedelta, date support for numpy arrays

arhamchopra · arhamchopra · commit af20b15fc30b · 2026-03-11T14:05:28.000-04:00
Signed-off-by: Arham Chopra &lt;arham.chopra@cubistsystematic.com&gt;
diff --git a/cpp/csp/python/adapters/ArrowCppNodes.cpp b/cpp/csp/python/adapters/ArrowCppNodes.cpp
@@ -78,6 +78,14 @@ static std::vector<int64_t> numpyShape( DialectGenericType ndarray )
 static int npyTypeFromPyType( DialectGenericType pyTypeObj )
 {
     auto & cspType = csp::python::pyTypeAsCspType( csp::python::toPythonBorrowed( pyTypeObj ) );
+    auto cspTypeEnum = cspType -> type();
+
+    // Temporal types don't have NPY_TYPE<> specializations — handle them explicitly
+    if( cspTypeEnum == CspType::Type::DATETIME || cspTypeEnum == CspType::Type::DATE )
+        return NPY_DATETIME;
+    if( cspTypeEnum == CspType::Type::TIMEDELTA )
+        return NPY_TIMEDELTA;
+
     return csp::PartialSwitchCspType<csp::CspType::Type::DOUBLE, csp::CspType::Type::INT64,
                                      csp::CspType::Type::BOOL, csp::CspType::Type::STRING>::invoke(
         cspType.get(),
diff --git a/cpp/csp/python/adapters/ArrowNumpyListReader.h b/cpp/csp/python/adapters/ArrowNumpyListReader.h
@@ -14,6 +14,8 @@
 #include <arrow/array.h>
 #include <arrow/type.h>
 
+#include <csp/core/Time.h>
+
 #include <codecvt>
 #include <cstring>
 #include <functional>
@@ -140,6 +142,117 @@ makeStringListReadValue()
     };
 }
 
+// Helper: compute nanosecond multiplier for a given arrow::TimeUnit
+inline int64_t listTimeUnitMultiplier( ::arrow::TimeUnit::type unit )
+{
+    switch( unit )
+    {
+        case ::arrow::TimeUnit::SECOND: return csp::NANOS_PER_SECOND;
+        case ::arrow::TimeUnit::MILLI:  return csp::NANOS_PER_MILLISECOND;
+        case ::arrow::TimeUnit::MICRO:  return csp::NANOS_PER_MICROSECOND;
+        case ::arrow::TimeUnit::NANO:   return 1LL;
+    }
+    CSP_THROW( TypeError, "Unexpected arrow TimeUnit: " << static_cast<int>( unit ) );
+}
+
+// Create a readValue lambda for datetime64[ns] from Arrow Timestamp/Date columns
+// ArrowValueArrayT is the typed Arrow array (TimestampArray, Date32Array, Date64Array)
+// multiplier converts from Arrow units to nanoseconds
+template<typename ArrowValueArrayT>
+std::function<DialectGenericType( const ::arrow::Array &, int64_t )>
+makeTimestampListReadValue( int64_t multiplier )
+{
+    // Lazily create and cache the datetime64[ns] descriptor
+    static PyArray_Descr * s_dtype = nullptr;
+    if( !s_dtype )
+    {
+        auto dtypeStr = PyPtr<PyObject>::own( PyUnicode_FromString( "<M8[ns]" ) );
+        PyArray_DescrConverter( dtypeStr.get(), &s_dtype );
+    }
+
+    return [multiplier]( const ::arrow::Array & arr, int64_t row ) -> DialectGenericType
+    {
+        auto & listArr = static_cast<const ::arrow::ListArray &>( arr );
+        auto values = std::static_pointer_cast<ArrowValueArrayT>( listArr.value_slice( row ) );
+
+        npy_intp size = values -> length();
+        Py_INCREF( s_dtype );
+        PyObject * pyArr = PyArray_SimpleNewFromDescr( 1, &size, s_dtype );
+        PyObjectPtr arrOwner{ PyObjectPtr::own( pyArr ) };
+
+        auto * buf = reinterpret_cast<int64_t *>( PyArray_DATA( reinterpret_cast<PyArrayObject *>( pyArr ) ) );
+
+        // Fast path: nanosecond units (multiplier==1), no nulls, AND 64-bit storage — bulk memcpy.
+        // Date32Array/Time32Array have 32-bit storage so memcpy is invalid for them;
+        // the constexpr sizeof check prevents this at compile time.
+        constexpr bool is64bit = sizeof( typename ArrowValueArrayT::value_type ) == sizeof( int64_t );
+        if constexpr( is64bit )
+        {
+            if( values -> null_count() == 0 && multiplier == 1 )
+            {
+                std::memcpy( buf, values -> raw_values(), sizeof( int64_t ) * size );
+                return fromPython<DialectGenericType>( pyArr );
+            }
+        }
+
+        for( int64_t i = 0; i < size; ++i )
+        {
+            if( values -> IsValid( i ) )
+                buf[i] = static_cast<int64_t>( values -> Value( i ) ) * multiplier;
+            else
+                buf[i] = NPY_DATETIME_NAT;
+        }
+
+        return fromPython<DialectGenericType>( pyArr );
+    };
+}
+
+// Create a readValue lambda for timedelta64[ns] from Arrow Duration/Time columns
+template<typename ArrowValueArrayT>
+std::function<DialectGenericType( const ::arrow::Array &, int64_t )>
+makeDurationListReadValue( int64_t multiplier )
+{
+    static PyArray_Descr * s_dtype = nullptr;
+    if( !s_dtype )
+    {
+        auto dtypeStr = PyPtr<PyObject>::own( PyUnicode_FromString( "<m8[ns]" ) );
+        PyArray_DescrConverter( dtypeStr.get(), &s_dtype );
+    }
+
+    return [multiplier]( const ::arrow::Array & arr, int64_t row ) -> DialectGenericType
+    {
+        auto & listArr = static_cast<const ::arrow::ListArray &>( arr );
+        auto values = std::static_pointer_cast<ArrowValueArrayT>( listArr.value_slice( row ) );
+
+        npy_intp size = values -> length();
+        Py_INCREF( s_dtype );
+        PyObject * pyArr = PyArray_SimpleNewFromDescr( 1, &size, s_dtype );
+        PyObjectPtr arrOwner{ PyObjectPtr::own( pyArr ) };
+
+        auto * buf = reinterpret_cast<int64_t *>( PyArray_DATA( reinterpret_cast<PyArrayObject *>( pyArr ) ) );
+
+        constexpr bool is64bit = sizeof( typename ArrowValueArrayT::value_type ) == sizeof( int64_t );
+        if constexpr( is64bit )
+        {
+            if( values -> null_count() == 0 && multiplier == 1 )
+            {
+                std::memcpy( buf, values -> raw_values(), sizeof( int64_t ) * size );
+                return fromPython<DialectGenericType>( pyArr );
+            }
+        }
+
+        for( int64_t i = 0; i < size; ++i )
+        {
+            if( values -> IsValid( i ) )
+                buf[i] = static_cast<int64_t>( values -> Value( i ) ) * multiplier;
+            else
+                buf[i] = NPY_DATETIME_NAT;
+        }
+
+        return fromPython<DialectGenericType>( pyArr );
+    };
+}
+
 // Dispatch on arrow list element type to create the appropriate readValue lambda
 inline std::function<DialectGenericType( const ::arrow::Array &, int64_t )>
 dispatchListReadValue( const std::shared_ptr<::arrow::ListType> & listType, const std::string & columnName )
@@ -162,6 +275,35 @@ dispatchListReadValue( const std::shared_ptr<::arrow::ListType> & listType, cons
             return makeStringListReadValue<::arrow::LargeStringArray>();
         case ::arrow::Type::LARGE_BINARY:
             return makeStringListReadValue<::arrow::LargeBinaryArray>();
+
+        // --- Temporal: datetime64[ns] ---
+        case ::arrow::Type::TIMESTAMP:
+        {
+            auto tsType = std::static_pointer_cast<::arrow::TimestampType>( listType -> value_type() );
+            return makeTimestampListReadValue<::arrow::TimestampArray>( listTimeUnitMultiplier( tsType -> unit() ) );
+        }
+        case ::arrow::Type::DATE32:
+            return makeTimestampListReadValue<::arrow::Date32Array>( csp::NANOS_PER_DAY );
+        case ::arrow::Type::DATE64:
+            return makeTimestampListReadValue<::arrow::Date64Array>( csp::NANOS_PER_MILLISECOND );
+
+        // --- Temporal: timedelta64[ns] ---
+        case ::arrow::Type::DURATION:
+        {
+            auto durType = std::static_pointer_cast<::arrow::DurationType>( listType -> value_type() );
+            return makeDurationListReadValue<::arrow::DurationArray>( listTimeUnitMultiplier( durType -> unit() ) );
+        }
+        case ::arrow::Type::TIME32:
+        {
+            auto timeType = std::static_pointer_cast<::arrow::Time32Type>( listType -> value_type() );
+            return makeDurationListReadValue<::arrow::Time32Array>( listTimeUnitMultiplier( timeType -> unit() ) );
+        }
+        case ::arrow::Type::TIME64:
+        {
+            auto timeType = std::static_pointer_cast<::arrow::Time64Type>( listType -> value_type() );
+            return makeDurationListReadValue<::arrow::Time64Array>( listTimeUnitMultiplier( timeType -> unit() ) );
+        }
+
         default:
             CSP_THROW( TypeError, "Unsupported list element type " << listType -> value_type() -> ToString()
                                    << " for list column '" << columnName << "'" );
diff --git a/cpp/csp/python/adapters/ArrowNumpyListWriter.h b/cpp/csp/python/adapters/ArrowNumpyListWriter.h
@@ -9,6 +9,7 @@
 #define _IN_CSP_PYTHON_ADAPTERS_ArrowNumpyListWriter_H
 
 #include <csp/adapters/arrow/ArrowFieldWriter.h>
+#include <csp/core/Time.h>
 #include <csp/python/Conversions.h>
 #include <csp/python/NumpyConversions.h>
 
@@ -174,6 +175,150 @@ class StringListWriter final : public csp::adapters::arrow::FieldWriter
     std::string                                         m_utf8Buf;            // reused buffer
 };
 
+// --- Temporal list writers (datetime64[ns] → Timestamp, timedelta64[ns] → Duration) ---
+
+// Writes numpy datetime64/timedelta64 arrays into Arrow List<Timestamp(ns,"UTC")> or List<Duration(ns)> columns.
+// Detects the numpy datetime unit at write time and converts to nanoseconds using scalingFromNumpyDtUnit().
+template<typename ArrowBuilderT>
+class TemporalListWriter final : public csp::adapters::arrow::FieldWriter
+{
+public:
+    TemporalListWriter( const std::string & columnName, const StructFieldPtr & structField,
+                        std::shared_ptr<::arrow::DataType> elementType )
+        : FieldWriter( { columnName }, { ::arrow::list( elementType ) }, structField )
+    {
+        m_valueBuilder = std::make_shared<ArrowBuilderT>( elementType, ::arrow::default_memory_pool() );
+        m_listBuilder = std::make_shared<::arrow::ListBuilder>( ::arrow::default_memory_pool(), m_valueBuilder );
+    }
+
+    void reserve( int64_t numRows ) override
+    {
+        ARROW_OK_OR_THROW_WRITER( m_listBuilder -> Reserve( numRows ), "Failed to reserve temporal list builder" );
+    }
+
+    void writeNull() override
+    {
+        ARROW_OK_OR_THROW_WRITER( m_listBuilder -> AppendNull(), "Failed to append null temporal list" );
+    }
+
+    std::vector<std::shared_ptr<::arrow::Array>> finish() override
+    {
+        std::shared_ptr<::arrow::Array> arr;
+        ARROW_OK_OR_THROW_WRITER( m_listBuilder -> Finish( &arr ), "Failed to finish temporal list array" );
+        return { arr };
+    }
+
+protected:
+    void doWrite( const Struct * s ) override
+    {
+        ARROW_OK_OR_THROW_WRITER( m_listBuilder -> Append(), "Failed to start temporal list" );
+        auto & dgt = m_field -> value<DialectGenericType>( s );
+        auto * pyArr = reinterpret_cast<PyArrayObject *>( csp::python::toPythonBorrowed( dgt ) );
+        npy_intp len = PyArray_SIZE( pyArr );
+
+        // PyArray_DATA requires C-contiguous layout for correct bulk reads.
+        // Non-contiguous arrays (slices, transposes) must be copied to contiguous form first.
+        PyArrayObject * contiguousArr = pyArr;
+        PyObjectPtr contiguousOwner;
+        if( !PyArray_IS_C_CONTIGUOUS( pyArr ) )
+        {
+            contiguousOwner = PyObjectPtr::own(
+                reinterpret_cast<PyObject *>( PyArray_GETCONTIGUOUS( pyArr ) ) );
+            contiguousArr = reinterpret_cast<PyArrayObject *>( contiguousOwner.get() );
+        }
+
+        auto * data = reinterpret_cast<const int64_t *>( PyArray_DATA( contiguousArr ) );
+
+        // Detect numpy datetime unit and compute scaling to nanoseconds
+        auto unit = datetimeUnitFromDescr( PyArray_DESCR( pyArr ) );
+        int64_t scaling = scalingFromNumpyDtUnit( unit );
+
+        if( scaling == 1 )
+        {
+            // Already nanoseconds — bulk append
+            ARROW_OK_OR_THROW_WRITER(
+                m_valueBuilder -> AppendValues( data, static_cast<int64_t>( len ) ),
+                "Failed to append temporal list elements" );
+        }
+        else
+        {
+            ARROW_OK_OR_THROW_WRITER(
+                m_valueBuilder -> Reserve( static_cast<int64_t>( len ) ),
+                "Failed to reserve temporal value builder" );
+            for( npy_intp i = 0; i < len; ++i )
+                m_valueBuilder -> UnsafeAppend( data[i] * scaling );
+        }
+    }
+
+private:
+    std::shared_ptr<ArrowBuilderT>          m_valueBuilder;
+    std::shared_ptr<::arrow::ListBuilder>   m_listBuilder;
+};
+
+// Date writer: numpy datetime64 → Arrow List<Date32> (nanoseconds → days since epoch)
+class DateListWriter final : public csp::adapters::arrow::FieldWriter
+{
+public:
+    DateListWriter( const std::string & columnName, const StructFieldPtr & structField )
+        : FieldWriter( { columnName }, { ::arrow::list( ::arrow::date32() ) }, structField )
+    {
+        m_valueBuilder = std::make_shared<::arrow::Date32Builder>();
+        m_listBuilder = std::make_shared<::arrow::ListBuilder>( ::arrow::default_memory_pool(), m_valueBuilder );
+    }
+
+    void reserve( int64_t numRows ) override
+    {
+        ARROW_OK_OR_THROW_WRITER( m_listBuilder -> Reserve( numRows ), "Failed to reserve date list builder" );
+    }
+
+    void writeNull() override
+    {
+        ARROW_OK_OR_THROW_WRITER( m_listBuilder -> AppendNull(), "Failed to append null date list" );
+    }
+
+    std::vector<std::shared_ptr<::arrow::Array>> finish() override
+    {
+        std::shared_ptr<::arrow::Array> arr;
+        ARROW_OK_OR_THROW_WRITER( m_listBuilder -> Finish( &arr ), "Failed to finish date list array" );
+        return { arr };
+    }
+
+protected:
+    void doWrite( const Struct * s ) override
+    {
+        ARROW_OK_OR_THROW_WRITER( m_listBuilder -> Append(), "Failed to start date list" );
+        auto & dgt = m_field -> value<DialectGenericType>( s );
+        auto * pyArr = reinterpret_cast<PyArrayObject *>( csp::python::toPythonBorrowed( dgt ) );
+        npy_intp len = PyArray_SIZE( pyArr );
+
+        // PyArray_DATA requires C-contiguous layout for correct bulk reads.
+        PyArrayObject * contiguousArr = pyArr;
+        PyObjectPtr contiguousOwner;
+        if( !PyArray_IS_C_CONTIGUOUS( pyArr ) )
+        {
+            contiguousOwner = PyObjectPtr::own(
+                reinterpret_cast<PyObject *>( PyArray_GETCONTIGUOUS( pyArr ) ) );
+            contiguousArr = reinterpret_cast<PyArrayObject *>( contiguousOwner.get() );
+        }
+
+        auto * data = reinterpret_cast<const int64_t *>( PyArray_DATA( contiguousArr ) );
+
+        // Detect numpy datetime unit and compute scaling to nanoseconds
+        auto unit = datetimeUnitFromDescr( PyArray_DESCR( pyArr ) );
+        int64_t scaling = scalingFromNumpyDtUnit( unit );
+
+        ARROW_OK_OR_THROW_WRITER(
+            m_valueBuilder -> Reserve( static_cast<int64_t>( len ) ),
+            "Failed to reserve date value builder" );
+        for( npy_intp i = 0; i < len; ++i )
+            m_valueBuilder -> UnsafeAppend( static_cast<int32_t>( data[i] * scaling / csp::NANOS_PER_DAY ) );
+    }
+
+private:
+    std::shared_ptr<::arrow::Date32Builder>  m_valueBuilder;
+    std::shared_ptr<::arrow::ListBuilder>    m_listBuilder;
+};
+
 // --- Dispatch by npy_type ---
 
 inline std::unique_ptr<csp::adapters::arrow::FieldWriter> dispatchListWriter(
@@ -193,6 +338,12 @@ inline std::unique_ptr<csp::adapters::arrow::FieldWriter> dispatchListWriter(
             return std::make_unique<NativeListWriter<bool, ::arrow::BooleanBuilder>>( columnName, structField, ::arrow::boolean() );
         case NPY_UNICODE:
             return std::make_unique<StringListWriter>( columnName, structField );
+        case NPY_DATETIME:
+            return std::make_unique<TemporalListWriter<::arrow::TimestampBuilder>>(
+                columnName, structField, std::make_shared<::arrow::TimestampType>( ::arrow::TimeUnit::NANO, "UTC" ) );
+        case NPY_TIMEDELTA:
+            return std::make_unique<TemporalListWriter<::arrow::DurationBuilder>>(
+                columnName, structField, std::make_shared<::arrow::DurationType>( ::arrow::TimeUnit::NANO ) );
         default:
             CSP_THROW( TypeError, "Unsupported numpy type " << npyType << " for list column '" << columnName << "'" );
     }
diff --git a/csp/tests/adapters/test_arrow_record_batches.py b/csp/tests/adapters/test_arrow_record_batches.py
