perf(read): skip the payload Concatenate when late materialization matches a single run

src/paimon/common/predicate/multi_literals_leaf_function.cpp

@@ -31,6 +31,7 @@
 #include "arrow/memory_pool.h"
 #include "arrow/type.h"
 #include "paimon/common/predicate/literal_converter.h"
+#include "paimon/common/utils/arrow/arrow_utils.h"
 #include "paimon/common/utils/arrow/status_utils.h"
 #include "paimon/common/utils/checked_cast.h"
 #include "paimon/data/decimal.h"
@@ -201,15 +202,9 @@ Result<std::vector<char>> ProbeInValueSet(const arrow::Array& array, const arrow
     // last owner at the end of the statement and leave that reference dangling.
     std::shared_ptr<arrow::Array> matches_array = matches.make_array();
     const auto& matched = checked_cast<const arrow::BooleanArray&>(*matches_array);
-    std::vector<char> is_valid(matched.length(), 0);
-    for (int64_t i = 0; i < matched.length(); i++) {
-        if (matched.IsNull(i)) {
-            // `IN` and `NOT IN` are both false on a null value, leave the row at 0.
-            continue;
-        }
-        is_valid[i] = static_cast<char>(matched.Value(i) != negate);
-    }
-    return is_valid;
+    // `IN` and `NOT IN` are both false on a null value, so the null rows unpack to 0 rather than to
+    // the negation of a value they do not hold.
+    return ArrowUtils::UnpackBooleansToBytes(matched, negate);
 }
 }  // namespace
 

src/paimon/common/predicate/null_false_leaf_binary_function.cpp

@@ -32,6 +32,7 @@
 #include "arrow/scalar.h"
 #include "arrow/type.h"
 #include "paimon/common/predicate/literal_converter.h"
+#include "paimon/common/utils/arrow/arrow_utils.h"
 #include "paimon/common/utils/arrow/status_utils.h"
 #include "paimon/common/utils/checked_cast.h"
 #include "paimon/common/utils/field_type_utils.h"
@@ -196,13 +197,9 @@ Result<std::vector<char>> ProbeComparison(const arrow::Array& array, const char*
     // last owner at the end of the statement and leave that reference dangling.
     std::shared_ptr<arrow::Array> matches_array = matches.make_array();
     const auto& matched = checked_cast<const arrow::BooleanArray&>(*matches_array);
-    std::vector<char> is_valid(matched.length(), 0);
-    for (int64_t i = 0; i < matched.length(); i++) {
-        // A kernel emits a null for a null row, and every `NullFalseLeafBinaryFunction` is false on
-        // one, so a null row is left at 0 rather than read as a value.
-        is_valid[i] = static_cast<char>(!matched.IsNull(i) && matched.Value(i));
-    }
-    return is_valid;
+    // A kernel emits a null for a null row, and every `NullFalseLeafBinaryFunction` is false on
+    // one, so the null rows unpack to 0 rather than to a value they do not hold.
+    return ArrowUtils::UnpackBooleansToBytes(matched, /*negate=*/false);
 }
 }  // namespace
 

src/paimon/common/reader/late_materializing_file_batch_reader.cpp

@@ -193,8 +193,19 @@ Result<FileBatchReader::ReadBatch> LateMaterializingFileBatchReader::ReadPayload
         // Compact the payload superset down to the matched rows (ascending file row order).
         PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector payload_slices,
                                ReaderUtils::GenerateFilteredArrayVector(payload_array, valid));
