Skip to content

Commit 04445fc

Browse files
arhamchopraarhamchopra
committed
Address comments
Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>
1 parent 9bf2603 commit 04445fc

13 files changed

Lines changed: 213 additions & 204 deletions

‎cpp/csp/adapters/parquet/CMakeLists.txt‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ set(PARQUET_HEADER_FILES
88
ParquetOutputFilenameAdapter.h
99
ParquetStatusUtils.h
1010
ParquetWriter.h
11-
RecordBatchFileSink.h
1211
RecordBatchSink.h
1312
RecordBatchStreamSource.h
1413
)
@@ -21,7 +20,7 @@ set(PARQUET_SOURCE_FILES
2120
ParquetOutputAdapterManager.cpp
2221
ParquetOutputFilenameAdapter.cpp
2322
ParquetWriter.cpp
24-
RecordBatchFileSink.cpp
23+
RecordBatchSink.cpp
2524
${PARQUET_HEADER_FILES}
2625
)
2726

‎cpp/csp/adapters/parquet/ParquetDictBasketOutputWriter.cpp‎

Lines changed: 24 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,11 @@ void ParquetDictBasketOutputWriter::start()
2121
m_indexSchema = ::arrow::schema( { ::arrow::field(
2222
m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 ) -> getColumnName(),
2323
m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 ) -> getDataType() ) }, getFileMetaData() );
24-
if( m_indexSink.onStart )
25-
m_indexSink.onStart( m_indexSchema );
24+
if( m_indexSink )
25+
m_indexSink -> onStart( m_indexSchema );
2626
auto & fileName = m_adapterMgr.getFileName();
27-
if( !fileName.empty() && m_indexSink.onFileChange )
28-
m_indexSink.onFileChange( fileName );
27+
if( !fileName.empty() && m_indexSink )
28+
m_indexSink -> onFileChange( fileName );
2929
}
3030

3131
void ParquetDictBasketOutputWriter::stop()
@@ -40,16 +40,9 @@ void ParquetDictBasketOutputWriter::stop()
4040
catch( ... ) { if( !firstError ) firstError = std::current_exception(); }
4141
};
4242

43-
auto && indexBuilder = m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 );
44-
if( indexBuilder -> length() > 0 && m_indexSink.onBatch )
45-
guard( [&]
46-
{
47-
auto arr = indexBuilder -> buildArray();
48-
auto rb = ::arrow::RecordBatch::Make( m_indexSchema, arr -> length(), { arr } );
49-
m_indexSink.onBatch( rb );
50-
} );
51-
if( m_indexSink.onStop )
52-
guard( [&] { m_indexSink.onStop(); } );
43+
guard( [&] { flushIndexBatch(); } );
44+
if( m_indexSink )
45+
guard( [&] { m_indexSink -> onStop(); } );
5346

5447
guard( [&] { ParquetWriter::stop(); } );
5548

@@ -73,12 +66,8 @@ void ParquetDictBasketOutputWriter::onEndCycle()
7366
m_cycleIndexOutputAdapter -> writeValue<std::uint16_t>( m_nextCycleIndex );
7467
auto && indexBuilder = m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 );
7568
indexBuilder -> handleRowFinished();
76-
if( indexBuilder -> length() >= getChunkSize() && m_indexSink.onBatch )
77-
{
78-
auto arr = indexBuilder -> buildArray();
79-
auto rb = ::arrow::RecordBatch::Make( m_indexSchema, arr -> length(), { arr } );
80-
m_indexSink.onBatch( rb );
81-
}
69+
if( indexBuilder -> length() >= getChunkSize() )
70+
flushIndexBatch();
8271
m_nextCycleIndex = 0;
8372
}
8473
else
@@ -101,22 +90,26 @@ void ParquetDictBasketOutputWriter::onFileNameChange( const std::string &fileNam
10190

10291
guard( [&] { ParquetWriter::onFileNameChange( fileName ); } );
10392

104-
// Flush any pending index data
105-
auto && indexBuilder = m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 );
106-
if( indexBuilder -> length() > 0 && m_indexSink.onBatch )
107-
guard( [&]
108-
{
109-
auto arr = indexBuilder -> buildArray();
110-
auto rb = ::arrow::RecordBatch::Make( m_indexSchema, arr -> length(), { arr } );
111-
m_indexSink.onBatch( rb );
112-
} );
113-
if( m_indexSink.onFileChange )
114-
guard( [&] { m_indexSink.onFileChange( fileName ); } );
93+
// Flush any pending index data, then rotate the index file.
94+
guard( [&] { flushIndexBatch(); } );
95+
if( m_indexSink )
96+
guard( [&] { m_indexSink -> onFileChange( fileName ); } );
11597

