Skip to content

Add arrow adapter to stream record batches zero-copy into csp #519

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 7 commits into
base: main
Choose a base branch
from
Draft

Add arrow adapter to stream record batches zero-copy into csp #519

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
f70df59
Add arrow adapter to stream record batches zero-copy into csp
arhamchopra Apr 4, 2025
b14e2c8
Add node to convert record batches into csp.Structs and fix arrow ada…
arhamchopra Apr 11, 2025
01efd9f
Bugfix in ParquetReader to read structs correctly for InMemoryTablePa…
arhamchopra Apr 11, 2025
abb25ed
Add node to convert csp.Structs into record batches and cleanup
arhamchopra Apr 12, 2025
e5523eb
Bugfix in PyCppNode.h move #endif to the end
arhamchopra Apr 30, 2025
92e5ff9
Support NumpyNDArray, Numpy1DArray in record_batches_to_struct/struct…
arhamchopra Apr 30, 2025
6297b23
Fix compilation errors
arhamchopra Apr 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions cpp/csp/adapters/parquet/ParquetDictBasketOutputWriter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,11 @@ void ParquetDictBasketOutputWriter::start()
m_indexFileWriterContainer = std::make_unique<MultipleFileWriterWrapperContainer>(
arrow::schema( { arrow::field( m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 ) -> getColumnName(),
m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 ) -> getDataType() ) } ),
m_adapterMgr.isWriteArrowBinary() );
if( !m_adapterMgr.getFileName().empty() )
m_adapterMgr -> isWriteArrowBinary() );
if( !m_adapterMgr -> getFileName().empty() )
{
m_indexFileWriterContainer -> open( m_adapterMgr.getFileName(),
m_adapterMgr.getCompression(), m_adapterMgr.isAllowOverwrite() );
m_indexFileWriterContainer -> open( m_adapterMgr -> getFileName(),
m_adapterMgr -> getCompression(), m_adapterMgr -> isAllowOverwrite() );

}
}
Expand All @@ -45,7 +45,7 @@ void ParquetDictBasketOutputWriter::stop()

void ParquetDictBasketOutputWriter::writeValue( const std::string &valueKey, const TimeSeriesProvider *ts )
{
m_adapterMgr.scheduleEndCycle();
m_adapterMgr -> scheduleEndCycle();
m_symbolOutputAdapter -> writeValue<std::string, StringArrayBuilder>( valueKey );
ParquetWriter::onEndCycle();
++m_nextCycleIndex;
Expand Down Expand Up @@ -86,7 +86,7 @@ void ParquetDictBasketOutputWriter::onFileNameChange( const std::string &fileNam
if(!fileName.empty())
{
m_indexFileWriterContainer
-> open( fileName, m_adapterMgr.getCompression(), m_adapterMgr.isAllowOverwrite() );
-> open( fileName, m_adapterMgr -> getCompression(), m_adapterMgr -> isAllowOverwrite() );
}

}
Expand Down
16 changes: 13 additions & 3 deletions cpp/csp/adapters/parquet/ParquetOutputAdapter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,10 +127,10 @@ ListColumnParquetOutputHandler::ListColumnParquetOutputHandler( Engine *engine,
listWriterInterface ) )
{
m_valueHandler = std::make_unique<ValueHandler>(
[ this ]( const TimeSeriesProvider *input )
[ this ]( const DialectGenericType& input )
{
static_cast<ListColumnArrayBuilder *>(this -> m_columnArrayBuilder.get())
-> setValue( input -> lastValueTyped<DialectGenericType>() );
-> setValue( input );
} );
}

Expand Down Expand Up @@ -176,7 +176,7 @@ std::shared_ptr<::arrow::ArrayBuilder> ListColumnParquetOutputHandler::createVal

void ListColumnParquetOutputAdapter::executeImpl()
{
( *m_valueHandler )( input() );
( *m_valueHandler )( input() -> lastValueTyped<DialectGenericType>() );
m_parquetWriter.scheduleEndCycleEvent();
}

Expand All @@ -192,6 +192,16 @@ StructParquetOutputHandler::StructParquetOutputHandler( Engine *engine, ParquetW
}
}

