mirror of
https://github.com/apache/arrow.git
synced 2026-10-29 09:05:47 +08:00
ARROW-25: [C++] Implement CSV reader
This includes:
- a CSV table reader written in C++
- a Python wrapper around the CSV table reader
- simple type inference for CSV values (null -> int64 -> float64 -> binary)
- generic null parsing using Pandas defaults as a baseline
("NA", "N/A", "NaN"...)
- some simple syntax parameters for CSV parsing
Not included:
- conversion and typing options
- performance tuning
This commit is contained in:
Antoine Pitrou committed
1 parent
d38cf86abb
commit
4ae93b2d55
33 files changed
+4553
-7
No files matched your search
@@ -30,6 +30,13 @@ set(ARROW_SRCS
|
||||
type.cc
|
||||
visitor.cc
|
||||
|
||||
csv/converter.cc
|
||||
csv/chunker.cc
|
||||
csv/column-builder.cc
|
||||
csv/options.cc
|
||||
csv/parser.cc
|
||||
csv/reader.cc
|
||||
|
||||
io/buffered.cc
|
||||
io/file.cc
|
||||
io/interfaces.cc
|
||||
@@ -43,6 +50,7 @@ set(ARROW_SRCS
|
||||
util/hash.cc
|
||||
util/io-util.cc
|
||||
util/key_value_metadata.cc
|
||||
util/task-group.cc
|
||||
util/thread-pool.cc
|
||||
util/logging.cc
|
||||
)
|
||||
@@ -217,5 +225,6 @@ ADD_ARROW_TEST(tensor-test)
|
||||
ADD_ARROW_BENCHMARK(builder-benchmark)
|
||||
ADD_ARROW_BENCHMARK(column-benchmark)
|
||||
|
||||
add_subdirectory(csv)
|
||||
add_subdirectory(io)
|
||||
add_subdirectory(util)
|
||||
@@ -0,0 +1,31 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
ADD_ARROW_TEST(csv-chunker-test)
|
||||
ADD_ARROW_TEST(csv-column-builder-test)
|
||||
ADD_ARROW_TEST(csv-converter-test)
|
||||
ADD_ARROW_TEST(csv-parser-test)
|
||||
|
||||
ADD_ARROW_BENCHMARK(csv-converter-benchmark)
|
||||
ADD_ARROW_BENCHMARK(csv-parser-benchmark)
|
||||
|
||||
# Headers: top level
|
||||
file(GLOB_RECURSE ARROW_CSV_HEADERS "*.h")
|
||||
|
||||
install(FILES
|
||||
${ARROW_CSV_HEADERS}
|
||||
DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/arrow/csv")
|
||||
@@ -0,0 +1,24 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef ARROW_CSV_API_H
|
||||
#define ARROW_CSV_API_H
|
||||
|
||||
#include "arrow/csv/options.h"
|
||||
#include "arrow/csv/reader.h"
|
||||
|
||||
#endif // ARROW_CSV_API_H
|
||||
@@ -0,0 +1,157 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include "arrow/csv/chunker.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/util/logging.h"
|
||||
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
Chunker::Chunker(ParseOptions options, int32_t max_num_rows)
|
||||
: options_(options), max_num_rows_(max_num_rows) {}
|
||||
|
||||
// NOTE: cvsmonkey (https://github.com/dw/csvmonkey) has optimization ideas
|
||||
|
||||
template <bool quoting, bool escaping>
|
||||
inline const char* Chunker::ReadLine(const char* data, const char* data_end) {
|
||||
DCHECK_EQ(quoting, options_.quoting);
|
||||
DCHECK_EQ(escaping, options_.escaping);
|
||||
|
||||
// The parsing state machine
|
||||
char c;
|
||||
|
||||
FieldStart:
|
||||
// At the start of a field
|
||||
// Quoting is only recognized at start of field
|
||||
if (quoting && ARROW_PREDICT_TRUE(data != data_end) && *data == options_.quote_char) {
|
||||
data++;
|
||||
goto InQuotedField;
|
||||
} else {
|
||||
goto InField;
|
||||
}
|
||||
|
||||
InField:
|
||||
// Inside a non-quoted part of a field
|
||||
if (ARROW_PREDICT_FALSE(data == data_end)) {
|
||||
goto AbortLine;
|
||||
}
|
||||
c = *data++;
|
||||
if (escaping && ARROW_PREDICT_FALSE(c == options_.escape_char)) {
|
||||
if (ARROW_PREDICT_FALSE(data == data_end)) {
|
||||
goto AbortLine;
|
||||
}
|
||||
data++;
|
||||
goto InField;
|
||||
}
|
||||
if (ARROW_PREDICT_FALSE(c == '\r')) {
|
||||
if (ARROW_PREDICT_TRUE(data != data_end) && *data == '\n') {
|
||||
data++;
|
||||
}
|
||||
goto LineEnd;
|
||||
}
|
||||
if (ARROW_PREDICT_FALSE(c == '\n')) {
|
||||
goto LineEnd;
|
||||
}
|
||||
if (ARROW_PREDICT_FALSE(c == options_.delimiter)) {
|
||||
goto FieldEnd;
|
||||
}
|
||||
goto InField;
|
||||
|
||||
InQuotedField:
|
||||
// Inside a quoted part of a field
|
||||
if (ARROW_PREDICT_FALSE(data == data_end)) {
|
||||
goto AbortLine;
|
||||
}
|
||||
c = *data++;
|
||||
if (escaping && ARROW_PREDICT_FALSE(c == options_.escape_char)) {
|
||||
if (data == data_end) {
|
||||
goto AbortLine;
|
||||
}
|
||||
data++;
|
||||
goto InQuotedField;
|
||||
}
|
||||
if (ARROW_PREDICT_FALSE(c == options_.quote_char)) {
|
||||
if (options_.double_quote && data != data_end && *data == options_.quote_char) {
|
||||
// Double-quoting
|
||||
data++;
|
||||
} else {
|
||||
// End of single-quoting
|
||||
goto InField;
|
||||
}
|
||||
}
|
||||
goto InQuotedField;
|
||||
|
||||
FieldEnd:
|
||||
// At the end of a field
|
||||
goto FieldStart;
|
||||
|
||||
LineEnd:
|
||||
// At the end of line, possibly in the middle of the newline separator
|
||||
// if (ARROW_PREDICT_TRUE(data < data_end) && data[-1] == '\r' && *data == '\n') {
|
||||
// data++;
|
||||
// }
|
||||
return data;
|
||||
|
||||
AbortLine:
|
||||
// Truncated line at end of block
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
template <bool quoting, bool escaping>
|
||||
Status Chunker::ProcessSpecialized(const char* start, uint32_t size, uint32_t* out_size) {
|
||||
DCHECK_EQ(quoting, options_.quoting);
|
||||
DCHECK_EQ(escaping, options_.escaping);
|
||||
|
||||
num_rows_ = 0;
|
||||
const char* data = start;
|
||||
const char* data_end = start + size;
|
||||
|
||||
while (data < data_end && num_rows_ < max_num_rows_) {
|
||||
const char* line_end = ReadLine<quoting, escaping>(data, data_end);
|
||||
if (line_end == nullptr) {
|
||||
// Cannot read any further
|
||||
break;
|
||||
}
|
||||
data = line_end;
|
||||
++num_rows_;
|
||||
}
|
||||
*out_size = static_cast<uint32_t>(data - start);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Chunker::Process(const char* start, uint32_t size, uint32_t* out_size) {
|
||||
if (options_.quoting) {
|
||||
if (options_.escaping) {
|
||||
return ProcessSpecialized<true, true>(start, size, out_size);
|
||||
} else {
|
||||
return ProcessSpecialized<true, false>(start, size, out_size);
|
||||
}
|
||||
} else {
|
||||
if (options_.escaping) {
|
||||
return ProcessSpecialized<false, true>(start, size, out_size);
|
||||
} else {
|
||||
return ProcessSpecialized<false, false>(start, size, out_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,77 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef ARROW_CSV_CHUNKER_H
|
||||
#define ARROW_CSV_CHUNKER_H
|
||||
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/csv/options.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/util/visibility.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
constexpr int32_t kMaxChunkerNumRows = 100000;
|
||||
|
||||
/// \class Chunker
|
||||
/// \brief A reusable block-based chunker for CSV data
|
||||
///
|
||||
/// The chunker takes a block of CSV data and finds a suitable place
|
||||
/// to cut it up without splitting a row.
|
||||
/// If the block is truncated (i.e. not all data can be chunked), it is up
|
||||
/// to the caller to arrange the next block to start with the trailing data.
|
||||
///
|
||||
/// Note: if the previous block ends with CR (0x0d) and a new block starts
|
||||
/// with LF (0x0a), the chunker will consider the leading newline as an empty line.
|
||||
class ARROW_EXPORT Chunker {
|
||||
public:
|
||||
explicit Chunker(ParseOptions options, int32_t max_num_rows = kMaxChunkerNumRows);
|
||||
|
||||
/// \brief Carve up a chunk in a block of data
|
||||
///
|
||||
/// Process a block of CSV data, reading up to max_num_rows rows.
|
||||
/// The number of bytes in the chunk is returned in out_size.
|
||||
Status Process(const char* data, uint32_t size, uint32_t* out_size);
|
||||
|
||||
int32_t num_rows() const { return num_rows_; }
|
||||
|
||||
protected:
|
||||
ARROW_DISALLOW_COPY_AND_ASSIGN(Chunker);
|
||||
|
||||
// Like Process(), but specialized for some parsing options
|
||||
template <bool quoting, bool escaping>
|
||||
Status ProcessSpecialized(const char* data, uint32_t size, uint32_t* out_size);
|
||||
|
||||
// Detect a single line from the data pointer. Return the line end,
|
||||
// or nullptr if the remaining line is truncated.
|
||||
template <bool quoting, bool escaping>
|
||||
inline const char* ReadLine(const char* data, const char* data_end);
|
||||
|
||||
ParseOptions options_;
|
||||
// The number of rows chunked from the block
|
||||
int32_t num_rows_;
|
||||
// The maximum number of rows to chunk from this block
|
||||
int32_t max_num_rows_;
|
||||
};
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
|
||||
#endif // ARROW_CSV_CHUNKER_H
|
||||
@@ -0,0 +1,335 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include <mutex>
|
||||
#include <utility>
|
||||
|
||||
#include "arrow/csv/column-builder.h"
|
||||
#include "arrow/util/logging.h"
|
||||
#include "arrow/util/macros.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
using internal::TaskGroup;
|
||||
|
||||
void ColumnBuilder::SetTaskGroup(const std::shared_ptr<internal::TaskGroup>& task_group) {
|
||||
task_group_ = task_group;
|
||||
}
|
||||
|
||||
void ColumnBuilder::Append(const std::shared_ptr<BlockParser>& parser) {
|
||||
Insert(static_cast<int64_t>(chunks_.size()), parser);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// Pre-typed column builder implementation
|
||||
|
||||
class TypedColumnBuilder : public ColumnBuilder {
|
||||
public:
|
||||
TypedColumnBuilder(const std::shared_ptr<DataType>& type, int32_t col_index,
|
||||
const ConvertOptions& options, MemoryPool* pool,
|
||||
const std::shared_ptr<internal::TaskGroup>& task_group)
|
||||
: ColumnBuilder(task_group),
|
||||
type_(type),
|
||||
col_index_(col_index),
|
||||
options_(options),
|
||||
pool_(pool) {}
|
||||
|
||||
Status Init();
|
||||
|
||||
void Insert(int64_t block_index, const std::shared_ptr<BlockParser>& parser) override;
|
||||
Status Finish(std::shared_ptr<ChunkedArray>* out) override;
|
||||
|
||||
protected:
|
||||
std::mutex mutex_;
|
||||
|
||||
std::shared_ptr<DataType> type_;
|
||||
int32_t col_index_;
|
||||
ConvertOptions options_;
|
||||
MemoryPool* pool_;
|
||||
|
||||
std::shared_ptr<Converter> converter_;
|
||||
};
|
||||
|
||||
Status TypedColumnBuilder::Init() {
|
||||
return Converter::Make(type_, options_, pool_, &converter_);
|
||||
}
|
||||
|
||||
void TypedColumnBuilder::Insert(int64_t block_index,
|
||||
const std::shared_ptr<BlockParser>& parser) {
|
||||
DCHECK_NE(converter_, nullptr);
|
||||
|
||||
// Create a null Array pointer at the back at the list
|
||||
// and spawn a task to initialize it after conversion
|
||||
size_t chunk_index = static_cast<size_t>(block_index);
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
if (chunks_.size() <= chunk_index) {
|
||||
chunks_.resize(chunk_index + 1);
|
||||
}
|
||||
}
|
||||
|
||||
// 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(converter_->Convert(*parser, col_index_, &res));
|
||||
|
||||
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);
|
||||
return Status::OK();
|
||||
});
|
||||
}
|
||||
|
||||
Status TypedColumnBuilder::Finish(std::shared_ptr<ChunkedArray>* out) {
|
||||
// Unnecessary iff all tasks have finished
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
|
||||
for (const auto& chunk : chunks_) {
|
||||
if (chunk == nullptr) {
|
||||
return Status::Invalid("a chunk failed converting for an unknown reason");
|
||||
}
|
||||
}
|
||||
*out = std::make_shared<ChunkedArray>(chunks_, type_);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// Type-inferring column builder implementation
|
||||
|
||||
class InferringColumnBuilder : public ColumnBuilder {
|
||||
public:
|
||||
InferringColumnBuilder(int32_t col_index, const ConvertOptions& options,
|
||||
MemoryPool* pool,
|
||||
const std::shared_ptr<internal::TaskGroup>& task_group)
|
||||
: ColumnBuilder(task_group),
|
||||
col_index_(col_index),
|
||||
options_(options),
|
||||
pool_(pool) {}
|
||||
|
||||
Status Init();
|
||||
|
||||
void Insert(int64_t block_index, const std::shared_ptr<BlockParser>& parser) override;
|
||||
Status Finish(std::shared_ptr<ChunkedArray>* out) override;
|
||||
|
||||
protected:
|
||||
Status LoosenType();
|
||||
Status UpdateType();
|
||||
Status TryConvertChunk(size_t chunk_index);
|
||||
// This must be called unlocked!
|
||||
void ScheduleConvertChunk(size_t chunk_index);
|
||||
|
||||
std::mutex mutex_;
|
||||
|
||||
int32_t col_index_;
|
||||
ConvertOptions options_;
|
||||
MemoryPool* pool_;
|
||||
std::shared_ptr<Converter> converter_;
|
||||
|
||||
// Current inference status
|
||||
enum class InferKind { Null, Integer, Real, Text };
|
||||
|
||||
std::shared_ptr<DataType> infer_type_;
|
||||
InferKind infer_kind_;
|
||||
bool can_loosen_type_;
|
||||
|
||||
// The parsers corresponding to each chunk (for reconverting)
|
||||
std::vector<std::shared_ptr<BlockParser>> parsers_;
|
||||
};
|
||||
|
||||
Status InferringColumnBuilder::Init() {
|
||||
infer_kind_ = InferKind::Null;
|
||||
return UpdateType();
|
||||
}
|
||||
|
||||
Status InferringColumnBuilder::LoosenType() {
|
||||
// We are locked
|
||||
|
||||
DCHECK(can_loosen_type_);
|
||||
switch (infer_kind_) {
|
||||
case InferKind::Null:
|
||||
infer_kind_ = InferKind::Integer;
|
||||
break;
|
||||
case InferKind::Integer:
|
||||
infer_kind_ = InferKind::Real;
|
||||
break;
|
||||
case InferKind::Real:
|
||||
infer_kind_ = InferKind::Text;
|
||||
break;
|
||||
case InferKind::Text:
|
||||
return Status::UnknownError("Shouldn't come here");
|
||||
}
|
||||
return UpdateType();
|
||||
}
|
||||
|
||||
Status InferringColumnBuilder::UpdateType() {
|
||||
// We are locked
|
||||
|
||||
switch (infer_kind_) {
|
||||
case InferKind::Null:
|
||||
infer_type_ = null();
|
||||
can_loosen_type_ = true;
|
||||
break;
|
||||
case InferKind::Integer:
|
||||
infer_type_ = int64();
|
||||
can_loosen_type_ = true;
|
||||
break;
|
||||
case InferKind::Real:
|
||||
infer_type_ = float64();
|
||||
can_loosen_type_ = true;
|
||||
break;
|
||||
case InferKind::Text:
|
||||
infer_type_ = binary();
|
||||
can_loosen_type_ = false;
|
||||
break;
|
||||
}
|
||||
return Converter::Make(infer_type_, options_, pool_, &converter_);
|
||||
}
|
||||
|
||||
void InferringColumnBuilder::ScheduleConvertChunk(size_t chunk_index) {
|
||||
// We're careful that all values in the closure outlive the Append() call
|
||||
task_group_->Append([=]() { return TryConvertChunk(chunk_index); });
|
||||
}
|
||||
|
||||
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 st = converter->Convert(*parser, col_index_, &res);
|
||||
lock.lock();
|
||||
|
||||
if (kind != infer_kind_) {
|
||||
// infer_kind_ was changed by another task, reconvert
|
||||
lock.unlock();
|
||||
ScheduleConvertChunk(chunk_index);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
if (st.ok()) {
|
||||
// Conversion succeeded
|
||||
chunks_[chunk_index] = std::move(res);
|
||||
if (!can_loosen_type_) {
|
||||
// We won't try to reconvert anymore
|
||||
parsers_[chunk_index].reset();
|
||||
}
|
||||
return Status::OK();
|
||||
} else if (can_loosen_type_) {
|
||||
// Conversion failed, try another type
|
||||
RETURN_NOT_OK(LoosenType());
|
||||
|
||||
// Reconvert past finished chunks
|
||||
// (unfinished chunks will notice by themselves if they need reconverting)
|
||||
size_t nchunks = chunks_.size();
|
||||
for (size_t i = 0; i < nchunks; ++i) {
|
||||
if (i != chunk_index && chunks_[i]) {
|
||||
// We're assuming the chunk was converted using the wrong type
|
||||
// (which should be true unless the executor reorders tasks)
|
||||
chunks_[i].reset();
|
||||
lock.unlock();
|
||||
ScheduleConvertChunk(i);
|
||||
lock.lock();
|
||||
}
|
||||
}
|
||||
|
||||
// Reconvert this chunk
|
||||
lock.unlock();
|
||||
ScheduleConvertChunk(chunk_index);
|
||||
|
||||
return Status::OK();
|
||||
} else {
|
||||
// Conversion failed but cannot loosen more
|
||||
return st;
|
||||
}
|
||||
}
|
||||
|
||||
void InferringColumnBuilder::Insert(int64_t block_index,
|
||||
const std::shared_ptr<BlockParser>& parser) {
|
||||
DCHECK_NE(converter_, nullptr);
|
||||
|
||||
// Create a slot for the new chunk and spawn a task to convert it
|
||||
size_t chunk_index = static_cast<size_t>(block_index);
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
if (chunks_.size() <= chunk_index) {
|
||||
chunks_.resize(chunk_index + 1);
|
||||
}
|
||||
if (parsers_.size() <= chunk_index) {
|
||||
parsers_.resize(chunk_index + 1);
|
||||
}
|
||||
// Should not insert an already converting chunk
|
||||
DCHECK_EQ(parsers_[chunk_index], nullptr);
|
||||
parsers_[chunk_index] = parser;
|
||||
}
|
||||
|
||||
ScheduleConvertChunk(chunk_index);
|
||||
}
|
||||
|
||||
Status InferringColumnBuilder::Finish(std::shared_ptr<ChunkedArray>* out) {
|
||||
// Unnecessary iff all tasks have finished
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
|
||||
for (const auto& chunk : chunks_) {
|
||||
if (chunk == nullptr) {
|
||||
return Status::Invalid("A chunk failed converting for an unknown reason");
|
||||
}
|
||||
// XXX Perhaps we must instead do a last equalization pass
|
||||
// in this serial step
|
||||
DCHECK_EQ(chunk->type()->id(), infer_type_->id())
|
||||
<< "Inference didn't equalize types!";
|
||||
}
|
||||
*out = std::make_shared<ChunkedArray>(chunks_, infer_type_);
|
||||
chunks_.clear();
|
||||
parsers_.clear();
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// Factory functions
|
||||
|
||||
Status ColumnBuilder::Make(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, default_memory_pool(), task_group);
|
||||
auto res = std::shared_ptr<ColumnBuilder>(ptr);
|
||||
RETURN_NOT_OK(ptr->Init());
|
||||
*out = res;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status ColumnBuilder::Make(int32_t col_index, const ConvertOptions& options,
|
||||
const std::shared_ptr<TaskGroup>& task_group,
|
||||
std::shared_ptr<ColumnBuilder>* out) {
|
||||
auto ptr =
|
||||
new InferringColumnBuilder(col_index, options, default_memory_pool(), task_group);
|
||||
auto res = std::shared_ptr<ColumnBuilder>(ptr);
|
||||
RETURN_NOT_OK(ptr->Init());
|
||||
*out = res;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,80 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef ARROW_CSV_COLUMN_BUILDER_H
|
||||
#define ARROW_CSV_COLUMN_BUILDER_H
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/array.h"
|
||||
#include "arrow/csv/converter.h"
|
||||
#include "arrow/csv/options.h"
|
||||
#include "arrow/memory_pool.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/table.h"
|
||||
#include "arrow/type.h"
|
||||
#include "arrow/util/task-group.h"
|
||||
#include "arrow/util/visibility.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
class ARROW_EXPORT ColumnBuilder {
|
||||
public:
|
||||
virtual ~ColumnBuilder() = default;
|
||||
|
||||
/// Spawn a task that will try to convert and append the given CSV block.
|
||||
/// All calls to Append() should happen on the same thread, otherwise
|
||||
/// call Insert() instead.
|
||||
virtual void Append(const std::shared_ptr<BlockParser>& parser);
|
||||
|
||||
/// Spawn a task that will try to convert and insert the given CSV block
|
||||
virtual void Insert(int64_t block_index,
|
||||
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;
|
||||
|
||||
/// Change the task group. The previous TaskGroup _must_ have finished!
|
||||
void SetTaskGroup(const std::shared_ptr<internal::TaskGroup>& task_group);
|
||||
|
||||
std::shared_ptr<internal::TaskGroup> task_group() { return task_group_; }
|
||||
|
||||
/// Construct a strictly-typed ColumnBuilder.
|
||||
static Status Make(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);
|
||||
|
||||
/// Construct a type-inferring ColumnBuilder.
|
||||
static Status Make(int32_t col_index, const ConvertOptions& options,
|
||||
const std::shared_ptr<internal::TaskGroup>& task_group,
|
||||
std::shared_ptr<ColumnBuilder>* out);
|
||||
|
||||
protected:
|
||||
explicit ColumnBuilder(const std::shared_ptr<internal::TaskGroup>& task_group)
|
||||
: task_group_(task_group) {}
|
||||
|
||||
std::shared_ptr<internal::TaskGroup> task_group_;
|
||||
ArrayVector chunks_;
|
||||
};
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
|
||||
#endif // ARROW_CSV_COLUMN_BUILDER_H
|
||||
@@ -0,0 +1,344 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include "arrow/csv/converter.h"
|
||||
#include "arrow/builder.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/type.h"
|
||||
#include "arrow/type_traits.h"
|
||||
#include "arrow/util/parsing.h"
|
||||
|
||||
#include <cstring>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
using internal::StringConverter;
|
||||
|
||||
namespace {
|
||||
|
||||
Status GenericConversionError(const std::shared_ptr<DataType>& type, const uint8_t* data,
|
||||
uint32_t size) {
|
||||
std::stringstream ss;
|
||||
ss << "CSV conversion error to " << type->ToString() << ": invalid value '"
|
||||
<< std::string(reinterpret_cast<const char*>(data), size) << "'";
|
||||
return Status::Invalid(ss.str());
|
||||
}
|
||||
|
||||
class ConcreteConverter : public Converter {
|
||||
public:
|
||||
using Converter::Converter;
|
||||
|
||||
protected:
|
||||
inline bool IsNull(const uint8_t* data, uint32_t size, bool quoted);
|
||||
};
|
||||
|
||||
// Recognize various spellings of null values. The list of possible spellings
|
||||
// is taken from Pandas read_csv() documentation.
|
||||
bool ConcreteConverter::IsNull(const uint8_t* data, uint32_t size, bool quoted) {
|
||||
if (quoted) {
|
||||
return false;
|
||||
}
|
||||
if (size == 0) {
|
||||
return true;
|
||||
}
|
||||
// No 1-character null value exists
|
||||
if (size == 1) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// XXX if the CSV parser guaranteed enough excess bytes at the end of the
|
||||
// parsed area, we wouldn't need to always check size before comparing characters.
|
||||
|
||||
auto chars = reinterpret_cast<const char*>(data);
|
||||
auto first = chars[0];
|
||||
auto second = chars[1];
|
||||
switch (first) {
|
||||
case 'N': {
|
||||
// "NA", "N/A", "NaN", "NULL"
|
||||
if (size == 2) {
|
||||
return second == 'A';
|
||||
}
|
||||
auto third = chars[2];
|
||||
if (size == 3) {
|
||||
return (second == '/' && third == 'A') || (second == 'a' && third == 'N');
|
||||
}
|
||||
if (size == 4) {
|
||||
return (second == 'U' && third == 'L' && chars[3] == 'L');
|
||||
}
|
||||
return false;
|
||||
}
|
||||
case 'n': {
|
||||
// "n/a", "nan", "null"
|
||||
if (size == 2) {
|
||||
return false;
|
||||
}
|
||||
auto third = chars[2];
|
||||
if (size == 3) {
|
||||
return (second == '/' && third == 'a') || (second == 'a' && third == 'n');
|
||||
}
|
||||
if (size == 4) {
|
||||
return (second == 'u' && third == 'l' && chars[3] == 'l');
|
||||
}
|
||||
return false;
|
||||
}
|
||||
case '1': {
|
||||
// '1.#IND', '1.#QNAN'
|
||||
if (size == 6) {
|
||||
// '#' is the most unlikely char here, check it first
|
||||
return (chars[2] == '#' && chars[1] == '.' && chars[3] == 'I' &&
|
||||
chars[4] == 'N' && chars[5] == 'D');
|
||||
}
|
||||
if (size == 7) {
|
||||
return (chars[2] == '#' && chars[1] == '.' && chars[3] == 'Q' &&
|
||||
chars[4] == 'N' && chars[5] == 'A' && chars[6] == 'N');
|
||||
}
|
||||
return false;
|
||||
}
|
||||
case '-': {
|
||||
switch (second) {
|
||||
case 'N':
|
||||
// "-NaN"
|
||||
return (size == 4 && chars[2] == 'a' && chars[3] == 'N');
|
||||
case 'n':
|
||||
// "-nan"
|
||||
return (size == 4 && chars[2] == 'a' && chars[3] == 'n');
|
||||
case '1':
|
||||
// "-1.#IND", "-1.#QNAN"
|
||||
if (size == 7) {
|
||||
return (chars[3] == '#' && chars[2] == '.' && chars[4] == 'I' &&
|
||||
chars[5] == 'N' && chars[6] == 'D');
|
||||
}
|
||||
if (size == 8) {
|
||||
return (chars[3] == '#' && chars[2] == '.' && chars[4] == 'Q' &&
|
||||
chars[5] == 'N' && chars[6] == 'A' && chars[7] == 'N');
|
||||
}
|
||||
return false;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
case '#': {
|
||||
// "#N/A", "#N/A N/A", "#NA"
|
||||
if (size < 3 || chars[1] != 'N') {
|
||||
return false;
|
||||
}
|
||||
auto third = chars[2];
|
||||
if (size == 3) {
|
||||
return third == 'A';
|
||||
}
|
||||
if (size == 4) {
|
||||
return third == '/' && chars[3] == 'A';
|
||||
}
|
||||
if (size == 8) {
|
||||
return std::memcmp(data + 2, "/A N/A", 5) == 0;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// Concrete Converter for null values
|
||||
|
||||
class NullConverter : public ConcreteConverter {
|
||||
public:
|
||||
using ConcreteConverter::ConcreteConverter;
|
||||
|
||||
Status Convert(const BlockParser& parser, int32_t col_index,
|
||||
std::shared_ptr<Array>* out) override;
|
||||
};
|
||||
|
||||
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
|
||||
|
||||
template <typename T>
|
||||
class VarSizeBinaryConverter : public ConcreteConverter {
|
||||
public:
|
||||
using ConcreteConverter::ConcreteConverter;
|
||||
|
||||
Status Convert(const BlockParser& parser, int32_t col_index,
|
||||
std::shared_ptr<Array>* out) override;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
Status VarSizeBinaryConverter<T>::Convert(const BlockParser& parser, int32_t col_index,
|
||||
std::shared_ptr<Array>* out) {
|
||||
using BuilderType = typename TypeTraits<T>::BuilderType;
|
||||
BuilderType builder(pool_);
|
||||
|
||||
// TODO handle nulls
|
||||
|
||||
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
|
||||
return builder.Append(data, size);
|
||||
};
|
||||
RETURN_NOT_OK(builder.Resize(parser.num_rows()));
|
||||
RETURN_NOT_OK(builder.ReserveData(parser.num_bytes()));
|
||||
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
|
||||
RETURN_NOT_OK(builder.Finish(out));
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// Concrete Converter for fixed-sized binary strings
|
||||
|
||||
class FixedSizeBinaryConverter : public ConcreteConverter {
|
||||
public:
|
||||
using ConcreteConverter::ConcreteConverter;
|
||||
|
||||
Status Convert(const BlockParser& parser, int32_t col_index,
|
||||
std::shared_ptr<Array>* out) override;
|
||||
};
|
||||
|
||||
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 handle nulls
|
||||
|
||||
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
|
||||
if (ARROW_PREDICT_FALSE(size != byte_width)) {
|
||||
std::stringstream ss;
|
||||
ss << "CSV conversion error to " << type_->ToString() << ": got a " << size
|
||||
<< "-byte long string";
|
||||
return Status::Invalid(ss.str());
|
||||
}
|
||||
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 numbers
|
||||
|
||||
template <typename T>
|
||||
class NumericConverter : public ConcreteConverter {
|
||||
public:
|
||||
using ConcreteConverter::ConcreteConverter;
|
||||
|
||||
Status Convert(const BlockParser& parser, int32_t col_index,
|
||||
std::shared_ptr<Array>* out) override;
|
||||
};
|
||||
|
||||
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;
|
||||
|
||||
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
|
||||
value_type value;
|
||||
if (IsNull(data, size, quoted)) {
|
||||
return builder.AppendNull();
|
||||
}
|
||||
if (ARROW_PREDICT_FALSE(
|
||||
!converter(reinterpret_cast<const char*>(data), size, &value))) {
|
||||
return GenericConversionError(type_, data, size);
|
||||
}
|
||||
return builder.Append(value);
|
||||
};
|
||||
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();
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// Base Converter class implementation
|
||||
|
||||
Converter::Converter(const std::shared_ptr<DataType>& type, ConvertOptions options,
|
||||
MemoryPool* pool)
|
||||
: options_(options), pool_(pool), type_(type) {}
|
||||
|
||||
Status Converter::Make(const std::shared_ptr<DataType>& type, ConvertOptions options,
|
||||
MemoryPool* pool, std::shared_ptr<Converter>* out) {
|
||||
Converter* result;
|
||||
|
||||
switch (type->id()) {
|
||||
#define CONVERTER_CASE(TYPE_ID, CONVERTER_TYPE) \
|
||||
case TYPE_ID: \
|
||||
result = new CONVERTER_TYPE(type, options, pool); \
|
||||
break;
|
||||
|
||||
CONVERTER_CASE(Type::NA, NullConverter)
|
||||
CONVERTER_CASE(Type::BINARY, VarSizeBinaryConverter<BinaryType>)
|
||||
CONVERTER_CASE(Type::STRING, VarSizeBinaryConverter<StringType>)
|
||||
CONVERTER_CASE(Type::FIXED_SIZE_BINARY, FixedSizeBinaryConverter)
|
||||
CONVERTER_CASE(Type::INT8, NumericConverter<Int8Type>)
|
||||
CONVERTER_CASE(Type::INT16, NumericConverter<Int16Type>)
|
||||
CONVERTER_CASE(Type::INT32, NumericConverter<Int32Type>)
|
||||
CONVERTER_CASE(Type::INT64, NumericConverter<Int64Type>)
|
||||
CONVERTER_CASE(Type::UINT8, NumericConverter<UInt8Type>)
|
||||
CONVERTER_CASE(Type::UINT16, NumericConverter<UInt16Type>)
|
||||
CONVERTER_CASE(Type::UINT32, NumericConverter<UInt32Type>)
|
||||
CONVERTER_CASE(Type::UINT64, NumericConverter<UInt64Type>)
|
||||
CONVERTER_CASE(Type::FLOAT, NumericConverter<FloatType>)
|
||||
CONVERTER_CASE(Type::DOUBLE, NumericConverter<DoubleType>)
|
||||
CONVERTER_CASE(Type::BOOL, NumericConverter<BooleanType>)
|
||||
|
||||
default: {
|
||||
std::stringstream ss;
|
||||
ss << "CSV conversion to " << type->ToString() << " is not supported";
|
||||
return Status::NotImplemented(ss.str());
|
||||
}
|
||||
|
||||
#undef CONVERTER_CASE
|
||||
}
|
||||
out->reset(result);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Converter::Make(const std::shared_ptr<DataType>& type, ConvertOptions options,
|
||||
std::shared_ptr<Converter>* out) {
|
||||
return Make(type, options, default_memory_pool(), out);
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,63 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef ARROW_CSV_CONVERTER_H
|
||||
#define ARROW_CSV_CONVERTER_H
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "arrow/csv/options.h"
|
||||
#include "arrow/csv/parser.h"
|
||||
#include "arrow/memory_pool.h"
|
||||
#include "arrow/util/visibility.h"
|
||||
|
||||
namespace arrow {
|
||||
|
||||
class Array;
|
||||
class DataType;
|
||||
class Status;
|
||||
|
||||
namespace csv {
|
||||
|
||||
class ARROW_EXPORT Converter {
|
||||
public:
|
||||
Converter(const std::shared_ptr<DataType>& type, ConvertOptions options,
|
||||
MemoryPool* pool);
|
||||
virtual ~Converter() = default;
|
||||
|
||||
virtual Status Convert(const BlockParser& parser, int32_t col_index,
|
||||
std::shared_ptr<Array>* out) = 0;
|
||||
|
||||
std::shared_ptr<DataType> type() const { return type_; }
|
||||
|
||||
static Status Make(const std::shared_ptr<DataType>& type, ConvertOptions options,
|
||||
std::shared_ptr<Converter>* out);
|
||||
static Status Make(const std::shared_ptr<DataType>& type, ConvertOptions options,
|
||||
MemoryPool* pool, std::shared_ptr<Converter>* out);
|
||||
|
||||
protected:
|
||||
ARROW_DISALLOW_COPY_AND_ASSIGN(Converter);
|
||||
|
||||
ConvertOptions options_;
|
||||
MemoryPool* pool_;
|
||||
std::shared_ptr<DataType> type_;
|
||||
};
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
|
||||
#endif // ARROW_CSV_CONVERTER_H
|
||||
@@ -0,0 +1,256 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <numeric>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/csv/chunker.h"
|
||||
#include "arrow/csv/test-common.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/test-util.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
void AssertChunkSize(Chunker& chunker, const std::string& str, uint32_t chunk_size,
|
||||
uint32_t num_rows) {
|
||||
uint32_t actual_chunk_size;
|
||||
ASSERT_OK(
|
||||
chunker.Process(str.data(), static_cast<uint32_t>(str.size()), &actual_chunk_size));
|
||||
ASSERT_EQ(actual_chunk_size, chunk_size);
|
||||
ASSERT_EQ(chunker.num_rows(), num_rows);
|
||||
}
|
||||
|
||||
template <typename IntContainer>
|
||||
void AssertChunking(Chunker& chunker, const std::string& str,
|
||||
const IntContainer& lengths) {
|
||||
uint32_t expected_chunk_size;
|
||||
uint32_t num_rows =
|
||||
static_cast<uint32_t>(std::distance(lengths.begin(), lengths.end()));
|
||||
|
||||
// First chunkize whole CSV block
|
||||
expected_chunk_size =
|
||||
static_cast<uint32_t>(std::accumulate(lengths.begin(), lengths.end(), 0ULL));
|
||||
AssertChunkSize(chunker, str, expected_chunk_size, num_rows);
|
||||
|
||||
// Then chunkize incomplete substrings of the block
|
||||
expected_chunk_size = 0;
|
||||
num_rows = 0;
|
||||
for (const auto length : lengths) {
|
||||
AssertChunkSize(chunker, str.substr(0, expected_chunk_size + length - 1),
|
||||
expected_chunk_size, num_rows);
|
||||
|
||||
++num_rows;
|
||||
expected_chunk_size += static_cast<uint32_t>(length);
|
||||
AssertChunkSize(chunker, str.substr(0, expected_chunk_size), expected_chunk_size,
|
||||
num_rows);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Chunker, Basics) {
|
||||
auto csv = MakeCSVData({"ab,c,\n", "def,,gh\n", ",ij,kl\n"});
|
||||
auto lengths = {6, 8, 7};
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
|
||||
TEST(Chunker, Empty) {
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
{
|
||||
auto csv = MakeCSVData({"\n"});
|
||||
auto lengths = {1};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({"\n\n"});
|
||||
auto lengths = {1, 1};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({",\n"});
|
||||
auto lengths = {2};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({",\n,\n"});
|
||||
auto lengths = {2, 2};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Chunker, Newlines) {
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
{
|
||||
auto csv = MakeCSVData({"a\n", "b\r", "c,d\r\n"});
|
||||
AssertChunkSize(chunker, csv, static_cast<uint32_t>(csv.size()), 3);
|
||||
// Trailing \n after \r is optional
|
||||
AssertChunkSize(chunker, csv.substr(0, csv.size() - 1),
|
||||
static_cast<uint32_t>(csv.size() - 1), 3);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Chunker, MaxNumRows) {
|
||||
Chunker chunker(ParseOptions::Defaults(), 3 /* max_num_rows */);
|
||||
auto csv = MakeCSVData({"ab\n", "c\n", "def\n", "g\n"});
|
||||
auto lengths = {3, 2, 4};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
|
||||
TEST(Chunker, QuotingSimple) {
|
||||
auto csv = MakeCSVData({"1,\",3,\",5\n"});
|
||||
{
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
auto lengths = {csv.size()};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.quoting = false;
|
||||
Chunker chunker(options);
|
||||
auto lengths = {csv.size()};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Chunker, QuotingNewline) {
|
||||
auto csv = MakeCSVData({"a,\"c \n d\",e\n"});
|
||||
{
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
auto lengths = {12};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.quoting = false;
|
||||
Chunker chunker(options);
|
||||
auto lengths = {6, 6};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Chunker, QuotingUnbalanced) {
|
||||
// Quote introduces a quoted field that doesn't end
|
||||
auto csv = MakeCSVData({"a,b\n", "1,\",3,,5\n", "c,d\n"});
|
||||
{
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
auto lengths = {4};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.quoting = false;
|
||||
Chunker chunker(options);
|
||||
auto lengths = {4, 9, 4};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Chunker, QuotingEmpty) {
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
{
|
||||
auto csv = MakeCSVData({"\"\"\n", "a\n"});
|
||||
auto lengths = {3, 2};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({",\"\"\n", "a\n"});
|
||||
auto lengths = {4, 2};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({"\"\",\n", "a\n"});
|
||||
auto lengths = {4, 2};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Chunker, QuotingDouble) {
|
||||
{
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
// 4 quotes is a quoted quote
|
||||
auto csv = MakeCSVData({"\"\"\"\"\n", "a\n"});
|
||||
auto lengths = {5, 2};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Chunker, QuotesSpecial) {
|
||||
// Some non-trivial cases
|
||||
{
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"a,b\"c,d\n", "e\n"});
|
||||
auto lengths = {8, 2};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"a,\"b\" \"c\",d\n", "e\n"});
|
||||
auto lengths = {12, 2};
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Chunker, Escaping) {
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.escaping = true;
|
||||
{
|
||||
auto csv = MakeCSVData({"a\\b,c\n", "d\n"});
|
||||
auto lengths = {6, 2};
|
||||
{
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
Chunker chunker(options);
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({"a\\,b,c\n", "d\n"});
|
||||
auto lengths = {7, 2};
|
||||
{
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
Chunker chunker(options);
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
{
|
||||
// Escaped newline
|
||||
auto csv = MakeCSVData({"a\\\nb\n", "c\n"});
|
||||
{
|
||||
auto lengths = {3, 2, 2};
|
||||
Chunker chunker(ParseOptions::Defaults());
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
{
|
||||
auto lengths = {5, 2};
|
||||
Chunker chunker(options);
|
||||
AssertChunking(chunker, csv, lengths);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,254 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/array.h"
|
||||
#include "arrow/csv/column-builder.h"
|
||||
#include "arrow/csv/test-common.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/test-util.h"
|
||||
#include "arrow/util/thread-pool.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
using internal::GetCpuThreadPool;
|
||||
using internal::TaskGroup;
|
||||
|
||||
void AssertBuilding(const std::shared_ptr<ColumnBuilder>& builder,
|
||||
const std::vector<std::vector<std::string>>& chunks,
|
||||
std::shared_ptr<ChunkedArray>* out) {
|
||||
for (const auto& chunk : chunks) {
|
||||
std::shared_ptr<BlockParser> parser;
|
||||
MakeColumnParser(chunk, &parser);
|
||||
builder->Append(parser);
|
||||
}
|
||||
ASSERT_OK(builder->task_group()->Finish());
|
||||
ASSERT_OK(builder->Finish(out));
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// Test functions begin here
|
||||
|
||||
TEST(ColumnBuilder, Empty) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(int32(), 0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {}, &actual);
|
||||
|
||||
ChunkedArray expected({}, int32());
|
||||
AssertChunkedEqual(*actual, expected);
|
||||
}
|
||||
|
||||
TEST(ColumnBuilder, Basics) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(int32(), 0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{"123", "-456"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<Int32Type>({{123, -456}}, &expected);
|
||||
AssertChunkedEqual(*actual, *expected);
|
||||
}
|
||||
|
||||
TEST(ColumnBuilder, Insert) {
|
||||
// Test ColumnBuilder::Insert()
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(int32(), 0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<BlockParser> parser;
|
||||
std::shared_ptr<ChunkedArray> actual, expected;
|
||||
MakeColumnParser({"456"}, &parser);
|
||||
builder->Insert(1, parser);
|
||||
MakeColumnParser({"123"}, &parser);
|
||||
builder->Insert(0, parser);
|
||||
ASSERT_OK(builder->task_group()->Finish());
|
||||
ASSERT_OK(builder->Finish(&actual));
|
||||
|
||||
ChunkedArrayFromVector<Int32Type>({{123}, {456}}, &expected);
|
||||
AssertChunkedEqual(*actual, *expected);
|
||||
}
|
||||
|
||||
TEST(ColumnBuilder, MultipleChunks) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(int32(), 0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{"1", "2", "3"}, {"4", "5"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<Int32Type>({{1, 2, 3}, {4, 5}}, &expected);
|
||||
AssertChunkedEqual(*actual, *expected);
|
||||
}
|
||||
|
||||
TEST(ColumnBuilder, MultipleChunksParallel) {
|
||||
auto tg = TaskGroup::MakeThreaded(GetCpuThreadPool());
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(int32(), 0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{"1", "2"}, {"3"}, {"4", "5"}, {"6", "7"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<Int32Type>({{1, 2}, {3}, {4, 5}, {6, 7}}, &expected);
|
||||
AssertChunkedEqual(*actual, *expected);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, Empty) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {}, &actual);
|
||||
|
||||
ASSERT_EQ(actual->type()->id(), Type::NA);
|
||||
ASSERT_EQ(actual->num_chunks(), 0);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, SingleChunkNull) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{"", "NA"}}, &actual);
|
||||
|
||||
ASSERT_EQ(actual->type()->id(), Type::NA);
|
||||
ASSERT_EQ(actual->length(), 2);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, MultipleChunkNull) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{"", "NA"}, {""}, {"NaN"}}, &actual);
|
||||
|
||||
ASSERT_EQ(actual->type()->id(), Type::NA);
|
||||
ASSERT_EQ(actual->length(), 4);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, SingleChunkInteger) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{"", "123", "456"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<Int64Type>({{false, true, true}}, {{0, 123, 456}}, &expected);
|
||||
AssertChunkedEqual(*expected, *actual);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, MultipleChunkInteger) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{""}, {"NA", "123", "456"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<Int64Type>({{false}, {false, true, true}}, {{0}, {0, 123, 456}},
|
||||
&expected);
|
||||
AssertChunkedEqual(*expected, *actual);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, SingleChunkReal) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{"", "0.0", "12.5"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<DoubleType>({{false, true, true}}, {{0.0, 0.0, 12.5}},
|
||||
&expected);
|
||||
AssertChunkedEqual(*expected, *actual);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, MultipleChunkReal) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{""}, {"008"}, {"NaN", "12.5"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<DoubleType>({{false}, {true}, {false, true}},
|
||||
{{0.0}, {8.0}, {0.0, 12.5}}, &expected);
|
||||
AssertChunkedEqual(*expected, *actual);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, SingleChunkBinary) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{"", "foo", "bar"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<BinaryType, std::string>({{true, true, true}},
|
||||
{{"", "foo", "bar"}}, &expected);
|
||||
AssertChunkedEqual(*expected, *actual);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, MultipleChunkBinary) {
|
||||
auto tg = TaskGroup::MakeSerial();
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{""}, {"008"}, {"NaN", "bar"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<BinaryType, std::string>(
|
||||
{{true}, {true}, {true, true}}, {{""}, {"008"}, {"NaN", "bar"}}, &expected);
|
||||
AssertChunkedEqual(*expected, *actual);
|
||||
}
|
||||
|
||||
TEST(InferringColumnBuilder, MultipleChunkIntegerParallel) {
|
||||
auto tg = TaskGroup::MakeThreaded(GetCpuThreadPool());
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
ASSERT_OK(ColumnBuilder::Make(0, ConvertOptions::Defaults(), tg, &builder));
|
||||
|
||||
std::shared_ptr<ChunkedArray> actual;
|
||||
AssertBuilding(builder, {{"1", "2"}, {"3"}, {"4", "5"}, {"6", "7"}}, &actual);
|
||||
|
||||
std::shared_ptr<ChunkedArray> expected;
|
||||
ChunkedArrayFromVector<Int64Type>({{1, 2}, {3}, {4, 5}, {6, 7}}, &expected);
|
||||
AssertChunkedEqual(*actual, *expected);
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,97 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include "benchmark/benchmark.h"
|
||||
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
#include "arrow/csv/converter.h"
|
||||
#include "arrow/csv/options.h"
|
||||
#include "arrow/csv/parser.h"
|
||||
#include "arrow/csv/test-common.h"
|
||||
#include "arrow/test-util.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
static std::shared_ptr<BlockParser> BuildInt64Data(int32_t num_rows) {
|
||||
const std::vector<std::string> base_rows = {"123\n", "4\n", "-317005557\n",
|
||||
"\n", "N/A\n", "0\n"};
|
||||
std::vector<std::string> rows;
|
||||
for (int32_t i = 0; i < num_rows; ++i) {
|
||||
rows.push_back(base_rows[i % base_rows.size()]);
|
||||
}
|
||||
|
||||
std::shared_ptr<BlockParser> result;
|
||||
MakeCSVParser(rows, &result);
|
||||
return result;
|
||||
}
|
||||
|
||||
static std::shared_ptr<BlockParser> BuildFloatData(int32_t num_rows) {
|
||||
const std::vector<std::string> base_rows = {"0\n", "123.456\n", "-3170.55766\n", "\n",
|
||||
"N/A\n"};
|
||||
std::vector<std::string> rows;
|
||||
for (int32_t i = 0; i < num_rows; ++i) {
|
||||
rows.push_back(base_rows[i % base_rows.size()]);
|
||||
}
|
||||
|
||||
std::shared_ptr<BlockParser> result;
|
||||
MakeCSVParser(rows, &result);
|
||||
return result;
|
||||
}
|
||||
|
||||
static void BenchmarkConversion(benchmark::State& state, // NOLINT non-const reference
|
||||
BlockParser& parser,
|
||||
const std::shared_ptr<DataType>& type,
|
||||
ConvertOptions options) {
|
||||
std::shared_ptr<Converter> converter;
|
||||
ABORT_NOT_OK(Converter::Make(type, options, &converter));
|
||||
|
||||
while (state.KeepRunning()) {
|
||||
std::shared_ptr<Array> result;
|
||||
ABORT_NOT_OK(converter->Convert(parser, 0 /* col_index */, &result));
|
||||
if (result->length() != parser.num_rows()) {
|
||||
std::cerr << "Conversion incomplete\n";
|
||||
std::abort();
|
||||
}
|
||||
}
|
||||
|
||||
state.SetItemsProcessed(state.iterations() * parser.num_rows());
|
||||
}
|
||||
|
||||
static void BM_Int64Conversion(benchmark::State& state) { // NOLINT non-const reference
|
||||
const int32_t num_rows = 10000;
|
||||
auto parser = BuildInt64Data(num_rows);
|
||||
auto options = ConvertOptions::Defaults();
|
||||
|
||||
BenchmarkConversion(state, *parser, int64(), options);
|
||||
}
|
||||
|
||||
static void BM_FloatConversion(benchmark::State& state) { // NOLINT non-const reference
|
||||
const int32_t num_rows = 10000;
|
||||
auto parser = BuildFloatData(num_rows);
|
||||
auto options = ConvertOptions::Defaults();
|
||||
|
||||
BenchmarkConversion(state, *parser, float64(), options);
|
||||
}
|
||||
|
||||
BENCHMARK(BM_Int64Conversion)->Repetitions(3);
|
||||
BENCHMARK(BM_FloatConversion)->Repetitions(3);
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,189 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/array.h"
|
||||
#include "arrow/csv/converter.h"
|
||||
#include "arrow/csv/test-common.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/test-util.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
// All recognized (non-empty) null values
|
||||
std::vector<std::string> AllNulls() {
|
||||
return {"#N/A\n", "#N/A N/A\n", "#NA\n", "-1.#IND\n", "-1.#QNAN\n", "-NaN\n",
|
||||
"-nan\n", "1.#IND\n", "1.#QNAN\n", "N/A\n", "NA\n", "NULL\n",
|
||||
"NaN\n", "n/a\n", "nan\n", "null\n"};
|
||||
}
|
||||
|
||||
template <typename DATA_TYPE, typename C_TYPE>
|
||||
void AssertConversion(const std::shared_ptr<DataType>& type,
|
||||
const std::vector<std::string>& csv_string,
|
||||
const std::vector<std::vector<C_TYPE>>& expected) {
|
||||
std::shared_ptr<BlockParser> parser;
|
||||
std::shared_ptr<Converter> converter;
|
||||
std::shared_ptr<Array> array, expected_array;
|
||||
|
||||
ASSERT_OK(Converter::Make(type, ConvertOptions::Defaults(), &converter));
|
||||
|
||||
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));
|
||||
ArrayFromVector<DATA_TYPE, C_TYPE>(type, expected[col_index], &expected_array);
|
||||
AssertArraysEqual(*expected_array, *array);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DATA_TYPE, typename C_TYPE>
|
||||
void AssertConversion(const std::shared_ptr<DataType>& type,
|
||||
const std::vector<std::string>& csv_string,
|
||||
const std::vector<std::vector<C_TYPE>>& expected,
|
||||
const std::vector<std::vector<bool>>& is_valid) {
|
||||
std::shared_ptr<BlockParser> parser;
|
||||
std::shared_ptr<Converter> converter;
|
||||
std::shared_ptr<Array> array, expected_array;
|
||||
|
||||
ASSERT_OK(Converter::Make(type, ConvertOptions::Defaults(), &converter));
|
||||
|
||||
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));
|
||||
ArrayFromVector<DATA_TYPE, C_TYPE>(type, is_valid[col_index], expected[col_index],
|
||||
&expected_array);
|
||||
AssertArraysEqual(*expected_array, *array);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DATA_TYPE, typename C_TYPE>
|
||||
void AssertConversionAllNulls(const std::shared_ptr<DataType>& type) {
|
||||
std::vector<std::string> nulls = AllNulls();
|
||||
std::vector<bool> is_valid(nulls.size(), false);
|
||||
std::vector<C_TYPE> values(nulls.size());
|
||||
AssertConversion<DATA_TYPE, C_TYPE>(type, nulls, {values}, {is_valid});
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// Test functions begin here
|
||||
|
||||
TEST(BinaryConversion, Basics) {
|
||||
AssertConversion<BinaryType, std::string>(binary(), {"ab,cde\n", ",gh\n"},
|
||||
{{"ab", ""}, {"cde", "gh"}});
|
||||
}
|
||||
|
||||
TEST(StringConversion, Basics) {
|
||||
AssertConversion<StringType, std::string>(utf8(), {"ab,cde\n", ",gh\n"},
|
||||
{{"ab", ""}, {"cde", "gh"}});
|
||||
}
|
||||
|
||||
TEST(FixedSizeBinaryConversion, Basics) {
|
||||
AssertConversion<FixedSizeBinaryType, std::string>(
|
||||
fixed_size_binary(2), {"ab,cd\n", "gh,ij\n"}, {{"ab", "gh"}, {"cd", "ij"}});
|
||||
}
|
||||
|
||||
TEST(FixedSizeBinaryConversion, Errors) {
|
||||
std::shared_ptr<BlockParser> parser;
|
||||
std::shared_ptr<Converter> converter;
|
||||
std::shared_ptr<Array> array;
|
||||
std::shared_ptr<DataType> type = fixed_size_binary(2);
|
||||
|
||||
ASSERT_OK(Converter::Make(type, ConvertOptions::Defaults(), &converter));
|
||||
|
||||
MakeCSVParser({"ab,cd\n", "g,ij\n"}, &parser);
|
||||
ASSERT_RAISES(Invalid, converter->Convert(*parser, 0, &array));
|
||||
ASSERT_OK(converter->Convert(*parser, 1, &array));
|
||||
}
|
||||
|
||||
TEST(NullConversion, Basics) {
|
||||
std::shared_ptr<BlockParser> parser;
|
||||
std::shared_ptr<Converter> converter;
|
||||
std::shared_ptr<Array> array;
|
||||
std::shared_ptr<DataType> type = null();
|
||||
|
||||
ASSERT_OK(Converter::Make(type, ConvertOptions::Defaults(), &converter));
|
||||
|
||||
MakeCSVParser({"NA,z\n", ",0\n"}, &parser);
|
||||
ASSERT_OK(converter->Convert(*parser, 0, &array));
|
||||
ASSERT_EQ(array->type()->id(), Type::NA);
|
||||
ASSERT_EQ(array->length(), 2);
|
||||
ASSERT_RAISES(Invalid, converter->Convert(*parser, 1, &array));
|
||||
}
|
||||
|
||||
TEST(IntegerConversion, Basics) {
|
||||
AssertConversion<Int8Type, int8_t>(int8(), {"12,34\n", "0,-128\n"},
|
||||
{{12, 0}, {34, -128}});
|
||||
AssertConversion<Int64Type, int64_t>(
|
||||
int64(), {"12,34\n", "9223372036854775807,-9223372036854775808\n"},
|
||||
{{12, 9223372036854775807LL}, {34, -9223372036854775807LL - 1}});
|
||||
|
||||
AssertConversion<UInt16Type, uint16_t>(uint16(), {"12,34\n", "0,65535\n"},
|
||||
{{12, 0}, {34, 65535}});
|
||||
AssertConversion<UInt64Type, uint64_t>(uint64(),
|
||||
{"12,34\n", "0,18446744073709551615\n"},
|
||||
{{12, 0}, {34, 18446744073709551615ULL}});
|
||||
}
|
||||
|
||||
TEST(IntegerConversion, Nulls) {
|
||||
AssertConversion<Int8Type, int8_t>(int8(), {"12,34\n", ",-128\n"},
|
||||
{{12, 0}, {34, -128}},
|
||||
{{true, false}, {true, true}});
|
||||
|
||||
AssertConversionAllNulls<Int8Type, int8_t>(int8());
|
||||
}
|
||||
|
||||
TEST(FloatingPointConversion, Basics) {
|
||||
AssertConversion<FloatType, float>(float32(), {"12,34.5\n", "0,-1e30\n"},
|
||||
{{12., 0.}, {34.5, -1e30f}});
|
||||
AssertConversion<DoubleType, double>(float64(), {"12,34.5\n", "0,-1e100\n"},
|
||||
{{12., 0.}, {34.5, -1e100}});
|
||||
}
|
||||
|
||||
TEST(FloatingPointConversion, Nulls) {
|
||||
AssertConversion<FloatType, float>(float32(), {"1.5,0.\n", ",-1e10\n"},
|
||||
{{1.5, 0.}, {0., -1e10f}},
|
||||
{{true, false}, {true, true}});
|
||||
|
||||
AssertConversionAllNulls<DoubleType, double>(float64());
|
||||
}
|
||||
|
||||
TEST(BooleanConversion, Basics) {
|
||||
// XXX we may want to accept more bool-like values
|
||||
AssertConversion<BooleanType, bool>(boolean(), {"true,false\n", "1,0\n"},
|
||||
{{true, true}, {false, false}});
|
||||
}
|
||||
|
||||
TEST(BooleanConversion, Nulls) {
|
||||
AssertConversion<BooleanType, bool>(boolean(), {"true,\n", "1,0\n"},
|
||||
{{true, true}, {false, false}},
|
||||
{{true, true}, {false, true}});
|
||||
}
|
||||
|
||||
TEST(DecimalConversion, NotImplemented) {
|
||||
std::shared_ptr<Converter> converter;
|
||||
ASSERT_RAISES(NotImplemented,
|
||||
Converter::Make(decimal(12, 3), ConvertOptions::Defaults(), &converter));
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,147 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include "benchmark/benchmark.h"
|
||||
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
#include "arrow/csv/chunker.h"
|
||||
#include "arrow/csv/options.h"
|
||||
#include "arrow/csv/parser.h"
|
||||
#include "arrow/test-util.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
static std::string BuildQuotedData(int32_t num_rows = 10000) {
|
||||
std::string one_row = "abc,\"d,f\",12.34,\n";
|
||||
std::stringstream ss;
|
||||
for (int32_t i = 0; i < num_rows; ++i) {
|
||||
ss << one_row;
|
||||
}
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
static std::string BuildEscapedData(int32_t num_rows = 10000) {
|
||||
std::string one_row = "abc,d\\,f,12.34,\n";
|
||||
std::stringstream ss;
|
||||
for (int32_t i = 0; i < num_rows; ++i) {
|
||||
ss << one_row;
|
||||
}
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
static void BenchmarkCSVChunking(benchmark::State& state, // NOLINT non-const reference
|
||||
const std::string& csv, int32_t num_rows,
|
||||
ParseOptions options) {
|
||||
Chunker chunker(options, num_rows + 1);
|
||||
|
||||
while (state.KeepRunning()) {
|
||||
uint32_t chunk_size;
|
||||
ABORT_NOT_OK(
|
||||
chunker.Process(csv.data(), static_cast<uint32_t>(csv.size()), &chunk_size));
|
||||
if (chunk_size != csv.size() || chunker.num_rows() != num_rows) {
|
||||
std::cerr << "Parsing incomplete\n";
|
||||
std::abort();
|
||||
}
|
||||
}
|
||||
state.SetBytesProcessed(state.iterations() * csv.size());
|
||||
}
|
||||
|
||||
static void BM_ChunkCSVQuotedBlock(
|
||||
benchmark::State& state) { // NOLINT non-const reference
|
||||
const int32_t num_rows = 5000;
|
||||
auto csv = BuildQuotedData(num_rows);
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.quoting = true;
|
||||
options.escaping = false;
|
||||
|
||||
BenchmarkCSVChunking(state, csv, num_rows, options);
|
||||
}
|
||||
|
||||
static void BM_ChunkCSVEscapedBlock(
|
||||
benchmark::State& state) { // NOLINT non-const reference
|
||||
const int32_t num_rows = 5000;
|
||||
auto csv = BuildEscapedData(num_rows);
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.quoting = false;
|
||||
options.escaping = true;
|
||||
|
||||
BenchmarkCSVChunking(state, csv, num_rows, options);
|
||||
}
|
||||
|
||||
static void BenchmarkCSVParsing(benchmark::State& state, // NOLINT non-const reference
|
||||
const std::string& csv, int32_t num_rows,
|
||||
ParseOptions options) {
|
||||
BlockParser parser(options, -1, num_rows + 1);
|
||||
|
||||
while (state.KeepRunning()) {
|
||||
uint32_t parsed_size;
|
||||
ABORT_NOT_OK(
|
||||
parser.Parse(csv.data(), static_cast<uint32_t>(csv.size()), &parsed_size));
|
||||
if (parsed_size != csv.size() || parser.num_rows() != num_rows) {
|
||||
std::cerr << "Parsing incomplete\n";
|
||||
std::abort();
|
||||
}
|
||||
// Include performance of visiting the parsed values, as that might
|
||||
// vary depending on the parser's internal data structures.
|
||||
bool dummy_quoted = false;
|
||||
uint32_t dummy_size = 0;
|
||||
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) {
|
||||
dummy_size += size;
|
||||
dummy_quoted ^= quoted;
|
||||
return Status::OK();
|
||||
};
|
||||
for (int32_t col = 0; col < parser.num_cols(); ++col) {
|
||||
ABORT_NOT_OK(parser.VisitColumn(col, visit));
|
||||
benchmark::DoNotOptimize(dummy_size);
|
||||
benchmark::DoNotOptimize(dummy_quoted);
|
||||
}
|
||||
}
|
||||
state.SetBytesProcessed(state.iterations() * csv.size());
|
||||
}
|
||||
|
||||
static void BM_ParseCSVQuotedBlock(
|
||||
benchmark::State& state) { // NOLINT non-const reference
|
||||
const int32_t num_rows = 5000;
|
||||
auto csv = BuildQuotedData(num_rows);
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.quoting = true;
|
||||
options.escaping = false;
|
||||
|
||||
BenchmarkCSVParsing(state, csv, num_rows, options);
|
||||
}
|
||||
|
||||
static void BM_ParseCSVEscapedBlock(
|
||||
benchmark::State& state) { // NOLINT non-const reference
|
||||
const int32_t num_rows = 5000;
|
||||
auto csv = BuildEscapedData(num_rows);
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.quoting = false;
|
||||
options.escaping = true;
|
||||
|
||||
BenchmarkCSVParsing(state, csv, num_rows, options);
|
||||
}
|
||||
|
||||
BENCHMARK(BM_ChunkCSVQuotedBlock)->Repetitions(3)->Unit(benchmark::kMicrosecond);
|
||||
BENCHMARK(BM_ChunkCSVEscapedBlock)->Repetitions(3)->Unit(benchmark::kMicrosecond);
|
||||
BENCHMARK(BM_ParseCSVQuotedBlock)->Repetitions(3)->Unit(benchmark::kMicrosecond);
|
||||
BENCHMARK(BM_ParseCSVEscapedBlock)->Repetitions(3)->Unit(benchmark::kMicrosecond);
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,421 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/csv/parser.h"
|
||||
#include "arrow/csv/test-common.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/test-util.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
// Read the column with the given index out of the BlockParser.
|
||||
void GetColumn(const BlockParser& parser, int32_t col_index,
|
||||
std::vector<std::string>* out, std::vector<bool>* out_quoted = nullptr) {
|
||||
std::vector<std::string> values;
|
||||
std::vector<bool> quoted_values;
|
||||
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
|
||||
values.push_back(std::string(reinterpret_cast<const char*>(data), size));
|
||||
if (out_quoted) {
|
||||
quoted_values.push_back(quoted);
|
||||
}
|
||||
return Status::OK();
|
||||
};
|
||||
ASSERT_OK(parser.VisitColumn(col_index, visit));
|
||||
*out = std::move(values);
|
||||
if (out_quoted) {
|
||||
*out_quoted = std::move(quoted_values);
|
||||
}
|
||||
}
|
||||
|
||||
Status Parse(BlockParser& parser, const std::string& str, uint32_t* out_size) {
|
||||
const char* data = str.data();
|
||||
uint32_t size = static_cast<uint32_t>(str.length());
|
||||
return parser.Parse(data, size, out_size);
|
||||
}
|
||||
|
||||
Status ParseFinal(BlockParser& parser, const std::string& str, uint32_t* out_size) {
|
||||
const char* data = str.data();
|
||||
uint32_t size = static_cast<uint32_t>(str.length());
|
||||
return parser.ParseFinal(data, size, out_size);
|
||||
}
|
||||
|
||||
void AssertParseOk(BlockParser& parser, const std::string& str) {
|
||||
uint32_t parsed_size = static_cast<uint32_t>(-1);
|
||||
ASSERT_OK(Parse(parser, str, &parsed_size));
|
||||
ASSERT_EQ(parsed_size, str.size());
|
||||
}
|
||||
|
||||
void AssertParseFinal(BlockParser& parser, const std::string& str) {
|
||||
uint32_t parsed_size = static_cast<uint32_t>(-1);
|
||||
ASSERT_OK(ParseFinal(parser, str, &parsed_size));
|
||||
ASSERT_EQ(parsed_size, str.size());
|
||||
}
|
||||
|
||||
void AssertParsePartial(BlockParser& parser, const std::string& str,
|
||||
uint32_t expected_size) {
|
||||
uint32_t parsed_size = static_cast<uint32_t>(-1);
|
||||
ASSERT_OK(Parse(parser, str, &parsed_size));
|
||||
ASSERT_EQ(parsed_size, expected_size);
|
||||
}
|
||||
|
||||
void AssertColumnEq(const BlockParser& parser, int32_t col_index,
|
||||
const std::vector<std::string> expected) {
|
||||
std::vector<std::string> values;
|
||||
GetColumn(parser, col_index, &values);
|
||||
ASSERT_EQ(parser.num_rows(), expected.size());
|
||||
ASSERT_EQ(values, expected);
|
||||
}
|
||||
|
||||
void AssertColumnEq(const BlockParser& parser, int32_t col_index,
|
||||
const std::vector<std::string> expected,
|
||||
const std::vector<bool> expected_quoted) {
|
||||
std::vector<std::string> values;
|
||||
std::vector<bool> quoted;
|
||||
GetColumn(parser, col_index, &values, "ed);
|
||||
ASSERT_EQ(parser.num_rows(), expected.size());
|
||||
ASSERT_EQ(values, expected);
|
||||
ASSERT_EQ(quoted, expected_quoted);
|
||||
}
|
||||
|
||||
void AssertColumnsEq(const BlockParser& parser,
|
||||
const std::vector<std::vector<std::string>> expected) {
|
||||
ASSERT_EQ(parser.num_cols(), expected.size());
|
||||
for (int32_t col_index = 0; col_index < parser.num_cols(); ++col_index) {
|
||||
AssertColumnEq(parser, col_index, expected[col_index]);
|
||||
}
|
||||
}
|
||||
|
||||
void AssertColumnsEq(const BlockParser& parser,
|
||||
const std::vector<std::vector<std::string>> expected,
|
||||
const std::vector<std::vector<bool>> quoted) {
|
||||
ASSERT_EQ(parser.num_cols(), expected.size());
|
||||
for (int32_t col_index = 0; col_index < parser.num_cols(); ++col_index) {
|
||||
AssertColumnEq(parser, col_index, expected[col_index], quoted[col_index]);
|
||||
}
|
||||
uint32_t total_bytes = 0;
|
||||
for (const auto& col : expected) {
|
||||
for (const auto& field : col) {
|
||||
total_bytes += static_cast<uint32_t>(field.size());
|
||||
}
|
||||
}
|
||||
ASSERT_EQ(total_bytes, parser.num_bytes());
|
||||
}
|
||||
|
||||
TEST(BlockParser, Basics) {
|
||||
auto csv = MakeCSVData({"ab,cd,\n", "ef,,gh\n", ",ij,kl\n"});
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"ab", "ef", ""}, {"cd", "", "ij"}, {"", "gh", "kl"}});
|
||||
}
|
||||
|
||||
TEST(BlockParser, Empty) {
|
||||
{
|
||||
auto csv = MakeCSVData({"\n"});
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{""}});
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({"\n\n"});
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"", ""}});
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({",\n"});
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{""}, {""}});
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({",\n,\n"});
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"", ""}, {"", ""}});
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BlockParser, Whitespace) {
|
||||
// Non-newline whitespace is preserved
|
||||
auto csv = MakeCSVData({"a b, cd, \n", " ef, \t,gh\n"});
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a b", " ef"}, {" cd", " \t"}, {" ", "gh"}});
|
||||
}
|
||||
|
||||
TEST(BlockParser, Newlines) {
|
||||
auto csv = MakeCSVData({"a,b\n", "c,d\r\n", "e,f\r", "g,h\r"});
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a", "c", "e", "g"}, {"b", "d", "f", "h"}});
|
||||
}
|
||||
|
||||
TEST(BlockParser, MaxNumRows) {
|
||||
auto csv = MakeCSVData({"a\n", "b\n", "c\n", "d\n"});
|
||||
BlockParser parser(ParseOptions::Defaults(), -1, 3 /* max_num_rows */);
|
||||
|
||||
AssertParsePartial(parser, csv, 6);
|
||||
AssertColumnsEq(parser, {{"a", "b", "c"}});
|
||||
|
||||
AssertParseOk(parser, csv.substr(6));
|
||||
AssertColumnsEq(parser, {{"d"}});
|
||||
|
||||
AssertParseOk(parser, csv.substr(8));
|
||||
AssertColumnsEq(parser, {{}});
|
||||
}
|
||||
|
||||
TEST(BlockParser, TruncatedData) {
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"a,b\n", "c,d\n"});
|
||||
for (auto trim : {1, 2, 3}) {
|
||||
AssertParsePartial(parser, csv.substr(0, csv.length() - trim), 4);
|
||||
AssertColumnsEq(parser, {{"a"}, {"b"}});
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BlockParser, Final) {
|
||||
// Tests for ParseFinal()
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"ab,cd\n", "ef,gh\n"});
|
||||
AssertParseFinal(parser, csv);
|
||||
AssertColumnsEq(parser, {{"ab", "ef"}, {"cd", "gh"}});
|
||||
|
||||
// Same without newline
|
||||
csv = MakeCSVData({"ab,cd\n", "ef,gh"});
|
||||
AssertParseFinal(parser, csv);
|
||||
AssertColumnsEq(parser, {{"ab", "ef"}, {"cd", "gh"}});
|
||||
|
||||
// Same with empty last item
|
||||
csv = MakeCSVData({"ab,cd\n", "ef,"});
|
||||
AssertParseFinal(parser, csv);
|
||||
AssertColumnsEq(parser, {{"ab", "ef"}, {"cd", ""}});
|
||||
|
||||
// Same with single line
|
||||
csv = MakeCSVData({"ab,cd"});
|
||||
AssertParseFinal(parser, csv);
|
||||
AssertColumnsEq(parser, {{"ab"}, {"cd"}});
|
||||
}
|
||||
|
||||
TEST(BlockParser, FinalTruncatedData) {
|
||||
// Test ParseFinal() with truncated data
|
||||
uint32_t out_size;
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"ab,cd\n", "ef"});
|
||||
Status st = ParseFinal(parser, csv, &out_size);
|
||||
ASSERT_RAISES(Invalid, st);
|
||||
}
|
||||
|
||||
TEST(BlockParser, QuotingSimple) {
|
||||
auto csv = MakeCSVData({"1,\",3,\",5\n"});
|
||||
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"1"}, {",3,"}, {"5"}},
|
||||
{{false}, {true}, {false}} /* quoted */);
|
||||
}
|
||||
{
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.quoting = false;
|
||||
BlockParser parser(options);
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"1"}, {"\""}, {"3"}, {"\""}, {"5"}},
|
||||
{{false}, {false}, {false}, {false}, {false}} /* quoted */);
|
||||
}
|
||||
{
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.quote_char = 'Z';
|
||||
BlockParser parser(options);
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"1"}, {"\""}, {"3"}, {"\""}, {"5"}},
|
||||
{{false}, {false}, {false}, {false}, {false}} /* quoted */);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BlockParser, QuotingNewline) {
|
||||
auto csv = MakeCSVData({"a,\"c \n d\",e\n"});
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a"}, {"c \n d"}, {"e"}},
|
||||
{{false}, {true}, {false}} /* quoted */);
|
||||
}
|
||||
|
||||
TEST(BlockParser, QuotingUnbalanced) {
|
||||
// Quote introduces a quoted field that doesn't end
|
||||
auto csv = MakeCSVData({"a,b\n", "1,\",3,,5\n"});
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParsePartial(parser, csv, 4);
|
||||
AssertColumnsEq(parser, {{"a"}, {"b"}}, {{false}, {false}} /* quoted */);
|
||||
}
|
||||
|
||||
TEST(BlockParser, QuotingEmpty) {
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"\"\"\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{""}}, {{true}} /* quoted */);
|
||||
}
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({",\"\"\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{""}, {""}}, {{false}, {true}} /* quoted */);
|
||||
}
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"\"\",\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{""}, {""}}, {{true}, {false}} /* quoted */);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BlockParser, QuotingDouble) {
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
// 4 quotes is a quoted quote
|
||||
auto csv = MakeCSVData({"\"\"\"\"\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"\""}}, {{true}} /* quoted */);
|
||||
}
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
// 4 quotes is a quoted quote
|
||||
auto csv = MakeCSVData({"a,\"\"\"\",b\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a"}, {"\""}, {"b"}},
|
||||
{{false}, {true}, {false}} /* quoted */);
|
||||
}
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
// 6 quotes is two quoted quotes
|
||||
auto csv = MakeCSVData({"\"\"\"\"\"\"\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"\"\""}}, {{true}} /* quoted */);
|
||||
}
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
// 6 quotes is two quoted quotes
|
||||
auto csv = MakeCSVData({"a,\"\"\"\"\"\",b\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a"}, {"\"\""}, {"b"}},
|
||||
{{false}, {true}, {false}} /* quoted */);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BlockParser, QuotesAndMore) {
|
||||
// There may be trailing data after the quoted part of a field
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"a,\"b\"c,d\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a"}, {"bc"}, {"d"}},
|
||||
{{false}, {true}, {false}} /* quoted */);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BlockParser, QuotesSpecial) {
|
||||
// Some non-trivial cases
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"a,b\"c,d\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a"}, {"b\"c"}, {"d"}},
|
||||
{{false}, {false}, {false}} /* quoted */);
|
||||
}
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"a,\"b\" \"c\",d\n"});
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a"}, {"b \"c\""}, {"d"}},
|
||||
{{false}, {true}, {false}} /* quoted */);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BlockParser, MismatchingNumColumns) {
|
||||
uint32_t out_size;
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
auto csv = MakeCSVData({"a,b\nc\n"});
|
||||
Status st = Parse(parser, csv, &out_size);
|
||||
ASSERT_RAISES(Invalid, st);
|
||||
}
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults(), 2 /* num_cols */);
|
||||
auto csv = MakeCSVData({"a\n"});
|
||||
Status st = Parse(parser, csv, &out_size);
|
||||
ASSERT_RAISES(Invalid, st);
|
||||
}
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults(), 2 /* num_cols */);
|
||||
auto csv = MakeCSVData({"a,b,c\n"});
|
||||
Status st = Parse(parser, csv, &out_size);
|
||||
ASSERT_RAISES(Invalid, st);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BlockParser, Escaping) {
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.escaping = true;
|
||||
|
||||
{
|
||||
auto csv = MakeCSVData({"a\\b,c\n"});
|
||||
{
|
||||
BlockParser parser(ParseOptions::Defaults());
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a\\b"}, {"c"}});
|
||||
}
|
||||
{
|
||||
BlockParser parser(options);
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"ab"}, {"c"}});
|
||||
}
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({"a\\,b,c\n"});
|
||||
BlockParser parser(options);
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a,b"}, {"c"}});
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BlockParser, QuotedEscape) {
|
||||
auto options = ParseOptions::Defaults();
|
||||
options.escaping = true;
|
||||
|
||||
{
|
||||
auto csv = MakeCSVData({"\"a\\,b\",c\n"});
|
||||
BlockParser parser(options);
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a,b"}, {"c"}}, {{true}, {false}} /* quoted */);
|
||||
}
|
||||
{
|
||||
auto csv = MakeCSVData({"\"a\\\"b\",c\n"});
|
||||
BlockParser parser(options);
|
||||
AssertParseOk(parser, csv);
|
||||
AssertColumnsEq(parser, {{"a\"b"}, {"c"}}, {{true}, {false}} /* quoted */);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,30 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include "arrow/csv/options.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
ParseOptions ParseOptions::Defaults() { return ParseOptions(); }
|
||||
|
||||
ConvertOptions ConvertOptions::Defaults() { return ConvertOptions(); }
|
||||
|
||||
ReadOptions ReadOptions::Defaults() { return ReadOptions(); }
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,71 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef ARROW_CSV_OPTIONS_H
|
||||
#define ARROW_CSV_OPTIONS_H
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
#include "arrow/util/visibility.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
struct ARROW_EXPORT ParseOptions {
|
||||
// Parsing options
|
||||
|
||||
// Field delimiter
|
||||
char delimiter = ',';
|
||||
// Whether quoting is used
|
||||
bool quoting = true;
|
||||
// Quoting character (if `quoting` is true)
|
||||
char quote_char = '"';
|
||||
// Whether a quote inside a value is double-quoted
|
||||
bool double_quote = true;
|
||||
// Whether escaping is used
|
||||
bool escaping = false;
|
||||
// Escaping character (if `escaping` is true)
|
||||
char escape_char = '\\';
|
||||
|
||||
// XXX Should this be in ReadOptions?
|
||||
// Number of header rows to skip
|
||||
int32_t header_rows = 1;
|
||||
|
||||
static ParseOptions Defaults();
|
||||
};
|
||||
|
||||
struct ARROW_EXPORT ConvertOptions {
|
||||
static ConvertOptions Defaults();
|
||||
};
|
||||
|
||||
struct ARROW_EXPORT ReadOptions {
|
||||
// Reader options
|
||||
|
||||
// Whether to use the global CPU thread pool
|
||||
bool use_threads = true;
|
||||
// Block size we request from the IO layer
|
||||
int32_t block_size = 1 << 20; // 1 MB
|
||||
// Max num rows per array chunk
|
||||
int32_t num_rows = 100000;
|
||||
|
||||
static ReadOptions Defaults();
|
||||
};
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
|
||||
#endif // ARROW_CSV_OPTIONS_H
|
||||
@@ -0,0 +1,240 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include "arrow/csv/parser.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/util/logging.h"
|
||||
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
static Status ParseError(const char* message) {
|
||||
std::stringstream ss;
|
||||
ss << "CSV parse error: " << message;
|
||||
return Status::Invalid(ss.str());
|
||||
}
|
||||
|
||||
static Status MismatchingColumns(int32_t expected, int32_t actual) {
|
||||
char s[50];
|
||||
snprintf(s, sizeof(s), "Expected %d columns, got %d", expected, actual);
|
||||
return ParseError(s);
|
||||
}
|
||||
|
||||
BlockParser::BlockParser(ParseOptions options, int32_t num_cols, int32_t max_num_rows)
|
||||
: options_(options), num_cols_(num_cols), max_num_rows_(max_num_rows) {}
|
||||
|
||||
Status BlockParser::ParseLine(const char* data, const char* data_end, bool is_final,
|
||||
const char** out_data) {
|
||||
int32_t num_cols = 0;
|
||||
char c;
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
bool quoted;
|
||||
#endif
|
||||
|
||||
auto saved_parsed_size = parsed_.size();
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
auto saved_values_size = values_.size();
|
||||
#else
|
||||
auto saved_offsets_size = offsets_.size();
|
||||
auto saved_quoted_size = quoted_.size();
|
||||
#endif
|
||||
|
||||
// Subroutines to manage parser state
|
||||
auto InitField = [&]() {};
|
||||
auto PushFieldChar = [&](char c) { parsed_.push_back(static_cast<uint8_t>(c)); };
|
||||
auto FinishField = [&]() {
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
ValueDesc v = {static_cast<uint32_t>(parsed_.size()) & 0x7fffffffU, quoted};
|
||||
values_.push_back(v);
|
||||
#else
|
||||
offsets_.push_back(parsed_.size());
|
||||
#endif
|
||||
++num_cols;
|
||||
};
|
||||
auto RewindState = [&]() {
|
||||
parsed_.resize(saved_parsed_size);
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
values_.resize(saved_values_size);
|
||||
#else
|
||||
offsets_.resize(saved_offsets_size);
|
||||
quoted_.resize(saved_quoted_size);
|
||||
#endif
|
||||
};
|
||||
|
||||
// The parsing state machine
|
||||
|
||||
FieldStart:
|
||||
// At the start of a field
|
||||
InitField();
|
||||
// Quoting is only recognized at start of field
|
||||
if (options_.quoting && data != data_end && *data == options_.quote_char) {
|
||||
++data;
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
quoted = true;
|
||||
#else
|
||||
quoted_.push_back(true);
|
||||
#endif
|
||||
goto InQuotedField;
|
||||
} else {
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
quoted = false;
|
||||
#else
|
||||
quoted_.push_back(false);
|
||||
#endif
|
||||
goto InField;
|
||||
}
|
||||
|
||||
InField:
|
||||
// Inside a non-quoted part of a field
|
||||
if (data == data_end) {
|
||||
goto AbortLine;
|
||||
}
|
||||
c = *data++;
|
||||
if (options_.escaping && c == options_.escape_char) {
|
||||
if (data == data_end) {
|
||||
goto AbortLine;
|
||||
}
|
||||
c = *data++;
|
||||
PushFieldChar(c);
|
||||
goto InField;
|
||||
}
|
||||
if (c == '\n' || c == '\r') {
|
||||
goto LineEnd;
|
||||
}
|
||||
if (c == options_.delimiter) {
|
||||
goto FieldEnd;
|
||||
}
|
||||
PushFieldChar(c);
|
||||
goto InField;
|
||||
|
||||
InQuotedField:
|
||||
// Inside a quoted part of a field
|
||||
if (data == data_end) {
|
||||
goto AbortLine;
|
||||
}
|
||||
c = *data++;
|
||||
if (options_.escaping && c == options_.escape_char) {
|
||||
if (data == data_end) {
|
||||
goto AbortLine;
|
||||
}
|
||||
c = *data++;
|
||||
PushFieldChar(c);
|
||||
goto InQuotedField;
|
||||
}
|
||||
if (c == options_.quote_char) {
|
||||
if (options_.double_quote && data < data_end && *data == options_.quote_char) {
|
||||
// Double-quoting
|
||||
++data;
|
||||
} else {
|
||||
// End of single-quoting
|
||||
goto InField;
|
||||
}
|
||||
}
|
||||
PushFieldChar(c);
|
||||
goto InQuotedField;
|
||||
|
||||
FieldEnd:
|
||||
// At the end of a field
|
||||
FinishField();
|
||||
goto FieldStart;
|
||||
|
||||
LineEnd:
|
||||
// At the end of line, possibly in the middle of the newline separator
|
||||
FinishField();
|
||||
if (data < data_end && data[-1] == '\r' && *data == '\n') {
|
||||
data++;
|
||||
}
|
||||
if (num_cols_ == -1) {
|
||||
num_cols_ = num_cols;
|
||||
} else if (num_cols != num_cols_) {
|
||||
return MismatchingColumns(num_cols_, num_cols);
|
||||
}
|
||||
*out_data = data;
|
||||
return Status::OK();
|
||||
|
||||
AbortLine:
|
||||
// Not a full line except perhaps if in final block
|
||||
if (is_final) {
|
||||
FinishField();
|
||||
if (num_cols_ == -1) {
|
||||
num_cols_ = num_cols;
|
||||
} else if (num_cols != num_cols_) {
|
||||
return MismatchingColumns(num_cols_, num_cols);
|
||||
}
|
||||
*out_data = data;
|
||||
return Status::OK();
|
||||
}
|
||||
// Truncated line at end of block, rewind parsed state
|
||||
RewindState();
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status BlockParser::DoParse(const char* start, uint32_t size, bool is_final,
|
||||
uint32_t* out_size) {
|
||||
num_rows_ = 0;
|
||||
// These don't shrink the allocated capacity, so reuses of BlockParser
|
||||
// avoid most allocations when appending data
|
||||
parsed_.clear();
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
values_.clear();
|
||||
values_.push_back({0, false});
|
||||
#else
|
||||
offsets_.clear();
|
||||
offsets_.push_back(0);
|
||||
quoted_.clear();
|
||||
#endif
|
||||
|
||||
const char* data = start;
|
||||
const char* data_end = start + size;
|
||||
|
||||
while (data < data_end && num_rows_ < max_num_rows_) {
|
||||
const char* line_end = data;
|
||||
RETURN_NOT_OK(ParseLine(data, data_end, is_final, &line_end));
|
||||
if (line_end == data) {
|
||||
// Cannot parse any further
|
||||
break;
|
||||
}
|
||||
data = line_end;
|
||||
++num_rows_;
|
||||
}
|
||||
if (num_cols_ == -1) {
|
||||
DCHECK_EQ(num_rows_, 0);
|
||||
}
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
DCHECK_EQ(values_.size(), 1 + static_cast<size_t>(num_rows_ * num_cols_));
|
||||
DCHECK_EQ(parsed_.size(), values_[values_.size() - 1].offset);
|
||||
#else
|
||||
DCHECK_EQ(offsets_.size(), 1 + static_cast<size_t>(num_rows_ * num_cols_));
|
||||
DCHECK_EQ(parsed_.size(), offsets_[offsets_.size() - 1]);
|
||||
#endif
|
||||
*out_size = static_cast<uint32_t>(data - start);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status BlockParser::Parse(const char* data, uint32_t size, uint32_t* out_size) {
|
||||
return DoParse(data, size, false /* is_final */, out_size);
|
||||
}
|
||||
|
||||
Status BlockParser::ParseFinal(const char* data, uint32_t size, uint32_t* out_size) {
|
||||
return DoParse(data, size, true /* is_final */, out_size);
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,136 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef ARROW_CSV_PARSER_H
|
||||
#define ARROW_CSV_PARSER_H
|
||||
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/csv/options.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/util/visibility.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
constexpr int32_t kMaxParserNumRows = 100000;
|
||||
|
||||
// Whether BlockParser will use bitfields for better memory consumption
|
||||
// and cache locality.
|
||||
#undef CSV_PARSER_USE_BITFIELD
|
||||
#define CSV_PARSER_USE_BITFIELD
|
||||
|
||||
/// \class BlockParser
|
||||
/// \brief A reusable block-based parser for CSV data
|
||||
///
|
||||
/// The parser takes a block of CSV data and delimits rows and fields,
|
||||
/// unquoting and unescaping them on the fly. Parsed data is own by the
|
||||
/// parser, so the original buffer can be discarded after Parse() returns.
|
||||
///
|
||||
/// If the block is truncated (i.e. not all data can be parsed), it is up
|
||||
/// to the caller to arrange the next block to start with the trailing data.
|
||||
/// Also, if the previous block ends with CR (0x0d) and a new block starts
|
||||
/// with LF (0x0a), the parser will consider the leading newline as an empty
|
||||
/// line; the caller should therefore strip it.
|
||||
class ARROW_EXPORT BlockParser {
|
||||
public:
|
||||
explicit BlockParser(ParseOptions options, int32_t num_cols = -1,
|
||||
int32_t max_num_rows = kMaxParserNumRows);
|
||||
|
||||
/// \brief Parse a block of data
|
||||
///
|
||||
/// Parse a block of CSV data, ingesting up to max_num_rows rows.
|
||||
/// The number of bytes actually parsed is returned in out_size.
|
||||
Status Parse(const char* data, uint32_t size, uint32_t* out_size);
|
||||
|
||||
/// \brief Parse the final block of data
|
||||
///
|
||||
/// Like Parse(), but called with the final block in a file.
|
||||
/// The last row may lack a trailing line separator.
|
||||
Status ParseFinal(const char* data, uint32_t size, uint32_t* out_size);
|
||||
|
||||
int32_t num_rows() const { return num_rows_; }
|
||||
int32_t num_cols() const { return num_cols_; }
|
||||
/// \brief Return the total size in bytes of parsed data
|
||||
int32_t num_bytes() const { return static_cast<int32_t>(parsed_.size()); }
|
||||
|
||||
/// \brief Visit parsed values in a column
|
||||
///
|
||||
/// The signature of the visitor is (const uint8_t* data, uint32_t size, bool quoted)
|
||||
template <typename Visitor>
|
||||
Status VisitColumn(int32_t col_index, Visitor&& visit) const {
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
for (int32_t pos = col_index; pos < num_cols_ * num_rows_; pos += num_cols_) {
|
||||
auto start = values_[pos].offset;
|
||||
auto stop = values_[pos + 1].offset;
|
||||
auto quoted = values_[pos + 1].quoted;
|
||||
ARROW_RETURN_NOT_OK(visit(parsed_.data() + start, stop - start, quoted));
|
||||
}
|
||||
#else
|
||||
for (int32_t pos = col_index; pos < num_cols_ * num_rows_; pos += num_cols_) {
|
||||
auto start = offsets_[pos];
|
||||
auto stop = offsets_[pos + 1];
|
||||
auto quoted = quoted_[pos];
|
||||
ARROW_RETURN_NOT_OK(visit(parsed_.data() + start, stop - start, quoted));
|
||||
}
|
||||
#endif
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
// XXX add a ClearColumn method to signal that a column won't be visited anymore?
|
||||
|
||||
protected:
|
||||
ARROW_DISALLOW_COPY_AND_ASSIGN(BlockParser);
|
||||
|
||||
Status DoParse(const char* data, uint32_t size, bool is_final, uint32_t* out_size);
|
||||
|
||||
// Parse a single line from the data pointer
|
||||
Status ParseLine(const char* data, const char* data_end, bool is_final,
|
||||
const char** out_data);
|
||||
|
||||
ParseOptions options_;
|
||||
// The number of rows parsed from the block
|
||||
int32_t num_rows_;
|
||||
// The number of columns (can be -1 at start)
|
||||
int32_t num_cols_;
|
||||
// The maximum number of rows to parse from this block
|
||||
int32_t max_num_rows_;
|
||||
|
||||
// Linear scratchpad for parsed values
|
||||
// XXX should we ensure it's padded with 8 or 16 excess zero bytes? it could help
|
||||
// with null parsing...
|
||||
// TODO should be able to presize scratch space
|
||||
std::vector<uint8_t> parsed_;
|
||||
#ifdef CSV_PARSER_USE_BITFIELD
|
||||
struct ValueDesc {
|
||||
uint32_t offset : 31;
|
||||
bool quoted : 1;
|
||||
};
|
||||
std::vector<ValueDesc> values_;
|
||||
#else
|
||||
// Value offsets inside the scratchpad
|
||||
std::vector<uint32_t> offsets_;
|
||||
// Whether each value was quoted or not
|
||||
std::vector<bool> quoted_;
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
|
||||
#endif // ARROW_CSV_PARSER_H
|
||||
@@ -0,0 +1,399 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/csv/chunker.h"
|
||||
#include "arrow/csv/column-builder.h"
|
||||
#include "arrow/csv/parser.h"
|
||||
#include "arrow/csv/reader.h"
|
||||
#include "arrow/io/interfaces.h"
|
||||
#include "arrow/io/readahead.h"
|
||||
#include "arrow/memory_pool.h"
|
||||
#include "arrow/record_batch.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/util/logging.h"
|
||||
#include "arrow/util/task-group.h"
|
||||
#include "arrow/util/thread-pool.h"
|
||||
#include "arrow/util/visibility.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
using internal::GetCpuThreadPool;
|
||||
using internal::ThreadPool;
|
||||
using io::internal::ReadaheadBuffer;
|
||||
using io::internal::ReadaheadSpooler;
|
||||
|
||||
static constexpr int64_t kDefaultLeftPadding = 2048; // 2 kB
|
||||
static constexpr int64_t kDefaultRightPadding = 16;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// Base class for common functionality
|
||||
|
||||
class BaseTableReader : public csv::TableReader {
|
||||
public:
|
||||
BaseTableReader(MemoryPool* pool, const ReadOptions& read_options,
|
||||
const ParseOptions& parse_options,
|
||||
const ConvertOptions& convert_options)
|
||||
: pool_(pool),
|
||||
read_options_(read_options),
|
||||
parse_options_(parse_options),
|
||||
convert_options_(convert_options) {}
|
||||
|
||||
protected:
|
||||
// Read a next data block, stitch it to trailing data
|
||||
Status ReadNextBlock() {
|
||||
bool trailing_data = cur_size_ > 0;
|
||||
ReadaheadBuffer rh;
|
||||
|
||||
if (trailing_data) {
|
||||
if (readahead_->GetLeftPadding() < cur_size_) {
|
||||
// Growth heuristic to try and ensure sufficient left padding
|
||||
// in subsequent reads
|
||||
readahead_->SetLeftPadding(cur_size_ * 3 / 2);
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_NOT_OK(readahead_->Read(&rh));
|
||||
if (!rh.buffer) {
|
||||
// EOF, let caller finish with existing data
|
||||
eof_ = true;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
std::shared_ptr<Buffer> new_block = rh.buffer;
|
||||
uint8_t* new_data = rh.buffer->mutable_data() + rh.left_padding;
|
||||
int64_t new_size = rh.buffer->size() - rh.left_padding - rh.right_padding;
|
||||
DCHECK_GT(new_size, 0); // ensured by ReadaheadSpooler
|
||||
|
||||
if (trailing_cr_ && new_data[0] == '\n') {
|
||||
// Skip '\r\n' line separator that started at the end of previous block
|
||||
++new_data;
|
||||
--new_size;
|
||||
}
|
||||
trailing_cr_ = (new_data[new_size - 1] == '\r');
|
||||
|
||||
if (trailing_data) {
|
||||
// Try to copy trailing data at the beginning of new block
|
||||
if (cur_size_ <= rh.left_padding) {
|
||||
// Can left-extend new block inside padding area
|
||||
new_data -= cur_size_;
|
||||
new_size += cur_size_;
|
||||
std::memcpy(new_data, cur_data_, cur_size_);
|
||||
} else {
|
||||
// Need to allocate bigger block and concatenate trailing + present data
|
||||
RETURN_NOT_OK(
|
||||
AllocateBuffer(pool_, cur_size_ + new_size + rh.right_padding, &new_block));
|
||||
std::memcpy(new_block->mutable_data(), cur_data_, cur_size_);
|
||||
std::memcpy(new_block->mutable_data() + cur_size_, new_data, new_size);
|
||||
std::memset(new_block->mutable_data() + cur_size_ + new_size, 0,
|
||||
rh.right_padding);
|
||||
new_data = new_block->mutable_data();
|
||||
new_size = cur_size_ + new_size;
|
||||
}
|
||||
}
|
||||
cur_block_ = new_block;
|
||||
cur_data_ = new_data;
|
||||
cur_size_ = new_size;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
// Read header and column names from current block, create column builders
|
||||
Status ProcessHeader() {
|
||||
DCHECK_GT(cur_size_, 0);
|
||||
if (parse_options_.header_rows == 0) {
|
||||
// TODO allow passing names and/or generate column numbers?
|
||||
return Status::Invalid("header_rows == 0 needs explicit column names");
|
||||
}
|
||||
|
||||
BlockParser parser(parse_options_, num_cols_, parse_options_.header_rows);
|
||||
|
||||
uint32_t parsed_size = 0;
|
||||
RETURN_NOT_OK(parser.Parse(reinterpret_cast<const char*>(cur_data_),
|
||||
static_cast<uint32_t>(cur_size_), &parsed_size));
|
||||
if (parser.num_rows() != parse_options_.header_rows) {
|
||||
return Status::Invalid(
|
||||
"Could not read header rows from CSV file, either "
|
||||
"file is too short or header is larger than block size");
|
||||
}
|
||||
if (parser.num_cols() == 0) {
|
||||
return Status::Invalid("No columns in CSV file");
|
||||
}
|
||||
num_cols_ = parser.num_cols();
|
||||
DCHECK_GT(num_cols_, 0);
|
||||
|
||||
for (int32_t col_index = 0; col_index < num_cols_; ++col_index) {
|
||||
auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status {
|
||||
if (names_.size() <= static_cast<uint32_t>(col_index)) {
|
||||
names_.emplace_back(reinterpret_cast<const char*>(data), size);
|
||||
}
|
||||
return Status::OK();
|
||||
};
|
||||
RETURN_NOT_OK(parser.VisitColumn(col_index, visit));
|
||||
std::shared_ptr<ColumnBuilder> builder;
|
||||
RETURN_NOT_OK(
|
||||
ColumnBuilder::Make(col_index, convert_options_, task_group_, &builder));
|
||||
column_builders_.push_back(builder);
|
||||
}
|
||||
|
||||
// Skip parsed header rows
|
||||
cur_data_ += parsed_size;
|
||||
cur_size_ -= parsed_size;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
// Trigger conversion of parsed block data
|
||||
Status ProcessData(const std::shared_ptr<BlockParser>& parser, int64_t block_index) {
|
||||
for (auto& builder : column_builders_) {
|
||||
builder->Insert(block_index, parser);
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status MakeTable(std::shared_ptr<Table>* out) {
|
||||
DCHECK_GT(num_cols_, 0);
|
||||
DCHECK_EQ(names_.size(), static_cast<uint32_t>(num_cols_));
|
||||
DCHECK_EQ(column_builders_.size(), static_cast<uint32_t>(num_cols_));
|
||||
|
||||
std::vector<std::shared_ptr<Field>> fields;
|
||||
std::vector<std::shared_ptr<Column>> columns;
|
||||
|
||||
for (int32_t i = 0; i < num_cols_; ++i) {
|
||||
std::shared_ptr<ChunkedArray> array;
|
||||
RETURN_NOT_OK(column_builders_[i]->Finish(&array));
|
||||
columns.push_back(std::make_shared<Column>(names_[i], array));
|
||||
fields.push_back(columns.back()->field());
|
||||
}
|
||||
*out = Table::Make(schema(fields), columns);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
MemoryPool* pool_;
|
||||
ReadOptions read_options_;
|
||||
ParseOptions parse_options_;
|
||||
ConvertOptions convert_options_;
|
||||
|
||||
int32_t num_cols_ = -1;
|
||||
std::shared_ptr<ReadaheadSpooler> readahead_;
|
||||
std::vector<std::string> names_;
|
||||
std::shared_ptr<internal::TaskGroup> task_group_;
|
||||
std::vector<std::shared_ptr<ColumnBuilder>> column_builders_;
|
||||
|
||||
// Current block and data pointer
|
||||
std::shared_ptr<Buffer> cur_block_;
|
||||
const uint8_t* cur_data_ = nullptr;
|
||||
int64_t cur_size_ = 0;
|
||||
// Index of current block inside data stream
|
||||
int64_t cur_block_index_ = 0;
|
||||
// Whether there was a trailing CR at the end of last parsed line
|
||||
bool trailing_cr_ = false;
|
||||
// Whether we reached input stream EOF. There may still be data left to
|
||||
// process in current block.
|
||||
bool eof_ = false;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// Serial TableReader implementation
|
||||
|
||||
class SerialTableReader : public BaseTableReader {
|
||||
public:
|
||||
SerialTableReader(MemoryPool* pool, std::shared_ptr<io::InputStream> input,
|
||||
const ReadOptions& read_options, const ParseOptions& parse_options,
|
||||
const ConvertOptions& convert_options)
|
||||
: BaseTableReader(pool, read_options, parse_options, convert_options) {
|
||||
// Since we're converting serially, no need to readahead more than one block
|
||||
int32_t block_queue_size = 1;
|
||||
readahead_ = std::make_shared<ReadaheadSpooler>(
|
||||
pool_, input, read_options_.block_size, block_queue_size, kDefaultLeftPadding,
|
||||
kDefaultRightPadding);
|
||||
}
|
||||
|
||||
Status Read(std::shared_ptr<Table>* out) {
|
||||
task_group_ = internal::TaskGroup::MakeSerial();
|
||||
|
||||
// First block
|
||||
RETURN_NOT_OK(ReadNextBlock());
|
||||
if (eof_) {
|
||||
return Status::Invalid("Empty CSV file");
|
||||
}
|
||||
RETURN_NOT_OK(ProcessHeader());
|
||||
|
||||
auto parser =
|
||||
std::make_shared<BlockParser>(parse_options_, num_cols_, read_options_.num_rows);
|
||||
while (!eof_) {
|
||||
// Consume current block
|
||||
uint32_t parsed_size = 0;
|
||||
RETURN_NOT_OK(parser->Parse(reinterpret_cast<const char*>(cur_data_),
|
||||
static_cast<uint32_t>(cur_size_), &parsed_size));
|
||||
if (parser->num_rows() > 0) {
|
||||
// Got some data
|
||||
RETURN_NOT_OK(ProcessData(parser, cur_block_index_++));
|
||||
cur_data_ += parsed_size;
|
||||
cur_size_ -= parsed_size;
|
||||
if (!task_group_->ok()) {
|
||||
// Conversion error => early exit
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
// Need to fetch more data to get at least one row
|
||||
RETURN_NOT_OK(ReadNextBlock());
|
||||
}
|
||||
}
|
||||
if (eof_ && cur_size_ > 0) {
|
||||
// Parse remaining data
|
||||
uint32_t parsed_size = 0;
|
||||
RETURN_NOT_OK(parser->ParseFinal(reinterpret_cast<const char*>(cur_data_),
|
||||
static_cast<uint32_t>(cur_size_), &parsed_size));
|
||||
if (parser->num_rows() > 0) {
|
||||
RETURN_NOT_OK(ProcessData(parser, cur_block_index_++));
|
||||
}
|
||||
}
|
||||
|
||||
// Finish conversion, create schema and table
|
||||
RETURN_NOT_OK(task_group_->Finish());
|
||||
return MakeTable(out);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// Parallel TableReader implementation
|
||||
|
||||
class ThreadedTableReader : public BaseTableReader {
|
||||
public:
|
||||
ThreadedTableReader(MemoryPool* pool, std::shared_ptr<io::InputStream> input,
|
||||
ThreadPool* thread_pool, const ReadOptions& read_options,
|
||||
const ParseOptions& parse_options,
|
||||
const ConvertOptions& convert_options)
|
||||
: BaseTableReader(pool, read_options, parse_options, convert_options),
|
||||
thread_pool_(thread_pool) {
|
||||
// Readahead one block per worker thread
|
||||
int32_t block_queue_size = thread_pool->GetCapacity();
|
||||
readahead_ = std::make_shared<ReadaheadSpooler>(
|
||||
pool_, input, read_options_.block_size, block_queue_size, kDefaultLeftPadding,
|
||||
kDefaultRightPadding);
|
||||
}
|
||||
|
||||
~ThreadedTableReader() {
|
||||
if (task_group_) {
|
||||
// In case of error, make sure all pending tasks are finished before
|
||||
// we start destroying BaseTableReader members
|
||||
ARROW_UNUSED(task_group_->Finish());
|
||||
}
|
||||
}
|
||||
|
||||
Status Read(std::shared_ptr<Table>* out) {
|
||||
task_group_ = internal::TaskGroup::MakeThreaded(thread_pool_);
|
||||
Chunker chunker(parse_options_, static_cast<int32_t>(read_options_.num_rows));
|
||||
|
||||
// Get first block and process header serially
|
||||
RETURN_NOT_OK(ReadNextBlock());
|
||||
if (eof_) {
|
||||
return Status::Invalid("Empty CSV file");
|
||||
}
|
||||
RETURN_NOT_OK(ProcessHeader());
|
||||
|
||||
while (!eof_ && task_group_->ok()) {
|
||||
// Consume current chunk
|
||||
uint32_t chunk_size = 0;
|
||||
RETURN_NOT_OK(chunker.Process(reinterpret_cast<const char*>(cur_data_),
|
||||
static_cast<uint32_t>(cur_size_), &chunk_size));
|
||||
int32_t chunk_rows = chunker.num_rows();
|
||||
if (chunk_rows > 0) {
|
||||
// Got a chunk of rows
|
||||
const uint8_t* chunk_data = cur_data_;
|
||||
std::shared_ptr<Buffer> chunk_buffer = cur_block_;
|
||||
int64_t chunk_index = cur_block_index_;
|
||||
|
||||
// "mutable" allows to modify captured by-copy chunk_buffer
|
||||
task_group_->Append([=]() mutable -> Status {
|
||||
auto parser =
|
||||
std::make_shared<BlockParser>(parse_options_, num_cols_, chunk_rows);
|
||||
uint32_t parsed_size = 0;
|
||||
RETURN_NOT_OK(parser->Parse(reinterpret_cast<const char*>(chunk_data),
|
||||
chunk_size, &parsed_size));
|
||||
if (parsed_size != chunk_size || chunk_rows != parser->num_rows()) {
|
||||
return Status::Invalid("Chunker and parser disagree on block size");
|
||||
}
|
||||
RETURN_NOT_OK(ProcessData(parser, chunk_index));
|
||||
// Keep chunk buffer alive within closure and release it at the end
|
||||
chunk_buffer.reset();
|
||||
return Status::OK();
|
||||
});
|
||||
cur_data_ += chunk_size;
|
||||
cur_size_ -= chunk_size;
|
||||
cur_block_index_++;
|
||||
} else {
|
||||
// Need to fetch more data to get at least one row
|
||||
RETURN_NOT_OK(ReadNextBlock());
|
||||
}
|
||||
}
|
||||
|
||||
// Finish all pending parallel tasks
|
||||
RETURN_NOT_OK(task_group_->Finish());
|
||||
|
||||
if (eof_ && cur_size_ > 0) {
|
||||
// Parse remaining data (serial)
|
||||
task_group_ = internal::TaskGroup::MakeSerial();
|
||||
for (auto& builder : column_builders_) {
|
||||
builder->SetTaskGroup(task_group_);
|
||||
}
|
||||
auto parser = std::make_shared<BlockParser>(parse_options_, num_cols_,
|
||||
read_options_.num_rows);
|
||||
uint32_t parsed_size = 0;
|
||||
RETURN_NOT_OK(parser->ParseFinal(reinterpret_cast<const char*>(cur_data_),
|
||||
static_cast<uint32_t>(cur_size_), &parsed_size));
|
||||
if (parser->num_rows() > 0) {
|
||||
RETURN_NOT_OK(ProcessData(parser, cur_block_index_++));
|
||||
}
|
||||
RETURN_NOT_OK(task_group_->Finish());
|
||||
}
|
||||
|
||||
// Create schema and table
|
||||
return MakeTable(out);
|
||||
}
|
||||
|
||||
protected:
|
||||
ThreadPool* thread_pool_;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// TableReader factory function
|
||||
|
||||
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<TableReader> result;
|
||||
if (read_options.use_threads) {
|
||||
result = std::make_shared<ThreadedTableReader>(
|
||||
pool, input, GetCpuThreadPool(), read_options, parse_options, convert_options);
|
||||
*out = result;
|
||||
return Status::OK();
|
||||
} else {
|
||||
result = std::make_shared<SerialTableReader>(pool, input, read_options, parse_options,
|
||||
convert_options);
|
||||
*out = result;
|
||||
return Status::OK();
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,53 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef ARROW_CSV_READER_H
|
||||
#define ARROW_CSV_READER_H
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "arrow/csv/options.h"
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/util/visibility.h"
|
||||
|
||||
namespace arrow {
|
||||
|
||||
class MemoryPool;
|
||||
class Table;
|
||||
|
||||
namespace io {
|
||||
class InputStream;
|
||||
} // namespace io
|
||||
|
||||
namespace csv {
|
||||
|
||||
class ARROW_EXPORT TableReader {
|
||||
public:
|
||||
virtual ~TableReader() = default;
|
||||
|
||||
virtual Status Read(std::shared_ptr<Table>* out) = 0;
|
||||
|
||||
// XXX pass optional schema?
|
||||
static Status Make(MemoryPool* pool, std::shared_ptr<io::InputStream> input,
|
||||
const ReadOptions&, const ParseOptions&, const ConvertOptions&,
|
||||
std::shared_ptr<TableReader>* out);
|
||||
};
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
|
||||
#endif // ARROW_CSV_READER_H
|
||||
@@ -0,0 +1,64 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef ARROW_CSV_TEST_COMMON_H
|
||||
#define ARROW_CSV_TEST_COMMON_H
|
||||
|
||||
#include <memory>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/csv/parser.h"
|
||||
#include "arrow/test-util.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace csv {
|
||||
|
||||
std::string MakeCSVData(std::vector<std::string> lines) {
|
||||
std::stringstream ss;
|
||||
for (const auto& line : lines) {
|
||||
ss << line;
|
||||
}
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
// Make a BlockParser from a vector of lines representing a CSV file
|
||||
void MakeCSVParser(std::vector<std::string> lines, std::shared_ptr<BlockParser>* out) {
|
||||
auto csv = MakeCSVData(lines);
|
||||
auto parser = std::make_shared<BlockParser>(ParseOptions::Defaults());
|
||||
uint32_t out_size;
|
||||
ASSERT_OK(parser->Parse(csv.data(), static_cast<uint32_t>(csv.size()), &out_size));
|
||||
ASSERT_EQ(out_size, csv.size()) << "trailing CSV data not parsed";
|
||||
*out = parser;
|
||||
}
|
||||
|
||||
// Make a BlockParser from a vector of strings representing a single CSV column
|
||||
void MakeColumnParser(std::vector<std::string> items, std::shared_ptr<BlockParser>* out) {
|
||||
std::vector<std::string> lines;
|
||||
for (const auto& item : items) {
|
||||
lines.push_back(item + '\n');
|
||||
}
|
||||
MakeCSVParser(lines, out);
|
||||
ASSERT_EQ((*out)->num_cols(), 1) << "Should have seen only 1 CSV column";
|
||||
ASSERT_EQ((*out)->num_rows(), items.size());
|
||||
}
|
||||
|
||||
} // namespace csv
|
||||
} // namespace arrow
|
||||
|
||||
#endif // ARROW_CSV_TEST_COMMON_H
|
||||
@@ -347,6 +347,10 @@ void AssertChunkedEqual(const ChunkedArray& expected, const ChunkedArray& actual
|
||||
}
|
||||
}
|
||||
|
||||
void AssertChunkedEqual(const ChunkedArray& actual, const ArrayVector& expected) {
|
||||
AssertChunkedEqual(ChunkedArray(expected, actual.type()), actual);
|
||||
}
|
||||
|
||||
void AssertBufferEqual(const Buffer& buffer, const std::vector<uint8_t>& expected) {
|
||||
ASSERT_EQ(buffer.size(), expected.size()) << "Mismatching buffer size";
|
||||
const uint8_t* buffer_data = buffer.data();
|
||||
@@ -410,11 +414,18 @@ void AssertTablesEqual(const Table& expected, const Table& actual,
|
||||
}
|
||||
}
|
||||
|
||||
template <typename TYPE, typename C_TYPE>
|
||||
template <typename TYPE, typename C_TYPE = typename TYPE::c_type>
|
||||
void ArrayFromVector(const std::shared_ptr<DataType>& type,
|
||||
const std::vector<bool>& is_valid, const std::vector<C_TYPE>& values,
|
||||
std::shared_ptr<Array>* out) {
|
||||
typename TypeTraits<TYPE>::BuilderType builder(type, default_memory_pool());
|
||||
DCHECK_EQ(TYPE::type_id, type->id())
|
||||
<< "template parameter and concrete DataType instance don't agree";
|
||||
|
||||
std::unique_ptr<ArrayBuilder> builder_ptr;
|
||||
ASSERT_OK(MakeBuilder(default_memory_pool(), type, &builder_ptr));
|
||||
// Get the concrete builder class to access its Append() specializations
|
||||
auto& builder = dynamic_cast<typename TypeTraits<TYPE>::BuilderType&>(*builder_ptr);
|
||||
|
||||
for (size_t i = 0; i < values.size(); ++i) {
|
||||
if (is_valid[i]) {
|
||||
ASSERT_OK(builder.Append(values[i]));
|
||||
@@ -425,20 +436,28 @@ void ArrayFromVector(const std::shared_ptr<DataType>& type,
|
||||
ASSERT_OK(builder.Finish(out));
|
||||
}
|
||||
|
||||
template <typename TYPE, typename C_TYPE>
|
||||
template <typename TYPE, typename C_TYPE = typename TYPE::c_type>
|
||||
void ArrayFromVector(const std::shared_ptr<DataType>& type,
|
||||
const std::vector<C_TYPE>& values, std::shared_ptr<Array>* out) {
|
||||
typename TypeTraits<TYPE>::BuilderType builder(type, default_memory_pool());
|
||||
DCHECK_EQ(TYPE::type_id, type->id())
|
||||
<< "template parameter and concrete DataType instance don't agree";
|
||||
|
||||
std::unique_ptr<ArrayBuilder> builder_ptr;
|
||||
ASSERT_OK(MakeBuilder(default_memory_pool(), type, &builder_ptr));
|
||||
// Get the concrete builder class to access its Append() specializations
|
||||
auto& builder = dynamic_cast<typename TypeTraits<TYPE>::BuilderType&>(*builder_ptr);
|
||||
|
||||
for (size_t i = 0; i < values.size(); ++i) {
|
||||
ASSERT_OK(builder.Append(values[i]));
|
||||
}
|
||||
ASSERT_OK(builder.Finish(out));
|
||||
}
|
||||
|
||||
template <typename TYPE, typename C_TYPE>
|
||||
template <typename TYPE, typename C_TYPE = typename TYPE::c_type>
|
||||
void ArrayFromVector(const std::vector<bool>& is_valid, const std::vector<C_TYPE>& values,
|
||||
std::shared_ptr<Array>* out) {
|
||||
typename TypeTraits<TYPE>::BuilderType builder;
|
||||
DCHECK_EQ(is_valid.size(), values.size());
|
||||
for (size_t i = 0; i < values.size(); ++i) {
|
||||
if (is_valid[i]) {
|
||||
ASSERT_OK(builder.Append(values[i]));
|
||||
@@ -449,7 +468,7 @@ void ArrayFromVector(const std::vector<bool>& is_valid, const std::vector<C_TYPE
|
||||
ASSERT_OK(builder.Finish(out));
|
||||
}
|
||||
|
||||
template <typename TYPE, typename C_TYPE>
|
||||
template <typename TYPE, typename C_TYPE = typename TYPE::c_type>
|
||||
void ArrayFromVector(const std::vector<C_TYPE>& values, std::shared_ptr<Array>* out) {
|
||||
typename TypeTraits<TYPE>::BuilderType builder;
|
||||
for (auto& value : values) {
|
||||
@@ -458,6 +477,32 @@ void ArrayFromVector(const std::vector<C_TYPE>& values, std::shared_ptr<Array>*
|
||||
ASSERT_OK(builder.Finish(out));
|
||||
}
|
||||
|
||||
template <typename TYPE, typename C_TYPE = typename TYPE::c_type>
|
||||
void ChunkedArrayFromVector(const std::vector<std::vector<bool>>& is_valid,
|
||||
const std::vector<std::vector<C_TYPE>>& values,
|
||||
std::shared_ptr<ChunkedArray>* out) {
|
||||
ArrayVector chunks;
|
||||
DCHECK_EQ(is_valid.size(), values.size());
|
||||
for (size_t i = 0; i < values.size(); ++i) {
|
||||
std::shared_ptr<Array> array;
|
||||
ArrayFromVector<TYPE, C_TYPE>(is_valid[i], values[i], &array);
|
||||
chunks.push_back(array);
|
||||
}
|
||||
*out = std::make_shared<ChunkedArray>(chunks);
|
||||
}
|
||||
|
||||
template <typename TYPE, typename C_TYPE = typename TYPE::c_type>
|
||||
void ChunkedArrayFromVector(const std::vector<std::vector<C_TYPE>>& values,
|
||||
std::shared_ptr<ChunkedArray>* out) {
|
||||
ArrayVector chunks;
|
||||
for (size_t i = 0; i < values.size(); ++i) {
|
||||
std::shared_ptr<Array> array;
|
||||
ArrayFromVector<TYPE, C_TYPE>(values[i], &array);
|
||||
chunks.push_back(array);
|
||||
}
|
||||
*out = std::make_shared<ChunkedArray>(chunks);
|
||||
}
|
||||
|
||||
template <class T, class Builder>
|
||||
Status MakeArray(const std::vector<uint8_t>& valid_bytes, const std::vector<T>& values,
|
||||
int64_t size, Builder* builder, std::shared_ptr<Array>* out) {
|
||||
|
||||
@@ -90,6 +90,7 @@ ADD_ARROW_TEST(rle-encoding-test)
|
||||
ADD_ARROW_TEST(parsing-util-test)
|
||||
ADD_ARROW_TEST(stl-util-test)
|
||||
ADD_ARROW_TEST(thread-pool-test)
|
||||
ADD_ARROW_TEST(task-group-test)
|
||||
ADD_ARROW_TEST(lazy-test)
|
||||
ADD_ARROW_TEST(logging-test)
|
||||
|
||||
|
||||
@@ -0,0 +1,264 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef _WIN32
|
||||
#include <sys/wait.h>
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <memory>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/test-util.h"
|
||||
#include "arrow/util/macros.h"
|
||||
#include "arrow/util/task-group.h"
|
||||
#include "arrow/util/thread-pool.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace internal {
|
||||
|
||||
static void sleep_for(double seconds) {
|
||||
std::this_thread::sleep_for(
|
||||
std::chrono::nanoseconds(static_cast<int64_t>(seconds * 1e9)));
|
||||
}
|
||||
|
||||
// Generate random sleep durations
|
||||
static std::vector<double> RandomSleepDurations(int nsleeps, double min_seconds,
|
||||
double max_seconds) {
|
||||
std::vector<double> sleeps;
|
||||
std::default_random_engine engine;
|
||||
std::uniform_real_distribution<> sleep_dist(min_seconds, max_seconds);
|
||||
for (int i = 0; i < nsleeps; ++i) {
|
||||
sleeps.push_back(sleep_dist(engine));
|
||||
}
|
||||
return sleeps;
|
||||
}
|
||||
|
||||
// Check TaskGroup behaviour with a bunch of all-successful tasks
|
||||
void TestTaskGroupSuccess(std::shared_ptr<TaskGroup> task_group) {
|
||||
const int NTASKS = 10;
|
||||
auto sleeps = RandomSleepDurations(NTASKS, 1e-3, 4e-3);
|
||||
|
||||
// Add NTASKS sleeps
|
||||
std::atomic<int> count(0);
|
||||
for (int i = 0; i < NTASKS; ++i) {
|
||||
task_group->Append([&, i]() {
|
||||
sleep_for(sleeps[i]);
|
||||
count += i;
|
||||
return Status::OK();
|
||||
});
|
||||
}
|
||||
ASSERT_TRUE(task_group->ok());
|
||||
|
||||
ASSERT_OK(task_group->Finish());
|
||||
ASSERT_TRUE(task_group->ok());
|
||||
ASSERT_EQ(count.load(), NTASKS * (NTASKS - 1) / 2);
|
||||
// Finish() is idempotent
|
||||
ASSERT_OK(task_group->Finish());
|
||||
}
|
||||
|
||||
// Check TaskGroup behaviour with some successful and some failing tasks
|
||||
void TestTaskGroupErrors(std::shared_ptr<TaskGroup> task_group) {
|
||||
const int NSUCCESSES = 2;
|
||||
const int NERRORS = 20;
|
||||
|
||||
std::atomic<int> count(0);
|
||||
|
||||
for (int i = 0; i < NSUCCESSES; ++i) {
|
||||
task_group->Append([&]() {
|
||||
count++;
|
||||
return Status::OK();
|
||||
});
|
||||
}
|
||||
ASSERT_TRUE(task_group->ok());
|
||||
for (int i = 0; i < NERRORS; ++i) {
|
||||
task_group->Append([&]() {
|
||||
sleep_for(1e-2);
|
||||
count++;
|
||||
return Status::Invalid("some message");
|
||||
});
|
||||
}
|
||||
|
||||
// Task error is propagated
|
||||
ASSERT_RAISES(Invalid, task_group->Finish());
|
||||
ASSERT_FALSE(task_group->ok());
|
||||
if (task_group->parallelism() == 1) {
|
||||
// Serial: exactly two successes and an error
|
||||
ASSERT_EQ(count.load(), 3);
|
||||
} else {
|
||||
// Parallel: at least two successes and an error
|
||||
ASSERT_GE(count.load(), 3);
|
||||
ASSERT_LE(count.load(), 2 * task_group->parallelism());
|
||||
}
|
||||
// Finish() is idempotent
|
||||
ASSERT_RAISES(Invalid, task_group->Finish());
|
||||
}
|
||||
|
||||
// Check TaskGroup behaviour with a bunch of all-successful tasks and task groups
|
||||
void TestTaskSubGroupsSuccess(std::shared_ptr<TaskGroup> task_group) {
|
||||
const int NTASKS = 50;
|
||||
const int NGROUPS = 7;
|
||||
|
||||
auto sleeps = RandomSleepDurations(NTASKS, 1e-4, 1e-3);
|
||||
std::vector<std::shared_ptr<TaskGroup>> groups = {task_group};
|
||||
|
||||
// Create some subgroups
|
||||
for (int i = 0; i < NGROUPS - 1; ++i) {
|
||||
groups.push_back(task_group->MakeSubGroup());
|
||||
}
|
||||
|
||||
// Add NTASKS sleeps amonst all groups
|
||||
std::atomic<int> count(0);
|
||||
for (int i = 0; i < NTASKS; ++i) {
|
||||
groups[i % NGROUPS]->Append([&, i]() {
|
||||
sleep_for(sleeps[i]);
|
||||
count += i;
|
||||
return Status::OK();
|
||||
});
|
||||
}
|
||||
ASSERT_TRUE(task_group->ok());
|
||||
|
||||
// Finish all subgroups first, then main group
|
||||
for (int i = NGROUPS - 1; i >= 0; --i) {
|
||||
ASSERT_OK(groups[i]->Finish());
|
||||
}
|
||||
ASSERT_TRUE(task_group->ok());
|
||||
ASSERT_EQ(count.load(), NTASKS * (NTASKS - 1) / 2);
|
||||
// Finish() is idempotent
|
||||
ASSERT_OK(task_group->Finish());
|
||||
}
|
||||
|
||||
// Check TaskGroup behaviour with both successful and failing tasks and task groups
|
||||
void TestTaskSubGroupsErrors(std::shared_ptr<TaskGroup> task_group) {
|
||||
const int NTASKS = 50;
|
||||
const int NGROUPS = 7;
|
||||
const int FAIL_EVERY = 17;
|
||||
std::vector<std::shared_ptr<TaskGroup>> groups = {task_group};
|
||||
|
||||
// Create some subgroups
|
||||
for (int i = 0; i < NGROUPS - 1; ++i) {
|
||||
groups.push_back(task_group->MakeSubGroup());
|
||||
}
|
||||
|
||||
// Add NTASKS sleeps amonst all groups
|
||||
for (int i = 0; i < NTASKS; ++i) {
|
||||
groups[i % NGROUPS]->Append([&, i]() {
|
||||
sleep_for(1e-3);
|
||||
// As NGROUPS > NTASKS / FAIL_EVERY, some subgroups are successful
|
||||
if (i % FAIL_EVERY == 0) {
|
||||
return Status::Invalid("some message");
|
||||
} else {
|
||||
return Status::OK();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Finish all subgroups first, then main group
|
||||
int nsuccessful = 0;
|
||||
for (int i = NGROUPS - 1; i > 0; --i) {
|
||||
Status st = groups[i]->Finish();
|
||||
if (st.ok()) {
|
||||
++nsuccessful;
|
||||
} else {
|
||||
ASSERT_RAISES(Invalid, st);
|
||||
}
|
||||
}
|
||||
ASSERT_RAISES(Invalid, task_group->Finish());
|
||||
ASSERT_FALSE(task_group->ok());
|
||||
// Finish() is idempotent
|
||||
ASSERT_RAISES(Invalid, task_group->Finish());
|
||||
}
|
||||
|
||||
// Check TaskGroup behaviour with tasks spawning other tasks
|
||||
void TestTasksSpawnTasks(std::shared_ptr<TaskGroup> task_group) {
|
||||
const int N = 6;
|
||||
|
||||
std::atomic<int> count(0);
|
||||
// Make a task that recursively spawns itself
|
||||
std::function<std::function<Status()>(int)> make_task = [&](int i) {
|
||||
return [&, i]() {
|
||||
count++;
|
||||
if (i > 0) {
|
||||
// Exercise parallelism by spawning two tasks at once and then sleeping
|
||||
task_group->Append(make_task(i - 1));
|
||||
task_group->Append(make_task(i - 1));
|
||||
sleep_for(1e-3);
|
||||
}
|
||||
return Status::OK();
|
||||
};
|
||||
};
|
||||
|
||||
task_group->Append(make_task(N));
|
||||
|
||||
ASSERT_OK(task_group->Finish());
|
||||
ASSERT_TRUE(task_group->ok());
|
||||
ASSERT_EQ(count.load(), (1 << (N + 1)) - 1);
|
||||
}
|
||||
|
||||
TEST(SerialTaskGroup, Success) { TestTaskGroupSuccess(TaskGroup::MakeSerial()); }
|
||||
|
||||
TEST(SerialTaskGroup, Errors) { TestTaskGroupErrors(TaskGroup::MakeSerial()); }
|
||||
|
||||
TEST(SerialTaskGroup, TasksSpawnTasks) { TestTasksSpawnTasks(TaskGroup::MakeSerial()); }
|
||||
|
||||
TEST(SerialTaskGroup, SubGroupsSuccess) {
|
||||
TestTaskSubGroupsSuccess(TaskGroup::MakeSerial());
|
||||
}
|
||||
|
||||
TEST(SerialTaskGroup, SubGroupsErrors) {
|
||||
TestTaskSubGroupsErrors(TaskGroup::MakeSerial());
|
||||
}
|
||||
|
||||
TEST(ThreadedTaskGroup, Success) {
|
||||
auto task_group = TaskGroup::MakeThreaded(GetCpuThreadPool());
|
||||
TestTaskGroupSuccess(task_group);
|
||||
}
|
||||
|
||||
TEST(ThreadedTaskGroup, Errors) {
|
||||
// Limit parallelism to ensure some tasks don't get started
|
||||
// after the first failing ones
|
||||
std::shared_ptr<ThreadPool> thread_pool;
|
||||
ASSERT_OK(ThreadPool::Make(4, &thread_pool));
|
||||
|
||||
TestTaskGroupErrors(TaskGroup::MakeThreaded(thread_pool.get()));
|
||||
}
|
||||
|
||||
TEST(ThreadedTaskGroup, TasksSpawnTasks) {
|
||||
auto task_group = TaskGroup::MakeThreaded(GetCpuThreadPool());
|
||||
TestTasksSpawnTasks(task_group);
|
||||
}
|
||||
|
||||
TEST(ThreadedTaskGroup, SubGroupsSuccess) {
|
||||
std::shared_ptr<ThreadPool> thread_pool;
|
||||
ASSERT_OK(ThreadPool::Make(4, &thread_pool));
|
||||
|
||||
TestTaskSubGroupsSuccess(TaskGroup::MakeThreaded(thread_pool.get()));
|
||||
}
|
||||
|
||||
TEST(ThreadedTaskGroup, SubGroupsErrors) {
|
||||
std::shared_ptr<ThreadPool> thread_pool;
|
||||
ASSERT_OK(ThreadPool::Make(4, &thread_pool));
|
||||
|
||||
TestTaskSubGroupsErrors(TaskGroup::MakeThreaded(thread_pool.get()));
|
||||
}
|
||||
|
||||
} // namespace internal
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,156 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#include "arrow/util/task-group.h"
|
||||
#include "arrow/util/logging.h"
|
||||
#include "arrow/util/thread-pool.h"
|
||||
|
||||
#include <mutex>
|
||||
|
||||
namespace arrow {
|
||||
namespace internal {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
// Serial TaskGroup implementation
|
||||
|
||||
class SerialTaskGroup : public TaskGroup {
|
||||
public:
|
||||
void AppendReal(std::function<Status()> task) override {
|
||||
DCHECK(!finished_);
|
||||
if (status_.ok()) {
|
||||
status_ &= task();
|
||||
}
|
||||
}
|
||||
|
||||
Status current_status() override { return status_; }
|
||||
|
||||
Status Finish() override {
|
||||
if (!finished_) {
|
||||
finished_ = true;
|
||||
if (parent_) {
|
||||
parent_->status_ &= status_;
|
||||
}
|
||||
}
|
||||
return status_;
|
||||
}
|
||||
|
||||
int parallelism() override { return 1; }
|
||||
|
||||
std::shared_ptr<TaskGroup> MakeSubGroup() override {
|
||||
auto child = new SerialTaskGroup();
|
||||
child->parent_ = this;
|
||||
return std::shared_ptr<TaskGroup>(child);
|
||||
}
|
||||
|
||||
protected:
|
||||
Status status_;
|
||||
bool finished_ = false;
|
||||
SerialTaskGroup* parent_ = nullptr;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
// Threaded TaskGroup implementation
|
||||
|
||||
class ThreadedTaskGroup : public TaskGroup {
|
||||
public:
|
||||
explicit ThreadedTaskGroup(ThreadPool* thread_pool) : thread_pool_(thread_pool) {}
|
||||
|
||||
~ThreadedTaskGroup() {
|
||||
// Make sure all pending tasks are finished, so that dangling references
|
||||
// to this don't persist.
|
||||
ARROW_UNUSED(Finish());
|
||||
}
|
||||
|
||||
void AppendReal(std::function<Status()> task) override {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
DCHECK(!finished_);
|
||||
|
||||
if (status_.ok()) {
|
||||
++nremaining_;
|
||||
status_ = thread_pool_->Spawn([&, task]() {
|
||||
std::unique_lock<std::mutex> lock(mutex_);
|
||||
if (status_.ok()) {
|
||||
lock.unlock();
|
||||
// XXX what about exceptions?
|
||||
Status st = task();
|
||||
lock.lock();
|
||||
status_ &= st;
|
||||
}
|
||||
OneTaskDone();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Status current_status() override {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
return status_;
|
||||
}
|
||||
|
||||
Status Finish() override {
|
||||
std::unique_lock<std::mutex> lock(mutex_);
|
||||
if (!finished_) {
|
||||
cv_.wait(lock, [&]() { return nremaining_ == 0; });
|
||||
// Current tasks may start other tasks, so only set this when done
|
||||
finished_ = true;
|
||||
if (parent_) {
|
||||
// Need to lock parent
|
||||
std::lock_guard<std::mutex> parent_lock(parent_->mutex_);
|
||||
parent_->OneTaskDone();
|
||||
}
|
||||
}
|
||||
return status_;
|
||||
}
|
||||
|
||||
int parallelism() override { return thread_pool_->GetCapacity(); }
|
||||
|
||||
std::shared_ptr<TaskGroup> MakeSubGroup() override {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
auto child = new ThreadedTaskGroup(thread_pool_);
|
||||
child->parent_ = this;
|
||||
nremaining_++;
|
||||
return std::shared_ptr<TaskGroup>(child);
|
||||
}
|
||||
|
||||
protected:
|
||||
void OneTaskDone() {
|
||||
// We are locked
|
||||
--nremaining_;
|
||||
DCHECK_GE(nremaining_, 0);
|
||||
if (nremaining_ == 0) {
|
||||
cv_.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
ThreadPool* thread_pool_;
|
||||
std::mutex mutex_;
|
||||
std::condition_variable cv_;
|
||||
Status status_;
|
||||
bool finished_ = false;
|
||||
int32_t nremaining_ = 0;
|
||||
ThreadedTaskGroup* parent_ = nullptr;
|
||||
};
|
||||
|
||||
std::shared_ptr<TaskGroup> TaskGroup::MakeSerial() {
|
||||
return std::shared_ptr<TaskGroup>(new SerialTaskGroup);
|
||||
}
|
||||
|
||||
std::shared_ptr<TaskGroup> TaskGroup::MakeThreaded(ThreadPool* thread_pool) {
|
||||
return std::shared_ptr<TaskGroup>(new ThreadedTaskGroup(thread_pool));
|
||||
}
|
||||
|
||||
} // namespace internal
|
||||
} // namespace arrow
|
||||
@@ -0,0 +1,90 @@
|
||||
// Licensed to the Apache Software Foundation (ASF) under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing,
|
||||
// software distributed under the License is distributed on an
|
||||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
// KIND, either express or implied. See the License for the
|
||||
// specific language governing permissions and limitations
|
||||
// under the License.
|
||||
|
||||
#ifndef ARROW_UTIL_TASK_GROUP_H
|
||||
#define ARROW_UTIL_TASK_GROUP_H
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "arrow/status.h"
|
||||
#include "arrow/util/macros.h"
|
||||
|
||||
namespace arrow {
|
||||
namespace internal {
|
||||
|
||||
class ThreadPool;
|
||||
|
||||
// TODO Simplify this. Subgroups don't seem necessary.
|
||||
|
||||
/// \brief A group of related tasks
|
||||
///
|
||||
/// A TaskGroup executes tasks with the signature `Status()`.
|
||||
/// Execution can be serial or parallel, depending on the TaskGroup
|
||||
/// implementation. When Finish() returns, it is guaranteed that all
|
||||
/// tasks have finished, or at least one has errored.
|
||||
///
|
||||
class ARROW_EXPORT TaskGroup {
|
||||
public:
|
||||
/// Add a Status-returning function to execute. Execution order is
|
||||
/// undefined. The function may be executed immediately or later.
|
||||
template <typename Function>
|
||||
void Append(Function&& func) {
|
||||
return AppendReal(std::forward<Function>(func));
|
||||
}
|
||||
|
||||
/// Wait for execution of all tasks (and subgroups) to be finished,
|
||||
/// or for at least one task (or subgroup) to error out.
|
||||
/// The returned Status propagates the error status of the first failing
|
||||
/// task (or subgroup).
|
||||
virtual Status Finish() = 0;
|
||||
|
||||
/// The current agregate error Status. Non-blocking, useful for stopping early.
|
||||
virtual Status current_status() = 0;
|
||||
|
||||
/// Whether some tasks have already failed. Non-blocking , useful for stopping early.
|
||||
bool ok() { return current_status().ok(); }
|
||||
|
||||
/// How many tasks can typically be executed in parallel.
|
||||
/// This is only a hint, useful for testing or debugging.
|
||||
virtual int parallelism() = 0;
|
||||
|
||||
/// Create a subgroup of this group. This group can only finish
|
||||
/// when all subgroups have finished (this means you must be
|
||||
/// be careful to call Finish() on subgroups before calling it
|
||||
/// on the main group).
|
||||
// XXX if a subgroup errors out, should it propagate immediately to the parent
|
||||
// and to children?
|
||||
virtual std::shared_ptr<TaskGroup> MakeSubGroup() = 0;
|
||||
|
||||
static std::shared_ptr<TaskGroup> MakeSerial();
|
||||
static std::shared_ptr<TaskGroup> MakeThreaded(internal::ThreadPool*);
|
||||
|
||||
virtual ~TaskGroup() = default;
|
||||
|
||||
protected:
|
||||
TaskGroup() = default;
|
||||
ARROW_DISALLOW_COPY_AND_ASSIGN(TaskGroup);
|
||||
|
||||
virtual void AppendReal(std::function<Status()> task) = 0;
|
||||
};
|
||||
|
||||
} // namespace internal
|
||||
} // namespace arrow
|
||||
|
||||
#endif // ARROW_UTIL_TASK_GROUP_H
|
||||
@@ -366,7 +366,7 @@ if (UNIX)
|
||||
endif()
|
||||
|
||||
set(CYTHON_EXTENSIONS
|
||||
lib
|
||||
lib _csv
|
||||
)
|
||||
|
||||
set(LINK_LIBS
|
||||
|
||||
@@ -0,0 +1,200 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
# cython: profile=False
|
||||
# distutils: language = c++
|
||||
# cython: embedsignature = True
|
||||
|
||||
from pyarrow.includes.common cimport *
|
||||
from pyarrow.includes.libarrow cimport *
|
||||
from pyarrow.lib cimport (check_status, MemoryPool, maybe_unbox_memory_pool,
|
||||
pyarrow_wrap_table, get_reader)
|
||||
|
||||
|
||||
cdef unsigned char _single_char(s) except 0:
|
||||
val = ord(s)
|
||||
if val == 0 or val > 127:
|
||||
raise ValueError("Expecting an ASCII character")
|
||||
return <unsigned char> val
|
||||
|
||||
|
||||
cdef class ReadOptions:
|
||||
cdef:
|
||||
CCSVReadOptions options
|
||||
|
||||
# Avoid mistakingly creating attributes
|
||||
__slots__ = ()
|
||||
|
||||
def __init__(self, use_threads=None, block_size=None, num_rows=None):
|
||||
self.options = CCSVReadOptions.Defaults()
|
||||
if use_threads is not None:
|
||||
self.use_threads = use_threads
|
||||
if block_size is not None:
|
||||
self.block_size = block_size
|
||||
if num_rows is not None:
|
||||
self.num_rows = num_rows
|
||||
|
||||
@property
|
||||
def use_threads(self):
|
||||
return self.options.use_threads
|
||||
|
||||
@use_threads.setter
|
||||
def use_threads(self, value):
|
||||
self.options.use_threads = value
|
||||
|
||||
@property
|
||||
def block_size(self):
|
||||
return self.options.block_size
|
||||
|
||||
@block_size.setter
|
||||
def block_size(self, value):
|
||||
self.options.block_size = value
|
||||
|
||||
@property
|
||||
def num_rows(self):
|
||||
return self.options.num_rows
|
||||
|
||||
@num_rows.setter
|
||||
def num_rows(self, value):
|
||||
self.options.num_rows = value
|
||||
|
||||
|
||||
cdef class ParseOptions:
|
||||
cdef:
|
||||
CCSVParseOptions options
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __init__(self, delimiter=None, quote_char=None, double_quote=None,
|
||||
escape_char=None, header_rows=None):
|
||||
self.options = CCSVParseOptions.Defaults()
|
||||
if delimiter is not None:
|
||||
self.delimiter = delimiter
|
||||
if quote_char is not None:
|
||||
self.quote_char = quote_char
|
||||
if double_quote is not None:
|
||||
self.double_quote = double_quote
|
||||
if escape_char is not None:
|
||||
self.escape_char = escape_char
|
||||
if header_rows is not None:
|
||||
self.header_rows = header_rows
|
||||
|
||||
@property
|
||||
def delimiter(self):
|
||||
return chr(self.options.delimiter)
|
||||
|
||||
@delimiter.setter
|
||||
def delimiter(self, value):
|
||||
self.options.delimiter = _single_char(value)
|
||||
|
||||
@property
|
||||
def quote_char(self):
|
||||
if self.options.quoting:
|
||||
return chr(self.options.quote_char)
|
||||
else:
|
||||
return False
|
||||
|
||||
@quote_char.setter
|
||||
def quote_char(self, value):
|
||||
if value is False:
|
||||
self.options.quoting = False
|
||||
else:
|
||||
self.options.quote_char = _single_char(value)
|
||||
self.options.quoting = True
|
||||
|
||||
@property
|
||||
def double_quote(self):
|
||||
return self.options.double_quote
|
||||
|
||||
@double_quote.setter
|
||||
def double_quote(self, value):
|
||||
self.options.double_quote = value
|
||||
|
||||
@property
|
||||
def escape_char(self):
|
||||
if self.options.escaping:
|
||||
return chr(self.options.escape_char)
|
||||
else:
|
||||
return False
|
||||
|
||||
@escape_char.setter
|
||||
def escape_char(self, value):
|
||||
if value is False:
|
||||
self.options.escaping = False
|
||||
else:
|
||||
self.options.escape_char = _single_char(value)
|
||||
self.options.escaping = True
|
||||
|
||||
@property
|
||||
def header_rows(self):
|
||||
return self.options.header_rows
|
||||
|
||||
@header_rows.setter
|
||||
def header_rows(self, value):
|
||||
self.options.header_rows = value
|
||||
|
||||
|
||||
cdef _get_reader(input_file, shared_ptr[InputStream]* out):
|
||||
cdef shared_ptr[RandomAccessFile] result
|
||||
use_memory_map = False
|
||||
get_reader(input_file, use_memory_map, &result)
|
||||
out[0] = <shared_ptr[InputStream]> result
|
||||
|
||||
|
||||
cdef _get_read_options(ReadOptions read_options, CCSVReadOptions* out):
|
||||
if read_options is None:
|
||||
out[0] = CCSVReadOptions.Defaults()
|
||||
else:
|
||||
out[0] = read_options.options
|
||||
|
||||
|
||||
cdef _get_parse_options(ParseOptions parse_options, CCSVParseOptions* out):
|
||||
if parse_options is None:
|
||||
out[0] = CCSVParseOptions.Defaults()
|
||||
else:
|
||||
out[0] = parse_options.options
|
||||
|
||||
|
||||
cdef _get_convert_options(convert_options, CCSVConvertOptions* out):
|
||||
if convert_options is None:
|
||||
out[0] = CCSVConvertOptions.Defaults()
|
||||
else:
|
||||
raise NotImplementedError("non-default convert options not supported")
|
||||
|
||||
|
||||
def read_csv(input_file, read_options=None, parse_options=None,
|
||||
convert_options=None, MemoryPool memory_pool=None):
|
||||
cdef:
|
||||
shared_ptr[InputStream] stream
|
||||
CCSVReadOptions c_read_options
|
||||
CCSVParseOptions c_parse_options
|
||||
CCSVConvertOptions c_convert_options
|
||||
shared_ptr[CCSVReader] reader
|
||||
shared_ptr[CTable] table
|
||||
|
||||
_get_reader(input_file, &stream)
|
||||
_get_read_options(read_options, &c_read_options)
|
||||
_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))
|
||||
with nogil:
|
||||
check_status(reader.get().Read(&table))
|
||||
|
||||
return pyarrow_wrap_table(table)
|
||||
@@ -0,0 +1,18 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
from pyarrow._csv import ReadOptions, ParseOptions, read_csv # noqa
|
||||
@@ -899,6 +899,41 @@ cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
|
||||
c_string GetColumnName(int i)
|
||||
|
||||
|
||||
cdef extern from "arrow/csv/api.h" namespace "arrow::csv" nogil:
|
||||
|
||||
cdef cppclass CCSVParseOptions" arrow::csv::ParseOptions":
|
||||
unsigned char delimiter
|
||||
c_bool quoting
|
||||
unsigned char quote_char
|
||||
c_bool double_quote
|
||||
c_bool escaping
|
||||
unsigned char escape_char
|
||||
int32_t header_rows
|
||||
|
||||
@staticmethod
|
||||
CCSVParseOptions Defaults()
|
||||
|
||||
cdef cppclass CCSVConvertOptions" arrow::csv::ConvertOptions":
|
||||
@staticmethod
|
||||
CCSVConvertOptions Defaults()
|
||||
|
||||
cdef cppclass CCSVReadOptions" arrow::csv::ReadOptions":
|
||||
c_bool use_threads
|
||||
int32_t block_size
|
||||
int32_t num_rows
|
||||
|
||||
@staticmethod
|
||||
CCSVReadOptions Defaults()
|
||||
|
||||
cdef cppclass CCSVReader" arrow::csv::TableReader":
|
||||
@staticmethod
|
||||
CStatus Make(CMemoryPool*, shared_ptr[InputStream],
|
||||
CCSVReadOptions, CCSVParseOptions, CCSVConvertOptions,
|
||||
shared_ptr[CCSVReader]* out)
|
||||
|
||||
CStatus Read(shared_ptr[CTable]* out)
|
||||
|
||||
|
||||
cdef extern from "arrow/compute/api.h" namespace "arrow::compute" nogil:
|
||||
|
||||
cdef cppclass CFunctionContext" arrow::compute::FunctionContext":
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import io
|
||||
import itertools
|
||||
import string
|
||||
import unittest
|
||||
|
||||
import pytest
|
||||
|
||||
import numpy as np
|
||||
|
||||
import pyarrow as pa
|
||||
from pyarrow.csv import read_csv, ReadOptions, ParseOptions
|
||||
|
||||
|
||||
def generate_col_names():
|
||||
# 'a', 'b'... 'z', then 'aa', 'ab'...
|
||||
letters = string.ascii_lowercase
|
||||
for letter in letters:
|
||||
yield letter
|
||||
for first in letter:
|
||||
for second in letter:
|
||||
yield first + second
|
||||
|
||||
|
||||
def make_random_csv(num_cols=2, num_rows=10, linesep=u'\r\n'):
|
||||
arr = np.random.RandomState(42).randint(0, 1000, size=(num_cols, num_rows))
|
||||
col_names = list(itertools.islice(generate_col_names(), num_cols))
|
||||
csv = io.StringIO()
|
||||
csv.write(u",".join(col_names))
|
||||
csv.write(linesep)
|
||||
for row in arr.T:
|
||||
csv.write(u",".join(map(str, row)))
|
||||
csv.write(linesep)
|
||||
csv = csv.getvalue().encode()
|
||||
columns = [pa.array(a, type=pa.int64()) for a in arr]
|
||||
expected = pa.Table.from_arrays(columns, col_names)
|
||||
return csv, expected
|
||||
|
||||
|
||||
def test_read_options():
|
||||
cls = ReadOptions
|
||||
opts = cls()
|
||||
|
||||
assert opts.block_size > 0
|
||||
opts.block_size = 12345
|
||||
assert opts.block_size == 12345
|
||||
|
||||
assert opts.use_threads is True
|
||||
opts.use_threads = False
|
||||
assert opts.use_threads is False
|
||||
|
||||
assert opts.num_rows > 0
|
||||
opts.num_rows = 456789
|
||||
assert opts.num_rows == 456789
|
||||
|
||||
opts = cls(block_size=1234, use_threads=False, num_rows=42)
|
||||
assert opts.block_size == 1234
|
||||
assert opts.use_threads is False
|
||||
assert opts.num_rows == 42
|
||||
|
||||
|
||||
def test_parse_options():
|
||||
cls = ParseOptions
|
||||
opts = cls()
|
||||
assert opts.delimiter == ','
|
||||
assert opts.quote_char == '"'
|
||||
assert opts.double_quote is True
|
||||
assert opts.escape_char is False
|
||||
assert opts.header_rows == 1
|
||||
|
||||
opts.delimiter = 'x'
|
||||
assert opts.delimiter == 'x'
|
||||
assert opts.quote_char == '"'
|
||||
|
||||
opts.escape_char = 'z'
|
||||
assert opts.escape_char == 'z'
|
||||
assert opts.quote_char == '"'
|
||||
|
||||
opts.quote_char = False
|
||||
assert opts.quote_char is False
|
||||
assert opts.escape_char == 'z'
|
||||
|
||||
opts.escape_char = False
|
||||
assert opts.escape_char is False
|
||||
assert opts.quote_char is False
|
||||
|
||||
opts.header_rows = 2
|
||||
assert opts.header_rows == 2
|
||||
|
||||
opts = cls(delimiter=';', quote_char='%', double_quote=False,
|
||||
escape_char='\\', header_rows=2)
|
||||
assert opts.delimiter == ';'
|
||||
assert opts.quote_char == '%'
|
||||
assert opts.double_quote is False
|
||||
assert opts.escape_char == '\\'
|
||||
assert opts.header_rows == 2
|
||||
|
||||
|
||||
class BaseTestCSVRead:
|
||||
|
||||
def read_bytes(self, b, **kwargs):
|
||||
return self.read_csv(pa.py_buffer(b), **kwargs)
|
||||
|
||||
def check_names(self, table, names):
|
||||
assert table.num_columns == len(names)
|
||||
assert [c.name for c in table.columns] == names
|
||||
|
||||
def test_header(self):
|
||||
rows = b"abc,def,gh\n"
|
||||
table = self.read_bytes(rows)
|
||||
assert isinstance(table, pa.Table)
|
||||
self.check_names(table, ["abc", "def", "gh"])
|
||||
assert table.num_rows == 0
|
||||
|
||||
def test_simple_ints(self):
|
||||
# Infer integer columns
|
||||
rows = b"a,b,c\n1,2,3\n4,5,6\n"
|
||||
table = self.read_bytes(rows)
|
||||
schema = pa.schema([('a', pa.int64()),
|
||||
('b', pa.int64()),
|
||||
('c', pa.int64())])
|
||||
assert table.schema == schema
|
||||
assert table.to_pydict() == {
|
||||
'a': [1, 4],
|
||||
'b': [2, 5],
|
||||
'c': [3, 6],
|
||||
}
|
||||
|
||||
def test_simple_varied(self):
|
||||
# Infer various kinds of data
|
||||
rows = b"a,b,c\n1,2,3\n4.0,-5,foo\n"
|
||||
table = self.read_bytes(rows)
|
||||
schema = pa.schema([('a', pa.float64()),
|
||||
('b', pa.int64()),
|
||||
('c', pa.binary())])
|
||||
assert table.schema == schema
|
||||
assert table.to_pydict() == {
|
||||
'a': [1.0, 4.0],
|
||||
'b': [2, -5],
|
||||
'c': [b"3", b"foo"],
|
||||
}
|
||||
|
||||
def test_simple_nulls(self):
|
||||
# Infer various kinds of data, with nulls
|
||||
rows = (b"a,b,c,d\n"
|
||||
b"1,2,,\n"
|
||||
b"nan,-5,foo,\n"
|
||||
b"4.5,#N/A,nan,\n")
|
||||
table = self.read_bytes(rows)
|
||||
schema = pa.schema([('a', pa.float64()),
|
||||
('b', pa.int64()),
|
||||
('c', pa.binary()),
|
||||
('d', pa.null())])
|
||||
assert table.schema == schema
|
||||
assert table.to_pydict() == {
|
||||
'a': [1.0, None, 4.5],
|
||||
'b': [2, -5, None],
|
||||
'c': [b"", b"foo", b"nan"],
|
||||
'd': [None, None, None]
|
||||
}
|
||||
|
||||
def test_no_ending_newline(self):
|
||||
# No \n after last line
|
||||
rows = b"a,b,c\n1,2,3\n4,5,6"
|
||||
table = self.read_bytes(rows)
|
||||
assert table.to_pydict() == {
|
||||
'a': [1, 4],
|
||||
'b': [2, 5],
|
||||
'c': [3, 6],
|
||||
}
|
||||
|
||||
def test_trivial(self):
|
||||
# A bit pointless, but at least it shouldn't crash
|
||||
rows = b"\n\n"
|
||||
table = self.read_bytes(rows)
|
||||
assert table.to_pydict() == {'': [None]}
|
||||
|
||||
def test_invalid_csv(self):
|
||||
# Various CSV errors
|
||||
rows = b"a,b,c\n1,2\n4,5,6\n"
|
||||
with pytest.raises(pa.ArrowInvalid, match="Expected 3 columns, got 2"):
|
||||
self.read_bytes(rows)
|
||||
rows = b"a,b,c\n1,2,3\n4"
|
||||
with pytest.raises(pa.ArrowInvalid, match="Expected 3 columns, got 1"):
|
||||
self.read_bytes(rows)
|
||||
rows = b""
|
||||
with pytest.raises(pa.ArrowInvalid, match="Empty CSV file"):
|
||||
self.read_bytes(rows)
|
||||
|
||||
def test_options_delimiter(self):
|
||||
rows = b"a;b,c\nde,fg;eh\n"
|
||||
table = self.read_bytes(rows)
|
||||
assert table.to_pydict() == {
|
||||
'a;b': [b'de'],
|
||||
'c': [b'fg;eh'],
|
||||
}
|
||||
opts = ParseOptions(delimiter=';')
|
||||
table = self.read_bytes(rows, parse_options=opts)
|
||||
assert table.to_pydict() == {
|
||||
'a': [b'de,fg'],
|
||||
'b,c': [b'eh'],
|
||||
}
|
||||
|
||||
def test_small_random_csv(self):
|
||||
csv, expected = make_random_csv(num_cols=2, num_rows=10)
|
||||
table = self.read_bytes(csv)
|
||||
assert table.schema == expected.schema
|
||||
assert table.equals(expected)
|
||||
assert table.to_pydict() == expected.to_pydict()
|
||||
|
||||
def test_stress_block_sizes(self):
|
||||
# Test a number of small block sizes to stress block stitching
|
||||
csv_base, expected = make_random_csv(num_cols=2, num_rows=500)
|
||||
block_sizes = [11, 12, 13, 17, 37, 111]
|
||||
csvs = [csv_base, csv_base.rstrip(b'\r\n')]
|
||||
for csv in csvs:
|
||||
for block_size in block_sizes:
|
||||
read_options = ReadOptions(block_size=block_size)
|
||||
table = self.read_bytes(csv, read_options=read_options)
|
||||
assert table.schema == expected.schema
|
||||
if not table.equals(expected):
|
||||
# Better error output
|
||||
assert table.to_pydict() == expected.to_pydict()
|
||||
|
||||
|
||||
class TestSerialCSVRead(BaseTestCSVRead, unittest.TestCase):
|
||||
|
||||
def read_csv(self, *args, **kwargs):
|
||||
read_options = kwargs.setdefault('read_options', ReadOptions())
|
||||
read_options.use_threads = False
|
||||
table = read_csv(*args, **kwargs)
|
||||
table._validate()
|
||||
return table
|
||||
|
||||
|
||||
class TestParallelCSVRead(BaseTestCSVRead, unittest.TestCase):
|
||||
|
||||
def read_csv(self, *args, **kwargs):
|
||||
read_options = kwargs.setdefault('read_options', ReadOptions())
|
||||
read_options.use_threads = True
|
||||
table = read_csv(*args, **kwargs)
|
||||
table._validate()
|
||||
return table
|
||||
@@ -156,6 +156,7 @@ class build_ext(_build_ext):
|
||||
|
||||
CYTHON_MODULE_NAMES = [
|
||||
'lib',
|
||||
'_csv',
|
||||
'_cuda',
|
||||
'_parquet',
|
||||
'_orc',
|
||||
|
||||
Reference in new issue
Block a user