11698
if( firstError )
11799
std::rethrow_exception( firstError );
118100
}
119101

102+
void ParquetDictBasketOutputWriter::flushIndexBatch()
103+
{
104+
auto && indexBuilder = m_cycleIndexOutputAdapter -> getColumnArrayBuilder( 0 );
105+
if( indexBuilder -> length() > 0 && m_indexSink )
106+
{
107+
auto arr = indexBuilder -> buildArray();
108+
auto rb = ::arrow::RecordBatch::Make( m_indexSchema, arr -> length(), { arr } );
109+
m_indexSink -> onBatch( rb );
110+
}
111+
}
112+
120113
SingleColumnParquetOutputHandler *ParquetDictBasketOutputWriter::createScalarOutputHandler( CspTypePtr type, const std::string &name )
121114
{
122115
m_allHandlers.push_back( std::make_unique<SingleColumnParquetOutputHandler>( m_engine, *this, type, name ) );

‎cpp/csp/adapters/parquet/ParquetDictBasketOutputWriter.h‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ class ParquetDictBasketOutputWriter : public ParquetWriter
1616
void start() override;
1717
void stop() override;
1818

19-
void setIndexSink( RecordBatchSink sink ) { m_indexSink = std::move( sink ); }
19+
void setIndexSink( RecordBatchSink sink ) { m_indexSink.emplace( std::move( sink ) ); }
2020

2121
virtual void writeValue( const std::string &valueKey, const TimeSeriesProvider *ts );
2222

@@ -27,11 +27,15 @@ class ParquetDictBasketOutputWriter : public ParquetWriter
2727
StructParquetOutputHandler *createStructOutputHandler( CspTypePtr type, const DictionaryPtr &fieldMap ) override;
2828

2929
private:
30+
// Build a RecordBatch from any pending index rows and hand it to the index sink. Shared by
31+
// onEndCycle (chunk full), onFileNameChange (rotation), and stop (final flush).
32+
void flushIndexBatch();
33+
3034
SingleColumnParquetOutputHandler *m_symbolOutputAdapter;
3135
SingleColumnParquetOutputHandler *m_cycleIndexOutputAdapter;
3236
std::uint16_t m_nextCycleIndex;
3337
std::vector<std::unique_ptr<ParquetOutputHandler>> m_allHandlers;
34-
RecordBatchSink m_indexSink;
38+
std::optional<RecordBatchSink> m_indexSink;
3539
std::shared_ptr<::arrow::Schema> m_indexSchema;
3640
};
3741

‎cpp/csp/adapters/parquet/ParquetOutputAdapter.cpp‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,15 @@ uint32_t ParquetOutputHandler::getChunkSize() const
1717

