Parquet output: harden teardown, basket metadata, and tests

Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>

Author: arhamchopra

cpp/csp/adapters/parquet/ParquetDictBasketOutputWriter.cpp

@@ -1,6 +1,7 @@
 #include <csp/adapters/parquet/ParquetDictBasketOutputWriter.h>
 #include <csp/adapters/parquet/ParquetOutputAdapter.h>
 #include <arrow/record_batch.h>
+#include <exception>
 
 namespace csp::adapters::parquet
 {
@@ -19,7 +20,7 @@ void ParquetDictBasketOutputWriter::start()
     ParquetWriter::start();
     m_indexSchema = ::arrow::schema( { ::arrow::field(
         m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 ) -> getColumnName(),
-        m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 ) -> getDataType() ) } );
+        m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 ) -> getDataType() ) }, getFileMetaData() );
     if( m_indexSink.onStart )
         m_indexSink.onStart( m_indexSchema );
     auto & fileName = m_adapterMgr.getFileName();
@@ -29,17 +30,31 @@ void ParquetDictBasketOutputWriter::start()
 
 void ParquetDictBasketOutputWriter::stop()
 {
+    // Exception-safe teardown: a failure writing/closing the index sink must not skip the value+symbol
+    // writer's stop() (which flushes the final batch and writes its footer) -- otherwise an index-side
+    // I/O error would leave the main basket data file unreadable. Attempt every step, rethrow first.
+    std::exception_ptr firstError;
+    auto guard = [&firstError]( auto && fn )
+    {
+        try { fn(); }
+        catch( ... ) { if( !firstError ) firstError = std::current_exception(); }
+    };
+
     auto && indexBuilder = m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 );
     if( indexBuilder -> length() > 0 && m_indexSink.onBatch )
-    {
-        auto arr = indexBuilder -> buildArray();
-        auto rb  = ::arrow::RecordBatch::Make( m_indexSchema, arr -> length(), { arr } );
-        m_indexSink.onBatch( rb );
-    }
+        guard( [&]
+        {
+            auto arr = indexBuilder -> buildArray();
+            auto rb  = ::arrow::RecordBatch::Make( m_indexSchema, arr -> length(), { arr } );
+            m_indexSink.onBatch( rb );
+        } );
     if( m_indexSink.onStop )
-        m_indexSink.onStop();
+        guard( [&] { m_indexSink.onStop(); } );
 
-    ParquetWriter::stop();
+    guard( [&] { ParquetWriter::stop(); } );
+
+    if( firstError )
+        std::rethrow_exception( firstError );
 }
 
 void ParquetDictBasketOutputWriter::writeValue( const std::string &valueKey, const TimeSeriesProvider *ts )
@@ -74,17 +89,32 @@ void ParquetDictBasketOutputWriter::onEndCycle()
 
 void ParquetDictBasketOutputWriter::onFileNameChange( const std::string &fileName )
 {
-    ParquetWriter::onFileNameChange( fileName );
+    // Rotate the main (value+symbol) writer, then flush pending index data and rotate the index file.
+    // Each step is guarded so a failure in one cannot leave the index file un-rotated and
+    // desynchronised from the main file; the first error is rethrown after all steps are attempted.
+    std::exception_ptr firstError;
+    auto guard = [&firstError]( auto && fn )
+    {
+        try { fn(); }
+        catch( ... ) { if( !firstError ) firstError = std::current_exception(); }
+    };
+
+    guard( [&] { ParquetWriter::onFileNameChange( fileName ); } );
+
     // Flush any pending index data
     auto && indexBuilder = m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 );
     if( indexBuilder -> length() > 0 && m_indexSink.onBatch )
-    {
-        auto arr = indexBuilder -> buildArray();
-        auto rb  = ::arrow::RecordBatch::Make( m_indexSchema, arr -> length(), { arr } );
-        m_indexSink.onBatch( rb );
-    }
+        guard( [&]
+        {
+            auto arr = indexBuilder -> buildArray();
+            auto rb  = ::arrow::RecordBatch::Make( m_indexSchema, arr -> length(), { arr } );
+            m_indexSink.onBatch( rb );
+        } );
     if( m_indexSink.onFileChange )
-        m_indexSink.onFileChange( fileName );
+        guard( [&] { m_indexSink.onFileChange( fileName ); } );
+
+    if( firstError )
+        std::rethrow_exception( firstError );
 }
 
 SingleColumnParquetOutputHandler *ParquetDictBasketOutputWriter::createScalarOutputHandler( CspTypePtr type, const std::string &name )

cpp/csp/adapters/parquet/ParquetOutputAdapterManager.cpp

@@ -4,6 +4,7 @@
 #include <csp/adapters/parquet/ParquetDictBasketOutputWriter.h>
 #include <csp/adapters/parquet/ParquetWriter.h>
 #include <csp/engine/Dictionary.h>
+#include <exception>
 
 namespace csp::adapters::parquet
 {
@@ -33,15 +34,26 @@ void ParquetOutputAdapterManager::start( DateTime starttime, DateTime endtime )
 
 void ParquetOutputAdapterManager::stop()
 {
-    // Stop all writers (flush + close their files) before destroying any of them, so a basket
-    // writer can never observe a destroyed sibling during teardown.
-    m_parquetWriter -> stop();
+    // Stop every writer (flush + close its files) before destroying any of them, so a basket writer
+    // can never observe a destroyed sibling during teardown. A failure stopping one writer must not
+    // skip the others (that would leave their files without footers / leak fds): attempt them all and
+    // rethrow the first error.
+    std::exception_ptr firstError;
+    auto stopWriter = [&firstError]( auto *writer )
+    {
+        try { writer -> stop(); }
+        catch( ... ) { if( !firstError ) firstError = std::current_exception(); }
+    };
+    if( m_parquetWriter )
+        stopWriter( m_parquetWriter.get() );
     for( auto &&writer:m_dictBasketWriters )
     {
-        writer -> stop();
+        stopWriter( writer.get() );
     }
     m_parquetWriter = nullptr;
     m_dictBasketWriters.clear();
+    if( firstError )
+        std::rethrow_exception( firstError );
 }
 
 DateTime ParquetOutputAdapterManager::processNextSimTimeSlice( DateTime time )
@@ -86,6 +98,10 @@ ParquetOutputAdapterManager::createDictOutputBasketWriter( const char *columnNam
 
     // Provide data sink via factory if available
     auto * writer = m_dictBasketWriters.back().get();
+    // Share the main writer's file_metadata, and claim any column_metadata the user attached to this
+    // basket's columns, so the basket files carry that metadata and the main writer does not later
+    // reject those keys as unmapped columns.
+    writer -> adoptMetadataFrom( *m_parquetWriter );
     if( m_sinkFactory )
     {
         writer -> setSink( m_sinkFactory( columnName ) );

cpp/csp/adapters/parquet/ParquetWriter.cpp

@@ -5,6 +5,7 @@
 #include <arrow/array.h>
 #include <arrow/builder.h>
 #include <arrow/util/key_value_metadata.h>
+#include <exception>
 
 namespace csp::adapters::parquet
 {
@@ -175,15 +176,51 @@ void ParquetWriter::start()
 
 void ParquetWriter::stop()
 {
+    // Exception-safe teardown: a failure flushing the final batch must not skip onStop() (which
+    // writes the file footer/closes the stream), and a failure in onStop() must not be masked.
+    // Attempt both and rethrow the first error.
+    std::exception_ptr firstError;
     if( m_curChunkSize > 0 )
-        flushBatch();
+    {
+        try { flushBatch(); }
+        catch( ... ) { firstError = std::current_exception(); }
+    }
     if( m_sink.onStop )
-        m_sink.onStop();
+    {
+        try { m_sink.onStop(); }
+        catch( ... ) { if( !firstError ) firstError = std::current_exception(); }
+    }
     m_fileOpen = false;
+    if( firstError )
+        std::rethrow_exception( firstError );
+}
+
+void ParquetWriter::adoptMetadataFrom( ParquetWriter & source )
+{
+    // Used by dict-basket writers: take the main writer's file-level metadata, and move any column
+    // metadata the user attached to one of THIS writer's columns out of the source writer (so the
+    // basket files carry it and the main writer no longer rejects those keys as "unmapped").
+    if( !m_fileMetaData )
+        m_fileMetaData = source.m_fileMetaData;
+    for( auto && columnName : m_publishedColumnNames )
+    {
+        auto it = source.m_columnMetaData.find( columnName );
+        if( it != source.m_columnMetaData.end() )
+        {
+            m_columnMetaData[ columnName ] = it -> second;
+            source.m_columnMetaData.erase( it );
+        }
+    }
 }
 
 void ParquetWriter::onEndCycle()
 {
+    // NOTE: when no file is open -- e.g. a filename_provider ticked "" to "pause" output -- this body is
+    // skipped, so handleRowFinished() does NOT run and each scratch column's isSet bit is never cleared
+    // for that cycle. If a column ticks *while paused* and then does NOT tick on the cycle the file
+    // reopens, the stale paused-window value is emitted instead of null (the dropped tick "leaks" into
+    // the first row after resume). This matches the previous writer's behaviour. If dropping ticks while
+    // paused should instead null them out, clear each builder's scratch/isSet here in the else path.
     if( isFileOpen() ) [[likely]]
     {
         DateTime now;
@@ -236,8 +273,11 @@ std::shared_ptr<::arrow::RecordBatch> ParquetWriter::buildRecordBatch()
     {
         auto column = builder -> buildArray();
         // Every column must contain exactly the rows accumulated this chunk; a mismatch would build a
-        // corrupt RecordBatch (debug-only guard, compiled out in release).
-        CSP_ASSERT( column -> length() == static_cast<std::int64_t>( m_curChunkSize ) );
+        // corrupt RecordBatch. Enforced at runtime (not just in debug) so a future column-builder bug
+        // can never silently produce a malformed parquet/arrow file.
+        CSP_TRUE_OR_THROW_RUNTIME( column -> length() == static_cast<std::int64_t>( m_curChunkSize ),
+            "ParquetWriter column '" << builder -> getColumnName() << "' produced " << column -> length()
+            << " rows but the current chunk has " << m_curChunkSize );
         columns.push_back( std::move( column ) );
     }
     return ::arrow::RecordBatch::Make( m_schema, m_curChunkSize, std::move( columns ) );

cpp/csp/adapters/parquet/ParquetWriter.h

@@ -36,6 +36,12 @@ class ParquetWriter : public EndCycleListener
 
     void setSink( RecordBatchSink sink ) { m_sink = std::move( sink ); }
 
+    // Dict-basket writers adopt the main writer's file-level metadata and any column metadata that
+    // targets their own columns, so basket files carry the same metadata and the main writer does not
+    // reject those keys as "unmapped".
+    const std::shared_ptr<::arrow::KeyValueMetadata> & getFileMetaData() const { return m_fileMetaData; }
+    void adoptMetadataFrom( ParquetWriter & source );
+
     SingleColumnParquetOutputAdapter *getScalarOutputAdapter( CspTypePtr &type, const std::string &columnName );
     StructParquetOutputAdapter *getStructOutputAdapter( CspTypePtr &type, csp::DictionaryPtr fieldMap );
     ListColumnParquetOutputAdapter *getListOutputAdapter( CspTypePtr &elemType, const std::string &columnName );

cpp/csp/adapters/parquet/RecordBatchFileSink.cpp

@@ -8,7 +8,6 @@
 #include <arrow/table.h>
 #include <arrow/type.h>
 #include <arrow/util/compression.h>
-#include <arrow/util/config.h>
 #include <parquet/arrow/writer.h>
 
 #include <algorithm>
@@ -77,11 +76,10 @@ FileWriter makeParquetFileWriter( const std::string & path, const std::shared_pt
 
     ::parquet::WriterProperties::Builder props;
     props.compression( resolveCompression( compression ) );
-#if ARROW_VERSION_MAJOR >= 20
+    // csp builds against libparquet >= 16.0.0, where PARQUET_2_6 (the modern default, available since
+    // Arrow 6.0) is always present, so there is no need to gate on the Arrow version or fall back to
+    // the deprecated PARQUET_2_0.
     props.version( ::parquet::ParquetVersion::PARQUET_2_6 );
-#else
-    props.version( ::parquet::ParquetVersion::PARQUET_2_0 );
-#endif
     ::parquet::ArrowWriterProperties::Builder arrowProps;
     arrowProps.store_schema(); // preserve arrow (file/column) schema metadata in the parquet file
 
@@ -128,9 +126,13 @@ FileWriter makeIpcFileWriter( const std::string & path, const std::shared_ptr<::
     std::shared_ptr<::arrow::ipc::RecordBatchWriter> writer = result.MoveValueUnsafe();
 
     return FileWriter{
-        [writer]( const std::shared_ptr<::arrow::RecordBatch> & rb )
+        [writer, compression]( const std::shared_ptr<::arrow::RecordBatch> & rb )
         {
-            STATUS_OK_OR_THROW_RUNTIME( writer -> WriteRecordBatch( *rb ), "Failed to write arrow record batch" );
+            // Arrow validates which codecs IPC allows here and appends its own reason (e.g. "Only
+            // LZ4_FRAME and ZSTD compression allowed for IPC"); surface which codec was requested
+            // rather than a bare "failed to write". The allowed set is Arrow's to define, not ours.
+            STATUS_OK_OR_THROW_RUNTIME( writer -> WriteRecordBatch( *rb ),
+                "Failed to write Arrow IPC record batch with compression '" << compression << "'" );
         },
         [writer, stream]()
         {
@@ -154,11 +156,26 @@ FileWriter makeSplitWriter( const std::string & dir, const std::shared_ptr<::arr
     subWriters -> reserve( schema -> num_fields() );
     colSchemas -> reserve( schema -> num_fields() );
 
-    for( int i = 0; i < schema -> num_fields(); ++i )
+    try
+    {
+        for( int i = 0; i < schema -> num_fields(); ++i )
+        {
+            auto colSchema = ::arrow::schema( { schema -> field( i ) }, schema -> metadata() );
+            colSchemas -> push_back( colSchema );
+            subWriters -> push_back( perColumn( dir + "/" + schema -> field( i ) -> name() + extension, colSchema ) );
+        }
+    }
+    catch( ... )
     {
-        auto colSchema = ::arrow::schema( { schema -> field( i ) }, schema -> metadata() );
-        colSchemas -> push_back( colSchema );
-        subWriters -> push_back( perColumn( dir + "/" + schema -> field( i ) -> name() + extension, colSchema ) );
+        // A later column failed to open (e.g. allow_overwrite=false and that file already exists).
+        // Close the sub-writers already opened so we don't leak file descriptors or leave half-open
+        // files behind, then rethrow the original error.
+        for( auto & w : *subWriters )
+        {
+            try { w.close(); }
+            catch( ... ) {}
+        }
+        throw;
     }
 
     return FileWriter{
@@ -230,8 +247,12 @@ RecordBatchSink makeFileSink( bool writeArrowBinary, bool splitColumns,
     sink.onStart      = [schemaHolder]( const std::shared_ptr<::arrow::Schema> & schema ) { *schemaHolder = schema; };
     sink.onBatch      = [current]( const std::shared_ptr<::arrow::RecordBatch> & rb )
     {
-        if( current -> has_value() )
-            ( *current ) -> writeBatch( rb );
+        // A batch must only ever arrive while a file is open (the writer guards this via isFileOpen()
+        // before flushing). If the writer and sink ever desync, fail loudly instead of silently
+        // dropping the batch's rows.
+        CSP_TRUE_OR_THROW_RUNTIME( current -> has_value(),
+            "RecordBatchFileSink received a record batch with no open output file" );
+        ( *current ) -> writeBatch( rb );
     };
     sink.onFileChange = [factory, schemaHolder, current, currentPath, closeCurrent]( const std::string & path )
     {

csp/adapters/output_adapters/parquet.py

@@ -77,6 +77,8 @@ def __init__(
         and at shutdown). It runs synchronously on the engine thread with the GIL held, so a slow visitor (e.g.
         uploading to remote storage) blocks the engine for its duration -- offload heavy work to a background
         queue/thread if that matters. An exception raised by the visitor propagates out of csp.run.
+        Note: file_visitor is NOT invoked for dict-basket files (the value/symbol/index files written by
+        publish_dict_basket); it only covers the main writer's files (and any regular published columns).
         """
         super().__init__()
         config = ParquetOutputConfig() if config is None else config.copy()