Fix tests arrow tests

alisw · Apr 20, 2020 · 582af8b · 582af8b
1 parent 6580d49
commit 582af8b
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/tree/dataframe/test/datasource_arrow.cxx b/tree/dataframe/test/datasource_arrow.cxx
@@ -30,6 +30,16 @@ std::shared_ptr<Schema> exampleSchema()
                   field("Married", arrow::boolean()), field("Babies", arrow::uint32())});
 }
 
+template <typename T>
+std::shared_ptr<T> makeColumn(std::shared_ptr<Field>, std::shared_ptr<arrow::Array> array) {
+  return std::make_shared<T>(field, array);
+}
+
+template <>
+std::shared_ptr<arrow::ChunkedArray> makeColumn<arrow::ChunkedArray>(std::shared_ptr<Field>, std::shared_ptr<arrow::Array> array) {
+  return std::make_shared<arrow::ChunkedArray>(array);
+}
+
 std::shared_ptr<Table> createTestTable()
 {
    auto schema_ = exampleSchema();
@@ -49,10 +59,14 @@ std::shared_ptr<Table> createTestTable()
    arrow::ArrayFromVector<BooleanType, bool>(marriageStatus, &arrays_[3]);
    arrow::ArrayFromVector<UInt32Type, unsigned int>(babies, &arrays_[4]);
 
-   std::vector<std::shared_ptr<Column>> columns_ = {
-      std::make_shared<Column>(schema_->field(0), arrays_[0]), std::make_shared<Column>(schema_->field(1), arrays_[1]),
-      std::make_shared<Column>(schema_->field(2), arrays_[2]), std::make_shared<Column>(schema_->field(3), arrays_[3]),
-      std::make_shared<Column>(schema_->field(4), arrays_[4])};
+   using ColumnType = typename decltype(std::declval<arrow::Table>().column(0))::element_type;
+
+   std::vector<std::shared_ptr<ColumnType>> columns_ = {
+      makeColumn<ColumnType>(schema_->field(0), arrays_[0]),
+      makeColumn<ColumnType>(schema_->field(1), arrays_[1]),
+      makeColumn<ColumnType>(schema_->field(2), arrays_[2]),
+      makeColumn<ColumnType>(schema_->field(3), arrays_[3]),
+      makeColumn<ColumnType>(schema_->field(4), arrays_[4])};
 
    auto table_ = Table::Make(schema_, columns_);
    return table_;