void StructParquetOutputHandler::writeValueFromArgs( const StructPtr input )
{
const Struct *structData = input.get();

for( auto &&valueHandler: m_valueHandlers )
{
valueHandler( structData );
}
}

void StructParquetOutputHandler::writeValueFromTs( const TimeSeriesProvider *input )
{
const Struct *structData = input -> lastValueTyped<StructPtr>().get();
Expand Down
10 changes: 8 additions & 2 deletions cpp/csp/adapters/parquet/ParquetOutputAdapter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,18 +93,23 @@ class ListColumnParquetOutputHandler : public ParquetOutputHandler
static_cast<ColumnBuilderType *>(this -> m_columnArrayBuilder.get()) -> setValue( value );
}

void writeValueFromTs( const TimeSeriesProvider *input ) override final
void writeValueFromArgs( const DialectGenericType& input )
{
( *m_valueHandler )( input );
}

void writeValueFromTs( const TimeSeriesProvider *input ) override final
{
( *m_valueHandler )( input -> lastValueTyped<DialectGenericType>() );
}

private:
std::shared_ptr<arrow::ArrayBuilder> createValueBuilder( const CspTypePtr &elemType,
DialectGenericListWriterInterface::Ptr &listWriterInterface );


protected :
using ValueHandler = std::function<void( const TimeSeriesProvider * )>;
using ValueHandler = std::function<void( const DialectGenericType& )>;

std::unique_ptr<ValueHandler> m_valueHandler;
std::shared_ptr<ListColumnArrayBuilder> m_columnArrayBuilder;
Expand Down Expand Up @@ -151,6 +156,7 @@ class StructParquetOutputHandler : public ParquetOutputHandler
}

void writeValueFromTs( const TimeSeriesProvider *input ) override final;
void writeValueFromArgs( const StructPtr input );

private:
using ValueHandler = std::function<void( const Struct * )>;
Expand Down
19 changes: 15 additions & 4 deletions cpp/csp/adapters/parquet/ParquetReader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,9 +192,17 @@ void SingleTableParquetReader::setColumnAdaptersFromCurrentTable()
columnAdapter = createColumnAdapter( *this, *field, getCurFileOrTableName(), &getStructColumnMeta() );
auto &fieldInfo = fieldsInfo[ index ];

for( std::size_t i = 0; i < fieldInfo.m_width; ++i )
if( isArrowIPC() )
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is way to check if we are using the InMemoryTableParquetReader

{
// Needed for all memory tables
m_neededColumnIndices.push_back( index );
}
else
{
m_neededColumnIndices.push_back( fieldInfo.m_startColumnIndex + i );
for( std::size_t i = 0; i < fieldInfo.m_width; ++i )
{
m_neededColumnIndices.push_back( fieldInfo.m_startColumnIndex + i );
}
}
}
else
Expand Down Expand Up @@ -382,10 +390,13 @@ void SingleFileParquetReader::clear()
}

InMemoryTableParquetReader::InMemoryTableParquetReader( GeneratorPtr generatorPtr, std::vector<std::string> columns,
bool allowMissingColumns, std::optional<std::string> symbolColumnName )
bool allowMissingColumns, std::optional<std::string> symbolColumnName, bool call_init )
: SingleTableParquetReader( columns, true, allowMissingColumns, symbolColumnName ), m_generatorPtr( generatorPtr )
{
init();
if( call_init )
{
init();
}
}

bool InMemoryTableParquetReader::openNextFile()
Expand Down
17 changes: 12 additions & 5 deletions cpp/csp/adapters/parquet/ParquetReader.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -424,21 +424,28 @@ class SingleFileParquetReader final : public SingleTableParquetReader
bool m_allowMissingFiles;
};

class InMemoryTableParquetReader final : public SingleTableParquetReader
class InMemoryTableParquetReader : public SingleTableParquetReader
{
public:
using GeneratorPtr = csp::Generator<std::shared_ptr<arrow::Table>, csp::DateTime, csp::DateTime>::Ptr;

InMemoryTableParquetReader( GeneratorPtr generatorPtr, std::vector<std::string> columns,
bool allowMissingColumns,
std::optional<std::string> symbolColumnName = {} );
std::optional<std::string> symbolColumnName = {},
bool call_init = true);
std::string getCurFileOrTableName() const override{ return "IN_MEMORY_TABLE"; }

