ARROW-3769: [C++] Add support for reading non-dictionary encoded binary Parquet columns directly as DictionaryArray

cpp/src/arrow/testing/random.cc

@@ -79,7 +79,7 @@ std::shared_ptr<Array> RandomArrayGenerator::Boolean(int64_t size, double probab
   // only calls the GenerateBitmap method.
   using GenOpt = GenerateOptions<int, std::uniform_int_distribution<int>>;
 
-  std::vector<std::shared_ptr<Buffer>> buffers{2};
+  BufferVector buffers{2};
   // Need 2 distinct generators such that probabilities are not shared.
   GenOpt value_gen(seed(), 0, 1, probability);
   GenOpt null_gen(seed(), 0, 1, null_probability);
@@ -100,7 +100,7 @@ static std::shared_ptr<NumericArray<ArrowType>> GenerateNumericArray(int64_t siz
                                                                      OptionType options) {
   using CType = typename ArrowType::c_type;
   auto type = TypeTraits<ArrowType>::type_singleton();
-  std::vector<std::shared_ptr<Buffer>> buffers{2};
+  BufferVector buffers{2};
 
   int64_t null_count = 0;
   ABORT_NOT_OK(AllocateEmptyBitmap(size, &buffers[0]));
@@ -145,5 +145,64 @@ PRIMITIVE_RAND_FLOAT_IMPL(Float64, double, DoubleType)
 #undef PRIMITIVE_RAND_FLOAT_IMPL
 #undef PRIMITIVE_RAND_IMPL
 
+std::shared_ptr<arrow::Array> RandomArrayGenerator::String(int64_t size,
+                                                           int32_t min_length,
+                                                           int32_t max_length,
+                                                           double null_probability) {
+  if (null_probability < 0 || null_probability > 1) {
+    ABORT_NOT_OK(Status::Invalid("null_probability must be between 0 and 1"));
+  }
+
+  auto int32_lengths = Int32(size, min_length, max_length, null_probability);
+  auto lengths = std::dynamic_pointer_cast<Int32Array>(int32_lengths);
+
+  // Visual Studio does not implement uniform_int_distribution for char types.
+  using GenOpt = GenerateOptions<uint8_t, std::uniform_int_distribution<uint16_t>>;
+  GenOpt options(seed(), static_cast<uint8_t>('A'), static_cast<uint8_t>('z'),
+                 /*null_probability=*/0);
+
+  std::vector<uint8_t> str_buffer(max_length);
+  StringBuilder builder;
+
+  for (int64_t i = 0; i < size; ++i) {
+    if (lengths->IsValid(i)) {
+      options.GenerateData(str_buffer.data(), lengths->Value(i));
+      ABORT_NOT_OK(builder.Append(str_buffer.data(), lengths->Value(i)));
+    } else {
+      ABORT_NOT_OK(builder.AppendNull());
+    }
+  }
+
+  std::shared_ptr<arrow::Array> result;
+  ABORT_NOT_OK(builder.Finish(&result));
+  return result;
+}
+
+std::shared_ptr<arrow::Array> RandomArrayGenerator::StringWithRepeats(
+    int64_t size, int64_t unique, int32_t min_length, int32_t max_length,
+    double null_probability) {
+  // Generate a random string dictionary without any nulls
+  auto array = String(unique, min_length, max_length, /*null_probability=*/0);
+  auto dictionary = std::dynamic_pointer_cast<StringArray>(array);
+
+  // Generate random indices to sample the dictionary with
+  auto id_array = Int64(size, 0, unique - 1, null_probability);
+  auto indices = std::dynamic_pointer_cast<Int64Array>(id_array);
+  StringBuilder builder;
+
+  for (int64_t i = 0; i < size; ++i) {
+    if (indices->IsValid(i)) {
+      const auto index = indices->Value(i);
+      const auto value = dictionary->GetView(index);
+      ABORT_NOT_OK(builder.Append(value));
+    } else {
+      ABORT_NOT_OK(builder.AppendNull());
+    }
+  }
+
+  std::shared_ptr<Array> result;
+  ABORT_NOT_OK(builder.Finish(&result));
+  return result;
+}
 }  // namespace random
 }  // namespace arrow

cpp/src/arrow/testing/random.h

@@ -198,6 +198,35 @@ class ARROW_EXPORT RandomArrayGenerator {
     }
   }
 