-        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> payload_compacted,
-                                          arrow::Concatenate(payload_slices, arrow_pool_.get()));
+        std::shared_ptr<arrow::Array> payload_compacted;
+        if (payload_slices.size() == 1) {
+            // A single run has nothing to merge, and `AssembleFullBatch` normalizes each column
+            // anyway, so hand the slice over rather than let `Concatenate` copy the payload a
+            // second time. A run at the batch head then stays zero-copy; a mid-batch run is still
+            // copied once by that normalization, so this is never worse. Where it stays zero-copy
+            // the slice keeps the batch's buffers alive until the consumer releases the assembled
+            // batch, bounded by the batch size times the prefetch queue depth.
+            payload_compacted = std::move(payload_slices.front());
+        } else {
+            PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+                payload_compacted, arrow::Concatenate(payload_slices, arrow_pool_.get()));
+        }
 
         auto card = static_cast<int64_t>(valid.Cardinality());
         if (probe_cursor_ + card > probe_data_->length()) {

src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp

@@ -29,6 +29,7 @@
 #include <vector>
 
 #include "arrow/api.h"
+#include "arrow/array/builder_dict.h"
 #include "arrow/array/builder_nested.h"
 #include "arrow/c/bridge.h"
 #include "gtest/gtest.h"
@@ -214,6 +215,31 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test {
         return array;
     }
 
+    // Build a struct with a dictionary-encoded payload column [k:int64, v:dictionary<int32,utf8>].
+    std::shared_ptr<arrow::Array> BuildDictionaryPayloadData(int32_t n) {
+        auto type =
+            arrow::struct_({arrow::field("k", arrow::int64()),
+                            arrow::field("v", arrow::dictionary(arrow::int32(), arrow::utf8()))});
+        // `StringDictionaryBuilder` picks the narrowest index type its dictionary fits in - int8
+        // for the three values below - and `StructBuilder::type()` takes each field type from its
+        // builder, so the column would come out as `dictionary<int8,utf8>` whatever `type` says.
+        // The 32 builder is what keeps the index width this fixture declares.
+        arrow::StructBuilder builder(type, arrow::default_memory_pool(),
+                                     {std::make_shared<arrow::Int64Builder>(),
+                                      std::make_shared<arrow::StringDictionary32Builder>()});
+        auto* k = checked_cast<arrow::Int64Builder*>(builder.field_builder(0));
+        auto* v = checked_cast<arrow::StringDictionary32Builder*>(builder.field_builder(1));
+        for (int32_t i = 0; i < n; ++i) {
+            EXPECT_TRUE(builder.Append().ok());
+            EXPECT_TRUE(k->Append(i).ok());
+            // Repeat a few values so the dictionary is shorter than the column it encodes.
+            EXPECT_TRUE(v->Append("d_" + std::to_string(i % 3)).ok());
+        }
+        std::shared_ptr<arrow::Array> array;
+        EXPECT_TRUE(builder.Finish(&array).ok());
+        return array;
+    }
+
  protected:
     std::shared_ptr<MemoryPool> pool_;
     std::shared_ptr<arrow::Field> k_field_;
@@ -533,6 +559,61 @@ TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) {
     }
 }
 