protected:
bool openNextFile() override;
bool readNextRowGroup() override;
virtual bool openNextFile() override;
virtual bool readNextRowGroup() override;
void setTable( std::shared_ptr<arrow::Table> table )
{
m_fullTable = table;
m_nextChunkIndex = 0;
m_curTable = nullptr;
}

void clear() override;
virtual void clear() override;

private:
GeneratorPtr m_generatorPtr;
Expand Down
18 changes: 18 additions & 0 deletions cpp/csp/adapters/parquet/ParquetReaderColumnAdapter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -749,6 +749,24 @@ void BaseListColumnAdapter<ArrowListArrayType, ValueArrayType, ValueType>::addSu
<< " in file " << m_parquetReader.getCurFileOrTableName() );
}

template< typename ArrowListArrayType, typename ValueArrayType, typename ValueType>
void BaseListColumnAdapter<ArrowListArrayType, ValueArrayType, ValueType>::addSubscriber( csp::adapters::utils::ValueDispatcher<const DialectGenericType&>::SubscriberType subscriber, std::optional<utils::Symbol> symbol, const DialectGenericListReaderInterface::Ptr &listReader )
{
CSP_TRUE_OR_THROW_RUNTIME( m_listReader == nullptr,
"Trying to subscribe list column in parquet reader more than once, this is not supported" );
CSP_TRUE_OR_THROW_RUNTIME( listReader != nullptr,
"Trying to subscribe list column in parquet reader with null listReader" );
m_listReader = std::dynamic_pointer_cast<TypedDialectGenericListReaderInterface<ValueType>>( listReader );
CSP_TRUE_OR_THROW_RUNTIME( m_listReader != nullptr,
"Subscribed to parquet column " << getColumnName() << " with type "
<< "NumpyArray[" << listReader -> getValueType() -> type().asString()
<< "] while "
<< " column type in file is NumpyArray["
<< getContainerValueType() -> type().asString() << "]"
<< " in file " << m_parquetReader.getCurFileOrTableName() );
m_dispatcher.addSubscriber( subscriber, symbol );
}

