mirror of
https://github.com/OPM/ResInsight.git
synced 2026-10-25 14:40:26 -05:00
Magne Sjaastad
ca5465a660
Adapts the code to handle different API versions of the Arrow library for reading Parquet files. It uses preprocessor directives to select the correct API based on the `ARROW_VERSION_MAJOR`. This ensures compatibility with both older and newer versions of the Arrow library.
121 lines
4.1 KiB
C++
121 lines
4.1 KiB
C++
#include "gtest/gtest.h"
|
|
|
|
#include "RiaTestDataDirectory.h"
|
|
|
|
#include "RifArrowTools.h"
|
|
|
|
#undef signals
|
|
#include <arrow/array/builder_primitive.h>
|
|
#include <arrow/csv/api.h>
|
|
#include <arrow/io/api.h>
|
|
#include <arrow/scalar.h>
|
|
#include <parquet/arrow/reader.h>
|
|
#define signals Q_SIGNALS
|
|
|
|
#include <QDir>
|
|
#include <QString>
|
|
|
|
TEST( RifParquetReaderTest, ReadValidFile )
|
|
{
|
|
QDir baseFolder( TEST_DATA_DIR );
|
|
|
|
QString filename( "RifParquetReader/example.parquet" );
|
|
QString filePath = baseFolder.absoluteFilePath( filename );
|
|
EXPECT_TRUE( QFile::exists( filePath ) );
|
|
|
|
arrow::MemoryPool* pool = arrow::default_memory_pool();
|
|
|
|
auto openResult = arrow::io::ReadableFile::Open( filePath.toStdString().c_str() );
|
|
EXPECT_TRUE( openResult.ok() );
|
|
|
|
std::shared_ptr<arrow::io::RandomAccessFile> input = std::move( openResult ).ValueOrDie();
|
|
|
|
// Open Parquet file reader
|
|
#if ARROW_VERSION_MAJOR >= 20
|
|
// New API: OpenFile returns arrow::Result
|
|
auto readerResult = parquet::arrow::OpenFile( input, pool );
|
|
EXPECT_TRUE( readerResult.ok() );
|
|
std::unique_ptr<parquet::arrow::FileReader> arrow_reader = std::move( readerResult ).ValueOrDie();
|
|
#else
|
|
// Old API: OpenFile takes output parameter
|
|
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
|
|
EXPECT_TRUE( parquet::arrow::OpenFile( input, pool, &arrow_reader ).ok() );
|
|
#endif
|
|
|
|
// Read entire file as a single Arrow table
|
|
std::shared_ptr<arrow::Table> table;
|
|
EXPECT_TRUE( arrow_reader->ReadTable( &table ).ok() );
|
|
|
|
// Expect one column named "col1"
|
|
EXPECT_EQ( table->columns().size(), 1 );
|
|
std::vector<std::string> expectedNames = { "col1" };
|
|
EXPECT_EQ( table->ColumnNames(), expectedNames );
|
|
|
|
// Expected the column to contain 100 int64 [0, 99]
|
|
auto columnData = table->column( 0 );
|
|
int expectedSize = 100;
|
|
EXPECT_EQ( columnData->length(), expectedSize );
|
|
for ( int i = 0; i < expectedSize; i++ )
|
|
{
|
|
std::shared_ptr<arrow::Scalar> scalar = columnData->GetScalar( i ).ValueOrDie();
|
|
std::shared_ptr<arrow::Int64Scalar> intScalar = std::dynamic_pointer_cast<arrow::Int64Scalar>( scalar );
|
|
EXPECT_TRUE( scalar->Equals( arrow::Int64Scalar( i ) ) );
|
|
}
|
|
}
|
|
|
|
TEST( RifParquetReaderTest, ConvertIntChunkedArrays )
|
|
{
|
|
arrow::Status status;
|
|
|
|
arrow::Int32Builder int_builder;
|
|
status = int_builder.Append( 1 );
|
|
status = int_builder.Append( 2 );
|
|
status = int_builder.Append( 3 );
|
|
|
|
std::shared_ptr<arrow::Array> int_array;
|
|
status = int_builder.Finish( &int_array );
|
|
|
|
auto int_chunked_array = std::make_shared<arrow::ChunkedArray>( int_array );
|
|
|
|
{
|
|
auto columnVector = RifArrowTools::chunkedArrayToVector<arrow::FloatArray, double>( int_chunked_array );
|
|
EXPECT_EQ( columnVector.size(), 3 );
|
|
}
|
|
{
|
|
auto columnVector = RifArrowTools::chunkedArrayToVector<arrow::Int32Array, int>( int_chunked_array );
|
|
EXPECT_EQ( columnVector.size(), 3 );
|
|
}
|
|
}
|
|
|
|
TEST( RifParquetReaderTest, ConvertFloatChunkedArrays )
|
|
{
|
|
arrow::Status status;
|
|
|
|
// Create an Arrow double array
|
|
std::vector<double> values = { 1.0, 2.0, 3.0, 4.0 };
|
|
std::shared_ptr<arrow::Array> array;
|
|
arrow::DoubleBuilder builder;
|
|
status = builder.AppendValues( values );
|
|
status = builder.Finish( &array );
|
|
|
|
// Create a chunked array from the Arrow array
|
|
std::shared_ptr<arrow::ChunkedArray> chunkedArray = std::make_shared<arrow::ChunkedArray>( array );
|
|
|
|
// Call the function under test
|
|
auto resultVector = RifArrowTools::chunkedArrayToVector<arrow::DoubleArray, double>( chunkedArray );
|
|
|
|
// Assert that the returned vector contains the expected values
|
|
ASSERT_EQ( resultVector.size(), values.size() );
|
|
for ( size_t i = 0; i < values.size(); ++i )
|
|
{
|
|
EXPECT_DOUBLE_EQ( resultVector[i], values[i] );
|
|
}
|
|
|
|
auto floatVector = RifArrowTools::chunkedArrayToVector<arrow::DoubleArray, float>( chunkedArray );
|
|
ASSERT_EQ( floatVector.size(), values.size() );
|
|
for ( size_t i = 0; i < values.size(); ++i )
|
|
{
|
|
EXPECT_DOUBLE_EQ( floatVector[i], values[i] );
|
|
}
|
|
}
|