update tests to test dictionary read of string columns

kszucs · Feb 23, 2020 · 7546851 · 7546851
1 parent 22caaa1
commit 7546851
Show file tree

Hide file tree

Showing 6 changed files with 32 additions and 8 deletions.
diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc
@@ -187,6 +187,8 @@ TEST_F(TestParquetFileFormat, ScanRecordBatchReader) {
 }
 
 TEST_F(TestParquetFileFormat, ScanRecordBatchReaderDictEncoded) {
+  schema_ = schema({field("utf8", utf8())});
+
   auto reader = GetRecordBatchReader();
   auto source = GetFileSource(reader.get());
 
@@ -198,7 +200,7 @@ TEST_F(TestParquetFileFormat, ScanRecordBatchReaderDictEncoded) {
   ASSERT_OK_AND_ASSIGN(auto scan_task_it, fragment->Scan(ctx_));
   int64_t row_count = 0;
 
-  Schema expected_schema({field("f64", dictionary(int32(), float64()))});
+  Schema expected_schema({field("utf8", dictionary(int32(), utf8()))});
 
   for (auto maybe_task : scan_task_it) {
     ASSERT_OK_AND_ASSIGN(auto task, std::move(maybe_task));
@@ -309,13 +311,15 @@ TEST_F(TestParquetFileFormat, Inspect) {
 }
 
 TEST_F(TestParquetFileFormat, InspectDictEncoded) {
+  schema_ = schema({field("utf8", utf8())});
+
   auto reader = GetRecordBatchReader();
   auto source = GetFileSource(reader.get());
 
   format_->read_dict_indices.insert(0);
   ASSERT_OK_AND_ASSIGN(auto actual, format_->Inspect(*source.get()));
 
-  Schema expected_schema({field("f64", dictionary(int32(), float64()))});
+  Schema expected_schema({field("utf8", dictionary(int32(), utf8()))});
   EXPECT_EQ(*actual, expected_schema);
 }
 

diff --git a/cpp/src/arrow/status.cc b/cpp/src/arrow/status.cc
@@ -47,9 +47,12 @@ std::string Status::CodeAsString() const {
   if (state_ == nullptr) {
     return "OK";
   }
+  return CodeAsString(code());
+}
 
+std::string Status::CodeAsString(StatusCode code) {
   const char* type;
-  switch (code()) {
+  switch (code) {
     case StatusCode::OK:
       type = "OK";
       break;

diff --git a/cpp/src/arrow/status.h b/cpp/src/arrow/status.h
@@ -302,6 +302,7 @@ class ARROW_EXPORT Status : public util::EqualityComparable<Status>,
   /// \brief Return a string representation of the status code, without the message
   /// text or POSIX code information.
   std::string CodeAsString() const;
+  static std::string CodeAsString(StatusCode);
 
   /// \brief Return the StatusCode value attached to this status.
   StatusCode code() const { return ok() ? StatusCode::OK : state_->code; }

diff --git a/cpp/src/arrow/testing/generator.cc b/cpp/src/arrow/testing/generator.cc
@@ -36,9 +36,9 @@ namespace arrow {
 
 template <typename ArrowType, typename CType = typename TypeTraits<ArrowType>::CType,
           typename BuilderType = typename TypeTraits<ArrowType>::BuilderType>
-static inline std::shared_ptr<Array> ConstantArray(int64_t size, CType value = 0) {
+static inline std::shared_ptr<Array> ConstantArray(int64_t size, CType value) {
   auto type = TypeTraits<ArrowType>::type_singleton();
-  auto builder_fn = [](BuilderType* builder) { builder->UnsafeAppend(CType(0)); };
+  auto builder_fn = [&](BuilderType* builder) { builder->UnsafeAppend(value); };
   return ArrayFromBuilderVisitor(type, size, builder_fn).ValueOrDie();
 }
 
@@ -90,4 +90,9 @@ std::shared_ptr<arrow::Array> ConstantArrayGenerator::Float64(int64_t size,
   return ConstantArray<DoubleType>(size, value);
 }
 
+std::shared_ptr<arrow::Array> ConstantArrayGenerator::String(int64_t size,
+                                                             std::string value) {
+  return ConstantArray<StringType>(size, value);
+}
+
 }  // namespace arrow
diff --git a/cpp/src/arrow/testing/generator.h b/cpp/src/arrow/testing/generator.h
@@ -19,6 +19,7 @@
 
 #include <cstdint>
 #include <memory>
+#include <string>
 #include <vector>
 
 #include "arrow/record_batch.h"
@@ -118,6 +119,14 @@ class ARROW_EXPORT ConstantArrayGenerator {
   /// \return a generated Array
   static std::shared_ptr<arrow::Array> Float64(int64_t size, double value = 0);
 
+  /// \brief Generates a constant StringArray
+  ///
+  /// \param[in] size the size of the array to generate
+  /// \param[in] value to repeat
+  ///
+  /// \return a generated Array
+  static std::shared_ptr<arrow::Array> String(int64_t size, std::string value = "");
+
   template <typename ArrowType, typename CType = typename ArrowType::c_type>
   static std::shared_ptr<arrow::Array> Numeric(int64_t size, CType value = 0) {
     switch (ArrowType::type_id) {
@@ -179,6 +188,8 @@ class ARROW_EXPORT ConstantArrayGenerator {
         return Float32(size);
       case Type::DOUBLE:
         return Float64(size);
+      case Type::STRING:
+        return String(size);
       default:
         return nullptr;
     }

diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc
@@ -2689,9 +2689,9 @@ void TryReadDataFile(const std::string& path,
     s = arrow_reader->ReadTable(&table);
   }
 
-  ASSERT_TRUE(s.code() == expected_code)
-      << "Expected reading file to return "
-      << arrow::Status(expected_code, "").CodeAsString() << ", but got " << s.ToString();
+  ASSERT_EQ(s.code(), expected_code)
+      << "Expected reading file to return " << arrow::Status::CodeAsString(expected_code)
+      << ", but got " << s.ToString();
 }
 
 TEST(TestArrowReaderAdHoc, Int96BadMemoryAccess) {