+// A payload column that the inner reader hands over dictionary-encoded has to stay that way in the
+// output. Compacting a batch down to a single matched run passes the run straight to the assembly
+// step rather than concatenating it, so on that path the encoding survives only because nothing
+// rewrites the array. Both shapes of a lone run are covered: one starting at the batch head, which
+// reaches the output untouched, and one starting mid-batch, which the per-column offset
+// normalization still copies.
+TEST_F(LateMaterializingFileBatchReaderTest, DictionaryPayloadColumnKeepsEncoding) {
+    auto data = BuildDictionaryPayloadData(8);
+    auto type = data->type();
+    auto mock = std::make_unique<MockFileBatchReader>(data, type, /*batch_size=*/3);
+    ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(
+                                          std::move(mock), GetArrowPool(pool_)));
+    // probe = {k}; payload = {v}. k >= 5 matches only row 5 of the batch holding rows 3..5, a run
+    // at batch offset 2, and both rows of the batch holding rows 6..7, a run at batch offset 0.
+    auto predicate =
+        PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l));
+    ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(type->fields()), predicate, std::nullopt));
+
+    std::vector<std::pair<int64_t, std::string>> rows;
+    while (true) {
+        ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap,
+                             reader->NextBatchWithBitmap());
+        if (BatchReader::IsEofBatch(batch_with_bitmap)) {
+            break;
+        }
+        auto& [batch, bitmap] = batch_with_bitmap;
+        auto& [c_array, c_schema] = batch;
+        arrow::Result<std::shared_ptr<arrow::Array>> imported =
+            arrow::ImportArray(c_array.get(), c_schema.get());
+        ASSERT_TRUE(imported.ok()) << imported.status().ToString();
+        auto struct_array = arrow::internal::checked_pointer_cast<arrow::StructArray>(*imported);
+        auto k = arrow::internal::checked_pointer_cast<arrow::Int64Array>(
+            struct_array->GetFieldByName("k"));
+        std::shared_ptr<arrow::Array> v = struct_array->GetFieldByName("v");
+        ASSERT_TRUE(k && v);
+        // The index width is checked before the casts below, which a release build resolves
+        // statically and would otherwise read the index buffer at the wrong width.
+        ASSERT_TRUE(v->type()->Equals(*arrow::dictionary(arrow::int32(), arrow::utf8())))
+            << v->type()->ToString();
+        auto dict = arrow::internal::checked_pointer_cast<arrow::DictionaryArray>(v);
+        auto indices = arrow::internal::checked_pointer_cast<arrow::Int32Array>(dict->indices());
+        auto values = arrow::internal::checked_pointer_cast<arrow::StringArray>(dict->dictionary());
+        ASSERT_TRUE(indices && values);
+        for (int64_t i = 0; i < struct_array->length(); ++i) {
+            rows.emplace_back(k->Value(i), values->GetString(indices->Value(i)));
+        }
+    }
+    ASSERT_EQ(rows.size(), 3u);  // k = 5,6,7
+    for (size_t i = 0; i < rows.size(); ++i) {
+        int64_t expected_k = 5 + static_cast<int64_t>(i);
+        EXPECT_EQ(rows[i].first, expected_k);
+        EXPECT_EQ(rows[i].second, "d_" + std::to_string(expected_k % 3));
+    }
+}
+
 // The late-materialization reader must work correctly as an inner reader driven by
 // PrefetchFileBatchReaderImpl (schema broadcast, range dispatch, seek, row-id tracking).
 TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) {

src/paimon/common/utils/arrow/arrow_utils.cpp

@@ -19,9 +19,16 @@
 
 #include "paimon/common/utils/arrow/arrow_utils.h"
 
+#include <algorithm>
+#include <array>
+#include <cstdint>
+#include <cstring>
+
 #include "arrow/array/array_base.h"
 #include "arrow/array/array_nested.h"
+#include "arrow/array/array_primitive.h"
 #include "arrow/array/concatenate.h"
+#include "arrow/array/data.h"
 #include "arrow/array/util.h"
 #include "arrow/buffer.h"
 #include "arrow/c/abi.h"
@@ -334,6 +341,31 @@ Result<std::shared_ptr<arrow::ArrayData>> RebaseToZeroOffset(
     return CopyToZeroOffset(data, pool);
 }
 
+/// The eight bytes one byte of a bitmap expands to: `kExpandedBytes[b].bytes[i]` is bit `i` of `b`.
+/// Naming the bytes one at a time rather than building a word keeps the expansion independent of
+/// the host's byte order.
+struct ExpandedByte {
+    char bytes[8];
+};
+
+constexpr ExpandedByte ExpandBitmapByte(uint8_t bits) {
+    ExpandedByte expanded{};
+    for (int32_t i = 0; i < 8; i++) {
+        expanded.bytes[i] = static_cast<char>((bits >> i) & uint8_t{1});
+    }
+    return expanded;
+}
+
+constexpr std::array<ExpandedByte, 256> MakeExpandedBytes() {
+    std::array<ExpandedByte, 256> table{};
+    for (int32_t bits = 0; bits < 256; bits++) {
+        table[static_cast<size_t>(bits)] = ExpandBitmapByte(static_cast<uint8_t>(bits));
+    }
+    return table;
+}
+
+constexpr std::array<ExpandedByte, 256> kExpandedBytes = MakeExpandedBytes();
+
 }  // namespace
 
 const char* ArrowUtils::kArrowSchemaMetadataKey = "ARROW:schema";
@@ -432,6 +464,50 @@ uint64_t ArrowUtils::GetArrayMemoryUsage(const std::shared_ptr<arrow::ArrayData>
     return result;
 }
 