1818
SingleColumnParquetOutputHandler::SingleColumnParquetOutputHandler( Engine *engine, ParquetWriter &parquetWriter, CspTypePtr &type,
1919
std::string columnName )
20-
: ParquetOutputHandler( parquetWriter, type )
20+
: ParquetOutputHandler( parquetWriter )
2121
{
2222
bool isBytes = false;
23-
CspTypePtr effectiveType = type;
24-
2523
if( type -> type() == CspType::TypeTraits::STRING )
2624
{
2725
isBytes = static_cast<const CspStringType &>( *type ).isBytes();
2826
}
2927

30-
auto arrowBuilder = createArrowBackedArrayBuilder( columnName, getChunkSize(), effectiveType, isBytes );
28+
auto arrowBuilder = createArrowBackedArrayBuilder( columnName, getChunkSize(), type, isBytes );
3129
auto * scratch = arrowBuilder -> scratch();
3230
auto field = arrowBuilder -> scratchField();
3331
m_columnArrayBuilder = arrowBuilder;
@@ -40,7 +38,7 @@ SingleColumnParquetOutputHandler::SingleColumnParquetOutputHandler( Engine *engi
4038
},
4139
[&]()
4240
{
43-
CSP_THROW( TypeError, "Writing of " << m_type -> type().asString() << " to parquet is not supported" );
41+
CSP_THROW( TypeError, "Writing of " << type -> type().asString() << " to parquet is not supported" );
4442
} );
4543
}
4644

