Fix arrow dictionary import.

Signed-off-by: ienkovich <[email protected]>

Author: ienkovich

omniscidb/ArrowStorage/ArrowStorage.cpp

@@ -730,6 +730,9 @@ void ArrowStorage::appendArrowTable(std::shared_ptr<arrow::Table> at, int table_
             CHECK(false);
         }
       } else if (col_type->isString()) {
+        if (col_arr->type()->id() == arrow::Type::DICTIONARY) {
+          col_arr = decodeArrowDictionary(col_arr);
+        }
       } else {
         col_arr =
             replaceNullValues(col_arr,
@@ -1118,6 +1121,19 @@ void ArrowStorage::compareSchemas(std::shared_ptr<arrow::Schema> lhs,
     auto lhs_type = lhs_fields[i]->type();
     auto rhs_type = rhs_fields[i]->type();
 
+    // For string and dictionary columns we allow both dictionary and plain strings on
+    // import.
+    if (lhs_type->id() == arrow::Type::DICTIONARY) {
+      lhs_type = static_cast<const arrow::DictionaryType*>(lhs_type.get())->value_type();
+    }
+    if (rhs_type->id() == arrow::Type::DICTIONARY) {
+      rhs_type = static_cast<const arrow::DictionaryType*>(rhs_type.get())->value_type();
+      if (rhs_type->id() != arrow::Type::STRING) {
+        throw std::runtime_error("Unsupported dictionary type: "s +
+                                 rhs_fields[i]->type()->ToString());
+      }
+    }
+
     if (!lhs_type->Equals(rhs_type) && (lhs_type->id() != arrow::Type::NA) &&
         (rhs_type->id() != arrow::Type::NA)) {
       throw std::runtime_error(

omniscidb/ArrowStorage/ArrowStorageUtils.cpp

@@ -17,6 +17,7 @@
 #include "IR/Context.h"
 #include "Shared/InlineNullValues.h"
 
+#include <arrow/compute/api.h>
 #include <tbb/parallel_for.h>
 #include <tbb/task_group.h>
 
@@ -1187,6 +1188,26 @@ std::shared_ptr<arrow::ChunkedArray> createDictionaryEncodedColumn(
   return nullptr;
 }
 
+namespace {
+
+template <typename INDEX_TYPE>
+void convertArrowDictionaryIndices(const std::vector<int>& indices_mapping,
+                                   std::shared_ptr<arrow::Array> indices,
+                                   int32_t* out_indices) {
+  using ArrowType = typename arrow::CTypeTraits<INDEX_TYPE>::ArrowType;
+  using ArrayType = typename arrow::TypeTraits<ArrowType>::ArrayType;
+  auto casted_indices = std::static_pointer_cast<ArrayType>(indices);
+  for (int i = 0; i < casted_indices->length(); i++) {
+    if (casted_indices->IsValid(i)) {
+      out_indices[i] = indices_mapping[casted_indices->Value(i)];
+    } else {
+      out_indices[i] = inline_int_null_value<int32_t>();
+    }
+  }
+}
+
+}  // namespace
+
 std::shared_ptr<arrow::ChunkedArray> convertArrowDictionary(
     StringDictionary* dict,
     std::shared_ptr<arrow::ChunkedArray> arr,
@@ -1195,35 +1216,67 @@ std::shared_ptr<arrow::ChunkedArray> convertArrowDictionary(
     throw std::runtime_error("Unsupported HDK dictionary for Arrow dictionary import: "s +
                              type->toString());
   }
+
+  // Create new arrow buffer to hold remapped indices
+  std::shared_ptr<arrow::Buffer> dict_indices_buf;
+  auto res = arrow::AllocateBuffer(arr->length() * sizeof(int32_t));
+  CHECK(res.ok());
+  dict_indices_buf = std::move(res).ValueOrDie();
+  auto cur_raw_data = reinterpret_cast<int32_t*>(dict_indices_buf->mutable_data());
+
   // TODO: allocate one big array and split it by fragments as it is done in
   // createDictionaryEncodedColumn
-  std::vector<std::shared_ptr<arrow::Array>> converted_chunks;
   for (auto& chunk : arr->chunks()) {
     auto dict_array = std::static_pointer_cast<arrow::DictionaryArray>(chunk);
     auto values = std::static_pointer_cast<arrow::StringArray>(dict_array->dictionary());
     std::vector<std::string_view> strings(values->length());
+    strings.reserve(values->length());
     for (int i = 0; i < values->length(); i++) {
       auto view = values->GetView(i);
       strings[i] = std::string_view(view.data(), view.length());
     }
-    auto arrow_indices =
-        std::static_pointer_cast<arrow::Int32Array>(dict_array->indices());
     std::vector<int> indices_mapping(values->length());
     dict->getOrAddBulk(strings, indices_mapping.data());
 
-    // create new arrow chunk with remapped indices
-    std::shared_ptr<arrow::Buffer> dict_indices_buf;
-    auto res = arrow::AllocateBuffer(arrow_indices->length() * sizeof(int32_t));
-    CHECK(res.ok());
-    dict_indices_buf = std::move(res).ValueOrDie();
-    auto raw_data = reinterpret_cast<int32_t*>(dict_indices_buf->mutable_data());
-
-    for (int i = 0; i < arrow_indices->length(); i++) {
-      raw_data[i] = indices_mapping[arrow_indices->Value(i)];
+    auto arrow_indices = dict_array->indices();
+    switch (arrow_indices->type_id()) {
+      case arrow::Type::INT8:
+        convertArrowDictionaryIndices<int8_t>(
+            indices_mapping, arrow_indices, cur_raw_data);
+        break;
+      case arrow::Type::INT16:
+        convertArrowDictionaryIndices<int16_t>(
+            indices_mapping, arrow_indices, cur_raw_data);
+        break;
+      case arrow::Type::INT32:
+        convertArrowDictionaryIndices<int32_t>(
+            indices_mapping, arrow_indices, cur_raw_data);
+        break;
+      case arrow::Type::INT64:
+        convertArrowDictionaryIndices<int64_t>(
+            indices_mapping, arrow_indices, cur_raw_data);
+        break;
+      default:
+        throw std::runtime_error("Unsupported Arrow dictionary for import: "s +
+                                 arr->type()->ToString());
     }
 
+    cur_raw_data += chunk->length();
+  }
+  std::vector<std::shared_ptr<arrow::Array>> converted_chunks;
+  converted_chunks.push_back(
+      std::make_shared<arrow::Int32Array>(arr->length(), dict_indices_buf));
+  return std::make_shared<arrow::ChunkedArray>(converted_chunks);
+}
+
+std::shared_ptr<arrow::ChunkedArray> decodeArrowDictionary(
+    std::shared_ptr<arrow::ChunkedArray> arr) {
+  std::vector<std::shared_ptr<arrow::Array>> converted_chunks;
+  for (auto& chunk : arr->chunks()) {
+    auto dict_arr = std::dynamic_pointer_cast<arrow::DictionaryArray>(chunk);
+    CHECK(dict_arr);
     converted_chunks.push_back(
-        std::make_shared<arrow::Int32Array>(arrow_indices->length(), dict_indices_buf));
+        arrow::compute::Take(*dict_arr->dictionary(), *dict_arr->indices()).ValueOrDie());
   }
   return std::make_shared<arrow::ChunkedArray>(converted_chunks);
 }

omniscidb/ArrowStorage/ArrowStorageUtils.h

@@ -44,3 +44,6 @@ std::shared_ptr<arrow::ChunkedArray> convertArrowDictionary(
     StringDictionary* dict,
     std::shared_ptr<arrow::ChunkedArray> arr,
     const hdk::ir::Type* type);
+
+std::shared_ptr<arrow::ChunkedArray> decodeArrowDictionary(
+    std::shared_ptr<arrow::ChunkedArray> arr);

omniscidb/Tests/ArrowStorageTest.cpp

@@ -323,21 +323,28 @@ void checkStringColumnData(ArrowStorage& storage,
   size_t end_row = std::min(row_count, start_row + fragment_size);
   size_t frag_rows = end_row - start_row;
   size_t chunk_size = 0;
+  bool has_nulls = false;
   for (size_t i = start_row; i < end_row; ++i) {
-    chunk_size += vals[i].size();
+    if (vals[i] != "<NULL>") {
+      chunk_size += vals[i].size();
+    } else {
+      has_nulls = true;
+    }
   }
   checkChunkMeta(chunk_meta_map.at(col_id),
                  storage.getColumnInfo(TEST_DB_ID, table_id, col_id)->type,
                  frag_rows,
                  chunk_size,
-                 false);
+                 has_nulls);
   std::vector<int8_t> expected_data(chunk_size);
   std::vector<uint32_t> expected_offset(frag_rows + 1);
   uint32_t data_offset = 0;
   for (size_t i = start_row; i < end_row; ++i) {
     expected_offset[i - start_row] = data_offset;
-    memcpy(expected_data.data() + data_offset, vals[i].data(), vals[i].size());
-    data_offset += vals[i].size();
+    if (vals[i] != "<NULL>") {
+      memcpy(expected_data.data() + data_offset, vals[i].data(), vals[i].size());
+      data_offset += vals[i].size();
+    }
   }
   expected_offset.back() = data_offset;
   checkFetchedData(storage, table_id, col_id, frag_idx + 1, expected_offset, {2});
@@ -361,17 +368,28 @@ void checkStringDictColumnData(ArrowStorage& storage,
   auto& dict = *storage.getDictMetadata(getDictId(col_info->type))->stringDict;
 
   std::vector<IndexType> expected_ids(frag_rows);
+  bool has_nulls = false;
+  IndexType min = std::numeric_limits<IndexType>::max();
+  IndexType max = std::numeric_limits<IndexType>::min();
   for (size_t i = start_row; i < end_row; ++i) {
-    expected_ids[i - start_row] = static_cast<IndexType>(dict.getIdOfString(expected[i]));
+    if (expected[i] == "<NULL>") {
+      expected_ids[i - start_row] = inline_int_null_value<IndexType>();
+      has_nulls = true;
+    } else {
+      expected_ids[i - start_row] =
+          static_cast<IndexType>(dict.getIdOfString(expected[i]));
+      min = std::min(min, expected_ids[i - start_row]);
+      max = std::max(max, expected_ids[i - start_row]);
+    }
   }
 
   checkChunkMeta(chunk_meta_map.at(col_id),
                  col_info->type,
                  frag_rows,
                  frag_rows * sizeof(IndexType),
-                 false,
-                 *std::min_element(expected_ids.begin(), expected_ids.end()),
-                 *std::max_element(expected_ids.begin(), expected_ids.end()));
+                 has_nulls,
+                 min,
+                 max);
 
   checkFetchedData(storage, table_id, col_id, frag_idx + 1, expected_ids);
 }
@@ -1754,6 +1772,109 @@ TEST_F(ArrowStorageTest, ImportParquet) {
             std::vector<double>({1.1, 2.2, 3.3, 4.4, 5.5}));
 }
 
+namespace {
+
+template <typename INDEX_TYPE,
+          bool NULL_INDICES = false,
+          bool NULL_VALUES = false,
+          bool TARGET_DICT = true>
+void TestImportArrowDict(ConfigPtr config) {
+  ArrowStorage storage(TEST_SCHEMA_ID, "test", TEST_DB_ID, config);
+  auto tinfo = storage.createTable(
+      "table1",
+      {{"col1",
+        TARGET_DICT ? static_cast<const hdk::ir::Type*>(ctx.extDict(ctx.text(), 0))
+                    : static_cast<const hdk::ir::Type*>(ctx.text())}});
+
+  using IndexArrowType = typename arrow::CTypeTraits<INDEX_TYPE>::ArrowType;
+  using IndexBuilder = typename arrow::TypeTraits<IndexArrowType>::BuilderType;
+
+  std::vector<std::shared_ptr<arrow::Array>> arrays;
+  IndexBuilder index_builder;
+  ARROW_THROW_NOT_OK(index_builder.Append(0));
+  ARROW_THROW_NOT_OK(index_builder.Append(1));
+  ARROW_THROW_NOT_OK(index_builder.Append(0));
+  ARROW_THROW_NOT_OK(index_builder.Append(1));
+  arrow::StringBuilder value_builder;
+  ARROW_THROW_NOT_OK(value_builder.Append("str1"));
+  ARROW_THROW_NOT_OK(value_builder.Append("str2"));
+  arrays.push_back(arrow::DictionaryArray::FromArrays(index_builder.Finish().ValueOrDie(),
+                                                      value_builder.Finish().ValueOrDie())
+                       .ValueOrDie());
+  ARROW_THROW_NOT_OK(index_builder.Append(0));
+  ARROW_THROW_NOT_OK(index_builder.Append(1));
+  if (NULL_INDICES) {
+    ARROW_THROW_NOT_OK(index_builder.AppendNull());
+  } else {
+    ARROW_THROW_NOT_OK(index_builder.Append(1));
+  }
+  ARROW_THROW_NOT_OK(index_builder.Append(2));
+  ARROW_THROW_NOT_OK(index_builder.Append(2));
+  if (NULL_VALUES) {
+    ARROW_THROW_NOT_OK(value_builder.AppendNull());
+  } else {
+    ARROW_THROW_NOT_OK(value_builder.Append("str2"));
+  }
+  ARROW_THROW_NOT_OK(value_builder.Append("str3"));
+  ARROW_THROW_NOT_OK(value_builder.Append("str4"));
+  arrays.push_back(arrow::DictionaryArray::FromArrays(index_builder.Finish().ValueOrDie(),
+                                                      value_builder.Finish().ValueOrDie())
+                       .ValueOrDie());
+  auto chunked_arr = std::make_shared<arrow::ChunkedArray>(arrays);
+
+  arrow::SchemaBuilder schema_builder;
+  ARROW_THROW_NOT_OK(schema_builder.AddField(
+      std::make_shared<arrow::Field>("col1", chunked_arr->type())));
+  auto schema = schema_builder.Finish().ValueOrDie();
+  auto at = arrow::Table::Make(schema, {chunked_arr});
+
+  storage.appendArrowTable(at, "table1");
+
+  checkData(storage,
+            tinfo->table_id,
+            9,
+            32'000'000,
+            std::vector<std::string>({"str1"s,
+                                      "str2"s,
+                                      "str1"s,
+                                      "str2"s,
+                                      NULL_VALUES ? "<NULL>"s : "str2"s,
+                                      "str3"s,
+                                      NULL_INDICES ? "<NULL>"s : "str3"s,
+                                      "str4"s,
+                                      "str4"s}));
+}
+
+}  // namespace
+
+TEST_F(ArrowStorageTest, ImportArrowDict8) {
+  TestImportArrowDict<int8_t>(config_);
+}
+
+TEST_F(ArrowStorageTest, ImportArrowDict16) {
+  TestImportArrowDict<int16_t>(config_);
+}
+
+TEST_F(ArrowStorageTest, ImportArrowDict32) {
+  TestImportArrowDict<int32_t>(config_);
+}
+
+TEST_F(ArrowStorageTest, ImportArrowDict64) {
+  TestImportArrowDict<int64_t>(config_);
+}
+
+TEST_F(ArrowStorageTest, ImportArrowDict_Null_Indices) {
+  TestImportArrowDict<int32_t, true, false>(config_);
+}
+
+TEST_F(ArrowStorageTest, ImportArrowDict_Null_Values) {
+  TestImportArrowDict<int32_t, false, true>(config_);
+}
+
+TEST_F(ArrowStorageTest, ImportArrowDictToPlainString) {
+  TestImportArrowDict<int32_t, true, true, false>(config_);
+}
+
 int main(int argc, char** argv) {
   TestHelpers::init_logger_stderr_only(argc, argv);
   testing::InitGoogleTest(&argc, argv);

python/tests/test_pyhdk_data_import.py

@@ -9,6 +9,8 @@
 import pyhdk
 import pyarrow
 
+from helpers import check_res
+
 
 class TestImport:
     def test_null_schema(self):
@@ -19,3 +21,25 @@ def test_null_schema(self):
         hdk = pyhdk.init()
         ht = hdk.import_arrow(table)
         hdk.drop_table(ht)
+
+    def test_dict_import(self):
+        hdk = pyhdk.init()
+        ht = hdk.create_table("table1", {"col1": "dict", "col2": "text"})
+
+        col1 = pyarrow.array(["str1", "str2"])
+        at = pyarrow.table([col1, col1], names=["col1", "col2"])
+        hdk.import_arrow(at, ht)
+
+        col1 = pyarrow.DictionaryArray.from_arrays([0, 1, 0, 1], ["str3", "str4"])
+        at = pyarrow.table([col1, col1], names=["col1", "col2"])
+        hdk.import_arrow(at, ht)
+
+        check_res(
+            ht.run(),
+            {
+                "col1": ["str1", "str2", "str3", "str4", "str3", "str4"],
+                "col2": ["str1", "str2", "str3", "str4", "str3", "str4"],
+            },
+        )
+
+        hdk.drop_table(ht)