template< typename ArrowListArrayType, typename ValueArrayType, typename ValueType >
void NativeListColumnAdapter<ArrowListArrayType, ValueArrayType, ValueType>::readCurValue()
{
Expand Down
12 changes: 12 additions & 0 deletions cpp/csp/adapters/parquet/ParquetReaderColumnAdapter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ class ParquetColumnAdapter
virtual void addSubscriber( ManagedSimInputAdapter *inputAdapter, std::optional<utils::Symbol> symbol = {} ) = 0;
// NOTE: This API is only defined for ListType Column Adapters
virtual void addSubscriber( ManagedSimInputAdapter *inputAdapter, std::optional<utils::Symbol> symbol, const DialectGenericListReaderInterface::Ptr &listReader ) = 0;
// NOTE: This API is only used to add subscriber for ListType column adapters in cases where there is no ManagedSimInputAdapter
virtual void addSubscriber( csp::adapters::utils::ValueDispatcher<const DialectGenericType&>::SubscriberType subscriber, std::optional<utils::Symbol> symbol, const DialectGenericListReaderInterface::Ptr &listReader ) = 0;

virtual void dispatchValue( const utils::Symbol *symbol ) = 0;

Expand Down Expand Up @@ -124,6 +126,10 @@ class MissingColumnAdapter : public ParquetColumnAdapter

virtual void addSubscriber( ManagedSimInputAdapter *inputAdapter, std::optional<utils::Symbol> symbol = {} ) override {};
virtual void addSubscriber( ManagedSimInputAdapter *inputAdapter, std::optional<utils::Symbol> symbol, const DialectGenericListReaderInterface::Ptr &listReader ) override {};
virtual void addSubscriber( csp::adapters::utils::ValueDispatcher<const DialectGenericType&>::SubscriberType subscriber, std::optional<utils::Symbol> symbol, const DialectGenericListReaderInterface::Ptr &listReader ) override
{
CSP_THROW(TypeError, "Trying to add DIALECT_GENERIC subscriber on non container type");
}

virtual void dispatchValue( const utils::Symbol *symbol ) override {};

Expand Down Expand Up @@ -172,6 +178,10 @@ class BaseTypedColumnAdapter : public ParquetColumnAdapter
void addSubscriber( ManagedSimInputAdapter *inputAdapter, std::optional<utils::Symbol> symbol = {} ) override;
void addSubscriber( ManagedSimInputAdapter *inputAdapter,
std::optional<utils::Symbol> symbol, const DialectGenericListReaderInterface::Ptr &listReader ) override;
virtual void addSubscriber( csp::adapters::utils::ValueDispatcher<const DialectGenericType&>::SubscriberType subscriber, std::optional<utils::Symbol> symbol, const DialectGenericListReaderInterface::Ptr &listReader ) override
{
CSP_THROW(TypeError, "Trying to add DIALECT_GENERIC subscriber on non container type");
}
void dispatchValue( const utils::Symbol *symbol ) override;

void ensureType( CspType::Ptr cspType ) override;
Expand Down Expand Up @@ -288,9 +298,11 @@ class BaseListColumnAdapter : public BaseTypedColumnAdapter<DialectGenericType,
public:
using BaseTypedColumnAdapter<DialectGenericType, ArrowListArrayType>::BaseTypedColumnAdapter;
using BaseTypedColumnAdapter<DialectGenericType, ArrowListArrayType>::getColumnName;
using BaseTypedColumnAdapter<DialectGenericType, ArrowListArrayType>::m_dispatcher;
void addSubscriber( ManagedSimInputAdapter *inputAdapter, std::optional<utils::Symbol> symbol = {} ) override;
void addSubscriber( ManagedSimInputAdapter *inputAdapter,
std::optional<utils::Symbol> symbol, const DialectGenericListReaderInterface::Ptr &listReader ) override;
void addSubscriber( csp::adapters::utils::ValueDispatcher<const DialectGenericType&>::SubscriberType subscriber, std::optional<utils::Symbol> symbol, const DialectGenericListReaderInterface::Ptr &listReader ) override;
CspTypePtr getNativeCspType() const override {return nullptr;}
bool isListType() const override{ return true; };
CspTypePtr getContainerValueType() const override{ return CspType::fromCType<ValueType>::type(); }
Expand Down
30 changes: 17 additions & 13 deletions cpp/csp/adapters/parquet/ParquetWriter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,12 @@
namespace csp::adapters::parquet
{

ParquetWriter::ParquetWriter()
: m_adapterMgr( nullptr ), m_engine( nullptr ), m_curChunkSize( 0 ), m_writeTimestampColumn( std::optional<bool>{} )
{}

ParquetWriter::ParquetWriter( ParquetOutputAdapterManager *mgr, std::optional<bool> writeTimestampColumn )
: m_adapterMgr( *mgr ), m_engine( mgr -> engine() ), m_curChunkSize( 0 ), m_writeTimestampColumn( writeTimestampColumn )
: m_adapterMgr( mgr ), m_engine( mgr -> engine() ), m_curChunkSize( 0 ), m_writeTimestampColumn( writeTimestampColumn )
{}

ParquetWriter::ParquetWriter( ParquetOutputAdapterManager *mgr, const Dictionary & properties ) : ParquetWriter( mgr, std::optional<bool>{} )
Expand Down Expand Up @@ -128,20 +132,20 @@ PushInputAdapter *ParquetWriter::getStatusAdapter()
void ParquetWriter::start()
{
std::vector<std::shared_ptr<arrow::Field>> arrowFields;
if( !m_writeTimestampColumn.has_value() && !m_adapterMgr.getTimestampColumnName().empty() )
if( !m_writeTimestampColumn.has_value() && !m_adapterMgr -> getTimestampColumnName().empty() )
{
m_writeTimestampColumn = true;
m_columnBuilders.push_back( std::make_shared<DatetimeArrayBuilder>( m_adapterMgr.getTimestampColumnName(), getChunkSize() ) );
m_columnBuilders.push_back( std::make_shared<DatetimeArrayBuilder>( m_adapterMgr -> getTimestampColumnName(), getChunkSize() ) );
std::shared_ptr<arrow::KeyValueMetadata> colMetaData;
auto colMetaIt = m_columnMetaData.find( m_adapterMgr.getTimestampColumnName() );
auto colMetaIt = m_columnMetaData.find( m_adapterMgr -> getTimestampColumnName() );
if( colMetaIt != m_columnMetaData.end() )
{
colMetaData = colMetaIt -> second;
m_columnMetaData.erase( colMetaIt );
}

arrowFields.push_back(
arrow::field( m_adapterMgr.getTimestampColumnName(), m_columnBuilders.back() -> getDataType(), colMetaData ) );
arrow::field( m_adapterMgr -> getTimestampColumnName(), m_columnBuilders.back() -> getDataType(), colMetaData ) );
}
else
{
Expand Down Expand Up @@ -194,7 +198,7 @@ void ParquetWriter::onEndCycle()
if( m_writeTimestampColumn.value() )
{
// Set the timestamp value it's always the first
now = m_adapterMgr.rootEngine() -> now();
now = m_adapterMgr -> rootEngine() -> now();
static_cast<DatetimeArrayBuilder *>(m_columnBuilders[ 0 ].get()) -> setValue( now );
}
for( auto &&columnBuilder:m_columnBuilders )
Expand All @@ -221,7 +225,7 @@ void ParquetWriter::onFileNameChange( const std::string &fileName )
if( !fileName.empty() )
{
m_fileWriterWrapperContainer
-> open( fileName, m_adapterMgr.getCompression(), m_adapterMgr.isAllowOverwrite() );
-> open( fileName, m_adapterMgr -> getCompression(), m_adapterMgr -> isAllowOverwrite() );
}
}

Expand All @@ -245,20 +249,20 @@ StructParquetOutputHandler *ParquetWriter::createStructOutputHandler( CspTypePtr

void ParquetWriter::initFileWriterContainer( std::shared_ptr<arrow::Schema> schema )
{
if( m_adapterMgr.isSplitColumnsToFiles() )
if( m_adapterMgr -> isSplitColumnsToFiles() )
{
m_fileWriterWrapperContainer = std::make_unique<MultipleFileWriterWrapperContainer>( schema,
m_adapterMgr.isWriteArrowBinary() );
m_adapterMgr -> isWriteArrowBinary() );
}
else
{
m_fileWriterWrapperContainer = std::make_unique<SingleFileWriterWrapperContainer>( schema,
m_adapterMgr.isWriteArrowBinary() );
m_adapterMgr -> isWriteArrowBinary() );
}
if( !m_adapterMgr.getFileName().empty() )
if( !m_adapterMgr -> getFileName().empty() )
{
m_fileWriterWrapperContainer -> open( m_adapterMgr.getFileName(),
m_adapterMgr.getCompression(), m_adapterMgr.isAllowOverwrite() );
m_fileWriterWrapperContainer -> open( m_adapterMgr -> getFileName(),
m_adapterMgr -> getCompression(), m_adapterMgr -> isAllowOverwrite() );
}
}

Expand Down
7 changes: 4 additions & 3 deletions cpp/csp/adapters/parquet/ParquetWriter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ class FileWriterWrapperContainer;
class ParquetWriter : public EndCycleListener
{
public:
ParquetWriter();
ParquetWriter( ParquetOutputAdapterManager *mgr, std::optional<bool> writeTimestampColumn = {} );
ParquetWriter( ParquetOutputAdapterManager *mgr, const Dictionary & properties );

Expand All @@ -53,11 +54,11 @@ class ParquetWriter : public EndCycleListener

void onEndCycle() override;

std::uint32_t getChunkSize() const{ return m_adapterMgr.getBatchSize(); }
virtual std::uint32_t getChunkSize() const{ return m_adapterMgr -> getBatchSize(); }

virtual void scheduleEndCycleEvent()
{
m_adapterMgr.scheduleEndCycle();
m_adapterMgr -> scheduleEndCycle();
}

bool isFileOpen() const;
Expand All @@ -76,7 +77,7 @@ class ParquetWriter : public EndCycleListener
using Adapters = std::vector<ParquetOutputHandler *>;
using PublishedColumnNames = std::unordered_set<std::string>;

ParquetOutputAdapterManager &m_adapterMgr;
ParquetOutputAdapterManager *m_adapterMgr;
Engine *m_engine;
private:
Adapters m_adapters;
Expand Down
Loading
Loading