@@ -64,7 +62,7 @@ template void SingleColumnParquetOutputHandler::writeValue<std::uint16_t, void>(
6462

6563
ListColumnParquetOutputHandler::ListColumnParquetOutputHandler( Engine *engine, ParquetWriter &parquetWriter, CspTypePtr &elemType,
6664
const std::string &columnName )
67-
: ParquetOutputHandler( parquetWriter, CspType::DIALECT_GENERIC() )
65+
: ParquetOutputHandler( parquetWriter )
6866
{
6967
auto [ valueBuilder, writeItemsFn ] = csp::adapters::arrow::createListFieldWriter( elemType );
7068
m_columnArrayBuilder = std::make_shared<ListColumnArrayBuilder>(
@@ -86,7 +84,7 @@ void ListColumnParquetOutputAdapter::executeImpl()
8684

8785
StructParquetOutputHandler::StructParquetOutputHandler( Engine *engine, ParquetWriter &parquetWriter, CspTypePtr &type,
8886
DictionaryPtr fieldMap )
89-
: ParquetOutputHandler( parquetWriter, type )
87+
: ParquetOutputHandler( parquetWriter )
9088
{
9189
auto structMetaPtr = std::static_pointer_cast<const CspStructType>( type ) -> meta().get();
9290

‎cpp/csp/adapters/parquet/ParquetOutputAdapter.h‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ class ParquetWriter;
2727
class ParquetOutputHandler
2828
{
2929
public:
30-
ParquetOutputHandler( ParquetWriter &parquetWriter, CspTypePtr &type )
31-
: m_type( type ), m_parquetWriter( parquetWriter )
30+
ParquetOutputHandler( ParquetWriter &parquetWriter )
31+
: m_parquetWriter( parquetWriter )
3232
{
3333
}
3434

@@ -41,7 +41,6 @@ class ParquetOutputHandler
4141
virtual void writeValueFromTs( const TimeSeriesProvider *input ) = 0;
4242

4343
protected:
44-
CspTypePtr &m_type;
4544
ParquetWriter &m_parquetWriter;
4645
};
4746

‎cpp/csp/adapters/parquet/ParquetOutputAdapterManager.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
#include <csp/core/Generator.h>
77
#include <csp/engine/AdapterManager.h>
88
#include <csp/engine/Dictionary.h>
9+
#include <functional>
10+
#include <memory>
911
#include <string>
1012
#include <unordered_map>
1113

‎cpp/csp/adapters/parquet/ParquetWriter.cpp‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -166,12 +166,18 @@ void ParquetWriter::start()
166166
CSP_THROW( ValueError, "parquet column metadata has unmapped column: '" << m_columnMetaData.begin() -> first << "'" );
167167

168168
m_schema = ::arrow::schema( arrowFields, m_fileMetaData );
169-
if( m_sink.onStart )
170-
m_sink.onStart( m_schema );
171-
auto & fileName = m_adapterMgr.getFileName();
172-
m_fileOpen = !fileName.empty();
173-
if( m_fileOpen && m_sink.onFileChange )
174-
m_sink.onFileChange( fileName );
169+
// Hand the schema to the sink, then open the first file (if any). m_fileOpen reflects whether a
170+
// file is actually open after onFileChange -- an empty filename (no file / paused) leaves it false.
171+
const std::string & fileName = m_adapterMgr.getFileName();
172+
if( m_sink )
173+
{
174+
m_sink -> onStart( m_schema );
175+
m_fileOpen = m_sink -> onFileChange( fileName );
176+
}
177+
else
178+
{
179+
m_fileOpen = !fileName.empty();
180+
}
175181
}
176182

177183
void ParquetWriter::stop()
@@ -185,9 +191,9 @@ void ParquetWriter::stop()
185191
try { flushBatch(); }
186192
catch( ... ) { firstError = std::current_exception(); }
187193
}
188-
if( m_sink.onStop )
194+
if( m_sink )
189195
{
190-
try { m_sink.onStop(); }
196+
try { m_sink -> onStop(); }
191197
catch( ... ) { if( !firstError ) firstError = std::current_exception(); }
192198
}
193199
m_fileOpen = false;
@@ -240,9 +246,10 @@ void ParquetWriter::onEndCycle()
240246
void ParquetWriter::onFileNameChange( const std::string &fileName )
241247
{
242248
flushBatch();
243-
if( m_sink.onFileChange )
244-
m_sink.onFileChange( fileName );
245-
m_fileOpen = !fileName.empty();
249+
if( m_sink )
250+
m_fileOpen = m_sink -> onFileChange( fileName );
251+
else
252+
m_fileOpen = !fileName.empty();
246253
}
247254

248255
SingleColumnParquetOutputHandler *ParquetWriter::createScalarOutputHandler( CspTypePtr type, const std::string &name )
@@ -286,8 +293,8 @@ void ParquetWriter::flushBatch()
286293
{
287294
if( !isFileOpen() ) [[unlikely]]
288295
CSP_THROW( csp::RuntimeException, "Trying to write to parquet/arrow file, when no file name was provided" );
289-
if( m_sink.onBatch )
290-
m_sink.onBatch( buildRecordBatch() );
296+
if( m_sink )
297+
m_sink -> onBatch( buildRecordBatch() );
291298
m_curChunkSize = 0;
292299
}
293300
}

‎cpp/csp/adapters/parquet/ParquetWriter.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include <csp/engine/Dictionary.h>
77
#include <csp/engine/Engine.h>
88
#include <arrow/record_batch.h>
9+
#include <memory>
910
#include <string>
1011
#include <optional>
1112
#include <unordered_map>
@@ -34,7 +35,7 @@ class ParquetWriter : public EndCycleListener
3435

3536
~ParquetWriter();
3637

37-
void setSink( RecordBatchSink sink ) { m_sink = std::move( sink ); }
38+
void setSink( RecordBatchSink sink ) { m_sink.emplace( std::move( sink ) ); }
3839

3940
// Dict-basket writers adopt the main writer's file-level metadata and any column metadata that
4041
// targets their own columns, so basket files carry the same metadata and the main writer does not
@@ -84,7 +85,7 @@ class ParquetWriter : public EndCycleListener
8485
Adapters m_adapters;
8586
PublishedColumnNames m_publishedColumnNames;
8687

87-
RecordBatchSink m_sink;
88+
std::optional<RecordBatchSink> m_sink;
8889
bool m_fileOpen = false;
8990

9091
std::shared_ptr<::arrow::Schema> m_schema;

‎cpp/csp/adapters/parquet/RecordBatchFileSink.h‎

Lines changed: 0 additions & 30 deletions
This file was deleted.

0 commit comments

Comments
 (0)