[18102873455] Handle sparse string columns on arrow read (#2700)

IvoDD · web-flow · commit 2cb7ff2d607f · 2025-10-13T15:04:01.000+03:00
#### Reference Issues/PRs Monday ref: 18102873455 #### What does this implement or fix? Change consists of two commits (which can be reviewed separately): First one (tests and benchmarks): - Adds a python test via ternary operator which reads sparse data - Adds a benchmark to verify type handler performance - Renames `test_arrow.py` to `test_arrow_read.py` Second one (actual fix): - Correctly handle sparse input string column and preserve the speed of dense handling (by filling the inverted bitset) #### Any other comments? This change slightly slows down (`503us` -> `594us`) string processing for arrow (most noticably if very few unique strings). Experimentation shows that 100% of the slowdown can be attributed to failure to inline [this lambda function](https://github.com/man-group/ArcticDB/pull/2700/files#diff-4e399a492a48cda55aa054cc5f6f30a73f35d8ff3aa11dbd32a1f8f2b17598b4R77-R94) argument to `for_each_enumerated`. If we unroll `for_each_enumerated` performance goes back to original level. This slowdown will be fixed in a follow up PR. (which might positively impact a lot of processing pipeline affected by a similar issue) Benchmark results: ``` benchmarks --benchmark_filter="BM_arrow_string.*" --benchmark_min_time=2000x --benchmark_repetitions=5 --benchmark_time_unit=us After first commit: ----------------------------------------------------------------------------------------- Benchmark Time CPU Iterations ----------------------------------------------------------------------------------------- BM_arrow_string_handler/10000/1/0_mean 59.7 us 65.1 us 5 BM_arrow_string_handler/10000/1/0_stddev 1.19 us 1.29 us 5 BM_arrow_string_handler/100000/1/0_mean 503 us 544 us 5 BM_arrow_string_handler/100000/1/0_stddev 8.52 us 3.18 us 5 BM_arrow_string_handler/10000/10000/0_mean 202 us 210 us 5 BM_arrow_string_handler/10000/10000/0_stddev 12.6 us 2.61 us 5 BM_arrow_string_handler/100000/100000/0_mean 2707 us 2906 us 5 BM_arrow_string_handler/100000/100000/0_stddev 52.1 us 13.0 us 5 After second commit: ----------------------------------------------------------------------------------------- Benchmark Time CPU Iterations ----------------------------------------------------------------------------------------- BM_arrow_string_handler/10000/1/0_mean 65.6 us 71.6 us 5 BM_arrow_string_handler/10000/1/0_stddev 1.01 us 1.10 us 5 BM_arrow_string_handler/100000/1/0_mean 594 us 645 us 5 BM_arrow_string_handler/100000/1/0_stddev 8.67 us 6.29 us 5 BM_arrow_string_handler/10000/10000/0_mean 203 us 215 us 5 BM_arrow_string_handler/10000/10000/0_stddev 0.838 us 0.894 us 5 BM_arrow_string_handler/100000/100000/0_mean 2850 us 3046 us 5 BM_arrow_string_handler/100000/100000/0_stddev 118 us 57.9 us 5 BM_arrow_string_handler/10000/1/5000_mean 69.2 us 75.5 us 5 BM_arrow_string_handler/10000/1/5000_stddev 0.815 us 0.892 us 5 BM_arrow_string_handler/100000/1/50000_mean 614 us 670 us 5 BM_arrow_string_handler/100000/1/50000_stddev 4.92 us 5.36 us 5 BM_arrow_string_handler/10000/1/10000_mean 6.91 us 7.53 us 5 BM_arrow_string_handler/10000/1/10000_stddev 0.161 us 0.179 us 5 BM_arrow_string_handler/100000/1/100000_mean 7.99 us 8.39 us 5 BM_arrow_string_handler/100000/1/100000_stddev 0.259 us 0.121 us 5 ``` #### Checklist <details> <summary> Checklist for code changes... </summary> - [ ] Have you updated the relevant docstrings, documentation and copyright notice? - [ ] Is this contribution tested against [all ArcticDB's features](../docs/mkdocs/docs/technical/contributing.md)? - [ ] Do all exceptions introduced raise appropriate [error messages](https://docs.arcticdb.io/error_messages/)? - [ ] Are API changes highlighted in the PR description? - [ ] Is the PR labelled as enhancement or bug so it appears in autogenerated release notes? </details>
diff --git a/cpp/arcticdb/CMakeLists.txt b/cpp/arcticdb/CMakeLists.txt
@@ -1105,6 +1105,7 @@ if(${TEST})
     set(benchmark_srcs
             stream/test/stream_test_common.cpp
             arrow/test/arrow_test_utils.hpp
+            arrow/test/benchmark_arrow_reads.cpp
             arrow/test/benchmark_arrow_writes.cpp
             column_store/test/benchmark_column.cpp
             column_store/test/benchmark_memory_segment.cpp
diff --git a/cpp/arcticdb/arrow/arrow_handlers.cpp b/cpp/arcticdb/arrow/arrow_handlers.cpp
@@ -11,6 +11,7 @@
 #include <arcticdb/util/decode_path_data.hpp>
 #include <arcticdb/pipeline/column_mapping.hpp>
 #include <arcticdb/column_store/string_pool.hpp>
+#include <arcticdb/column_store/column_algorithms.hpp>
 
 namespace arcticdb {
 
@@ -51,7 +52,6 @@ void ArrowStringHandler::convert_type(
         std::any&, const std::shared_ptr<StringPool>& string_pool
 ) const {
     using ArcticStringColumnTag = ScalarTagType<DataTypeTag<DataType::UTF_DYNAMIC64>>;
-    auto input_data = source_column.data();
     struct DictEntry {
         int32_t offset_buffer_pos_;
         int64_t string_buffer_pos_;
@@ -65,49 +65,49 @@ void ArrowStringHandler::convert_type(
     unique_offsets.reserve(source_column.row_count());
     int64_t bytes = 0;
     int32_t unique_offset_count = 0;
-    auto dest_ptr = reinterpret_cast<int32_t*>(
-            dest_column.bytes_at(mapping.offset_bytes_, source_column.row_count() * sizeof(int32_t))
-    );
+    auto dest_ptr = reinterpret_cast<int32_t*>(dest_column.bytes_at(mapping.offset_bytes_, mapping.dest_bytes_));
 
-    util::BitSet bitset;
-    util::BitSet::bulk_insert_iterator inserter(bitset);
-    const auto end = input_data.cend<ArcticStringColumnTag, IteratorType::ENUMERATED>();
-    // First go through the source column once to compute the size of offset and string buffers.
-    // TODO: This can't be right if the column was sparse as it has only been decoded, not expanded
-    for (auto en = input_data.cbegin<ArcticStringColumnTag, IteratorType::ENUMERATED>(); en != end; ++en) {
-        if (is_a_string(en->value())) {
+    util::BitSet dest_bitset;
+    util::BitSet::bulk_insert_iterator inserter(dest_bitset);
+    // For dense columns we populate an inverted bitmap (i.e. set the bits where values are missing)
+    // This makes processing fully dense columns faster because it doesn't need to insert often into the bitmap.
+    // Benchmarks in benchmark_arrow_reads.cpp show a 20% speedup with this approach for a dense string column with few
+    // unique strings.
+    bool populate_inverted_bitset = !source_column.opt_sparse_map().has_value();
+    for_each_enumerated<ArcticStringColumnTag>(source_column, [&](const auto& en) {
+        if (is_a_string(en.value())) {
             auto [entry, is_emplaced] = unique_offsets.try_emplace(
-                    en->value(),
-                    DictEntry{
-                            static_cast<int32_t>(unique_offset_count), bytes, string_pool->get_const_view(en->value())
-                    }
+                    en.value(), DictEntry{unique_offset_count, bytes, string_pool->get_const_view(en.value())}
             );
             if (is_emplaced) {
                 bytes += entry->second.strv.size();
-                unique_offsets_in_order.push_back(en->value());
+                unique_offsets_in_order.push_back(en.value());
                 ++unique_offset_count;
             }
-            *dest_ptr = entry->second.offset_buffer_pos_;
-        } else {
-            inserter = en->idx();
+            dest_ptr[en.idx()] = entry->second.offset_buffer_pos_;
+            if (!populate_inverted_bitset) {
+                inserter = en.idx();
+            }
+        } else if (populate_inverted_bitset) {
+            inserter = en.idx();
         }
-        ++dest_ptr;
-    }
+    });
     inserter.flush();
-    // At this point bitset has ones where the source column contained None or NaN
-    // Inverting and shrinking to the source column size it then makes a sparse map for the input data
-    bitset.invert();
-    // TODO: row_count() here won't be right when the original data was sparse, but we don't support sparse
-    // string columns yet anyway
-    bitset.resize(source_column.row_count());
-    if (bitset.count() != bitset.size()) {
-        handle_truncation(bitset, mapping.truncate_);
-        create_dense_bitmap(mapping.offset_bytes_, bitset, dest_column, AllocationType::DETACHABLE);
-    } // else there weren't any Nones or NaNs
-    // bitset.count() == 0 is the special case where all of the rows contained None or NaN. In this case, do not create
+    if (populate_inverted_bitset) {
+        // For dense columns at this point bitset has ones where the source column is missing values.
+        // We invert it back, so we can use it for the sparse map on the output column.
+        dest_bitset.invert();
+    }
+    dest_bitset.resize(mapping.num_rows_);
+
+    if (dest_bitset.count() != dest_bitset.size()) {
+        handle_truncation(dest_bitset, mapping.truncate_);
+        create_dense_bitmap(mapping.offset_bytes_, dest_bitset, dest_column, AllocationType::DETACHABLE);
+    } // else there weren't any missing values
+    // bitset.count() == 0 is the special case where all the rows were missing. In this case, do not create
     // the extra string and offset buffers. string_dict_from_block will then do the right thing and call
     // minimal_strings_dict
-    if (bitset.count() > 0) {
+    if (dest_bitset.count() > 0) {
         auto& string_buffer = dest_column.create_extra_buffer(
                 mapping.offset_bytes_, ExtraBufferType::STRING, bytes, AllocationType::DETACHABLE
         );
diff --git a/cpp/arcticdb/arrow/test/arrow_test_utils.hpp b/cpp/arcticdb/arrow/test/arrow_test_utils.hpp
@@ -5,10 +5,12 @@
  * As of the Change Date specified in that file, in accordance with the Business Source License, use of this software
  * will be governed by the Apache License, version 2.0.
  */
+#pragma once
 
 #include <sparrow/record_batch.hpp>
 
 #include <arcticdb/util/allocator.hpp>
+#include <arcticdb/column_store/column.hpp>
 
 using namespace arcticdb;
 
@@ -30,10 +32,33 @@ sparrow::array create_array(const std::vector<T>& data) {
     }
 }
 
-sparrow::record_batch create_record_batch(const std::vector<std::pair<std::string, sparrow::array>>& columns) {
+inline sparrow::record_batch create_record_batch(const std::vector<std::pair<std::string, sparrow::array>>& columns) {
     sparrow::record_batch record_batch{};
     for (const auto& column : columns) {
         record_batch.add_column(column.first, column.second);
     }
     return record_batch;
 }
+
+template<typename RawType>
+void allocate_and_fill_chunked_column(
+        Column& column, size_t num_rows, size_t chunk_size, std::optional<std::span<RawType>> values = std::nullopt
+) {
+    // Allocate column in chunks
+    for (size_t row = 0; row < num_rows; row += chunk_size) {
+        auto data_size = data_type_size(column.type(), OutputFormat::ARROW, DataTypeMode::EXTERNAL);
+        auto current_block_size = std::min(chunk_size, num_rows - row);
+        auto bytes = current_block_size * data_size;
+        column.allocate_data(bytes);
+        column.advance_data(bytes);
+    }
+
+    // Actually fill the data
+    for (size_t row = 0; row < num_rows; ++row) {
+        if (values.has_value()) {
+            column.reference_at<RawType>(row) = values.value()[row];
+        } else {
+            column.reference_at<RawType>(row) = static_cast<RawType>(row);
+        }
+    }
+}
diff --git a/cpp/arcticdb/arrow/test/benchmark_arrow_reads.cpp b/cpp/arcticdb/arrow/test/benchmark_arrow_reads.cpp
@@ -0,0 +1,78 @@
+/* Copyright 2025 Man Group Operations Limited
+ *
+ * Use of this software is governed by the Business Source License 1.1 included in the file licenses/BSL.txt.
+ *
+ * As of the Change Date specified in that file, in accordance with the Business Source License, use of this software
+ * will be governed by the Apache License, version 2.0.
+ */
+
+#include <benchmark/benchmark.h>
+#include <fmt/format.h>
+
+#include <arcticdb/arrow/test/arrow_test_utils.hpp>
+#include <arcticdb/arrow/arrow_handlers.hpp>
+#include <arcticdb/column_store/string_pool.hpp>
+#include <arcticdb/pipeline/column_mapping.hpp>
+
+using namespace arcticdb;
+
+// run like: --benchmark_time_unit=ms --benchmark_filter=.* --benchmark_min_time=5x
+
+static void BM_arrow_string_handler(benchmark::State& state) {
+    const auto num_rows = state.range(0);
+    const auto num_different_strings = state.range(1);
+    const auto num_sparse = state.range(2);
+
+    const double sparsity_ratio = num_sparse / num_rows;
+    auto source_type_desc = TypeDescriptor{DataType::UTF_DYNAMIC64, Dimension::Dim0};
+    auto dest_type_desc = TypeDescriptor{DataType::UTF_DYNAMIC32, Dimension::Dim0};
+    auto dest_size = data_type_size(dest_type_desc, OutputFormat::ARROW, DataTypeMode::EXTERNAL);
+    auto sparsity = num_sparse == 0 ? Sparsity::NOT_PERMITTED : Sparsity::PERMITTED;
+
+    // Fill up source column
+    auto source_column = Column(source_type_desc, num_rows, AllocationType::DYNAMIC, sparsity);
+    auto string_pool = std::make_shared<StringPool>();
+    for (auto i = 0u, num_set = 0u; i < num_rows; i++) {
+        auto offset = string_pool->get(fmt::format("str_{}", i % num_different_strings)).offset();
+        auto expected_set = std::round((i + 1) * (1 - sparsity_ratio));
+        if (num_set < expected_set) {
+            source_column.set_scalar(i, offset);
+            ++num_set;
+        }
+    }
+    auto mapping = ColumnMapping(
+            source_type_desc,
+            dest_type_desc,
+            FieldWrapper(dest_type_desc, "col").field(),
+            dest_size,
+            num_rows,
+            0,
+            0,
+            dest_size * num_different_strings,
+            0
+    );
+    auto handler = ArrowStringHandler();
+    auto handler_data = std::any{};
+
+    for (auto _ : state) {
+        state.PauseTiming();
+        auto dest_column = Column(dest_type_desc, 0, AllocationType::DETACHABLE, sparsity);
+        allocate_and_fill_chunked_column<int32_t>(dest_column, num_rows, num_rows);
+        state.ResumeTiming();
+        handler.convert_type(source_column, dest_column, mapping, DecodePathData{}, handler_data, string_pool);
+    }
+}
+
+BENCHMARK(BM_arrow_string_handler)
+        // Not sparse, small string and offset buffers
+        ->Args({10'000, 1, 0})
+        ->Args({100'000, 1, 0})
+        // Not sparse, large string and offset buffers
+        ->Args({10'000, 10'000, 0})
+        ->Args({100'000, 100'000, 0})
+        // Half sparse
+        ->Args({10'000, 1, 5'000})
+        ->Args({100'000, 1, 50'000})
+        // Fully sparse
+        ->Args({10'000, 1, 10'000})
+        ->Args({100'000, 1, 100'000});
diff --git a/cpp/arcticdb/arrow/test/test_arrow_read.cpp b/cpp/arcticdb/arrow/test/test_arrow_read.cpp
@@ -12,34 +12,12 @@
 
 #include <arcticdb/pipeline/column_mapping.hpp>
 #include <arcticdb/stream/test/stream_test_common.hpp>
+#include <arcticdb/arrow/test/arrow_test_utils.hpp>
 #include <arcticdb/arrow/arrow_utils.hpp>
 #include <arcticdb/arrow/arrow_handlers.hpp>
 
 using namespace arcticdb;
 
-template<typename RawType>
-void allocate_and_fill_chunked_column(
-        Column& column, size_t num_rows, size_t chunk_size, std::optional<std::span<RawType>> values = std::nullopt
-) {
-    // Allocate column in chunks
-    for (size_t row = 0; row < num_rows; row += chunk_size) {
-        auto data_size = data_type_size(column.type(), OutputFormat::ARROW, DataTypeMode::EXTERNAL);
-        auto current_block_size = std::min(chunk_size, num_rows - row);
-        auto bytes = current_block_size * data_size;
-        column.allocate_data(bytes);
-        column.advance_data(bytes);
-    }
-
-    // Actually fill the data
-    for (size_t row = 0; row < num_rows; ++row) {
-        if (values.has_value()) {
-            column.reference_at<RawType>(row) = values.value()[row];
-        } else {
-            column.reference_at<RawType>(row) = static_cast<RawType>(row);
-        }
-    }
-}
-
 SegmentInMemory get_detachable_segment(
         StreamId symbol, std::span<const FieldRef> fields, size_t num_rows, size_t chunk_size
 ) {
diff --git a/python/tests/conftest.py b/python/tests/conftest.py
@@ -1123,6 +1123,14 @@ def lmdb_version_store_dynamic_schema(
         raise ValueError(f"Unexpected encoding version: {encoding_version}")
 
 
+@pytest.fixture
+def lmdb_version_store_dynamic_schema_arrow(lmdb_version_store_dynamic_schema_v1) -> NativeVersionStore:
+    store = lmdb_version_store_dynamic_schema_v1
+    store.set_output_format(OutputFormat.EXPERIMENTAL_ARROW)
+    store._set_allow_arrow_input()
+    return store
+
+
 @pytest.fixture
 def lmdb_version_store_empty_types_v1(version_store_factory, lib_name) -> NativeVersionStore:
     library_name = lib_name + "_v1"
diff --git a/python/tests/unit/arcticdb/version_store/test_arrow_read.py b/python/tests/unit/arcticdb/version_store/test_arrow_read.py