+std::vector<char> ArrowUtils::UnpackBooleansToBytes(const arrow::BooleanArray& array, bool negate) {
+    const int64_t length = array.length();
+    std::vector<char> is_valid(static_cast<size_t>(length));
+    if (length == 0) {
+        return is_valid;
+    }
+
+    const std::shared_ptr<arrow::ArrayData>& data = array.data();
+    // `BooleanArray::Value(i)` and `Array::IsValid(i)` both read their bitmap at bit `i + offset`
+    // from the start of the buffer, so one bit index serves the value and the validity alike.
+    const int64_t bit_offset = data->offset;
+    const auto values = data->GetValuesSafe<uint8_t>(/*i=*/1, /*absolute_offset=*/0);
+    const uint8_t* validity = array.null_bitmap_data();
+
+    auto unpack_row = [&](int64_t row) -> char {
+        const int64_t bit = bit_offset + row;
+        const bool holds_value = validity == nullptr || arrow::bit_util::GetBit(validity, bit);
+        return static_cast<char>(holds_value && (arrow::bit_util::GetBit(values, bit) != negate));
+    };
+
+    char* out = is_valid.data();
+    int64_t row = 0;
+    // The rows sharing the leading partial byte, which the byte at a time body cannot take whole.
+    // None when the bit offset is already on a byte boundary, as it is for the batch a kernel has
+    // just written.
+    const int64_t head = std::min<int64_t>(length, (8 - (bit_offset & 7)) & 7);
+    for (; row < head; row++) {
+        out[row] = unpack_row(row);
+    }
+    for (; row + 8 <= length; row += 8) {
+        const int64_t byte = (bit_offset + row) >> 3;
+        // Every row holds a value when there is no validity bitmap, so all eight bits pass.
+        const uint8_t valid_byte = validity == nullptr ? uint8_t{0xFF} : validity[byte];
+        const uint8_t value_byte = values[byte];
+        const uint8_t bits = negate ? static_cast<uint8_t>(~value_byte & valid_byte)
+                                    : static_cast<uint8_t>(value_byte & valid_byte);
+        std::memcpy(out + row, kExpandedBytes[bits].bytes, sizeof(ExpandedByte));
+    }
+    for (; row < length; row++) {
+        out[row] = unpack_row(row);
+    }
+    return is_valid;
+}
+
 bool ArrowUtils::EqualsIgnoreNullable(const std::shared_ptr<arrow::DataType>& type,
                                       const std::shared_ptr<arrow::DataType>& other_type) {
     if (type->id() != other_type->id() || type->num_fields() != other_type->num_fields()) {

src/paimon/common/utils/arrow/arrow_utils.h

@@ -53,6 +53,22 @@ class PAIMON_EXPORT ArrowUtils {
 
     static uint64_t GetArrayMemoryUsage(const std::shared_ptr<arrow::ArrayData>& data);
 
+    /// Expands the booleans of `array` into one byte per row, the shape `LeafFunction::Test`
+    /// returns. Row `i` becomes 1 when it holds a value and that value, once `negate` has been
+    /// applied to it, is true; a null row becomes 0.
+    ///
+    /// A kernel a leaf function probes with writes a boolean bitmap, one bit per row, so the rows
+    /// it selects have to be spread over a byte each before the caller can use them. Reading that
+    /// bitmap a byte at a time and writing eight rows at once keeps the spread off the per row
+    /// path, where reading it a bit at a time costs a shift, a mask and a byte store per row.
+    ///
+    /// @param array The bitmap to expand. Its offset and its validity are honoured exactly as
+    ///              `BooleanArray::Value()` and `Array::IsValid()` honour them.
+    /// @param negate Whether to take the rows whose bit is clear rather than the rows whose bit is
+    ///               set, which is what `NOT IN` needs. A null row stays 0 either way.
+    /// @return One byte per row of `array`.
+    static std::vector<char> UnpackBooleansToBytes(const arrow::BooleanArray& array, bool negate);
+
     static Result<std::shared_ptr<arrow::StructArray>> RemoveFieldFromStructArray(
         const std::shared_ptr<arrow::StructArray>& struct_array, const std::string& field_name);