+  /// \brief Generates a random StringArray
+  ///
+  /// \param[in] size the size of the array to generate
+  /// \param[in] min_length the lower bound of the string length
+  ///            determined by the uniform distribution
+  /// \param[in] max_length the upper bound of the string length
+  ///            determined by the uniform distribution
+  /// \param[in] null_probability the probability of a row being null
+  ///
+  /// \return a generated Array
+  std::shared_ptr<arrow::Array> String(int64_t size, int32_t min_length,
+                                       int32_t max_length, double null_probability);
+
+  /// \brief Generates a random StringArray with repeated values
+  ///
+  /// \param[in] size the size of the array to generate
+  /// \param[in] unique the number of unique string values used
+  ///            to populate the array
+  /// \param[in] min_length the lower bound of the string length
+  ///            determined by the uniform distribution
+  /// \param[in] max_length the upper bound of the string length
+  ///            determined by the uniform distribution
+  /// \param[in] null_probability the probability of a row being null
+  ///
+  /// \return a generated Array
+  std::shared_ptr<arrow::Array> StringWithRepeats(int64_t size, int64_t unique,
+                                                  int32_t min_length, int32_t max_length,
+                                                  double null_probability);
+
  private:
   SeedType seed() { return seed_distribution_(seed_rng_); }
 

cpp/src/parquet/arrow/arrow-reader-writer-test.cc

@@ -30,6 +30,7 @@
 #include <vector>
 
 #include "arrow/api.h"
+#include "arrow/testing/random.h"
 #include "arrow/testing/util.h"
 #include "arrow/type_traits.h"
 #include "arrow/util/decimal.h"
@@ -2439,6 +2440,80 @@ TEST(TestArrowWriterAdHoc, SchemaMismatch) {
   ASSERT_RAISES(Invalid, writer->WriteTable(*tbl, 1));
 }
 
+// ----------------------------------------------------------------------
+// Tests for directly reading DictionaryArray
+class TestArrowReadDictionary : public ::testing::TestWithParam<double> {
+ public:
+  void SetUp() override {
+    GenerateData(GetParam());
+    ASSERT_NO_FATAL_FAILURE(
+        WriteTableToBuffer(expected_dense_, expected_dense_->num_rows() / 2,
+                           default_arrow_writer_properties(), &buffer_));
+
+    properties_ = default_arrow_reader_properties();
+  }
+
+  void GenerateData(double null_probability) {
+    constexpr int num_unique = 100;
+    constexpr int repeat = 10;
+    constexpr int64_t min_length = 2;
+    constexpr int64_t max_length = 10;
+    ::arrow::random::RandomArrayGenerator rag(0);
+    auto dense_array = rag.StringWithRepeats(repeat * num_unique, num_unique, min_length,
+                                             max_length, null_probability);
+    expected_dense_ = MakeSimpleTable(dense_array, /*nullable=*/true);
+
+    ::arrow::StringDictionaryBuilder builder(default_memory_pool());
+    const auto& string_array = static_cast<const ::arrow::StringArray&>(*dense_array);
+    ASSERT_OK(builder.AppendArray(string_array));
+
+    std::shared_ptr<::arrow::Array> dict_array;
+    ASSERT_OK(builder.Finish(&dict_array));
+    expected_dict_ = MakeSimpleTable(dict_array, /*nullable=*/true);
+
+    // TODO(hatemhelal): Figure out if we can use the following to init the expected_dict_
+    // Currently fails due to DataType mismatch for indices array.
+    //    Datum out;
+    //    FunctionContext ctx(default_memory_pool());
+    //    ASSERT_OK(DictionaryEncode(&ctx, Datum(dense_array), &out));
+    //    expected_dict_ = MakeSimpleTable(out.make_array(), /*nullable=*/true);
+  }
+
+  void TearDown() override {}
+
+  void CheckReadWholeFile(const Table& expected) {
+    std::unique_ptr<FileReader> reader;
+    ASSERT_OK_NO_THROW(OpenFile(std::make_shared<BufferReader>(buffer_),
+                                ::arrow::default_memory_pool(), properties_, &reader));
+
+    std::shared_ptr<Table> actual;
+    ASSERT_OK_NO_THROW(reader->ReadTable(&actual));
+    ::arrow::AssertTablesEqual(*actual, expected, /*same_chunk_layout=*/false);
+  }
+
+  static std::vector<double> null_probabilites() { return {0.0, 0.5, 1}; }
+
+ protected:
+  std::shared_ptr<Table> expected_dense_;
+  std::shared_ptr<Table> expected_dict_;
+  std::shared_ptr<Buffer> buffer_;
+  ArrowReaderProperties properties_;
+};
+
+TEST_P(TestArrowReadDictionary, ReadWholeFileDict) {
+  properties_.set_read_dictionary(0, true);
+  CheckReadWholeFile(*expected_dict_);
+}
+
+TEST_P(TestArrowReadDictionary, ReadWholeFileDense) {
+  properties_.set_read_dictionary(0, false);
+  CheckReadWholeFile(*expected_dense_);
+}
+
+INSTANTIATE_TEST_CASE_P(
+    ReadDictionary, TestArrowReadDictionary,
+    ::testing::ValuesIn(TestArrowReadDictionary::null_probabilites()));
+
 }  // namespace arrow
 
 }  // namespace parquet