mirrors/arrow
0
0 Fork 0
mirror of https://github.com/apache/arrow.git synced 2026-10-29 17:15:39 +08:00
Code Activity

ARROW-7236: [C++] Add Result<T> APIs to arrow/csv

Browse Source
This commit is contained in:
Antoine Pitrou committed 2019-12-02 18:21:18 +01:00
1 parent b218a7fdae
commit a49de2c697
11 files changed
+260 -273
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cpp/src/arrow/csv/column_builder.cc
+38 -46
View File
@@ -64,7 +64,7 @@ class NullColumnBuilder : public ColumnBuilder {
Status Init();
void Insert(int64_t block_index, const std::shared_ptr<BlockParser>& parser) override;
Status Finish(std::shared_ptr<ChunkedArray>* out) override;
Result<std::shared_ptr<ChunkedArray>> Finish() override;
// While NullColumnBuilder is so cheap that it doesn't need parallelization
// in itself, the CSV reader doesn't know this and can still call it from
@@ -106,7 +106,7 @@ void NullColumnBuilder::Insert(int64_t block_index,
});
}
Status NullColumnBuilder::Finish(std::shared_ptr<ChunkedArray>* out) {
Result<std::shared_ptr<ChunkedArray>> NullColumnBuilder::Finish() {
// Unnecessary iff all tasks have finished
std::lock_guard<std::mutex> lock(mutex_);
@@ -115,8 +115,7 @@ Status NullColumnBuilder::Finish(std::shared_ptr<ChunkedArray>* out) {
return Status::Invalid("a chunk failed allocating for an unknown reason");
}
}
*out = std::make_shared<ChunkedArray>(chunks_, type_);
return Status::OK();
return std::make_shared<ChunkedArray>(chunks_, type_);
}
//////////////////////////////////////////////////////////////////////////
@@ -136,7 +135,7 @@ class TypedColumnBuilder : public ColumnBuilder {
Status Init();
void Insert(int64_t block_index, const std::shared_ptr<BlockParser>& parser) override;
Status Finish(std::shared_ptr<ChunkedArray>* out) override;
Result<std::shared_ptr<ChunkedArray>> Finish() override;
protected:
Status WrapConversionError(const Status& st) {
@@ -182,18 +181,20 @@ void TypedColumnBuilder::Insert(int64_t block_index,
// We're careful that all references in the closure outlive the Append() call
task_group_->Append([=]() -> Status {
std::shared_ptr<Array> res;
RETURN_NOT_OK(WrapConversionError(converter_->Convert(*parser, col_index_, &res)));
auto maybe_array = converter_->Convert(*parser, col_index_);
if (!maybe_array.ok()) {
return WrapConversionError(maybe_array.status());
}
std::lock_guard<std::mutex> lock(mutex_);
// Should not insert an already converted chunk
DCHECK_EQ(chunks_[chunk_index], nullptr);
chunks_[chunk_index] = std::move(res);
chunks_[chunk_index] = *std::move(maybe_array);
return Status::OK();
});
}
Status TypedColumnBuilder::Finish(std::shared_ptr<ChunkedArray>* out) {
Result<std::shared_ptr<ChunkedArray>> TypedColumnBuilder::Finish() {
// Unnecessary iff all tasks have finished
std::lock_guard<std::mutex> lock(mutex_);
@@ -202,8 +203,7 @@ Status TypedColumnBuilder::Finish(std::shared_ptr<ChunkedArray>* out) {
return Status::Invalid("a chunk failed converting for an unknown reason");
}
}
*out = std::make_shared<ChunkedArray>(chunks_, type_);
return Status::OK();
return std::make_shared<ChunkedArray>(chunks_, type_);
}
//////////////////////////////////////////////////////////////////////////
@@ -222,7 +222,7 @@ class InferringColumnBuilder : public ColumnBuilder {
Status Init();
void Insert(int64_t block_index, const std::shared_ptr<BlockParser>& parser) override;
Status Finish(std::shared_ptr<ChunkedArray>* out) override;
Result<std::shared_ptr<ChunkedArray>> Finish() override;
protected:
Status LoosenType(const Status& conversion_error);
@@ -366,13 +366,12 @@ Status InferringColumnBuilder::TryConvertChunk(size_t chunk_index) {
std::unique_lock<std::mutex> lock(mutex_);
std::shared_ptr<Converter> converter = converter_;
std::shared_ptr<BlockParser> parser = parsers_[chunk_index];
std::shared_ptr<Array> res;
InferKind kind = infer_kind_;
DCHECK_NE(parser, nullptr);
lock.unlock();
Status conversion_status = converter->Convert(*parser, col_index_, &res);
auto maybe_array = converter->Convert(*parser, col_index_);
lock.lock();
if (kind != infer_kind_) {
@@ -382,9 +381,9 @@ Status InferringColumnBuilder::TryConvertChunk(size_t chunk_index) {
return Status::OK();
}
if (conversion_status.ok()) {
if (maybe_array.ok()) {
// Conversion succeeded
chunks_[chunk_index] = std::move(res);
chunks_[chunk_index] = std::move(*maybe_array);
if (!can_loosen_type_) {
// We won't try to reconvert anymore
parsers_[chunk_index].reset();
@@ -392,7 +391,7 @@ Status InferringColumnBuilder::TryConvertChunk(size_t chunk_index) {
return Status::OK();
} else if (can_loosen_type_) {
// Conversion failed, try another type
RETURN_NOT_OK(LoosenType(conversion_status));
RETURN_NOT_OK(LoosenType(maybe_array.status()));
// Reconvert past finished chunks
// (unfinished chunks will notice by themselves if they need reconverting)
@@ -415,7 +414,7 @@ Status InferringColumnBuilder::TryConvertChunk(size_t chunk_index) {
return Status::OK();
} else {
// Conversion failed but cannot loosen more
return conversion_status;
return maybe_array.status();
}
}
@@ -441,7 +440,7 @@ void InferringColumnBuilder::Insert(int64_t block_index,
ScheduleConvertChunk(chunk_index);
}
Status InferringColumnBuilder::Finish(std::shared_ptr<ChunkedArray>* out) {
Result<std::shared_ptr<ChunkedArray>> InferringColumnBuilder::Finish() {
// Unnecessary iff all tasks have finished
std::lock_guard<std::mutex> lock(mutex_);
@@ -454,46 +453,39 @@ Status InferringColumnBuilder::Finish(std::shared_ptr<ChunkedArray>* out) {
DCHECK_EQ(chunk->type()->id(), infer_type_->id())
<< "Inference didn't equalize types!";
}
*out = std::make_shared<ChunkedArray>(chunks_, infer_type_);
auto ptr = std::make_shared<ChunkedArray>(chunks_, infer_type_);
chunks_.clear();
parsers_.clear();
return Status::OK();
return ptr;
}
//////////////////////////////////////////////////////////////////////////
// Factory functions
Status ColumnBuilder::Make(MemoryPool* pool, const std::shared_ptr<DataType>& type,
int32_t col_index, const ConvertOptions& options,
const std::shared_ptr<TaskGroup>& task_group,
std::shared_ptr<ColumnBuilder>* out) {
auto ptr = new TypedColumnBuilder(type, col_index, options, pool, task_group);
auto res = std::shared_ptr<ColumnBuilder>(ptr);
Result<std::shared_ptr<ColumnBuilder>> ColumnBuilder::Make(
MemoryPool* pool, const std::shared_ptr<DataType>& type, int32_t col_index,
const ConvertOptions& options, const std::shared_ptr<TaskGroup>& task_group) {
auto ptr =
std::make_shared<TypedColumnBuilder>(type, col_index, options, pool, task_group);
RETURN_NOT_OK(ptr->Init());
*out = res;
return Status::OK();
return ptr;
}
Status ColumnBuilder::Make(MemoryPool* pool, int32_t col_index,
const ConvertOptions& options,
const std::shared_ptr<TaskGroup>& task_group,
std::shared_ptr<ColumnBuilder>* out) {
// XXX
auto ptr = new InferringColumnBuilder(col_index, options, pool, task_group);
auto res = std::shared_ptr<ColumnBuilder>(ptr);
Result<std::shared_ptr<ColumnBuilder>> ColumnBuilder::Make(
MemoryPool* pool, int32_t col_index, const ConvertOptions& options,
const std::shared_ptr<TaskGroup>& task_group) {
auto ptr =
std::make_shared<InferringColumnBuilder>(col_index, options, pool, task_group);
RETURN_NOT_OK(ptr->Init());
*out = res;
return Status::OK();
return ptr;
}
Status ColumnBuilder::MakeNull(MemoryPool* pool, const std::shared_ptr<DataType>& type,
const std::shared_ptr<internal::TaskGroup>& task_group,
std::shared_ptr<ColumnBuilder>* out) {
auto res = std::make_shared<NullColumnBuilder>(type, pool, task_group);
RETURN_NOT_OK(res->Init());
*out = std::move(res);
return Status::OK();
Result<std::shared_ptr<ColumnBuilder>> ColumnBuilder::MakeNull(
MemoryPool* pool, const std::shared_ptr<DataType>& type,
const std::shared_ptr<internal::TaskGroup>& task_group) {
auto ptr = std::make_shared<NullColumnBuilder>(type, pool, task_group);
RETURN_NOT_OK(ptr->Init());
return ptr;
}
} // namespace csv
cpp/src/arrow/csv/column_builder.h
+14 -19
View File
@@ -15,21 +15,18 @@
// specific language governing permissions and limitations
// under the License.
#ifndef ARROW_CSV_COLUMN_BUILDER_H
#define ARROW_CSV_COLUMN_BUILDER_H
#pragma once
#include <cstdint>
#include <memory>
#include "arrow/array.h"
#include "arrow/result.h"
#include "arrow/status.h"
#include "arrow/type_fwd.h"
#include "arrow/util/visibility.h"
namespace arrow {
class ChunkedArray;
class DataType;
namespace internal {
class TaskGroup;
@@ -55,7 +52,7 @@ class ARROW_EXPORT ColumnBuilder {
const std::shared_ptr<BlockParser>& parser) = 0;
/// Return the final chunked array. The TaskGroup _must_ have finished!
virtual Status Finish(std::shared_ptr<ChunkedArray>* out) = 0;
virtual Result<std::shared_ptr<ChunkedArray>> Finish() = 0;
/// Change the task group. The previous TaskGroup _must_ have finished!
void SetTaskGroup(const std::shared_ptr<internal::TaskGroup>& task_group);
@@ -63,21 +60,21 @@ class ARROW_EXPORT ColumnBuilder {
std::shared_ptr<internal::TaskGroup> task_group() { return task_group_; }
/// Construct a strictly-typed ColumnBuilder.
static Status Make(MemoryPool* pool, const std::shared_ptr<DataType>& type,
int32_t col_index, const ConvertOptions& options,
const std::shared_ptr<internal::TaskGroup>& task_group,
std::shared_ptr<ColumnBuilder>* out);
static Result<std::shared_ptr<ColumnBuilder>> Make(
MemoryPool* pool, const std::shared_ptr<DataType>& type, int32_t col_index,
const ConvertOptions& options,
const std::shared_ptr<internal::TaskGroup>& task_group);
/// Construct a type-inferring ColumnBuilder.
static Status Make(MemoryPool* pool, int32_t col_index, const ConvertOptions& options,
const std::shared_ptr<internal::TaskGroup>& task_group,
std::shared_ptr<ColumnBuilder>* out);
static Result<std::shared_ptr<ColumnBuilder>> Make(
MemoryPool* pool, int32_t col_index, const ConvertOptions& options,
const std::shared_ptr<internal::TaskGroup>& task_group);
/// Construct a ColumnBuilder for a column of nulls
/// (i.e. not present in the CSV file).
static Status MakeNull(MemoryPool* pool, const std::shared_ptr<DataType>& type,
const std::shared_ptr<internal::TaskGroup>& task_group,
std::shared_ptr<ColumnBuilder>* out);
static Result<std::shared_ptr<ColumnBuilder>> MakeNull(
MemoryPool* pool, const std::shared_ptr<DataType>& type,
const std::shared_ptr<internal::TaskGroup>& task_group);
protected:
explicit ColumnBuilder(const std::shared_ptr<internal::TaskGroup>& task_group)
@@ -89,5 +86,3 @@ class ARROW_EXPORT ColumnBuilder {
} // namespace csv
} // namespace arrow
#endif // ARROW_CSV_COLUMN_BUILDER_H
cpp/src/arrow/csv/column_builder_test.cc
+19 -16
View File
@@ -51,7 +51,7 @@ void AssertBuilding(const std::shared_ptr<ColumnBuilder>& builder,
builder->Append(parser);
}
ASSERT_OK(builder->task_group()->Finish());
ASSERT_OK(builder->Finish(out));
ASSERT_OK_AND_ASSIGN(*out, builder->Finish());
ASSERT_OK((*out)->ValidateFull());
}
@@ -60,7 +60,8 @@ void CheckInferred(const std::shared_ptr<TaskGroup>& tg, const ChunkData& csv_da
std::shared_ptr<ChunkedArray> expected) {
std::shared_ptr<ColumnBuilder> builder;
std::shared_ptr<ChunkedArray> actual;
ASSERT_OK(ColumnBuilder::Make(default_memory_pool(), 0, options, tg, &builder));
ASSERT_OK_AND_ASSIGN(builder,
ColumnBuilder::Make(default_memory_pool(), 0, options, tg));
AssertBuilding(builder, csv_data, &actual);
AssertChunkedEqual(*actual, *expected);
}
@@ -77,7 +78,8 @@ void CheckFixedType(const std::shared_ptr<TaskGroup>& tg,
std::shared_ptr<ChunkedArray> expected) {
std::shared_ptr<ColumnBuilder> builder;
std::shared_ptr<ChunkedArray> actual;
ASSERT_OK(ColumnBuilder::Make(default_memory_pool(), type, 0, options, tg, &builder));
ASSERT_OK_AND_ASSIGN(builder,
ColumnBuilder::Make(default_memory_pool(), type, 0, options, tg));
AssertBuilding(builder, csv_data, &actual);
AssertChunkedEqual(*actual, *expected);
}
@@ -100,7 +102,7 @@ TEST(NullColumnBuilder, Empty) {
auto tg = TaskGroup::MakeSerial();
std::shared_ptr<ColumnBuilder> builder;
ASSERT_OK(ColumnBuilder::MakeNull(default_memory_pool(), type, tg, &builder));
ASSERT_OK_AND_ASSIGN(builder, ColumnBuilder::MakeNull(default_memory_pool(), type, tg));
std::shared_ptr<ChunkedArray> actual;
AssertBuilding(builder, {}, &actual);
@@ -115,7 +117,7 @@ TEST(NullColumnBuilder, InsertNull) {
auto tg = TaskGroup::MakeSerial();
std::shared_ptr<ColumnBuilder> builder;
ASSERT_OK(ColumnBuilder::MakeNull(default_memory_pool(), type, tg, &builder));
ASSERT_OK_AND_ASSIGN(builder, ColumnBuilder::MakeNull(default_memory_pool(), type, tg));
std::shared_ptr<BlockParser> parser;
std::shared_ptr<ChunkedArray> actual, expected;
@@ -125,7 +127,7 @@ TEST(NullColumnBuilder, InsertNull) {
MakeColumnParser({"123"}, &parser);
builder->Insert(0, parser);
ASSERT_OK(builder->task_group()->Finish());
ASSERT_OK(builder->Finish(&actual));
ASSERT_OK_AND_ASSIGN(actual, builder->Finish());
ASSERT_OK(actual->ValidateFull());
auto chunks =
@@ -140,7 +142,7 @@ TEST(NullColumnBuilder, InsertTyped) {
auto tg = TaskGroup::MakeSerial();
std::shared_ptr<ColumnBuilder> builder;
ASSERT_OK(ColumnBuilder::MakeNull(default_memory_pool(), type, tg, &builder));
ASSERT_OK_AND_ASSIGN(builder, ColumnBuilder::MakeNull(default_memory_pool(), type, tg));
std::shared_ptr<BlockParser> parser;
std::shared_ptr<ChunkedArray> actual, expected;
@@ -150,7 +152,7 @@ TEST(NullColumnBuilder, InsertTyped) {
MakeColumnParser({"jkl"}, &parser);
builder->Insert(0, parser);
ASSERT_OK(builder->task_group()->Finish());
ASSERT_OK(builder->Finish(&actual));
ASSERT_OK_AND_ASSIGN(actual, builder->Finish());
ASSERT_OK(actual->ValidateFull());
auto chunks = ArrayVector{ArrayFromJSON(type, "[null]"),
@@ -164,7 +166,7 @@ TEST(NullColumnBuilder, EmptyChunks) {
auto tg = TaskGroup::MakeSerial();
std::shared_ptr<ColumnBuilder> builder;
ASSERT_OK(ColumnBuilder::MakeNull(default_memory_pool(), type, tg, &builder));
ASSERT_OK_AND_ASSIGN(builder, ColumnBuilder::MakeNull(default_memory_pool(), type, tg));
std::shared_ptr<BlockParser> parser;
std::shared_ptr<ChunkedArray> actual, expected;
@@ -176,7 +178,7 @@ TEST(NullColumnBuilder, EmptyChunks) {
MakeColumnParser({}, &parser);
builder->Insert(2, parser);
ASSERT_OK(builder->task_group()->Finish());
ASSERT_OK(builder->Finish(&actual));
ASSERT_OK_AND_ASSIGN(actual, builder->Finish());
ASSERT_OK(actual->ValidateFull());
auto chunks =
@@ -193,8 +195,8 @@ TEST(ColumnBuilder, Empty) {
auto options = ConvertOptions::Defaults();
auto tg = TaskGroup::MakeSerial();
std::shared_ptr<ColumnBuilder> builder;
ASSERT_OK(
ColumnBuilder::Make(default_memory_pool(), int32(), 0, options, tg, &builder));
ASSERT_OK_AND_ASSIGN(
builder, ColumnBuilder::Make(default_memory_pool(), int32(), 0, options, tg));
std::shared_ptr<ChunkedArray> actual;
AssertBuilding(builder, {}, &actual);
@@ -216,8 +218,8 @@ TEST(ColumnBuilder, Insert) {
auto options = ConvertOptions::Defaults();
auto tg = TaskGroup::MakeSerial();
std::shared_ptr<ColumnBuilder> builder;
ASSERT_OK(
ColumnBuilder::Make(default_memory_pool(), int32(), 0, options, tg, &builder));
ASSERT_OK_AND_ASSIGN(
builder, ColumnBuilder::Make(default_memory_pool(), int32(), 0, options, tg));
std::shared_ptr<BlockParser> parser;
std::shared_ptr<ChunkedArray> actual, expected;
@@ -226,7 +228,7 @@ TEST(ColumnBuilder, Insert) {
MakeColumnParser({"123"}, &parser);
builder->Insert(0, parser);
ASSERT_OK(builder->task_group()->Finish());
ASSERT_OK(builder->Finish(&actual));
ASSERT_OK_AND_ASSIGN(actual, builder->Finish());
ASSERT_OK(actual->ValidateFull());
ChunkedArrayFromVector<Int32Type>({{123}, {456}}, &expected);
@@ -428,7 +430,8 @@ void CheckAutoDictEncoded(const std::shared_ptr<TaskGroup>& tg, const ChunkData&
std::vector<std::shared_ptr<Array>> expected_dictionaries) {
std::shared_ptr<ColumnBuilder> builder;
std::shared_ptr<ChunkedArray> actual;
ASSERT_OK(ColumnBuilder::Make(default_memory_pool(), 0, options, tg, &builder));
ASSERT_OK_AND_ASSIGN(builder,
ColumnBuilder::Make(default_memory_pool(), 0, options, tg));
AssertBuilding(builder, csv_data, &actual);
ASSERT_EQ(actual->num_chunks(), static_cast<int>(csv_data.size()));
for (int i = 0; i < actual->num_chunks(); ++i) {
cpp/src/arrow/csv/converter.cc
+115 -123
View File
@@ -133,27 +133,24 @@ class NullConverter : public ConcreteConverter {
public:
using ConcreteConverter::ConcreteConverter;
Status Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) override;
Result<std::shared_ptr<Array>> Convert(const BlockParser& parser,
int32_t col_index) override {
NullBuilder builder(pool_);
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
if (ARROW_PREDICT_TRUE(IsNull(data, size, quoted))) {
return builder.AppendNull();
} else {
return GenericConversionError(type_, data, size);
}
};
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
std::shared_ptr<Array> res;
RETURN_NOT_OK(builder.Finish(&res));
return res;
}
};
Status NullConverter::Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) {
NullBuilder builder(pool_);
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
if (ARROW_PREDICT_TRUE(IsNull(data, size, quoted))) {
return builder.AppendNull();
} else {
return GenericConversionError(type_, data, size);
}
};
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
RETURN_NOT_OK(builder.Finish(out));
return Status::OK();
}
/////////////////////////////////////////////////////////////////////////
// Concrete Converter for var-sized binary strings
@@ -162,8 +159,8 @@ class BinaryConverter : public ConcreteConverter {
public:
using ConcreteConverter::ConcreteConverter;
Status Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) override {
Result<std::shared_ptr<Array>> Convert(const BlockParser& parser,
int32_t col_index) override {
using BuilderType = typename TypeTraits<T>::BuilderType;
BuilderType builder(pool_);
@@ -193,9 +190,9 @@ class BinaryConverter : public ConcreteConverter {
RETURN_NOT_OK(parser.VisitColumn(col_index, visit_non_null));
}
RETURN_NOT_OK(builder.Finish(out));
return Status::OK();
std::shared_ptr<Array> res;
RETURN_NOT_OK(builder.Finish(&res));
return res;
}
protected:
@@ -210,8 +207,8 @@ class DictionaryBinaryConverter : public ConcreteDictionaryConverter {
public:
using ConcreteDictionaryConverter::ConcreteDictionaryConverter;
Status Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) override {
Result<std::shared_ptr<Array>> Convert(const BlockParser& parser,
int32_t col_index) override {
// We use a fixed index width so that all column chunks get the same index type
using BuilderType = Dictionary32Builder<T>;
BuilderType builder(type_, pool_);
@@ -244,9 +241,9 @@ class DictionaryBinaryConverter : public ConcreteDictionaryConverter {
RETURN_NOT_OK(parser.VisitColumn(col_index, visit_non_null));
}
RETURN_NOT_OK(builder.Finish(out));
return Status::OK();
std::shared_ptr<Array> res;
RETURN_NOT_OK(builder.Finish(&res));
return res;
}
void SetMaxCardinality(int32_t max_length) override { max_cardinality_ = max_length; }
@@ -267,31 +264,29 @@ class FixedSizeBinaryConverter : public ConcreteConverter {
public:
using ConcreteConverter::ConcreteConverter;
Status Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) override;
Result<std::shared_ptr<Array>> Convert(const BlockParser& parser,
int32_t col_index) override {
FixedSizeBinaryBuilder builder(type_, pool_);
const uint32_t byte_width = static_cast<uint32_t>(builder.byte_width());
// TODO do we accept nulls here?
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
if (ARROW_PREDICT_FALSE(size != byte_width)) {
return Status::Invalid("CSV conversion error to ", type_->ToString(), ": got a ",
size, "-byte long string");
}
return builder.Append(data);
};
RETURN_NOT_OK(builder.Resize(parser.num_rows()));
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
std::shared_ptr<Array> res;
RETURN_NOT_OK(builder.Finish(&res));
return res;
}
};
Status FixedSizeBinaryConverter::Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) {
FixedSizeBinaryBuilder builder(type_, pool_);
const uint32_t byte_width = static_cast<uint32_t>(builder.byte_width());
// TODO do we accept nulls here?
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
if (ARROW_PREDICT_FALSE(size != byte_width)) {
return Status::Invalid("CSV conversion error to ", type_->ToString(), ": got a ",
size, "-byte long string");
}
return builder.Append(data);
};
RETURN_NOT_OK(builder.Resize(parser.num_rows()));
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
RETURN_NOT_OK(builder.Finish(out));
return Status::OK();
}
/////////////////////////////////////////////////////////////////////////
// Concrete Converter for booleans
@@ -299,8 +294,35 @@ class BooleanConverter : public ConcreteConverter {
public:
using ConcreteConverter::ConcreteConverter;
Status Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) override;
Result<std::shared_ptr<Array>> Convert(const BlockParser& parser,
int32_t col_index) override {
BooleanBuilder builder(type_, pool_);
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
// XXX should quoted values be allowed at all?
if (IsNull(data, size, quoted)) {
builder.UnsafeAppendNull();
return Status::OK();
}
if (false_trie_.Find(
util::string_view(reinterpret_cast<const char*>(data), size)) >= 0) {
builder.UnsafeAppend(false);
return Status::OK();
}
if (true_trie_.Find(util::string_view(reinterpret_cast<const char*>(data), size)) >=
0) {
builder.UnsafeAppend(true);
return Status::OK();
}
return GenericConversionError(type_, data, size);
};
RETURN_NOT_OK(builder.Resize(parser.num_rows()));
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
std::shared_ptr<Array> res;
RETURN_NOT_OK(builder.Finish(&res));
return res;
}
protected:
Status Initialize() override {
@@ -314,35 +336,6 @@ class BooleanConverter : public ConcreteConverter {
Trie false_trie_;
};
Status BooleanConverter::Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) {
BooleanBuilder builder(type_, pool_);
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
// XXX should quoted values be allowed at all?
if (IsNull(data, size, quoted)) {
builder.UnsafeAppendNull();
return Status::OK();
}
if (false_trie_.Find(util::string_view(reinterpret_cast<const char*>(data), size)) >=
0) {
builder.UnsafeAppend(false);
return Status::OK();
}
if (true_trie_.Find(util::string_view(reinterpret_cast<const char*>(data), size)) >=
0) {
builder.UnsafeAppend(true);
return Status::OK();
}
return GenericConversionError(type_, data, size);
};
RETURN_NOT_OK(builder.Resize(parser.num_rows()));
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
RETURN_NOT_OK(builder.Finish(out));
return Status::OK();
}
/////////////////////////////////////////////////////////////////////////
// Concrete Converter for numbers
@@ -351,42 +344,39 @@ class NumericConverter : public ConcreteConverter {
public:
using ConcreteConverter::ConcreteConverter;
Status Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) override;
};
Result<std::shared_ptr<Array>> Convert(const BlockParser& parser,
int32_t col_index) override {
using BuilderType = typename TypeTraits<T>::BuilderType;
using value_type = typename StringConverter<T>::value_type;
template <typename T>
Status NumericConverter<T>::Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) {
using BuilderType = typename TypeTraits<T>::BuilderType;
using value_type = typename StringConverter<T>::value_type;
BuilderType builder(type_, pool_);
StringConverter<T> converter;
BuilderType builder(type_, pool_);
StringConverter<T> converter;
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
// XXX should quoted values be allowed at all?
value_type value;
if (IsNull(data, size, quoted)) {
builder.UnsafeAppendNull();
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
// XXX should quoted values be allowed at all?
value_type value;
if (IsNull(data, size, quoted)) {
builder.UnsafeAppendNull();
return Status::OK();
}
if (!std::is_same<BooleanType, T>::value) {
TrimWhiteSpace(&data, &size);
}
if (ARROW_PREDICT_FALSE(
!converter(reinterpret_cast<const char*>(data), size, &value))) {
return GenericConversionError(type_, data, size);
}
builder.UnsafeAppend(value);
return Status::OK();
}
if (!std::is_same<BooleanType, T>::value) {
TrimWhiteSpace(&data, &size);
}
if (ARROW_PREDICT_FALSE(
!converter(reinterpret_cast<const char*>(data), size, &value))) {
return GenericConversionError(type_, data, size);
}
builder.UnsafeAppend(value);
return Status::OK();
};
RETURN_NOT_OK(builder.Resize(parser.num_rows()));
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
RETURN_NOT_OK(builder.Finish(out));
};
RETURN_NOT_OK(builder.Resize(parser.num_rows()));
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
return Status::OK();
}
std::shared_ptr<Array> res;
RETURN_NOT_OK(builder.Finish(&res));
return res;
}
};
/////////////////////////////////////////////////////////////////////////
// Concrete Converter for timestamps
@@ -395,8 +385,8 @@ class TimestampConverter : public ConcreteConverter {
public:
using ConcreteConverter::ConcreteConverter;
Status Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) override {
Result<std::shared_ptr<Array>> Convert(const BlockParser& parser,
int32_t col_index) override {
using value_type = TimestampType::c_type;
TimestampBuilder builder(type_, pool_);
@@ -417,9 +407,10 @@ class TimestampConverter : public ConcreteConverter {
};
RETURN_NOT_OK(builder.Resize(parser.num_rows()));
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
RETURN_NOT_OK(builder.Finish(out));
return Status::OK();
std::shared_ptr<Array> res;
RETURN_NOT_OK(builder.Finish(&res));
return res;
}
};
@@ -430,8 +421,8 @@ class DecimalConverter : public ConcreteConverter {
public:
using ConcreteConverter::ConcreteConverter;
Status Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) override {
Result<std::shared_ptr<Array>> Convert(const BlockParser& parser,
int32_t col_index) override {
Decimal128Builder builder(type_, pool_);
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
@@ -460,9 +451,10 @@ class DecimalConverter : public ConcreteConverter {
};
RETURN_NOT_OK(builder.Resize(parser.num_rows()));
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
RETURN_NOT_OK(builder.Finish(out));
return Status::OK();
std::shared_ptr<Array> res;
RETURN_NOT_OK(builder.Finish(&res));
return res;
}
};
cpp/src/arrow/csv/converter.h
+3 -6
View File
@@ -15,8 +15,7 @@
// specific language governing permissions and limitations
// under the License.
#ifndef ARROW_CSV_CONVERTER_H
#define ARROW_CSV_CONVERTER_H
#pragma once
#include <cstdint>
#include <memory>
@@ -38,8 +37,8 @@ class ARROW_EXPORT Converter {
MemoryPool* pool);
virtual ~Converter() = default;
virtual Status Convert(const BlockParser& parser, int32_t col_index,
std::shared_ptr<Array>* out) = 0;
virtual Result<std::shared_ptr<Array>> Convert(const BlockParser& parser,
int32_t col_index) = 0;
std::shared_ptr<DataType> type() const { return type_; }
@@ -74,5 +73,3 @@ class ARROW_EXPORT DictionaryConverter : public Converter {
} // namespace csv
} // namespace arrow
#endif // ARROW_CSV_CONVERTER_H
cpp/src/arrow/csv/converter_test.cc
+7 -10
View File
@@ -61,7 +61,7 @@ void AssertConversion(const std::shared_ptr<DataType>& type,
MakeCSVParser(csv_string, &parser);
for (int32_t col_index = 0; col_index < static_cast<int32_t>(expected.size());
++col_index) {
ASSERT_OK(converter->Convert(*parser, col_index, &array));
ASSERT_OK_AND_ASSIGN(array, converter->Convert(*parser, col_index));
ASSERT_OK(array->ValidateFull());
ArrayFromVector<DATA_TYPE, C_TYPE>(type, expected[col_index], &expected_array);
AssertArraysEqual(*expected_array, *array);
@@ -83,7 +83,7 @@ void AssertConversion(const std::shared_ptr<DataType>& type,
MakeCSVParser(csv_string, &parser);
for (int32_t col_index = 0; col_index < static_cast<int32_t>(expected.size());
++col_index) {
ASSERT_OK(converter->Convert(*parser, col_index, &array));
ASSERT_OK_AND_ASSIGN(array, converter->Convert(*parser, col_index));
ASSERT_OK(array->ValidateFull());
ArrayFromVector<DATA_TYPE, C_TYPE>(type, is_valid[col_index], expected[col_index],
&expected_array);
@@ -107,9 +107,7 @@ Result<std::shared_ptr<Array>> DictConversion(
MakeCSVParser({csv_string}, parse_options, &parser);
const int32_t col_index = 0;
std::shared_ptr<Array> array;
RETURN_NOT_OK(converter->Convert(*parser, col_index, &array));
return array;
return converter->Convert(*parser, col_index);
}
void AssertDictConversion(const std::string& csv_string,
@@ -146,16 +144,15 @@ void AssertConversionError(const std::shared_ptr<DataType>& type,
ConvertOptions options = ConvertOptions::Defaults()) {
std::shared_ptr<BlockParser> parser;
std::shared_ptr<Converter> converter;
std::shared_ptr<Array> array;
ASSERT_OK_AND_ASSIGN(converter, Converter::Make(type, options));
MakeCSVParser(csv_string, &parser);
for (int32_t i = 0; i < parser->num_cols(); ++i) {
if (invalid_columns.find(i) == invalid_columns.end()) {
ASSERT_OK(converter->Convert(*parser, i, &array));
ASSERT_OK(converter->Convert(*parser, i));
} else {
ASSERT_RAISES(Invalid, converter->Convert(*parser, i, &array));
ASSERT_RAISES(Invalid, converter->Convert(*parser, i));
}
}
}
@@ -244,10 +241,10 @@ TEST(NullConversion, Basics) {
ASSERT_OK_AND_ASSIGN(converter, Converter::Make(type, ConvertOptions::Defaults()));
MakeCSVParser({"NA,z\n", ",0\n"}, &parser);
ASSERT_OK(converter->Convert(*parser, 0, &array));
ASSERT_OK_AND_ASSIGN(array, converter->Convert(*parser, 0));
ASSERT_EQ(array->type()->id(), Type::NA);
ASSERT_EQ(array->length(), 2);
ASSERT_RAISES(Invalid, converter->Convert(*parser, 1, &array));
ASSERT_RAISES(Invalid, converter->Convert(*parser, 1));
}
TEST(IntegerConversion, Basics) {
cpp/src/arrow/csv/parser.h
+1 -4
View File
@@ -15,8 +15,7 @@
// specific language governing permissions and limitations
// under the License.
#ifndef ARROW_CSV_PARSER_H
#define ARROW_CSV_PARSER_H
#pragma once
#include <cstddef>
#include <cstdint>
@@ -179,5 +178,3 @@ class ARROW_EXPORT BlockParser {
} // namespace csv
} // namespace arrow
#endif // ARROW_CSV_PARSER_H
cpp/src/arrow/csv/reader.cc
+40 -31
View File
@@ -174,10 +174,9 @@ class BaseTableReader : public csv::TableReader {
// Make column builders, assuming inclusion of all columns in CSV file order
Status MakeColumnBuilders() {
for (int32_t col_index = 0; col_index < num_csv_cols_; ++col_index) {
std::shared_ptr<ColumnBuilder> builder;
const auto& col_name = column_names_[col_index];
RETURN_NOT_OK(MakeCSVColumnBuilder(col_name, col_index, &builder));
ARROW_ASSIGN_OR_RAISE(auto builder, MakeCSVColumnBuilder(col_name, col_index));
column_builders_.push_back(builder);
builder_names_.push_back(col_name);
}
@@ -199,11 +198,11 @@ class BaseTableReader : public csv::TableReader {
auto it = col_indices.find(col_name);
if (it != col_indices.end()) {
auto col_index = it->second;
RETURN_NOT_OK(MakeCSVColumnBuilder(col_name, col_index, &builder));
ARROW_ASSIGN_OR_RAISE(builder, MakeCSVColumnBuilder(col_name, col_index));
} else {
// Column not in the CSV file
if (convert_options_.include_missing_columns) {
RETURN_NOT_OK(MakeNullColumnBuilder(col_name, &builder));
ARROW_ASSIGN_OR_RAISE(builder, MakeNullColumnBuilder(col_name));
} else {
return Status::KeyError("Column '", col_name,
"' in include_columns "
@@ -217,21 +216,21 @@ class BaseTableReader : public csv::TableReader {
}
// Make a column builder for the given CSV column name and index
Status MakeCSVColumnBuilder(const std::string& col_name, int32_t col_index,
std::shared_ptr<ColumnBuilder>* out) {
Result<std::shared_ptr<ColumnBuilder>> MakeCSVColumnBuilder(const std::string& col_name,
int32_t col_index) {
// Does the named column have a fixed type?
auto it = convert_options_.column_types.find(col_name);
if (it == convert_options_.column_types.end()) {
return ColumnBuilder::Make(pool_, col_index, convert_options_, task_group_, out);
return ColumnBuilder::Make(pool_, col_index, convert_options_, task_group_);
} else {
return ColumnBuilder::Make(pool_, it->second, col_index, convert_options_,
task_group_, out);
task_group_);
}
}
// Make a column builder for a column of nulls
Status MakeNullColumnBuilder(const std::string& col_name,
std::shared_ptr<ColumnBuilder>* out) {
Result<std::shared_ptr<ColumnBuilder>> MakeNullColumnBuilder(
const std::string& col_name) {
std::shared_ptr<DataType> type;
// If the named column have a fixed type, use it, otherwise use null()
auto it = convert_options_.column_types.find(col_name);
@@ -240,7 +239,7 @@ class BaseTableReader : public csv::TableReader {
} else {
type = null();
}
return ColumnBuilder::MakeNull(pool_, type, task_group_, out);
return ColumnBuilder::MakeNull(pool_, type, task_group_);
}
std::vector<std::string> GenerateColumnNames(int32_t num_cols) {
@@ -296,20 +295,18 @@ class BaseTableReader : public csv::TableReader {
return Status::OK();
}
Status MakeTable(std::shared_ptr<Table>* out) {
Result<std::shared_ptr<Table>> MakeTable() {
DCHECK_EQ(column_builders_.size(), builder_names_.size());
std::vector<std::shared_ptr<Field>> fields;
std::vector<std::shared_ptr<ChunkedArray>> columns;
for (int32_t i = 0; i < static_cast<int32_t>(builder_names_.size()); ++i) {
std::shared_ptr<ChunkedArray> array;
RETURN_NOT_OK(column_builders_[i]->Finish(&array));
ARROW_ASSIGN_OR_RAISE(auto array, column_builders_[i]->Finish());
fields.push_back(::arrow::field(builder_names_[i], array->type()));
columns.emplace_back(std::move(array));
}
*out = Table::Make(schema(fields), columns);
return Status::OK();
return Table::Make(schema(fields), columns);
}
MemoryPool* pool_;
@@ -350,7 +347,7 @@ class SerialTableReader : public BaseTableReader {
&block_iterator_);
}
Status Read(std::shared_ptr<Table>* out) override {
Result<std::shared_ptr<Table>> Read() override {
task_group_ = internal::TaskGroup::MakeSerial();
// First block
@@ -393,7 +390,7 @@ class SerialTableReader : public BaseTableReader {
// Finish conversion, create schema and table
RETURN_NOT_OK(task_group_->Finish());
return MakeTable(out);
return MakeTable();
}
};
@@ -426,7 +423,7 @@ class ThreadedTableReader : public BaseTableReader {
&block_iterator_);
}
Status Read(std::shared_ptr<Table>* out) override {
Result<std::shared_ptr<Table>> Read() override {
task_group_ = internal::TaskGroup::MakeThreaded(thread_pool_);
// Read first block and process header serially
@@ -473,7 +470,7 @@ class ThreadedTableReader : public BaseTableReader {
// Finish conversion, create schema and table
RETURN_NOT_OK(task_group_->Finish());
return MakeTable(out);
return MakeTable();
}
protected:
@@ -483,23 +480,35 @@ class ThreadedTableReader : public BaseTableReader {
/////////////////////////////////////////////////////////////////////////
// TableReader factory function
Result<std::shared_ptr<TableReader>> TableReader::Make(
MemoryPool* pool, std::shared_ptr<io::InputStream> input,
const ReadOptions& read_options, const ParseOptions& parse_options,
const ConvertOptions& convert_options) {
std::shared_ptr<BaseTableReader> reader;
if (read_options.use_threads) {
reader = std::make_shared<ThreadedTableReader>(
pool, input, read_options, parse_options, convert_options, GetCpuThreadPool());
} else {
reader = std::make_shared<SerialTableReader>(pool, input, read_options, parse_options,
convert_options);
}
RETURN_NOT_OK(reader->Init());
return reader;
}
/////////////////////////////////////////////////////////////////////////
// Deprecated API(s)
Status TableReader::Make(MemoryPool* pool, std::shared_ptr<io::InputStream> input,
const ReadOptions& read_options,
const ParseOptions& parse_options,
const ConvertOptions& convert_options,
std::shared_ptr<TableReader>* out) {
std::shared_ptr<BaseTableReader> result;
if (read_options.use_threads) {
result = std::make_shared<ThreadedTableReader>(
pool, input, read_options, parse_options, convert_options, GetCpuThreadPool());
} else {
result = std::make_shared<SerialTableReader>(pool, input, read_options, parse_options,
convert_options);
}
RETURN_NOT_OK(result->Init());
*out = result;
return Status::OK();
return Make(pool, std::move(input), read_options, parse_options, convert_options)
.Value(out);
}
Status TableReader::Read(std::shared_ptr<Table>* out) { return Read().Value(out); }
} // namespace csv
} // namespace arrow
cpp/src/arrow/csv/reader.h
+14 -10
View File
@@ -15,20 +15,16 @@
// specific language governing permissions and limitations
// under the License.
#ifndef ARROW_CSV_READER_H
#define ARROW_CSV_READER_H
#pragma once
#include <memory>
#include "arrow/csv/options.h" // IWYU pragma: keep
#include "arrow/status.h"
#include "arrow/result.h"
#include "arrow/type_fwd.h"
#include "arrow/util/visibility.h"
namespace arrow {
class MemoryPool;
class Table;
namespace io {
class InputStream;
} // namespace io
@@ -41,9 +37,19 @@ class ARROW_EXPORT TableReader {
virtual ~TableReader() = default;
/// Read the entire CSV file and convert it to a Arrow Table
virtual Status Read(std::shared_ptr<Table>* out) = 0;
virtual Result<std::shared_ptr<Table>> Read() = 0;
ARROW_DEPRECATED("Use Result-returning overload")
virtual Status Read(std::shared_ptr<Table>* out);
/// Create a TableReader instance
static Result<std::shared_ptr<TableReader>> Make(MemoryPool* pool,
std::shared_ptr<io::InputStream> input,
const ReadOptions&,
const ParseOptions&,
const ConvertOptions&);
ARROW_DEPRECATED("Use Result-returning overload")
static Status Make(MemoryPool* pool, std::shared_ptr<io::InputStream> input,
const ReadOptions&, const ParseOptions&, const ConvertOptions&,
std::shared_ptr<TableReader>* out);
@@ -51,5 +57,3 @@ class ARROW_EXPORT TableReader {
} // namespace csv
} // namespace arrow
#endif // ARROW_CSV_READER_H
python/pyarrow/_csv.pyx
+5 -4
View File
@@ -575,10 +575,11 @@ def read_csv(input_file, read_options=None, parse_options=None,
_get_parse_options(parse_options, &c_parse_options)
_get_convert_options(convert_options, &c_convert_options)
check_status(CCSVReader.Make(maybe_unbox_memory_pool(memory_pool),
stream, c_read_options, c_parse_options,
c_convert_options, &reader))
reader = GetResultValue(CCSVReader.Make(
maybe_unbox_memory_pool(memory_pool), stream,
c_read_options, c_parse_options, c_convert_options))
with nogil:
check_status(reader.get().Read(&table))
table = GetResultValue(reader.get().Read())
return pyarrow_wrap_table(table)
python/pyarrow/includes/libarrow.pxd
+4 -4
View File
@@ -1354,11 +1354,11 @@ cdef extern from "arrow/csv/api.h" namespace "arrow::csv" nogil:
cdef cppclass CCSVReader" arrow::csv::TableReader":
@staticmethod
CStatus Make(CMemoryPool*, shared_ptr[CInputStream],
CCSVReadOptions, CCSVParseOptions, CCSVConvertOptions,
shared_ptr[CCSVReader]* out)
CResult[shared_ptr[CCSVReader]] Make(
CMemoryPool*, shared_ptr[CInputStream],
CCSVReadOptions, CCSVParseOptions, CCSVConvertOptions)
CStatus Read(shared_ptr[CTable]* out)
CResult[shared_ptr[CTable]] Read()
cdef extern from "arrow/json/options.h" nogil: