diff --git a/Cargo.lock b/Cargo.lock index 367ba7805..34d0a5f5f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5372,7 +5372,6 @@ dependencies = [ [[package]] name = "lance-encoding" version = "13.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-arith", "arrow-array", @@ -5891,6 +5890,7 @@ dependencies = [ "env_logger", "futures", "lance-core", + "lance-encoding", "lance-io", "lance-namespace", "lance-namespace-impls", diff --git a/Cargo.toml b/Cargo.toml index c64305048..52bedc5ea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,6 @@ [workspace] members = ["rust/lancedb", "nodejs", "python"] +exclude = ["vendor/lance-encoding"] resolver = "2" [workspace.package] @@ -79,6 +80,9 @@ tonic = { version = "0.14", features = ["tls-native-roots", "tls-ring"] } uuid = { version = "1.7.0", features = ["v4", "v7"] } chrono = { version = "0.4", default-features = false, features = ["clock"] } +[patch."https://github.com/lance-format/lance.git"] +lance-encoding = { path = "vendor/lance-encoding" } + [profile.ci] debug = "line-tables-only" inherits = "dev" diff --git a/python/Cargo.toml b/python/Cargo.toml index 17ff3f4fc..ef5a5e0f2 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -20,6 +20,9 @@ arrow = { workspace = true, features = ["pyarrow"] } async-trait.workspace = true bytes.workspace = true lancedb.workspace = true +# Maturin includes direct path dependencies in source distributions. The +# workspace patch also applies to transitive Lance writer dependencies. +lance-encoding = { path = "../vendor/lance-encoding", version = "=13.0.0-beta.12" } datafusion-common.workspace = true lance-core.workspace = true lance-namespace.workspace = true diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 5c56fea38..513c696ae 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -6,6 +6,7 @@ import ctypes import gc import json import os +import random import sys import threading import warnings @@ -412,6 +413,21 @@ def test_create_table_rejects_single_dictionary(mem_db: DBConnection): ) +def test_create_table_with_negative_int64_dictionary_tail(tmp_db: DBConnection): + # A dictionary of negative values can have a full 64-bit width and a + # partial 1024-value chunk. The writer must not choose out-of-line + # bitpacking, which requires at least one bit of savings. + rng = random.Random(42) + values = list(range(-800, 800)) + rows = [{"f0": values[rng.randrange(len(values))]} for _ in range(9216)] + schema = pa.schema([pa.field("f0", pa.int64(), nullable=True)]) + + table = tmp_db.create_table("negative_int64_tail", data=rows, schema=schema) + + assert table.count_rows() == len(rows) + assert table.to_arrow().column("f0").to_pylist() == [row["f0"] for row in rows] + + def test_empty_table(mem_db: DBConnection): schema = pa.schema( [ diff --git a/vendor/lance-encoding/Cargo.toml b/vendor/lance-encoding/Cargo.toml new file mode 100644 index 000000000..3819fa3d2 --- /dev/null +++ b/vendor/lance-encoding/Cargo.toml @@ -0,0 +1,55 @@ +[package] +name = "lance-encoding" +version = "13.0.0-beta.12" +edition = "2024" +license = "Apache-2.0" +repository = "https://github.com/lance-format/lance" +description = "Encoders and decoders for the Lance file format" +rust-version = "1.91.0" + +# A temporary local copy of the Lance encoding crate, pinned to the same tag as +# the rest of the Lance workspace dependencies. This makes the selector fix +# available until it is released upstream. +[workspace] + +[dependencies] +lance-arrow = { version = "=13.0.0-beta.12", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.12" } +lance-core = { version = "=13.0.0-beta.12", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.12" } +lance-bitpacking = { version = "=13.0.0-beta.12", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.12", optional = true } +fsst = { version = "=13.0.0-beta.12", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.12" } +arrow-arith = "58.0.0" +arrow-array = "58.0.0" +arrow-buffer = "58.0.0" +arrow-cast = "58.0.0" +arrow-data = "58.0.0" +arrow-schema = "58.0.0" +arrow-select = "58.0.0" +bytes = "1.11.1" +futures = "0.3" +hex = "0.4.3" +itertools = "0.14" +log = "0.4" +num-traits = "0.2" +prost = "0.14.1" +hyperloglogplus = { version = "0.4.1", features = ["const-loop"] } +tokio = { version = "1.23", features = ["rt-multi-thread", "macros", "fs", "sync"] } +tracing = "0.1" +xxhash-rust = { version = "0.8.15", features = ["xxh3"] } +bytemuck = { version = "1", default-features = false, features = ["extern_crate_alloc"] } +byteorder = "1.5" +lz4 = { version = "1", optional = true } +zstd = { version = "0.13", optional = true } + +[build-dependencies] +prost-build = "0.14.1" +protobuf-src = { version = "2.1", optional = true } + +[features] +default = ["lz4", "zstd", "bitpacking"] +protoc = ["dep:protobuf-src"] +bitpacking = ["dep:lance-bitpacking"] +lz4 = ["dep:lz4"] +zstd = ["dep:zstd"] + +[package.metadata.docs.rs] +features = ["protoc"] diff --git a/vendor/lance-encoding/LICENSE b/vendor/lance-encoding/LICENSE new file mode 100644 index 000000000..cd4c38213 --- /dev/null +++ b/vendor/lance-encoding/LICENSE @@ -0,0 +1,255 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +-------------------------------------------------------------------------------- + +This project includes code from Ritchie Vink's Polars project, which is licensed +under the MIT license: + + Copyright (c) 2020 Ritchie Vink + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. + +https://github.com/pola-rs/polars/blob/main/LICENSE + +-------------------------------------------------------------------------------- + +This project includes code adapted from the quickwit-oss/bitpacking crate, which +is licensed under the MIT license: + + Copyright (c) 2016 Paul Masurel + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. + +https://github.com/quickwit-oss/bitpacking/blob/main/LICENSE diff --git a/vendor/lance-encoding/PATCH.md b/vendor/lance-encoding/PATCH.md new file mode 100644 index 000000000..44a0ca0e0 --- /dev/null +++ b/vendor/lance-encoding/PATCH.md @@ -0,0 +1,11 @@ +This is `rust/lance-encoding` from +[`lance-format/lance` v13.0.0-beta.12](https://github.com/lance-format/lance/tree/5d2946fa452c61cb018144f3200892b6dd4a8b4f/rust/lance-encoding), +with the crate's build protos copied into `protos/`. Its manifest uses explicit +dependencies so it can be built outside the Lance workspace. + +The only source change is in `src/compression.rs`: skip out-of-line bitpacking +when the selected width is the full width of the values. That encoder requires +positive bit savings for a partial 1024-value chunk and panics in debug builds +for signed dictionaries containing negative integers. + +Remove this patch and the root Cargo patch entry once Lance releases the guard. diff --git a/vendor/lance-encoding/README.md b/vendor/lance-encoding/README.md new file mode 100644 index 000000000..3e6f7c554 --- /dev/null +++ b/vendor/lance-encoding/README.md @@ -0,0 +1,6 @@ +# lance-encoding + +`lance-encoding` is an internal sub-crate, containing encoders and decoders +for the Lance file format. + +**Important Note**: This crate is **not intended for external usage**. diff --git a/vendor/lance-encoding/build.rs b/vendor/lance-encoding/build.rs new file mode 100644 index 000000000..92fe03589 --- /dev/null +++ b/vendor/lance-encoding/build.rs @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::io::Result; + +fn main() -> Result<()> { + println!("cargo:rerun-if-changed=protos"); + + #[cfg(feature = "protoc")] + // Use vendored protobuf compiler if requested. + unsafe { + std::env::set_var("PROTOC", protobuf_src::protoc()); + } + + let mut prost_build = prost_build::Config::new(); + prost_build.protoc_arg("--experimental_allow_proto3_optional"); + prost_build.enable_type_names(); + prost_build.bytes(["."]); // Enable Bytes type for all messages to avoid Vec clones. + prost_build.compile_protos(&["./protos/encodings_v2_0.proto"], &["./protos"])?; + prost_build.compile_protos(&["./protos/encodings_v2_1.proto"], &["./protos"])?; + + Ok(()) +} diff --git a/vendor/lance-encoding/protos/encodings_v2_0.proto b/vendor/lance-encoding/protos/encodings_v2_0.proto new file mode 100644 index 000000000..45776fb34 --- /dev/null +++ b/vendor/lance-encoding/protos/encodings_v2_0.proto @@ -0,0 +1,358 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +syntax = "proto3"; + +package lance.encodings; + +import "google/protobuf/empty.proto"; + +/* This file contains a specification for encodings that can be used + * to store and load Arrow data into a Lance file for the 2.0 format. It + * has been superseded by encodings21.proto which is used for the 2.1 format. + * + * # Types + * + * This file assumes the user wants to load data into Arrow arrays and + * explains how to map Arrow arrays into Lance files. Encodings are divided + * into "array encoding" (which maps to an Arrow array and may contain multiple + * buffers) and "buffer encoding" (which encodes a single buffer of data). + * + * # Encoding Tree + * + * Most encodings are layered on top of each other. These form a tree of + * encodings with a single root node. To encode an array you will typically + * start with the root node and then take the output from that root encoding + * and feed it into child encodings. The decoding process works in reverse. + * + * # Multi-column Encodings + * + * Some Arrow arrays will map to more than one column of Lance data. For + * example, struct arrays and list arrays. This file only contains encodings + * for a single column. However, it does describe how multi-column arrays can + * be encoded. + */ + +/* A pointer to a buffer in a Lance file + * + * A writer can place a buffer in three different locations. The buffer + * can go in the data page, in the column metadata, or in the file metadata. + * The writer is free to choose whatever is most appropriate (for example, a dictionary + * that is shared across all pages in a column will probably go in the column + * metadata). This specification does not dictate where the buffer should go. + */ +message Buffer { + // The index of the buffer in the collection of buffers + uint32 buffer_index = 1; + // The collection holding the buffer + enum BufferType { + // The buffer is stored in the data page itself + page = 0; + // The buffer is stored in the column metadata + column = 1; + // The buffer is stored in the file metadata + file = 2; + }; + BufferType buffer_type = 2; +} + +/* An encoding that adds nullability to another array encoding + * + * This can wrap any array encoding and add nullability information + */ +message Nullable { + message NoNull { + ArrayEncoding values = 1; + } + message AllNull {} + message SomeNull { + ArrayEncoding validity = 1; + ArrayEncoding values = 2; + } + oneof nullability { + // The array has no nulls and there is a single buffer needed + NoNull no_nulls = 1; + // The array may have nulls and we need two buffers + SomeNull some_nulls = 2; + // All values are null (no buffers needed) + AllNull all_nulls = 3; + } +} + +// An array encoding for variable-length list fields +message List { + /* An array containing the offsets into an items array. + * + * This array will have num_rows items and will never + * have nulls. + * + * If the list at index i is not null then offsets[i] will + * contain `base + len(list)` where `base` is defined as: + * i == 0: 0 + * i > 0: (offsets[i-1] % null_offset_adjustment) + * + * To help understand we can consider the following example list: + * [ [A, B], null, [], [C, D, E] ] + * + * The offsets will be [2, ?, 2, 5] + * + * If the incoming list at index i IS null then offsets[i] will + * contain `base + len(list) + null_offset_adjustment` where `base` + * is defined the same as above. + * + * To complete the above example let's assume that `null_offset_adjustment` + * is 7. Then the offsets will be [2, 9, 2, 5] + * + * If there are no nulls then the offsets we write here are exactly the + * same as the offsets in an Arrow list array (except we omit the leading + * 0 which is redundant) + * + * The reason we do this is so that reading a single list at index i only + * requires us to load the indices at i and i-1. + * + * If the offset at index i is greater than `null_offset_adjustment`` + * then the list at index i is null. + * + * Otherwise the length of the list is `offsets[i] - base` where + * base is defined the same as above. + * + * Let's consider our example offsets: [2, 9, 2, 5] + * + * We can take any range of lists and determine how many list items are + * referenced by the sublist. + * + * 0..3: [_, 5] -> items 0..5 (base = 0* and end is 5) + * 0..2: [_, 2] -> items 0..2 (base = 0* and end is 2) + * 0..1: [_, 9] -> items 0..2 (base = 0* and end is 9 % 7) + * 1..3: [2, 5] -> items 2..5 (base = 2 and end is 5) + * 1..2: [2, 2] -> items 2..2 (base = 2 and end is 2) + * 2..3: [9, 5] -> items 2..5 (base = 9 % 7 and end is 5) + * + * * When the start of our range is the 0th item the base is always 0 and we only + * need to load a single index from disk to determine the range. + * + * The data type of the offsets array is flexible and does not need + * to match the data type of the destination array. Please note that the offsets + * array is very likely to be efficiently encoded by bit packing deltas. + */ + ArrayEncoding offsets = 1; + /* If a list is null then we add this value to the offset + * + * This value must be greater than the length of the items so that + * (offset + null_offset_adjustment) is never used by a non-null list. + * + * Note that this value cannot be equal to the length of the items + * because then a page with a single list would store [ X ] and we + * couldn't know if that is a null list or a list with X items. + * + * Therefore, the best choice for this value is 1 + # of items. + * Choosing this will maximize the bit packing that we can apply to the offsets. + */ + uint64 null_offset_adjustment = 2; + /* How many items are referenced by these offsets. This is needed in + * order to determine which items pages map to this offsets page. + */ + uint64 num_items = 3; +} + +// An array encoding for fixed-size list fields +message FixedSizeList { + /// The number of items in each list + uint32 dimension = 1; + /// True if the list is nullable + bool has_validity = 3; + /// The items in the list + ArrayEncoding items = 2; +} + +message Compression { + string scheme = 1; + optional int32 level = 2; +} + +// Fixed width items placed contiguously in a buffer +message Flat { + /* the number of bits per value, must be greater than 0, does + * not need to be a multiple of 8 + */ + uint64 bits_per_value = 1; + // the buffer of values + Buffer buffer = 2; + /* The Compression message can specify the compression scheme (e.g. zstd) and any + * other information that is needed for decompression. + * + * If this array is compressed then the bits_per_value refers to the uncompressed + * data. + */ + Compression compression = 3; +} + +// Compression algorithm where all values have a constant value +message Constant { + // The value (TODO: define encoding for literals?) + bytes value = 1; +} + +// Items are bitpacked in a buffer +message Bitpacked { + // the number of bits used for a value in the buffer + uint64 compressed_bits_per_value = 1; + + // the number of bits of the uncompressed value. e.g. for a u32, this will be 32 + uint64 uncompressed_bits_per_value = 2; + + // The items in the list + Buffer buffer = 3; + + // Whether or not a sign bit is included in the bitpacked value + bool signed = 4; +} + +// Items are bitpacked in a buffer +message BitpackedForNonNeg { + // the number of bits used for a value in the buffer + uint64 compressed_bits_per_value = 1; + + // the number of bits of the uncompressed value. e.g. for a u32, this will be 32 + uint64 uncompressed_bits_per_value = 2; + + // The items in the list + Buffer buffer = 3; +} + +// Opaque bitpacking variant where the bits per value are stored inline in the chunks themselves +message InlineBitpacking { + // the number of bits of the uncompressed value. e.g. for a u32, this will be 32 + uint64 uncompressed_bits_per_value = 2; +} + +// Transparent bitpacking variant where the number of bits per value is fixed through the whole buffer +message OutOfLineBitpacking { + // the number of bits of the uncompressed value. e.g. for a u32, this will be 32 + uint64 uncompressed_bits_per_value = 2; + // The number of compressed bits per value, fixed across the entire buffer + uint64 compressed_bits_per_value = 3; +} + +/* An array encoding for shredded structs that will never be null + * + * There is no actual data in this column. + * + * TODO: Struct validity bitmaps will be placed here. + */ +message SimpleStruct {} + +// An array encoding for binary fields +message Binary { + ArrayEncoding indices = 1; + ArrayEncoding bytes = 2; + uint64 null_adjustment = 3; +} + +message Variable { + uint32 bits_per_offset = 1; +} + +message Fsst { + ArrayEncoding binary = 1; + bytes symbol_table = 2; +} + +// An array encoding for dictionary-encoded fields +message Dictionary { + ArrayEncoding indices = 1; + ArrayEncoding items = 2; + uint32 num_dictionary_items = 3; +} + +message PackedStruct { + repeated ArrayEncoding inner = 1; + Buffer buffer = 2; +} + +message PackedStructFixedWidthMiniBlock { + ArrayEncoding Flat = 1; + repeated uint32 bits_per_values = 2; +} + +message FixedSizeBinary { + ArrayEncoding bytes = 1; + uint32 byte_width = 2; +} + +message Block { + string scheme = 1; +} + +// Run-Length Encoding for miniblock format +message Rle { + // Number of bits per value (8, 16, 32, 64, or 128) + uint64 bits_per_value = 1; +} + +// Byte Stream Split encoding for floating point values +message ByteStreamSplit { + // Number of bits per value (32 for float, 64 for double) + uint64 bits_per_value = 1; +} + +// General miniblock encoding - wraps another miniblock encoding with compression +message GeneralMiniBlock { + // The inner miniblock encoding (e.g., Rle, Bitpacked, etc.) + ArrayEncoding inner = 1; + // The compression scheme to apply to the miniblock buffers + Compression compression = 2; +} + +// Encodings that decode into an Arrow array +message ArrayEncoding { + oneof array_encoding { + Flat flat = 1; + Nullable nullable = 2; + FixedSizeList fixed_size_list = 3; + List list = 4; + SimpleStruct struct = 5; + Binary binary = 6; + Dictionary dictionary = 7; + Fsst fsst = 8; + PackedStruct packed_struct = 9; + Bitpacked bitpacked = 10; + FixedSizeBinary fixed_size_binary = 11; + BitpackedForNonNeg bitpacked_for_non_neg = 12; + Constant constant = 13; + InlineBitpacking inline_bitpacking = 14; + OutOfLineBitpacking out_of_line_bitpacking = 15; + Variable variable = 16; + PackedStructFixedWidthMiniBlock packed_struct_fixed_width_mini_block = 17; + Block block = 18; + Rle rle = 19; + GeneralMiniBlock general_mini_block = 20; + ByteStreamSplit byte_stream_split = 21; + } +} + +/* Wraps a column with a zone map index that can be used + * to apply pushdown filters + */ +message ZoneIndex { + uint32 rows_per_zone = 1; + Buffer zone_map_buffer = 2; + ColumnEncoding inner = 3; +} + +/* Marks a column as blob data. It will contain a packed struct + * with fields position and size (u64) + */ +message Blob { + ColumnEncoding inner = 1; +} + +// Encodings that describe a column of values +message ColumnEncoding { + oneof column_encoding { + // No special encoding, just column values + google.protobuf.Empty values = 1; + ZoneIndex zone_index = 2; + Blob blob = 3; + } +} diff --git a/vendor/lance-encoding/protos/encodings_v2_1.proto b/vendor/lance-encoding/protos/encodings_v2_1.proto new file mode 100644 index 000000000..643aa1f9d --- /dev/null +++ b/vendor/lance-encoding/protos/encodings_v2_1.proto @@ -0,0 +1,675 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +syntax = "proto3"; + +package lance.encodings21; + +/* This file contains a specification for encodings that can be used + * to store and load Arrow data into a Lance file for the 2.1 format. + * + * # Types + * + * This file assumes the user wants to load data into Arrow arrays and + * explains how to map Arrow arrays into Lance files. Encodings are divided + * into "structural encodings" (which are used to encode the structure of the + * data such as any list or struct layers) and "compressive encodings" (which + * are used to compress the actual data values). + * + * # Standardized Interpretation of Counting Terms + * + * When working with 2.1 encodings we have a number of different "counting terms" and it can be + * difficult to understand what we mean when we are talking about a "number of values". Here is + * a standard interpretation of these terms: + * + * To understand these definitions consider a data type FIXED_SIZE_LIST>. + * + * A "value" is an abstract term when we aren't being specific. + * + * - num_rows: This is the highest level counting term. A single row includes everything in the + * fixed size list. This is what the user asks for when they asks for a range of rows. + * - num_elements: The number of elements is the number of rows multiplied by the dimension of any + * fixed size list wrappers. This is what you get when you flatten the FSL layer and + * is the starting point for structural encoding. Note that an element can be a list + * value or a single primitive value. + * - num_items: The number of items is the number of values in the repetition and definition vectors + * after everything has been flattened. + * - num_visible_items: The number of visible items is the number of items after invisible items + * have been removed. Invisible items are rep/def levels that don't correspond to an + * actual value. + */ + + +/* # Structural Encodings + * + * The following message are used to describe the structural encoding of the + * data. In this document, we refer to these structural encodings as layouts. + */ + +/* Repetition and definition levels are described in more detail elsewhere. As we peel through + * the structure of an array we will encounter layers of struct and list. Each of these layers + * potentially adds a new level to the repetition and definition levels. This message describes + * the meaning of each layer. + */ +enum RepDefLayer { + // Should never be used, included for debugging purporses and general protobuf best practice + REPDEF_UNSPECIFIED = 0; + // All values are valid (can be primitive or struct) + REPDEF_ALL_VALID_ITEM = 1; + // All list values are valid + REPDEF_ALL_VALID_LIST = 2; + // There are one or more null items (can be primitive or struct) + REPDEF_NULLABLE_ITEM = 3; + // A list layer with null lists but no empty lists + REPDEF_NULLABLE_LIST = 4; + // A list layer with empty lists but no null lists + REPDEF_EMPTYABLE_LIST = 5; + // A list layer with both empty lists and null lists + REPDEF_NULL_AND_EMPTY_LIST = 6; +} + +/* A layout used for pages where the data is small + * + * In this case we can fit many values into a single disk sector and transposing buffers is + * expensive. As a result, we do not transpose the buffers but compress the data into small + * chunks (called mini blocks) which are roughly the size of a disk sector. + * + * The end result is a small amount of read amplification (since we must read an entire page + * at a time) but we have more flexibility in compression and do less work per value when + * compressing and decompressing in bulk. + */ +message MiniBlockLayout { + /* Description of the compression of repetition levels (e.g. how many bits per rep) + * + * Optional, if there is no repetition then this field is not present + */ + CompressiveEncoding rep_compression = 1; + /* Description of the compression of definition levels (e.g. how many bits per def) + * + * Optional, if there is no definition then this field is not present + */ + CompressiveEncoding def_compression = 2; + // Description of the compression of values + CompressiveEncoding value_compression = 3; + /* Description of the compression of the dictionary data + * + * Optional, if there is no dictionary then this field is not present + */ + CompressiveEncoding dictionary = 4; + // Number of items in the dictionary + uint64 num_dictionary_items = 5; + // The meaning of each repdef layer, used to interpret repdef buffers correctly + repeated RepDefLayer layers = 6; + /* The number of buffers in each mini-block, this is determined by the compression and does + * NOT include the repetition or definition buffers (the presence of these buffers can be determined + * by looking at the rep_compression and def_compression fields) + */ + uint64 num_buffers = 7; + /* The depth of the repetition index. + * + * If there is repetition then the depth must be at least 1. If there are many layers + * of repetition then deeper repetition indices will support deeper nested random access. For + * example, given 5 layers of repetition then the repetition index depth must be at least + * 3 to support access like `rows[50][17][3]`. + * + * We require `repetition_index_depth + 1` u64 values per mini-block to store the repetition + * index if the `repetition_index_depth` is greater than 0. The +1 is because we need to store + * the number of "leftover items" at the end of the chunk. Otherwise, we wouldn't have any way + * to know if the final item in a chunk is valid or not. + */ + uint32 repetition_index_depth = 8; + /* The page already records how many rows are in the page. For mini-block we also need to know how + * many "items" are in the page. A row and an item are the same thing unless the page has lists. + */ + uint64 num_items = 9; + + // Since Lance 2.2, miniblocks have larger chunk sizes (>= 64KB) + bool has_large_chunk = 10; +} + +/* A layout used for pages where the data is large + * + * In this case the cost of transposing the data is relatively small (compared to the cost of writing the data) + * and so we just zip the buffers together + */ +message FullZipLayout { + // The number of bits of repetition info (0 if there is no repetition) + uint32 bits_rep = 1; + // The number of bits of definition info (0 if there is no definition) + uint32 bits_def = 2; + /* The number of bits of value info + * + * Note: we use bits here (and not bytes) for consistency with other encodings. However, in practice, + * there is never a reason to use a bits per value that is not a multiple of 8. The complexity is not + * worth the small savings in space since this encoding is typically used with large values already. + */ + oneof details { + // If this is a fixed width block then we need to have a fixed number of bits per value + uint32 bits_per_value = 3; + // If this is a variable width block then we need to have a fixed number of bits per offset + uint32 bits_per_offset = 4; + } + // The number of items in the page + uint32 num_items = 5; + // The number of visible items in the page + uint32 num_visible_items = 6; + // Description of the compression of values + CompressiveEncoding value_compression = 7; + // The meaning of each repdef layer, used to interpret repdef buffers correctly + repeated RepDefLayer layers = 8; +} + +/* A layout used for sparse flat or nested pages where Arrow structure is represented directly + * in layer-local slot domains instead of as dense repetition / definition events. + * + * Structural layers are ordered from outer-most to inner-most. Values remain mini-block + * compressed and are split into independently readable chunks. + */ +message SparseLayout { + // Description of the compression of values. + CompressiveEncoding value_compression = 1; + // Number of value buffers in each mini-block chunk. This does not include structural buffers. + uint64 num_buffers = 2; + /* Number of entries in the equivalent dense repetition / definition stream. This equals + * num_visible_items plus one structural placeholder for every list slot without children. + * Null leaf slots count as visible items because they still occupy positions in Arrow's + * leaf value buffer. For example, a nullable primitive with 100 slots, 30 of them null, + * has num_items = num_visible_items = 100. + */ + uint64 num_items = 3; + // Number of leaf value slots encoded in the value chunks, including null leaf slots. + uint64 num_visible_items = 4; + // If true, chunk-local value buffer sizes use u32. Otherwise they use u16. + bool has_large_chunk = 5; + /* Structural layers ordered from outer-most to inner-most. This may be empty for a flat, + * non-nullable leaf page whose scheduling domain equals num_visible_items. + */ + repeated SparseStructuralLayer structural_layers = 6; +} + +/* A domain is a layer-local integer coordinate space [0, num_slots). A slot is one + * element in that space. The outer-most domain is the page's top-level rows; each + * layer's child domain is the next layer's parent domain, and the terminal child + * domain contains num_visible_items leaf value slots. + */ +message SparseStructuralLayer { + // Exactly one layer kind is required. + oneof layer { + SparseValidityLayer validity = 1; + SparseListLayer list = 2; + SparseFixedSizeListLayer fixed_size_list = 3; + } +} + +message SparseValidityLayer { + // Number of nullable item or struct slots in this layer's parent and child domain. + uint64 num_slots = 1; + // Validity for the slots in this layer. + SparseValiditySet validity = 2; +} + +message SparseListLayer { + // Number of list, large-list, or map slots in this layer's parent domain. + uint64 num_slots = 1; + // Number of slots in this layer's child domain. + uint64 num_child_slots = 2; + // Non-empty parent slots. Valid parent slots absent from this set are empty lists. + SparsePositionSet non_empty_positions = 3; + // Positive child counts corresponding one-for-one with non_empty_positions. + SparseCountSet counts = 4; + // Validity for the parent slots in this layer. + SparseValiditySet validity = 5; +} + +message SparseFixedSizeListLayer { + // Number of fixed-size-list slots in this layer's parent domain. + uint64 num_slots = 1; + // Number of children per parent slot. The child domain has num_slots * dimension slots. + uint64 dimension = 2; + // Validity for the parent slots in this layer. + SparseValiditySet validity = 3; +} + +message SparseValiditySet { + enum Meaning { + SPARSE_VALIDITY_UNSPECIFIED = 0; + // Stored positions are null; all other positions are valid. + SPARSE_VALIDITY_NULL_POSITIONS = 1; + // Stored positions are valid; all other positions are null. + SPARSE_VALIDITY_VALID_POSITIONS = 2; + } + + Meaning meaning = 1; + SparsePositionSet positions = 2; +} + +message SparsePositionEmpty {} + +message SparsePositionAll {} + +message SparsePositionRange { + uint64 start = 1; + uint64 length = 2; +} + +message SparsePositionSet { + oneof positions { + // Delta-compressed u64 positions. Cardinality is num_positions. + CompressiveEncoding explicit = 1; + // One contiguous, non-empty range. + SparsePositionRange range = 2; + // Every position in the domain. + SparsePositionAll all = 3; + // No positions in the domain. + SparsePositionEmpty empty = 4; + } + // Semantic cardinality of this set. + uint64 num_positions = 5; +} + +message SparseCountEmpty {} + +message SparseCountConstant { + // Child count shared by every non-empty list slot. + uint64 value = 1; +} + +message SparseCountSet { + oneof counts { + // Compressed u64 child counts. Cardinality comes from the containing position set. + CompressiveEncoding explicit = 1; + // One positive child count shared by every non-empty list slot. + SparseCountConstant constant = 2; + // No counts; valid only when there are no non-empty list slots. + SparseCountEmpty empty = 3; + } +} + +/* A layout used for pages where all (visible) values are the same scalar value. + * + * This generalizes the prior AllNullLayout semantics for file_version >= 2.2. + * + * There may be buffers of repetition and definition information if required in order + * to interpret what kind of nulls are present / which items are visible. + */ +message ConstantLayout { + // The meaning of each repdef layer, used to interpret repdef buffers correctly + repeated RepDefLayer layers = 5; + + /* Inline fixed-width scalar value bytes. + * + * This MUST only be used for types where a single non-null element is represented by a single + * fixed-width Arrow value buffer (i.e. no offsets buffer, no child data). + * + * Constraints: + * - MUST be absent for an all-null page + * - MUST be <= 32 bytes if present + */ + optional bytes inline_value = 6; + + /* Optional compression algorithm used for the repetition buffer. + * If absent, repetition levels are stored as raw u16 values. + */ + CompressiveEncoding rep_compression = 7; + /* Optional compression algorithm used for the definition buffer. + * If absent, definition levels are stored as raw u16 values. + */ + CompressiveEncoding def_compression = 8; + // Number of values in repetition buffer after decompression. + uint64 num_rep_values = 9; + // Number of values in definition buffer after decompression. + uint64 num_def_values = 10; +} + +/* A layout where large binary data is encoded externally and only + * the descriptions (position + size) are placed in the page + * + * Repdef information is stored in the descriptions. A description with a size of + * 0 and a position of 0 is an empty value. A description with a size of 0 and a + * non-zero position is a null value and the position is the repdef value. + */ +message BlobLayout { + // The inner layout used to store the descriptions + PageLayout inner_layout = 1; + /* The meaning of each repdef layer, used to interpret repdef buffers correctly + * + * The inner layout's repdef layers will always be 1 all valid item layer + */ + repeated RepDefLayer layers = 2; +} + +// Describes the structural encoding of a page +message PageLayout { + oneof layout { + // A layout used for pages where the data is small + MiniBlockLayout mini_block_layout = 1; + // A layout used for pages where all (visible) values are the same scalar value or null. + ConstantLayout constant_layout = 2; + // A layout used for pages where the data is large + FullZipLayout full_zip_layout = 3; + /* A layout where large binary data is encoded externally + * and only the descriptions are put in the page + */ + BlobLayout blob_layout = 4; + // A sparse structural layout. This variant requires file version 2.3 or later. + SparseLayout sparse_layout = 5; + } +} + +/* # Compressive Encodings + * + * These encodings describe how an array is compressed. An encoding may split an + * array into multiple buffers. The buffers can then be compressed further (and split + * into yet more buffers). The entire process forms a tree of encodings with the root + * of the tree being the initial array and the leaves being the final compressed buffers. + * + * # Data blocks and buffers + * + * Data blocks are a simplified version of arrays and represent a collection of buffers grouped + * with some kind of interpretation. Data blocks are the input and output of compressive encodings. + * There are different kinds of data blocks: + * - Fixed width data blocks (e.g. u8, u16, ...) + * - Variable width data blocks (e.g. strings, binary) + * - Struct data blocks (note: this is for packed structs, normal structs are encoded in the structural encoding) + * + * In addition, leaf encodings may output "buffers". These are fully compressed buffers of data that + * are stored in the page and no longer compressed. + */ + +enum CompressionScheme { + COMPRESSION_ALGORITHM_UNSPECIFIED = 0; + COMPRESSION_ALGORITHM_LZ4 = 1; + COMPRESSION_ALGORITHM_ZSTD = 2; +} + +/* Compression applied to a single buffer of data + * + * A buffer is the leaf of the compression tree. Unlike data blocks, which can + * be further compressed with a variety of techniques, a buffer cannot be understood + * in any particular way. + * + * A general compression scheme may be applied to a buffer. This is something like + * zstd, lz4, etc. The entire buffer is compressed as a single unit. If this happens + * then any parent encoding becomes opaque, even if it would normally be transparent. + * + * This is a leaf, no further compression is applied to the data. + */ +message BufferCompression { + // A general compression scheme to apply to the buffer + CompressionScheme scheme = 1; + /* The compression level + * + * Optional, if not present a scheme-specific default value will be used. + * + * Interpretation of this value depends on the compression scheme. Generally, larger + * values indicate more compression at the expense of more CPU time. + */ + optional int32 level = 2; +} + +/* Fixed width items placed contiguously in a single buffer + * + * This is a leaf encoding, there is no compression applied to the data. + * + * This is a transparent encoding by definition. + * + * The input is a fixed-width data block. + * The output is a single buffer. + */ +message Flat { + /* the number of bits per value, must be greater than 0, does + * not need to be a multiple of 8 + */ + uint64 bits_per_value = 1; + // The compression applied to the data + optional BufferCompression data = 2; +} + +/* Variable width items have the values stored in one buffer and the + * offsets are output as a data block that may be further compressed. + * + * This is a partial leaf encoding. Values are not compressed but + * the offsets may be further compressed. + * + * This is a transparent encoding by definition. + * + * The input is a variable-width data block. + * The output is a single fixed-width data block (the offsets) and + * a single buffer (the values) + */ +message Variable { + // Describes how the offsets data block is compressed + CompressiveEncoding offsets = 1; + // The compression applied to the values + optional BufferCompression values = 2; +} + +/* Compression algorithm where all values have a constant value (encoded in the description) + * + * This is a leaf encoding, there is no compression applied to the data. + * + * The input can be any kind of data block. + * There is no output. + */ +message Constant { + // The value (TODO: define encoding for literals?) + optional bytes value = 1; +} + +/* A compression scheme in which a single fixed-width block is "packed" into + * a smaller fixed-width block values where each value has fewer bits. + * + * This is typically done by throwing away the most significant bits of each value when + * those bits are all the same. + * + * In this scheme the number of bits per value is fixed across the entire buffer and stored + * in this message. + * + * This is a transparent encoding. + * + * The input is a fixed-width data block. + * The output is a single fixed-width data block. + */ +message OutOfLineBitpacking { + // the number of bits of the uncompressed value. e.g. for a u32, this will be 32 + uint64 uncompressed_bits_per_value = 1; + // The compression used to store the bitpacked values data block + CompressiveEncoding values = 3; +} + +/* Bitpacking variant where the bits per value are stored inline in the chunks themselves + * + * This variation of bitpacking allows for the number of bits per value to change throughout the + * buffer, which makes the compression more robust to outliers. + * + * This is an opaque encoding. + * + * The input is a fixed-width data block. + * The output is a single buffer. + */ +message InlineBitpacking { + // the number of bits of the uncompressed value. e.g. for a u32, this will be 32 + uint64 uncompressed_bits_per_value = 1; + // The compression applied to the values + optional BufferCompression values = 2; +} + +/* A compression scheme for variable-width data + * + * A small dictionary (referred to as a "symbol table") is used to compress the values. + * In this scheme there is a single symbol table for the entire page and it is stored in the + * encoding description itself. + * + * This is a transparent encoding. + * + * The input is a variable-width data block. + * The output is a single variable-width data block. + */ +message Fsst { + // The FSST symbol table + bytes symbol_table = 1; + // The compression used to store the compressed values data block + CompressiveEncoding values = 2; +} + +/* A compression scheme where common values are stored in a dictionary and the values are + * encoded as indices into the dictionary. + * + * This is an opaque encoding unless the dictionary is considered metadata. + * + * The input is a any kind of data block. + * There are two outputs: + * - A data block of the same kind as the input (the dictionary) + * - A fixed-width data block containing the indices into the dictionary. + */ +message Dictionary { + // The compression used to store the indices data block + CompressiveEncoding indices = 1; + // The compression used to store the dictionary items data block + CompressiveEncoding items = 2; + // The number of items in the dictionary + uint32 num_dictionary_items = 3; +} + +/* A compression scheme where runs of common values are encoded as a single value and a count + * + * This is an opaque encoding unless the run lengths are considered metadata. + * + * The input is a single data block of any kind. + * There are two outputs: + * - A data block of the same kind as the input (the run values) + * - A fixed-width data block containing the lengths of the runs + */ +message Rle { + // The compression used to store the run values data block + CompressiveEncoding values = 1; + // The compression used to store the run lengths data block + CompressiveEncoding run_lengths = 2; +} + +/* Converts a fixed-size-list of values into a flattened list of values + * + * This encoding does not actually compress the data, it just flattens out the FSL layers. + * + * This is a transparent encoding. + * + * The input is a single block of fixed-width data (with a wide width and few items) + * The output is a single block of fixed-width data (with a narrow width and many items) + */ +message FixedSizeList { + // The number of items in this layer of FSL + uint64 items_per_value = 1; + // Whether or not there is a validity buffer + bool has_validity = 3; + // The compression used to store the flattened values data block + CompressiveEncoding values = 2; +} + +/* Packs a struct containing only fixed-width children into a single fixed-width data block + * + * The children are concatenated row by row and stored as a single fixed-width buffer. This is + * the legacy packed struct representation and remains available for backwards compatibility. + */ +message PackedStruct { + // The number of bits contributed by each child field in the packed row + repeated uint64 bits_per_value = 1; + // The compression used to store the packed fixed-width values + CompressiveEncoding values = 2; +} + +/* Variable-width packed struct encoding (2.2 extension) + * + * Each child value is compressed independently before being transposed into + * a row-major layout. This preserves per-field compression boundaries at the + * cost of disabling mini-block compression. Readers must prefer this field + * when present and fall back to the legacy encoding otherwise. + */ +message VariablePackedStruct { + // Per-field encoding metadata in struct order + repeated FieldEncoding fields = 1; + + // Encoding description for a single child field + message FieldEncoding { + // Compression applied to individual field values before transposition + CompressiveEncoding value = 1; + oneof layout { + // Bit width of each compressed value (when fixed width) + uint64 bits_per_value = 2; + // Bit width of the length prefix for variable-width compressed values + uint64 bits_per_length = 3; + } + } +} + +/* A compression scheme that wraps the underlying data with general compression + * + * Note: The application of wrapped compression will depend on the layout of the data. + * If we apply it to mini-block data then we compress entire mini-blocks. If we apply + * it to full-zip data then we compress each value individually. + * + * Note: Wrapped compression is somewhat unique at the moment as it is applied to the + * output of the inner encoding and not the input like all other compressive encodings. + * + * Note: General compression can usually be applied in two spots. We can apply + * it to individual buffers or we can apply it here, to the entire array. + * + * For example, let's say we are storing mini-blocks of strings and we are using + * FSST and bitpacking the offsets. We have something like this... + * + * WRAPPED(†3) -> FSST -> VARIABLE -(offsets)-> INLINE_BITPACKING -(data)-> FLAT -> BUFFER (†1) + * -(data)-> BUFFER (†2) + * + * General compression can be applied at †1, †2, or †3 (or any combination of these). + * + * If we apply it at †1 then we apply it just to the bitpacked offsets + * If we apply it at †2 then we apply it just to the FSST compressed data + * If we apply it at †3 then we apply it to the entire mini-block (both offsets and data) + * + * The input is a single data block of any kind. + * The output is a single data block of the same kind as the input. + */ +message General { + // The compression to apply to the values + BufferCompression compression = 1; + // The compression used to store the output data block + CompressiveEncoding values = 3; +} + +/* A compression scheme where fixed-width values are transposed into a series of byte streams + * + * This is commonly used for floating point values where the upper bits (the mantissa) have a + * significantly different meaning than the lower bits. By splitting the values into byte streams + * we group the mantissa bits together and the exponent bits together. The end result is typically + * more compressible. + * + * Note that this encoding is mostly useful when combined with other encodings. It does not do any + * compression on its own. + * + * This is an opaque encoding. + * + * The input is a fixed-width data block + * The output is a single fixed-width data block + */ +message ByteStreamSplit { + // The compression used to store the values + CompressiveEncoding values = 1; +} + +// An encoding that compresses a data block into buffers +message CompressiveEncoding { + oneof compression { + Flat flat = 1; + Variable variable = 2; + Constant constant = 3; + OutOfLineBitpacking out_of_line_bitpacking = 4; + InlineBitpacking inline_bitpacking = 5; + Fsst fsst = 6; + Dictionary dictionary = 7; + Rle rle = 8; + ByteStreamSplit byte_stream_split = 9; + General general = 10; + FixedSizeList fixed_size_list = 11; + PackedStruct packed_struct = 12; + VariablePackedStruct variable_packed_struct = 13; + } +} diff --git a/vendor/lance-encoding/src/array_encoding.rs b/vendor/lance-encoding/src/array_encoding.rs new file mode 100644 index 000000000..a76c6a186 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Encoding and decoding mechanisms described by [`crate::format::pb::ArrayEncoding`]. +//! +//! File versions decide which mechanisms to compose and accept. This module +//! contains only the reusable implementation of that persisted grammar. + +pub mod logical; +pub mod physical; +mod strategy; + +pub use strategy::ArrayFieldEncodingStrategy; diff --git a/vendor/lance-encoding/src/array_encoding/logical.rs b/vendor/lance-encoding/src/array_encoding/logical.rs new file mode 100644 index 000000000..bdc981b2c --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/logical.rs @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +pub mod binary; +pub mod blob; +pub mod list; +pub mod primitive; +pub mod r#struct; diff --git a/vendor/lance-encoding/src/array_encoding/logical/binary.rs b/vendor/lance-encoding/src/array_encoding/logical/binary.rs new file mode 100644 index 000000000..4a4ea45a5 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/logical/binary.rs @@ -0,0 +1,262 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_array::{ + Array, ArrayRef, GenericByteArray, GenericListArray, + cast::AsArray, + types::{BinaryType, ByteArrayType, LargeBinaryType, LargeUtf8Type, UInt8Type, Utf8Type}, +}; + +use arrow_schema::DataType; +use futures::{FutureExt, future::BoxFuture}; +use lance_core::{Error, Result}; +use log::trace; + +use crate::{ + decoder::{ + DecodeArrayTask, FilterExpression, MessageType, NextDecodeTask, PriorityRange, + ScheduledScanLine, SchedulerContext, + }, + decoder::{DecoderReady, DrainLimit, FieldScheduler, LogicalPageDecoder, SchedulingJob}, +}; + +/// Wraps a varbin scheduler and uses a BinaryPageDecoder to cast +/// the result to the appropriate type +#[derive(Debug)] +pub struct BinarySchedulingJob<'a> { + scheduler: &'a BinaryFieldScheduler, + inner: Box, +} + +impl SchedulingJob for BinarySchedulingJob<'_> { + fn schedule_next( + &mut self, + context: &mut SchedulerContext, + priority: &dyn PriorityRange, + ) -> Result { + let inner_scan = self.inner.schedule_next(context, priority)?; + let wrapped_decoders = inner_scan + .decoders + .into_iter() + .map(|message| { + let decoder = message.into_array(); + MessageType::DecoderReady(DecoderReady { + path: decoder.path, + decoder: Box::new(BinaryPageDecoder { + inner: decoder.decoder, + data_type: self.scheduler.data_type.clone(), + }), + }) + }) + .collect::>(); + Ok(ScheduledScanLine { + decoders: wrapped_decoders, + rows_scheduled: inner_scan.rows_scheduled, + }) + } + + fn num_rows(&self) -> u64 { + self.inner.num_rows() + } +} + +/// A logical scheduler for utf8/binary pages which assumes the data are encoded as `List` +#[derive(Debug)] +pub struct BinaryFieldScheduler { + varbin_scheduler: Arc, + data_type: DataType, +} + +impl BinaryFieldScheduler { + // Create a new ListPageScheduler + pub fn new(varbin_scheduler: Arc, data_type: DataType) -> Self { + Self { + varbin_scheduler, + data_type, + } + } +} + +impl FieldScheduler for BinaryFieldScheduler { + fn schedule_ranges<'a>( + &'a self, + ranges: &[std::ops::Range], + filter: &FilterExpression, + ) -> Result> { + trace!("Scheduling binary for {} ranges", ranges.len()); + let varbin_job = self.varbin_scheduler.schedule_ranges(ranges, filter)?; + Ok(Box::new(BinarySchedulingJob { + scheduler: self, + inner: varbin_job, + })) + } + + fn num_rows(&self) -> u64 { + self.varbin_scheduler.num_rows() + } + + fn initialize<'a>( + &'a self, + _filter: &'a FilterExpression, + _context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + // 2.0 schedulers do not need to initialize + std::future::ready(Ok(())).boxed() + } +} + +#[derive(Debug)] +pub struct BinaryPageDecoder { + inner: Box, + data_type: DataType, +} + +impl LogicalPageDecoder for BinaryPageDecoder { + fn wait_for_loaded(&mut self, num_rows: u64) -> BoxFuture<'_, Result<()>> { + self.inner.wait_for_loaded(num_rows) + } + + fn drain(&mut self, num_rows: u64) -> Result { + let inner_task = self.inner.drain(num_rows)?; + Ok(NextDecodeTask { + num_rows: inner_task.num_rows, + task: Box::new(BinaryArrayDecoder { + inner: inner_task.task, + data_type: self.data_type.clone(), + }), + }) + } + + fn max_rows_to_drain(&self, num_rows: u64, byte_budget: u64) -> Result { + self.inner.max_rows_to_drain(num_rows, byte_budget) + } + + fn data_type(&self) -> &DataType { + &self.data_type + } + + fn rows_loaded(&self) -> u64 { + self.inner.rows_loaded() + } + + fn num_rows(&self) -> u64 { + self.inner.num_rows() + } + + fn rows_drained(&self) -> u64 { + self.inner.rows_drained() + } +} + +pub struct BinaryArrayDecoder { + inner: Box, + data_type: DataType, +} + +impl BinaryArrayDecoder { + fn from_list_array(array: &GenericListArray) -> Result { + let values = array + .values() + .as_primitive::() + .values() + .inner() + .clone(); + let offsets = array.offsets().clone(); + let array = GenericByteArray::::try_new(offsets, values, array.nulls().cloned())?; + Ok(Arc::new(array)) + } +} + +impl DecodeArrayTask for BinaryArrayDecoder { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + let data_type = self.data_type; + let (arr, _) = self.inner.decode()?; + let result = match data_type { + DataType::Binary => Self::from_list_array::(arr.as_list::())?, + DataType::LargeBinary => { + Self::from_list_array::(arr.as_list::())? + } + DataType::Utf8 => Self::from_list_array::(arr.as_list::())?, + DataType::LargeUtf8 => Self::from_list_array::(arr.as_list::())?, + other => { + return Err(Error::internal(format!( + "Binary decoder does not support data type {other}" + ))); + } + }; + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array + // v2.0 path does not need it so we return 0. + Ok((result, 0)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{ListArray, UInt8Array}; + use arrow_buffer::OffsetBuffer; + use arrow_schema::Field; + + use super::*; + use crate::decoder::DecodeArrayTask; + + struct StubDecodeTask { + array: ArrayRef, + } + + impl DecodeArrayTask for StubDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + Ok((self.array, 0)) + } + } + + fn make_single_byte_list(value: u8) -> ListArray { + let offsets = OffsetBuffer::from_lengths([1_usize]); + let values: ArrayRef = Arc::new(UInt8Array::from(vec![value])); + ListArray::try_new( + Arc::new(Field::new("item", DataType::UInt8, false)), + offsets, + values, + None, + ) + .unwrap() + } + + #[test] + fn logical_utf8_decode_preserves_non_overflow_arrow_error() { + let list = make_single_byte_list(0xFF_u8); + let decoder = BinaryArrayDecoder { + inner: Box::new(StubDecodeTask { + array: Arc::new(list), + }), + data_type: DataType::Utf8, + }; + + let error = Box::new(decoder).decode().unwrap_err(); + let message = error.to_string(); + assert!(!message.contains("more than 2GiB of string/binary data")); + assert!(message.to_lowercase().contains("utf")); + } + + #[test] + fn logical_binary_decode_returns_internal_error_for_unsupported_type() { + let list = make_single_byte_list(b'x'); + let decoder = BinaryArrayDecoder { + inner: Box::new(StubDecodeTask { + array: Arc::new(list), + }), + data_type: DataType::Int32, + }; + + let error = Box::new(decoder).decode().unwrap_err(); + assert!(matches!(error, Error::Internal { .. })); + assert!( + error + .to_string() + .contains("Binary decoder does not support data type Int32") + ); + } +} diff --git a/vendor/lance-encoding/src/array_encoding/logical/blob.rs b/vendor/lance-encoding/src/array_encoding/logical/blob.rs new file mode 100644 index 000000000..6789125ac --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/logical/blob.rs @@ -0,0 +1,544 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::VecDeque, sync::Arc, vec}; + +use arrow_array::{ + Array, ArrayRef, LargeBinaryArray, PrimitiveArray, StructArray, UInt64Array, cast::AsArray, + types::UInt64Type, +}; +use arrow_buffer::{ + BooleanBuffer, BooleanBufferBuilder, Buffer, NullBuffer, OffsetBuffer, ScalarBuffer, +}; +use arrow_schema::DataType; +use bytes::Bytes; +use futures::{FutureExt, future::BoxFuture}; + +use lance_core::{Error, Result, datatypes::BLOB_DESC_FIELDS}; + +use crate::{ + EncodingsIo, + buffer::LanceBuffer, + decoder::{ + DecodeArrayTask, FilterExpression, MessageType, NextDecodeTask, PriorityRange, + ScheduledScanLine, SchedulerContext, + }, + decoder::{DecoderReady, FieldScheduler, LogicalPageDecoder, SchedulingJob}, + encoder::{EncodeTask, FieldEncoder, OutOfLineBuffers}, + format::pb::{Blob, ColumnEncoding, column_encoding}, + repdef::RepDefBuilder, +}; + +/// A field scheduler for large binary data +/// +/// Large binary data (1MiB+) can be inefficient if we store as a regular primitive. We +/// essentially end up with 1 page per row (or a few rows) and the overhead of the +/// metadata can be significant. +/// +/// At the same time the benefits of using pages (contiguous arrays) are pretty small since +/// we can generally perform random access at these sizes without much penalty. +/// +/// This encoder gives up the random access and stores the large binary data out of line. This +/// keeps the metadata small. +#[derive(Debug)] +pub struct BlobFieldScheduler { + descriptions_scheduler: Arc, +} + +impl BlobFieldScheduler { + pub fn new(descriptions_scheduler: Arc) -> Self { + Self { + descriptions_scheduler, + } + } +} + +#[derive(Debug)] +struct BlobFieldSchedulingJob<'a> { + descriptions_job: Box, +} + +impl SchedulingJob for BlobFieldSchedulingJob<'_> { + fn schedule_next( + &mut self, + context: &mut SchedulerContext, + priority: &dyn PriorityRange, + ) -> Result { + let next_descriptions = self.descriptions_job.schedule_next(context, priority)?; + let mut priority = priority.current_priority(); + let decoders = next_descriptions.decoders.into_iter().map(|decoder| { + let decoder = decoder.into_array(); + let path = decoder.path; + let mut decoder = decoder.decoder; + let num_rows = decoder.num_rows(); + let descriptions_fut = async move { + decoder + .wait_for_loaded(decoder.num_rows() - 1) + .await + .unwrap(); + let descriptions_task = decoder.drain(decoder.num_rows()).unwrap(); + descriptions_task.task.decode().map(|(arr, _)| arr) + } + .boxed(); + let decoder = Box::new(BlobFieldDecoder { + io: context.io().clone(), + unloaded_descriptions: Some(descriptions_fut), + positions: PrimitiveArray::::from_iter_values(vec![]), + sizes: PrimitiveArray::::from_iter_values(vec![]), + num_rows, + loaded: VecDeque::new(), + validity: VecDeque::new(), + rows_loaded: 0, + rows_drained: 0, + base_priority: priority, + }); + priority += num_rows; + MessageType::DecoderReady(DecoderReady { decoder, path }) + }); + Ok(ScheduledScanLine { + decoders: decoders.collect(), + rows_scheduled: next_descriptions.rows_scheduled, + }) + } + + fn num_rows(&self) -> u64 { + self.descriptions_job.num_rows() + } +} + +impl FieldScheduler for BlobFieldScheduler { + fn schedule_ranges<'a>( + &'a self, + ranges: &[std::ops::Range], + filter: &FilterExpression, + ) -> Result> { + let descriptions_job = self + .descriptions_scheduler + .schedule_ranges(ranges, filter)?; + Ok(Box::new(BlobFieldSchedulingJob { descriptions_job })) + } + + fn num_rows(&self) -> u64 { + self.descriptions_scheduler.num_rows() + } + + fn initialize<'a>( + &'a self, + filter: &'a FilterExpression, + context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + self.descriptions_scheduler.initialize(filter, context) + } +} + +pub struct BlobFieldDecoder { + io: Arc, + unloaded_descriptions: Option>>, + positions: PrimitiveArray, + sizes: PrimitiveArray, + num_rows: u64, + loaded: VecDeque, + validity: VecDeque, + rows_loaded: u64, + rows_drained: u64, + base_priority: u64, +} + +impl BlobFieldDecoder { + fn drain_validity(&mut self, num_values: usize) -> Result> { + let mut validity = BooleanBufferBuilder::new(num_values); + let mut remaining = num_values; + while remaining > 0 { + let next = self.validity.front_mut().unwrap(); + if remaining < next.len() { + let slice = next.slice(0, remaining); + validity.append_buffer(&slice); + *next = next.slice(remaining, next.len() - remaining); + remaining = 0; + } else { + validity.append_buffer(next); + remaining -= next.len(); + self.validity.pop_front(); + } + } + let nulls = NullBuffer::new(validity.finish()); + if nulls.null_count() == 0 { + Ok(None) + } else { + Ok(Some(nulls)) + } + } +} + +impl std::fmt::Debug for BlobFieldDecoder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BlobFieldDecoder") + .field("num_rows", &self.num_rows) + .field("rows_loaded", &self.rows_loaded) + .field("rows_drained", &self.rows_drained) + .finish() + } +} + +impl LogicalPageDecoder for BlobFieldDecoder { + fn wait_for_loaded(&mut self, loaded_need: u64) -> BoxFuture<'_, Result<()>> { + async move { + if self.unloaded_descriptions.is_some() { + let descriptions = self.unloaded_descriptions.take().unwrap().await?; + let descriptions = descriptions.as_struct(); + self.positions = descriptions.column(0).as_primitive().clone(); + self.sizes = descriptions.column(1).as_primitive().clone(); + } + let start = self.rows_loaded as usize; + let end = (loaded_need + 1).min(self.num_rows) as usize; + let positions = self.positions.values().slice(start, end - start); + let sizes = self.sizes.values().slice(start, end - start); + let ranges = positions + .iter() + .zip(sizes.iter()) + .map(|(position, size)| *position..(*position + *size)) + .collect::>(); + let validity = positions + .iter() + .zip(sizes.iter()) + .map(|(p, s)| *p != 1 || *s != 0) + .collect::(); + // Run the I/O before mutating decoder state, so a failed load + // leaves the decoder untouched and `wait_for_loaded` is the single, + // clean point of failure. + let bytes = self + .io + .submit_request(ranges, self.base_priority + start as u64) + .await?; + self.validity.push_back(validity); + self.loaded.extend(bytes); + self.rows_loaded = end as u64; + Ok(()) + } + .boxed() + } + + fn rows_loaded(&self) -> u64 { + self.rows_loaded + } + + fn num_rows(&self) -> u64 { + self.num_rows + } + + fn rows_drained(&self) -> u64 { + self.rows_drained + } + + fn drain(&mut self, num_rows: u64) -> Result { + if num_rows as usize > self.loaded.len() { + // `loaded` is populated by `wait_for_loaded`; guard the + // load-before-drain contract so a violation surfaces as an error + // rather than an out-of-bounds drain panic. + return Err(Error::internal(format!( + "BlobFieldDecoder was asked to drain {num_rows} rows but only \ + {} are loaded", + self.loaded.len(), + ))); + } + let bytes = self.loaded.drain(0..num_rows as usize).collect::>(); + let validity = self.drain_validity(num_rows as usize)?; + self.rows_drained += num_rows; + Ok(NextDecodeTask { + num_rows, + task: Box::new(BlobArrayDecodeTask::new(bytes, validity)), + }) + } + + fn data_type(&self) -> &DataType { + &DataType::LargeBinary + } +} + +struct BlobArrayDecodeTask { + bytes: Vec, + validity: Option, +} + +impl BlobArrayDecodeTask { + fn new(bytes: Vec, validity: Option) -> Self { + Self { bytes, validity } + } +} + +impl DecodeArrayTask for BlobArrayDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + let num_bytes = self.bytes.iter().map(|b| b.len()).sum::(); + let offsets = self + .bytes + .iter() + .scan(0, |state, b| { + let start = *state; + *state += b.len(); + Some(start as i64) + }) + .chain(std::iter::once(num_bytes as i64)) + .collect::>(); + let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets)); + let mut buffer = Vec::with_capacity(num_bytes); + for bytes in self.bytes { + buffer.extend_from_slice(&bytes); + } + let data_buf = Buffer::from_vec(buffer); + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array + // v2.0 path does not need it so we return 0. + Ok(( + Arc::new(LargeBinaryArray::new(offsets, data_buf, self.validity)), + 0, + )) + } +} + +pub struct BlobFieldEncoder { + description_encoder: Box, +} + +impl BlobFieldEncoder { + pub fn new(description_encoder: Box) -> Self { + Self { + description_encoder, + } + } + + fn write_bins(array: ArrayRef, external_buffers: &mut OutOfLineBuffers) -> Result { + let binarray = array.as_binary_opt::().ok_or_else(|| { + Error::invalid_input_source( + format!("Expected large_binary and received {}", array.data_type()).into(), + ) + })?; + let mut positions = Vec::with_capacity(array.len()); + let mut sizes = Vec::with_capacity(array.len()); + let data = binarray.values(); + let nulls = binarray + .nulls() + .cloned() + .unwrap_or(NullBuffer::new_valid(binarray.len())); + for (w, is_valid) in binarray.value_offsets().windows(2).zip(&nulls) { + if is_valid { + let start = w[0] as u64; + let end = w[1] as u64; + let size = end - start; + if size > 0 { + let val = data.slice_with_length(start as usize, size as usize); + let position = external_buffers.add_buffer(LanceBuffer::from(val)); + positions.push(position); + sizes.push(size); + } else { + // Empty values are always (0,0) + positions.push(0); + sizes.push(0); + } + } else { + // Null values are always (1, 0) + positions.push(1); + sizes.push(0); + } + } + let positions = Arc::new(UInt64Array::from(positions)); + let sizes = Arc::new(UInt64Array::from(sizes)); + let descriptions = Arc::new(StructArray::new( + BLOB_DESC_FIELDS.clone(), + vec![positions, sizes], + None, + )); + Ok(descriptions) + } +} + +impl FieldEncoder for BlobFieldEncoder { + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + let descriptions = Self::write_bins(array, external_buffers)?; + self.description_encoder.maybe_encode( + descriptions, + external_buffers, + repdef, + row_number, + num_rows, + ) + } + + // If there is any data left in the buffer then create an encode task from it + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + self.description_encoder.flush(external_buffers) + } + + fn num_columns(&self) -> u32 { + self.description_encoder.num_columns() + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + let inner_finished = self.description_encoder.finish(external_buffers); + async move { + let mut cols = inner_finished.await?; + assert_eq!(cols.len(), 1); + let encoding = std::mem::take(&mut cols[0].encoding); + let wrapped_encoding = ColumnEncoding { + column_encoding: Some(column_encoding::ColumnEncoding::Blob(Box::new(Blob { + inner: Some(Box::new(encoding)), + }))), + }; + cols[0].encoding = wrapped_encoding; + Ok(cols) + } + .boxed() + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::{HashMap, VecDeque}, + ops::Range, + sync::{Arc, LazyLock}, + }; + + use arrow_array::{LargeBinaryArray, PrimitiveArray, types::UInt64Type}; + use arrow_schema::{DataType, Field}; + use bytes::Bytes; + use futures::{FutureExt, future::BoxFuture}; + use lance_arrow::BLOB_META_KEY; + use lance_core::{Error, Result}; + + use super::BlobFieldDecoder; + use crate::{ + EncodingsIo, + decoder::LogicalPageDecoder, + format::pb::column_encoding, + testing::TestEncoding, + testing::{TestCases, check_round_trip_encoding_of_data, check_specific_random}, + }; + + static BLOB_META: LazyLock> = LazyLock::new(|| { + [(BLOB_META_KEY.to_string(), "true".to_string())] + .iter() + .cloned() + .collect::>() + }); + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_basic_blob( + #[values(TestEncoding::Array, TestEncoding::StructuralU16)] encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let field = Field::new("", DataType::LargeBinary, false).with_metadata(BLOB_META.clone()); + check_specific_random( + field, + TestCases::basic() + .with_encoding(encoding) + .with_page_sizes(vec![page_size]) + .with_slicing_modes([use_slicing]), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_blob() { + let val1: &[u8] = &[1, 2, 3]; + let val2: &[u8] = &[7, 8, 9]; + let array = Arc::new(LargeBinaryArray::from(vec![Some(val1), None, Some(val2)])); + let test_cases = TestCases::default() + .with_array_and_u16_encodings() + .with_expected_encoding("packed_struct") + .with_verify_encoding(Arc::new(|cols, encoding| { + if *encoding == TestEncoding::Array { + // In 2.0 we used a special "column encoding" to mark blob fields. In 2.1 we + // don't do this and just rely on the regular page encoding. + assert_eq!(cols.len(), 1); + let col = &cols[0]; + assert!(matches!( + col.encoding.column_encoding.as_ref().unwrap(), + column_encoding::ColumnEncoding::Blob(_) + )); + } + })); + // Use blob encoding if requested + check_round_trip_encoding_of_data(vec![array.clone()], &test_cases, BLOB_META.clone()) + .await; + + let test_cases = TestCases::default() + .with_structural_encodings() + .with_verify_encoding(Arc::new(|cols, encoding| { + if *encoding == TestEncoding::Array { + assert_eq!(cols.len(), 1); + let col = &cols[0]; + assert!(!matches!( + col.encoding.column_encoding.as_ref().unwrap(), + column_encoding::ColumnEncoding::Blob(_) + )); + } + })); + // Don't use blob encoding if not requested + check_round_trip_encoding_of_data(vec![array], &test_cases, Default::default()).await; + } + + /// An `EncodingsIo` that rejects every request, simulating cloud storage + /// returning a retryable error (e.g. an exhausted HTTP 503 retry budget). + #[derive(Debug)] + struct FailingScheduler; + + impl EncodingsIo for FailingScheduler { + fn submit_request( + &self, + _ranges: Vec>, + _priority: u64, + ) -> BoxFuture<'static, Result>> { + std::future::ready(Err(Error::io("simulated HTTP 503 from cloud storage"))).boxed() + } + } + + /// A failed blob load must surface through `wait_for_loaded` and leave the + /// decoder untouched -- never half-advanced into a state where a later + /// `drain` reads rows that never loaded (which panicked with "range end + /// index N out of range for slice of length 0"). + #[test_log::test(tokio::test)] + async fn test_io_failure_leaves_blob_decoder_consistent() { + let num_rows = 8u64; + // `positions`/`sizes` only need `num_rows` entries; the failing + // scheduler rejects the request regardless of the ranges it is given. + let descs = PrimitiveArray::::from_iter_values(std::iter::repeat_n( + 0u64, + num_rows as usize, + )); + + let mut decoder = BlobFieldDecoder { + io: Arc::new(FailingScheduler), + unloaded_descriptions: None, + positions: descs.clone(), + sizes: descs, + num_rows, + loaded: VecDeque::new(), + validity: VecDeque::new(), + rows_loaded: 0, + rows_drained: 0, + base_priority: 0, + }; + + // `wait_for_loaded` propagates the I/O failure... + assert!(decoder.wait_for_loaded(num_rows - 1).await.is_err()); + + // ...and leaves no half-loaded state behind. + assert_eq!(decoder.rows_loaded, 0); + assert!(decoder.loaded.is_empty()); + assert!(decoder.validity.is_empty()); + + // A drain in this state errors instead of panicking on the empty buffer. + assert!(decoder.drain(num_rows).is_err()); + } +} diff --git a/vendor/lance-encoding/src/array_encoding/logical/list.rs b/vendor/lance-encoding/src/array_encoding/logical/list.rs new file mode 100644 index 000000000..453e0c744 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/logical/list.rs @@ -0,0 +1,1505 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ + collections::VecDeque, + ops::Range, + sync::{Arc, OnceLock}, +}; + +use arrow_array::{ + Array, ArrayRef, BooleanArray, Int32Array, Int64Array, LargeListArray, ListArray, UInt64Array, + cast::AsArray, + new_empty_array, + types::{Int32Type, Int64Type, UInt64Type}, +}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, Buffer, NullBuffer, OffsetBuffer}; +use arrow_schema::{DataType, Field, Fields}; +use futures::{FutureExt, future::BoxFuture}; +use lance_core::{Error, Result, cache::LanceCache, utils::parse::str_is_truthy}; +use log::trace; +use tokio::task::JoinHandle; + +use crate::{ + EncodingsIo, + array_encoding::logical::r#struct::{SimpleStructDecoder, SimpleStructScheduler}, + buffer::LanceBuffer, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + decoder::{ + DecodeArrayTask, DecodeBatchScheduler, DrainLimit, FieldScheduler, FilterExpression, + ListPriorityRange, LogicalPageDecoder, MessageType, NextDecodeTask, PageEncoding, + PriorityRange, ScheduledScanLine, SchedulerContext, SchedulingJob, + }, + encoder::{ + ArrayEncoder, EncodeTask, EncodedArray, EncodedColumn, EncodedPage, FieldEncoder, + OutOfLineBuffers, + }, + format::pb, + repdef::RepDefBuilder, + utils::accumulation::AccumulationQueue, +}; + +/// When set, indirect I/O in the 2.0 list scheduler bypasses the backpressure system. +/// +/// This can be a blunt instrument to avoid deadlocks in 2.0 scenarios +/// Set LANCE_BYPASS_INDIRECT_IO_BACKPRESSURE=1 to enable. +static BYPASS_INDIRECT_IO_BACKPRESSURE: OnceLock = OnceLock::new(); + +fn bypass_indirect_io_backpressure() -> bool { + *BYPASS_INDIRECT_IO_BACKPRESSURE.get_or_init(|| { + std::env::var("LANCE_BYPASS_INDIRECT_IO_BACKPRESSURE") + .map(|val| str_is_truthy(&val)) + .unwrap_or(false) + }) +} + +// Scheduling lists is tricky. Imagine the following scenario: +// +// * There are 2000 offsets per offsets page +// * The user requests range 8000..8500 +// +// First, since 8000 matches the start of an offsets page, we don't need to read an extra offset. +// +// Since this range matches the start of a page, we know we will get an offsets array like +// [0, ...] +// +// We need to restore nulls, which relies on a null offset adjustment, which is unique to each offsets +// page. +// +// We need to map this to [X, ...] where X is the sum of the number of items in the 0-2000, 2000-4000, +// and 4000-6000 pages. +// +// This gets even trickier if a range spans multiple offsets pages. For example, given the same +// scenario but the user requests 7999..8500. In this case the first page read will include an +// extra offset (e.g. we need to read 7998..8000), the null adjustment will be different between the +// two, and the items offset will be different. +// +// To handle this, we take the incoming row requests, look at the page info, and then calculate +// list requests. + +#[derive(Debug)] +struct ListRequest { + /// How many lists this request maps to + num_lists: u64, + /// Did this request include an extra offset + includes_extra_offset: bool, + /// The null offset adjustment for this request + null_offset_adjustment: u64, + /// items offset to apply + items_offset: u64, +} + +#[derive(Debug)] +struct ListRequestsIter { + // The bool triggers whether we need to skip an offset or not + list_requests: VecDeque, + offsets_requests: Vec>, +} + +impl ListRequestsIter { + // TODO: This logic relies on row_ranges being ordered and may be a problem when we + // add proper support for out-of-order take + fn new(row_ranges: &[Range], page_infos: &[OffsetPageInfo]) -> Self { + let mut items_offset = 0; + let mut offsets_offset = 0; + let mut page_infos_iter = page_infos.iter(); + let mut cur_page_info = page_infos_iter.next().unwrap(); + let mut list_requests = VecDeque::new(); + let mut offsets_requests = Vec::new(); + + // Each row range maps to at least one list request. It may map to more if the + // range spans multiple offsets pages. + for range in row_ranges { + let mut range = range.clone(); + + // Skip any offsets pages that are before the range + while offsets_offset + (cur_page_info.offsets_in_page) <= range.start { + trace!("Skipping null offset adjustment chunk {:?}", offsets_offset); + offsets_offset += cur_page_info.offsets_in_page; + items_offset += cur_page_info.num_items_referenced_by_page; + cur_page_info = page_infos_iter.next().unwrap(); + } + + // If the range starts at the beginning of an offsets page we don't need + // to read an extra offset + let mut includes_extra_offset = range.start != offsets_offset; + if includes_extra_offset { + offsets_requests.push(range.start - 1..range.end); + } else { + offsets_requests.push(range.clone()); + } + + // At this point our range overlaps the current page (cur_page_info) and + // we can start slicing it into list requests + while !range.is_empty() { + // The end of the list request is the min of the end of the range + // and the end of the current page + let end = offsets_offset + cur_page_info.offsets_in_page; + let last = end >= range.end; + let end = end.min(range.end); + list_requests.push_back(ListRequest { + num_lists: end - range.start, + includes_extra_offset, + null_offset_adjustment: cur_page_info.null_offset_adjustment, + items_offset, + }); + + includes_extra_offset = false; + range.start = end; + // If there is still more data in the range, we need to move to the + // next page + if !last { + offsets_offset += cur_page_info.offsets_in_page; + items_offset += cur_page_info.num_items_referenced_by_page; + cur_page_info = page_infos_iter.next().unwrap(); + } + } + } + Self { + list_requests, + offsets_requests, + } + } + + // Given a page of offset data, grab the corresponding list requests + fn next(&mut self, mut num_offsets: u64) -> Vec { + let mut list_requests = Vec::new(); + while num_offsets > 0 { + let req = self.list_requests.front_mut().unwrap(); + // If the request did not start at zero then we need to read an extra offset + if req.includes_extra_offset { + num_offsets -= 1; + debug_assert_ne!(num_offsets, 0); + } + if num_offsets >= req.num_lists { + num_offsets -= req.num_lists; + list_requests.push(self.list_requests.pop_front().unwrap()); + } else { + let sub_req = ListRequest { + num_lists: num_offsets, + includes_extra_offset: req.includes_extra_offset, + null_offset_adjustment: req.null_offset_adjustment, + items_offset: req.items_offset, + }; + + list_requests.push(sub_req); + req.includes_extra_offset = false; + req.num_lists -= num_offsets; + num_offsets = 0; + } + } + list_requests + } +} + +/// Given a list of offsets and a list of requested list row ranges we need to rewrite the offsets so that +/// they appear as expected for a list array. This involves a number of tasks: +/// +/// * Nulls in the offsets are represented by oversize values and these need to be converted to +/// the appropriate length +/// * For each range we (usually) load N + 1 offsets, so if we have 5 ranges we have 5 extra values +/// and we need to drop 4 of those. +/// * Ranges may not start at 0 and, while we don't strictly need to, we want to go ahead and normalize +/// the offsets so that the first offset is 0. +/// +/// Throughout the comments we will consider the following example case: +/// +/// The user requests the following ranges of lists (list_row_ranges): [0..3, 5..6] +/// +/// This is a total of 4 lists. The loaded offsets are [10, 20, 120, 150, 60]. The last valid offset is 99. +/// The null_offset_adjustment will be 100. +/// +/// Our desired output offsets are going to be [0, 10, 20, 20, 30] and the item ranges are [0..20] and [50..60] +/// The validity array is [true, true, false, true] +fn decode_offsets( + offsets: &dyn Array, + list_requests: &[ListRequest], + null_offset_adjustment: u64, +) -> (VecDeque>, Vec, BooleanBuffer) { + // In our example this is [10, 20, 120, 50, 60] + let numeric_offsets = offsets.as_primitive::(); + // In our example there are 4 total lists + let total_num_lists = list_requests.iter().map(|req| req.num_lists).sum::() as u32; + let mut normalized_offsets = Vec::with_capacity(total_num_lists as usize); + let mut validity_buffer = BooleanBufferBuilder::new(total_num_lists as usize); + // The first output offset is always 0 no matter what + normalized_offsets.push(0); + let mut last_normalized_offset = 0; + let offsets_values = numeric_offsets.values(); + + let mut item_ranges = VecDeque::new(); + let mut offsets_offset: u32 = 0; + // All ranges should be non-empty + debug_assert!(list_requests.iter().all(|r| r.num_lists > 0)); + for req in list_requests { + // The # of lists in this particular range + let num_lists = req.num_lists; + + // Because we know the first offset is always 0 we don't store that. This means we have special + // logic if a range starts at 0 (we didn't need to read an extra offset value in that case) + // In our example we enter this special case on the first range (0..3) but not the second (5..6) + // This means the first range, which has 3 lists, maps to 3 values in our offsets array [10, 20, 120] + // However, the second range, which has 1 list, maps to 2 values in our offsets array [150, 60] + let (items_range, offsets_to_norm_start, num_offsets_to_norm) = + if !req.includes_extra_offset { + // In our example items start is 0 and items_end is 20 + let first_offset_idx = 0_usize; + let num_offsets = num_lists as usize; + let items_start = 0; + let items_end = offsets_values[num_offsets - 1] % null_offset_adjustment; + let items_range = items_start..items_end; + (items_range, first_offset_idx, num_offsets) + } else { + // In our example, offsets_offset will be 3, items_start will be 50, and items_end will + // be 60 + let first_offset_idx = offsets_offset as usize; + let num_offsets = num_lists as usize + 1; + let items_start = offsets_values[first_offset_idx] % null_offset_adjustment; + let items_end = + offsets_values[first_offset_idx + num_offsets - 1] % null_offset_adjustment; + let items_range = items_start..items_end; + (items_range, first_offset_idx, num_offsets) + }; + + // TODO: Maybe consider writing whether there are nulls or not as part of the + // page description. Then we can skip all validity work. Not clear if that will + // be any benefit though. + + // We calculate validity from all elements but the first (or all elements + // if this is the special zero-start case) + // + // So, in our first pass through, we consider [10, 20, 120] (1 null) + // In our second pass through we only consider [60] (0 nulls) + // Note that the 150 is null but we only loaded it to know where the 50-60 list started + // and it doesn't actually correspond to a list (e.g. list 4 is null but we aren't loading it + // here) + let validity_start = if !req.includes_extra_offset { + 0 + } else { + offsets_to_norm_start + 1 + }; + for off in offsets_values + .slice(validity_start, num_lists as usize) + .iter() + { + validity_buffer.append(*off < null_offset_adjustment); + } + + // In our special case we need to account for the offset 0-first_item + if !req.includes_extra_offset { + let first_item = offsets_values[0] % null_offset_adjustment; + normalized_offsets.push(first_item); + last_normalized_offset = first_item; + } + + // Finally, we go through and shift the offsets. If we just returned them as is (taking care of + // nulls) we would get [0, 10, 20, 20, 60] but our last list only has 10 items, not 40 and so we + // need to shift that 60 to a 40. + normalized_offsets.extend( + offsets_values + .slice(offsets_to_norm_start, num_offsets_to_norm) + .windows(2) + .map(|w| { + let start = w[0] % null_offset_adjustment; + let end = w[1] % null_offset_adjustment; + if end < start { + panic!("End is less than start in window {:?} with null_offset_adjustment={} we get start={} and end={}", w, null_offset_adjustment, start, end); + } + let length = end - start; + last_normalized_offset += length; + last_normalized_offset + }), + ); + trace!( + "List offsets range of {} lists maps to item range {:?}", + num_lists, items_range + ); + offsets_offset += num_offsets_to_norm as u32; + if !items_range.is_empty() { + let items_range = + items_range.start + req.items_offset..items_range.end + req.items_offset; + item_ranges.push_back(items_range); + } + } + + let validity = validity_buffer.finish(); + (item_ranges, normalized_offsets, validity) +} + +/// After scheduling the offsets we immediately launch this task as a new tokio task +/// This task waits for the offsets to arrive, decodes them, and then schedules the I/O +/// for the items. +/// +/// This task does not wait for the items data. That happens on the main decode loop (unless +/// we have list of list of ... in which case it happens in the outer indirect decode loop) +#[allow(clippy::too_many_arguments)] +async fn indirect_schedule_task( + mut offsets_decoder: Box, + list_requests: Vec, + null_offset_adjustment: u64, + items_scheduler: Arc, + items_type: DataType, + io: Arc, + cache: Arc, + priority: Box, +) -> Result { + let num_offsets = offsets_decoder.num_rows(); + // We know the offsets are a primitive array and thus will not need additional + // pages. We can use a dummy receiver to match the decoder API + offsets_decoder.wait_for_loaded(num_offsets - 1).await?; + let decode_task = offsets_decoder.drain(num_offsets)?; + let (offsets, _) = decode_task.task.decode()?; + + let (item_ranges, offsets, validity) = + decode_offsets(offsets.as_ref(), &list_requests, null_offset_adjustment); + + trace!( + "Indirectly scheduling items ranges {:?} from list items column with {} rows (and priority {:?})", + item_ranges, + items_scheduler.num_rows(), + priority + ); + let offsets: Arc<[u64]> = offsets.into(); + + // All requested lists are empty + if item_ranges.is_empty() { + debug_assert!(item_ranges.iter().all(|r| r.start == r.end)); + return Ok(IndirectlyLoaded { + root_decoder: None, + offsets, + validity, + }); + } + let item_ranges = item_ranges.into_iter().collect::>(); + let num_items = item_ranges.iter().map(|r| r.end - r.start).sum::(); + + // Create a new root scheduler, which has one column, which is our items data + let root_fields = Fields::from(vec![Field::new("item", items_type, true)]); + let indirect_root_scheduler = + SimpleStructScheduler::new(vec![items_scheduler], root_fields.clone(), num_items); + #[allow(deprecated)] + let mut indirect_scheduler = DecodeBatchScheduler::from_scheduler( + Arc::new(indirect_root_scheduler), + root_fields.clone(), + cache, + ); + let mut root_decoder = SimpleStructDecoder::new(root_fields, num_items); + + let priority = Box::new(ListPriorityRange::new(priority, offsets.clone())); + + let indirect_messages = indirect_scheduler.schedule_ranges_to_vec( + &item_ranges, + // Can't push filters into list items + &FilterExpression::no_filter(), + io, + Some(priority), + )?; + + for message in indirect_messages { + for decoder in message.decoders { + let decoder = decoder.into_array(); + if !decoder.path.is_empty() { + root_decoder.accept_child(decoder)?; + } + } + } + + Ok(IndirectlyLoaded { + offsets, + validity, + root_decoder: Some(root_decoder), + }) +} + +#[derive(Debug)] +struct ListFieldSchedulingJob<'a> { + scheduler: &'a ListFieldScheduler, + offsets: Box, + num_rows: u64, + list_requests_iter: ListRequestsIter, +} + +impl<'a> ListFieldSchedulingJob<'a> { + fn try_new( + scheduler: &'a ListFieldScheduler, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result { + let list_requests_iter = ListRequestsIter::new(ranges, &scheduler.offset_page_info); + let num_rows = ranges.iter().map(|r| r.end - r.start).sum::(); + let offsets = scheduler + .offsets_scheduler + .schedule_ranges(&list_requests_iter.offsets_requests, filter)?; + Ok(Self { + scheduler, + offsets, + list_requests_iter, + num_rows, + }) + } +} + +impl SchedulingJob for ListFieldSchedulingJob<'_> { + fn schedule_next( + &mut self, + context: &mut SchedulerContext, + priority: &dyn PriorityRange, + ) -> Result { + let next_offsets = self.offsets.schedule_next(context, priority)?; + let offsets_scheduled = next_offsets.rows_scheduled; + let list_reqs = self.list_requests_iter.next(offsets_scheduled); + trace!( + "Scheduled {} offsets which maps to list requests: {:?}", + offsets_scheduled, list_reqs + ); + let null_offset_adjustment = list_reqs[0].null_offset_adjustment; + // It shouldn't be possible for `list_reqs` to span more than one offsets page and so it shouldn't + // be possible for the null_offset_adjustment to change + debug_assert!( + list_reqs + .iter() + .all(|req| req.null_offset_adjustment == null_offset_adjustment) + ); + let num_rows = list_reqs.iter().map(|req| req.num_lists).sum::(); + // offsets is a uint64 which is guaranteed to create one decoder on each call to schedule_next + let next_offsets_decoder = next_offsets + .decoders + .into_iter() + .next() + .unwrap() + .into_array() + .decoder; + + let items_scheduler = self.scheduler.items_scheduler.clone(); + let items_type = self.scheduler.items_field.data_type().clone(); + let base_io = context.io().clone(); + let io = if bypass_indirect_io_backpressure() { + base_io.with_bypass_backpressure().unwrap_or(base_io) + } else { + base_io + }; + let cache = context.cache().clone(); + + // Immediately spawn the indirect scheduling + let indirect_fut = tokio::spawn(indirect_schedule_task( + next_offsets_decoder, + list_reqs, + null_offset_adjustment, + items_scheduler, + items_type, + io, + cache, + priority.box_clone(), + )); + + // Return a decoder + let decoder = Box::new(ListPageDecoder { + offsets: Arc::new([]), + validity: BooleanBuffer::new(Buffer::from_vec(Vec::::default()), 0, 0), + item_decoder: None, + rows_drained: 0, + rows_loaded: 0, + items_field: self.scheduler.items_field.clone(), + num_rows, + unloaded: Some(indirect_fut), + offset_type: self.scheduler.offset_type.clone(), + data_type: self.scheduler.list_type.clone(), + }); + #[allow(deprecated)] + let decoder = context.locate_decoder(decoder); + Ok(ScheduledScanLine { + decoders: vec![MessageType::DecoderReady(decoder)], + rows_scheduled: num_rows, + }) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +/// A page scheduler for list fields that encodes offsets in one field and items in another +/// +/// The list scheduler is somewhat unique because it requires indirect I/O. We cannot know the +/// ranges we need simply by looking at the metadata. This means that list scheduling doesn't +/// fit neatly into the two-thread schedule-loop / decode-loop model. To handle this, when a +/// list page is scheduled, we only schedule the I/O for the offsets and then we immediately +/// launch a new tokio task. This new task waits for the offsets, decodes them, and then +/// schedules the I/O for the items. Keep in mind that list items can be lists themselves. If +/// that is the case then this indirection will continue. The decode task that is returned will +/// only finish `wait`ing when all of the I/O has completed. +/// +/// Whenever we schedule follow-up I/O like this the priority is based on the top-level row +/// index. This helps ensure that earlier rows get finished completely (including follow up +/// tasks) before we perform I/O for later rows. +#[derive(Debug)] +pub struct ListFieldScheduler { + offsets_scheduler: Arc, + items_scheduler: Arc, + items_field: Arc, + offset_type: DataType, + list_type: DataType, + offset_page_info: Vec, +} + +/// The offsets are stored in a uint64 encoded column. For each page we +/// store some supplementary data that helps us understand the offsets. +/// This is needed to construct the scheduler +#[derive(Debug)] +pub struct OffsetPageInfo { + pub offsets_in_page: u64, + pub null_offset_adjustment: u64, + pub num_items_referenced_by_page: u64, +} + +impl ListFieldScheduler { + // Create a new ListPageScheduler + pub fn new( + offsets_scheduler: Arc, + items_scheduler: Arc, + items_field: Arc, + // Should be int32 or int64 + offset_type: DataType, + offset_page_info: Vec, + ) -> Self { + let list_type = match &offset_type { + DataType::Int32 => DataType::List(items_field.clone()), + DataType::Int64 => DataType::LargeList(items_field.clone()), + _ => panic!("Unexpected offset type {}", offset_type), + }; + Self { + offsets_scheduler, + items_scheduler, + items_field, + offset_type, + offset_page_info, + list_type, + } + } +} + +impl FieldScheduler for ListFieldScheduler { + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result> { + Ok(Box::new(ListFieldSchedulingJob::try_new( + self, ranges, filter, + )?)) + } + + fn num_rows(&self) -> u64 { + self.offsets_scheduler.num_rows() + } + + fn initialize<'a>( + &'a self, + _filter: &'a FilterExpression, + _context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + // 2.0 schedulers do not need to initialize + std::future::ready(Ok(())).boxed() + } +} + +/// As soon as the first call to decode comes in we wait for all indirect I/O to +/// complete. +/// +/// Once the indirect I/O is finished we pull items out of `unawaited`, wait them +/// (this wait should return immediately) and then push them into `item_decoders`. +/// +/// We then drain from `item_decoders`, popping item pages off as we finish with +/// them. +/// +/// TODO: Test the case where a single list page has multiple items pages +#[derive(Debug)] +struct ListPageDecoder { + unloaded: Option>>, + // offsets and validity will have already been decoded as part of the indirect I/O + offsets: Arc<[u64]>, + validity: BooleanBuffer, + item_decoder: Option, + num_rows: u64, + rows_drained: u64, + rows_loaded: u64, + items_field: Arc, + offset_type: DataType, + data_type: DataType, +} + +struct ListDecodeTask { + offsets: Vec, + validity: BooleanBuffer, + // Will be None if there are no items (all empty / null lists) + items: Option>, + items_field: Arc, + offset_type: DataType, +} + +fn oversized_batch_error( + items_field: &Field, + requested_range: Range, + decodable_prefix_end: u64, + num_items: u64, +) -> Error { + let prefix_detail = if decodable_prefix_end > requested_range.start { + format!( + "rows {}..{} fit, but requesting rows {}..{} would require {num_items} {}", + requested_range.start, + decodable_prefix_end, + requested_range.start, + requested_range.end, + if items_field.data_type() == &DataType::UInt8 { + "bytes" + } else { + "items" + } + ) + } else { + format!( + "requesting rows {}..{} would require {num_items} {}", + requested_range.start, + requested_range.end, + if items_field.data_type() == &DataType::UInt8 { + "bytes" + } else { + "items" + } + ) + }; + if items_field.data_type() == &DataType::UInt8 { + Error::not_supported(format!( + "Could not create array with more than 2GiB of string/binary data in a single batch \ + ({prefix_detail}). Please reduce the batch_size, set LANCE_DEFAULT_BATCH_SIZE to a \ + smaller value, or convert the column to large_string/large_binary." + )) + } else { + Error::not_supported(format!( + "Could not create a list array with more than i32::MAX items in a single batch \ + ({prefix_detail}). Please reduce the batch_size." + )) + } +} + +impl DecodeArrayTask for ListDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + let items = self + .items + .map(|items| { + // When we run the indirect I/O we wrap things in a struct array with a single field + // named "item". We can unwrap that now. + let (wrapped_items, _) = items.decode()?; + Result::Ok(wrapped_items.as_struct().column(0).clone()) + }) + .unwrap_or_else(|| Ok(new_empty_array(self.items_field.data_type())))?; + + // The offsets are already decoded but they need to be shifted back to 0 and cast + // to the appropriate type + // + // Although, in some cases, the shift IS strictly required since the unshifted offsets + // may cross i32::MAX even though the shifted offsets do not + let offsets = UInt64Array::from(self.offsets); + let validity = NullBuffer::new(self.validity); + let validity = if validity.null_count() == 0 { + None + } else { + Some(validity) + }; + let min_offset = UInt64Array::new_scalar(offsets.value(0)); + let offsets = arrow_arith::numeric::sub(&offsets, &min_offset)?; + let array: ArrayRef = match &self.offset_type { + DataType::Int32 => { + let offsets = arrow_cast::cast(&offsets, &DataType::Int32)?; + let offsets_i32 = offsets.as_primitive::(); + let offsets = OffsetBuffer::new(offsets_i32.values().clone()); + + Arc::new(ListArray::try_new( + self.items_field.clone(), + offsets, + items, + validity, + )?) + } + DataType::Int64 => { + let offsets = arrow_cast::cast(&offsets, &DataType::Int64)?; + let offsets_i64 = offsets.as_primitive::(); + let offsets = OffsetBuffer::new(offsets_i64.values().clone()); + + Arc::new(LargeListArray::try_new( + self.items_field.clone(), + offsets, + items, + validity, + )?) + } + _ => panic!("ListDecodeTask with data type that is not i32 or i64"), + }; + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array + // v2.0 path does not need it so we return 0. + Ok((array, 0)) + } +} + +// Helper method that performs binary search. However, once the +// target is found it walks past any duplicates. E.g. if the +// input list is [0, 3, 5, 5, 5, 7] then this will only return +// 0, 1, 4, or 5. +fn binary_search_to_end(to_search: &[u64], target: u64) -> u64 { + let mut result = match to_search.binary_search(&target) { + Ok(idx) => idx, + Err(idx) => idx - 1, + }; + while result < (to_search.len() - 1) && to_search[result + 1] == target { + result += 1; + } + result as u64 +} + +impl LogicalPageDecoder for ListPageDecoder { + fn wait_for_loaded(&mut self, loaded_need: u64) -> BoxFuture<'_, Result<()>> { + async move { + // wait for the indirect I/O to finish, run the scheduler for the indirect + // I/O and then wait for enough items to arrive + if self.unloaded.is_some() { + trace!("List scheduler needs to wait for indirect I/O to complete"); + let indirectly_loaded = self.unloaded.take().unwrap().await; + if let Err(err) = indirectly_loaded { + match err.try_into_panic() { + Ok(err) => std::panic::resume_unwind(err), + Err(err) => panic!("{:?}", err), + }; + } + let indirectly_loaded = indirectly_loaded.unwrap()?; + + self.offsets = indirectly_loaded.offsets; + self.validity = indirectly_loaded.validity; + self.item_decoder = indirectly_loaded.root_decoder; + } + if self.rows_loaded > loaded_need { + return Ok(()); + } + + let boundary = loaded_need as usize; + debug_assert!(boundary < self.num_rows as usize); + // We need more than X lists which means we need at least X+1 lists which means + // we need at least offsets[X+1] items which means we need more than offsets[X+1]-1 items. + let items_needed = self.offsets[boundary + 1].saturating_sub(1); + trace!( + "List decoder is waiting for more than {} rows to be loaded and {}/{} are already loaded. To satisfy this we need more than {} loaded items", + loaded_need, + self.rows_loaded, + self.num_rows, + items_needed, + ); + + let items_loaded = if let Some(item_decoder) = self.item_decoder.as_mut() { + item_decoder.wait_for_loaded(items_needed).await?; + item_decoder.rows_loaded() + } else { + 0 + }; + + self.rows_loaded = binary_search_to_end(&self.offsets, items_loaded); + trace!("List decoder now has {} loaded rows", self.rows_loaded); + + Ok(()) + } + .boxed() + } + + fn drain(&mut self, num_rows: u64) -> Result { + // We already have the offsets but need to drain the item pages + let mut actual_num_rows = num_rows; + let item_start = self.offsets[self.rows_drained as usize]; + if self.offset_type != DataType::Int64 { + // We might not be able to drain `num_rows` because that request might contain more than 2^31 items + // so we need to figure out how many rows we can actually drain. + while actual_num_rows > 0 { + let num_items = + self.offsets[(self.rows_drained + actual_num_rows) as usize] - item_start; + if num_items <= i32::MAX as u64 { + break; + } + // TODO: This could be slow. Maybe faster to start from zero or do binary search. Investigate when + // actually adding support for smaller than requested batches + actual_num_rows -= 1; + } + } + if actual_num_rows < num_rows { + let requested_range = self.rows_drained..self.rows_drained + num_rows; + let num_items = self.offsets[requested_range.end as usize] - item_start; + return Err(oversized_batch_error( + self.items_field.as_ref(), + requested_range, + self.rows_drained + actual_num_rows, + num_items, + )); + } + let offsets = self.offsets + [self.rows_drained as usize..(self.rows_drained + actual_num_rows + 1) as usize] + .to_vec(); + let validity = self + .validity + .slice(self.rows_drained as usize, actual_num_rows as usize); + let start = offsets[0]; + let end = offsets[offsets.len() - 1]; + let num_items_to_drain = end - start; + + let item_decode = if num_items_to_drain == 0 { + None + } else { + self.item_decoder + .as_mut() + .map(|item_decoder| Result::Ok(item_decoder.drain(num_items_to_drain)?.task)) + .transpose()? + }; + + self.rows_drained += num_rows; + Ok(NextDecodeTask { + num_rows, + task: Box::new(ListDecodeTask { + offsets, + validity, + items_field: self.items_field.clone(), + items: item_decode, + offset_type: self.offset_type.clone(), + }) as Box, + }) + } + + fn max_rows_to_drain(&self, num_rows: u64, byte_budget: u64) -> Result { + if num_rows == 0 { + return Ok(DrainLimit { rows: 0, bytes: 0 }); + } + let start_index = self.rows_drained as usize; + let end_index = start_index + num_rows as usize; + let start_offset = self.offsets[start_index]; + // The list's own offsets buffer counts items; when it is i32, the cumulative + // item count of one output array must stay within the budget. Offsets are + // rebased per output batch, so only the delta from `start_offset` counts. + let mut safe_rows = num_rows; + if self.offset_type != DataType::Int64 { + let max_offset = start_offset.saturating_add(byte_budget); + safe_rows = self.offsets[start_index + 1..=end_index] + .partition_point(|offset| *offset <= max_offset) as u64; + } + + // The item array has its own offset buffers with their own budget; shrink to + // the largest whole number of lists whose items also fit. + let mut item_bytes = 0u64; + let mut safe_end_index = start_index + safe_rows as usize; + let requested_items = self.offsets[safe_end_index] - start_offset; + if requested_items > 0 + && let Some(item_decoder) = &self.item_decoder + { + let item_limit = item_decoder.max_rows_to_drain(requested_items, byte_budget)?; + item_bytes = item_limit.bytes; + if item_limit.rows < requested_items { + let max_item_offset = start_offset + item_limit.rows; + safe_rows = self.offsets[start_index + 1..=safe_end_index] + .partition_point(|offset| *offset <= max_item_offset) + as u64; + safe_end_index = start_index + safe_rows as usize; + } + } + let items_consumed = self.offsets[safe_end_index] - start_offset; + let bytes = if self.offset_type != DataType::Int64 { + // One shared budget covers both nesting levels conservatively: charge + // whichever level consumes more. + items_consumed.max(item_bytes) + } else { + item_bytes + }; + Ok(DrainLimit { + rows: safe_rows, + bytes, + }) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } + + fn rows_loaded(&self) -> u64 { + self.rows_loaded + } + + fn rows_drained(&self) -> u64 { + self.rows_drained + } + + fn data_type(&self) -> &DataType { + &self.data_type + } +} + +struct IndirectlyLoaded { + offsets: Arc<[u64]>, + validity: BooleanBuffer, + root_decoder: Option, +} + +impl std::fmt::Debug for IndirectlyLoaded { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("IndirectlyLoaded") + .field("offsets", &self.offsets) + .field("validity", &self.validity) + .finish() + } +} + +/// An encoder for list offsets that "stitches" offsets and encodes nulls into the offsets +/// +/// If we need to encode several list arrays into a single page then we need to "stitch" the offsets +/// For example, imagine we have list arrays [[0, 1], [2]] and [[3, 4, 5]]. +/// +/// We will have offset arrays [0, 2, 3] and [0, 3]. We don't want to encode [0, 2, 3, 0, 3]. What +/// we want is [0, 2, 3, 6] +/// +/// This encoder also handles validity by converting a null value into an oversized offset. For example, +/// if we have four lists with offsets [0, 20, 20, 20, 30] and the list at index 2 is null (note that +/// the list at index 1 is empty) then we turn this into offsets [0, 20, 20, 51, 30]. We replace a null +/// offset with previous_offset + max_offset + 1. This makes it possible to load a single item from the +/// list array. +/// +/// These offsets are always stored on disk as a u64 array. First, this is because its simply much more +/// likely than one expects that this is needed, even if our lists are not massive. This is because we +/// only write an offsets page when we have enough data. This means we will probably accumulate a million +/// offsets or more before we bother to write a page. If our lists have a few thousand items a piece then +/// we end up passing the u32::MAX boundary. +/// +/// The second reason is that list offsets are very easily compacted with delta + bit packing and so those +/// u64 offsets should easily be shrunk down before being put on disk. +/// +/// This encoder can encode both lists and large lists. It can decode the resulting column into either type +/// as well. (TODO: Test and enable large lists) +/// +/// You can even write as a large list and decode as a regular list (as long as no single list has more than +/// 2^31 items) or vice versa. You could even encode a mixed stream of list and large list (but unclear that +/// would ever be useful) +#[derive(Debug)] +struct ListOffsetsEncoder { + // An accumulation queue, we insert both offset arrays and validity arrays into this queue + accumulation_queue: AccumulationQueue, + // The inner encoder of offset values + inner_encoder: Arc, + column_index: u32, +} + +impl ListOffsetsEncoder { + fn new( + cache_bytes: u64, + keep_original_array: bool, + column_index: u32, + inner_encoder: Arc, + ) -> Self { + Self { + accumulation_queue: AccumulationQueue::new( + cache_bytes, + column_index, + keep_original_array, + ), + inner_encoder, + column_index, + } + } + + /// Given a list array, return the offsets as a standalone ArrayRef (either an Int32Array or Int64Array) + fn extract_offsets(list_arr: &dyn Array) -> ArrayRef { + match list_arr.data_type() { + DataType::List(_) => { + let offsets = list_arr.as_list::().offsets().clone(); + Arc::new(Int32Array::new(offsets.into_inner(), None)) + } + DataType::LargeList(_) => { + let offsets = list_arr.as_list::().offsets().clone(); + Arc::new(Int64Array::new(offsets.into_inner(), None)) + } + _ => panic!(), + } + } + + /// Converts the validity of a list array into a boolean array. If there is no validity information + /// then this is an empty boolean array. + fn extract_validity(list_arr: &dyn Array) -> ArrayRef { + if let Some(validity) = list_arr.nulls() { + Arc::new(BooleanArray::new(validity.inner().clone(), None)) + } else { + // We convert None validity into an empty array because the accumulation queue can't + // handle Option + new_empty_array(&DataType::Boolean) + } + } + + fn make_encode_task(&self, arrays: Vec) -> EncodeTask { + let inner_encoder = self.inner_encoder.clone(); + let column_idx = self.column_index; + // At this point we should have 2*N arrays where the even-indexed arrays are integer offsets + // and the odd-indexed arrays are boolean validity bitmaps + let offset_arrays = arrays.iter().step_by(2).cloned().collect::>(); + let validity_arrays = arrays.into_iter().skip(1).step_by(2).collect::>(); + + tokio::task::spawn(async move { + let num_rows = + offset_arrays.iter().map(|arr| arr.len()).sum::() - offset_arrays.len(); + let num_rows = num_rows as u64; + let mut buffer_index = 0; + let array = Self::do_encode( + offset_arrays, + validity_arrays, + &mut buffer_index, + num_rows, + inner_encoder, + )?; + let (data, description) = array.into_buffers(); + Ok(EncodedPage { + data, + description: PageEncoding::Legacy(description), + num_rows, + column_idx, + row_number: 0, // V2.0 encoders do not use + }) + }) + .map(|res_res| res_res.unwrap()) + .boxed() + } + + fn maybe_encode_offsets_and_validity(&mut self, list_arr: &dyn Array) -> Option { + let offsets = Self::extract_offsets(list_arr); + let validity = Self::extract_validity(list_arr); + let num_rows = offsets.len() as u64; + // Either inserting the offsets OR inserting the validity could cause the + // accumulation queue to fill up + if let Some(mut arrays) = self + .accumulation_queue + .insert(offsets, /*row_number=*/ 0, num_rows) + { + arrays.0.push(validity); + Some(self.make_encode_task(arrays.0)) + } else if let Some(arrays) = self + .accumulation_queue + .insert(validity, /*row_number=*/ 0, num_rows) + { + Some(self.make_encode_task(arrays.0)) + } else { + None + } + } + + fn flush(&mut self) -> Option { + if let Some(arrays) = self.accumulation_queue.flush() { + Some(self.make_encode_task(arrays.0)) + } else { + None + } + } + + // Get's the total number of items covered by an array of offsets (keeping in + // mind that the first offset may not be zero) + fn get_offset_span(array: &dyn Array) -> u64 { + match array.data_type() { + DataType::Int32 => { + let arr_i32 = array.as_primitive::(); + (arr_i32.value(arr_i32.len() - 1) - arr_i32.value(0)) as u64 + } + DataType::Int64 => { + let arr_i64 = array.as_primitive::(); + (arr_i64.value(arr_i64.len() - 1) - arr_i64.value(0)) as u64 + } + _ => panic!(), + } + } + + // This is where we do the work to actually shift the offsets and encode nulls + // Note that the output is u64 and the input could be i32 OR i64. + fn extend_offsets_vec_u64( + dest: &mut Vec, + offsets: &dyn Array, + validity: Option<&BooleanArray>, + // The offset of this list into the destination + base: u64, + null_offset_adjustment: u64, + ) { + match offsets.data_type() { + DataType::Int32 => { + let offsets_i32 = offsets.as_primitive::(); + let start = offsets_i32.value(0) as u64; + // If we want to take a list from start..X and change it into + // a list from end..X then we need to add (base - start) to all elements + // Note that `modifier` may be negative but (item + modifier) will always be >= 0 + let modifier = base as i64 - start as i64; + if let Some(validity) = validity { + dest.extend( + offsets_i32 + .values() + .iter() + .skip(1) + .zip(validity.values().iter()) + .map(|(&off, valid)| { + (off as i64 + modifier) as u64 + + (!valid as u64 * null_offset_adjustment) + }), + ); + } else { + dest.extend( + offsets_i32 + .values() + .iter() + .skip(1) + // Subtract by `start` so offsets start at 0 + .map(|&v| (v as i64 + modifier) as u64), + ); + } + } + DataType::Int64 => { + let offsets_i64 = offsets.as_primitive::(); + let start = offsets_i64.value(0) as u64; + // If we want to take a list from start..X and change it into + // a list from end..X then we need to add (base - start) to all elements + // Note that `modifier` may be negative but (item + modifier) will always be >= 0 + let modifier = base as i64 - start as i64; + if let Some(validity) = validity { + dest.extend( + offsets_i64 + .values() + .iter() + .skip(1) + .zip(validity.values().iter()) + .map(|(&off, valid)| { + (off + modifier) as u64 + (!valid as u64 * null_offset_adjustment) + }), + ) + } else { + dest.extend( + offsets_i64 + .values() + .iter() + .skip(1) + .map(|&v| (v + modifier) as u64), + ); + } + } + _ => panic!("Invalid list offsets data type {:?}", offsets.data_type()), + } + } + + fn do_encode_u64( + offset_arrays: Vec, + validity: Vec>, + num_offsets: u64, + null_offset_adjustment: u64, + buffer_index: &mut u32, + inner_encoder: Arc, + ) -> Result { + let mut offsets = Vec::with_capacity(num_offsets as usize); + for (offsets_arr, validity_arr) in offset_arrays.iter().zip(validity) { + let last_prev_offset = offsets.last().copied().unwrap_or(0) % null_offset_adjustment; + Self::extend_offsets_vec_u64( + &mut offsets, + &offsets_arr, + validity_arr, + last_prev_offset, + null_offset_adjustment, + ); + } + let offsets_data = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(offsets), + num_values: num_offsets, + block_info: BlockInfo::new(), + }); + inner_encoder.encode(offsets_data, &DataType::UInt64, buffer_index) + } + + fn do_encode( + offset_arrays: Vec, + validity_arrays: Vec, + buffer_index: &mut u32, + num_offsets: u64, + inner_encoder: Arc, + ) -> Result { + let validity_arrays = validity_arrays + .iter() + .map(|v| { + if v.is_empty() { + None + } else { + Some(v.as_boolean()) + } + }) + .collect::>(); + debug_assert_eq!(offset_arrays.len(), validity_arrays.len()); + let total_span = offset_arrays + .iter() + .map(|arr| Self::get_offset_span(arr.as_ref())) + .sum::(); + // See encodings.proto for reasoning behind this value + let null_offset_adjustment = total_span + 1; + let encoded_offsets = Self::do_encode_u64( + offset_arrays, + validity_arrays, + num_offsets, + null_offset_adjustment, + buffer_index, + inner_encoder, + )?; + Ok(EncodedArray { + data: encoded_offsets.data, + encoding: pb::ArrayEncoding { + array_encoding: Some(pb::array_encoding::ArrayEncoding::List(Box::new( + pb::List { + offsets: Some(Box::new(encoded_offsets.encoding)), + null_offset_adjustment, + num_items: total_span, + }, + ))), + }, + }) + } +} + +pub struct ListFieldEncoder { + offsets_encoder: ListOffsetsEncoder, + items_encoder: Box, +} + +impl ListFieldEncoder { + pub fn new( + items_encoder: Box, + inner_offsets_encoder: Arc, + cache_bytes_per_columns: u64, + keep_original_array: bool, + column_index: u32, + ) -> Self { + Self { + offsets_encoder: ListOffsetsEncoder::new( + cache_bytes_per_columns, + keep_original_array, + column_index, + inner_offsets_encoder, + ), + items_encoder, + } + } + + fn combine_tasks( + offsets_tasks: Vec, + item_tasks: Vec, + ) -> Result> { + let mut all_tasks = offsets_tasks; + let item_tasks = item_tasks; + all_tasks.extend(item_tasks); + Ok(all_tasks) + } +} + +impl FieldEncoder for ListFieldEncoder { + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + // The list may have an offset / shorter length which means the underlying + // values array could be longer than what we need to encode and so we need + // to slice down to the region of interest. + let items = match array.data_type() { + DataType::List(_) => { + let list_arr = array.as_list::(); + let items_start = list_arr.value_offsets()[list_arr.offset()] as usize; + let items_end = + list_arr.value_offsets()[list_arr.offset() + list_arr.len()] as usize; + list_arr + .values() + .slice(items_start, items_end - items_start) + } + DataType::LargeList(_) => { + let list_arr = array.as_list::(); + let items_start = list_arr.value_offsets()[list_arr.offset()] as usize; + let items_end = + list_arr.value_offsets()[list_arr.offset() + list_arr.len()] as usize; + list_arr + .values() + .slice(items_start, items_end - items_start) + } + _ => panic!(), + }; + let offsets_tasks = self + .offsets_encoder + .maybe_encode_offsets_and_validity(array.as_ref()) + .map(|task| vec![task]) + .unwrap_or_default(); + let mut item_tasks = self.items_encoder.maybe_encode( + items, + external_buffers, + repdef, + row_number, + num_rows, + )?; + if !offsets_tasks.is_empty() && item_tasks.is_empty() { + // An items page cannot currently be shared by two different offsets pages. This is + // a limitation in the current scheduler and could be addressed in the future. As a result + // we always need to encode the items page if we encode the offsets page. + // + // In practice this isn't usually too bad unless we are targeting very small pages. + item_tasks = self.items_encoder.flush(external_buffers)?; + } + Self::combine_tasks(offsets_tasks, item_tasks) + } + + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + let offsets_tasks = self + .offsets_encoder + .flush() + .map(|task| vec![task]) + .unwrap_or_default(); + let item_tasks = self.items_encoder.flush(external_buffers)?; + Self::combine_tasks(offsets_tasks, item_tasks) + } + + fn num_columns(&self) -> u32 { + self.items_encoder.num_columns() + 1 + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + let inner_columns = self.items_encoder.finish(external_buffers); + async move { + let mut columns = vec![EncodedColumn::default()]; + let inner_columns = inner_columns.await?; + columns.extend(inner_columns); + Ok(columns) + } + .boxed() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_buffer::BooleanBuffer; + use arrow_schema::{DataType, Field}; + + use super::*; + use crate::decoder::LogicalPageDecoder; + + #[test] + fn oversized_binary_batch_error_is_actionable() { + let error = oversized_batch_error( + &Field::new("item", DataType::UInt8, false), + 32..160, + 96, + i32::MAX as u64 + 1, + ); + assert!( + error + .to_string() + .contains("more than 2GiB of string/binary data") + ); + assert!(error.to_string().contains("rows 32..96 fit")); + assert!(error.to_string().contains("requesting rows 32..160")); + assert!(error.to_string().contains("batch_size")); + assert!(error.to_string().contains("LANCE_DEFAULT_BATCH_SIZE")); + assert!(error.to_string().contains("large_string/large_binary")); + } + + #[test] + fn oversized_list_batch_error_is_actionable() { + let error = oversized_batch_error( + &Field::new("item", DataType::Int32, false), + 8..32, + 16, + i32::MAX as u64 + 1, + ); + let message = error.to_string(); + assert!(message.contains("list array")); + assert!(message.contains("more than i32::MAX items")); + assert!(message.contains("rows 8..16 fit")); + assert!(message.contains("requesting rows 8..32")); + } + + #[test] + fn list_decoder_overflow_reports_row_span() { + let mut decoder = ListPageDecoder { + unloaded: None, + offsets: Arc::<[u64]>::from(vec![0_u64, 1, 2, 3, 3 + i32::MAX as u64 + 1]), + validity: BooleanBuffer::from_iter([true, true, true, true]), + item_decoder: None, + num_rows: 4, + rows_drained: 2, + rows_loaded: 4, + items_field: Arc::new(Field::new("item", DataType::UInt8, false)), + offset_type: DataType::Int32, + data_type: DataType::List(Arc::new(Field::new("item", DataType::UInt8, false))), + }; + + let Err(error) = decoder.drain(2) else { + panic!("expected overflow error"); + }; + let message = error.to_string(); + assert!(message.contains("rows 2..3 fit")); + assert!(message.contains("requesting rows 2..4")); + assert!(message.contains(&(i32::MAX as u64 + 2).to_string())); + assert!(message.contains("batch_size")); + } + + fn list_page_decoder(offsets: Vec) -> ListPageDecoder { + let items_field = Arc::new(Field::new("item", DataType::Int32, true)); + ListPageDecoder { + unloaded: None, + validity: BooleanBuffer::new_set(offsets.len() - 1), + num_rows: offsets.len() as u64 - 1, + offsets: offsets.into(), + item_decoder: None, + rows_drained: 0, + rows_loaded: 0, + data_type: DataType::List(items_field.clone()), + items_field, + offset_type: DataType::Int32, + } + } + + const FULL_BUDGET: u64 = i32::MAX as u64; + + #[test] + fn test_list_page_limits_i32_offset_batches() { + let mut decoder = list_page_decoder(vec![0, i32::MAX as u64, i32::MAX as u64 + 1]); + let limit = decoder.max_rows_to_drain(2, FULL_BUDGET).unwrap(); + assert_eq!(limit.rows, 1); + assert_eq!(limit.bytes, i32::MAX as u64); + + // Offsets are rebased for each output batch, so the following row is safe. + decoder.rows_drained = 1; + let limit = decoder.max_rows_to_drain(1, FULL_BUDGET).unwrap(); + assert_eq!(limit.rows, 1); + assert_eq!(limit.bytes, 1); + } + + #[test] + fn test_list_page_stays_within_reduced_budget() { + // Mid-batch a page sees only the budget the batch has left; splitting is + // by actual consumption, not page boundaries. + let decoder = list_page_decoder(vec![0, 10, 20, 30, 40]); + let limit = decoder.max_rows_to_drain(4, 25).unwrap(); + assert_eq!(limit.rows, 2); + assert_eq!(limit.bytes, 20); + let limit = decoder.max_rows_to_drain(4, FULL_BUDGET).unwrap(); + assert_eq!(limit.rows, 4); + assert_eq!(limit.bytes, 40); + } + + #[test] + fn test_list_page_reports_zero_rows_for_single_row_over_i32_offset_limit() { + // The stream root emits a zero-row limit as a single-row batch, letting a + // genuinely oversized row fail at the arrow layer as it always did. + let decoder = list_page_decoder(vec![0, i32::MAX as u64 + 1]); + let limit = decoder.max_rows_to_drain(1, FULL_BUDGET).unwrap(); + assert_eq!(limit.rows, 0); + } +} diff --git a/vendor/lance-encoding/src/array_encoding/logical/primitive.rs b/vendor/lance-encoding/src/array_encoding/logical/primitive.rs new file mode 100644 index 000000000..66a6bbdca --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/logical/primitive.rs @@ -0,0 +1,589 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{fmt::Debug, ops::Range, sync::Arc, vec}; + +use arrow_array::{Array, ArrayRef, cast::AsArray, make_array}; +use arrow_buffer::bit_util; +use arrow_schema::DataType; +use futures::{FutureExt, future::BoxFuture}; +use log::trace; + +use crate::decoder::{ColumnBuffers, PageBuffers}; +use crate::decoder::{ + DrainLimit, FieldScheduler, I32_OFFSET_BYTE_BUDGET, LogicalPageDecoder, SchedulingJob, + has_i32_offsets, +}; +use crate::encoder::ArrayEncodingStrategy; +use crate::utils::accumulation::AccumulationQueue; +use crate::{array_encoding::physical::decoder_from_array_encoding, data::DataBlock}; +use lance_core::{Error, Result, datatypes::Field}; + +use crate::{ + decoder::{ + DecodeArrayTask, FilterExpression, MessageType, NextDecodeTask, PageEncoding, PageInfo, + PageScheduler, PrimitivePageDecoder, PriorityRange, ScheduledScanLine, SchedulerContext, + }, + encoder::{ + EncodeTask, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, OutOfLineBuffers, + }, + repdef::RepDefBuilder, +}; + +#[derive(Debug)] +struct PrimitivePage { + scheduler: Box, + num_rows: u64, + page_index: u32, +} + +/// A field scheduler for primitive fields +/// +/// This maps to exactly one column and it assumes that the top-level +/// encoding of each page is "basic". The basic encoding decodes into an +/// optional buffer of validity and a fixed-width buffer of values +/// which is exactly what we need to create a primitive array. +/// +/// Note: we consider booleans and fixed-size-lists of primitive types to be +/// primitive types. This is slightly different than arrow-rs's definition +#[derive(Debug)] +pub struct PrimitiveFieldScheduler { + data_type: DataType, + page_schedulers: Vec, + num_rows: u64, + should_validate: bool, + column_index: u32, +} + +impl PrimitiveFieldScheduler { + pub fn new( + column_index: u32, + data_type: DataType, + pages: Arc<[PageInfo]>, + buffers: ColumnBuffers, + should_validate: bool, + ) -> Self { + let page_schedulers = pages + .iter() + .enumerate() + // Buggy versions of Lance could sometimes create empty pages + .filter(|(page_index, page)| { + log::trace!("Skipping empty page with index {}", page_index); + page.num_rows > 0 + }) + .map(|(page_index, page)| { + let page_buffers = PageBuffers { + column_buffers: buffers, + positions_and_sizes: &page.buffer_offsets_and_sizes, + }; + let scheduler = decoder_from_array_encoding( + page.encoding.as_legacy(), + &page_buffers, + &data_type, + ); + PrimitivePage { + scheduler, + num_rows: page.num_rows, + page_index: page_index as u32, + } + }) + .collect::>(); + let num_rows = page_schedulers.iter().map(|p| p.num_rows).sum(); + Self { + data_type, + page_schedulers, + num_rows, + should_validate, + column_index, + } + } +} + +#[derive(Debug)] +struct PrimitiveFieldSchedulingJob<'a> { + scheduler: &'a PrimitiveFieldScheduler, + ranges: Vec>, + page_idx: usize, + range_idx: usize, + range_offset: u64, + global_row_offset: u64, +} + +impl<'a> PrimitiveFieldSchedulingJob<'a> { + pub fn new(scheduler: &'a PrimitiveFieldScheduler, ranges: Vec>) -> Self { + Self { + scheduler, + ranges, + page_idx: 0, + range_idx: 0, + range_offset: 0, + global_row_offset: 0, + } + } +} + +impl SchedulingJob for PrimitiveFieldSchedulingJob<'_> { + fn schedule_next( + &mut self, + context: &mut SchedulerContext, + priority: &dyn PriorityRange, + ) -> Result { + debug_assert!(self.range_idx < self.ranges.len()); + // Get our current range + let mut range = self.ranges[self.range_idx].clone(); + range.start += self.range_offset; + + let mut cur_page = &self.scheduler.page_schedulers[self.page_idx]; + trace!( + "Current range is {:?} and current page has {} rows", + range, cur_page.num_rows + ); + // Skip entire pages until we have some overlap with our next range + while cur_page.num_rows + self.global_row_offset <= range.start { + self.global_row_offset += cur_page.num_rows; + self.page_idx += 1; + trace!("Skipping entire page of {} rows", cur_page.num_rows); + cur_page = &self.scheduler.page_schedulers[self.page_idx]; + } + + // Now the cur_page has overlap with range. Continue looping through ranges + // until we find a range that exceeds the current page + + let mut ranges_in_page = Vec::new(); + while cur_page.num_rows + self.global_row_offset > range.start { + range.start = range.start.max(self.global_row_offset); + let start_in_page = range.start - self.global_row_offset; + let end_in_page = start_in_page + (range.end - range.start); + let end_in_page = end_in_page.min(cur_page.num_rows); + let last_in_range = (end_in_page + self.global_row_offset) >= range.end; + + ranges_in_page.push(start_in_page..end_in_page); + if last_in_range { + self.range_idx += 1; + if self.range_idx == self.ranges.len() { + break; + } + range = self.ranges[self.range_idx].clone(); + } else { + break; + } + } + + let num_rows_in_next = ranges_in_page.iter().map(|r| r.end - r.start).sum(); + trace!( + "Scheduling {} rows across {} ranges from page with {} rows (priority={}, column_index={}, page_index={})", + num_rows_in_next, + ranges_in_page.len(), + cur_page.num_rows, + priority.current_priority(), + self.scheduler.column_index, + cur_page.page_index, + ); + + self.global_row_offset += cur_page.num_rows; + self.page_idx += 1; + + let physical_decoder = cur_page.scheduler.schedule_ranges( + &ranges_in_page, + context.io(), + priority.current_priority(), + ); + + let logical_decoder = PrimitiveFieldDecoder { + data_type: self.scheduler.data_type.clone(), + column_index: self.scheduler.column_index, + unloaded_physical_decoder: Some(physical_decoder), + physical_decoder: None, + rows_drained: 0, + num_rows: num_rows_in_next, + should_validate: self.scheduler.should_validate, + page_index: cur_page.page_index, + }; + + let decoder = Box::new(logical_decoder); + #[allow(deprecated)] + let decoder_ready = context.locate_decoder(decoder); + Ok(ScheduledScanLine { + decoders: vec![MessageType::DecoderReady(decoder_ready)], + rows_scheduled: num_rows_in_next, + }) + } + + fn num_rows(&self) -> u64 { + self.ranges.iter().map(|r| r.end - r.start).sum() + } +} + +impl FieldScheduler for PrimitiveFieldScheduler { + fn num_rows(&self) -> u64 { + self.num_rows + } + + fn schedule_ranges<'a>( + &'a self, + ranges: &[std::ops::Range], + // TODO: Could potentially use filter to simplify decode, something of a micro-optimization probably + _filter: &FilterExpression, + ) -> Result> { + Ok(Box::new(PrimitiveFieldSchedulingJob::new( + self, + ranges.to_vec(), + ))) + } + + fn initialize<'a>( + &'a self, + _filter: &'a FilterExpression, + _context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + // 2.0 schedulers do not need to initialize + std::future::ready(Ok(())).boxed() + } +} + +pub struct PrimitiveFieldDecoder { + data_type: DataType, + unloaded_physical_decoder: Option>>>, + physical_decoder: Option>, + should_validate: bool, + num_rows: u64, + rows_drained: u64, + column_index: u32, + page_index: u32, +} + +impl PrimitiveFieldDecoder { + pub fn new_from_data( + physical_decoder: Arc, + data_type: DataType, + num_rows: u64, + should_validate: bool, + ) -> Self { + Self { + data_type, + unloaded_physical_decoder: None, + physical_decoder: Some(physical_decoder), + should_validate, + num_rows, + rows_drained: 0, + column_index: u32::MAX, + page_index: u32::MAX, + } + } +} + +impl Debug for PrimitiveFieldDecoder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PrimitiveFieldDecoder") + .field("data_type", &self.data_type) + .field("num_rows", &self.num_rows) + .field("rows_drained", &self.rows_drained) + .finish() + } +} + +struct PrimitiveFieldDecodeTask { + rows_to_skip: u64, + rows_to_take: u64, + should_validate: bool, + physical_decoder: Arc, + data_type: DataType, +} + +impl DecodeArrayTask for PrimitiveFieldDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + let block = self + .physical_decoder + .decode(self.rows_to_skip, self.rows_to_take)?; + + let array = make_array(block.into_arrow(self.data_type.clone(), self.should_validate)?); + + // This is a bit of a hack to work around https://github.com/apache/arrow-rs/issues/6302 + // + // We change from nulls-in-dictionary (storage format) to nulls-in-indices (arrow-rs preferred + // format) + // + // The calculation of logical_nulls is not free and would be good to avoid in the future + if let DataType::Dictionary(_, _) = self.data_type { + let dict = array.as_any_dictionary(); + if let Some(nulls) = array.logical_nulls() { + let new_indices = dict.keys().to_data(); + let new_array = make_array( + new_indices + .into_builder() + .nulls(Some(nulls)) + .add_child_data(dict.values().to_data()) + .data_type(dict.data_type().clone()) + .build()?, + ); + return Ok((new_array, 0)); + } + } + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array + // v2.0 path does not need it so we return 0. + Ok((array, 0)) + } +} + +impl LogicalPageDecoder for PrimitiveFieldDecoder { + // TODO: In the future, at some point, we may consider partially waiting for primitive pages by + // breaking up large I/O into smaller I/O as a way to accelerate the "time-to-first-decode" + fn wait_for_loaded(&mut self, loaded_need: u64) -> BoxFuture<'_, Result<()>> { + log::trace!( + "primitive wait for more than {} rows on column {} and page {} (page has {} rows)", + loaded_need, + self.column_index, + self.page_index, + self.num_rows + ); + async move { + let physical_decoder = self.unloaded_physical_decoder.take().unwrap().await?; + self.physical_decoder = Some(Arc::from(physical_decoder)); + Ok(()) + } + .boxed() + } + + fn drain(&mut self, num_rows: u64) -> Result { + if self.physical_decoder.as_ref().is_none() { + return Err(lance_core::Error::internal(format!( + "drain was called on primitive field decoder for data type {} on column {} but the decoder was never awaited", + self.data_type, self.column_index + ))); + } + + let rows_to_skip = self.rows_drained; + let rows_to_take = num_rows; + + self.rows_drained += rows_to_take; + + let task = Box::new(PrimitiveFieldDecodeTask { + rows_to_skip, + rows_to_take, + should_validate: self.should_validate, + physical_decoder: self.physical_decoder.as_ref().unwrap().clone(), + data_type: self.data_type.clone(), + }); + + Ok(NextDecodeTask { + task, + num_rows: rows_to_take, + }) + } + + fn max_rows_to_drain(&self, num_rows: u64, byte_budget: u64) -> Result { + if num_rows == 0 || !has_i32_offsets(&self.data_type) { + return Ok(DrainLimit { + rows: num_rows, + bytes: 0, + }); + } + let Some(physical_decoder) = self.physical_decoder.as_ref() else { + return Err(lance_core::Error::internal(format!( + "max_rows_to_drain was called on primitive field decoder for data type {} on column {} but the decoder was never awaited", + self.data_type, self.column_index + ))); + }; + match physical_decoder.variable_width_bytes(self.rows_drained, num_rows)? { + Some(bytes) if bytes <= byte_budget => Ok(DrainLimit { + rows: num_rows, + bytes, + }), + Some(_) => { + // Largest prefix of the request that still fits the budget. + // `variable_width_bytes` is monotonically non-decreasing in the row + // count; zero rows always fit. + let mut fits = 0u64; + let mut too_big = num_rows; + while fits + 1 < too_big { + let mid = fits + (too_big - fits) / 2; + let mid_bytes = physical_decoder + .variable_width_bytes(self.rows_drained, mid)? + .unwrap_or(0); + if mid_bytes <= byte_budget { + fits = mid; + } else { + too_big = mid; + } + } + let bytes = physical_decoder + .variable_width_bytes(self.rows_drained, fits)? + .unwrap_or(0); + Ok(DrainLimit { rows: fits, bytes }) + } + // The encoding cannot report sizes. Only accept this page's rows when + // the batch has accumulated no variable-width bytes yet (the budget is + // still whole), and charge the entire budget so nothing stacks after + // them; otherwise contribute nothing so the batch ends before them. + None if byte_budget >= I32_OFFSET_BYTE_BUDGET => Ok(DrainLimit { + rows: num_rows, + bytes: byte_budget, + }), + None => Ok(DrainLimit { rows: 0, bytes: 0 }), + } + } + + fn rows_loaded(&self) -> u64 { + if self.unloaded_physical_decoder.is_some() { + 0 + } else { + self.num_rows + } + } + + fn rows_drained(&self) -> u64 { + if self.unloaded_physical_decoder.is_some() { + 0 + } else { + self.rows_drained + } + } + + fn num_rows(&self) -> u64 { + self.num_rows + } + + fn data_type(&self) -> &DataType { + &self.data_type + } +} + +pub struct PrimitiveFieldEncoder { + accumulation_queue: AccumulationQueue, + array_encoding_strategy: Arc, + column_index: u32, + field: Field, + max_page_bytes: u64, +} + +impl PrimitiveFieldEncoder { + pub fn try_new( + options: &EncodingOptions, + array_encoding_strategy: Arc, + column_index: u32, + field: Field, + ) -> Result { + Ok(Self { + accumulation_queue: AccumulationQueue::new( + options.cache_bytes_per_column, + column_index, + options.keep_original_array, + ), + column_index, + max_page_bytes: options.max_page_bytes, + array_encoding_strategy, + field, + }) + } + + fn create_encode_task(&mut self, arrays: Vec) -> Result { + let encoder = self + .array_encoding_strategy + .create_array_encoder(&arrays, &self.field)?; + let column_idx = self.column_index; + let data_type = self.field.data_type(); + + Ok(tokio::task::spawn(async move { + let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); + let data = DataBlock::from_arrays(&arrays, num_values); + let mut buffer_index = 0; + let array = encoder.encode(data, &data_type, &mut buffer_index)?; + let (data, description) = array.into_buffers(); + Ok(EncodedPage { + data, + description: PageEncoding::Legacy(description), + num_rows: num_values, + column_idx, + row_number: 0, // v2.0 encoders do not use + }) + }) + .map(|res_res| { + res_res.unwrap_or_else(|err| { + Err(Error::internal(format!( + "Encoding task failed with error: {:?}", + err + ))) + }) + }) + .boxed()) + } + + // Creates an encode task, consuming all buffered data + fn do_flush(&mut self, arrays: Vec) -> Result> { + DataBlock::validate_arrays(&arrays, &self.field.name)?; + if arrays.len() == 1 { + let array = arrays.into_iter().next().unwrap(); + let size_bytes = array.get_buffer_memory_size(); + let num_parts = bit_util::ceil(size_bytes, self.max_page_bytes as usize); + // Can't slice it finer than 1 page per row + let num_parts = num_parts.min(array.len()); + if num_parts <= 1 { + // One part and it fits in a page + Ok(vec![self.create_encode_task(vec![array])?]) + } else { + // One part and it needs to be sliced into multiple pages + + // This isn't perfect (items in the array might not all have the same size) + // but it's a reasonable stab for now) + let mut tasks = Vec::with_capacity(num_parts); + let mut offset = 0; + let part_size = bit_util::ceil(array.len(), num_parts); + for _ in 0..num_parts { + let avail = array.len() - offset; + if avail == 0 { + break; + } + let chunk_size = avail.min(part_size); + let part = array.slice(offset, chunk_size); + let task = self.create_encode_task(vec![part])?; + tasks.push(task); + offset += chunk_size; + } + Ok(tasks) + } + } else { + // Multiple parts that (presumably) all fit in a page + // + // TODO: Could check here if there are any jumbo parts in the mix that need splitting + Ok(vec![self.create_encode_task(arrays)?]) + } + } +} + +impl FieldEncoder for PrimitiveFieldEncoder { + // Buffers data, if there is enough to write a page then we create an encode task + fn maybe_encode( + &mut self, + array: ArrayRef, + _external_buffers: &mut OutOfLineBuffers, + _repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + if let Some(arrays) = self.accumulation_queue.insert(array, row_number, num_rows) { + Ok(self.do_flush(arrays.0)?) + } else { + Ok(vec![]) + } + } + + // If there is any data left in the buffer then create an encode task from it + fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result> { + if let Some(arrays) = self.accumulation_queue.flush() { + Ok(self.do_flush(arrays.0)?) + } else { + Ok(vec![]) + } + } + + fn num_columns(&self) -> u32 { + 1 + } + + fn finish( + &mut self, + _external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + std::future::ready(Ok(vec![EncodedColumn::default()])).boxed() + } +} diff --git a/vendor/lance-encoding/src/array_encoding/logical/struct.rs b/vendor/lance-encoding/src/array_encoding/logical/struct.rs new file mode 100644 index 000000000..6153f1a09 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/logical/struct.rs @@ -0,0 +1,661 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ + collections::{BinaryHeap, VecDeque}, + ops::Range, + sync::Arc, +}; + +use crate::{ + decoder::{ + DecodeArrayTask, DrainLimit, FilterExpression, MessageType, NextDecodeTask, PriorityRange, + ScheduledScanLine, SchedulerContext, + }, + decoder::{DecoderReady, FieldScheduler, LogicalPageDecoder, SchedulingJob}, +}; +use arrow_array::{ArrayRef, StructArray}; +use arrow_schema::{DataType, Field, Fields}; +use futures::{FutureExt, StreamExt, TryStreamExt, future::BoxFuture, stream::FuturesUnordered}; +use lance_core::{Error, Result}; +use log::trace; + +#[derive(Debug)] +struct SchedulingJobWithStatus<'a> { + col_idx: u32, + col_name: &'a str, + job: Box, + rows_scheduled: u64, + rows_remaining: u64, +} + +impl PartialEq for SchedulingJobWithStatus<'_> { + fn eq(&self, other: &Self) -> bool { + self.col_idx == other.col_idx + } +} + +impl Eq for SchedulingJobWithStatus<'_> {} + +impl PartialOrd for SchedulingJobWithStatus<'_> { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for SchedulingJobWithStatus<'_> { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + // Note this is reversed to make it min-heap + other.rows_scheduled.cmp(&self.rows_scheduled) + } +} + +#[derive(Debug)] +struct EmptyStructDecodeTask { + num_rows: u64, +} + +impl DecodeArrayTask for EmptyStructDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array + // v2.0 path does not need it so we return 0. + Ok(( + Arc::new(StructArray::new_empty_fields(self.num_rows as usize, None)), + 0, + )) + } +} + +#[derive(Debug)] +struct EmptyStructDecoder { + num_rows: u64, + rows_drained: u64, + data_type: DataType, +} + +impl EmptyStructDecoder { + fn new(num_rows: u64) -> Self { + Self { + num_rows, + rows_drained: 0, + data_type: DataType::Struct(Fields::from(Vec::::default())), + } + } +} + +impl LogicalPageDecoder for EmptyStructDecoder { + fn wait_for_loaded(&mut self, _loaded_need: u64) -> BoxFuture<'_, Result<()>> { + Box::pin(std::future::ready(Ok(()))) + } + fn rows_loaded(&self) -> u64 { + self.num_rows + } + fn rows_unloaded(&self) -> u64 { + 0 + } + fn num_rows(&self) -> u64 { + self.num_rows + } + fn rows_drained(&self) -> u64 { + self.rows_drained + } + fn drain(&mut self, num_rows: u64) -> Result { + self.rows_drained += num_rows; + Ok(NextDecodeTask { + num_rows, + task: Box::new(EmptyStructDecodeTask { num_rows }), + }) + } + fn data_type(&self) -> &DataType { + &self.data_type + } +} + +#[derive(Debug)] +struct EmptyStructSchedulerJob { + num_rows: u64, +} + +impl SchedulingJob for EmptyStructSchedulerJob { + fn schedule_next( + &mut self, + context: &mut SchedulerContext, + _priority: &dyn PriorityRange, + ) -> Result { + let empty_decoder = Box::new(EmptyStructDecoder::new(self.num_rows)); + #[allow(deprecated)] + let struct_decoder = context.locate_decoder(empty_decoder); + Ok(ScheduledScanLine { + decoders: vec![MessageType::DecoderReady(struct_decoder)], + rows_scheduled: self.num_rows, + }) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +/// Scheduling job for struct data +/// +/// The order in which we schedule the children is important. We want to schedule the child +/// with the least amount of data first. +/// +/// This allows us to decode entire rows as quickly as possible +#[derive(Debug)] +struct SimpleStructSchedulerJob<'a> { + scheduler: &'a SimpleStructScheduler, + /// A min-heap whose key is the # of rows currently scheduled + children: BinaryHeap>, + rows_scheduled: u64, + num_rows: u64, + initialized: bool, +} + +impl<'a> SimpleStructSchedulerJob<'a> { + fn new( + scheduler: &'a SimpleStructScheduler, + children: Vec>, + num_rows: u64, + ) -> Self { + let children = children + .into_iter() + .enumerate() + .map(|(idx, job)| SchedulingJobWithStatus { + col_idx: idx as u32, + col_name: scheduler.child_fields[idx].name(), + job, + rows_scheduled: 0, + rows_remaining: num_rows, + }) + .collect::>(); + Self { + scheduler, + children, + rows_scheduled: 0, + num_rows, + initialized: false, + } + } +} + +impl SchedulingJob for SimpleStructSchedulerJob<'_> { + fn schedule_next( + &mut self, + mut context: &mut SchedulerContext, + priority: &dyn PriorityRange, + ) -> Result { + let mut decoders = Vec::new(); + if !self.initialized { + // Send info to the decoder thread so it knows a struct is here. In the future we will also + // send validity info here. + let struct_decoder = Box::new(SimpleStructDecoder::new( + self.scheduler.child_fields.clone(), + self.num_rows, + )); + #[allow(deprecated)] + let struct_decoder = context.locate_decoder(struct_decoder); + decoders.push(MessageType::DecoderReady(struct_decoder)); + self.initialized = true; + } + let old_rows_scheduled = self.rows_scheduled; + // Schedule as many children as we need to until we have scheduled at least one + // complete row + while old_rows_scheduled == self.rows_scheduled { + let mut next_child = self.children.pop().unwrap(); + trace!("Scheduling more rows for child {}", next_child.col_idx); + let scoped = context.push(next_child.col_name, next_child.col_idx); + let child_scan = next_child.job.schedule_next(scoped.context, priority)?; + trace!( + "Scheduled {} rows for child {}", + child_scan.rows_scheduled, next_child.col_idx + ); + next_child.rows_scheduled += child_scan.rows_scheduled; + next_child.rows_remaining -= child_scan.rows_scheduled; + decoders.extend(child_scan.decoders); + self.children.push(next_child); + self.rows_scheduled = self.children.peek().unwrap().rows_scheduled; + context = scoped.pop(); + } + let struct_rows_scheduled = self.rows_scheduled - old_rows_scheduled; + Ok(ScheduledScanLine { + decoders, + rows_scheduled: struct_rows_scheduled, + }) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +/// A scheduler for structs +/// +/// The implementation is actually a bit more tricky than one might initially think. We can't just +/// go through and schedule each column one after the other. This would mean our decode can't start +/// until nearly all the data has arrived (since we need data from each column) +/// +/// Instead, we schedule in row-major fashion +/// +/// Note: this scheduler is the starting point for all decoding. This is because we treat the top-level +/// record batch as a non-nullable struct. +#[derive(Debug)] +pub struct SimpleStructScheduler { + children: Vec>, + child_fields: Fields, + num_rows: u64, +} + +impl SimpleStructScheduler { + pub fn new( + children: Vec>, + child_fields: Fields, + num_rows: u64, + ) -> Self { + let num_rows = children + .first() + .map(|child| child.num_rows()) + .unwrap_or(num_rows); + debug_assert!(children.iter().all(|child| child.num_rows() == num_rows)); + Self { + children, + child_fields, + num_rows, + } + } +} + +impl FieldScheduler for SimpleStructScheduler { + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result> { + if self.children.is_empty() { + return Ok(Box::new(EmptyStructSchedulerJob { + num_rows: ranges.iter().map(|r| r.end - r.start).sum(), + })); + } + let child_schedulers = self + .children + .iter() + .map(|child| child.schedule_ranges(ranges, filter)) + .collect::>>()?; + let num_rows = child_schedulers[0].num_rows(); + Ok(Box::new(SimpleStructSchedulerJob::new( + self, + child_schedulers, + num_rows, + ))) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } + + fn initialize<'a>( + &'a self, + _filter: &'a FilterExpression, + _context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + let futures = self + .children + .iter() + .map(|child| child.initialize(_filter, _context)) + .collect::>(); + async move { + futures + .map(|res| res.map(|_| ())) + .try_collect::>() + .await?; + Ok(()) + } + .boxed() + } +} + +#[derive(Debug)] +struct ChildState { + // As child decoders are scheduled they are added to this queue + // Once the decoder is fully drained it is popped from this queue + // + // TODO: It may be a minor perf optimization, in some rare cases, if we have a separate + // "fully awaited but not yet drained" queue so we don't loop through fully awaited pages + // during each call to wait. + // + // Note: This queue may have more than one page in it if the batch size is very large + // or pages are very small + // TODO: Test this case + scheduled: VecDeque>, + // Rows that have been awaited + rows_loaded: u64, + // Rows that have drained + rows_drained: u64, + // Rows that have been popped (the decoder has been completely drained and removed from `scheduled`) + rows_popped: u64, + // Total number of rows in the struct + num_rows: u64, + // The field index in the struct (used for debugging / logging) + field_index: u32, +} + +impl ChildState { + fn new(num_rows: u64, field_index: u32) -> Self { + Self { + scheduled: VecDeque::new(), + rows_loaded: 0, + rows_drained: 0, + rows_popped: 0, + num_rows, + field_index, + } + } + + // Wait for the next set of rows to arrive + // + // Wait until we have at least `loaded_need` loaded and stop as soon as we + // go above that limit. + async fn wait_for_loaded(&mut self, loaded_need: u64) -> Result<()> { + trace!( + "Struct child {} waiting for more than {} rows to be loaded and {} are fully loaded already", + self.field_index, loaded_need, self.rows_loaded, + ); + let mut fully_loaded = self.rows_popped; + for (page_idx, next_decoder) in self.scheduled.iter_mut().enumerate() { + if next_decoder.rows_unloaded() > 0 { + let mut current_need = loaded_need; + current_need -= fully_loaded; + let rows_in_page = next_decoder.num_rows(); + let need_for_page = (rows_in_page - 1).min(current_need); + trace!( + "Struct child {} page {} will wait until more than {} rows loaded from page with {} rows", + self.field_index, page_idx, need_for_page, rows_in_page, + ); + // We might only await part of a page. This is important for things + // like the struct> case where we have one outer page, one + // middle page, and then a bunch of inner pages. If we await the entire + // middle page then we will have to wait for all the inner pages to arrive + // before we can start decoding. + next_decoder.wait_for_loaded(need_for_page).await?; + let now_loaded = next_decoder.rows_loaded(); + fully_loaded += now_loaded; + trace!( + "Struct child {} page {} await and now has {} loaded rows and we have {} fully loaded", + self.field_index, page_idx, now_loaded, fully_loaded + ); + } else { + fully_loaded += next_decoder.num_rows(); + } + if fully_loaded > loaded_need { + break; + } + } + self.rows_loaded = fully_loaded; + trace!( + "Struct child {} loaded {} new rows and now {} are loaded", + self.field_index, fully_loaded, self.rows_loaded + ); + Ok(()) + } + + fn drain(&mut self, num_rows: u64) -> Result { + trace!("Struct draining {} rows", num_rows); + + trace!( + "Draining {} rows from struct page with {} rows already drained", + num_rows, self.rows_drained + ); + let mut remaining = num_rows; + let mut composite = CompositeDecodeTask { + tasks: Vec::new(), + num_rows: 0, + has_more: true, + }; + while remaining > 0 { + let next = self.scheduled.front_mut().unwrap(); + let rows_to_take = remaining.min(next.rows_left()); + let next_task = next.drain(rows_to_take)?; + if next.rows_left() == 0 { + trace!("Completely drained page"); + self.rows_popped += next.num_rows(); + self.scheduled.pop_front(); + } + remaining -= rows_to_take; + composite.tasks.push(next_task.task); + composite.num_rows += next_task.num_rows; + } + self.rows_drained += num_rows; + composite.has_more = self.rows_drained != self.num_rows; + Ok(composite) + } + + fn max_rows_to_drain(&self, num_rows: u64, byte_budget: u64) -> Result { + let mut safe_rows = 0; + let mut bytes = 0u64; + let mut budget = byte_budget; + let mut remaining = num_rows; + for decoder in &self.scheduled { + if remaining == 0 { + break; + } + let rows_in_page = remaining.min(decoder.rows_left()); + let page_limit = decoder.max_rows_to_drain(rows_in_page, budget)?; + safe_rows += page_limit.rows; + bytes = bytes.saturating_add(page_limit.bytes); + if page_limit.rows < rows_in_page { + break; + } + budget = budget.saturating_sub(page_limit.bytes); + remaining -= rows_in_page; + } + Ok(DrainLimit { + rows: safe_rows, + bytes, + }) + } +} + +// Wrapper around ChildState that orders using rows_unawaited +struct WaitOrder<'a>(&'a mut ChildState); + +impl Eq for WaitOrder<'_> {} +impl PartialEq for WaitOrder<'_> { + fn eq(&self, other: &Self) -> bool { + self.0.rows_loaded == other.0.rows_loaded + } +} +impl Ord for WaitOrder<'_> { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + // Note: this is inverted so we have a min-heap + other.0.rows_loaded.cmp(&self.0.rows_loaded) + } +} +impl PartialOrd for WaitOrder<'_> { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +#[derive(Debug)] +pub struct SimpleStructDecoder { + children: Vec, + child_fields: Fields, + data_type: DataType, + num_rows: u64, +} + +impl SimpleStructDecoder { + pub fn new(child_fields: Fields, num_rows: u64) -> Self { + let data_type = DataType::Struct(child_fields.clone()); + Self { + children: child_fields + .iter() + .enumerate() + .map(|(idx, _)| ChildState::new(num_rows, idx as u32)) + .collect(), + child_fields, + data_type, + num_rows, + } + } + + async fn do_wait_for_loaded(&mut self, loaded_need: u64) -> Result<()> { + let mut wait_orders = self + .children + .iter_mut() + .filter_map(|child| { + if child.rows_loaded <= loaded_need { + Some(WaitOrder(child)) + } else { + None + } + }) + .collect::>(); + while !wait_orders.is_empty() { + let next_waiter = wait_orders.pop().unwrap(); + let next_highest = wait_orders + .peek() + .map(|w| w.0.rows_loaded) + .unwrap_or(u64::MAX); + // Wait until you have the number of rows needed, or at least more than the + // next highest waiter + let limit = loaded_need.min(next_highest); + next_waiter.0.wait_for_loaded(limit).await?; + log::trace!( + "Struct child {} finished await pass and now {} are loaded", + next_waiter.0.field_index, + next_waiter.0.rows_loaded + ); + if next_waiter.0.rows_loaded <= loaded_need { + wait_orders.push(next_waiter); + } + } + Ok(()) + } +} + +impl LogicalPageDecoder for SimpleStructDecoder { + fn accept_child(&mut self, mut child: DecoderReady) -> Result<()> { + // children with empty path should not be delivered to this method + let child_idx = child.path.pop_front().unwrap(); + if child.path.is_empty() { + // This decoder is intended for us + self.children[child_idx as usize] + .scheduled + .push_back(child.decoder); + } else { + // This decoder is intended for one of our children + let intended = self.children[child_idx as usize].scheduled.back_mut().ok_or_else(|| Error::internal(format!("Decoder scheduled for child at index {} but we don't have any child at that index yet", child_idx)))?; + intended.accept_child(child)?; + } + Ok(()) + } + + fn wait_for_loaded(&mut self, loaded_need: u64) -> BoxFuture<'_, Result<()>> { + self.do_wait_for_loaded(loaded_need).boxed() + } + + fn drain(&mut self, num_rows: u64) -> Result { + let child_tasks = self + .children + .iter_mut() + .map(|child| child.drain(num_rows)) + .collect::>>()?; + let num_rows = child_tasks[0].num_rows; + debug_assert!(child_tasks.iter().all(|task| task.num_rows == num_rows)); + Ok(NextDecodeTask { + task: Box::new(SimpleStructDecodeTask { + children: child_tasks, + child_fields: self.child_fields.clone(), + }), + num_rows, + }) + } + + fn max_rows_to_drain(&self, num_rows: u64, byte_budget: u64) -> Result { + // Each child column concatenates into its own Arrow array, so every child + // gets the full budget; the most constrained child bounds the batch. + self.children.iter().try_fold( + DrainLimit { + rows: num_rows, + bytes: 0, + }, + |limit, child| { + child + .max_rows_to_drain(num_rows, byte_budget) + .map(|child_limit| DrainLimit { + rows: limit.rows.min(child_limit.rows), + bytes: limit.bytes.max(child_limit.bytes), + }) + }, + ) + } + + fn rows_loaded(&self) -> u64 { + self.children.iter().map(|c| c.rows_loaded).min().unwrap() + } + + fn rows_drained(&self) -> u64 { + // All children should have the same number of rows drained + debug_assert!( + self.children + .iter() + .all(|c| c.rows_drained == self.children[0].rows_drained) + ); + self.children[0].rows_drained + } + + fn num_rows(&self) -> u64 { + self.num_rows + } + + fn data_type(&self) -> &DataType { + &self.data_type + } +} + +struct CompositeDecodeTask { + // One per child + tasks: Vec>, + num_rows: u64, + has_more: bool, +} + +impl CompositeDecodeTask { + fn decode(self) -> Result { + let arrays = self + .tasks + .into_iter() + .map(|task| task.decode().map(|(arr, _)| arr)) + .collect::>>()?; + let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::>(); + // TODO: If this is a primitive column we should be able to avoid this + // allocation + copy with "page bridging" which could save us a few CPU + // cycles. + // + // This optimization is probably most important for super fast storage like NVME + // where the page size can be smaller. + Ok(arrow_select::concat::concat(&array_refs)?) + } +} + +struct SimpleStructDecodeTask { + children: Vec, + child_fields: Fields, +} + +impl DecodeArrayTask for SimpleStructDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + let child_arrays = self + .children + .into_iter() + .map(|child| child.decode()) + .collect::>>()?; + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array + // v2.0 path does not need it so we return 0. + Ok(( + Arc::new(StructArray::try_new(self.child_fields, child_arrays, None)?), + 0, + )) + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical.rs b/vendor/lance-encoding/src/array_encoding/physical.rs new file mode 100644 index 000000000..66af9e806 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical.rs @@ -0,0 +1,329 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow_schema::DataType; +use lance_arrow::DataTypeExt; + +use crate::{ + array_encoding::physical::{ + basic::BasicPageScheduler, binary::BinaryPageScheduler, bitmap::DenseBitmapScheduler, + dictionary::DictionaryPageScheduler, fixed_size_list::FixedListScheduler, + fsst::FsstPageScheduler, packed_struct::PackedStructPageScheduler, + value::ValuePageScheduler, + }, + buffer::LanceBuffer, + decoder::{PageBuffers, PageScheduler}, + encodings::physical::block::{CompressionConfig, CompressionScheme}, + format::pb::{self, PackedStruct}, +}; + +pub mod basic; +pub mod binary; +pub mod bitmap; +#[cfg(feature = "bitpacking")] +pub mod bitpack; +pub mod block; +pub mod dictionary; +pub mod fixed_size_binary; +pub mod fixed_size_list; +pub mod fsst; +pub mod packed_struct; +pub mod value; + +// Translate a protobuf buffer description into a position in the file. This could be a page +// buffer, a column buffer, or a file buffer. +fn get_buffer(buffer_desc: &pb::Buffer, buffers: &PageBuffers) -> (u64, u64) { + let index = buffer_desc.buffer_index as usize; + + match pb::buffer::BufferType::try_from(buffer_desc.buffer_type).unwrap() { + pb::buffer::BufferType::Page => buffers.positions_and_sizes[index], + pb::buffer::BufferType::Column => buffers.column_buffers.positions_and_sizes[index], + pb::buffer::BufferType::File => { + buffers.column_buffers.file_buffers.positions_and_sizes[index] + } + } +} + +/// Convert a protobuf buffer encoding into a physical page scheduler +fn get_buffer_decoder(encoding: &pb::Flat, buffers: &PageBuffers) -> Box { + let (buffer_offset, buffer_size) = get_buffer(encoding.buffer.as_ref().unwrap(), buffers); + let compression_config: CompressionConfig = match encoding.compression.as_ref() { + None => CompressionConfig::new(CompressionScheme::None, None), + Some(compression) => CompressionConfig::new( + compression.scheme.as_str().parse().unwrap(), + compression.level, + ), + }; + match encoding.bits_per_value { + 1 => Box::new(DenseBitmapScheduler::new(buffer_offset)), + bits_per_value => { + if bits_per_value % 8 != 0 { + todo!( + "bits_per_value ({}) that is not a multiple of 8", + bits_per_value + ); + } + Box::new(ValuePageScheduler::new( + bits_per_value / 8, + buffer_offset, + buffer_size, + compression_config, + )) + } + } +} + +#[cfg(feature = "bitpacking")] +fn get_bitpacked_buffer_decoder( + encoding: &pb::Bitpacked, + buffers: &PageBuffers, +) -> Box { + let (buffer_offset, _buffer_size) = get_buffer(encoding.buffer.as_ref().unwrap(), buffers); + + Box::new(bitpack::BitpackedScheduler::new( + encoding.compressed_bits_per_value, + encoding.uncompressed_bits_per_value, + buffer_offset, + encoding.signed, + )) +} + +#[cfg(feature = "bitpacking")] +fn get_bitpacked_for_non_neg_buffer_decoder( + encoding: &pb::BitpackedForNonNeg, + buffers: &PageBuffers, +) -> Box { + let (buffer_offset, _buffer_size) = get_buffer(encoding.buffer.as_ref().unwrap(), buffers); + + Box::new(bitpack::BitpackedForNonNegScheduler::new( + encoding.compressed_bits_per_value, + encoding.uncompressed_bits_per_value, + buffer_offset, + )) +} + +fn decoder_from_packed_struct( + packed_struct: &PackedStruct, + buffers: &PageBuffers, + data_type: &DataType, +) -> Box { + let inner_encodings = &packed_struct.inner; + let fields = match data_type { + DataType::Struct(fields) => Some(fields), + _ => None, + } + .unwrap(); + + let inner_datatypes = fields + .iter() + .map(|field| field.data_type()) + .collect::>(); + + let mut inner_schedulers = Vec::with_capacity(fields.len()); + for i in 0..fields.len() { + let inner_encoding = &inner_encodings[i]; + let inner_datatype = inner_datatypes[i]; + let inner_scheduler = decoder_from_array_encoding(inner_encoding, buffers, inner_datatype); + inner_schedulers.push(inner_scheduler); + } + + let packed_buffer = packed_struct.buffer.as_ref().unwrap(); + let (buffer_offset, _) = get_buffer(packed_buffer, buffers); + + Box::new(PackedStructPageScheduler::new( + inner_schedulers, + data_type.clone(), + buffer_offset, + )) +} + +/// Convert a protobuf array encoding into a physical page scheduler +pub fn decoder_from_array_encoding( + encoding: &pb::ArrayEncoding, + buffers: &PageBuffers, + data_type: &DataType, +) -> Box { + match encoding.array_encoding.as_ref().unwrap() { + pb::array_encoding::ArrayEncoding::Nullable(basic) => { + match basic.nullability.as_ref().unwrap() { + pb::nullable::Nullability::NoNulls(no_nulls) => Box::new( + BasicPageScheduler::new_non_nullable(decoder_from_array_encoding( + no_nulls.values.as_ref().unwrap(), + buffers, + data_type, + )), + ), + pb::nullable::Nullability::SomeNulls(some_nulls) => { + Box::new(BasicPageScheduler::new_nullable( + decoder_from_array_encoding( + some_nulls.validity.as_ref().unwrap(), + buffers, + data_type, + ), + decoder_from_array_encoding( + some_nulls.values.as_ref().unwrap(), + buffers, + data_type, + ), + )) + } + pb::nullable::Nullability::AllNulls(_) => { + Box::new(BasicPageScheduler::new_all_null()) + } + } + } + #[cfg(feature = "bitpacking")] + pb::array_encoding::ArrayEncoding::Bitpacked(bitpacked) => { + get_bitpacked_buffer_decoder(bitpacked, buffers) + } + #[cfg(not(feature = "bitpacking"))] + pb::array_encoding::ArrayEncoding::Bitpacked(_) => { + panic!("Runtime built without bitpacking support") + } + pb::array_encoding::ArrayEncoding::Flat(flat) => get_buffer_decoder(flat, buffers), + pb::array_encoding::ArrayEncoding::FixedSizeList(fixed_size_list) => { + let item_encoding = fixed_size_list.items.as_ref().unwrap(); + let item_scheduler = decoder_from_array_encoding(item_encoding, buffers, data_type); + Box::new(FixedListScheduler::new( + item_scheduler, + fixed_size_list.dimension, + )) + } + // This is a column containing the list offsets. This wrapper is superfluous at the moment + // since we know it is a list based on the schema. In the future there may be different ways + // of storing the list offsets. + pb::array_encoding::ArrayEncoding::List(list) => { + decoder_from_array_encoding(list.offsets.as_ref().unwrap(), buffers, data_type) + } + pb::array_encoding::ArrayEncoding::Binary(binary) => { + let indices_encoding = binary.indices.as_ref().unwrap(); + let bytes_encoding = binary.bytes.as_ref().unwrap(); + + let indices_scheduler = + decoder_from_array_encoding(indices_encoding, buffers, data_type); + let bytes_scheduler = decoder_from_array_encoding(bytes_encoding, buffers, data_type); + + let offset_type = match data_type { + DataType::LargeBinary | DataType::LargeUtf8 => DataType::Int64, + _ => DataType::Int32, + }; + + Box::new(BinaryPageScheduler::new( + indices_scheduler.into(), + bytes_scheduler.into(), + offset_type, + binary.null_adjustment, + )) + } + pb::array_encoding::ArrayEncoding::Fsst(fsst) => { + let inner = + decoder_from_array_encoding(fsst.binary.as_ref().unwrap(), buffers, data_type); + + Box::new(FsstPageScheduler::new( + inner, + LanceBuffer::from_bytes(fsst.symbol_table.clone(), 1), + )) + } + pb::array_encoding::ArrayEncoding::Dictionary(dictionary) => { + let indices_encoding = dictionary.indices.as_ref().unwrap(); + let items_encoding = dictionary.items.as_ref().unwrap(); + let num_dictionary_items = dictionary.num_dictionary_items; + + // We can get here in 2 ways. The data is dictionary encoded and the user wants a dictionary or + // the data is dictionary encoded, as an optimization, and the user wants the value type. Figure + // out the value type. + let value_type = if let DataType::Dictionary(_, value_type) = data_type { + value_type + } else { + data_type + }; + + // Note: we don't actually know the indices type here, passing down `data_type` works ok because + // the dictionary indices are always integers and we don't need the data_type to figure out how + // to decode integers. + let indices_scheduler = + decoder_from_array_encoding(indices_encoding, buffers, data_type); + + let items_scheduler = decoder_from_array_encoding(items_encoding, buffers, value_type); + + let should_decode_dict = !data_type.is_dictionary(); + + Box::new(DictionaryPageScheduler::new( + indices_scheduler.into(), + items_scheduler.into(), + num_dictionary_items, + should_decode_dict, + )) + } + pb::array_encoding::ArrayEncoding::FixedSizeBinary(fixed_size_binary) => { + let bytes_encoding = fixed_size_binary.bytes.as_ref().unwrap(); + let bytes_scheduler = decoder_from_array_encoding(bytes_encoding, buffers, data_type); + let bytes_per_offset = match data_type { + DataType::LargeBinary | DataType::LargeUtf8 => 8, + DataType::Binary | DataType::Utf8 => 4, + _ => panic!("FixedSizeBinary only supports binary and utf8 types"), + }; + + Box::new(fixed_size_binary::FixedSizeBinaryPageScheduler::new( + bytes_scheduler, + fixed_size_binary.byte_width, + bytes_per_offset, + )) + } + pb::array_encoding::ArrayEncoding::PackedStruct(packed_struct) => { + decoder_from_packed_struct(packed_struct, buffers, data_type) + } + #[cfg(feature = "bitpacking")] + pb::array_encoding::ArrayEncoding::BitpackedForNonNeg(bitpacked) => { + get_bitpacked_for_non_neg_buffer_decoder(bitpacked, buffers) + } + #[cfg(not(feature = "bitpacking"))] + pb::array_encoding::ArrayEncoding::BitpackedForNonNeg(_) => { + panic!("Runtime built without bitpacking support") + } + // Currently there is no way to encode struct nullability and structs are encoded with a "header" column + // (that has no data). We never actually decode that column and so this branch is never actually encountered. + // + // This will change in the future when we add support for struct nullability. + pb::array_encoding::ArrayEncoding::Struct(_) => unreachable!(), + // 2.1 only + _ => unreachable!("Unsupported array encoding: {:?}", encoding), + } +} + +#[cfg(test)] +mod tests { + use crate::array_encoding::physical::get_buffer_decoder; + use crate::decoder::{ColumnBuffers, FileBuffers, PageBuffers}; + use crate::format::pb; + + #[test] + fn test_get_buffer_decoder_for_compressed_buffer() { + let page_scheduler = get_buffer_decoder( + &pb::Flat { + buffer: Some(pb::Buffer { + buffer_index: 0, + buffer_type: pb::buffer::BufferType::File as i32, + }), + bits_per_value: 8, + compression: Some(pb::Compression { + scheme: "zstd".to_string(), + level: Some(0), + }), + }, + &PageBuffers { + column_buffers: ColumnBuffers { + file_buffers: FileBuffers { + positions_and_sizes: &[(0, 100)], + }, + positions_and_sizes: &[], + }, + positions_and_sizes: &[], + }, + ); + assert_eq!( + format!("{:?}", page_scheduler).as_str(), + "ValuePageScheduler { bytes_per_value: 1, buffer_offset: 0, buffer_size: 100, compression_config: CompressionConfig { scheme: Zstd, level: Some(0) } }" + ); + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/basic.rs b/vendor/lance-encoding/src/array_encoding/physical/basic.rs new file mode 100644 index 000000000..fa4c58247 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/basic.rs @@ -0,0 +1,246 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_schema::DataType; +use futures::{FutureExt, future::BoxFuture}; +use log::trace; + +use crate::{ + EncodingsIo, + data::{AllNullDataBlock, BlockInfo, DataBlock, NullableDataBlock}, + decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, + format::ProtobufUtils, +}; + +use lance_core::Result; + +struct DataDecoders { + validity: Box, + values: Box, +} + +enum DataNullStatus { + // Neither validity nor values + All, + // Values only + None(Box), + // Validity and values + Some(DataDecoders), +} + +#[derive(Debug)] +struct DataSchedulers { + validity: Box, + values: Box, +} + +#[derive(Debug)] +enum SchedulerNullStatus { + // Values only + None(Box), + // Validity and values + Some(DataSchedulers), + // Neither validity nor values + All, +} + +impl SchedulerNullStatus { + fn values_scheduler(&self) -> Option<&dyn PageScheduler> { + match self { + Self::All => None, + Self::None(values) => Some(values.as_ref()), + Self::Some(schedulers) => Some(schedulers.values.as_ref()), + } + } +} + +/// A physical scheduler for "basic" fields. These are fields that have an optional +/// validity bitmap and some kind of values buffer. +/// +/// No actual decoding happens here, we are simply aggregating the two buffers. +/// +/// If everything is null then there are no data buffers at all. +// TODO: Add support/tests for primitive nulls +// TODO: Add tests for the all-null case +// +// Right now this is always present on primitive fields. In the future we may use a +// sentinel encoding instead. +#[derive(Debug)] +pub struct BasicPageScheduler { + mode: SchedulerNullStatus, +} + +impl BasicPageScheduler { + /// Creates a new instance that expects a validity bitmap + pub fn new_nullable( + validity_decoder: Box, + values_decoder: Box, + ) -> Self { + Self { + mode: SchedulerNullStatus::Some(DataSchedulers { + validity: validity_decoder, + values: values_decoder, + }), + } + } + + /// Create a new instance that does not need a validity bitmap because no item is null + pub fn new_non_nullable(values_decoder: Box) -> Self { + Self { + mode: SchedulerNullStatus::None(values_decoder), + } + } + + /// Create a new instance where all values are null + /// + /// It may seem strange we need `values_decoder` here but Arrow requires that value + /// buffers still be allocated / sized even if everything is null. So we need the value + /// decoder to calculate the capacity of the garbage buffer. + pub fn new_all_null() -> Self { + Self { + mode: SchedulerNullStatus::All, + } + } +} + +impl PageScheduler for BasicPageScheduler { + fn schedule_ranges( + &self, + ranges: &[std::ops::Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + let validity_future = match &self.mode { + SchedulerNullStatus::None(_) | SchedulerNullStatus::All => None, + SchedulerNullStatus::Some(schedulers) => Some(schedulers.validity.schedule_ranges( + ranges, + scheduler, + top_level_row, + )), + }; + + let values_future = if let Some(values_scheduler) = self.mode.values_scheduler() { + Some( + values_scheduler + .schedule_ranges(ranges, scheduler, top_level_row) + .boxed(), + ) + } else { + trace!("No values fetch needed since values all null"); + None + }; + + async move { + let mode = match (values_future, validity_future) { + (None, None) => DataNullStatus::All, + (Some(values_future), None) => DataNullStatus::None(values_future.await?), + (Some(values_future), Some(validity_future)) => { + DataNullStatus::Some(DataDecoders { + values: values_future.await?, + validity: validity_future.await?, + }) + } + _ => unreachable!(), + }; + Ok(Box::new(BasicPageDecoder { mode }) as Box) + } + .boxed() + } +} + +struct BasicPageDecoder { + mode: DataNullStatus, +} + +impl PrimitivePageDecoder for BasicPageDecoder { + fn variable_width_bytes(&self, rows_to_skip: u64, num_rows: u64) -> Result> { + match &self.mode { + DataNullStatus::Some(decoders) => { + decoders.values.variable_width_bytes(rows_to_skip, num_rows) + } + DataNullStatus::All => Ok(Some(0)), + DataNullStatus::None(values) => values.variable_width_bytes(rows_to_skip, num_rows), + } + } + + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + match &self.mode { + DataNullStatus::Some(decoders) => { + let validity = decoders.validity.decode(rows_to_skip, num_rows)?; + let validity = validity.as_fixed_width().unwrap(); + let values = decoders.values.decode(rows_to_skip, num_rows)?; + Ok(DataBlock::Nullable(NullableDataBlock { + data: Box::new(values), + nulls: validity.data, + block_info: BlockInfo::new(), + })) + } + DataNullStatus::All => Ok(DataBlock::AllNull(AllNullDataBlock { + num_values: num_rows, + })), + DataNullStatus::None(values) => values.decode(rows_to_skip, num_rows), + } + } +} + +#[derive(Debug)] +pub struct BasicEncoder { + values_encoder: Box, +} + +impl BasicEncoder { + pub fn new(values_encoder: Box) -> Self { + Self { values_encoder } + } +} + +impl ArrayEncoder for BasicEncoder { + fn encode( + &self, + data: DataBlock, + data_type: &DataType, + buffer_index: &mut u32, + ) -> Result { + match data { + DataBlock::AllNull(_) => { + let encoding = ProtobufUtils::basic_all_null_encoding(); + Ok(EncodedArray { data, encoding }) + } + DataBlock::Nullable(nullable) => { + let validity_buffer_index = *buffer_index; + *buffer_index += 1; + + let validity_desc = ProtobufUtils::flat_encoding( + 1, + validity_buffer_index, + /*compression=*/ None, + ); + let encoded_values = + self.values_encoder + .encode(*nullable.data, data_type, buffer_index)?; + let encoding = + ProtobufUtils::basic_some_null_encoding(validity_desc, encoded_values.encoding); + let encoded = DataBlock::Nullable(NullableDataBlock { + data: Box::new(encoded_values.data), + nulls: nullable.nulls, + block_info: BlockInfo::new(), + }); + Ok(EncodedArray { + data: encoded, + encoding, + }) + } + _ => { + let encoded_values = self.values_encoder.encode(data, data_type, buffer_index)?; + let encoding = ProtobufUtils::basic_no_null_encoding(encoded_values.encoding); + Ok(EncodedArray { + data: encoded_values.data, + encoding, + }) + } + } + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/binary.rs b/vendor/lance-encoding/src/array_encoding/physical/binary.rs new file mode 100644 index 000000000..f614b1e2e --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/binary.rs @@ -0,0 +1,873 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use core::panic; +use std::sync::Arc; + +use arrow_array::ArrayRef; +use arrow_array::cast::AsArray; +use arrow_array::types::UInt64Type; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer, bit_util}; +use futures::TryFutureExt; + +use futures::{FutureExt, future::BoxFuture}; + +use crate::array_encoding::logical::primitive::PrimitiveFieldDecoder; +use crate::buffer::LanceBuffer; +use crate::data::{ + BlockInfo, DataBlock, FixedWidthDataBlock, NullableDataBlock, VariableWidthBlock, +}; +use crate::decoder::LogicalPageDecoder; +use crate::encoder::{ArrayEncoder, EncodedArray}; +use crate::encodings::physical::block::{ + BufferCompressor, CompressionConfig, GeneralBufferCompressor, +}; +use crate::format::ProtobufUtils; +use crate::{ + EncodingsIo, + decoder::{PageScheduler, PrimitivePageDecoder}, +}; + +use arrow_array::{PrimitiveArray, UInt64Array}; +use arrow_schema::DataType; +use lance_core::{Error, Result}; + +fn oversized_binary_batch_error(num_rows: u64, num_bytes: u64) -> Error { + Error::not_supported(format!( + "Could not create array with more than 2GiB of string/binary data in a single batch \ + ({} rows would require {} bytes). Please reduce the batch_size, set \ + LANCE_DEFAULT_BATCH_SIZE to a smaller value, or convert the column to \ + large_string/large_binary.", + num_rows, num_bytes + )) +} + +fn oversized_large_binary_batch_error(num_rows: u64, num_bytes: u64) -> Error { + Error::not_supported(format!( + "Could not create large_string/large_binary array in a single batch because {} rows \ + would require {} bytes, which exceeds i64::MAX. Please reduce the batch_size or set \ + LANCE_DEFAULT_BATCH_SIZE to a smaller value.", + num_rows, num_bytes + )) +} + +struct IndicesNormalizer { + indices: Vec, + validity: BooleanBufferBuilder, + null_adjustment: u64, +} + +impl IndicesNormalizer { + fn new(num_rows: u64, null_adjustment: u64) -> Self { + let mut indices = Vec::with_capacity(num_rows as usize); + indices.push(0); + Self { + indices, + validity: BooleanBufferBuilder::new(num_rows as usize), + null_adjustment, + } + } + + fn normalize(&self, val: u64) -> (bool, u64) { + if val >= self.null_adjustment { + (false, val - self.null_adjustment) + } else { + (true, val) + } + } + + fn extend(&mut self, new_indices: &PrimitiveArray, is_start: bool) -> Result<()> { + let mut last = *self.indices.last().unwrap(); + if is_start { + let (is_valid, val) = self.normalize(new_indices.value(0)); + self.indices.push(val); + self.validity.append(is_valid); + last += val; + } + let mut prev = self.normalize(*new_indices.values().first().unwrap()).1; + for (i, w) in new_indices.values().windows(2).enumerate() { + let (is_valid, val) = self.normalize(w[1]); + let next = match val.checked_sub(prev) { + Some(delta) => delta + last, + None => { + return Err(lance_core::Error::invalid_input(format!( + "corrupt binary page: normalized offset {} is less than previous offset {} \ + at index {}, null_adjustment={}, raw values were [{}, {}]. \ + This usually indicates the file data has been corrupted.", + val, prev, i, self.null_adjustment, w[0], w[1] + ))); + } + }; + self.indices.push(next); + self.validity.append(is_valid); + prev = val; + last = next; + } + Ok(()) + } + + fn into_parts(mut self) -> (Vec, BooleanBuffer) { + (self.indices, self.validity.finish()) + } +} + +#[derive(Debug)] +pub struct BinaryPageScheduler { + indices_scheduler: Arc, + bytes_scheduler: Arc, + offsets_type: DataType, + null_adjustment: u64, +} + +impl BinaryPageScheduler { + pub fn new( + indices_scheduler: Arc, + bytes_scheduler: Arc, + offsets_type: DataType, + null_adjustment: u64, + ) -> Self { + Self { + indices_scheduler, + bytes_scheduler, + offsets_type, + null_adjustment, + } + } + + fn decode_indices(decoder: Arc, num_rows: u64) -> Result { + let mut primitive_wrapper = + PrimitiveFieldDecoder::new_from_data(decoder, DataType::UInt64, num_rows, false); + let drained_task = primitive_wrapper.drain(num_rows)?; + let indices_decode_task = drained_task.task; + indices_decode_task.decode().map(|(arr, _)| arr) + } +} + +struct IndirectData { + decoded_indices: UInt64Array, + offsets_type: DataType, + validity: BooleanBuffer, + bytes_decoder_fut: BoxFuture<'static, Result>>, +} + +impl PageScheduler for BinaryPageScheduler { + fn schedule_ranges( + &self, + ranges: &[std::ops::Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + // ranges corresponds to row ranges that the user wants to fetch. + // if user wants row range a..b + // Case 1: if a != 0, we need indices a-1..b to decode + // Case 2: if a = 0, we need indices 0..b to decode + let indices_ranges = ranges + .iter() + .map(|range| { + if range.start != 0 { + (range.start - 1)..range.end + } else { + 0..range.end + } + }) + .collect::>>(); + + // We schedule all the indices for decoding together + // This is more efficient compared to scheduling them one by one (reduces speed significantly for random access) + let indices_page_decoder = + self.indices_scheduler + .schedule_ranges(&indices_ranges, scheduler, top_level_row); + + let num_rows = ranges.iter().map(|r| r.end - r.start).sum::(); + let indices_num_rows = indices_ranges.iter().map(|r| r.end - r.start).sum::(); + + let ranges = ranges.to_vec(); + let copy_scheduler = scheduler.clone(); + let copy_bytes_scheduler = self.bytes_scheduler.clone(); + let null_adjustment = self.null_adjustment; + let offsets_type = self.offsets_type.clone(); + + tokio::spawn(async move { + // For the following data: + // "abcd", "hello", "abcd", "apple", "hello", "abcd" + // 4, 9, 13, 18, 23, 27 + // e.g. want to scan rows 0, 2, 4 + // i.e. offsets are 4 | 9, 13 | 18, 23 + // Normalization is required for decoding later on + // Normalize each part: 0, 4 | 0, 4 | 0, 5 + // Remove leading zeros except first one: 0, 4 | 4 | 5 + // Cumulative sum: 0, 4 | 8 | 13 + // These are the normalized offsets stored in decoded_indices + // Rest of the workflow is continued later in BinaryPageDecoder + let indices_decoder = Arc::from(indices_page_decoder.await?); + let indices = Self::decode_indices(indices_decoder, indices_num_rows)?; + let decoded_indices = indices.as_primitive::(); + + let mut indices_builder = IndicesNormalizer::new(num_rows, null_adjustment); + let mut bytes_ranges = Vec::new(); + let mut curr_offset_index = 0; + + for curr_row_range in ranges.iter() { + let row_start = curr_row_range.start; + let curr_range_len = (curr_row_range.end - row_start) as usize; + + let curr_indices; + + if row_start == 0 { + curr_indices = decoded_indices.slice(0, curr_range_len); + curr_offset_index = curr_range_len; + } else { + curr_indices = decoded_indices.slice(curr_offset_index, curr_range_len + 1); + curr_offset_index += curr_range_len + 1; + } + + let first = if row_start == 0 { + 0 + } else { + indices_builder + .normalize(*curr_indices.values().first().unwrap()) + .1 + }; + let last = indices_builder + .normalize(*curr_indices.values().last().unwrap()) + .1; + + if first != last { + bytes_ranges.push(first..last); + } + + indices_builder.extend(&curr_indices, row_start == 0)?; + } + + let (indices, validity) = indices_builder.into_parts(); + let decoded_indices = UInt64Array::from(indices); + + // In the indirect task we schedule the bytes, but we do not await them. We don't want to + // await the bytes until the decoder is ready for them so that we don't release the backpressure + // too early + let bytes_decoder_fut = + copy_bytes_scheduler.schedule_ranges(&bytes_ranges, ©_scheduler, top_level_row); + + Ok(IndirectData { + decoded_indices, + validity, + offsets_type, + bytes_decoder_fut, + }) + }) + // Propagate join panic + .map(|join_handle| join_handle.unwrap()) + .and_then(|indirect_data| { + async move { + // Later, this will be called once the decoder actually starts polling. At that point + // we await the bytes (releasing the backpressure) + let bytes_decoder = indirect_data.bytes_decoder_fut.await?; + Ok(Box::new(BinaryPageDecoder { + decoded_indices: indirect_data.decoded_indices, + offsets_type: indirect_data.offsets_type, + validity: indirect_data.validity, + bytes_decoder, + }) as Box) + } + }) + .boxed() + } +} + +struct BinaryPageDecoder { + decoded_indices: UInt64Array, + offsets_type: DataType, + validity: BooleanBuffer, + bytes_decoder: Box, +} + +impl PrimitivePageDecoder for BinaryPageDecoder { + fn variable_width_bytes(&self, rows_to_skip: u64, num_rows: u64) -> Result> { + if num_rows == 0 { + return Ok(Some(0)); + } + // `decoded_indices` holds one cumulative byte offset per row plus a final + // sentinel, so the value bytes for the requested rows are the difference + // between the bounding entries. Only value bytes count: Arrow's i32 + // limit constrains the final offset value, not the offset buffer size. + let value_bytes = self + .decoded_indices + .value((rows_to_skip + num_rows) as usize) + - self.decoded_indices.value(rows_to_skip as usize); + Ok(Some(value_bytes)) + } + + // Continuing the example from BinaryPageScheduler + // Suppose batch_size = 2. Then first, rows_to_skip=0, num_rows=2 + // Need to scan 2 rows + // First row will be 4-0=4 bytes, second also 8-4=4 bytes. + // Allocate 8 bytes capacity. + // Next rows_to_skip=2, num_rows=1 + // Skip 8 bytes. Allocate 5 bytes capacity. + // + // The normalized offsets are [0, 4, 8, 13] + // We only need [8, 13] to decode in this case. + // These need to be normalized in order to build the string later + // So return [0, 5] + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + // STEP 1: validity buffer + let target_validity = self + .validity + .slice(rows_to_skip as usize, num_rows as usize); + let has_nulls = target_validity.count_set_bits() < target_validity.len(); + + let validity_buffer = if has_nulls { + let num_validity_bits = arrow_buffer::bit_util::ceil(num_rows as usize, 8); + let mut validity_buffer = Vec::with_capacity(num_validity_bits); + + if rows_to_skip == 0 { + validity_buffer.extend_from_slice(target_validity.inner().as_slice()); + } else { + // Need to copy the buffer because there may be a bit offset in first byte + let target_validity = BooleanBuffer::from_iter(target_validity.iter()); + validity_buffer.extend_from_slice(target_validity.inner().as_slice()); + } + Some(validity_buffer) + } else { + None + }; + + // STEP 2: offsets buffer + // Currently we always do a copy here, we need to cast to the appropriate type + // and we go ahead and normalize so the starting offset is 0 (though we could skip + // this) + let bytes_per_offset = match self.offsets_type { + DataType::Int32 => 4, + DataType::Int64 => 8, + _ => panic!("Unsupported offsets type"), + }; + + let target_offsets = self + .decoded_indices + .slice(rows_to_skip as usize, (num_rows + 1) as usize); + + // Normalize and cast (TODO: could fuse these into one pass for micro-optimization) + let target_vec = target_offsets.values(); + let start = target_vec[0]; + let end = *target_vec.last().unwrap(); + let num_bytes = end - start; + let offsets_buffer = match bytes_per_offset { + 4 => { + if num_bytes > i32::MAX as u64 { + return Err(oversized_binary_batch_error(num_rows, num_bytes)); + } + ScalarBuffer::from_iter(target_vec.iter().map(|&offset| (offset - start) as i32)) + .into_inner() + } + 8 => { + if num_bytes > i64::MAX as u64 { + return Err(oversized_large_binary_batch_error(num_rows, num_bytes)); + } + ScalarBuffer::from_iter(target_vec.iter().map(|&offset| (offset - start) as i64)) + .into_inner() + } + _ => panic!("Unsupported offsets type"), + }; + + let bytes_to_skip = self.decoded_indices.value(rows_to_skip as usize); + let num_bytes = self + .decoded_indices + .value((rows_to_skip + num_rows) as usize) + - bytes_to_skip; + + let bytes = self.bytes_decoder.decode(bytes_to_skip, num_bytes)?; + let bytes = bytes.as_fixed_width().unwrap(); + debug_assert_eq!(bytes.bits_per_value, 8); + + let string_data = DataBlock::VariableWidth(VariableWidthBlock { + bits_per_offset: bytes_per_offset * 8, + data: bytes.data, + num_values: num_rows, + offsets: LanceBuffer::from(offsets_buffer), + block_info: BlockInfo::new(), + }); + if let Some(validity) = validity_buffer { + Ok(DataBlock::Nullable(NullableDataBlock { + data: Box::new(string_data), + nulls: LanceBuffer::from(validity), + block_info: BlockInfo::new(), + })) + } else { + Ok(string_data) + } + } +} + +#[derive(Debug)] +pub struct BinaryEncoder { + indices_encoder: Box, + compression_config: Option, + buffer_compressor: Option>, +} + +impl BinaryEncoder { + pub fn try_new( + indices_encoder: Box, + compression_config: Option, + ) -> Result { + let buffer_compressor = compression_config + .map(GeneralBufferCompressor::get_compressor) + .transpose()?; + Ok(Self { + indices_encoder, + compression_config, + buffer_compressor, + }) + } + + // In 2.1 we will materialize nulls higher up (in the primitive encoder). Unfortunately, + // in 2.0 we actually need to write the offsets. + fn all_null_variable_width(data_type: &DataType, num_values: u64) -> VariableWidthBlock { + if matches!(data_type, DataType::Binary | DataType::Utf8) { + VariableWidthBlock { + bits_per_offset: 32, + data: LanceBuffer::empty(), + num_values, + offsets: LanceBuffer::reinterpret_vec(vec![0_u32; num_values as usize + 1]), + block_info: BlockInfo::new(), + } + } else { + VariableWidthBlock { + bits_per_offset: 64, + data: LanceBuffer::empty(), + num_values, + offsets: LanceBuffer::reinterpret_vec(vec![0_u64; num_values as usize + 1]), + block_info: BlockInfo::new(), + } + } + } +} + +// Creates indices arrays from string arrays +// Strings are a vector of arrays corresponding to each record batch +// Zero offset is removed from the start of the offsets array +// The indices array is computed across all arrays in the vector +fn get_indices_from_string_arrays( + offsets: LanceBuffer, + bits_per_offset: u8, + nulls: Option, + num_rows: usize, +) -> (DataBlock, u64) { + let mut indices = Vec::with_capacity(num_rows); + let mut last_offset = 0_u64; + if bits_per_offset == 32 { + let offsets = offsets.borrow_to_typed_slice::(); + indices.extend(offsets.as_ref().windows(2).map(|w| { + let strlen = (w[1] - w[0]) as u64; + last_offset += strlen; + last_offset + })); + } else if bits_per_offset == 64 { + let offsets = offsets.borrow_to_typed_slice::(); + indices.extend(offsets.as_ref().windows(2).map(|w| { + let strlen = (w[1] - w[0]) as u64; + last_offset += strlen; + last_offset + })); + } + + if indices.is_empty() { + return ( + DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::empty(), + num_values: 0, + block_info: BlockInfo::new(), + }), + 0, + ); + } + + let last_offset = *indices.last().expect("Indices array is empty"); + // 8 exabytes in a single array seems unlikely but...just in case + assert!( + last_offset < u64::MAX / 2, + "Indices array with strings up to 2^63 is too large for this encoding" + ); + let null_adjustment: u64 = *indices.last().expect("Indices array is empty") + 1; + + if let Some(nulls) = nulls { + let nulls = NullBuffer::new(BooleanBuffer::new(nulls.into_buffer(), 0, num_rows)); + indices + .iter_mut() + .zip(nulls.iter()) + .for_each(|(index, is_valid)| { + if !is_valid { + *index += null_adjustment; + } + }); + } + let indices = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(indices), + num_values: num_rows as u64, + block_info: BlockInfo::new(), + }); + (indices, null_adjustment) +} + +impl ArrayEncoder for BinaryEncoder { + fn encode( + &self, + data: DataBlock, + data_type: &DataType, + buffer_index: &mut u32, + ) -> Result { + let (mut data, nulls) = match data { + DataBlock::Nullable(nullable) => { + let data = nullable.data.as_variable_width().unwrap(); + (data, Some(nullable.nulls)) + } + DataBlock::VariableWidth(variable) => (variable, None), + DataBlock::AllNull(all_null) => { + let data = Self::all_null_variable_width(data_type, all_null.num_values); + let validity = + LanceBuffer::all_unset(bit_util::ceil(all_null.num_values as usize, 8)); + (data, Some(validity)) + } + _ => panic!("Expected variable width data block but got {}", data.name()), + }; + + let (indices, null_adjustment) = get_indices_from_string_arrays( + data.offsets, + data.bits_per_offset, + nulls, + data.num_values as usize, + ); + let encoded_indices = + self.indices_encoder + .encode(indices, &DataType::UInt64, buffer_index)?; + + let encoded_indices_data = encoded_indices.data.as_fixed_width().unwrap(); + + assert!(encoded_indices_data.bits_per_value <= 64); + + if let Some(buffer_compressor) = &self.buffer_compressor { + let mut compressed_data = Vec::with_capacity(data.data.len()); + buffer_compressor.compress(&data.data, &mut compressed_data)?; + data.data = LanceBuffer::from(compressed_data); + } + + let data = DataBlock::VariableWidth(VariableWidthBlock { + bits_per_offset: encoded_indices_data.bits_per_value as u8, + offsets: encoded_indices_data.data, + data: data.data, + num_values: data.num_values, + block_info: BlockInfo::new(), + }); + + let bytes_buffer_index = *buffer_index; + *buffer_index += 1; + + let bytes_encoding = ProtobufUtils::flat_encoding( + /*bits_per_value=*/ 8, + bytes_buffer_index, + self.compression_config, + ); + + let encoding = + ProtobufUtils::binary(encoded_indices.encoding, bytes_encoding, null_adjustment); + + Ok(EncodedArray { data, encoding }) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::StringArray; + + use super::*; + + #[derive(Debug)] + struct EmptyBytesDecoder; + + impl PrimitivePageDecoder for EmptyBytesDecoder { + fn decode(&self, _rows_to_skip: u64, _num_rows: u64) -> Result { + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 8, + data: LanceBuffer::empty(), + num_values: 0, + block_info: BlockInfo::new(), + })) + } + } + + #[derive(Debug)] + struct BytesDecoder { + bytes: Vec, + } + + impl PrimitivePageDecoder for BytesDecoder { + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + let start = rows_to_skip as usize; + let end = start + num_rows as usize; + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 8, + data: LanceBuffer::from(self.bytes[start..end].to_vec()), + num_values: num_rows, + block_info: BlockInfo::new(), + })) + } + } + + #[test] + fn test_encode_indices_adjusts_nulls() { + // Null entries in string arrays should be adjusted + let string_array = Arc::new(StringArray::from(vec![ + None, + Some("foo"), + Some("foo"), + None, + None, + None, + ])) as ArrayRef; + let string_data = DataBlock::from(string_array).as_nullable().unwrap(); + let nulls = string_data.nulls; + let string_data = string_data.data.as_variable_width().unwrap(); + + let (indices, null_adjustment) = get_indices_from_string_arrays( + string_data.offsets, + string_data.bits_per_offset, + Some(nulls), + string_data.num_values as usize, + ); + + let indices = indices.as_fixed_width().unwrap(); + assert_eq!(indices.bits_per_value, 64); + assert_eq!( + indices.data, + LanceBuffer::reinterpret_vec(vec![7_u64, 3, 6, 13, 13, 13]) + ); + assert_eq!(null_adjustment, 7); + } + + #[test] + fn test_binary_overflow_error_is_actionable() { + let num_rows = 1; + let start = 100_u64; + let end = start + i32::MAX as u64 + 1; + let decoded_indices = UInt64Array::from(vec![start, end]); + let decoder = BinaryPageDecoder { + decoded_indices, + validity: BooleanBuffer::from_iter([true]), + offsets_type: DataType::Int32, + bytes_decoder: Box::new(EmptyBytesDecoder), + }; + + let error = decoder.decode(0, num_rows).unwrap_err(); + assert!(matches!(error, Error::NotSupported { .. })); + let message = error.to_string(); + assert!(message.contains("more than 2GiB of string/binary data")); + assert!(message.contains("batch_size")); + assert!(message.contains("LANCE_DEFAULT_BATCH_SIZE")); + assert!(message.contains("large_string/large_binary")); + } + + #[test] + fn test_large_binary_overflow_error_is_actionable() { + let num_rows = 1; + let start = 100_u64; + let end = start + i64::MAX as u64 + 1; + let decoded_indices = UInt64Array::from(vec![start, end]); + let decoder = BinaryPageDecoder { + decoded_indices, + validity: BooleanBuffer::from_iter([true]), + offsets_type: DataType::Int64, + bytes_decoder: Box::new(EmptyBytesDecoder), + }; + + let error = decoder.decode(0, num_rows).unwrap_err(); + assert!(matches!(error, Error::NotSupported { .. })); + let message = error.to_string(); + assert!(message.contains("large_string/large_binary")); + assert!(message.contains("exceeds i64::MAX")); + assert!(message.contains("batch_size")); + assert!(message.contains("LANCE_DEFAULT_BATCH_SIZE")); + } + + #[test] + fn test_large_binary_decode_success_path() { + let decoded_indices = UInt64Array::from(vec![100_u64, 102, 105]); + let mut bytes = vec![0_u8; 100]; + bytes.extend_from_slice(b"abcde"); + let decoder = BinaryPageDecoder { + decoded_indices, + validity: BooleanBuffer::from_iter([true, true]), + offsets_type: DataType::Int64, + bytes_decoder: Box::new(BytesDecoder { bytes }), + }; + + let data = decoder.decode(0, 2).unwrap(); + let variable = data.as_variable_width().unwrap(); + assert_eq!(variable.bits_per_offset, 64); + assert_eq!(variable.data.as_ref(), b"abcde"); + assert_eq!( + variable.offsets.borrow_to_typed_slice::().as_ref(), + &[0_i64, 2, 5] + ); + } + + #[derive(Debug)] + struct NeverDecodedStub; + + impl PrimitivePageDecoder for NeverDecodedStub { + fn decode(&self, _rows_to_skip: u64, _num_rows: u64) -> Result { + unreachable!("byte accounting must not decode any values") + } + } + + /// A real physical binary page over strings of the given byte lengths. + fn binary_page(value_lens: &[u64]) -> BinaryPageDecoder { + let mut indices = vec![0u64]; + for len in value_lens { + indices.push(indices.last().unwrap() + len); + } + BinaryPageDecoder { + decoded_indices: UInt64Array::from(indices), + offsets_type: DataType::Int32, + validity: BooleanBuffer::new_set(value_lens.len()), + bytes_decoder: Box::new(NeverDecodedStub), + } + } + + /// Variable-width value bytes of `n` rows starting at `skip` (what the + /// output array's i32 offsets index into). + fn expected_bytes(value_lens: &[u64], skip: usize, n: usize) -> u64 { + value_lens[skip..skip + n].iter().sum::() + } + + #[test] + fn test_unknown_size_pages_split_at_page_boundaries() { + use crate::array_encoding::logical::primitive::PrimitiveFieldDecoder; + use crate::array_encoding::logical::r#struct::SimpleStructDecoder; + use crate::decoder::{DecoderReady, I32_OFFSET_BYTE_BUDGET, LogicalPageDecoder}; + use arrow_schema::{Field as ArrowField, Fields}; + use std::collections::VecDeque; + + /// A variable-width page that cannot report sizes (default + /// `variable_width_bytes` returns `None`). + #[derive(Debug)] + struct UnknownSizeStub; + + impl PrimitivePageDecoder for UnknownSizeStub { + fn decode(&self, _rows_to_skip: u64, num_rows: u64) -> Result { + Ok(DataBlock::VariableWidth(VariableWidthBlock { + bits_per_offset: 32, + data: LanceBuffer::empty(), + offsets: LanceBuffer::reinterpret_vec(vec![0_i32; num_rows as usize + 1]), + num_values: num_rows, + block_info: BlockInfo::new(), + })) + } + } + + let fields = Fields::from(vec![ArrowField::new("value", DataType::Utf8, false)]); + let mut root = SimpleStructDecoder::new(fields, 6); + for _ in 0..2 { + root.accept_child(DecoderReady { + decoder: Box::new(PrimitiveFieldDecoder::new_from_data( + Arc::new(UnknownSizeStub), + DataType::Utf8, + 3, + false, + )), + path: VecDeque::from([0]), + }) + .unwrap(); + } + + // Each page's size is unknown, so pages must not stack in one batch: + // the batch takes the first page alone and stops at its boundary. + let limit = root.max_rows_to_drain(6, I32_OFFSET_BYTE_BUDGET).unwrap(); + assert_eq!(limit.rows, 3); + } + + #[test] + fn test_physical_binary_page_reports_variable_width_bytes() { + let page = binary_page(&[4, 4, 5]); + assert_eq!( + page.variable_width_bytes(0, 3).unwrap(), + Some(expected_bytes(&[4, 4, 5], 0, 3)) + ); + assert_eq!( + page.variable_width_bytes(1, 2).unwrap(), + Some(expected_bytes(&[4, 4, 5], 1, 2)) + ); + assert_eq!(page.variable_width_bytes(0, 0).unwrap(), Some(0)); + } + + #[test] + fn test_primitive_field_decoder_truncates_to_byte_budget() { + use crate::array_encoding::logical::primitive::PrimitiveFieldDecoder; + use crate::decoder::LogicalPageDecoder; + + let lens = [4u64, 4, 5]; + let decoder = PrimitiveFieldDecoder::new_from_data( + Arc::new(binary_page(&lens)), + DataType::Utf8, + 3, + false, + ); + let full = expected_bytes(&lens, 0, 3); + + let limit = decoder.max_rows_to_drain(3, full).unwrap(); + assert_eq!((limit.rows, limit.bytes), (3, full)); + + // One byte short of the full request: only two rows fit. + let limit = decoder.max_rows_to_drain(3, full - 1).unwrap(); + assert_eq!((limit.rows, limit.bytes), (2, expected_bytes(&lens, 0, 2))); + + // Nothing fits: zero rows, so the stream root can raise an error. + let limit = decoder.max_rows_to_drain(3, 0).unwrap(); + assert_eq!(limit.rows, 0); + } + + #[test] + fn test_physical_binary_pages_accumulate_budget_across_pages() { + use crate::array_encoding::logical::primitive::PrimitiveFieldDecoder; + use crate::array_encoding::logical::r#struct::SimpleStructDecoder; + use crate::decoder::{DecoderReady, LogicalPageDecoder}; + use arrow_schema::{Field as ArrowField, Fields}; + use std::collections::VecDeque; + + let lens = [4u64, 4, 5]; + let page_bytes = expected_bytes(&lens, 0, 3); + let fields = Fields::from(vec![ArrowField::new("value", DataType::Utf8, false)]); + let mut root = SimpleStructDecoder::new(fields, 6); + for _ in 0..2 { + root.accept_child(DecoderReady { + decoder: Box::new(PrimitiveFieldDecoder::new_from_data( + Arc::new(binary_page(&lens)), + DataType::Utf8, + 3, + false, + )), + path: VecDeque::from([0]), + }) + .unwrap(); + } + + // Both pages fit: the batch spans the page boundary. + let limit = root.max_rows_to_drain(6, page_bytes * 2).unwrap(); + assert_eq!((limit.rows, limit.bytes), (6, page_bytes * 2)); + + // Page 1 plus the first two rows of page 2 fit. + let second_page_prefix = expected_bytes(&lens, 0, 2); + let limit = root + .max_rows_to_drain(6, page_bytes + second_page_prefix) + .unwrap(); + assert_eq!( + (limit.rows, limit.bytes), + (5, page_bytes + second_page_prefix) + ); + + // Budget for page 1 only: the batch stops at the page boundary. + let limit = root.max_rows_to_drain(6, page_bytes).unwrap(); + assert_eq!((limit.rows, limit.bytes), (3, page_bytes)); + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/bitmap.rs b/vendor/lance-encoding/src/array_encoding/physical/bitmap.rs new file mode 100644 index 000000000..2168aff24 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/bitmap.rs @@ -0,0 +1,223 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ops::Range, sync::Arc}; + +use arrow_buffer::BooleanBufferBuilder; +use bytes::Bytes; + +use futures::{FutureExt, future::BoxFuture}; +use lance_core::Result; +use log::trace; + +use crate::{ + EncodingsIo, + buffer::LanceBuffer, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + decoder::{PageScheduler, PrimitivePageDecoder}, +}; + +/// A physical scheduler for bitmap buffers encoded densely as 1 bit per value +/// with bit-endianness(e.g. what Arrow uses for validity bitmaps and boolean arrays) +/// +/// This decoder decodes from one buffer of disk data into one buffer of memory data +#[derive(Debug, Clone, Copy)] +pub struct DenseBitmapScheduler { + buffer_offset: u64, +} + +impl DenseBitmapScheduler { + pub fn new(buffer_offset: u64) -> Self { + Self { buffer_offset } + } +} + +impl PageScheduler for DenseBitmapScheduler { + fn schedule_ranges( + &self, + ranges: &[Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + let mut min = u64::MAX; + let mut max = 0; + let chunk_reqs = ranges + .iter() + .map(|range| { + debug_assert_ne!(range.start, range.end); + let start = self.buffer_offset + range.start / 8; + let bit_offset = range.start % 8; + let end = self.buffer_offset + range.end.div_ceil(8); + let byte_range = start..end; + min = min.min(start); + max = max.max(end); + (byte_range, bit_offset, range.end - range.start) + }) + .collect::>(); + + let byte_ranges = chunk_reqs + .iter() + .map(|(range, _, _)| range.clone()) + .collect::>(); + trace!( + "Scheduling I/O for {} ranges across byte range {}..{}", + byte_ranges.len(), + min, + max + ); + let bytes = scheduler.submit_request(byte_ranges, top_level_row); + + async move { + let bytes = bytes.await?; + let chunks = bytes + .into_iter() + .zip(chunk_reqs) + .map(|(bytes, (_, bit_offset, length))| BitmapData { + data: bytes, + bit_offset, + length, + }) + .collect::>(); + Ok(Box::new(BitmapDecoder { chunks }) as Box) + } + .boxed() + } +} + +struct BitmapData { + data: Bytes, + bit_offset: u64, + length: u64, +} + +struct BitmapDecoder { + chunks: Vec, +} + +impl PrimitivePageDecoder for BitmapDecoder { + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + let mut rows_to_skip = rows_to_skip; + let mut dest_builder = BooleanBufferBuilder::new(num_rows as usize); + + let mut rows_remaining = num_rows; + for chunk in &self.chunks { + if chunk.length <= rows_to_skip { + rows_to_skip -= chunk.length; + } else { + let start = rows_to_skip + chunk.bit_offset; + let num_vals_to_take = rows_remaining.min(chunk.length - rows_to_skip); + let end = start + num_vals_to_take; + dest_builder.append_packed_range(start as usize..end as usize, &chunk.data); + rows_to_skip = 0; + rows_remaining -= num_vals_to_take; + } + } + + let bool_buffer = dest_builder.finish().into_inner(); + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(bool_buffer), + bits_per_value: 1, + num_values: num_rows, + block_info: BlockInfo::new(), + })) + } +} + +#[cfg(test)] +mod tests { + + use arrow_array::BooleanArray; + use arrow_schema::{DataType, Field}; + use bytes::Bytes; + use std::{collections::HashMap, sync::Arc}; + + use crate::array_encoding::physical::bitmap::BitmapData; + use crate::data::{DataBlock, FixedWidthDataBlock}; + use crate::decoder::PrimitivePageDecoder; + use crate::testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data}; + + use super::BitmapDecoder; + + #[test_log::test(tokio::test)] + async fn test_bitmap_boolean() { + let field = Field::new("", DataType::Boolean, false); + check_basic_random(field).await; + } + + #[test_log::test(tokio::test)] + async fn test_fsl_bitmap_boolean() { + let field = Field::new("", DataType::Boolean, true); + let field = Field::new("", DataType::FixedSizeList(Arc::new(field), 3), true); + check_basic_random(field).await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_boolean() { + let array = BooleanArray::from(vec![ + Some(false), + Some(true), + None, + Some(false), + Some(true), + None, + Some(false), + None, + None, + ]); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(1..9) + .with_indices(vec![0, 1, 3, 4]); + check_round_trip_encoding_of_data(vec![Arc::new(array)], &test_cases, HashMap::default()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_tiny_boolean() { + // Test case for a tiny boolean array that is technically smaller than 1 byte + let array = BooleanArray::from(vec![Some(false), Some(true), None]); + + let test_cases = TestCases::default() + .with_range(0..1) + .with_range(1..3) + .with_indices(vec![0, 2]); + check_round_trip_encoding_of_data(vec![Arc::new(array)], &test_cases, HashMap::default()) + .await; + } + + #[test] + fn test_bitmap_decoder_edge_cases() { + // Regression for a case where the row skip and the bit offset + // require us to read from the second Bytes instead of the first + let decoder = BitmapDecoder { + chunks: vec![ + BitmapData { + data: Bytes::from_static(&[0b11111111]), + bit_offset: 4, + length: 4, + }, + BitmapData { + data: Bytes::from_static(&[0b00000000]), + bit_offset: 4, + length: 4, + }, + ], + }; + + // Read from first and second chunk + let result = decoder.decode(2, 4).unwrap(); + let DataBlock::FixedWidth(FixedWidthDataBlock { data, .. }) = result else { + panic!("expected fixed width data block"); + }; + assert_eq!(data.as_ref(), &[0b00000011]); + + // Read from second chunk + let result = decoder.decode(5, 1).unwrap(); + let DataBlock::FixedWidth(FixedWidthDataBlock { data, .. }) = result else { + panic!("expected fixed width data block"); + }; + assert_eq!(data.as_ref(), &[0b00000000]); + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/bitpack.rs b/vendor/lance-encoding/src/array_encoding/physical/bitpack.rs new file mode 100644 index 000000000..a76add68a --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/bitpack.rs @@ -0,0 +1,728 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_buffer::bit_util::ceil; +use bytes::Bytes; +use futures::future::{BoxFuture, FutureExt}; +use log::trace; + +use lance_bitpacking::BitPacking; +use lance_core::{Error, Result}; + +use crate::buffer::LanceBuffer; +use crate::data::BlockInfo; +use crate::data::{DataBlock, FixedWidthDataBlock}; +use crate::decoder::{PageScheduler, PrimitivePageDecoder}; +use bytemuck::cast_slice; + +const LOG_ELEMS_PER_CHUNK: u8 = 10; +const ELEMS_PER_CHUNK: u64 = 1 << LOG_ELEMS_PER_CHUNK; + +#[derive(Debug)] +pub struct BitpackedForNonNegScheduler { + compressed_bit_width: u64, + uncompressed_bits_per_value: u64, + buffer_offset: u64, +} + +impl BitpackedForNonNegScheduler { + pub fn new( + compressed_bit_width: u64, + uncompressed_bits_per_value: u64, + buffer_offset: u64, + ) -> Self { + Self { + compressed_bit_width, + uncompressed_bits_per_value, + buffer_offset, + } + } + + fn locate_chunk_start(&self, relative_row_num: u64) -> u64 { + let chunk_size = ELEMS_PER_CHUNK * self.compressed_bit_width / 8; + self.buffer_offset + (relative_row_num / ELEMS_PER_CHUNK * chunk_size) + } + + fn locate_chunk_end(&self, relative_row_num: u64) -> u64 { + let chunk_size = ELEMS_PER_CHUNK * self.compressed_bit_width / 8; + self.buffer_offset + (relative_row_num / ELEMS_PER_CHUNK * chunk_size) + chunk_size + } +} + +impl PageScheduler for BitpackedForNonNegScheduler { + fn schedule_ranges( + &self, + ranges: &[std::ops::Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + assert!(!ranges.is_empty()); + + let mut byte_ranges = vec![]; + + // map one bytes to multiple ranges, one bytes has at least one range corresponding to it + let mut bytes_idx_to_range_indices = vec![]; + let first_byte_range = std::ops::Range { + start: self.locate_chunk_start(ranges[0].start), + end: self.locate_chunk_end(ranges[0].end - 1), + }; // the ranges are half-open + byte_ranges.push(first_byte_range); + bytes_idx_to_range_indices.push(vec![ranges[0].clone()]); + + for (i, range) in ranges.iter().enumerate().skip(1) { + let this_start = self.locate_chunk_start(range.start); + let this_end = self.locate_chunk_end(range.end - 1); + + // when the current range start is in the same chunk as the previous range's end, we colaesce this two bytes ranges + // when the current range start is not in the same chunk as the previous range's end, we create a new bytes range + if this_start == self.locate_chunk_start(ranges[i - 1].end - 1) { + byte_ranges.last_mut().unwrap().end = this_end; + bytes_idx_to_range_indices + .last_mut() + .unwrap() + .push(range.clone()); + } else { + byte_ranges.push(this_start..this_end); + bytes_idx_to_range_indices.push(vec![range.clone()]); + } + } + + trace!( + "Scheduling I/O for {} ranges spread across byte range {}..{}", + byte_ranges.len(), + byte_ranges[0].start, + byte_ranges.last().unwrap().end + ); + + let bytes = scheduler.submit_request(byte_ranges.clone(), top_level_row); + + // copy the necessary data from `self` to move into the async block + let compressed_bit_width = self.compressed_bit_width; + let uncompressed_bits_per_value = self.uncompressed_bits_per_value; + let num_rows = ranges.iter().map(|range| range.end - range.start).sum(); + + async move { + let bytes = bytes.await?; + let decompressed_output = bitpacked_for_non_neg_decode( + compressed_bit_width, + uncompressed_bits_per_value, + &bytes, + &bytes_idx_to_range_indices, + num_rows, + ); + Ok(Box::new(BitpackedForNonNegPageDecoder { + uncompressed_bits_per_value, + decompressed_buf: decompressed_output, + }) as Box) + } + .boxed() + } +} + +#[derive(Debug)] +struct BitpackedForNonNegPageDecoder { + // number of bits in the uncompressed value. E.g. this will be 32 for DataType::UInt32 + uncompressed_bits_per_value: u64, + + decompressed_buf: LanceBuffer, +} + +impl PrimitivePageDecoder for BitpackedForNonNegPageDecoder { + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + if ![8, 16, 32, 64].contains(&self.uncompressed_bits_per_value) { + return Err(Error::invalid_input_source("BitpackedForNonNegPageDecoder should only has uncompressed_bits_per_value of 8, 16, 32, or 64".into())); + } + + let elem_size_in_bytes = self.uncompressed_bits_per_value / 8; + + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: self.decompressed_buf.slice_with_length( + (rows_to_skip * elem_size_in_bytes) as usize, + (num_rows * elem_size_in_bytes) as usize, + ), + bits_per_value: self.uncompressed_bits_per_value, + num_values: num_rows, + block_info: BlockInfo::new(), + })) + } +} + +macro_rules! bitpacked_decode { + ($uncompressed_type:ty, $compressed_bit_width:expr, $data:expr, $bytes_idx_to_range_indices:expr, $num_rows:expr) => {{ + let mut decompressed: Vec<$uncompressed_type> = Vec::with_capacity($num_rows as usize); + let packed_chunk_size_in_byte: usize = (ELEMS_PER_CHUNK * $compressed_bit_width) as usize / 8; + let mut decompress_chunk_buf = vec![0 as $uncompressed_type; ELEMS_PER_CHUNK as usize]; + + for (i, bytes) in $data.iter().enumerate() { + let mut ranges_idx = 0; + let mut curr_range_start = $bytes_idx_to_range_indices[i][0].start; + let mut chunk_num = 0; + + while chunk_num * packed_chunk_size_in_byte < bytes.len() { + // Copy for memory alignment + // TODO: This copy should not be needed + let chunk_in_u8: Vec = bytes[chunk_num * packed_chunk_size_in_byte..] + [..packed_chunk_size_in_byte] + .to_vec(); + chunk_num += 1; + let chunk = cast_slice(&chunk_in_u8); + unsafe { + BitPacking::unchecked_unpack( + $compressed_bit_width as usize, + chunk, + &mut decompress_chunk_buf, + ); + } + + loop { + // Case 1: All the elements after (curr_range_start % ELEMS_PER_CHUNK) inside this chunk are needed. + let elems_after_curr_range_start_in_this_chunk = + ELEMS_PER_CHUNK - curr_range_start % ELEMS_PER_CHUNK; + if curr_range_start + elems_after_curr_range_start_in_this_chunk + <= $bytes_idx_to_range_indices[i][ranges_idx].end + { + decompressed.extend_from_slice( + &decompress_chunk_buf[(curr_range_start % ELEMS_PER_CHUNK) as usize..], + ); + curr_range_start += elems_after_curr_range_start_in_this_chunk; + break; + } else { + // Case 2: Only part of the elements after (curr_range_start % ELEMS_PER_CHUNK) inside this chunk are needed. + let elems_this_range_needed_in_this_chunk = + ($bytes_idx_to_range_indices[i][ranges_idx].end - curr_range_start) + .min(ELEMS_PER_CHUNK - curr_range_start % ELEMS_PER_CHUNK); + decompressed.extend_from_slice( + &decompress_chunk_buf[(curr_range_start % ELEMS_PER_CHUNK) as usize..] + [..elems_this_range_needed_in_this_chunk as usize], + ); + if curr_range_start + elems_this_range_needed_in_this_chunk + == $bytes_idx_to_range_indices[i][ranges_idx].end + { + ranges_idx += 1; + if ranges_idx == $bytes_idx_to_range_indices[i].len() { + break; + } + curr_range_start = $bytes_idx_to_range_indices[i][ranges_idx].start; + } else { + curr_range_start += elems_this_range_needed_in_this_chunk; + } + } + } + } + } + + LanceBuffer::reinterpret_vec(decompressed) + }}; +} + +fn bitpacked_for_non_neg_decode( + compressed_bit_width: u64, + uncompressed_bits_per_value: u64, + data: &[Bytes], + bytes_idx_to_range_indices: &[Vec>], + num_rows: u64, +) -> LanceBuffer { + match uncompressed_bits_per_value { + 8 => bitpacked_decode!( + u8, + compressed_bit_width, + data, + bytes_idx_to_range_indices, + num_rows + ), + 16 => bitpacked_decode!( + u16, + compressed_bit_width, + data, + bytes_idx_to_range_indices, + num_rows + ), + 32 => bitpacked_decode!( + u32, + compressed_bit_width, + data, + bytes_idx_to_range_indices, + num_rows + ), + 64 => bitpacked_decode!( + u64, + compressed_bit_width, + data, + bytes_idx_to_range_indices, + num_rows + ), + _ => unreachable!( + "bitpacked_for_non_neg_decode only supports 8, 16, 32, 64 uncompressed_bits_per_value" + ), + } +} + +// A physical scheduler for bitpacked buffers +#[derive(Debug, Clone, Copy)] +pub struct BitpackedScheduler { + bits_per_value: u64, + uncompressed_bits_per_value: u64, + buffer_offset: u64, + signed: bool, +} + +impl BitpackedScheduler { + pub fn new( + bits_per_value: u64, + uncompressed_bits_per_value: u64, + buffer_offset: u64, + signed: bool, + ) -> Self { + Self { + bits_per_value, + uncompressed_bits_per_value, + buffer_offset, + signed, + } + } +} + +impl PageScheduler for BitpackedScheduler { + fn schedule_ranges( + &self, + ranges: &[std::ops::Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + let mut min = u64::MAX; + let mut max = 0; + + let mut buffer_bit_start_offsets: Vec = vec![]; + let mut buffer_bit_end_offsets: Vec> = vec![]; + let byte_ranges = ranges + .iter() + .map(|range| { + let start_byte_offset = range.start * self.bits_per_value / 8; + let mut end_byte_offset = range.end * self.bits_per_value / 8; + if !(range.end * self.bits_per_value).is_multiple_of(8) { + // If the end of the range is not byte-aligned, we need to read one more byte + end_byte_offset += 1; + + let end_bit_offset = range.end * self.bits_per_value % 8; + buffer_bit_end_offsets.push(Some(end_bit_offset as u8)); + } else { + buffer_bit_end_offsets.push(None); + } + + let start_bit_offset = range.start * self.bits_per_value % 8; + buffer_bit_start_offsets.push(start_bit_offset as u8); + + let start = self.buffer_offset + start_byte_offset; + let end = self.buffer_offset + end_byte_offset; + min = min.min(start); + max = max.max(end); + + start..end + }) + .collect::>(); + + trace!( + "Scheduling I/O for {} ranges spread across byte range {}..{}", + byte_ranges.len(), + min, + max + ); + + let bytes = scheduler.submit_request(byte_ranges, top_level_row); + + let bits_per_value = self.bits_per_value; + let uncompressed_bits_per_value = self.uncompressed_bits_per_value; + let signed = self.signed; + async move { + let bytes = bytes.await?; + Ok(Box::new(BitpackedPageDecoder { + buffer_bit_start_offsets, + buffer_bit_end_offsets, + bits_per_value, + uncompressed_bits_per_value, + signed, + data: bytes, + }) as Box) + } + .boxed() + } +} + +#[derive(Debug)] +struct BitpackedPageDecoder { + // bit offsets of the first value within each buffer + buffer_bit_start_offsets: Vec, + + // bit offsets of the last value within each buffer. e.g. if there was a buffer + // with 2 values, packed into 5 bits, this would be [Some(3)], indicating that + // the bits from the 3rd->8th bit in the last byte shouldn't be decoded. + buffer_bit_end_offsets: Vec>, + + // the number of bits used to represent a compressed value. E.g. if the max value + // in the page was 7 (0b111), then this will be 3 + bits_per_value: u64, + + // number of bits in the uncompressed value. E.g. this will be 32 for u32 + uncompressed_bits_per_value: u64, + + // whether or not to use the msb as a sign bit during decoding + signed: bool, + + data: Vec, +} + +impl PrimitivePageDecoder for BitpackedPageDecoder { + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + let num_bytes = self.uncompressed_bits_per_value / 8 * num_rows; + let mut dest = vec![0; num_bytes as usize]; + + // current maximum supported bits per value = 64 + debug_assert!(self.bits_per_value <= 64); + + let mut rows_to_skip = rows_to_skip; + let mut rows_taken = 0; + let byte_len = self.uncompressed_bits_per_value / 8; + let mut dst_idx = 0; // index for current byte being written to destination buffer + + // create bit mask for source bits + let mask = u64::MAX >> (64 - self.bits_per_value); + + for i in 0..self.data.len() { + let src = &self.data[i]; + let (mut src_idx, mut src_offset) = match compute_start_offset( + rows_to_skip, + src.len(), + self.bits_per_value, + self.buffer_bit_start_offsets[i], + self.buffer_bit_end_offsets[i], + ) { + StartOffset::SkipFull(rows_to_skip_here) => { + rows_to_skip -= rows_to_skip_here; + continue; + } + StartOffset::SkipSome(buffer_start_offset) => ( + buffer_start_offset.index, + buffer_start_offset.bit_offset as u64, + ), + }; + + while src_idx < src.len() && rows_taken < num_rows { + rows_taken += 1; + let mut curr_mask = mask; // copy mask + + // current source byte being written to destination + let mut curr_src = src[src_idx] & (curr_mask << src_offset) as u8; + + // how many bits from the current source value have been written to destination + let mut src_bits_written = 0; + + // the offset within the current destination byte to write to + let mut dst_offset = 0; + + let is_negative = is_encoded_item_negative( + src, + src_idx, + src_offset, + self.bits_per_value as usize, + ); + + while src_bits_written < self.bits_per_value { + // write bits from current source byte into destination + dest[dst_idx] += (curr_src >> src_offset) << dst_offset; + let bits_written = (self.bits_per_value - src_bits_written) + .min(8 - src_offset) + .min(8 - dst_offset); + src_bits_written += bits_written; + dst_offset += bits_written; + src_offset += bits_written; + curr_mask >>= bits_written; + + if dst_offset == 8 { + dst_idx += 1; + dst_offset = 0; + } + + if src_offset == 8 { + src_idx += 1; + src_offset = 0; + if src_idx == src.len() { + break; + } + curr_src = src[src_idx] & curr_mask as u8; + } + } + + // if the type is signed, need to pad out the rest of the byte with 1s + let mut negative_padded_current_byte = false; + if self.signed && is_negative && dst_offset > 0 { + negative_padded_current_byte = true; + while dst_offset < 8 { + dest[dst_idx] |= 1 << dst_offset; + dst_offset += 1; + } + } + + // advance destination offset to the next location + // note that we don't need to do this if we wrote the full number of bits + // because source index would have been advanced by the inner loop above + if self.uncompressed_bits_per_value != self.bits_per_value { + let partial_bytes_written = ceil(self.bits_per_value as usize, 8); + + // we also want to move one location to the next location in destination, + // unless we wrote something byte-aligned in which case the logic above + // would have already advanced dst_idx + let mut to_next_byte = 1; + if self.bits_per_value.is_multiple_of(8) { + to_next_byte = 0; + } + let next_dst_idx = + dst_idx + byte_len as usize - partial_bytes_written + to_next_byte; + + // pad remaining bytes with 1 for negative signed numbers + if self.signed && is_negative { + if !negative_padded_current_byte { + dest[dst_idx] = 0xFF; + } + for i in dest.iter_mut().take(next_dst_idx).skip(dst_idx + 1) { + *i = 0xFF; + } + } + + dst_idx = next_dst_idx; + } + + // If we've reached the last byte, there may be some extra bits from the + // next value outside the range. We don't want to be taking those. + if let Some(buffer_bit_end_offset) = self.buffer_bit_end_offsets[i] + && src_idx == src.len() - 1 + && src_offset >= buffer_bit_end_offset as u64 + { + break; + } + } + } + + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(dest), + bits_per_value: self.uncompressed_bits_per_value, + num_values: num_rows, + block_info: BlockInfo::new(), + })) + } +} + +fn is_encoded_item_negative(src: &Bytes, src_idx: usize, src_offset: u64, num_bits: usize) -> bool { + let mut last_byte_idx = src_idx + ((src_offset as usize + num_bits) / 8); + let shift_amount = (src_offset as usize + num_bits) % 8; + let shift_amount = if shift_amount == 0 { + last_byte_idx -= 1; + 7 + } else { + shift_amount - 1 + }; + let last_byte = src[last_byte_idx]; + let sign_bit_mask = 1 << shift_amount; + let sign_bit = last_byte & sign_bit_mask; + + sign_bit > 0 +} + +#[derive(Debug, PartialEq)] +struct BufferStartOffset { + index: usize, + bit_offset: u8, +} + +#[derive(Debug, PartialEq)] +enum StartOffset { + // skip the full buffer. The value is how many rows are skipped + // by skipping the full buffer (e.g., # rows in buffer) + SkipFull(u64), + + // skip to some start offset in the buffer + SkipSome(BufferStartOffset), +} + +/// compute how far ahead in this buffer should we skip ahead and start reading +/// +/// * `rows_to_skip` - how many rows to skip +/// * `buffer_len` - length buf buffer (in bytes) +/// * `bits_per_value` - number of bits used to represent a single bitpacked value +/// * `buffer_start_bit_offset` - offset of the start of the first value within the +/// buffer's first byte +/// * `buffer_end_bit_offset` - end bit of the last value within the buffer. Can be +/// `None` if the end of the last value is byte aligned with end of buffer. +fn compute_start_offset( + rows_to_skip: u64, + buffer_len: usize, + bits_per_value: u64, + buffer_start_bit_offset: u8, + buffer_end_bit_offset: Option, +) -> StartOffset { + let rows_in_buffer = rows_in_buffer( + buffer_len, + bits_per_value, + buffer_start_bit_offset, + buffer_end_bit_offset, + ); + if rows_to_skip >= rows_in_buffer { + return StartOffset::SkipFull(rows_in_buffer); + } + + let start_bit = rows_to_skip * bits_per_value + buffer_start_bit_offset as u64; + let start_byte = start_bit / 8; + + StartOffset::SkipSome(BufferStartOffset { + index: start_byte as usize, + bit_offset: (start_bit % 8) as u8, + }) +} + +/// calculates the number of rows in a buffer +fn rows_in_buffer( + buffer_len: usize, + bits_per_value: u64, + buffer_start_bit_offset: u8, + buffer_end_bit_offset: Option, +) -> u64 { + let mut bits_in_buffer = (buffer_len * 8) as u64 - buffer_start_bit_offset as u64; + + // if the end of the last value of the buffer isn't byte aligned, subtract the + // end offset from the total number of bits in buffer + if let Some(buffer_end_bit_offset) = buffer_end_bit_offset { + bits_in_buffer -= (8 - buffer_end_bit_offset) as u64; + } + + bits_in_buffer / bits_per_value +} + +#[cfg(test)] +pub mod test { + use super::*; + use crate::BufferScheduler; + + use arrow_buffer::ArrowNativeType; + + fn fixed_width_values(block: DataBlock) -> Vec { + let DataBlock::FixedWidth(FixedWidthDataBlock { data, .. }) = block else { + panic!("expected fixed-width data"); + }; + data.borrow_to_typed_slice::().to_vec() + } + + #[test] + fn test_rows_in_buffer() { + let test_cases = vec![ + (5usize, 5u64, 0u8, None, 8u64), + (2, 3, 0, Some(5), 4), + (2, 3, 7, Some(6), 2), + ]; + + for ( + buffer_len, + bits_per_value, + buffer_start_bit_offset, + buffer_end_bit_offset, + expected, + ) in test_cases + { + let result = rows_in_buffer( + buffer_len, + bits_per_value, + buffer_start_bit_offset, + buffer_end_bit_offset, + ); + assert_eq!(expected, result); + } + } + + #[test] + fn test_compute_start_offset() { + let result = compute_start_offset(0, 5, 5, 0, None); + assert_eq!( + StartOffset::SkipSome(BufferStartOffset { + index: 0, + bit_offset: 0 + }), + result + ); + + let result = compute_start_offset(10, 5, 5, 0, None); + assert_eq!(StartOffset::SkipFull(8), result); + } + + #[test_log::test(tokio::test)] + async fn test_bitpacked_scheduler_non_byte_aligned_ranges() { + // Legacy 5-bit LSB-first encoding of values 1..=10, with a two-byte prefix. + let data = Bytes::from_static(&[0xFA, 0xCE, 0x41, 0x0C, 0x52, 0xCC, 0x41, 0x49, 0x01]); + let io: Arc = Arc::new(BufferScheduler::new(data)); + let scheduler = BitpackedScheduler::new(5, 16, 2, false); + + let decoder = scheduler + .schedule_ranges(&[1..3, 5..9], &io, 0) + .await + .unwrap(); + let decoded = decoder.decode(3, 3).unwrap(); + + // The scheduled rows are [2, 3, 6, 7, 8, 9]. Skipping across the first + // Bytes response must begin in the middle of the second response. + assert_eq!(fixed_width_values::(decoded), vec![7, 8, 9]); + } + + #[test_log::test(tokio::test)] + async fn test_bitpacked_scheduler_signed_byte_aligned() { + // Legacy 16-bit little-endian encoding of [513, -1, 4660, -32768]. + let data = Bytes::from_static(&[0x01, 0x02, 0xFF, 0xFF, 0x34, 0x12, 0x00, 0x80]); + let io: Arc = Arc::new(BufferScheduler::new(data)); + let scheduler = BitpackedScheduler::new(16, 32, 0, true); + + let decoder = scheduler.schedule_ranges(&[1..4], &io, 0).await.unwrap(); + let decoded = decoder.decode(0, 3).unwrap(); + + assert_eq!(fixed_width_values::(decoded), vec![-1, 4660, -32768]); + } + + #[test_log::test(tokio::test)] + async fn test_bitpacked_for_non_negative_scheduler_chunk_boundary() { + const BIT_WIDTH: usize = 7; + const NUM_CHUNKS: usize = 2; + const WORDS_PER_CHUNK: usize = ELEMS_PER_CHUNK as usize * BIT_WIDTH / u32::BITS as usize; + + let values = (0..ELEMS_PER_CHUNK as usize * NUM_CHUNKS) + .map(|index| ((index * 13 + 7) % 127) as u32) + .collect::>(); + let mut packed = vec![0_u32; WORDS_PER_CHUNK * NUM_CHUNKS]; + for chunk_index in 0..NUM_CHUNKS { + let value_start = chunk_index * ELEMS_PER_CHUNK as usize; + let word_start = chunk_index * WORDS_PER_CHUNK; + // SAFETY: Both slices have the exact input and output lengths required + // for one 1,024-value chunk at this bit width. + unsafe { + BitPacking::unchecked_pack( + BIT_WIDTH, + &values[value_start..value_start + ELEMS_PER_CHUNK as usize], + &mut packed[word_start..word_start + WORDS_PER_CHUNK], + ); + } + } + + let mut data = vec![0xFA, 0xCE, 0x01]; + data.extend_from_slice(cast_slice(&packed)); + let io: Arc = Arc::new(BufferScheduler::new(Bytes::from(data))); + let scheduler = BitpackedForNonNegScheduler::new(BIT_WIDTH as u64, 32, 3); + + let decoder = scheduler + .schedule_ranges(&[1020..1028, 1530..1533], &io, 0) + .await + .unwrap(); + let decoded = decoder.decode(2, 7).unwrap(); + let expected = (1022..1028) + .chain(1530..1531) + .map(|index| values[index]) + .collect::>(); + + assert_eq!(fixed_width_values::(decoded), expected); + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/block.rs b/vendor/lance-encoding/src/array_encoding/physical/block.rs new file mode 100644 index 000000000..3bbd966e2 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/block.rs @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow_schema::DataType; + +use crate::{ + data::{BlockInfo, DataBlock, OpaqueBlock}, + encoder::{ArrayEncoder, EncodedArray}, + encodings::physical::block::{CompressedBufferEncoder, CompressionConfig, CompressionScheme}, + format::ProtobufUtils, +}; + +use lance_core::Result; + +impl ArrayEncoder for CompressedBufferEncoder { + fn encode( + &self, + data: DataBlock, + _data_type: &DataType, + buffer_index: &mut u32, + ) -> Result { + let uncompressed_data = data.as_fixed_width().unwrap(); + + let mut compressed_buf = Vec::with_capacity(uncompressed_data.data.len()); + self.compressor + .compress(&uncompressed_data.data, &mut compressed_buf)?; + + let compressed_data = DataBlock::Opaque(OpaqueBlock { + buffers: vec![compressed_buf.into()], + num_values: uncompressed_data.num_values, + block_info: BlockInfo::new(), + }); + + let comp_buf_index = *buffer_index; + *buffer_index += 1; + + let encoding = ProtobufUtils::flat_encoding( + uncompressed_data.bits_per_value, + comp_buf_index, + Some(CompressionConfig::new(CompressionScheme::Zstd, None)), + ); + + Ok(EncodedArray { + data: compressed_data, + encoding, + }) + } +} + +#[cfg(test)] +mod tests { + use crate::{buffer::LanceBuffer, data::FixedWidthDataBlock}; + + use super::*; + + #[test] + fn test_compressed_buffer_encoder() { + let encoder = CompressedBufferEncoder::default(); + let data = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(vec![0, 1, 2, 3, 4, 5, 6, 7]), + num_values: 8, + block_info: BlockInfo::new(), + }); + + let mut buffer_index = 0; + let encoded_array_result = encoder.encode(data, &DataType::Int64, &mut buffer_index); + assert!(encoded_array_result.is_ok(), "{:?}", encoded_array_result); + let encoded_array = encoded_array_result.unwrap(); + assert_eq!(encoded_array.data.num_values(), 8); + let buffers = encoded_array.data.into_buffers(); + assert_eq!(buffers.len(), 1); + assert!(buffers[0].len() < 64 * 8); + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/dictionary.rs b/vendor/lance-encoding/src/array_encoding/physical/dictionary.rs new file mode 100644 index 000000000..9d02ed782 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/dictionary.rs @@ -0,0 +1,651 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; +use std::vec; + +use arrow_array::builder::{ArrayBuilder, StringBuilder}; +use arrow_array::cast::AsArray; +use arrow_array::types::UInt8Type; +use arrow_array::{ + Array, ArrayRef, DictionaryArray, StringArray, UInt8Array, make_array, new_null_array, +}; +use arrow_schema::DataType; +use futures::{FutureExt, future::BoxFuture}; +use lance_arrow::DataTypeExt; +use lance_core::{Error, Result}; +use std::collections::HashMap; + +use crate::array_encoding::logical::primitive::PrimitiveFieldDecoder; +use crate::buffer::LanceBuffer; +use crate::data::{ + BlockInfo, DataBlock, DictionaryDataBlock, FixedWidthDataBlock, NullableDataBlock, + VariableWidthBlock, +}; +use crate::decoder::LogicalPageDecoder; +use crate::format::ProtobufUtils; +use crate::{ + EncodingsIo, + decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, +}; + +#[derive(Debug)] +pub struct DictionaryPageScheduler { + indices_scheduler: Arc, + items_scheduler: Arc, + // The number of items in the dictionary + num_dictionary_items: u32, + // If true, decode the dictionary items. If false, leave them dictionary encoded (e.g. the + // output type is probably a dictionary type) + should_decode_dict: bool, +} + +impl DictionaryPageScheduler { + pub fn new( + indices_scheduler: Arc, + items_scheduler: Arc, + num_dictionary_items: u32, + should_decode_dict: bool, + ) -> Self { + Self { + indices_scheduler, + items_scheduler, + num_dictionary_items, + should_decode_dict, + } + } +} + +impl PageScheduler for DictionaryPageScheduler { + fn schedule_ranges( + &self, + ranges: &[std::ops::Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + // We want to decode indices and items + // e.g. indices [0, 1, 2, 0, 1, 0] + // items (dictionary) ["abcd", "hello", "apple"] + // This will map to ["abcd", "hello", "apple", "abcd", "hello", "abcd"] + // We decode all the items during scheduling itself + // These are used to rebuild the string later + + // Schedule indices for decoding + let indices_page_decoder = + self.indices_scheduler + .schedule_ranges(ranges, scheduler, top_level_row); + + // Schedule items for decoding + let items_range = 0..(self.num_dictionary_items as u64); + let items_page_decoder = self.items_scheduler.schedule_ranges( + std::slice::from_ref(&items_range), + scheduler, + top_level_row, + ); + + let copy_size = self.num_dictionary_items as u64; + + if self.should_decode_dict { + tokio::spawn(async move { + let items_decoder: Arc = + Arc::from(items_page_decoder.await?); + + let mut primitive_wrapper = PrimitiveFieldDecoder::new_from_data( + items_decoder.clone(), + DataType::Utf8, + copy_size, + false, + ); + + // Decode all items + let drained_task = primitive_wrapper.drain(copy_size)?; + let items_decode_task = drained_task.task; + let (decoded_dict, _) = items_decode_task.decode()?; + + let indices_decoder: Box = indices_page_decoder.await?; + + Ok(Box::new(DictionaryPageDecoder { + decoded_dict, + indices_decoder, + }) as Box) + }) + .map(|join_handle| join_handle.unwrap()) + .boxed() + } else { + let num_dictionary_items = self.num_dictionary_items; + tokio::spawn(async move { + let items_decoder: Arc = + Arc::from(items_page_decoder.await?); + + let decoded_dict = items_decoder + .decode(0, num_dictionary_items as u64)? + .clone(); + + let indices_decoder = indices_page_decoder.await?; + + Ok(Box::new(DirectDictionaryPageDecoder { + decoded_dict, + indices_decoder, + }) as Box) + }) + .map(|join_handle| join_handle.unwrap()) + .boxed() + } + } +} + +struct DirectDictionaryPageDecoder { + decoded_dict: DataBlock, + indices_decoder: Box, +} + +impl PrimitivePageDecoder for DirectDictionaryPageDecoder { + fn variable_width_bytes(&self, _rows_to_skip: u64, _num_rows: u64) -> Result> { + // The decoded batch shares this page's dictionary values; Arrow + // concatenation may merge each page's values into one array, so charge + // the whole dictionary once per page. + Ok(Some(self.decoded_dict.data_size())) + } + + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + let indices = self + .indices_decoder + .decode(rows_to_skip, num_rows)? + .as_fixed_width() + .unwrap(); + let dict = self.decoded_dict.clone(); + Ok(DataBlock::Dictionary(DictionaryDataBlock { + indices, + dictionary: Box::new(dict), + })) + } +} + +struct DictionaryPageDecoder { + decoded_dict: Arc, + indices_decoder: Box, +} + +impl PrimitivePageDecoder for DictionaryPageDecoder { + fn variable_width_bytes(&self, rows_to_skip: u64, num_rows: u64) -> Result> { + // Decoding materializes each row's dictionary value into a plain string + // array. The u8 indices are cheap to decode, so compute the exact value + // bytes by summing each requested row's dictionary entry length (index 0 + // is the in-band null and contributes nothing). + let Some(strings) = self.decoded_dict.as_any().downcast_ref::() else { + return Ok(None); + }; + let indices = self.indices_decoder.decode(rows_to_skip, num_rows)?; + let indices = match &indices { + DataBlock::FixedWidth(fixed) => fixed, + DataBlock::Nullable(nullable) => { + let Some(fixed) = nullable.data.as_fixed_width_ref() else { + return Ok(None); + }; + fixed + } + _ => return Ok(None), + }; + if indices.bits_per_value != 8 { + return Ok(None); + } + let mut bytes = 0u64; + for &code in indices.data.as_ref() { + if code != 0 { + bytes += strings.value_length((code - 1) as usize) as u64; + } + } + Ok(Some(bytes)) + } + + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + // Decode the indices + let indices_data = self.indices_decoder.decode(rows_to_skip, num_rows)?; + + let indices_array = make_array(indices_data.into_arrow(DataType::UInt8, false)?); + let indices_array = indices_array.as_primitive::(); + + let dictionary = self.decoded_dict.clone(); + + let adjusted_indices: UInt8Array = indices_array + .iter() + .map(|x| match x { + Some(0) => None, + Some(x) => Some(x - 1), + None => None, + }) + .collect(); + + // Build dictionary array using indices and items + let dict_array = + DictionaryArray::::try_new(adjusted_indices, dictionary).unwrap(); + let string_array = arrow_cast::cast(&dict_array, &DataType::Utf8).unwrap(); + let string_array = string_array.as_any().downcast_ref::().unwrap(); + + let null_buffer = string_array.nulls().map(|n| n.buffer().clone()); + let offsets_buffer = string_array.offsets().inner().inner().clone(); + let bytes_buffer = string_array.values().clone(); + + let string_data = DataBlock::VariableWidth(VariableWidthBlock { + bits_per_offset: 32, + data: LanceBuffer::from(bytes_buffer), + offsets: LanceBuffer::from(offsets_buffer), + num_values: num_rows, + block_info: BlockInfo::new(), + }); + if let Some(nulls) = null_buffer { + Ok(DataBlock::Nullable(NullableDataBlock { + data: Box::new(string_data), + nulls: LanceBuffer::from(nulls), + block_info: BlockInfo::new(), + })) + } else { + Ok(string_data) + } + } +} + +/// An encoder for data that is already dictionary encoded. Stores the +/// data as a dictionary encoding. +#[derive(Debug)] +pub struct AlreadyDictionaryEncoder { + indices_encoder: Box, + items_encoder: Box, +} + +impl AlreadyDictionaryEncoder { + pub fn new( + indices_encoder: Box, + items_encoder: Box, + ) -> Self { + Self { + indices_encoder, + items_encoder, + } + } +} + +impl ArrayEncoder for AlreadyDictionaryEncoder { + fn encode( + &self, + data: DataBlock, + data_type: &DataType, + buffer_index: &mut u32, + ) -> Result { + let DataType::Dictionary(key_type, value_type) = data_type else { + panic!("Expected dictionary type"); + }; + + let dict_data = match data { + DataBlock::Dictionary(dict_data) => dict_data, + DataBlock::AllNull(all_null) => { + // In 2.1 this won't happen, kind of annoying to materialize a bunch of nulls + let indices = UInt8Array::from(vec![0; all_null.num_values as usize]); + let indices = arrow_cast::cast(&indices, key_type.as_ref()).unwrap(); + let indices = indices.into_data(); + let values = new_null_array(value_type, 1); + DictionaryDataBlock { + indices: FixedWidthDataBlock { + bits_per_value: key_type.byte_width() as u64 * 8, + data: LanceBuffer::from(indices.buffers()[0].clone()), + num_values: all_null.num_values, + block_info: BlockInfo::new(), + }, + dictionary: Box::new(DataBlock::from_array(values)), + } + } + _ => panic!("Expected dictionary data"), + }; + let declared_key_bits = key_type.byte_width() as u64 * 8; + if dict_data.indices.bits_per_value != declared_key_bits { + return Err(Error::invalid_input(format!( + "dictionary indices use {} bits but the declared {} key type uses {} bits; the normalized dictionary has {} values", + dict_data.indices.bits_per_value, + key_type, + declared_key_bits, + dict_data.dictionary.num_values() + ))); + } + let num_dictionary_items = dict_data.dictionary.num_values() as u32; + + let encoded_indices = self.indices_encoder.encode( + DataBlock::FixedWidth(dict_data.indices), + key_type, + buffer_index, + )?; + let encoded_items = + self.items_encoder + .encode(*dict_data.dictionary, value_type, buffer_index)?; + + let encoded = DataBlock::Dictionary(DictionaryDataBlock { + dictionary: Box::new(encoded_items.data), + indices: encoded_indices.data.as_fixed_width().unwrap(), + }); + + let encoding = ProtobufUtils::dict_encoding( + encoded_indices.encoding, + encoded_items.encoding, + num_dictionary_items, + ); + + Ok(EncodedArray { + data: encoded, + encoding, + }) + } +} + +#[derive(Debug)] +pub struct DictionaryEncoder { + indices_encoder: Box, + items_encoder: Box, +} + +impl DictionaryEncoder { + pub fn new( + indices_encoder: Box, + items_encoder: Box, + ) -> Self { + Self { + indices_encoder, + items_encoder, + } + } +} + +fn encode_dict_indices_and_items(string_array: &StringArray) -> (ArrayRef, ArrayRef) { + let mut arr_hashmap: HashMap<&str, u8> = HashMap::new(); + // We start with a dict index of 1 because the value 0 is reserved for nulls + // The dict indices are adjusted by subtracting 1 later during decode + let mut curr_dict_index = 1; + let total_capacity = string_array.len(); + + let mut dict_indices = Vec::with_capacity(total_capacity); + let mut dict_builder = StringBuilder::new(); + + for i in 0..string_array.len() { + if !string_array.is_valid(i) { + // null value + dict_indices.push(0); + continue; + } + + let st = string_array.value(i); + + let hashmap_entry = *arr_hashmap.entry(st).or_insert(curr_dict_index); + dict_indices.push(hashmap_entry); + + // if item didn't exist in the hashmap, add it to the dictionary + // and increment the dictionary index + if hashmap_entry == curr_dict_index { + dict_builder.append_value(st); + curr_dict_index += 1; + } + } + + let array_dict_indices = Arc::new(UInt8Array::from(dict_indices)) as ArrayRef; + + // If there is an empty dictionary: + // Either there is an array of nulls or an empty array altogether + // In this case create the dictionary with a single null element + // Because decoding [] is not currently supported by the binary decoder + if dict_builder.is_empty() { + dict_builder.append_option(Option::<&str>::None); + } + + let dict_elements = dict_builder.finish(); + let array_dict_elements = arrow_cast::cast(&dict_elements, &DataType::Utf8).unwrap(); + + (array_dict_indices, array_dict_elements) +} + +impl ArrayEncoder for DictionaryEncoder { + fn encode( + &self, + data: DataBlock, + data_type: &DataType, + buffer_index: &mut u32, + ) -> Result { + if !matches!(data_type, DataType::Utf8) { + return Err(Error::invalid_input_source( + format!( + "DictionaryEncoder only supports string arrays but got {}", + data_type + ) + .into(), + )); + } + // We only support string arrays for now + let str_data = make_array(data.into_arrow(DataType::Utf8, false)?); + + let (index_array, items_array) = encode_dict_indices_and_items(str_data.as_string()); + let dict_size = items_array.len() as u32; + let index_data = DataBlock::from(index_array); + let items_data = DataBlock::from(items_array); + + let encoded_indices = + self.indices_encoder + .encode(index_data, &DataType::UInt8, buffer_index)?; + + let encoded_items = self + .items_encoder + .encode(items_data, &DataType::Utf8, buffer_index)?; + + let encoded_data = DataBlock::Dictionary(DictionaryDataBlock { + indices: encoded_indices.data.as_fixed_width().unwrap(), + dictionary: Box::new(encoded_items.data), + }); + + let encoding = ProtobufUtils::dict_encoding( + encoded_indices.encoding, + encoded_items.encoding, + dict_size, + ); + + Ok(EncodedArray { + data: encoded_data, + encoding, + }) + } +} + +#[cfg(test)] +mod tests { + + use arrow_array::{ + ArrayRef, DictionaryArray, StringArray, UInt8Array, + builder::{LargeStringBuilder, StringBuilder}, + }; + use arrow_schema::{DataType, Field}; + use std::{collections::HashMap, sync::Arc, vec}; + + use crate::testing::{ + TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data, + }; + use rstest::rstest; + + use super::encode_dict_indices_and_items; + + // These tests cover the case where we opportunistically convert some (or all) pages of + // a string column into dictionaries (and decode on read) + + #[test] + fn test_encode_dict_nulls() { + // Null entries in string arrays should be adjusted + let string_array = Arc::new(StringArray::from(vec![ + None, + Some("foo"), + Some("bar"), + Some("bar"), + None, + Some("foo"), + None, + None, + ])); + let (dict_indices, dict_items) = encode_dict_indices_and_items(&string_array); + + let expected_indices = Arc::new(UInt8Array::from(vec![0, 1, 2, 2, 0, 1, 0, 0])) as ArrayRef; + let expected_items = Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef; + assert_eq!(&dict_indices, &expected_indices); + assert_eq!(&dict_items, &expected_items); + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_random_dictionary( + #[values( + DataType::Utf8, + DataType::Binary, + DataType::LargeBinary, + DataType::LargeUtf8 + )] + data_type: DataType, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let nullable = matches!(data_type, DataType::LargeBinary | DataType::LargeUtf8); + let field = Field::new("", data_type, nullable); + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_utf8() { + let string_array = StringArray::from(vec![Some("abc"), Some("de"), None, Some("fgh")]); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(1..3) + .with_indices(vec![1, 3]); + check_round_trip_encoding_of_data( + vec![Arc::new(string_array)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_sliced_utf8() { + let string_array = StringArray::from(vec![Some("abc"), Some("de"), None, Some("fgh")]); + let string_array = string_array.slice(1, 3); + + let test_cases = TestCases::default() + .with_range(0..1) + .with_range(0..2) + .with_range(1..2); + check_round_trip_encoding_of_data( + vec![Arc::new(string_array)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_empty_strings() { + // Scenario 1: Some strings are empty + + let values = [Some("abc"), Some(""), None]; + // Test empty list at beginning, middle, and end + for order in [[0, 1, 2], [1, 0, 2], [2, 0, 1]] { + let mut string_builder = StringBuilder::new(); + for idx in order { + string_builder.append_option(values[idx]); + } + let string_array = Arc::new(string_builder.finish()); + let test_cases = TestCases::default() + .with_indices(vec![1]) + .with_indices(vec![0]) + .with_indices(vec![2]); + check_round_trip_encoding_of_data( + vec![string_array.clone()], + &test_cases, + HashMap::new(), + ) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()) + .await; + } + + // Scenario 2: All strings are empty + + // When encoding an array of empty strings there are no bytes to encode + // which is strange and we want to ensure we handle it + let string_array = Arc::new(StringArray::from(vec![Some(""), None, Some("")])); + + let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]); + check_round_trip_encoding_of_data(vec![string_array.clone()], &test_cases, HashMap::new()) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await; + } + + #[test_log::test(tokio::test)] + #[ignore] // This test is quite slow in debug mode + async fn test_jumbo_string() { + // This is an overflow test. We have a list of lists where each list + // has 1Mi items. We encode 5000 of these lists and so we have over 4Gi in the + // offsets range + let mut string_builder = LargeStringBuilder::new(); + // a 1 MiB string + let giant_string = String::from_iter((0..(1024 * 1024)).map(|_| '0')); + for _ in 0..5000 { + string_builder.append_option(Some(&giant_string)); + } + let giant_array = Arc::new(string_builder.finish()) as ArrayRef; + let arrs = vec![giant_array]; + + // // We can't validate because our validation relies on concatenating all input arrays + let test_cases = TestCases::default().without_validation(); + check_round_trip_encoding_of_data(arrs, &test_cases, HashMap::new()).await; + } + + // These tests cover the case where the input is already dictionary encoded + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_random_dictionary_input( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let dict_field = Field::new( + "", + DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)), + false, + ); + check_basic_random_case(dict_field, encoding, page_size, use_slicing).await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_already_dictionary() { + let values = StringArray::from_iter_values(["a", "bb", "ccc"]); + let indices = UInt8Array::from(vec![0, 1, 2, 0, 1, 2, 0, 1, 2]); + let dict_array = DictionaryArray::new(indices, Arc::new(values)); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(1..3) + .with_range(2..4) + .with_indices(vec![1]) + .with_indices(vec![2]); + check_round_trip_encoding_of_data(vec![Arc::new(dict_array)], &test_cases, HashMap::new()) + .await; + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/fixed_size_binary.rs b/vendor/lance-encoding/src/array_encoding/physical/fixed_size_binary.rs new file mode 100644 index 000000000..a5ed222e1 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/fixed_size_binary.rs @@ -0,0 +1,366 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_buffer::ScalarBuffer; +use futures::{FutureExt, future::BoxFuture}; +use lance_core::Result; + +use crate::{ + EncodingsIo, + buffer::LanceBuffer, + data::{BlockInfo, DataBlock, VariableWidthBlock}, + decoder::{PageScheduler, PrimitivePageDecoder}, +}; + +/// A scheduler for fixed size binary data +#[derive(Debug)] +pub struct FixedSizeBinaryPageScheduler { + bytes_scheduler: Box, + byte_width: u32, + bytes_per_offset: u32, +} + +impl FixedSizeBinaryPageScheduler { + pub fn new( + bytes_scheduler: Box, + byte_width: u32, + bytes_per_offset: u32, + ) -> Self { + Self { + bytes_scheduler, + byte_width, + bytes_per_offset, + } + } +} + +impl PageScheduler for FixedSizeBinaryPageScheduler { + fn schedule_ranges( + &self, + ranges: &[std::ops::Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + let expanded_ranges = ranges + .iter() + .map(|range| { + (range.start * self.byte_width as u64)..(range.end * self.byte_width as u64) + }) + .collect::>(); + + let bytes_page_decoder = + self.bytes_scheduler + .schedule_ranges(&expanded_ranges, scheduler, top_level_row); + + let byte_width = self.byte_width as u64; + let bytes_per_offset = self.bytes_per_offset; + + async move { + let bytes_decoder = bytes_page_decoder.await?; + Ok(Box::new(FixedSizeBinaryDecoder { + bytes_decoder, + byte_width, + bytes_per_offset, + }) as Box) + } + .boxed() + } +} + +pub struct FixedSizeBinaryDecoder { + bytes_decoder: Box, + byte_width: u64, + bytes_per_offset: u32, +} + +impl PrimitivePageDecoder for FixedSizeBinaryDecoder { + fn variable_width_bytes(&self, _rows_to_skip: u64, num_rows: u64) -> Result> { + Ok(Some(num_rows * self.byte_width)) + } + + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + let rows_to_skip = rows_to_skip * self.byte_width; + let num_bytes = num_rows * self.byte_width; + let bytes = self.bytes_decoder.decode(rows_to_skip, num_bytes)?; + let bytes = bytes.as_fixed_width().unwrap(); + debug_assert_eq!(bytes.bits_per_value, self.byte_width * 8); + + let offsets_buffer = match self.bytes_per_offset { + 8 => { + let offsets_vec = (0..(num_rows + 1)) + .map(|i| i * self.byte_width) + .collect::>(); + + ScalarBuffer::from(offsets_vec).into_inner() + } + 4 => { + let offsets_vec = (0..(num_rows as u32 + 1)) + .map(|i| i * self.byte_width as u32) + .collect::>(); + + ScalarBuffer::from(offsets_vec).into_inner() + } + _ => panic!("Unsupported offsets type"), + }; + + let string_data = DataBlock::VariableWidth(VariableWidthBlock { + bits_per_offset: (self.bytes_per_offset * 8) as u8, + data: bytes.data, + num_values: num_rows, + offsets: LanceBuffer::from(offsets_buffer), + block_info: BlockInfo::new(), + }); + + Ok(string_data) + } +} + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use arrow_array::{ + Array, ArrayRef, FixedSizeBinaryArray, LargeStringArray, StringArray, + builder::LargeStringBuilder, + }; + use arrow_buffer::Buffer; + use arrow_data::ArrayData; + use arrow_schema::{DataType, Field}; + + use crate::array_encoding::physical::fixed_size_binary::FixedSizeBinaryDecoder; + use crate::data::{DataBlock, FixedWidthDataBlock}; + use crate::decoder::PrimitivePageDecoder; + use crate::testing::{ + TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data, + }; + use rstest::rstest; + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_fixed_size_random( + #[values( + DataType::Utf8, + DataType::Binary, + DataType::LargeBinary, + DataType::LargeUtf8 + )] + data_type: DataType, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let nullable = matches!(data_type, DataType::LargeBinary | DataType::LargeUtf8); + let field = Field::new("", data_type, nullable); + // This test only generates fixed-size binary arrays for Utf8 and Binary. + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_fixed_size_utf8() { + let string_array = StringArray::from(vec![ + Some("abc"), + Some("def"), + Some("ghi"), + Some("jkl"), + Some("mno"), + ]); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(1..3) + .with_indices(vec![0, 1, 3, 4]); + + check_round_trip_encoding_of_data( + vec![Arc::new(string_array)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_fixed_size_with_nulls_utf8() { + let string_array = + LargeStringArray::from(vec![Some("abc"), None, Some("ghi"), None, Some("mno")]); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(1..3) + .with_indices(vec![0, 1, 3, 4]); + + check_round_trip_encoding_of_data( + vec![Arc::new(string_array)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_fixed_size_sliced_utf8() { + let string_array = StringArray::from(vec![Some("abc"), Some("def"), None, Some("fgh")]); + let string_array = string_array.slice(1, 3); + + let test_cases = TestCases::default() + .with_range(0..1) + .with_range(0..2) + .with_range(1..2); + check_round_trip_encoding_of_data( + vec![Arc::new(string_array)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_fixed_size_empty_strings() { + // All strings are empty + + // When encoding an array of empty strings there are no bytes to encode + // which is strange and we want to ensure we handle it + let string_array = Arc::new(StringArray::from(vec![Some(""), None, Some("")])); + + let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]); + check_round_trip_encoding_of_data(vec![string_array.clone()], &test_cases, HashMap::new()) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await; + } + + #[test_log::test(tokio::test)] + #[ignore] // This test is quite slow in debug mode + async fn test_jumbo_string() { + // This is an overflow test. We have a list of lists where each list + // has 1Mi items. We encode 5000 of these lists and so we have over 4Gi in the + // offsets range + let mut string_builder = LargeStringBuilder::new(); + // a 1 MiB string + let giant_string = String::from_iter((0..(1024 * 1024)).map(|_| '0')); + for _ in 0..5000 { + string_builder.append_option(Some(&giant_string)); + } + let giant_array = Arc::new(string_builder.finish()) as ArrayRef; + let arrs = vec![giant_array]; + + // // We can't validate because our validation relies on concatenating all input arrays + let test_cases = TestCases::default().without_validation(); + check_round_trip_encoding_of_data(arrs, &test_cases, HashMap::new()).await; + } + + struct FixedWidthCloningPageDecoder { + data_block: FixedWidthDataBlock, + } + + impl PrimitivePageDecoder for FixedWidthCloningPageDecoder { + // clone the given data block as decoded data block + fn decode( + &self, + _rows_to_skip: u64, + _num_rows: u64, + ) -> lance_core::error::Result { + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: self.data_block.data.deep_copy(), + bits_per_value: self.data_block.bits_per_value, + num_values: self.data_block.num_values, + block_info: self.data_block.block_info.clone(), + })) + } + } + + #[test] + fn test_fixed_size_binary_decoder() { + let values: [u8; 6] = *b"aaabbb"; + let num_values = 2u64; + let byte_width = 3; + let array_data = ArrayData::builder(DataType::FixedSizeBinary(byte_width)) + .len(num_values as usize) + .add_buffer(Buffer::from(&values[..])) + .build() + .unwrap(); + let fixed_size_binary_array = FixedSizeBinaryArray::from(array_data); + let arrays = vec![Arc::new(fixed_size_binary_array) as ArrayRef]; + let fixed_width_data_block = DataBlock::from_arrays(&arrays, num_values); + assert_eq!(fixed_width_data_block.name(), "FixedWidth"); + + let bytes_decoder = FixedWidthCloningPageDecoder { + data_block: fixed_width_data_block.as_fixed_width().unwrap(), + }; + let decoder = FixedSizeBinaryDecoder { + bytes_decoder: Box::new(bytes_decoder), + byte_width: byte_width as u64, + bytes_per_offset: 4, // 32-bits offset binary + }; + + let decoded_binary = decoder.decode(0, num_values).unwrap(); + let maybe_data = decoded_binary.into_arrow(DataType::Utf8, true); + assert!(maybe_data.is_ok()); + let data = maybe_data.unwrap(); + let string_array = StringArray::from(data); + assert_eq!(string_array.len(), num_values as usize); + assert_eq!(string_array.value(0), "aaa"); + assert_eq!(string_array.value(1), "bbb"); + } + + #[test] + fn test_fixed_size_binary_pages_split_by_byte_budget() { + use crate::array_encoding::logical::primitive::PrimitiveFieldDecoder; + use crate::array_encoding::logical::r#struct::SimpleStructDecoder; + use crate::decoder::{DecoderReady, DrainLimit, LogicalPageDecoder}; + use arrow_schema::{Field as ArrowField, Fields}; + use lance_core::Result; + use std::collections::VecDeque; + + #[derive(Debug)] + struct NeverDecodedStub; + + impl PrimitivePageDecoder for NeverDecodedStub { + fn decode(&self, _rows_to_skip: u64, _num_rows: u64) -> Result { + unreachable!("byte accounting must not decode any values") + } + } + + // Two pages of 3 rows x 10 bytes each; each page fits the budget alone + // but the pair does not. + let fields = Fields::from(vec![ArrowField::new("value", DataType::Utf8, false)]); + let mut root = SimpleStructDecoder::new(fields, 6); + for _ in 0..2 { + root.accept_child(DecoderReady { + decoder: Box::new(PrimitiveFieldDecoder::new_from_data( + Arc::new(FixedSizeBinaryDecoder { + bytes_decoder: Box::new(NeverDecodedStub), + byte_width: 10, + bytes_per_offset: 4, + }), + DataType::Utf8, + 3, + false, + )), + path: VecDeque::from([0]), + }) + .unwrap(); + } + + // Budget for page 1 plus two rows of page 2. + let limit = root.max_rows_to_drain(6, 50).unwrap(); + assert_eq!(limit, DrainLimit { rows: 5, bytes: 50 }); + + // Budget for exactly one page: the batch stops at the page boundary. + let limit = root.max_rows_to_drain(6, 30).unwrap(); + assert_eq!(limit, DrainLimit { rows: 3, bytes: 30 }); + + // Both pages fit a large enough budget. + let limit = root.max_rows_to_drain(6, 60).unwrap(); + assert_eq!(limit, DrainLimit { rows: 6, bytes: 60 }); + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/fixed_size_list.rs b/vendor/lance-encoding/src/array_encoding/physical/fixed_size_list.rs new file mode 100644 index 000000000..49d2dbff1 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/fixed_size_list.rs @@ -0,0 +1,275 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_schema::DataType; +use futures::{FutureExt, future::BoxFuture}; +use lance_core::Result; +use log::trace; + +use crate::{ + EncodingsIo, + data::{DataBlock, FixedSizeListBlock}, + decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, + format::ProtobufUtils, +}; + +/// A scheduler for fixed size lists of primitive values +/// +/// This scheduler is, itself, primitive +#[derive(Debug)] +pub struct FixedListScheduler { + items_scheduler: Box, + dimension: u32, +} + +impl FixedListScheduler { + pub fn new(items_scheduler: Box, dimension: u32) -> Self { + Self { + items_scheduler, + dimension, + } + } +} + +impl PageScheduler for FixedListScheduler { + fn schedule_ranges( + &self, + ranges: &[std::ops::Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + let expanded_ranges = ranges + .iter() + .map(|range| (range.start * self.dimension as u64)..(range.end * self.dimension as u64)) + .collect::>(); + trace!( + "Expanding {} fsl ranges across {}..{} to item ranges across {}..{}", + ranges.len(), + ranges[0].start, + ranges[ranges.len() - 1].end, + expanded_ranges[0].start, + expanded_ranges[expanded_ranges.len() - 1].end + ); + let inner_page_decoder = + self.items_scheduler + .schedule_ranges(&expanded_ranges, scheduler, top_level_row); + let dimension = self.dimension; + async move { + let items_decoder = inner_page_decoder.await?; + Ok(Box::new(FixedListDecoder { + items_decoder, + dimension: dimension as u64, + }) as Box) + } + .boxed() + } +} + +pub struct FixedListDecoder { + items_decoder: Box, + dimension: u64, +} + +impl PrimitivePageDecoder for FixedListDecoder { + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + let rows_to_skip = rows_to_skip * self.dimension; + let num_child_rows = num_rows * self.dimension; + let child_data = self.items_decoder.decode(rows_to_skip, num_child_rows)?; + Ok(DataBlock::FixedSizeList(FixedSizeListBlock { + child: Box::new(child_data), + dimension: self.dimension, + })) + } +} + +#[derive(Debug)] +pub struct FslEncoder { + items_encoder: Box, + dimension: u32, +} + +impl FslEncoder { + pub fn new(items_encoder: Box, dimension: u32) -> Self { + Self { + items_encoder, + dimension, + } + } +} + +impl ArrayEncoder for FslEncoder { + fn encode( + &self, + data: DataBlock, + data_type: &DataType, + buffer_index: &mut u32, + ) -> Result { + let inner_type = match data_type { + DataType::FixedSizeList(inner_field, _) => inner_field.data_type().clone(), + _ => panic!("Expected fixed size list data type and got {}", data_type), + }; + let data = data.as_fixed_size_list().unwrap(); + let child = *data.child; + + let encoded_data = self + .items_encoder + .encode(child, &inner_type, buffer_index)?; + + let data = DataBlock::FixedSizeList(FixedSizeListBlock { + child: Box::new(encoded_data.data), + dimension: self.dimension as u64, + }); + + let encoding = + ProtobufUtils::fsl_encoding(self.dimension as u64, encoded_data.encoding, false); + Ok(EncodedArray { data, encoding }) + } +} + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use arrow_array::{FixedSizeListArray, Int32Array, types::Int32Type}; + use arrow_buffer::{BooleanBuffer, NullBuffer}; + use arrow_schema::{DataType, Field}; + use lance_datagen::{ArrayGeneratorExt, RowCount, array, gen_array}; + + use crate::testing::{ + TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data, + }; + use rstest::rstest; + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_value_fsl_primitive( + #[values(DataType::Int8, DataType::Float32, DataType::Float64)] data_type: DataType, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let inner_field = Field::new("item", data_type, true); + let data_type = DataType::FixedSizeList(Arc::new(inner_field), 16); + let field = Field::new("", data_type, false); + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_fsl() { + // [0, NULL], NULL, [4, 5] + let items = Arc::new(Int32Array::from(vec![ + Some(0), + None, + Some(2), + Some(3), + Some(4), + Some(5), + ])); + let items_field = Arc::new(Field::new("item", DataType::Int32, true)); + let list_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, false, true])); + let list = Arc::new(FixedSizeListArray::new( + items_field, + 2, + items, + Some(list_nulls), + )); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_range(0..2) + .with_range(1..3) + .with_indices(vec![0, 1, 2]) + .with_indices(vec![1]) + .with_indices(vec![2]) + .with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![list], &test_cases, HashMap::default()).await; + } + + #[test_log::test(tokio::test)] + #[ignore] + async fn test_simple_wide_fsl() { + let items = gen_array(array::rand::().with_random_nulls(0.1)) + .into_array_rows(RowCount::from(4096)) + .unwrap(); + let items_field = Arc::new(Field::new("item", DataType::Int32, true)); + let list_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, false, true, false])); + let list = Arc::new(FixedSizeListArray::new( + items_field, + 1024, + items, + Some(list_nulls), + )); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_range(0..2) + .with_range(1..3) + .with_indices(vec![0, 1, 2]) + .with_indices(vec![1]) + .with_indices(vec![2]) + .with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![list], &test_cases, HashMap::default()).await; + } + + #[test_log::test(tokio::test)] + async fn test_nested_fsl() { + // [[0, 1], NULL], NULL, [[8, 9], [NULL, 11]] + let items = Arc::new(Int32Array::from(vec![ + Some(0), + Some(1), + None, + None, + None, + None, + None, + None, + Some(8), + Some(9), + None, + Some(11), + ])); + let items_field = Arc::new(Field::new("item", DataType::Int32, true)); + let inner_list_nulls = NullBuffer::new(BooleanBuffer::from(vec![ + true, false, false, false, true, true, + ])); + let inner_list = Arc::new(FixedSizeListArray::new( + items_field.clone(), + 2, + items, + Some(inner_list_nulls), + )); + let inner_list_field = Arc::new(Field::new( + "item", + DataType::FixedSizeList(items_field, 2), + true, + )); + let outer_list_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, false, true])); + let outer_list = Arc::new(FixedSizeListArray::new( + inner_list_field, + 2, + inner_list, + Some(outer_list_nulls), + )); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_range(0..2) + .with_range(1..3) + .with_indices(vec![0, 1, 2]) + .with_indices(vec![2]) + .with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![outer_list], &test_cases, HashMap::default()).await; + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/fsst.rs b/vendor/lance-encoding/src/array_encoding/physical/fsst.rs new file mode 100644 index 000000000..d6e8868ef --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/fsst.rs @@ -0,0 +1,208 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ops::Range, sync::Arc}; + +use arrow_buffer::ScalarBuffer; +use arrow_schema::DataType; +use futures::{FutureExt, future::BoxFuture}; + +use lance_core::Result; + +use crate::{ + EncodingsIo, + buffer::LanceBuffer, + data::{BlockInfo, DataBlock, NullableDataBlock, VariableWidthBlock}, + decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, + format::ProtobufUtils, +}; + +#[derive(Debug)] +pub struct FsstPageScheduler { + inner_scheduler: Box, + symbol_table: LanceBuffer, +} + +impl FsstPageScheduler { + pub fn new(inner_scheduler: Box, symbol_table: LanceBuffer) -> Self { + Self { + inner_scheduler, + symbol_table, + } + } +} + +impl PageScheduler for FsstPageScheduler { + fn schedule_ranges( + &self, + ranges: &[Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + let inner_decoder = self + .inner_scheduler + .schedule_ranges(ranges, scheduler, top_level_row); + let symbol_table = self.symbol_table.clone(); + + async move { + let inner_decoder = inner_decoder.await?; + Ok(Box::new(FsstPageDecoder { + inner_decoder, + symbol_table, + }) as Box) + } + .boxed() + } +} + +struct FsstPageDecoder { + inner_decoder: Box, + symbol_table: LanceBuffer, +} + +impl PrimitivePageDecoder for FsstPageDecoder { + fn variable_width_bytes(&self, rows_to_skip: u64, num_rows: u64) -> Result> { + // FSST decompression expands by at most 8x, so the inner (compressed) + // byte count times 8 is a safe upper bound on the decoded size. + Ok(self + .inner_decoder + .variable_width_bytes(rows_to_skip, num_rows)? + .map(|compressed_bytes| compressed_bytes.saturating_mul(8))) + } + + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + let compressed_data = self.inner_decoder.decode(rows_to_skip, num_rows)?; + let (string_data, nulls) = match compressed_data { + DataBlock::Nullable(nullable) => { + let data = nullable.data.as_variable_width().unwrap(); + Result::Ok((data, Some(nullable.nulls))) + } + DataBlock::VariableWidth(variable) => Ok((variable, None)), + _ => panic!("Received non-variable width data from inner decoder"), + }?; + + let offsets = ScalarBuffer::::from(string_data.offsets.into_buffer()); + let bytes = string_data.data.into_buffer(); + + let mut decompressed_offsets = vec![0_i32; offsets.len()]; + let mut decompressed_bytes = vec![0_u8; bytes.len() * 8]; + // Safety: Exposes uninitialized memory but we're about to clobber it + unsafe { + decompressed_bytes.set_len(decompressed_bytes.capacity()); + } + fsst::fsst::decompress( + &self.symbol_table, + &bytes, + &offsets, + &mut decompressed_bytes, + &mut decompressed_offsets, + ) + .map_err(crate::encodings::physical::fsst::map_fsst_error)?; + + // TODO: Change PrimitivePageDecoder to use Vec instead of BytesMut + // since there is no way to get BytesMut from Vec but these copies should be avoidable + // This is not the first time this has happened + let mut offsets_as_bytes_mut = Vec::with_capacity(decompressed_offsets.len()); + let decompressed_offsets = ScalarBuffer::::from(decompressed_offsets); + offsets_as_bytes_mut.extend_from_slice(decompressed_offsets.inner().as_slice()); + + let mut bytes_as_bytes_mut = Vec::with_capacity(decompressed_bytes.len()); + bytes_as_bytes_mut.extend_from_slice(&decompressed_bytes); + + let new_string_data = DataBlock::VariableWidth(VariableWidthBlock { + bits_per_offset: 32, + data: LanceBuffer::from(bytes_as_bytes_mut), + num_values: num_rows, + offsets: LanceBuffer::from(offsets_as_bytes_mut), + block_info: BlockInfo::new(), + }); + + if let Some(nulls) = nulls { + Ok(DataBlock::Nullable(NullableDataBlock { + data: Box::new(new_string_data), + nulls, + block_info: BlockInfo::new(), + })) + } else { + Ok(new_string_data) + } + } +} + +#[derive(Debug)] +pub struct FsstArrayEncoder { + inner_encoder: Box, +} + +impl FsstArrayEncoder { + pub fn new(inner_encoder: Box) -> Self { + Self { inner_encoder } + } +} + +impl ArrayEncoder for FsstArrayEncoder { + fn encode( + &self, + data: DataBlock, + data_type: &DataType, + buffer_index: &mut u32, + ) -> lance_core::Result { + let (data, nulls) = match data { + DataBlock::Nullable(nullable) => { + let data = nullable.data.as_variable_width().unwrap(); + (data, Some(nullable.nulls)) + } + DataBlock::VariableWidth(variable) => (variable, None), + _ => panic!("Expected variable width data block"), + }; + assert_eq!(data.bits_per_offset, 32); + let num_values = data.num_values; + let offsets = data.offsets.borrow_to_typed_slice::(); + let offsets_slice = offsets.as_ref(); + let bytes_data = data.data.into_buffer(); + + let mut dest_offsets = vec![0_i32; offsets_slice.len() * 2]; + let mut dest_values = vec![0_u8; bytes_data.len() * 2]; + let mut symbol_table = vec![0_u8; fsst::fsst::FSST_SYMBOL_TABLE_SIZE]; + + fsst::fsst::compress( + &mut symbol_table, + bytes_data.as_slice(), + offsets_slice, + &mut dest_values, + &mut dest_offsets, + )?; + + let dest_offset = LanceBuffer::reinterpret_vec(dest_offsets); + let dest_values = LanceBuffer::from(dest_values); + let dest_data = DataBlock::VariableWidth(VariableWidthBlock { + bits_per_offset: 32, + data: dest_values, + num_values, + offsets: dest_offset, + block_info: BlockInfo::new(), + }); + + let data_block = if let Some(nulls) = nulls { + DataBlock::Nullable(NullableDataBlock { + data: Box::new(dest_data), + nulls, + block_info: BlockInfo::new(), + }) + } else { + dest_data + }; + + let inner_encoded = self + .inner_encoder + .encode(data_block, data_type, buffer_index)?; + + let encoding = ProtobufUtils::fsst(inner_encoded.encoding, symbol_table); + + Ok(EncodedArray { + data: inner_encoded.data, + encoding, + }) + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/packed_struct.rs b/vendor/lance-encoding/src/array_encoding/physical/packed_struct.rs new file mode 100644 index 000000000..aadd440c9 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/packed_struct.rs @@ -0,0 +1,350 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_schema::{DataType, Fields}; +use bytes::Bytes; +use bytes::BytesMut; +use futures::{FutureExt, future::BoxFuture}; +use lance_arrow::DataTypeExt; +use lance_core::{Error, Result}; + +use crate::data::BlockInfo; +use crate::data::FixedSizeListBlock; +use crate::format::ProtobufUtils; +use crate::{ + EncodingsIo, + buffer::LanceBuffer, + data::{DataBlock, FixedWidthDataBlock, StructDataBlock}, + decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, +}; + +#[derive(Debug)] +pub struct PackedStructPageScheduler { + // We don't actually need these schedulers right now since we decode all the field bytes directly + // But they can be useful if we actually need to use the decoders for the inner fields later + // e.g. once bitpacking is added + _inner_schedulers: Vec>, + fields: Fields, + buffer_offset: u64, +} + +impl PackedStructPageScheduler { + pub fn new( + _inner_schedulers: Vec>, + struct_datatype: DataType, + buffer_offset: u64, + ) -> Self { + let DataType::Struct(fields) = struct_datatype else { + panic!("Struct datatype expected"); + }; + Self { + _inner_schedulers, + fields, + buffer_offset, + } + } +} + +impl PageScheduler for PackedStructPageScheduler { + fn schedule_ranges( + &self, + ranges: &[std::ops::Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + let mut total_bytes_per_row: u64 = 0; + + for field in &self.fields { + let bytes_per_field = field.data_type().byte_width() as u64; + total_bytes_per_row += bytes_per_field; + } + + // Parts of the arrays in a page may be encoded in different encoding tasks + // In that case decoding two different sets of rows can result in the same ranges parameter being passed in + // e.g. we may get ranges[0..2] and ranges[0..2] to decode 4 rows through 2 tasks + // So to get the correct byte ranges we need to know the position of the buffer in the page (i.e. the buffer offset) + // This is computed directly from the buffer stored in the protobuf + let byte_ranges = ranges + .iter() + .map(|range| { + let start = self.buffer_offset + (range.start * total_bytes_per_row); + let end = self.buffer_offset + (range.end * total_bytes_per_row); + start..end + }) + .collect::>(); + + // Directly creates a future to decode the bytes + let bytes = scheduler.submit_request(byte_ranges, top_level_row); + + let copy_struct_fields = self.fields.clone(); + + tokio::spawn(async move { + let bytes = bytes.await?; + + let mut combined_bytes = BytesMut::default(); + for byte_slice in bytes { + combined_bytes.extend_from_slice(&byte_slice); + } + + Ok(Box::new(PackedStructPageDecoder { + data: combined_bytes.freeze(), + fields: copy_struct_fields, + total_bytes_per_row: total_bytes_per_row as usize, + }) as Box) + }) + .map(|join_handle| join_handle.unwrap()) + .boxed() + } +} + +struct PackedStructPageDecoder { + data: Bytes, + fields: Fields, + total_bytes_per_row: usize, +} + +impl PrimitivePageDecoder for PackedStructPageDecoder { + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + // Decoding workflow: + // rows 0-2: {x: [1, 2, 3], y: [4, 5, 6], z: [7, 8, 9]} + // rows 3-5: {x: [10, 11, 12], y: [13, 14, 15], z: [16, 17, 18]} + // packed encoding: [ + // [1, 4, 7, 2, 5, 8, 3, 6, 9], + // [10, 13, 16, 11, 14, 17, 12, 15, 18] + // ] + // suppose bytes_per_field=1, 4, 8 for fields x, y, and z, respectively. + // Then total_bytes_per_row = 13 + // Suppose rows_to_skip=1 and num_rows=2. Then we will slice bytes 13 to 39. + // Now we have [2, 5, 8, 3, 6, 9] + // We rearrange this to get [BytesMut(2, 3), BytesMut(5, 6), BytesMut(8, 9)] as a Vec + // This is used to reconstruct the struct array later + + let bytes_to_skip = (rows_to_skip as usize) * self.total_bytes_per_row; + + let mut children = Vec::with_capacity(self.fields.len()); + + let mut start_index = 0; + + for field in &self.fields { + let bytes_per_field = field.data_type().byte_width(); + let mut field_bytes = Vec::with_capacity(bytes_per_field * num_rows as usize); + + let mut byte_index = start_index; + + for _ in 0..num_rows { + let start = bytes_to_skip + byte_index; + field_bytes.extend_from_slice(&self.data[start..(start + bytes_per_field)]); + byte_index += self.total_bytes_per_row; + } + + start_index += bytes_per_field; + let child_block = FixedWidthDataBlock { + data: LanceBuffer::from(field_bytes), + bits_per_value: bytes_per_field as u64 * 8, + num_values: num_rows, + block_info: BlockInfo::new(), + }; + let child_block = FixedSizeListBlock::from_flat(child_block, field.data_type()); + children.push(child_block); + } + Ok(DataBlock::Struct(StructDataBlock { + children, + block_info: BlockInfo::default(), + validity: None, + })) + } +} + +#[derive(Debug)] +pub struct PackedStructEncoder { + inner_encoders: Vec>, +} + +impl PackedStructEncoder { + pub fn new(inner_encoders: Vec>) -> Self { + Self { inner_encoders } + } +} + +impl ArrayEncoder for PackedStructEncoder { + fn encode( + &self, + data: DataBlock, + data_type: &DataType, + buffer_index: &mut u32, + ) -> Result { + let struct_data = data.as_struct().unwrap(); + + let DataType::Struct(child_types) = data_type else { + panic!("Struct datatype expected"); + }; + + // Encode individual fields + let mut encoded_fields = Vec::with_capacity(struct_data.children.len()); + for ((child, encoder), child_type) in struct_data + .children + .into_iter() + .zip(&self.inner_encoders) + .zip(child_types) + { + encoded_fields.push(encoder.encode(child, child_type.data_type(), &mut 0)?); + } + + let (encoded_data_vec, child_encodings): (Vec<_>, Vec<_>) = encoded_fields + .into_iter() + .map(|field| (field.data, field.encoding)) + .unzip(); + + // Zip together encoded data + // + // We can currently encode both FixedWidth and FixedSizeList. In order + // to encode the latter we "flatten" it converting a FixedSizeList into + // a FixedWidth with very wide items. + let fixed_fields = encoded_data_vec + .into_iter() + .map(|child| match child { + DataBlock::FixedWidth(fixed) => Ok(fixed), + DataBlock::FixedSizeList(fixed_size_list) => { + let flattened = fixed_size_list.try_into_flat().ok_or_else(|| { + Error::invalid_input( + "Packed struct encoder cannot pack nullable fixed-width data blocks", + ) + })?; + Ok(flattened) + } + _ => Err(Error::invalid_input( + "Packed struct encoder currently only implemented for fixed-width data blocks", + )), + }) + .collect::>>()?; + let total_bits_per_value = fixed_fields.iter().map(|f| f.bits_per_value).sum::(); + + let num_values = fixed_fields[0].num_values; + debug_assert!( + fixed_fields + .iter() + .all(|field| field.num_values == num_values) + ); + + let zipped_input = fixed_fields + .into_iter() + .map(|field| (field.data, field.bits_per_value)) + .collect::>(); + let zipped = LanceBuffer::zip_into_one(zipped_input, num_values)?; + + // Create encoding protobuf + let index = *buffer_index; + *buffer_index += 1; + + let packed_data = DataBlock::FixedWidth(FixedWidthDataBlock { + data: zipped, + bits_per_value: total_bits_per_value, + num_values, + block_info: BlockInfo::new(), + }); + + let encoding = ProtobufUtils::packed_struct(child_encodings, index); + + Ok(EncodedArray { + data: packed_data, + encoding, + }) + } +} + +#[cfg(test)] +mod tests { + + use arrow_array::{ArrayRef, Int32Array, StructArray, UInt8Array, UInt64Array}; + use arrow_schema::{DataType, Field, Fields}; + use std::{collections::HashMap, sync::Arc, vec}; + + use crate::testing::{ + TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data, + }; + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_random_packed_struct( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let data_type = DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::UInt64, false), + Field::new("b", DataType::UInt32, false), + ])); + let mut metadata = HashMap::new(); + metadata.insert("packed".to_string(), "true".to_string()); + + let field = Field::new("", data_type, false).with_metadata(metadata); + + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[test_log::test(tokio::test)] + async fn test_specific_packed_struct() { + let array1 = Arc::new(UInt64Array::from(vec![1, 2, 3, 4])); + let array2 = Arc::new(Int32Array::from(vec![5, 6, 7, 8])); + let array3 = Arc::new(UInt8Array::from(vec![9, 10, 11, 12])); + + let struct_array1 = Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("x", DataType::UInt64, false)), + array1.clone() as ArrayRef, + ), + ( + Arc::new(Field::new("y", DataType::Int32, false)), + array2.clone() as ArrayRef, + ), + ( + Arc::new(Field::new("z", DataType::UInt8, false)), + array3.clone() as ArrayRef, + ), + ])); + + let array4 = Arc::new(UInt64Array::from(vec![13, 14, 15, 16])); + let array5 = Arc::new(Int32Array::from(vec![17, 18, 19, 20])); + let array6 = Arc::new(UInt8Array::from(vec![21, 22, 23, 24])); + + let struct_array2 = Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("x", DataType::UInt64, false)), + array4.clone() as ArrayRef, + ), + ( + Arc::new(Field::new("y", DataType::Int32, false)), + array5.clone() as ArrayRef, + ), + ( + Arc::new(Field::new("z", DataType::UInt8, false)), + array6.clone() as ArrayRef, + ), + ])); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..6) + .with_range(1..4) + .with_indices(vec![1, 3, 7]); + + let mut metadata = HashMap::new(); + metadata.insert("packed".to_string(), "true".to_string()); + + check_round_trip_encoding_of_data( + vec![struct_array1, struct_array2], + &test_cases, + metadata, + ) + .await; + } +} diff --git a/vendor/lance-encoding/src/array_encoding/physical/value.rs b/vendor/lance-encoding/src/array_encoding/physical/value.rs new file mode 100644 index 000000000..9f594a002 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/physical/value.rs @@ -0,0 +1,243 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow_schema::DataType; +use bytes::Bytes; +use futures::{FutureExt, future::BoxFuture}; +use log::trace; +use std::ops::Range; +use std::sync::{Arc, Mutex}; + +use crate::buffer::LanceBuffer; +use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock}; +use crate::encodings::physical::block::{ + CompressionConfig, CompressionScheme, GeneralBufferCompressor, +}; +use crate::encodings::physical::value::ValueEncoder; +use crate::format::ProtobufUtils; +use crate::{ + EncodingsIo, + decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, +}; + +use lance_core::{Error, Result}; + +/// Scheduler for a simple encoding where buffers of fixed-size items are stored as-is on disk +#[derive(Debug, Clone, Copy)] +pub struct ValuePageScheduler { + // TODO: do we really support values greater than 2^32 bytes per value? + // I think we want to, in theory, but will need to test this case. + bytes_per_value: u64, + buffer_offset: u64, + buffer_size: u64, + compression_config: CompressionConfig, +} + +impl ValuePageScheduler { + pub fn new( + bytes_per_value: u64, + buffer_offset: u64, + buffer_size: u64, + compression_config: CompressionConfig, + ) -> Self { + Self { + bytes_per_value, + buffer_offset, + buffer_size, + compression_config, + } + } +} + +impl PageScheduler for ValuePageScheduler { + fn schedule_ranges( + &self, + ranges: &[std::ops::Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>> { + let (mut min, mut max) = (u64::MAX, 0); + let byte_ranges = if self.compression_config.scheme == CompressionScheme::None { + ranges + .iter() + .map(|range| { + let start = self.buffer_offset + (range.start * self.bytes_per_value); + let end = self.buffer_offset + (range.end * self.bytes_per_value); + min = min.min(start); + max = max.max(end); + start..end + }) + .collect::>() + } else { + min = self.buffer_offset; + max = self.buffer_offset + self.buffer_size; + // for compressed page, the ranges are always the entire page, + // and it is guaranteed that only one range is passed + vec![Range { + start: min, + end: max, + }] + }; + + trace!( + "Scheduling I/O for {} ranges spread across byte range {}..{}", + byte_ranges.len(), + min, + max + ); + let bytes = scheduler.submit_request(byte_ranges, top_level_row); + let bytes_per_value = self.bytes_per_value; + + let range_offsets = if self.compression_config.scheme != CompressionScheme::None { + ranges + .iter() + .map(|range| { + let start = (range.start * bytes_per_value) as usize; + let end = (range.end * bytes_per_value) as usize; + start..end + }) + .collect::>() + } else { + vec![] + }; + + let compression_config = self.compression_config; + async move { + let bytes = bytes.await?; + + Ok(Box::new(ValuePageDecoder { + bytes_per_value, + data: bytes, + uncompressed_data: Arc::new(Mutex::new(None)), + uncompressed_range_offsets: range_offsets, + compression_config, + }) as Box) + } + .boxed() + } +} + +struct ValuePageDecoder { + bytes_per_value: u64, + data: Vec, + uncompressed_data: Arc>>>, + uncompressed_range_offsets: Vec>, + compression_config: CompressionConfig, +} + +impl ValuePageDecoder { + fn decompress(&self) -> Result> { + // for compressed page, it is guaranteed that only one range is passed + let bytes_u8: Vec = self.data[0].to_vec(); + let buffer_compressor = GeneralBufferCompressor::get_compressor(self.compression_config)?; + let mut uncompressed_bytes: Vec = Vec::new(); + buffer_compressor.decompress(&bytes_u8, &mut uncompressed_bytes)?; + + let mut bytes_in_ranges: Vec = + Vec::with_capacity(self.uncompressed_range_offsets.len()); + for range in &self.uncompressed_range_offsets { + let start = range.start; + let end = range.end; + bytes_in_ranges.push(Bytes::from(uncompressed_bytes[start..end].to_vec())); + } + Ok(bytes_in_ranges) + } + + fn get_uncompressed_bytes(&self) -> Result>>>> { + let mut uncompressed_bytes = self.uncompressed_data.lock().unwrap(); + if uncompressed_bytes.is_none() { + *uncompressed_bytes = Some(self.decompress()?); + } + Ok(Arc::clone(&self.uncompressed_data)) + } + + fn is_compressed(&self) -> bool { + !self.uncompressed_range_offsets.is_empty() + } + + fn decode_buffers<'a>( + &'a self, + buffers: impl IntoIterator, + mut bytes_to_skip: u64, + mut bytes_to_take: u64, + ) -> LanceBuffer { + let mut dest: Option> = None; + + for buf in buffers.into_iter() { + let buf_len = buf.len() as u64; + if bytes_to_skip > buf_len { + bytes_to_skip -= buf_len; + } else { + let bytes_to_take_here = (buf_len - bytes_to_skip).min(bytes_to_take); + bytes_to_take -= bytes_to_take_here; + let start = bytes_to_skip as usize; + let end = start + bytes_to_take_here as usize; + let slice = buf.slice(start..end); + match (&mut dest, bytes_to_take) { + (None, 0) => { + // The entire request is contained in one buffer so we can maybe zero-copy + // if the slice is aligned properly + return LanceBuffer::from_bytes(slice, self.bytes_per_value); + } + (None, _) => { + dest.replace(Vec::with_capacity(bytes_to_take as usize)); + } + _ => {} + } + dest.as_mut().unwrap().extend_from_slice(&slice); + bytes_to_skip = 0; + } + } + LanceBuffer::from(dest.unwrap_or_default()) + } +} + +impl PrimitivePageDecoder for ValuePageDecoder { + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + let bytes_to_skip = rows_to_skip * self.bytes_per_value; + let bytes_to_take = num_rows * self.bytes_per_value; + + let data_buffer = if self.is_compressed() { + let decoding_data = self.get_uncompressed_bytes()?; + let buffers = decoding_data.lock().unwrap(); + self.decode_buffers(buffers.as_ref().unwrap(), bytes_to_skip, bytes_to_take) + } else { + self.decode_buffers(&self.data, bytes_to_skip, bytes_to_take) + }; + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.bytes_per_value * 8, + data: data_buffer, + num_values: num_rows, + block_info: BlockInfo::new(), + })) + } +} + +impl ArrayEncoder for ValueEncoder { + fn encode( + &self, + data: DataBlock, + _data_type: &DataType, + buffer_index: &mut u32, + ) -> Result { + let index = *buffer_index; + *buffer_index += 1; + + let encoding = match &data { + DataBlock::FixedWidth(fixed_width) => Ok(ProtobufUtils::flat_encoding( + fixed_width.bits_per_value, + index, + None, + )), + _ => Err(Error::invalid_input_source( + format!( + "Cannot encode a data block of type {} with ValueEncoder", + data.name() + ) + .into(), + )), + }?; + Ok(EncodedArray { data, encoding }) + } +} diff --git a/vendor/lance-encoding/src/array_encoding/strategy.rs b/vendor/lance-encoding/src/array_encoding/strategy.rs new file mode 100644 index 000000000..20439bad7 --- /dev/null +++ b/vendor/lance-encoding/src/array_encoding/strategy.rs @@ -0,0 +1,1101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::HashMap, env, hash::RandomState, sync::Arc}; + +use arrow_array::{ + Array, ArrayRef, GenericListArray, OffsetSizeTrait, UInt8Array, cast::AsArray, make_array, +}; +use arrow_buffer::{BooleanBuffer, NullBuffer}; +use arrow_data::ArrayData; +use arrow_schema::DataType; +use hyperloglogplus::{HyperLogLog, HyperLogLogPlus}; +use lance_arrow::BLOB_META_KEY; + +use crate::{ + array_encoding::{ + logical::{ + blob::BlobFieldEncoder, list::ListFieldEncoder, primitive::PrimitiveFieldEncoder, + }, + physical::{ + basic::BasicEncoder, + binary::BinaryEncoder, + dictionary::{AlreadyDictionaryEncoder, DictionaryEncoder}, + fixed_size_list::FslEncoder, + fsst::FsstArrayEncoder, + packed_struct::PackedStructEncoder, + }, + }, + constants::{ + COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, PACKED_STRUCT_LEGACY_META_KEY, + PACKED_STRUCT_META_KEY, + }, + encoder::{ + ArrayEncoder, ArrayEncodingStrategy, ColumnIndexSequence, EncodeTask, EncodedColumn, + FieldEncoder, FieldEncodingContext, FieldEncodingStrategy, OutOfLineBuffers, + }, + encodings::{ + logical::r#struct::StructFieldEncoder, + physical::{ + block::{CompressionConfig, CompressionScheme}, + value::ValueEncoder, + }, + }, +}; + +use lance_core::datatypes::{BLOB_DESC_FIELD, Field}; +use lance_core::utils::parse::str_is_truthy; +use lance_core::{Error, Result}; + +/// Field-to-column composition for the `pb::ArrayEncoding` grammar. +#[derive(Debug)] +pub struct ArrayFieldEncodingStrategy { + array_encoding_strategy: Arc, +} + +struct ValidatingFieldEncoder { + inner: Box, + field: Field, + prepared_array: Option, +} + +impl ValidatingFieldEncoder { + fn new(inner: Box, field: Field) -> Self { + Self { + inner, + field, + prepared_array: None, + } + } +} + +impl FieldEncoder for ValidatingFieldEncoder { + fn prepare_array(&mut self, array: ArrayRef) -> Result { + ArrayFieldEncodingStrategy::validate_v2_0_array(array.as_ref(), &self.field, true)?; + let array = ArrayFieldEncodingStrategy::clear_unreachable_v2_0_fsl_struct_validity(array)?; + self.prepared_array = Some(array.clone()); + Ok(array) + } + + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + repdef: crate::repdef::RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + let array = match self.prepared_array.take() { + Some(prepared) if Arc::ptr_eq(&prepared, &array) => prepared, + _ => { + // Direct field encoders have historically accepted arrays whose top-level + // nullability differs from the declared field. Keep that API behavior while + // still making the v2.0 struct-validity invariant unavoidable. + ArrayFieldEncodingStrategy::validate_v2_0_array( + array.as_ref(), + &self.field, + false, + )?; + ArrayFieldEncodingStrategy::clear_unreachable_v2_0_fsl_struct_validity(array)? + } + }; + self.inner + .maybe_encode(array, external_buffers, repdef, row_number, num_rows) + } + + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + self.inner.flush(external_buffers) + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> futures::future::BoxFuture<'_, Result>> { + self.inner.finish(external_buffers) + } + + fn num_columns(&self) -> u32 { + self.inner.num_columns() + } +} + +impl ArrayFieldEncodingStrategy { + /// Create the field strategy for the `pb::ArrayEncoding` grammar. + /// + /// ``` + /// use lance_encoding::encoder::ArrayFieldEncodingStrategy; + /// + /// let strategy = ArrayFieldEncodingStrategy::new(); + /// ``` + pub fn new() -> Self { + Self { + array_encoding_strategy: Arc::new(ArrayStrategy), + } + } + + fn is_primitive_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Boolean + | DataType::Date32 + | DataType::Date64 + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Duration(_) + | DataType::Float16 + | DataType::Float32 + | DataType::Float64 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Int8 + | DataType::Interval(_) + | DataType::Null + | DataType::Time32(_) + | DataType::Time64(_) + | DataType::Timestamp(_, _) + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::UInt8 + | DataType::FixedSizeBinary(_) + | DataType::FixedSizeList(_, _) + | DataType::Binary + | DataType::LargeBinary + | DataType::Utf8 + | DataType::LargeUtf8, + ) + } + + fn validate_v2_0_array( + array: &dyn Array, + field: &Field, + enforce_field_nullability: bool, + ) -> Result<()> { + Self::validate_v2_0_array_reachability(array, field, enforce_field_nullability, None) + } + + fn field_requires_v2_0_validation(field: &Field, enforce_field_nullability: bool) -> bool { + field.logical_type.is_struct() + || (enforce_field_nullability && !field.nullable) + || field + .children + .iter() + .any(|child| Self::field_requires_v2_0_validation(child, enforce_field_nullability)) + } + + fn is_reachable(reachable: Option<&BooleanBuffer>, index: usize) -> bool { + reachable.map(|mask| mask.value(index)).unwrap_or(true) + } + + fn positional_child_reachability( + array: &dyn Array, + reachable: Option<&BooleanBuffer>, + ) -> Option { + if reachable.is_none() && array.null_count() == 0 { + return None; + } + Some(BooleanBuffer::from_iter((0..array.len()).map(|index| { + Self::is_reachable(reachable, index) && array.is_valid(index) + }))) + } + + fn list_child_reachability( + array: &GenericListArray, + reachable: Option<&BooleanBuffer>, + ) -> Option { + let values_len = array.values().len(); + let offsets = array.offsets(); + if reachable.is_none() + && array.null_count() == 0 + && offsets.first().map(|offset| offset.as_usize()) == Some(0) + && offsets.last().map(|offset| offset.as_usize()) == Some(values_len) + { + return None; + } + + let mut child_reachable = vec![false; values_len]; + for index in 0..array.len() { + if Self::is_reachable(reachable, index) && array.is_valid(index) { + let start = offsets[index].as_usize(); + let end = offsets[index + 1].as_usize(); + for is_reachable in child_reachable.iter_mut().take(end).skip(start) { + *is_reachable = true; + } + } + } + Some(BooleanBuffer::from_iter(child_reachable)) + } + + fn map_child_reachability( + array: &arrow_array::MapArray, + reachable: Option<&BooleanBuffer>, + ) -> Option { + let values_len = array.entries().len(); + let offsets = array.offsets(); + if reachable.is_none() + && array.null_count() == 0 + && offsets.first().copied() == Some(0) + && offsets.last().copied() == Some(values_len as i32) + { + return None; + } + + let mut child_reachable = vec![false; values_len]; + for index in 0..array.len() { + if Self::is_reachable(reachable, index) && array.is_valid(index) { + let start = offsets[index] as usize; + let end = offsets[index + 1] as usize; + for is_reachable in child_reachable.iter_mut().take(end).skip(start) { + *is_reachable = true; + } + } + } + Some(BooleanBuffer::from_iter(child_reachable)) + } + + fn fixed_size_list_child_reachability( + array: &arrow_array::FixedSizeListArray, + reachable: Option<&BooleanBuffer>, + ) -> Option { + let values_len = array.values().len(); + let dimension = array.value_length() as usize; + let (first_value, final_value) = if array.is_empty() { + (0, 0) + } else { + ( + array.value_offset(0) as usize, + array.value_offset(array.len() - 1) as usize + dimension, + ) + }; + if reachable.is_none() + && array.null_count() == 0 + && first_value == 0 + && final_value == values_len + { + return None; + } + + let mut child_reachable = vec![false; values_len]; + for index in 0..array.len() { + if Self::is_reachable(reachable, index) && array.is_valid(index) { + let start = array.value_offset(index) as usize; + for is_reachable in child_reachable.iter_mut().skip(start).take(dimension) { + *is_reachable = true; + } + } + } + Some(BooleanBuffer::from_iter(child_reachable)) + } + + fn validate_v2_0_array_reachability( + array: &dyn Array, + field: &Field, + enforce_field_nullability: bool, + reachable: Option<&BooleanBuffer>, + ) -> Result<()> { + if !Self::field_requires_v2_0_validation(field, enforce_field_nullability) { + return Ok(()); + } + + let reachable_null_count = if array.null_count() == 0 { + 0 + } else if let Some(reachable) = reachable { + (0..array.len()) + .filter(|index| reachable.value(*index) && array.is_null(*index)) + .count() + } else { + array.null_count() + }; + + if enforce_field_nullability && !field.nullable && reachable_null_count > 0 { + return Err(Error::invalid_input(format!( + "The field `{}` contained null values even though the field is marked non-null in the schema", + field.name + ))); + } + if field.logical_type.is_struct() && reachable_null_count > 0 { + return Err(Error::invalid_input(format!( + "The struct field `{}` contains {} null value(s), but Lance file version 2.0 does not encode struct validity; use file version 2.1 or later", + field.name, reachable_null_count + ))); + } + + match array.data_type() { + DataType::Struct(_) => { + let child_reachable = Self::positional_child_reachability(array, reachable); + for (child_field, child_array) in + field.children.iter().zip(array.as_struct().columns()) + { + Self::validate_v2_0_array_reachability( + child_array.as_ref(), + child_field, + enforce_field_nullability, + child_reachable.as_ref(), + )?; + } + } + DataType::List(_) => { + let list_array = array.as_list::(); + if let Some(child_field) = field.children.first() { + let child_reachable = Self::list_child_reachability(list_array, reachable); + Self::validate_v2_0_array_reachability( + list_array.values().as_ref(), + child_field, + enforce_field_nullability, + child_reachable.as_ref(), + )?; + } + } + DataType::LargeList(_) => { + let list_array = array.as_list::(); + if let Some(child_field) = field.children.first() { + let child_reachable = Self::list_child_reachability(list_array, reachable); + Self::validate_v2_0_array_reachability( + list_array.values().as_ref(), + child_field, + enforce_field_nullability, + child_reachable.as_ref(), + )?; + } + } + DataType::Map(_, _) => { + let map_array = array.as_map(); + if let Some(child_field) = field.children.first() { + let child_reachable = Self::map_child_reachability(map_array, reachable); + Self::validate_v2_0_array_reachability( + map_array.entries(), + child_field, + enforce_field_nullability, + child_reachable.as_ref(), + )?; + } + } + DataType::FixedSizeList(_, _) => { + let list_array = array.as_fixed_size_list(); + if let Some(child_field) = field.children.first() { + let child_reachable = + Self::fixed_size_list_child_reachability(list_array, reachable); + Self::validate_v2_0_array_reachability( + list_array.values().as_ref(), + child_field, + enforce_field_nullability, + child_reachable.as_ref(), + )?; + } + } + _ => { + let array_data = array.to_data(); + for (child_field, child_data) in field.children.iter().zip(array_data.child_data()) + { + let child_array = make_array(child_data.clone()); + let child_reachable = (child_array.len() == array.len()) + .then(|| Self::positional_child_reachability(array, reachable)) + .flatten(); + Self::validate_v2_0_array_reachability( + child_array.as_ref(), + child_field, + enforce_field_nullability, + child_reachable.as_ref(), + )?; + } + } + } + Ok(()) + } + + fn clear_unreachable_v2_0_fsl_struct_validity(array: ArrayRef) -> Result { + fn force_valid_under_null_struct( + data: ArrayData, + struct_nulls: &NullBuffer, + ) -> Result { + let Some(child_nulls) = data.nulls() else { + return Ok(data); + }; + if !(0..data.len()) + .any(|index| struct_nulls.is_null(index) && child_nulls.is_null(index)) + { + return Ok(data); + } + + let nulls = + NullBuffer::new(BooleanBuffer::from_iter((0..data.len()).map(|index| { + struct_nulls.is_null(index) || child_nulls.is_valid(index) + }))); + let nulls = (nulls.null_count() > 0).then_some(nulls); + Ok(data.into_builder().nulls(nulls).build()?) + } + + fn strip_struct_validity(data: ArrayData) -> Result<(ArrayData, bool)> { + let struct_fields = match data.data_type() { + DataType::Struct(fields) => Some(fields.clone()), + _ => None, + }; + let struct_nulls = struct_fields + .as_ref() + .and_then(|_| data.nulls()) + .filter(|nulls| nulls.null_count() > 0) + .cloned(); + let mut children_changed = false; + let children = data + .child_data() + .iter() + .cloned() + .enumerate() + .map(|(index, child)| { + let (mut child, changed) = strip_struct_validity(child)?; + children_changed |= changed; + if let (Some(fields), Some(struct_nulls)) = (&struct_fields, &struct_nulls) + && !fields[index].is_nullable() + { + let original_null_count = child.null_count(); + child = force_valid_under_null_struct(child, struct_nulls)?; + children_changed |= child.null_count() != original_null_count; + } + Ok(child) + }) + .collect::>>()?; + if struct_nulls.is_none() && !children_changed { + return Ok((data, false)); + } + + let mut builder = data.into_builder().child_data(children); + if struct_nulls.is_some() { + builder = builder.nulls(None); + } + Ok((builder.build()?, true)) + } + + fn normalize(data: ArrayData) -> Result<(ArrayData, bool)> { + let is_fixed_size_list = matches!(data.data_type(), DataType::FixedSizeList(_, _)); + let mut children_changed = false; + let children = data + .child_data() + .iter() + .cloned() + .map(|child| { + let (child, changed) = if is_fixed_size_list { + strip_struct_validity(child)? + } else { + normalize(child)? + }; + children_changed |= changed; + Ok(child) + }) + .collect::>>()?; + if !children_changed { + return Ok((data, false)); + } + Ok((data.into_builder().child_data(children).build()?, true)) + } + + let (data, changed) = normalize(array.to_data())?; + if changed { + Ok(make_array(data)) + } else { + Ok(array) + } + } + + fn create_field_encoder_raw( + &self, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result> { + let options = context.options; + let data_type = field.data_type(); + if Self::is_primitive_type(&data_type) { + let column_index = column_index.next_column_index(field.id as u32); + if field.metadata.contains_key(BLOB_META_KEY) { + let mut packed_meta = HashMap::new(); + packed_meta.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string()); + let desc_field = + Field::try_from(BLOB_DESC_FIELD.clone().with_metadata(packed_meta)).unwrap(); + let desc_encoder = Box::new(PrimitiveFieldEncoder::try_new( + options, + self.array_encoding_strategy.clone(), + column_index, + desc_field, + )?); + Ok(Box::new(BlobFieldEncoder::new(desc_encoder))) + } else { + Ok(Box::new(PrimitiveFieldEncoder::try_new( + options, + self.array_encoding_strategy.clone(), + column_index, + field.clone(), + )?)) + } + } else { + match data_type { + DataType::List(_child) | DataType::LargeList(_child) => { + let list_idx = column_index.next_column_index(field.id as u32); + let inner_encoding = + self.create_field_encoder_raw(&field.children[0], column_index, context)?; + let offsets_encoder = + Arc::new(BasicEncoder::new(Box::new(ValueEncoder::default()))); + Ok(Box::new(ListFieldEncoder::new( + inner_encoding, + offsets_encoder, + options.cache_bytes_per_column, + options.keep_original_array, + list_idx, + ))) + } + DataType::Struct(_) => { + let field_metadata = &field.metadata; + if field_metadata + .get(PACKED_STRUCT_LEGACY_META_KEY) + .map(|v| str_is_truthy(v)) + .unwrap_or(field_metadata.contains_key(PACKED_STRUCT_META_KEY)) + { + Ok(Box::new(PrimitiveFieldEncoder::try_new( + options, + self.array_encoding_strategy.clone(), + column_index.next_column_index(field.id as u32), + field.clone(), + )?)) + } else { + let header_idx = column_index.next_column_index(field.id as u32); + let children_encoders = field + .children + .iter() + .map(|field| { + self.create_field_encoder_raw(field, column_index, context) + }) + .collect::>>()?; + Ok(Box::new(StructFieldEncoder::new( + children_encoders, + header_idx, + ))) + } + } + DataType::Dictionary(_, value_type) => { + if Self::is_primitive_type(&value_type) { + Ok(Box::new(PrimitiveFieldEncoder::try_new( + options, + self.array_encoding_strategy.clone(), + column_index.next_column_index(field.id as u32), + field.clone(), + )?)) + } else { + Err(Error::not_supported_source(format!( + "cannot encode a dictionary column whose value type is a logical type ({})", + value_type + ).into())) + } + } + _ => Err(Error::not_supported_source( + format!( + "Lance v2.0 has no field encoding for '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )), + } + } + } +} + +impl Default for ArrayFieldEncodingStrategy { + fn default() -> Self { + Self::new() + } +} + +impl FieldEncodingStrategy for ArrayFieldEncodingStrategy { + fn validate_array(&self, array: &dyn Array, field: &Field) -> Result<()> { + Self::validate_v2_0_array(array, field, true) + } + + fn create_field_encoder( + &self, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result> { + let encoder = self.create_field_encoder_raw(field, column_index, context)?; + Ok(Box::new(ValidatingFieldEncoder::new( + encoder, + field.clone(), + ))) + } +} + +/// Page-encoding selection for the `pb::ArrayEncoding` grammar. +#[derive(Debug)] +struct ArrayStrategy; + +impl ArrayStrategy { + fn get_field_compression(field_meta: &HashMap) -> Option { + let compression = field_meta.get(COMPRESSION_META_KEY)?; + let compression_scheme = compression.parse::(); + match compression_scheme { + Ok(compression_scheme) => Some(CompressionConfig::new( + compression_scheme, + field_meta + .get(COMPRESSION_LEVEL_META_KEY) + .and_then(|level| level.parse().ok()), + )), + Err(_) => None, + } + } + + fn default_binary_encoder( + arrays: &[ArrayRef], + field_meta: Option<&HashMap>, + data_size: u64, + ) -> Result> { + let bin_indices_encoder = + Self::choose_array_encoder(arrays, &DataType::UInt64, data_size, false, None)?; + + if let Some(compression) = field_meta.and_then(Self::get_field_compression) { + if compression.scheme() == CompressionScheme::Fsst { + // User requested FSST + let raw_encoder = Box::new(BinaryEncoder::try_new(bin_indices_encoder, None)?); + Ok(Box::new(FsstArrayEncoder::new(raw_encoder))) + } else { + // Generic compression + Ok(Box::new(BinaryEncoder::try_new( + bin_indices_encoder, + Some(compression), + )?)) + } + } else { + Ok(Box::new(BinaryEncoder::try_new(bin_indices_encoder, None)?)) + } + } + + fn choose_array_encoder( + arrays: &[ArrayRef], + data_type: &DataType, + data_size: u64, + use_dict_encoding: bool, + field_meta: Option<&HashMap>, + ) -> Result> { + match data_type { + DataType::FixedSizeList(inner, dimension) => { + Ok(Box::new(BasicEncoder::new(Box::new(FslEncoder::new( + Self::choose_array_encoder( + arrays, + inner.data_type(), + data_size, + use_dict_encoding, + None, + )?, + *dimension as u32, + ))))) + } + DataType::Dictionary(key_type, value_type) => { + let key_encoder = + Self::choose_array_encoder(arrays, key_type, data_size, false, None)?; + let value_encoder = + Self::choose_array_encoder(arrays, value_type, data_size, false, None)?; + + Ok(Box::new(AlreadyDictionaryEncoder::new( + key_encoder, + value_encoder, + ))) + } + DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => { + if use_dict_encoding { + let dict_indices_encoder = Self::choose_array_encoder( + // We need to pass arrays to this method to figure out what kind of compression to + // use but we haven't actually calculated the indices yet. For now, we just assume + // worst case and use the full range. In the future maybe we can pass in statistics + // instead of the actual data + &[Arc::new(UInt8Array::from_iter_values(0_u8..255_u8))], + &DataType::UInt8, + data_size, + false, + None, + )?; + let dict_items_encoder = Self::choose_array_encoder( + arrays, + &DataType::Utf8, + data_size, + false, + None, + )?; + + Ok(Box::new(DictionaryEncoder::new( + dict_indices_encoder, + dict_items_encoder, + ))) + } else { + Self::default_binary_encoder(arrays, field_meta, data_size) + } + } + DataType::Struct(fields) => { + let num_fields = fields.len(); + let mut inner_encoders = Vec::new(); + + for i in 0..num_fields { + let inner_datatype = fields[i].data_type(); + let inner_encoder = Self::choose_array_encoder( + arrays, + inner_datatype, + data_size, + use_dict_encoding, + None, + )?; + inner_encoders.push(inner_encoder); + } + + Ok(Box::new(PackedStructEncoder::new(inner_encoders))) + } + DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => Ok( + Box::new(BasicEncoder::new(Box::new(ValueEncoder::default()))), + ), + + // TODO: for signed integers, I intend to make it a cascaded encoding, a sparse array for the negative values and very wide(bit-width) values, + // then a bitpacked array for the narrow(bit-width) values, I need `BitpackedForNeg` to be merged first, I am + // thinking about putting this sparse array in the metadata so bitpacking remain using one page buffer only. + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Box::new( + BasicEncoder::new(Box::new(ValueEncoder::default())), + )), + _ => Ok(Box::new(BasicEncoder::new(Box::new( + ValueEncoder::default(), + )))), + } + } +} + +fn get_dict_encoding_threshold() -> u64 { + env::var("LANCE_DICT_ENCODING_THRESHOLD") + .ok() + .and_then(|val| val.parse().ok()) + .unwrap_or(100) +} + +// check whether we want to use dictionary encoding or not +// by applying a threshold on cardinality +// returns true if cardinality < threshold but false if the total number of rows is less than the threshold +// The choice to use 100 is just a heuristic for now +// hyperloglog is used for cardinality estimation +// error rate = 1.04 / sqrt(2^p), where p is the precision +// and error rate is 1.04 / sqrt(2^12) = 1.56% +fn check_dict_encoding(arrays: &[ArrayRef], threshold: u64) -> bool { + let num_total_rows = arrays.iter().map(|arr| arr.len()).sum::(); + if num_total_rows < threshold as usize { + return false; + } + const PRECISION: u8 = 12; + + let mut hll: HyperLogLogPlus = + HyperLogLogPlus::new(PRECISION, RandomState::new()).unwrap(); + + for arr in arrays { + let string_array = arrow_array::cast::as_string_array(arr); + for value in string_array.iter().flatten() { + hll.insert(value); + let estimated_cardinality = hll.count() as u64; + if estimated_cardinality >= threshold { + return false; + } + } + } + + true +} + +#[cfg(test)] +fn check_fixed_size_encoding(arrays: &[ArrayRef]) -> Option { + if arrays.is_empty() { + return None; + } + + // make sure no array has an empty string + if !arrays.iter().all(|arr| { + if let Some(arr) = arr.as_string_opt::() { + arr.iter().flatten().all(|s| !s.is_empty()) + } else if let Some(arr) = arr.as_binary_opt::() { + arr.iter().flatten().all(|s| !s.is_empty()) + } else if let Some(arr) = arr.as_string_opt::() { + arr.iter().flatten().all(|s| !s.is_empty()) + } else if let Some(arr) = arr.as_binary_opt::() { + arr.iter().flatten().all(|s| !s.is_empty()) + } else { + panic!("wrong dtype"); + } + }) { + return None; + } + + let lengths = arrays + .iter() + .flat_map(|arr| { + if let Some(arr) = arr.as_string_opt::() { + let offsets = arr.offsets().inner(); + offsets + .windows(2) + .map(|w| (w[1] - w[0]) as u64) + .collect::>() + } else if let Some(arr) = arr.as_binary_opt::() { + let offsets = arr.offsets().inner(); + offsets + .windows(2) + .map(|w| (w[1] - w[0]) as u64) + .collect::>() + } else if let Some(arr) = arr.as_string_opt::() { + let offsets = arr.offsets().inner(); + offsets + .windows(2) + .map(|w| (w[1] - w[0]) as u64) + .collect::>() + } else if let Some(arr) = arr.as_binary_opt::() { + let offsets = arr.offsets().inner(); + offsets + .windows(2) + .map(|w| (w[1] - w[0]) as u64) + .collect::>() + } else { + panic!("wrong dtype"); + } + }) + .collect::>(); + + // find first non-zero value in lengths + let first_non_zero = lengths.iter().position(|&x| x != 0); + if let Some(first_non_zero) = first_non_zero { + // make sure all lengths are equal to first_non_zero length or zero + if !lengths + .iter() + .all(|&x| x == 0 || x == lengths[first_non_zero]) + { + return None; + } + + // set the byte width + Some(lengths[first_non_zero]) + } else { + None + } +} + +impl ArrayEncodingStrategy for ArrayStrategy { + fn create_array_encoder( + &self, + arrays: &[ArrayRef], + field: &Field, + ) -> Result> { + let data_size = arrays + .iter() + .map(|arr| arr.get_buffer_memory_size() as u64) + .sum::(); + let data_type = arrays[0].data_type(); + + let use_dict_encoding = data_type == &DataType::Utf8 + && check_dict_encoding(arrays, get_dict_encoding_threshold()); + + Self::choose_array_encoder( + arrays, + data_type, + data_size, + use_dict_encoding, + Some(&field.metadata), + ) + } +} + +#[cfg(test)] +mod tests { + use super::{ + ArrayEncodingStrategy, ArrayFieldEncodingStrategy, ArrayStrategy, check_dict_encoding, + check_fixed_size_encoding, + }; + use crate::constants::{COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY}; + use crate::encoder::{BatchEncoder, EncodingOptions}; + use arrow_array::{ArrayRef, StringArray}; + use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema}; + use lance_core::{Error, datatypes::Schema}; + use std::collections::HashMap; + use std::sync::Arc; + + #[test] + fn test_unsupported_field_type_returns_error() { + let entries = Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + ])), + false, + ); + let arrow_schema = ArrowSchema::new(vec![Field::new( + "attributes", + DataType::Map(Arc::new(entries), false), + true, + )]); + let schema = Schema::try_from(&arrow_schema).unwrap(); + + let error = BatchEncoder::try_new( + &schema, + &ArrayFieldEncodingStrategy::new(), + &EncodingOptions::default(), + ) + .err() + .unwrap(); + + assert!(matches!(error, Error::NotSupported { .. })); + assert!(error.to_string().contains("attributes")); + assert!(error.to_string().contains("Map")); + } + + fn is_dict_encoding_applicable(arr: Vec>, threshold: u64) -> bool { + let arr = StringArray::from(arr); + let arr = Arc::new(arr) as ArrayRef; + check_dict_encoding(&[arr], threshold) + } + + #[test] + fn test_dict_encoding_should_be_applied_if_cardinality_less_than_threshold() { + assert!(is_dict_encoding_applicable( + vec![Some("a"), Some("b"), Some("a"), Some("b")], + 3, + )); + } + + #[test] + fn test_dict_encoding_should_not_be_applied_if_cardinality_larger_than_threshold() { + assert!(!is_dict_encoding_applicable( + vec![Some("a"), Some("b"), Some("c"), Some("d")], + 3, + )); + } + + #[test] + fn test_dict_encoding_should_not_be_applied_if_cardinality_equal_to_threshold() { + assert!(!is_dict_encoding_applicable( + vec![Some("a"), Some("b"), Some("c"), Some("a")], + 3, + )); + } + + #[test] + fn test_dict_encoding_should_not_be_applied_for_empty_arrays() { + assert!(!is_dict_encoding_applicable(vec![], 3)); + } + + #[test] + fn test_dict_encoding_should_not_be_applied_for_smaller_than_threshold_arrays() { + assert!(!is_dict_encoding_applicable(vec![Some("a"), Some("a")], 3)); + } + + fn is_fixed_size_encoding_applicable(arrays: Vec>>) -> bool { + let mut final_arrays = Vec::new(); + for arr in arrays { + let arr = StringArray::from(arr); + let arr = Arc::new(arr) as ArrayRef; + final_arrays.push(arr); + } + + check_fixed_size_encoding(&final_arrays).is_some() + } + + #[test] + fn test_fixed_size_binary_encoding_applicable() { + assert!(!is_fixed_size_encoding_applicable(vec![vec![]])); + + assert!(is_fixed_size_encoding_applicable(vec![vec![ + Some("a"), + Some("b") + ]])); + + assert!(!is_fixed_size_encoding_applicable(vec![vec![ + Some("abc"), + Some("de") + ]])); + + assert!(is_fixed_size_encoding_applicable(vec![vec![ + Some("pqr"), + None + ]])); + + assert!(!is_fixed_size_encoding_applicable(vec![vec![ + Some("pqr"), + Some("") + ]])); + + assert!(!is_fixed_size_encoding_applicable(vec![vec![ + Some(""), + Some("") + ]])); + } + + #[test] + fn test_fixed_size_binary_encoding_applicable_multiple_arrays() { + assert!(is_fixed_size_encoding_applicable(vec![ + vec![Some("a"), Some("b")], + vec![Some("c"), Some("d")] + ])); + + assert!(!is_fixed_size_encoding_applicable(vec![ + vec![Some("ab"), Some("bc")], + vec![Some("c"), Some("d")] + ])); + + assert!(!is_fixed_size_encoding_applicable(vec![ + vec![Some("ab"), None], + vec![None, Some("d")] + ])); + + assert!(is_fixed_size_encoding_applicable(vec![ + vec![Some("a"), None], + vec![None, Some("d")] + ])); + + assert!(!is_fixed_size_encoding_applicable(vec![ + vec![Some(""), None], + vec![None, Some("")] + ])); + + assert!(!is_fixed_size_encoding_applicable(vec![ + vec![None, None], + vec![None, None] + ])); + } + + fn verify_array_encoder( + array: ArrayRef, + field_meta: Option>, + expected_encoder: &str, + ) { + let encoding_strategy = ArrayStrategy; + let mut field = Field::new("test_field", array.data_type().clone(), true); + if let Some(field_meta) = field_meta { + field.set_metadata(field_meta); + } + let lance_field = lance_core::datatypes::Field::try_from(field).unwrap(); + let encoder_result = encoding_strategy.create_array_encoder(&[array], &lance_field); + assert!(encoder_result.is_ok()); + let encoder = encoder_result.unwrap(); + assert_eq!(format!("{:?}", encoder).as_str(), expected_encoder); + } + + #[test] + fn test_choose_encoder_for_zstd_compressed_string_field() { + verify_array_encoder( + Arc::new(StringArray::from(vec!["a", "bb", "ccc"])), + Some(HashMap::from([( + COMPRESSION_META_KEY.to_string(), + "zstd".to_string(), + )])), + "BinaryEncoder { indices_encoder: BasicEncoder { values_encoder: ValueEncoder }, compression_config: Some(CompressionConfig { scheme: Zstd, level: None }), buffer_compressor: Some(ZstdBufferCompressor { compression_level: 0 }) }", + ); + } + + #[test] + fn test_choose_encoder_for_zstd_compression_level() { + verify_array_encoder( + Arc::new(StringArray::from(vec!["a", "bb", "ccc"])), + Some(HashMap::from([ + (COMPRESSION_META_KEY.to_string(), "zstd".to_string()), + (COMPRESSION_LEVEL_META_KEY.to_string(), "22".to_string()), + ])), + "BinaryEncoder { indices_encoder: BasicEncoder { values_encoder: ValueEncoder }, compression_config: Some(CompressionConfig { scheme: Zstd, level: Some(22) }), buffer_compressor: Some(ZstdBufferCompressor { compression_level: 22 }) }", + ); + } +} diff --git a/vendor/lance-encoding/src/buffer.rs b/vendor/lance-encoding/src/buffer.rs new file mode 100644 index 000000000..646f4f515 --- /dev/null +++ b/vendor/lance-encoding/src/buffer.rs @@ -0,0 +1,503 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Utilities for byte arrays + +use std::{ops::Deref, panic::RefUnwindSafe, ptr::NonNull, sync::Arc}; + +use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, ScalarBuffer}; +use lance_core::{Error, Result, utils::bit::is_pwr_two}; +use std::borrow::Cow; + +/// A copy-on-write byte buffer. +/// +/// It wraps arrow_buffer::Buffer which provides: +/// - Cheap cloning (reference counted) +/// - Zero-copy slicing +/// - Automatic memory alignment +/// +/// LanceBuffer is designed to be used in situations where you might need to +/// pass around byte buffers efficiently without worrying about ownership. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LanceBuffer(Buffer); + +impl LanceBuffer { + /// Convert into an Arrow buffer. Never copies data. + pub fn into_buffer(self) -> Buffer { + self.0 + } + + /// Returns a buffer of the given size with all bits set to 0 + pub fn all_unset(len: usize) -> Self { + Self(Buffer::from_vec(vec![0; len])) + } + + /// Returns a buffer of the given size with all bits set to 1 + pub fn all_set(len: usize) -> Self { + Self(Buffer::from_vec(vec![0xff; len])) + } + + /// Creates an empty buffer + pub fn empty() -> Self { + Self(Buffer::from_vec(Vec::::new())) + } + + /// Converts the buffer into a hex string + pub fn as_hex(&self) -> String { + hex::encode_upper(self) + } + + /// Combine multiple buffers into a single buffer + /// + /// This does involve a data copy (and allocation of a new buffer) + pub fn concat(buffers: &[Self]) -> Self { + let total_len = buffers.iter().map(|b| b.len()).sum(); + let mut data = Vec::with_capacity(total_len); + for buffer in buffers { + data.extend_from_slice(buffer.as_ref()); + } + Self(Buffer::from_vec(data)) + } + + /// Converts the buffer into a hex string, inserting a space + /// between words + pub fn as_spaced_hex(&self, bytes_per_word: u32) -> String { + let hex = self.as_hex(); + let chars_per_word = bytes_per_word as usize * 2; + let num_words = hex.len() / chars_per_word; + let mut spaced_hex = String::with_capacity(hex.len() + num_words); + for (i, c) in hex.chars().enumerate() { + if i % chars_per_word == 0 && i != 0 { + spaced_hex.push(' '); + } + spaced_hex.push(c); + } + spaced_hex + } + + /// Create a LanceBuffer from a bytes::Bytes object + /// + /// The alignment must be specified (as `bytes_per_value`) since we want to make + /// sure we can safely reinterpret the buffer. + /// + /// If the buffer is properly aligned this will be zero-copy. If not, a copy + /// will be made. + /// + /// If `bytes_per_value` is not a power of two, then we assume the buffer is + /// never going to be reinterpret into another type and we can safely + /// ignore the alignment. + /// + /// This is a zero-copy operation when the buffer is properly aligned. + pub fn from_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self { + if is_pwr_two(bytes_per_value) && bytes.as_ptr().align_offset(bytes_per_value as usize) != 0 + { + // The original buffer is not aligned, cannot zero-copy + let mut buf = Vec::with_capacity(bytes.len()); + buf.extend_from_slice(&bytes); + Self(Buffer::from_vec(buf)) + } else { + // The original buffer is aligned, can zero-copy + // SAFETY: the alignment is correct we can make this conversion + unsafe { + Self(Buffer::from_custom_allocation( + NonNull::new(bytes.as_ptr() as _).expect("should be a valid pointer"), + bytes.len(), + Arc::new(bytes), + )) + } + } + } + + /// Make an owned copy of the buffer (always does a copy of the data) + pub fn deep_copy(&self) -> Self { + Self(Buffer::from_vec(self.0.to_vec())) + } + + /// Reinterprets a `Vec` as a LanceBuffer + /// + /// This is a zero-copy operation. We can safely reinterpret `Vec` into `&[u8]` which is what happens here. + /// However, we cannot safely reinterpret a `Vec` into a `Vec` in rust due to alignment constraints + /// from [`Vec::from_raw_parts`]: + /// + /// > `T` needs to have the same alignment as what `ptr` was allocated with. + /// > (`T` having a less strict alignment is not sufficient, the alignment really + /// > needs to be equal to satisfy the `dealloc` requirement that memory must be + /// > allocated and deallocated with the same layout.) + pub fn reinterpret_vec(vec: Vec) -> Self { + Self(Buffer::from_vec(vec)) + } + + /// Reinterprets `Arc<[T]>` as a LanceBuffer + /// + /// This is similar to [`Self::reinterpret_vec`] but for `Arc<[T]>` instead of `Vec` + /// + /// The same alignment constraints apply + pub fn reinterpret_slice(arc: Arc<[T]>) -> Self { + let slice = arc.as_ref(); + let data = NonNull::new(slice.as_ptr() as _).unwrap_or(NonNull::dangling()); + let len = std::mem::size_of_val(slice); + // SAFETY: the ptr will be valid for len items if the Arc<[T]> is valid + let buffer = unsafe { Buffer::from_custom_allocation(data, len, Arc::new(arc)) }; + Self(buffer) + } + + /// Reinterprets a LanceBuffer into a `Vec` + /// + /// If the underlying buffer is not properly aligned, this will involve a copy of the data + /// + /// Note: doing this sort of re-interpretation generally makes assumptions about the endianness + /// of the data. Lance does not support big-endian machines so this is safe. However, if we end + /// up supporting big-endian machines in the future, then any use of this method will need to be + /// carefully reviewed. + pub fn borrow_to_typed_slice(&self) -> ScalarBuffer { + let align = std::mem::align_of::(); + let is_aligned = self.as_ptr().align_offset(align) == 0; + if !self.len().is_multiple_of(std::mem::size_of::()) { + panic!( + "attempt to borrow_to_typed_slice to data type of size {} but we have {} bytes which isn't evenly divisible", + std::mem::size_of::(), + self.len() + ); + } + + if is_aligned { + ScalarBuffer::::from(self.clone().into_buffer()) + } else { + let num_values = self.len() / std::mem::size_of::(); + let vec = Vec::::with_capacity(num_values); + let mut bytes = MutableBuffer::from(vec); + bytes.extend_from_slice(self); + ScalarBuffer::::from(Buffer::from(bytes)) + } + } + + /// Reinterprets a LanceBuffer into a `&[T]` + /// + /// Unlike [`Self::borrow_to_typed_slice`], this function returns a `Cow<'_, [T]>` instead of an owned + /// buffer. It saves the cost of Arc creation and destruction, which can be really helpful when + /// we borrow data and just drop it without reusing it. + /// + /// Caller should decide which way to use based on their own needs. + /// + /// If the underlying buffer is not properly aligned, this will involve a copy of the data + /// + /// Note: doing this sort of re-interpretation generally makes assumptions about the endianness + /// of the data. Lance does not support big-endian machines so this is safe. However, if we end + /// up supporting big-endian machines in the future, then any use of this method will need to be + /// carefully reviewed. + pub fn borrow_to_typed_view(&self) -> Cow<'_, [T]> { + let align = std::mem::align_of::(); + if !self.len().is_multiple_of(std::mem::size_of::()) { + panic!( + "attempt to view data type of size {} but we have {} bytes which isn't evenly divisible", + std::mem::size_of::(), + self.len() + ); + } + + if self.as_ptr().align_offset(align) == 0 { + Cow::Borrowed(bytemuck::cast_slice(&self.0)) + } else { + Cow::Owned(bytemuck::pod_collect_to_vec(self.0.as_slice())) + } + } + + /// Concatenates multiple buffers into a single buffer, consuming the input buffers + /// + /// If there is only one buffer, it will be returned as is + pub fn concat_into_one(buffers: Vec) -> Self { + if buffers.len() == 1 { + return buffers.into_iter().next().unwrap(); + } + + let mut total_len = 0; + for buffer in &buffers { + total_len += buffer.len(); + } + + let mut data = Vec::with_capacity(total_len); + for buffer in buffers { + data.extend_from_slice(buffer.as_ref()); + } + + Self(Buffer::from_vec(data)) + } + + /// Zips multiple buffers into a single buffer, consuming the input buffers + /// + /// Unlike concat_into_one this "zips" the buffers, interleaving the values + pub fn zip_into_one(buffers: Vec<(Self, u64)>, num_values: u64) -> Result { + let bytes_per_value = buffers.iter().map(|(_, bits_per_value)| { + if bits_per_value % 8 == 0 { + Ok(bits_per_value / 8) + } else { + Err(Error::invalid_input_source(format!("LanceBuffer::zip_into_one only supports full-byte buffers currently and received a buffer with {} bits per value", bits_per_value).into())) + } + }).collect::>>()?; + let total_bytes_per_value = bytes_per_value.iter().sum::(); + let total_bytes = (total_bytes_per_value * num_values) as usize; + + let mut zipped = vec![0_u8; total_bytes]; + let mut buffer_ptrs = buffers + .iter() + .zip(bytes_per_value) + .map(|((buffer, _), bytes_per_value)| (buffer.as_ptr(), bytes_per_value as usize)) + .collect::>(); + + let mut zipped_ptr = zipped.as_mut_ptr(); + unsafe { + let end = zipped_ptr.add(total_bytes); + while zipped_ptr < end { + for (buf, bytes_per_value) in buffer_ptrs.iter_mut() { + std::ptr::copy_nonoverlapping(*buf, zipped_ptr, *bytes_per_value); + zipped_ptr = zipped_ptr.add(*bytes_per_value); + *buf = buf.add(*bytes_per_value); + } + } + } + + Ok(Self(Buffer::from_vec(zipped))) + } + + /// Create a LanceBuffer from a slice + /// + /// This is NOT a zero-copy operation. We can't create a borrowed buffer because + /// we have no way of extending the lifetime of the slice. + pub fn copy_slice(slice: &[u8]) -> Self { + Self(Buffer::from_vec(slice.to_vec())) + } + + /// Create a LanceBuffer from an array (fixed-size slice) + /// + /// This is NOT a zero-copy operation. The slice memory could be on the stack and + /// thus we can't forget it. + pub fn copy_array(array: [u8; N]) -> Self { + Self(Buffer::from_vec(Vec::from(array))) + } + + #[allow(clippy::len_without_is_empty)] + pub fn len(&self) -> usize { + self.0.len() + } + + /// Returns a new [LanceBuffer] that is a slice of this buffer starting at `offset`, + /// with `length` bytes. + /// Doing so allows the same memory region to be shared between lance buffers. + /// # Panics + /// Panics if `(offset + length)` is larger than the existing length. + pub fn slice_with_length(&self, offset: usize, length: usize) -> Self { + let original_buffer_len = self.len(); + assert!( + offset.saturating_add(length) <= original_buffer_len, + "the offset + length of the sliced Buffer cannot exceed the existing length" + ); + Self(self.0.slice_with_length(offset, length)) + } + + /// Returns a new [LanceBuffer] that is a slice of this buffer starting at bit `offset` + /// with `length` bits. + /// + /// Unlike `slice_with_length`, this method allows for slicing at a bit level but always + /// requires a copy of the data (unless offset is byte-aligned) + /// + /// This method performs the bit slice using the Arrow convention of *bitwise* little-endian + /// + /// This means, given the bit buffer 0bABCDEFGH_HIJKLMNOP and the slice starting at bit 3 and + /// with length 8, the result will be 0bNOPABCDE + pub fn bit_slice_le_with_length(&self, offset: usize, length: usize) -> Self { + let sliced = self.0.bit_slice(offset, length); + Self(sliced) + } + + /// Get a pointer to the underlying data + pub fn as_ptr(&self) -> *const u8 { + self.0.as_ptr() + } +} + +impl AsRef<[u8]> for LanceBuffer { + fn as_ref(&self) -> &[u8] { + self.0.as_slice() + } +} + +impl Deref for LanceBuffer { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + self.as_ref() + } +} + +// All `From` implementations are zero-copy + +impl From> for LanceBuffer { + fn from(buffer: Vec) -> Self { + Self(Buffer::from_vec(buffer)) + } +} + +impl From for LanceBuffer { + fn from(buffer: Buffer) -> Self { + Self(buffer) + } +} + +// An iterator that keeps a clone of a borrowed LanceBuffer so we +// can have a 'static lifetime +pub struct LanceBufferIter { + buffer: Buffer, + index: usize, +} + +impl Iterator for LanceBufferIter { + type Item = u8; + + fn next(&mut self) -> Option { + if self.index >= self.buffer.len() { + None + } else { + // SAFETY: we just checked that index is in bounds + let byte = unsafe { self.buffer.get_unchecked(self.index) }; + self.index += 1; + Some(*byte) + } + } +} + +impl IntoIterator for LanceBuffer { + type Item = u8; + type IntoIter = LanceBufferIter; + + fn into_iter(self) -> Self::IntoIter { + LanceBufferIter { + buffer: self.0, + index: 0, + } + } +} + +#[cfg(test)] +mod tests { + use arrow_buffer::Buffer; + + use super::LanceBuffer; + + #[test] + fn test_eq() { + let buf = LanceBuffer::from(Buffer::from_vec(vec![1_u8, 2, 3])); + let buf2 = LanceBuffer::from(vec![1, 2, 3]); + assert_eq!(buf, buf2); + } + + #[test] + fn test_reinterpret_vec() { + let vec = vec![1_u32, 2, 3]; + let buf = LanceBuffer::reinterpret_vec(vec); + + let mut expected = Vec::with_capacity(12); + expected.extend_from_slice(&1_u32.to_ne_bytes()); + expected.extend_from_slice(&2_u32.to_ne_bytes()); + expected.extend_from_slice(&3_u32.to_ne_bytes()); + let expected = LanceBuffer::from(expected); + + assert_eq!(expected, buf); + assert_eq!(buf.borrow_to_typed_slice::().as_ref(), vec![1, 2, 3]); + } + + #[test] + fn test_concat() { + let buf1 = LanceBuffer::from(vec![1_u8, 2, 3]); + let buf2 = LanceBuffer::from(vec![4_u8, 5, 6]); + let buf3 = LanceBuffer::from(vec![7_u8, 8, 9]); + + let expected = LanceBuffer::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9]); + assert_eq!( + expected, + LanceBuffer::concat_into_one(vec![buf1, buf2, buf3]) + ); + + let empty = LanceBuffer::empty(); + assert_eq!( + LanceBuffer::empty(), + LanceBuffer::concat_into_one(vec![empty]) + ); + + let expected = LanceBuffer::from(vec![1, 2, 3]); + assert_eq!( + expected, + LanceBuffer::concat_into_one(vec![expected.deep_copy(), LanceBuffer::empty()]) + ); + } + + #[test] + fn test_zip() { + let buf1 = LanceBuffer::from(vec![1_u8, 2, 3]); + let buf2 = LanceBuffer::reinterpret_vec(vec![1_u16, 2, 3]); + let buf3 = LanceBuffer::reinterpret_vec(vec![1_u32, 2, 3]); + + let zipped = LanceBuffer::zip_into_one(vec![(buf1, 8), (buf2, 16), (buf3, 32)], 3).unwrap(); + + assert_eq!(zipped.len(), 21); + + let mut expected = Vec::with_capacity(21); + for i in 1..4 { + expected.push(i as u8); + expected.extend_from_slice(&(i as u16).to_ne_bytes()); + expected.extend_from_slice(&(i as u32).to_ne_bytes()); + } + let expected = LanceBuffer::from(expected); + + assert_eq!(expected, zipped); + } + + #[test] + fn test_hex() { + let buf = LanceBuffer::from(vec![1, 2, 15, 20]); + assert_eq!("01020F14", buf.as_hex()); + } + + #[test] + #[should_panic] + fn test_to_typed_slice_invalid() { + let buf = LanceBuffer::from(vec![0, 1, 2]); + buf.borrow_to_typed_slice::(); + } + + #[test] + fn test_to_typed_slice() { + // Buffer is aligned, no copy will be made, both calls + // should get same ptr + let buf = LanceBuffer::from(vec![0, 1]); + let borrow = buf.borrow_to_typed_slice::(); + let view_ptr = borrow.as_ref().as_ptr(); + let borrow2 = buf.borrow_to_typed_slice::(); + let view_ptr2 = borrow2.as_ref().as_ptr(); + + assert_eq!(view_ptr, view_ptr2); + + let bytes = bytes::Bytes::from(vec![0, 1, 2]); + let sliced = bytes.slice(1..3); + // Intentionally LYING about alignment here to trigger test + let buf = LanceBuffer::from_bytes(sliced, 1); + let borrow = buf.borrow_to_typed_slice::(); + let view_ptr = borrow.as_ref().as_ptr(); + let borrow2 = buf.borrow_to_typed_slice::(); + let view_ptr2 = borrow2.as_ref().as_ptr(); + + assert_ne!(view_ptr, view_ptr2); + } + + #[test] + fn test_bit_slice_le() { + let buf = LanceBuffer::from(vec![0x0F, 0x0B]); + + // Keep in mind that validity buffers are *bitwise* little-endian + assert_eq!(buf.bit_slice_le_with_length(0, 4).as_ref(), &[0x0F]); + assert_eq!(buf.bit_slice_le_with_length(4, 4).as_ref(), &[0x00]); + assert_eq!(buf.bit_slice_le_with_length(3, 8).as_ref(), &[0x61]); + assert_eq!(buf.bit_slice_le_with_length(0, 8).as_ref(), &[0x0F]); + assert_eq!(buf.bit_slice_le_with_length(4, 8).as_ref(), &[0xB0]); + assert_eq!(buf.bit_slice_le_with_length(4, 12).as_ref(), &[0xB0, 0x00]); + } +} diff --git a/vendor/lance-encoding/src/compression.rs b/vendor/lance-encoding/src/compression.rs new file mode 100644 index 000000000..af8a0dff7 --- /dev/null +++ b/vendor/lance-encoding/src/compression.rs @@ -0,0 +1,2806 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Compression traits and definitions for Lance 2.1 +//! +//! In 2.1 the first step of encoding is structural encoding, where we shred inputs into +//! leaf arrays and take care of the validity / offsets structure. Then we pick a structural +//! encoding (mini-block or full-zip) and then we compress the data. +//! +//! This module defines the traits for the compression step. Each structural encoding has its +//! own compression strategy. +//! +//! Miniblock compression is a block based approach for small data. Since we introduce some read +//! amplification and decompress entire blocks we are able to use opaque compression. +//! +//! Fullzip compression is a per-value approach where we require that values are transparently +//! compressed so that we can locate them later. + +#[cfg(feature = "bitpacking")] +use crate::encodings::physical::bitpacking::{InlineBitpacking, OutOfLineBitpacking}; +use crate::{ + buffer::LanceBuffer, + compression_config::{BssMode, CompressionFieldParams}, + constants::{ + BSS_META_KEY, COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, RLE_THRESHOLD_META_KEY, + }, + data::{DataBlock, FixedWidthDataBlock, VariableWidthBlock}, + encodings::{ + logical::primitive::{ + fullzip::PerValueCompressor, + miniblock::{MAX_MINIBLOCK_VALUES, MiniBlockCompressor}, + }, + physical::{ + binary::{ + BinaryBlockDecompressor, BinaryMiniBlockDecompressor, BinaryMiniBlockEncoder, + VariableDecoder, VariableEncoder, + }, + block::{ + CompressedBufferEncoder, CompressionConfig, CompressionScheme, + GeneralBlockDecompressor, + }, + byte_stream_split::{ + ByteStreamSplitDecompressor, ByteStreamSplitEncoder, should_use_bss, + }, + constant::ConstantDecompressor, + fsst::{ + FsstMiniBlockDecompressor, FsstMiniBlockEncoder, FsstPerValueDecompressor, + FsstPerValueEncoder, + }, + general::{GeneralMiniBlockCompressor, GeneralMiniBlockDecompressor}, + packed::{ + PackedStructFixedPerValueDecompressor, PackedStructFixedPerValueEncoder, + PackedStructFixedWidthMiniBlockDecompressor, + PackedStructFixedWidthMiniBlockEncoder, PackedStructVariablePerValueDecompressor, + PackedStructVariablePerValueEncoder, VariablePackedStructFieldDecoder, + VariablePackedStructFieldKind, + }, + rle::{ + RleChildDecompressor, RleDecompressor, RleEncoder, RunLengthWidth, + rle_encoded_size, select_run_length_width, + }, + value::{ValueDecompressor, ValueEncoder}, + }, + }, + format::pb21::{CompressiveEncoding, compressive_encoding::Compression}, + statistics::{GetStat, Stat}, +}; + +use arrow_array::{cast::AsArray, types::UInt64Type}; +use arrow_schema::DataType; +use fsst::fsst::{FSST_LEAST_INPUT_MAX_LENGTH, FSST_LEAST_INPUT_SIZE}; +use lance_core::{Error, Result, datatypes::Field, error::LanceOptionExt}; +use std::{str::FromStr, sync::Arc}; + +/// Default threshold for RLE compression selection when the user explicitly provides a threshold. +/// +/// If no threshold is provided, we use a size model instead of a fixed run ratio. +/// This preserves existing behavior for users relying on the default, while making +/// the default selection more type-aware. +const DEFAULT_RLE_COMPRESSION_THRESHOLD: f64 = 0.5; + +// Minimum block size (32kb) to trigger general block compression +const MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION: u64 = 32 * 1024; +const RLE_BLOCK_HEADER_BYTES: u128 = std::mem::size_of::() as u128; + +/// Trait for compression algorithms that compress an entire block of data into one opaque +/// and self-described chunk. +/// +/// This is actually a _third_ compression strategy used in a few corner cases today (TODO: remove?) +/// +/// This is the most general type of compression. There are no constraints on the method +/// of compression it is assumed that the entire block of data will be present at decompression. +/// +/// This is the least appropriate strategy for random access because we must load the entire +/// block to access any single value. This should only be used for cases where random access is never +/// required (e.g. when encoding metadata buffers like a dictionary or for encoding rep/def +/// mini-block chunks) +pub trait BlockCompressor: std::fmt::Debug + Send + Sync { + /// Compress the data into zero or one buffers and describe the codec used. + /// + /// `None` represents a metadata-only codec. `Some` represents a physical + /// payload, including a zero-byte payload. + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)>; +} + +/// A trait to pick which compression to use for given data +/// +/// There are several different kinds of compression. +/// +/// - Block compression is the most generic, but most difficult to use efficiently +/// - Per-value compression results in either a fixed width data block or a variable +/// width data block. In other words, there is some number of bits per value. +/// In addition, each value should be independently decompressible. +/// - Mini-block compression results in a small block of opaque data for chunks +/// of rows. Each block is somewhere between 0 and 16KiB in size. This is +/// used for narrow data types (both fixed and variable length) where we can +/// fit many values into an 16KiB block. +pub trait CompressionStrategy: Send + Sync + std::fmt::Debug { + /// Create a block compressor for the given data. + fn create_block_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result>; + + /// Create a per-value compressor for the given data + fn create_per_value( + &self, + field: &Field, + data: &DataBlock, + ) -> Result>; + + /// Create a mini-block compressor for the given data + fn create_miniblock_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result>; +} + +pub(crate) fn compress_required_block( + strategy: &dyn CompressionStrategy, + field: &Field, + data: DataBlock, +) -> Result<(LanceBuffer, CompressiveEncoding)> { + let compressor = strategy.create_block_compressor(field, &data)?; + let (payload, encoding) = compressor.compress(data)?; + let payload = payload.ok_or_else(|| { + Error::internal("Required block compressor selected a metadata-only codec".to_string()) + })?; + Ok((payload, encoding)) +} + +fn try_bss_for_mini_block( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, +) -> Option> { + // BSS requires general compression to be effective + // If compression is not set or explicitly disabled, skip BSS + if params.compression.is_none() || params.compression.as_deref() == Some("none") { + return None; + } + + let mode = params.bss.unwrap_or(BssMode::Auto); + // should_use_bss already checks for supported bit widths (32/64) + if should_use_bss(data, mode) { + return Some(Box::new(ByteStreamSplitEncoder::new( + data.bits_per_value as usize, + ))); + } + None +} + +fn rle_is_applicable(data: &FixedWidthDataBlock, params: &CompressionFieldParams) -> Option { + let bits = data.bits_per_value; + if !matches!(bits, 8 | 16 | 32 | 64) { + return None; + } + + let type_size = bits / 8; + let run_count = data.expect_single_stat::(Stat::RunCount); + let threshold = params + .rle_threshold + .unwrap_or(DEFAULT_RLE_COMPRESSION_THRESHOLD); + + // If the user explicitly provided a threshold then honor it as an additional guard. + // A lower threshold makes RLE harder to trigger and can be used to avoid CPU overhead. + let passes_threshold = match params.rle_threshold { + Some(_) => (run_count as f64) < (data.num_values as f64) * threshold, + None => true, + }; + + if !passes_threshold { + return None; + } + + Some((data.num_values as u128) * (type_size as u128)) +} + +fn rle_beats_raw_and_bitpacking( + data: &FixedWidthDataBlock, + encoded_bytes: u128, + raw_bytes: u128, +) -> bool { + if encoded_bytes >= raw_bytes { + return false; + } + + #[cfg(feature = "bitpacking")] + { + if let Some(bitpack_bytes) = estimate_inline_bitpacking_bytes(data).map(u128::from) + && bitpack_bytes < encoded_bytes + { + return false; + } + } + true +} + +fn try_fixed_u8_rle_for_mini_block( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, +) -> Option> { + let raw_bytes = rle_is_applicable(data, params)?; + let rle_bytes = estimate_rle_size_for_width_from_data( + data, + Some(*MAX_MINIBLOCK_VALUES), + RunLengthWidth::U8, + ) + .ok()?; + rle_beats_raw_and_bitpacking(data, rle_bytes, raw_bytes) + .then(|| Box::new(RleEncoder::with_run_length_width(RunLengthWidth::U8)) as _) +} + +fn try_child_rle_for_mini_block( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, +) -> Option> { + let raw_bytes = rle_is_applicable(data, params)?; + let (run_length_width, estimated_bytes) = + estimate_rle_width_and_size_from_data(data, Some(*MAX_MINIBLOCK_VALUES)).ok()?; + let child_compression = rle_child_compression_config(params); + let encoder = || { + RleEncoder::with_child_encoding( + run_length_width, + child_compression, + child_compression, + true, + ) + }; + + #[cfg(feature = "bitpacking")] + let bitpack_bytes = estimate_inline_bitpacking_bytes(data).map(u128::from); + #[cfg(not(feature = "bitpacking"))] + let bitpack_bytes = None::; + + let should_measure_children = (child_compression.is_some() || cfg!(feature = "bitpacking")) + && (estimated_bytes >= raw_bytes + || bitpack_bytes.is_some_and(|bytes| bytes < estimated_bytes)); + let selected_bytes = if should_measure_children { + encoder().selected_payload_size(data).ok()? + } else { + estimated_bytes + }; + + rle_beats_raw_and_bitpacking(data, selected_bytes, raw_bytes).then(|| Box::new(encoder()) as _) +} + +fn rle_child_compression_config(params: &CompressionFieldParams) -> Option { + let raw = params.compression.as_deref()?; + if matches!(raw, "none" | "fsst") { + return None; + } + let scheme = CompressionScheme::from_str(raw).ok()?; + Some(CompressionConfig::new(scheme, params.compression_level)) +} + +fn try_rle_for_block_with_width( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, + run_length_width: RunLengthWidth, + rle_payload_bytes: u128, +) -> Result>> { + let bits = data.bits_per_value; + if !matches!(bits, 8 | 16 | 32 | 64) { + return Ok(None); + } + + let run_count = data.expect_single_stat::(Stat::RunCount); + let threshold = params + .rle_threshold + .unwrap_or(DEFAULT_RLE_COMPRESSION_THRESHOLD); + + let passes_threshold = match params.rle_threshold { + Some(_) => (run_count as f64) < (data.num_values as f64) * threshold, + None => true, + }; + + if !passes_threshold { + return Ok(None); + } + + let raw_bytes = (data.num_values as u128) * ((bits / 8) as u128); + let rle_bytes = rle_payload_bytes.saturating_add(RLE_BLOCK_HEADER_BYTES); + + if rle_bytes >= raw_bytes { + return Ok(None); + } + + #[cfg(feature = "bitpacking")] + { + if let Some(bitpack_bytes) = estimate_block_bitpacking_bytes(data) + && bitpack_bytes < rle_bytes + { + return Ok(None); + } + } + + Ok(Some(Box::new(RleEncoder::with_run_length_width( + run_length_width, + )))) +} + +fn try_fixed_u8_rle_for_block( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, +) -> Result>> { + if !matches!(data.bits_per_value, 8 | 16 | 32 | 64) { + return Ok(None); + } + let encoded_bytes = estimate_rle_size_for_width_from_data(data, None, RunLengthWidth::U8)?; + try_rle_for_block_with_width(data, params, RunLengthWidth::U8, encoded_bytes) +} + +fn try_variable_rle_for_block( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, +) -> Result>> { + if !matches!(data.bits_per_value, 8 | 16 | 32 | 64) { + return Ok(None); + } + let (width, encoded_bytes) = estimate_rle_width_and_size_from_data(data, None)?; + try_rle_for_block_with_width(data, params, width, encoded_bytes) +} + +fn estimate_rle_width_and_size_from_data( + data: &FixedWidthDataBlock, + max_segment_values: Option, +) -> Result<(RunLengthWidth, u128)> { + select_run_length_width( + &data.data, + data.num_values, + data.bits_per_value, + max_segment_values, + ) +} + +fn estimate_rle_size_for_width_from_data( + data: &FixedWidthDataBlock, + max_segment_values: Option, + run_length_width: RunLengthWidth, +) -> Result { + rle_encoded_size( + &data.data, + data.num_values, + data.bits_per_value, + max_segment_values, + run_length_width, + ) +} + +fn try_bitpack_for_mini_block(_data: &FixedWidthDataBlock) -> Option> { + #[cfg(feature = "bitpacking")] + { + let bits = _data.bits_per_value; + if estimate_inline_bitpacking_bytes(_data).is_some() { + return Some(Box::new(InlineBitpacking::new(bits))); + } + None + } + #[cfg(not(feature = "bitpacking"))] + { + None + } +} + +#[cfg(feature = "bitpacking")] +fn estimate_inline_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option { + use arrow_array::cast::AsArray; + + let bits = data.bits_per_value; + if !matches!(bits, 8 | 16 | 32 | 64) { + return None; + } + if data.num_values == 0 { + return None; + } + + let bit_widths = data.expect_stat(Stat::BitWidth); + let widths = bit_widths.as_primitive::(); + + let words_per_chunk: u128 = 1; + let word_bytes: u128 = (bits / 8) as u128; + let mut total_words: u128 = 0; + for i in 0..widths.len() { + let bit_width = widths.value(i) as u128; + let packed_words = (1024u128 * bit_width) / (bits as u128); + total_words = total_words.saturating_add(words_per_chunk.saturating_add(packed_words)); + } + + let estimated_bytes = total_words.saturating_mul(word_bytes); + let raw_bytes = data.data_size() as u128; + + if estimated_bytes >= raw_bytes { + return None; + } + + u64::try_from(estimated_bytes).ok() +} + +fn try_bitpack_for_block(data: &FixedWidthDataBlock) -> Option> { + let bits = data.bits_per_value; + if !matches!(bits, 8 | 16 | 32 | 64) { + return None; + } + + let bit_widths = data.expect_stat(Stat::BitWidth); + let widths = bit_widths.as_primitive::(); + let max_bit_width = *widths.values().iter().max().unwrap(); + + // Out-of-line bitpacking requires a narrower width, including for its raw + // tail. A full-width chunk can appear when signed values are negative. + if max_bit_width >= bits { + return None; + } + + let too_small = + widths.len() == 1 && InlineBitpacking::min_size_bytes(widths.value(0)) >= data.data_size(); + + if too_small { + return None; + } + + if data.num_values <= 1024 { + Some(Box::new(InlineBitpacking::new(bits))) + } else { + Some(Box::new(OutOfLineBitpacking::new(max_bit_width, bits))) + } +} + +#[cfg(feature = "bitpacking")] +fn estimate_block_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option { + let bits = data.bits_per_value; + if !matches!(bits, 8 | 16 | 32 | 64) || data.num_values == 0 { + return None; + } + + let bit_widths = data.expect_stat(Stat::BitWidth); + let widths = bit_widths.as_primitive::(); + let max_bit_width = *widths.values().iter().max()?; + let word_bytes = (bits / 8) as u128; + + let bitpacked_words = if data.num_values <= 1024 { + 1 + (1024u128 * (max_bit_width as u128)) / (bits as u128) + } else { + estimate_out_of_line_bitpacking_words(data.num_values, max_bit_width, bits)? + }; + let bitpacked_bytes = bitpacked_words.saturating_mul(word_bytes); + if bitpacked_bytes >= data.data_size() as u128 { + return None; + } + + Some(bitpacked_bytes) +} + +#[cfg(feature = "bitpacking")] +fn estimate_out_of_line_bitpacking_words( + num_values: u64, + compressed_bits_per_value: u64, + bits_per_value: u64, +) -> Option { + let num_values = usize::try_from(num_values).ok()?; + let compressed_bits_per_value = usize::try_from(compressed_bits_per_value).ok()?; + let bits_per_value = usize::try_from(bits_per_value).ok()?; + if compressed_bits_per_value >= bits_per_value { + return None; + } + + let elems_per_chunk = 1024usize; + let num_chunks = num_values.div_ceil(elems_per_chunk); + let words_per_chunk = (elems_per_chunk * compressed_bits_per_value).div_ceil(bits_per_value); + let last_chunk_is_runt = !num_values.is_multiple_of(elems_per_chunk); + + if !last_chunk_is_runt { + return Some((num_chunks * words_per_chunk) as u128); + } + + let num_whole_chunks = num_chunks - 1; + let remaining_items = num_values - num_whole_chunks * elems_per_chunk; + let tail_bit_savings = bits_per_value - compressed_bits_per_value; + let padding_cost = compressed_bits_per_value * (elems_per_chunk - remaining_items); + let tail_pack_savings = tail_bit_savings * remaining_items; + let tail_words = if padding_cost < tail_pack_savings { + words_per_chunk + } else { + remaining_items + }; + + Some((num_whole_chunks * words_per_chunk + tail_words) as u128) +} + +fn maybe_wrap_general_for_mini_block( + inner: Box, + params: &CompressionFieldParams, +) -> Result> { + match params.compression.as_deref() { + None | Some("none") | Some("fsst") => Ok(inner), + Some(raw) => { + let scheme = CompressionScheme::from_str(raw) + .map_err(|_| Error::invalid_input(format!("Unknown compression scheme: {raw}")))?; + let cfg = CompressionConfig::new(scheme, params.compression_level); + Ok(Box::new(GeneralMiniBlockCompressor::new(inner, cfg))) + } + } +} + +fn try_general_compression( + field_params: &CompressionFieldParams, + data: &DataBlock, +) -> Result, CompressionConfig)>> { + // Explicitly disable general compression. + if field_params.compression.as_deref() == Some("none") { + return Ok(None); + } + + // User-requested compression (unused today but perhaps still used + // in the future someday) + if let Some(compression_scheme) = &field_params.compression { + let scheme: CompressionScheme = compression_scheme.parse()?; + let config = CompressionConfig::new(scheme, field_params.compression_level); + let compressor = Box::new(CompressedBufferEncoder::try_new(config)?); + return Ok(Some((compressor, config))); + } + + // Automatic compression for large blocks + if data.data_size() > MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION { + let compressor = Box::new(CompressedBufferEncoder::default()); + let config = compressor.compressor.config(); + return Ok(Some((compressor, config))); + } + + Ok(None) +} + +/// Parse field-level compression metadata without applying format-specific constraints. +pub fn field_metadata_params(field: &Field) -> CompressionFieldParams { + let mut params = CompressionFieldParams::default(); + + if let Some(compression) = field.metadata.get(COMPRESSION_META_KEY) { + params.compression = Some(compression.clone()); + } + if let Some(level) = field.metadata.get(COMPRESSION_LEVEL_META_KEY) { + params.compression_level = level.parse().ok(); + } + if let Some(threshold) = field.metadata.get(RLE_THRESHOLD_META_KEY) { + params.rle_threshold = threshold.parse().ok(); + } + if let Some(bss_str) = field.metadata.get(BSS_META_KEY) { + match BssMode::parse(bss_str) { + Some(mode) => params.bss = Some(mode), + None => log::warn!("Invalid BSS mode '{}', using default", bss_str), + } + } + if let Some(minichunk_size_str) = field + .metadata + .get(super::constants::MINICHUNK_SIZE_META_KEY) + { + if let Ok(minichunk_size) = minichunk_size_str.parse::() { + params.minichunk_size = Some(minichunk_size); + } else { + log::warn!("Invalid minichunk_size '{}', skipping", minichunk_size_str); + } + } + + params +} + +/// Apply general-purpose compression requested for a fixed-width miniblock. +pub fn finalize_miniblock_compressor( + data: &DataBlock, + compressor: Box, + params: &CompressionFieldParams, +) -> Result> { + if matches!(data, DataBlock::FixedWidth(_)) { + maybe_wrap_general_for_mini_block(compressor, params) + } else { + Ok(compressor) + } +} + +/// Honor an explicit `compression = none` request for fixed-width miniblocks. +pub fn try_uncompressed_fixed_width_miniblock( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Option> { + (matches!(data, DataBlock::FixedWidth(_)) && params.compression.as_deref() == Some("none")) + .then(|| Box::new(ValueEncoder::default()) as _) +} + +/// Select byte-stream-split compression for an applicable fixed-width miniblock. +pub fn try_byte_stream_split_miniblock( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Option> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_bss_for_mini_block(data, params) +} + +/// Select the original fixed-u8 RLE miniblock grammar. +pub fn try_fixed_u8_rle_miniblock( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Option> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_fixed_u8_rle_for_mini_block(data, params) +} + +/// Select variable-width RLE with independently encoded children. +pub fn try_child_rle_miniblock( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Option> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_child_rle_for_mini_block(data, params) +} + +/// Select inline bitpacking for applicable fixed-width miniblocks. +pub fn try_bitpacking_miniblock(data: &DataBlock) -> Option> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_bitpack_for_mini_block(data) +} + +/// Store fixed-width miniblock values without a value codec. +pub fn try_raw_fixed_width_miniblock(data: &DataBlock) -> Option> { + matches!(data, DataBlock::FixedWidth(_)).then(|| Box::new(ValueEncoder::default()) as _) +} + +/// Encode variable-width miniblocks with binary or FSST encoding. +pub fn try_variable_width_miniblock( + field: &Field, + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result>> { + let DataBlock::VariableWidth(data) = data else { + return Ok(None); + }; + if data.bits_per_offset != 32 && data.bits_per_offset != 64 { + return Err(Error::invalid_input(format!( + "Variable width compression not supported for {} bit offsets", + data.bits_per_offset + ))); + } + + let compression = params.compression.as_deref(); + let data_size = data.expect_single_stat::(Stat::DataSize); + let max_len = data.expect_single_stat::(Stat::MaxLength); + if compression == Some("none") { + return Ok(Some(Box::new(BinaryMiniBlockEncoder::new( + params.minichunk_size, + )))); + } + + let use_fsst = compression == Some("fsst") + || (compression.is_none() + && !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) + && max_len >= FSST_LEAST_INPUT_MAX_LENGTH + && data_size >= FSST_LEAST_INPUT_SIZE as u64); + let mut encoder: Box = if use_fsst { + Box::new(FsstMiniBlockEncoder::new(params.minichunk_size)) + } else { + Box::new(BinaryMiniBlockEncoder::new(params.minichunk_size)) + }; + if let Some(compression_scheme) = compression.filter(|scheme| *scheme != "fsst") { + let scheme: CompressionScheme = compression_scheme.parse()?; + let config = CompressionConfig::new(scheme, params.compression_level); + encoder = Box::new(GeneralMiniBlockCompressor::new(encoder, config)); + } + Ok(Some(encoder)) +} + +/// Encode fixed-width packed structs as miniblocks. +pub fn try_fixed_packed_struct_miniblock( + data: &DataBlock, +) -> Result>> { + let DataBlock::Struct(data) = data else { + return Ok(None); + }; + if data.has_variable_width_child() { + return Err(Error::invalid_input( + "Packed struct mini-block encoding supports only fixed-width children", + )); + } + Ok(Some(Box::new( + PackedStructFixedWidthMiniBlockEncoder::default(), + ))) +} + +/// Store fixed-size-list miniblocks without a value codec. +pub fn try_raw_fixed_size_list_miniblock(data: &DataBlock) -> Option> { + matches!(data, DataBlock::FixedSizeList(_)).then(|| Box::new(ValueEncoder::default()) as _) +} + +/// Store fixed-width and fixed-size-list values directly in full-zip pages. +pub fn try_raw_per_value(data: &DataBlock) -> Option> { + matches!(data, DataBlock::FixedWidth(_) | DataBlock::FixedSizeList(_)) + .then(|| Box::new(ValueEncoder::default()) as _) +} + +fn validate_packed_struct(field: &Field, data: &DataBlock) -> Result> { + let DataBlock::Struct(data) = data else { + return Ok(None); + }; + if field.children.len() != data.children.len() { + return Err(Error::invalid_input( + "Struct field metadata does not match data block children", + )); + } + Ok(Some(data.has_variable_width_child())) +} + +/// Reject variable-width packed structs while preserving the fixed-width error. +pub fn reject_packed_struct_per_value( + field: &Field, + data: &DataBlock, +) -> Result>> { + let Some(has_variable_child) = validate_packed_struct(field, data)? else { + return Ok(None); + }; + if has_variable_child { + return Err(Error::not_supported_source( + "Variable packed struct encoding is not enabled by the selected file format".into(), + )); + } + Err(Error::invalid_input( + "Packed struct per-value compression should not be used for fixed-width-only structs", + )) +} + +/// Encode variable-width packed structs with the exact strategy recursively. +pub fn try_variable_packed_struct_per_value( + strategy: Arc, + field: &Field, + data: &DataBlock, +) -> Result>> { + let Some(has_variable_child) = validate_packed_struct(field, data)? else { + return Ok(None); + }; + if !has_variable_child { + return Err(Error::invalid_input( + "Packed struct per-value compression should not be used for fixed-width-only structs", + )); + } + Ok(Some(Box::new(PackedStructVariablePerValueEncoder::new( + strategy, + field.children.clone(), + )))) +} + +/// Encode all packed structs with the exact strategy recursively. +pub fn try_packed_struct_per_value( + strategy: Arc, + field: &Field, + data: &DataBlock, +) -> Result>> { + let Some(has_variable_child) = validate_packed_struct(field, data)? else { + return Ok(None); + }; + if has_variable_child { + return Ok(Some(Box::new(PackedStructVariablePerValueEncoder::new( + strategy, + field.children.clone(), + )))); + } + + Ok(Some(Box::new(PackedStructFixedPerValueEncoder::new( + field.children.clone(), + )))) +} + +/// Encode variable-width values directly, with FSST or per-value compression +/// when applicable. +pub fn try_variable_width_per_value( + field: &Field, + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result>> { + let DataBlock::VariableWidth(data) = data else { + return Ok(None); + }; + let compression = params.compression.as_deref(); + if compression == Some("none") { + return Ok(Some(Box::new(VariableEncoder::default()))); + } + + let max_len = data.expect_single_stat::(Stat::MaxLength); + let data_size = data.expect_single_stat::(Stat::DataSize); + let per_value_requested = compression.is_some_and(|compression| compression != "fsst"); + if (max_len > 32 * 1024 || per_value_requested) && data_size >= FSST_LEAST_INPUT_SIZE as u64 { + if compression == Some("zstd") { + let config = CompressionConfig::new(CompressionScheme::Zstd, params.compression_level); + return Ok(Some(Box::new(CompressedBufferEncoder::try_new(config)?))); + } + return Ok(Some(Box::new(CompressedBufferEncoder::default()))); + } + + if data.bits_per_offset != 32 && data.bits_per_offset != 64 { + return Err(Error::invalid_input(format!( + "Per-value compression does not support variable-width data with {}-bit offsets", + data.bits_per_offset + ))); + } + let encoder = Box::new(VariableEncoder::default()); + let use_fsst = compression == Some("fsst") + || (compression.is_none() + && !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) + && max_len >= FSST_LEAST_INPUT_MAX_LENGTH + && data_size >= FSST_LEAST_INPUT_SIZE as u64); + Ok(Some(if use_fsst { + Box::new(FsstPerValueEncoder::new(encoder)) + } else { + encoder + })) +} + +/// Select fixed-u8 RLE for block compression. +pub fn try_fixed_u8_rle_block( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result>> { + let DataBlock::FixedWidth(data) = data else { + return Ok(None); + }; + try_fixed_u8_rle_for_block(data, params) +} + +/// Select variable-width RLE for block compression. +pub fn try_variable_rle_block( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result>> { + let DataBlock::FixedWidth(data) = data else { + return Ok(None); + }; + try_variable_rle_for_block(data, params) +} + +/// Select block bitpacking for applicable fixed-width values. +pub fn try_bitpacking_block(data: &DataBlock) -> Option> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_bitpack_for_block(data) +} + +/// Select explicitly requested or automatic general-purpose block compression. +pub fn try_general_block( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result>> { + let Some((compressor, _config)) = try_general_compression(params, data)? else { + return Ok(None); + }; + Ok(Some(compressor)) +} + +/// Store fixed- and variable-width block values without block compression. +pub fn try_raw_block(data: &DataBlock) -> Option> { + match data { + DataBlock::FixedWidth(_) => { + Some(Box::new(ValueEncoder::default()) as Box) + } + DataBlock::VariableWidth(_) => { + Some(Box::new(VariableEncoder::default()) as Box) + } + _ => None, + } +} + +pub trait MiniBlockDecompressor: std::fmt::Debug + Send + Sync { + fn decompress(&self, data: Vec, num_values: u64) -> Result; + + /// Returns the exact aggregate decoded size when it is determined solely by the value count. + /// + /// Implementations should only return `Some` when this aggregate estimate can be used by + /// [`DataBlockBuilder`](crate::data::DataBlockBuilder) to preallocate the decoded output + /// exactly. Outputs with multiple buffers or whose layout-dependent allocation cannot be + /// represented by one aggregate estimate should return `None`. + fn decoded_size_bytes(&self, _num_values: u64) -> Option { + None + } +} + +pub trait FixedPerValueDecompressor: std::fmt::Debug + Send + Sync { + /// Decompress one or more values + fn decompress(&self, data: FixedWidthDataBlock, num_values: u64) -> Result; + /// The number of bits in each value + /// + /// Currently (and probably long term) this must be a multiple of 8 + fn bits_per_value(&self) -> u64; + + /// Returns the exact aggregate decoded size when it is determined solely by the value count. + /// + /// Implementations should only return `Some` when this aggregate estimate can be used by + /// [`DataBlockBuilder`](crate::data::DataBlockBuilder) to preallocate the decoded output + /// exactly. Outputs with multiple buffers or whose layout-dependent allocation cannot be + /// represented by one aggregate estimate should return `None`. + fn decoded_size_bytes(&self, _num_values: u64) -> Option { + None + } +} + +pub trait VariablePerValueDecompressor: std::fmt::Debug + Send + Sync { + /// Decompress one or more values + fn decompress(&self, data: VariableWidthBlock) -> Result; +} + +pub trait BlockDecompressor: std::fmt::Debug + Send + Sync { + fn decompress(&self, data: Option, num_values: u64) -> Result; + + /// Whether this codec consumes one payload buffer. + fn requires_payload(&self) -> bool { + true + } + + /// Inspect a block for an exact payload-derived value count when supported. + /// + /// This must not materialize the decoded values. `None` means the encoding requires an + /// external count or cannot safely prove one from this payload. + fn infer_num_values(&self, _data: &LanceBuffer) -> Result> { + Ok(None) + } +} + +pub(crate) fn require_block_payload(data: Option, codec: &str) -> Result { + data.ok_or_else(|| Error::invalid_input(format!("{codec} requires one payload"))) +} + +pub(crate) fn require_no_block_payload(data: Option, codec: &str) -> Result<()> { + if data.is_some() { + return Err(Error::invalid_input(format!("{codec} expects no payload"))); + } + Ok(()) +} + +pub trait DecompressionStrategy: std::fmt::Debug + Send + Sync { + fn create_miniblock_decompressor( + &self, + description: &CompressiveEncoding, + decompression_strategy: &dyn DecompressionStrategy, + ) -> Result>; + + fn create_fixed_per_value_decompressor( + &self, + description: &CompressiveEncoding, + ) -> Result>; + + fn create_variable_per_value_decompressor( + &self, + description: &CompressiveEncoding, + ) -> Result>; + + fn create_block_decompressor( + &self, + description: &CompressiveEncoding, + ) -> Result>; +} + +#[derive(Debug, Default)] +pub struct DefaultDecompressionStrategy {} + +impl DecompressionStrategy for DefaultDecompressionStrategy { + fn create_miniblock_decompressor( + &self, + description: &CompressiveEncoding, + decompression_strategy: &dyn DecompressionStrategy, + ) -> Result> { + match description.compression.as_ref().unwrap() { + Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))), + #[cfg(feature = "bitpacking")] + Compression::InlineBitpacking(description) => { + Ok(Box::new(InlineBitpacking::from_description(description))) + } + #[cfg(not(feature = "bitpacking"))] + Compression::InlineBitpacking(_) => Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )), + Compression::Variable(variable) => { + let Compression::Flat(offsets) = variable + .offsets + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + else { + panic!("Variable compression only supports flat offsets") + }; + Ok(Box::new(BinaryMiniBlockDecompressor::new( + offsets.bits_per_value as u8, + ))) + } + Compression::Fsst(description) => { + let inner_decompressor = decompression_strategy.create_miniblock_decompressor( + description.values.as_ref().unwrap(), + decompression_strategy, + )?; + Ok(Box::new(FsstMiniBlockDecompressor::new( + description, + inner_decompressor, + ))) + } + Compression::PackedStruct(description) => Ok(Box::new( + PackedStructFixedWidthMiniBlockDecompressor::new(description), + )), + Compression::VariablePackedStruct(_) => Err(Error::not_supported_source( + "variable packed struct decoding is not yet implemented".into(), + )), + Compression::FixedSizeList(fsl) => { + // In the future, we might need to do something more complex here if FSL supports + // compression. + Ok(Box::new(ValueDecompressor::from_fsl(fsl)?)) + } + Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor( + rle, + decompression_strategy, + )?)), + Compression::ByteStreamSplit(bss) => { + let Compression::Flat(values) = + bss.values.as_ref().unwrap().compression.as_ref().unwrap() + else { + panic!("ByteStreamSplit compression only supports flat values") + }; + Ok(Box::new(ByteStreamSplitDecompressor::new( + values.bits_per_value as usize, + ))) + } + Compression::General(general) => { + // Create inner decompressor + let inner_decompressor = self.create_miniblock_decompressor( + general.values.as_ref().ok_or_else(|| { + Error::invalid_input("GeneralMiniBlock missing inner encoding") + })?, + decompression_strategy, + )?; + + // Parse compression config + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input("GeneralMiniBlock missing compression config") + })?; + + let scheme = compression.scheme().try_into()?; + + let compression_config = CompressionConfig::new(scheme, compression.level); + + Ok(Box::new(GeneralMiniBlockDecompressor::new( + inner_decompressor, + compression_config, + ))) + } + _ => todo!(), + } + } + + fn create_fixed_per_value_decompressor( + &self, + description: &CompressiveEncoding, + ) -> Result> { + match description.compression.as_ref().unwrap() { + Compression::Constant(constant) => Ok(Box::new(ConstantDecompressor::new( + constant + .value + .as_ref() + .map(|v| LanceBuffer::from_bytes(v.clone(), 1)), + ))), + Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))), + Compression::FixedSizeList(fsl) => Ok(Box::new(ValueDecompressor::from_fsl(fsl)?)), + Compression::PackedStruct(description) => Ok(Box::new( + PackedStructFixedPerValueDecompressor::new(description)?, + )), + _ => todo!("fixed-per-value decompressor for {:?}", description), + } + } + + fn create_variable_per_value_decompressor( + &self, + description: &CompressiveEncoding, + ) -> Result> { + match description.compression.as_ref().unwrap() { + Compression::Variable(variable) => { + let Compression::Flat(offsets) = variable + .offsets + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + else { + panic!("Variable compression only supports flat offsets") + }; + assert!(offsets.bits_per_value < u8::MAX as u64); + Ok(Box::new(VariableDecoder::default())) + } + Compression::Fsst(fsst) => Ok(Box::new(FsstPerValueDecompressor::new( + LanceBuffer::from_bytes(fsst.symbol_table.clone(), 1), + Box::new(VariableDecoder::default()), + ))), + Compression::General(general) => Ok(Box::new(CompressedBufferEncoder::from_scheme( + general.compression.as_ref().expect_ok()?.scheme(), + )?)), + Compression::VariablePackedStruct(description) => { + let mut fields = Vec::with_capacity(description.fields.len()); + for field in &description.fields { + let value_encoding = field.value.as_ref().ok_or_else(|| { + Error::invalid_input("VariablePackedStruct field is missing value encoding") + })?; + let decoder = match field.layout.as_ref().ok_or_else(|| { + Error::invalid_input("VariablePackedStruct field is missing layout details") + })? { + crate::format::pb21::variable_packed_struct::field_encoding::Layout::BitsPerValue( + bits_per_value, + ) => { + let decompressor = + self.create_fixed_per_value_decompressor(value_encoding)?; + VariablePackedStructFieldDecoder { + kind: VariablePackedStructFieldKind::Fixed { + bits_per_value: *bits_per_value, + decompressor: Arc::from(decompressor), + }, + } + } + crate::format::pb21::variable_packed_struct::field_encoding::Layout::BitsPerLength( + bits_per_length, + ) => { + let decompressor = + self.create_variable_per_value_decompressor(value_encoding)?; + VariablePackedStructFieldDecoder { + kind: VariablePackedStructFieldKind::Variable { + bits_per_length: *bits_per_length, + decompressor: Arc::from(decompressor), + }, + } + } + }; + fields.push(decoder); + } + Ok(Box::new(PackedStructVariablePerValueDecompressor::new( + fields, + ))) + } + _ => todo!("variable-per-value decompressor for {:?}", description), + } + } + + fn create_block_decompressor( + &self, + description: &CompressiveEncoding, + ) -> Result> { + match description.compression.as_ref().unwrap() { + Compression::InlineBitpacking(inline_bitpacking) => Ok(Box::new( + InlineBitpacking::from_description(inline_bitpacking), + )), + Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))), + Compression::Constant(constant) => { + let scalar = constant + .value + .as_ref() + .map(|v| LanceBuffer::from_bytes(v.clone(), 1)); + Ok(Box::new(ConstantDecompressor::new(scalar))) + } + Compression::Variable(_) => Ok(Box::new(BinaryBlockDecompressor::default())), + Compression::FixedSizeList(fsl) => { + Ok(Box::new(ValueDecompressor::from_fsl(fsl.as_ref())?)) + } + Compression::OutOfLineBitpacking(out_of_line) => { + // Extract the compressed bit width from the values encoding + let compressed_bit_width = match out_of_line + .values + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + { + Compression::Flat(flat) => flat.bits_per_value, + _ => { + return Err(Error::invalid_input_source( + "OutOfLineBitpacking values must use Flat encoding".into(), + )); + } + }; + Ok(Box::new(OutOfLineBitpacking::new( + compressed_bit_width, + out_of_line.uncompressed_bits_per_value, + ))) + } + Compression::General(general) => { + let inner_desc = general + .values + .as_ref() + .ok_or_else(|| { + Error::invalid_input("General compression missing inner encoding") + })? + .as_ref(); + let inner_decompressor = self.create_block_decompressor(inner_desc)?; + + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input("General compression missing compression config") + })?; + let scheme = compression.scheme().try_into()?; + let config = CompressionConfig::new(scheme, compression.level); + let general_decompressor = + GeneralBlockDecompressor::try_new(inner_decompressor, config)?; + + Ok(Box::new(general_decompressor)) + } + Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(rle, self)?)), + _ => todo!(), + } + } +} +pub(crate) fn create_rle_decompressor( + rle: &crate::format::pb21::Rle, + decompression_strategy: &dyn DecompressionStrategy, +) -> Result { + let values = rle + .values + .as_ref() + .ok_or_else(|| Error::invalid_input("RLE compression missing values encoding"))?; + let run_lengths = rle + .run_lengths + .as_ref() + .ok_or_else(|| Error::invalid_input("RLE compression missing run lengths encoding"))?; + + let values = create_rle_child_decompressor(values, "values", decompression_strategy)?; + let run_lengths = + create_rle_child_decompressor(run_lengths, "run lengths", decompression_strategy)?; + + if !matches!(values.bits_per_value(), 8 | 16 | 32 | 64) { + return Err(Error::invalid_input(format!( + "RLE compression only supports 8, 16, 32, or 64-bit values, got {}", + values.bits_per_value() + ))); + } + + let run_length_width = + RunLengthWidth::from_bits(run_lengths.bits_per_value()).ok_or_else(|| { + Error::invalid_input(format!( + "RLE compression only supports 8, 16, or 32-bit run lengths, got {}", + run_lengths.bits_per_value() + )) + })?; + + if values.requires_num_values() && run_lengths.requires_num_values() { + return Err(Error::invalid_input( + "RLE values and run lengths child encodings cannot both require the run count", + )); + } + + if values.is_identity() && run_lengths.is_identity() { + return Ok(RleDecompressor::with_run_length_width( + values.bits_per_value(), + run_length_width, + )); + } + + Ok(RleDecompressor::with_child_decompressors( + values.bits_per_value(), + run_length_width, + values, + run_lengths, + )) +} + +fn create_rle_child_decompressor( + encoding: &CompressiveEncoding, + role: &str, + decompression_strategy: &dyn DecompressionStrategy, +) -> Result { + let compression = encoding + .compression + .as_ref() + .ok_or_else(|| Error::invalid_input(format!("RLE {role} missing child compression")))?; + let (bits_per_value, requires_num_values, needs_decompressor) = + validate_rle_child_compression(compression, role)?; + + if needs_decompressor { + Ok(RleChildDecompressor::block( + bits_per_value, + decompression_strategy.create_block_decompressor(encoding)?, + requires_num_values, + )) + } else { + Ok(RleChildDecompressor::flat(bits_per_value)) + } +} + +fn validate_rle_child_compression( + compression: &Compression, + role: &str, +) -> Result<(u64, bool, bool)> { + match compression { + Compression::Flat(flat) => Ok((flat.bits_per_value, false, false)), + Compression::General(general) => { + general.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} general child missing compression config" + )) + })?; + let values = general.values.as_ref().ok_or_else(|| { + Error::invalid_input(format!("RLE {role} general child missing inner encoding")) + })?; + let inner = values.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} general child missing inner compression" + )) + })?; + let (bits_per_value, requires_num_values) = + validate_rle_block_child_inner(inner, role)?; + Ok((bits_per_value, requires_num_values, true)) + } + Compression::OutOfLineBitpacking(out_of_line) => { + let values = out_of_line.values.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} bitpacking child missing values encoding" + )) + })?; + let Compression::Flat(_) = values.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} bitpacking child missing values compression" + )) + })? + else { + return Err(Error::invalid_input(format!( + "RLE {role} bitpacking child only supports flat values" + ))); + }; + Ok((out_of_line.uncompressed_bits_per_value, true, true)) + } + other => Err(Error::invalid_input(format!( + "RLE {role} only supports flat, general, or out-of-line bitpacking child encodings, got {}", + compression_name(other) + ))), + } +} + +fn validate_rle_block_child_inner(compression: &Compression, role: &str) -> Result<(u64, bool)> { + match compression { + Compression::Flat(flat) => Ok((flat.bits_per_value, false)), + Compression::OutOfLineBitpacking(out_of_line) => { + let values = out_of_line.values.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} bitpacking child missing values encoding" + )) + })?; + let Compression::Flat(_) = values.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} bitpacking child missing values compression" + )) + })? + else { + return Err(Error::invalid_input(format!( + "RLE {role} bitpacking child only supports flat values" + ))); + }; + Ok((out_of_line.uncompressed_bits_per_value, true)) + } + other => Err(Error::invalid_input(format!( + "RLE {role} general child only supports flat or out-of-line bitpacking inner encodings, got {}", + compression_name(other) + ))), + } +} + +fn compression_name(compression: &Compression) -> &'static str { + match compression { + Compression::Flat(_) => "flat", + Compression::Variable(_) => "variable", + Compression::Fsst(_) => "fsst", + Compression::OutOfLineBitpacking(_) => "out-of-line bitpacking", + Compression::InlineBitpacking(_) => "inline bitpacking", + Compression::General(_) => "general", + Compression::Constant(_) => "constant", + Compression::Dictionary(_) => "dictionary", + Compression::ByteStreamSplit(_) => "byte stream split", + Compression::PackedStruct(_) => "packed struct", + Compression::FixedSizeList(_) => "fixed-size list", + Compression::VariablePackedStruct(_) => "variable packed struct", + Compression::Rle(_) => "rle", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::buffer::LanceBuffer; + use crate::compression_config::CompressionParams; + use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock}; + use crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext; + use crate::statistics::ComputeStat; + use crate::testing::{TestEncoding, extract_array_encoding_chain, test_compression_strategy}; + use arrow_schema::{DataType, Field as ArrowField}; + use std::collections::HashMap; + + fn strategy(encoding: TestEncoding, params: CompressionParams) -> Arc { + test_compression_strategy(encoding, params) + } + + fn baseline_strategy(params: CompressionParams) -> Arc { + strategy(TestEncoding::StructuralU16, params) + } + + fn selected_block_codec( + strategy: &Arc, + field: &Field, + data: &DataBlock, + ) -> (Box, CompressiveEncoding) { + let compressor = strategy.create_block_compressor(field, data).unwrap(); + let (_, encoding) = compressor.compress(data.clone()).unwrap(); + (compressor, encoding) + } + + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) + } + + fn create_test_field(name: &str, data_type: DataType) -> Field { + let arrow_field = ArrowField::new(name, data_type, true); + let mut field = Field::try_from(&arrow_field).unwrap(); + field.id = -1; + field + } + + fn create_fixed_width_block_with_stats( + bits_per_value: u64, + num_values: u64, + run_count: u64, + ) -> DataBlock { + // Create varied data to avoid low entropy + let bytes_per_value = (bits_per_value / 8) as usize; + let total_bytes = bytes_per_value * num_values as usize; + let mut data = vec![0u8; total_bytes]; + + // Create data with specified run count + let values_per_run = (num_values / run_count).max(1); + let mut run_value = 0u8; + + for i in 0..num_values as usize { + if i % values_per_run as usize == 0 { + run_value = run_value.wrapping_add(17); // Use prime to get varied values + } + // Fill all bytes of the value to create high entropy + for j in 0..bytes_per_value { + let byte_offset = i * bytes_per_value + j; + if byte_offset < data.len() { + data[byte_offset] = run_value.wrapping_add(j as u8); + } + } + } + + let mut block = FixedWidthDataBlock { + bits_per_value, + data: LanceBuffer::reinterpret_vec(data), + num_values, + block_info: BlockInfo::default(), + }; + + // Compute all statistics including BytePositionEntropy + use crate::statistics::ComputeStat; + block.compute_stat(); + + DataBlock::FixedWidth(block) + } + + fn create_fixed_width_block(bits_per_value: u64, num_values: u64) -> DataBlock { + // Create data with some variety to avoid always triggering BSS + let bytes_per_value = (bits_per_value / 8) as usize; + let total_bytes = bytes_per_value * num_values as usize; + let mut data = vec![0u8; total_bytes]; + + // Add some variation to the data to make it more realistic + for i in 0..num_values as usize { + let byte_offset = i * bytes_per_value; + if byte_offset < data.len() { + data[byte_offset] = (i % 256) as u8; + } + } + + let mut block = FixedWidthDataBlock { + bits_per_value, + data: LanceBuffer::reinterpret_vec(data), + num_values, + block_info: BlockInfo::default(), + }; + + // Compute all statistics including BytePositionEntropy + use crate::statistics::ComputeStat; + block.compute_stat(); + + DataBlock::FixedWidth(block) + } + + fn rle_run_length_bits(encoding: &CompressiveEncoding) -> u64 { + let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { + panic!("expected RLE encoding"); + }; + let Compression::Flat(run_lengths) = rle + .run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + else { + panic!("expected flat run lengths"); + }; + run_lengths.bits_per_value + } + + fn expect_rle_encoding(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle { + match encoding.compression.as_ref().unwrap() { + Compression::Rle(rle) => rle, + Compression::General(general) => { + let inner = general.values.as_ref().unwrap(); + let Compression::Rle(rle) = inner.compression.as_ref().unwrap() else { + panic!("expected wrapped RLE encoding"); + }; + rle + } + other => panic!("expected RLE encoding, got {}", compression_name(other)), + } + } + + fn create_variable_width_block( + bits_per_offset: u8, + num_values: u64, + avg_value_size: usize, + ) -> DataBlock { + use crate::statistics::ComputeStat; + + // Create offsets buffer (num_values + 1 offsets) + let mut offsets = Vec::with_capacity((num_values + 1) as usize); + let mut current_offset = 0i64; + offsets.push(current_offset); + + // Generate offsets with varying value sizes + for i in 0..num_values { + let value_size = if avg_value_size == 0 { + 1 + } else { + ((avg_value_size as i64 + (i as i64 % 8) - 4).max(1) as usize) + .min(avg_value_size * 2) + }; + current_offset += value_size as i64; + offsets.push(current_offset); + } + + // Create data buffer with realistic content + let total_data_size = current_offset as usize; + let mut data = vec![0u8; total_data_size]; + + // Fill data with varied content + for i in 0..num_values { + let start_offset = offsets[i as usize] as usize; + let end_offset = offsets[(i + 1) as usize] as usize; + + let content = (i % 256) as u8; + for j in 0..end_offset - start_offset { + data[start_offset + j] = content.wrapping_add(j as u8); + } + } + + // Convert offsets to appropriate lance buffer + let offsets_buffer = match bits_per_offset { + 32 => { + let offsets_32: Vec = offsets.iter().map(|&o| o as i32).collect(); + LanceBuffer::reinterpret_vec(offsets_32) + } + 64 => LanceBuffer::reinterpret_vec(offsets), + _ => panic!("Unsupported bits_per_offset: {}", bits_per_offset), + }; + + let mut block = VariableWidthBlock { + data: LanceBuffer::from(data), + offsets: offsets_buffer, + bits_per_offset, + num_values, + block_info: BlockInfo::default(), + }; + + block.compute_stat(); + DataBlock::VariableWidth(block) + } + + fn create_fsst_candidate_variable_width_block() -> DataBlock { + create_variable_width_block(32, 4096, FSST_LEAST_INPUT_MAX_LENGTH as usize + 16) + } + + #[test] + fn test_parameter_based_compression() { + let mut params = CompressionParams::new(); + + // Configure RLE for ID columns with BSS explicitly disabled + params.columns.insert( + "*_id".to_string(), + CompressionFieldParams { + rle_threshold: Some(0.3), + compression: Some("lz4".to_string()), + compression_level: None, + bss: Some(BssMode::Off), // Explicitly disable BSS to test RLE + minichunk_size: None, + }, + ); + + let strategy = baseline_strategy(params); + let field = create_test_field("user_id", DataType::Int32); + + // Create data with low run count for RLE + // Use create_fixed_width_block_with_stats which properly sets run count + let data = create_fixed_width_block_with_stats(32, 1000, 100); // 100 runs out of 1000 values + + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + // Should use RLE due to low threshold (0.3) and low run count (100/1000 = 0.1) + let debug_str = format!("{:?}", compressor); + + // The compressor should be RLE wrapped in general compression + assert!(debug_str.contains("GeneralMiniBlockCompressor")); + assert!(debug_str.contains("RleEncoder")); + } + + #[test] + fn test_type_level_parameters() { + let mut params = CompressionParams::new(); + + // Configure all Int32 to use specific settings + params.types.insert( + "Int32".to_string(), + CompressionFieldParams { + rle_threshold: Some(0.1), // Very low threshold + compression: Some("zstd".to_string()), + compression_level: Some(3), + bss: Some(BssMode::Off), // Disable BSS to test RLE + minichunk_size: None, + }, + ); + + let strategy = baseline_strategy(params); + let field = create_test_field("some_column", DataType::Int32); + // Create data with very low run count (50 runs for 1000 values = 0.05 ratio) + let data = create_fixed_width_block_with_stats(32, 1000, 50); + + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + // Should use RLE due to very low threshold + assert!(format!("{:?}", compressor).contains("RleEncoder")); + } + + // Regression for #6626: an all-zero stat segment (e.g. rep/def for a long + // run of empty lists) used to disable block bitpacking entirely. + #[test] + #[cfg(feature = "bitpacking")] + fn test_block_bitpacks_with_zero_segment() { + let strategy = baseline_strategy(CompressionParams::default()); + let field = create_test_field("levels", DataType::UInt16); + + // First 1024 zeros, then 1024 ones; max bit width is 1. + let mut values: Vec = vec![0; 1024]; + values.extend(std::iter::repeat_n(1u16, 1024)); + let mut block = FixedWidthDataBlock { + bits_per_value: 16, + data: LanceBuffer::reinterpret_vec(values), + num_values: 2048, + block_info: BlockInfo::default(), + }; + block.compute_stat(); + let data = DataBlock::FixedWidth(block); + + let compressor = strategy.create_block_compressor(&field, &data).unwrap(); + let debug_str = format!("{:?}", compressor); + assert!( + debug_str.contains("OutOfLineBitpacking"), + "expected OutOfLineBitpacking, got: {debug_str}" + ); + } + + #[test] + fn test_rle_block_accounts_for_header_before_selecting() { + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); + let field = create_test_field("small_constant", DataType::Int32); + let values = vec![42i32; 2]; + let mut block = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 2, + block_info: BlockInfo::default(), + }; + block.compute_stat(); + let data = DataBlock::FixedWidth(block); + + let (compressor, encoding) = selected_block_codec(&strategy, &field, &data); + + assert!(format!("{compressor:?}").contains("ValueEncoder")); + assert!(matches!( + encoding.compression.as_ref(), + Some(Compression::Flat(_)) + )); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_block_prefers_bitpacking_when_smaller() { + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); + let field = create_test_field("levels", DataType::UInt16); + + let mut values = Vec::with_capacity(2048); + for run_idx in 0..1024 { + values.extend(std::iter::repeat_n((run_idx % 2) as u16, 2)); + } + let mut block = FixedWidthDataBlock { + bits_per_value: 16, + data: LanceBuffer::reinterpret_vec(values), + num_values: 2048, + block_info: BlockInfo::default(), + }; + block.compute_stat(); + let data = DataBlock::FixedWidth(block); + + let (compressor, encoding) = selected_block_codec(&strategy, &field, &data); + let debug_str = format!("{compressor:?}"); + assert!( + debug_str.contains("OutOfLineBitpacking"), + "expected OutOfLineBitpacking, got: {debug_str}" + ); + assert!(matches!( + encoding.compression.as_ref(), + Some(Compression::OutOfLineBitpacking(_)) + )); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_low_cardinality_prefers_bitpacking_over_rle() { + let strategy = baseline_strategy(CompressionParams::default()); + let field = create_test_field("int_score", DataType::Int64); + + // Low cardinality values (3/4/5) but with moderate run count: + // RLE compresses vs raw, yet bitpacking should be smaller. + let mut values: Vec = Vec::with_capacity(256); + for run_idx in 0..64 { + let value = match run_idx % 3 { + 0 => 3u64, + 1 => 4u64, + _ => 5u64, + }; + values.extend(std::iter::repeat_n(value, 4)); + } + + let mut block = FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(values), + num_values: 256, + block_info: BlockInfo::default(), + }; + + use crate::statistics::ComputeStat; + block.compute_stat(); + + let data = DataBlock::FixedWidth(block); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{:?}", compressor); + assert!( + debug_str.contains("InlineBitpacking"), + "expected InlineBitpacking, got: {debug_str}" + ); + assert!( + !debug_str.contains("RleEncoder"), + "expected RLE to be skipped when bitpacking is smaller, got: {debug_str}" + ); + } + + fn check_uncompressed_encoding(encoding: &CompressiveEncoding, variable: bool) { + let chain = extract_array_encoding_chain(encoding); + if variable { + assert_eq!(chain.len(), 2); + assert_eq!(chain.first().unwrap().as_str(), "variable"); + assert_eq!(chain.get(1).unwrap().as_str(), "flat"); + } else { + assert_eq!(chain.len(), 1); + assert_eq!(chain.first().unwrap().as_str(), "flat"); + } + } + + #[test] + fn test_none_compression() { + let mut params = CompressionParams::new(); + + // Disable compression for embeddings + params.columns.insert( + "embeddings".to_string(), + CompressionFieldParams { + compression: Some("none".to_string()), + ..Default::default() + }, + ); + + let strategy = baseline_strategy(params); + let field = create_test_field("embeddings", DataType::Float32); + let fixed_data = create_fixed_width_block(32, 1000); + let variable_data = create_variable_width_block(32, 10, 32 * 1024); + + // Test miniblock + let compressor = strategy + .create_miniblock_compressor(&field, &fixed_data) + .unwrap(); + let (_block, encoding) = compressor + .compress(miniblock_context(), fixed_data.clone()) + .unwrap(); + check_uncompressed_encoding(&encoding, false); + let compressor = strategy + .create_miniblock_compressor(&field, &variable_data) + .unwrap(); + let (_block, encoding) = compressor + .compress(miniblock_context(), variable_data.clone()) + .unwrap(); + check_uncompressed_encoding(&encoding, true); + + // Test pervalue + let compressor = strategy.create_per_value(&field, &fixed_data).unwrap(); + let (_block, encoding) = compressor.compress(fixed_data).unwrap(); + check_uncompressed_encoding(&encoding, false); + let compressor = strategy.create_per_value(&field, &variable_data).unwrap(); + let (_block, encoding) = compressor.compress(variable_data).unwrap(); + check_uncompressed_encoding(&encoding, true); + } + + #[test] + fn test_field_metadata_none_compression() { + // Prepare field with metadata for none compression + let mut arrow_field = ArrowField::new("simple_col", DataType::Binary, true); + let mut metadata = HashMap::new(); + metadata.insert(COMPRESSION_META_KEY.to_string(), "none".to_string()); + arrow_field = arrow_field.with_metadata(metadata); + let field = Field::try_from(&arrow_field).unwrap(); + + let strategy = baseline_strategy(CompressionParams::new()); + + // Test miniblock + let fixed_data = create_fixed_width_block(32, 1000); + let variable_data = create_variable_width_block(32, 10, 32 * 1024); + + let compressor = strategy + .create_miniblock_compressor(&field, &fixed_data) + .unwrap(); + let (_block, encoding) = compressor + .compress(miniblock_context(), fixed_data.clone()) + .unwrap(); + check_uncompressed_encoding(&encoding, false); + + let compressor = strategy + .create_miniblock_compressor(&field, &variable_data) + .unwrap(); + let (_block, encoding) = compressor + .compress(miniblock_context(), variable_data.clone()) + .unwrap(); + check_uncompressed_encoding(&encoding, true); + + // Test pervalue + let compressor = strategy.create_per_value(&field, &fixed_data).unwrap(); + let (_block, encoding) = compressor.compress(fixed_data).unwrap(); + check_uncompressed_encoding(&encoding, false); + + let compressor = strategy.create_per_value(&field, &variable_data).unwrap(); + let (_block, encoding) = compressor.compress(variable_data).unwrap(); + check_uncompressed_encoding(&encoding, true); + } + + #[test] + fn test_auto_fsst_disabled_for_binary_fields() { + let strategy = baseline_strategy(CompressionParams::default()); + let field = create_test_field("bytes", DataType::Binary); + let variable_data = create_fsst_candidate_variable_width_block(); + + let miniblock = strategy + .create_miniblock_compressor(&field, &variable_data) + .unwrap(); + let miniblock_debug = format!("{:?}", miniblock); + assert!( + miniblock_debug.contains("BinaryMiniBlockEncoder"), + "expected BinaryMiniBlockEncoder, got: {miniblock_debug}" + ); + assert!( + !miniblock_debug.contains("FsstMiniBlockEncoder"), + "did not expect FsstMiniBlockEncoder, got: {miniblock_debug}" + ); + + let per_value = strategy.create_per_value(&field, &variable_data).unwrap(); + let per_value_debug = format!("{:?}", per_value); + assert!( + per_value_debug.contains("VariableEncoder"), + "expected VariableEncoder, got: {per_value_debug}" + ); + assert!( + !per_value_debug.contains("FsstPerValueEncoder"), + "did not expect FsstPerValueEncoder, got: {per_value_debug}" + ); + } + + #[test] + fn test_auto_fsst_still_enabled_for_utf8_fields() { + let strategy = baseline_strategy(CompressionParams::default()); + let field = create_test_field("text", DataType::Utf8); + let variable_data = create_fsst_candidate_variable_width_block(); + + let miniblock = strategy + .create_miniblock_compressor(&field, &variable_data) + .unwrap(); + let miniblock_debug = format!("{:?}", miniblock); + assert!( + miniblock_debug.contains("FsstMiniBlockEncoder"), + "expected FsstMiniBlockEncoder, got: {miniblock_debug}" + ); + + let per_value = strategy.create_per_value(&field, &variable_data).unwrap(); + let per_value_debug = format!("{:?}", per_value); + assert!( + per_value_debug.contains("FsstPerValueEncoder"), + "expected FsstPerValueEncoder, got: {per_value_debug}" + ); + } + + #[test] + fn test_explicit_fsst_still_supported_for_binary_fields() { + let mut params = CompressionParams::new(); + params.columns.insert( + "bytes".to_string(), + CompressionFieldParams { + compression: Some("fsst".to_string()), + ..Default::default() + }, + ); + + let strategy = baseline_strategy(params); + let field = create_test_field("bytes", DataType::Binary); + let variable_data = create_fsst_candidate_variable_width_block(); + + let miniblock = strategy + .create_miniblock_compressor(&field, &variable_data) + .unwrap(); + let miniblock_debug = format!("{:?}", miniblock); + assert!( + miniblock_debug.contains("FsstMiniBlockEncoder"), + "expected FsstMiniBlockEncoder, got: {miniblock_debug}" + ); + + let per_value = strategy.create_per_value(&field, &variable_data).unwrap(); + let per_value_debug = format!("{:?}", per_value); + assert!( + per_value_debug.contains("FsstPerValueEncoder"), + "expected FsstPerValueEncoder, got: {per_value_debug}" + ); + } + + #[test] + #[cfg(feature = "zstd")] + fn test_compression_level_honored_for_large_per_value() { + let mut params = CompressionParams::new(); + params.columns.insert( + "html".to_string(), + CompressionFieldParams { + compression: Some("zstd".to_string()), + compression_level: Some(19), + ..Default::default() + }, + ); + let strategy = baseline_strategy(params); + let field = create_test_field("html", DataType::Utf8); + let large = create_variable_width_block(32, 64, 40 * 1024); + + let per_value = strategy.create_per_value(&field, &large).unwrap(); + let debug = format!("{per_value:?}"); + assert!( + debug.contains("ZstdBufferCompressor") && debug.contains("compression_level: 19"), + "expected zstd level 19 to reach the per-value compressor, got: {debug}" + ); + } + + #[test] + fn test_parameter_merge_priority() { + let mut params = CompressionParams::new(); + + // Set type-level + params.types.insert( + "Int32".to_string(), + CompressionFieldParams { + rle_threshold: Some(0.5), + compression: Some("lz4".to_string()), + ..Default::default() + }, + ); + + // Set column-level (highest priority) + params.columns.insert( + "user_id".to_string(), + CompressionFieldParams { + rle_threshold: Some(0.2), + compression: Some("zstd".to_string()), + compression_level: Some(6), + bss: None, + minichunk_size: None, + }, + ); + + // Get merged params + let merged = params.get_field_params("user_id", &DataType::Int32); + + // Column params should override type params + assert_eq!(merged.rle_threshold, Some(0.2)); + assert_eq!(merged.compression, Some("zstd".to_string())); + assert_eq!(merged.compression_level, Some(6)); + + // Test field with only type params + let merged = params.get_field_params("other_field", &DataType::Int32); + assert_eq!(merged.rle_threshold, Some(0.5)); + assert_eq!(merged.compression, Some("lz4".to_string())); + assert_eq!(merged.compression_level, None); + } + + #[test] + fn test_pattern_matching() { + let mut params = CompressionParams::new(); + + // Configure pattern for log files + params.columns.insert( + "log_*".to_string(), + CompressionFieldParams { + compression: Some("zstd".to_string()), + compression_level: Some(6), + ..Default::default() + }, + ); + + // Should match pattern + let merged = params.get_field_params("log_messages", &DataType::Utf8); + assert_eq!(merged.compression, Some("zstd".to_string())); + assert_eq!(merged.compression_level, Some(6)); + + // Should not match + let merged = params.get_field_params("messages_log", &DataType::Utf8); + assert_eq!(merged.compression, None); + } + + #[test] + fn test_legacy_metadata_support() { + let params = CompressionParams::new(); + let strategy = baseline_strategy(params); + + // Test field with "none" compression metadata + let mut metadata = HashMap::new(); + metadata.insert(COMPRESSION_META_KEY.to_string(), "none".to_string()); + let mut field = create_test_field("some_column", DataType::Int32); + field.metadata = metadata; + + let data = create_fixed_width_block(32, 1000); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + + // Should respect metadata and use ValueEncoder + assert!(format!("{:?}", compressor).contains("ValueEncoder")); + } + + #[test] + fn test_default_behavior() { + // Empty params should fall back to default behavior + let params = CompressionParams::new(); + let strategy = baseline_strategy(params); + + let field = create_test_field("random_column", DataType::Int32); + // Create data with high run count that won't trigger RLE (600 runs for 1000 values = 0.6 ratio) + let data = create_fixed_width_block_with_stats(32, 1000, 600); + + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + // Should use default strategy's decision + let debug_str = format!("{:?}", compressor); + assert!(debug_str.contains("ValueEncoder") || debug_str.contains("InlineBitpacking")); + } + + #[test] + fn test_field_metadata_compression() { + let params = CompressionParams::new(); + let strategy = baseline_strategy(params); + + // Test field with compression metadata + let mut metadata = HashMap::new(); + metadata.insert(COMPRESSION_META_KEY.to_string(), "zstd".to_string()); + metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), "6".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let data = create_fixed_width_block(32, 1000); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + + // Should use zstd with level 6 + let debug_str = format!("{:?}", compressor); + assert!(debug_str.contains("GeneralMiniBlockCompressor")); + } + + #[test] + fn test_field_metadata_rle_threshold() { + let params = CompressionParams::new(); + let strategy = baseline_strategy(params); + + // Test field with RLE threshold metadata + let mut metadata = HashMap::new(); + metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "0.8".to_string()); + metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); // Disable BSS to test RLE + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + // Create data with low run count (e.g., 100 runs for 1000 values = 0.1 ratio) + // This ensures run_count (100) < num_values * threshold (1000 * 0.8 = 800) + let data = create_fixed_width_block_with_stats(32, 1000, 100); + + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + + // Should use RLE because run_count (100) < num_values * threshold (800) + let debug_str = format!("{:?}", compressor); + assert!(debug_str.contains("RleEncoder")); + } + + #[test] + fn test_rle_v2_miniblock_selects_u16_run_lengths() { + let mut metadata = HashMap::new(); + metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string()); + metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let values = vec![7i32; 1000]; + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 1000, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 16); + } + + #[test] + fn test_rle_v2_miniblock_keeps_u8_run_lengths_before_v2_3() { + for version in [TestEncoding::StructuralU16, TestEncoding::StructuralU32] { + let mut metadata = HashMap::new(); + metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string()); + metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let values = vec![7i32; 1000]; + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 1000, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = strategy(version, CompressionParams::default()); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 8, "version={version}"); + } + } + + #[test] + fn test_rle_v2_uses_selected_width_cost_before_bitpacking() { + let mut metadata = HashMap::new(); + metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string()); + metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let values = vec![0i32; 4096]; + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 4096, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{compressor:?}"); + assert!(debug_str.contains("RleEncoder")); + + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 16); + } + + #[test] + fn test_rle_v2_sorted_dictionary_indices_select_u16_run_lengths() { + let field = create_test_field("dict_indices", DataType::Int32); + + let mut values = Vec::with_capacity(1_200); + for value in 0..4 { + values.extend(std::iter::repeat_n(value, 300)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 1_200, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 16); + } + + #[test] + fn test_rle_v2_short_runs_keep_u8_run_lengths() { + let field = create_test_field("dict_indices", DataType::Int32); + + let mut values = Vec::with_capacity(1_280); + for value in 0..10 { + values.extend(std::iter::repeat_n(value, 128)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 1_280, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 8); + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_rle_miniblock_released_versions_keep_flat_children_when_compression_requested() { + for version in [TestEncoding::StructuralU16, TestEncoding::StructuralU32] { + let mut params = CompressionParams::new(); + params.columns.insert( + "dict_indices".to_string(), + CompressionFieldParams { + compression: Some( + if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string(), + ), + rle_threshold: Some(1.0), + bss: Some(BssMode::Off), + ..Default::default() + }, + ); + let strategy = strategy(version, params); + let field = create_test_field("dict_indices", DataType::UInt32); + + let mut values = Vec::with_capacity(8192 * 4); + for value in 0..8192u32 { + values.extend(std::iter::repeat_n(value, 4)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 8192 * 4, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); + let rle = expect_rle_encoding(&encoding); + + assert!( + matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + Compression::Flat(_) + ), + "version={version}" + ); + assert!( + matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + Compression::Flat(_) + ), + "version={version}" + ); + } + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_strategy_bitpacks_child_values_when_smaller() { + let field = create_test_field("dict_indices", DataType::Int32); + + let mut values = Vec::with_capacity(8192 * 4); + for value in 0..8192 { + values.extend(std::iter::repeat_n(value, 4)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 8192 * 4, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{compressor:?}"); + assert!(debug_str.contains("RleEncoder")); + + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); + let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { + panic!("expected RLE encoding"); + }; + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + Compression::OutOfLineBitpacking(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + Compression::Flat(_) + )); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_keeps_child_bitpacked_rle_when_smaller_than_inline_bitpacking() { + let field = create_test_field("int_score", DataType::UInt64); + + let mut values = Vec::with_capacity(8192 * 8); + for run_idx in 0..8192 { + let value = match run_idx % 3 { + 0 => 3u64, + 1 => 4u64, + _ => 5u64, + }; + values.extend(std::iter::repeat_n(value, 8)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(values), + num_values: 8192 * 8, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{compressor:?}"); + assert!( + debug_str.contains("RleEncoder"), + "expected RLE to beat inline bitpacking after child selection, got: {debug_str}" + ); + + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); + let rle = expect_rle_encoding(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + Compression::OutOfLineBitpacking(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + Compression::Flat(_) + )); + } + + #[test] + fn test_field_metadata_override_params() { + // Set up params with one configuration + let mut params = CompressionParams::new(); + params.columns.insert( + "test_column".to_string(), + CompressionFieldParams { + rle_threshold: Some(0.3), + compression: Some("lz4".to_string()), + compression_level: None, + bss: None, + minichunk_size: None, + }, + ); + + let strategy = baseline_strategy(params); + + // Field metadata should override params + let mut metadata = HashMap::new(); + metadata.insert(COMPRESSION_META_KEY.to_string(), "none".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let data = create_fixed_width_block(32, 1000); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + + // Should use none compression (from metadata) instead of lz4 (from params) + assert!(format!("{:?}", compressor).contains("ValueEncoder")); + } + + #[test] + fn test_field_metadata_mixed_configuration() { + // Configure type-level params + let mut params = CompressionParams::new(); + params.types.insert( + "Int32".to_string(), + CompressionFieldParams { + rle_threshold: Some(0.5), + compression: Some("lz4".to_string()), + ..Default::default() + }, + ); + + let strategy = baseline_strategy(params); + + // Field metadata provides partial override + let mut metadata = HashMap::new(); + metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), "3".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let data = create_fixed_width_block(32, 1000); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + + // Should use lz4 (from type params) with level 3 (from metadata) + let debug_str = format!("{:?}", compressor); + assert!(debug_str.contains("GeneralMiniBlockCompressor")); + } + + #[test] + fn test_bss_field_metadata() { + let params = CompressionParams::new(); + let strategy = baseline_strategy(params); + + // Test BSS "on" mode with compression enabled (BSS requires compression to be effective) + let mut metadata = HashMap::new(); + metadata.insert(BSS_META_KEY.to_string(), "on".to_string()); + metadata.insert(COMPRESSION_META_KEY.to_string(), "lz4".to_string()); + let arrow_field = + ArrowField::new("temperature", DataType::Float32, false).with_metadata(metadata); + let field = Field::try_from(&arrow_field).unwrap(); + + // Create float data + let data = create_fixed_width_block(32, 100); + + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{:?}", compressor); + assert!(debug_str.contains("ByteStreamSplitEncoder")); + } + + #[test] + fn test_bss_with_compression() { + let params = CompressionParams::new(); + let strategy = baseline_strategy(params); + + // Test BSS with LZ4 compression + let mut metadata = HashMap::new(); + metadata.insert(BSS_META_KEY.to_string(), "on".to_string()); + metadata.insert(COMPRESSION_META_KEY.to_string(), "lz4".to_string()); + let arrow_field = + ArrowField::new("sensor_data", DataType::Float64, false).with_metadata(metadata); + let field = Field::try_from(&arrow_field).unwrap(); + + // Create double data + let data = create_fixed_width_block(64, 100); + + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{:?}", compressor); + // Should have BSS wrapped in general compression + assert!(debug_str.contains("GeneralMiniBlockCompressor")); + assert!(debug_str.contains("ByteStreamSplitEncoder")); + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_general_block_decompression_fixed_width_v2_2() { + // Request general compression via the write path (2.2 requirement) and ensure the read path mirrors it. + let mut params = CompressionParams::new(); + params.columns.insert( + "dict_values".to_string(), + CompressionFieldParams { + compression: Some(if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string()), + ..Default::default() + }, + ); + + let strategy = strategy(TestEncoding::StructuralU32, params); + + let field = create_test_field("dict_values", DataType::FixedSizeBinary(3)); + let data = create_fixed_width_block(24, 1024); + let DataBlock::FixedWidth(expected_block) = &data else { + panic!("expected fixed width block"); + }; + let expected_bits = expected_block.bits_per_value; + let expected_num_values = expected_block.num_values; + let num_values = expected_num_values; + + let compressor = strategy + .create_block_compressor(&field, &data) + .expect("general compression should be selected"); + let (compressed_buffer, encoding) = compressor + .compress(data.clone()) + .expect("write path general compression should succeed"); + match encoding.compression.as_ref() { + Some(Compression::General(_)) => {} + other => panic!("expected general compression, got {:?}", other), + } + + let decompressor = DefaultDecompressionStrategy::default() + .create_block_decompressor(&encoding) + .expect("general block decompressor should be created"); + + let decoded = decompressor + .decompress(compressed_buffer, num_values) + .expect("decompression should succeed"); + + match decoded { + DataBlock::FixedWidth(block) => { + assert_eq!(block.bits_per_value, expected_bits); + assert_eq!(block.num_values, expected_num_values); + assert_eq!(block.data.as_ref(), expected_block.data.as_ref()); + } + _ => panic!("expected fixed width block"), + } + } + + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn assert_general_block_preserves_compression_level( + compression: &str, + expected_scheme: crate::format::pb21::CompressionScheme, + compression_level: Option, + ) { + let mut params = CompressionParams::new(); + params.columns.insert( + "dict_values".to_string(), + CompressionFieldParams { + compression: Some(compression.to_string()), + compression_level, + ..Default::default() + }, + ); + let strategy = strategy(TestEncoding::StructuralU32, params); + let field = create_test_field("dict_values", DataType::FixedSizeBinary(3)); + let data = create_fixed_width_block(24, 1024); + + let compressor = strategy.create_block_compressor(&field, &data).unwrap(); + let (_, encoding) = compressor.compress(data).unwrap(); + let Some(Compression::General(general)) = encoding.compression.as_ref() else { + panic!("expected general compression"); + }; + + assert_eq!( + general.compression.as_ref(), + Some(&crate::format::pb21::BufferCompression { + scheme: expected_scheme as i32, + level: compression_level, + }) + ); + } + + #[test] + #[cfg(feature = "zstd")] + fn test_general_block_preserves_absent_zstd_level() { + assert_general_block_preserves_compression_level( + "zstd", + crate::format::pb21::CompressionScheme::CompressionAlgorithmZstd, + None, + ); + } + + #[test] + #[cfg(feature = "lz4")] + fn test_general_block_preserves_explicit_lz4_level() { + assert_general_block_preserves_compression_level( + "lz4", + crate::format::pb21::CompressionScheme::CompressionAlgorithmLz4, + Some(7), + ); + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_general_compression_not_selected_for_v2_1_even_if_requested() { + let mut params = CompressionParams::new(); + params.columns.insert( + "dict_values".to_string(), + CompressionFieldParams { + compression: Some(if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string()), + ..Default::default() + }, + ); + + let strategy = strategy(TestEncoding::StructuralU16, params); + let field = create_test_field("dict_values", DataType::FixedSizeBinary(3)); + let data = create_fixed_width_block(24, 1024); + + let compressor = strategy + .create_block_compressor(&field, &data) + .expect("block compressor selection should succeed"); + let (_, encoding) = compressor.compress(data).unwrap(); + + assert!( + !matches!(encoding.compression.as_ref(), Some(Compression::General(_))), + "general compression should not be selected for V2.1" + ); + } + + #[test] + fn test_none_compression_disables_auto_general_block_compression() { + let mut params = CompressionParams::new(); + params.columns.insert( + "dict_values".to_string(), + CompressionFieldParams { + compression: Some("none".to_string()), + ..Default::default() + }, + ); + + let strategy = strategy(TestEncoding::StructuralU32, params); + let field = create_test_field("dict_values", DataType::FixedSizeBinary(3)); + let data = create_fixed_width_block(24, 20_000); + + assert!( + data.data_size() > MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION, + "test requires block size above automatic general compression threshold" + ); + + let compressor = strategy + .create_block_compressor(&field, &data) + .expect("block compressor selection should succeed"); + let (_, encoding) = compressor.compress(data).unwrap(); + + assert!( + !matches!(encoding.compression.as_ref(), Some(Compression::General(_))), + "compression=none should disable automatic block general compression" + ); + } + + #[test] + fn test_rle_v2_block_selects_u32_run_lengths() { + let field = create_test_field("dict_indices", DataType::Int32); + let expected_values = vec![42i32; 70_000]; + let mut block = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(expected_values.clone()), + num_values: expected_values.len() as u64, + block_info: BlockInfo::default(), + }; + block.compute_stat(); + let data = DataBlock::FixedWidth(block); + + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::new()); + let compressor = strategy.create_block_compressor(&field, &data).unwrap(); + let (compressed, encoding) = compressor.compress(data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 32); + let decompressor = DefaultDecompressionStrategy::default() + .create_block_decompressor(&encoding) + .unwrap(); + let decoded = decompressor + .decompress(compressed, expected_values.len() as u64) + .unwrap(); + + match decoded { + DataBlock::FixedWidth(block) => { + let values = block.data.borrow_to_typed_slice::(); + assert_eq!(values.as_ref(), expected_values); + } + _ => panic!("expected fixed-width block"), + } + } + + #[test] + fn test_rle_v2_block_keeps_u8_run_lengths_for_v2_2() { + let field = create_test_field("dict_indices", DataType::Int32); + let values = vec![42i32; 70_000]; + let mut block = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 70_000, + block_info: BlockInfo::default(), + }; + block.compute_stat(); + let data = DataBlock::FixedWidth(block); + + let strategy = strategy(TestEncoding::StructuralU32, CompressionParams::new()); + let compressor = strategy.create_block_compressor(&field, &data).unwrap(); + let (_, encoding) = compressor.compress(data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 8); + } + + #[test] + fn test_rle_block_used_for_version_v2_2() { + let field = create_test_field("test_repdef", DataType::UInt16); + + // Create highly repetitive data + let num_values = 1000u64; + let mut data = Vec::with_capacity(num_values as usize); + for i in 0..10 { + for _ in 0..100 { + data.push(i as u16); + } + } + + let mut block = FixedWidthDataBlock { + bits_per_value: 16, + data: LanceBuffer::reinterpret_vec(data), + num_values, + block_info: BlockInfo::default(), + }; + + block.compute_stat(); + + let data_block = DataBlock::FixedWidth(block); + + let strategy = strategy(TestEncoding::StructuralU32, CompressionParams::new()); + + let compressor = strategy + .create_block_compressor(&field, &data_block) + .unwrap(); + + let debug_str = format!("{:?}", compressor); + assert!(debug_str.contains("RleEncoder")); + } + + #[test] + fn test_rle_block_not_used_for_version_v2_1() { + let field = create_test_field("test_repdef", DataType::UInt16); + + // Create highly repetitive data + let num_values = 1000u64; + let mut data = Vec::with_capacity(num_values as usize); + for i in 0..10 { + for _ in 0..100 { + data.push(i as u16); + } + } + + let mut block = FixedWidthDataBlock { + bits_per_value: 16, + data: LanceBuffer::reinterpret_vec(data), + num_values, + block_info: BlockInfo::default(), + }; + + block.compute_stat(); + + let data_block = DataBlock::FixedWidth(block); + + let strategy = strategy(TestEncoding::StructuralU16, CompressionParams::new()); + + let compressor = strategy + .create_block_compressor(&field, &data_block) + .unwrap(); + + let debug_str = format!("{:?}", compressor); + assert!( + !debug_str.contains("RleEncoder"), + "RLE should not be used for V2.1" + ); + } +} diff --git a/vendor/lance-encoding/src/compression_config.rs b/vendor/lance-encoding/src/compression_config.rs new file mode 100644 index 000000000..4aee75b21 --- /dev/null +++ b/vendor/lance-encoding/src/compression_config.rs @@ -0,0 +1,299 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Configuration for compression parameters +//! +//! This module provides types for configuring compression strategies +//! on a per-column or per-type basis using a parameter-driven approach. + +use std::collections::HashMap; + +use arrow_schema::DataType; + +/// Byte stream split encoding mode +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum BssMode { + /// Never use BSS + Off, + /// Always use BSS for floating point data + On, + /// Automatically decide based on data characteristics + Auto, +} + +impl BssMode { + /// Convert to internal sensitivity value + pub fn to_sensitivity(&self) -> f32 { + match self { + Self::Off => 0.0, + Self::On => 1.0, + Self::Auto => 0.5, // Default sensitivity for auto mode + } + } + + /// Parse from string + pub fn parse(s: &str) -> Option { + match s.to_lowercase().as_str() { + "off" => Some(Self::Off), + "on" => Some(Self::On), + "auto" => Some(Self::Auto), + _ => None, + } + } +} + +/// Compression parameter configuration +#[derive(Debug, Clone, PartialEq)] +pub struct CompressionParams { + /// Column-level parameters: column name/pattern -> parameters + pub columns: HashMap, + + /// Type-level parameters: data type name -> parameters + pub types: HashMap, +} + +/// Field-level compression parameters +#[derive(Debug, Clone, PartialEq, Default)] +pub struct CompressionFieldParams { + /// RLE threshold (0.0-1.0) + /// When run_count < num_values * threshold, RLE will be used + pub rle_threshold: Option, + + /// General compression scheme: "lz4", "zstd", "none" + pub compression: Option, + + /// Compression level (only for schemes that support it, e.g., zstd) + pub compression_level: Option, + + /// Byte stream split mode for floating point data + pub bss: Option, + + /// Minichunk size threshold for encoding + pub minichunk_size: Option, +} + +impl CompressionParams { + /// Create empty compression parameters + pub fn new() -> Self { + Self { + columns: HashMap::new(), + types: HashMap::new(), + } + } + + /// Get effective parameters for a field (merging type params and column params) + pub fn get_field_params( + &self, + field_name: &str, + data_type: &DataType, + ) -> CompressionFieldParams { + let mut params = CompressionFieldParams::default(); + + // Apply type-level parameters + let type_name = data_type.to_string(); + if let Some(type_params) = self.types.get(&type_name) { + params.merge(type_params); + } + + // Apply column-level parameters (highest priority) + // First check exact match + if let Some(col_params) = self.columns.get(field_name) { + params.merge(col_params); + } else { + // Check pattern matching + for (pattern, col_params) in &self.columns { + if matches_pattern(field_name, pattern) { + params.merge(col_params); + break; // Use first matching pattern + } + } + } + + params + } +} + +impl Default for CompressionParams { + fn default() -> Self { + Self::new() + } +} + +impl CompressionFieldParams { + /// Merge another CompressionFieldParams, non-None values will override + pub fn merge(&mut self, other: &Self) { + if other.rle_threshold.is_some() { + self.rle_threshold = other.rle_threshold; + } + if other.compression.is_some() { + self.compression = other.compression.clone(); + } + if other.compression_level.is_some() { + self.compression_level = other.compression_level; + } + if other.bss.is_some() { + self.bss = other.bss; + } + if other.minichunk_size.is_some() { + self.minichunk_size = other.minichunk_size; + } + } +} + +/// Check if a name matches a pattern (supports wildcards) +fn matches_pattern(name: &str, pattern: &str) -> bool { + if pattern == "*" { + return true; + } + + if let Some(prefix) = pattern.strip_suffix('*') { + return name.starts_with(prefix); + } + + if let Some(suffix) = pattern.strip_prefix('*') { + return name.ends_with(suffix); + } + + if pattern.contains('*') { + // Simple glob pattern matching (only supports single * in middle) + if let Some(pos) = pattern.find('*') { + let prefix = &pattern[..pos]; + let suffix = &pattern[pos + 1..]; + return name.starts_with(prefix) + && name.ends_with(suffix) + && name.len() >= pattern.len() - 1; + } + } + + name == pattern +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_pattern_matching() { + assert!(matches_pattern("user_id", "*_id")); + assert!(matches_pattern("product_id", "*_id")); + assert!(!matches_pattern("identity", "*_id")); + + assert!(matches_pattern("log_message", "log_*")); + assert!(matches_pattern("log_level", "log_*")); + assert!(!matches_pattern("message_log", "log_*")); + + assert!(matches_pattern("test_field_name", "test_*_name")); + assert!(matches_pattern("test_column_name", "test_*_name")); + assert!(!matches_pattern("test_name", "test_*_name")); + + assert!(matches_pattern("anything", "*")); + assert!(matches_pattern("exact_match", "exact_match")); + } + + #[test] + fn test_field_params_merge() { + let mut params = CompressionFieldParams::default(); + assert_eq!(params.rle_threshold, None); + assert_eq!(params.compression, None); + assert_eq!(params.compression_level, None); + assert_eq!(params.bss, None); + + let other = CompressionFieldParams { + rle_threshold: Some(0.3), + compression: Some("lz4".to_string()), + compression_level: None, + bss: Some(BssMode::On), + minichunk_size: None, + }; + + params.merge(&other); + assert_eq!(params.rle_threshold, Some(0.3)); + assert_eq!(params.compression, Some("lz4".to_string())); + assert_eq!(params.compression_level, None); + assert_eq!(params.bss, Some(BssMode::On)); + + let another = CompressionFieldParams { + rle_threshold: None, + compression: Some("zstd".to_string()), + compression_level: Some(3), + bss: Some(BssMode::Auto), + minichunk_size: None, + }; + + params.merge(&another); + assert_eq!(params.rle_threshold, Some(0.3)); // Not overridden + assert_eq!(params.compression, Some("zstd".to_string())); // Overridden + assert_eq!(params.compression_level, Some(3)); // New value + assert_eq!(params.bss, Some(BssMode::Auto)); // Overridden + } + + #[test] + fn test_get_field_params() { + let mut params = CompressionParams::new(); + + // Set type-level params + params.types.insert( + "Int32".to_string(), + CompressionFieldParams { + rle_threshold: Some(0.5), + compression: Some("lz4".to_string()), + ..Default::default() + }, + ); + + // Set column-level params + params.columns.insert( + "*_id".to_string(), + CompressionFieldParams { + rle_threshold: Some(0.3), + compression: Some("zstd".to_string()), + compression_level: Some(3), + bss: None, + minichunk_size: None, + }, + ); + + // Test no match (should get default) + let field_params = params.get_field_params("some_field", &DataType::Float32); + assert_eq!(field_params.compression, None); + assert_eq!(field_params.rle_threshold, None); + + // Test type match only + let field_params = params.get_field_params("some_field", &DataType::Int32); + assert_eq!(field_params.compression, Some("lz4".to_string())); // From type + assert_eq!(field_params.rle_threshold, Some(0.5)); // From type + + // Test column override (pattern match) + let field_params = params.get_field_params("user_id", &DataType::Int32); + assert_eq!(field_params.compression, Some("zstd".to_string())); // From column + assert_eq!(field_params.compression_level, Some(3)); // From column + assert_eq!(field_params.rle_threshold, Some(0.3)); // From column (overrides type) + } + + #[test] + fn test_exact_match_priority() { + let mut params = CompressionParams::new(); + + // Add pattern + params.columns.insert( + "*_id".to_string(), + CompressionFieldParams { + compression: Some("lz4".to_string()), + ..Default::default() + }, + ); + + // Add exact match + params.columns.insert( + "user_id".to_string(), + CompressionFieldParams { + compression: Some("zstd".to_string()), + ..Default::default() + }, + ); + + // Exact match should win + let field_params = params.get_field_params("user_id", &DataType::Int32); + assert_eq!(field_params.compression, Some("zstd".to_string())); + } +} diff --git a/vendor/lance-encoding/src/constants.rs b/vendor/lance-encoding/src/constants.rs new file mode 100644 index 000000000..0bd31d676 --- /dev/null +++ b/vendor/lance-encoding/src/constants.rs @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Constants for Lance encoding metadata keys +//! +//! These constants define the metadata keys used in Arrow field metadata +//! to configure various encoding behaviors in Lance. + +// Compression-related metadata keys +/// Metadata key for specifying compression scheme (e.g., "lz4", "zstd", "none") +pub const COMPRESSION_META_KEY: &str = "lance-encoding:compression"; +/// Metadata key for specifying compression level (applies to schemes that support levels) +pub const COMPRESSION_LEVEL_META_KEY: &str = "lance-encoding:compression-level"; +/// Metadata key for specifying RLE (Run-Length Encoding) threshold +pub const RLE_THRESHOLD_META_KEY: &str = "lance-encoding:rle-threshold"; +/// Metadata key for specifying minichunk size +pub const MINICHUNK_SIZE_META_KEY: &str = "lance-encoding:minichunk-size"; + +// Dictionary encoding metadata keys +/// Metadata key for specifying dictionary encoding threshold divisor +/// Set to a large value to discourage dictionary encoding +/// Set to a small value to encourage dictionary encoding +pub const DICT_DIVISOR_META_KEY: &str = "lance-encoding:dict-divisor"; +/// Metadata key for dictionary encoding size ratio threshold (0.0-1.0] +/// If estimated_dict_size/raw_size < ratio, use dictionary encoding. +/// Example: 0.8 means use dict if encoded size < 80% of raw size +/// Default: 0.8 +pub const DICT_SIZE_RATIO_META_KEY: &str = "lance-encoding:dict-size-ratio"; +/// Metadata key for selecting general compression scheme for dictionary values +/// Valid values: "lz4", "zstd", "none" +pub const DICT_VALUES_COMPRESSION_META_KEY: &str = "lance-encoding:dict-values-compression"; +/// Metadata key for selecting compression level for dictionary values +/// Applies to schemes that support levels (e.g. zstd) +pub const DICT_VALUES_COMPRESSION_LEVEL_META_KEY: &str = + "lance-encoding:dict-values-compression-level"; + +/// Environment variable for selecting general compression scheme for dictionary values +pub const DICT_VALUES_COMPRESSION_ENV_VAR: &str = "LANCE_ENCODING_DICT_VALUES_COMPRESSION"; +/// Environment variable for selecting compression level for dictionary values +pub const DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR: &str = + "LANCE_ENCODING_DICT_VALUES_COMPRESSION_LEVEL"; + +// NOTE: BLOB_META_KEY is defined in lance-core to avoid circular dependency + +// Packed struct encoding metadata keys +/// Legacy metadata key for packed struct encoding (deprecated) +pub const PACKED_STRUCT_LEGACY_META_KEY: &str = "packed"; +/// Metadata key for packed struct encoding +pub const PACKED_STRUCT_META_KEY: &str = "lance-encoding:packed"; + +// Structural encoding metadata keys +/// Metadata key for specifying structural encoding type +pub const STRUCTURAL_ENCODING_META_KEY: &str = "lance-encoding:structural-encoding"; +/// Value for miniblock structural encoding +pub const STRUCTURAL_ENCODING_MINIBLOCK: &str = "miniblock"; +/// Value for fullzip structural encoding +pub const STRUCTURAL_ENCODING_FULLZIP: &str = "fullzip"; +/// Value for sparse structural encoding +pub const STRUCTURAL_ENCODING_SPARSE: &str = "sparse"; + +// Byte stream split metadata keys +/// Metadata key for byte stream split encoding configuration +pub const BSS_META_KEY: &str = "lance-encoding:bss"; +/// Default BSS mode +pub const DEFAULT_BSS_MODE: &str = "auto"; diff --git a/vendor/lance-encoding/src/data.rs b/vendor/lance-encoding/src/data.rs new file mode 100644 index 000000000..84ea0300f --- /dev/null +++ b/vendor/lance-encoding/src/data.rs @@ -0,0 +1,2898 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Data layouts to represent encoded data in a sub-Arrow format +//! +//! These [`DataBlock`] structures represent physical layouts. They fill a gap somewhere +//! between [`arrow_data::ArrayData`] (which, as a collection of buffers, is too +//! generic because it doesn't give us enough information about what those buffers represent) +//! and [`arrow_array::array::Array`] (which is too specific, because it cares about the +//! logical data type). +//! +//! In addition, the layouts represented here are slightly stricter than Arrow's layout rules. +//! For example, offset buffers MUST start with 0. These additional restrictions impose a +//! slight penalty on encode (to normalize arrow data) but make the development of encoders +//! and decoders easier (since they can rely on a normalized representation) + +use std::{ + ops::Range, + sync::{Arc, RwLock}, +}; + +use arrow_array::{ + Array, ArrayRef, OffsetSizeTrait, UInt64Array, + cast::AsArray, + new_empty_array, new_null_array, + types::{ArrowDictionaryKeyType, UInt8Type, UInt16Type, UInt32Type, UInt64Type}, +}; +use arrow_buffer::{ArrowNativeType, BooleanBuffer, BooleanBufferBuilder, NullBuffer}; +use arrow_data::{ArrayData, ArrayDataBuilder}; +use arrow_schema::DataType; +use lance_arrow::DataTypeExt; + +use lance_core::{Error, Result}; + +use crate::{ + buffer::LanceBuffer, + statistics::{ComputeStat, Stat}, +}; + +/// A data block with no buffers where everything is null +/// +/// Note: this data block should not be used for future work. It will be deprecated +/// in the 2.1 version of the format where nullability will be handled by the structural +/// encoders. +#[derive(Debug, Clone)] +pub struct AllNullDataBlock { + /// The number of values represented by this block + pub num_values: u64, +} + +impl AllNullDataBlock { + fn into_arrow(self, data_type: DataType, _validate: bool) -> Result { + Ok(ArrayData::new_null(&data_type, self.num_values as usize)) + } + + fn into_buffers(self) -> Vec { + vec![] + } +} + +use std::collections::HashMap; + +// `BlockInfo` stores the statistics of this `DataBlock`, such as `NullCount` for `NullableDataBlock`, +// `BitWidth` for `FixedWidthDataBlock`, `Cardinality` for all `DataBlock` +#[derive(Debug, Clone)] +pub struct BlockInfo(pub Arc>>>); + +impl Default for BlockInfo { + fn default() -> Self { + Self::new() + } +} + +impl BlockInfo { + pub fn new() -> Self { + Self(Arc::new(RwLock::new(HashMap::new()))) + } +} + +impl PartialEq for BlockInfo { + fn eq(&self, other: &Self) -> bool { + let self_info = self.0.read().unwrap(); + let other_info = other.0.read().unwrap(); + *self_info == *other_info + } +} + +/// Wraps a data block and adds nullability information to it +/// +/// Note: this data block should not be used for future work. It will be deprecated +/// in the 2.1 version of the format where nullability will be handled by the structural +/// encoders. +#[derive(Debug, Clone)] +pub struct NullableDataBlock { + /// The underlying data + pub data: Box, + /// A bitmap of validity for each value + pub nulls: LanceBuffer, + + pub block_info: BlockInfo, +} + +impl NullableDataBlock { + fn into_arrow(self, data_type: DataType, _validate: bool) -> Result { + let nulls = self.nulls.into_buffer(); + let data = self.data.into_arrow_impl(data_type, true)?.into_builder(); + let data = data.null_bit_buffer(Some(nulls)); + Ok(data.build()?) + } + + fn into_buffers(self) -> Vec { + let mut buffers = vec![self.nulls]; + buffers.extend(self.data.into_buffers()); + buffers + } + + pub fn data_size(&self) -> u64 { + self.data.data_size() + self.nulls.len() as u64 + } +} + +/// A block representing the same constant value repeated many times +#[derive(Debug, PartialEq, Clone)] +pub struct ConstantDataBlock { + /// Data buffer containing the value + pub data: LanceBuffer, + /// The number of values + pub num_values: u64, +} + +impl ConstantDataBlock { + fn into_buffers(self) -> Vec { + vec![self.data] + } + + fn into_arrow(self, _data_type: DataType, _validate: bool) -> Result { + // We don't need this yet but if we come up with some way of serializing + // scalars to/from bytes then we could implement it. + todo!() + } + + pub fn try_clone(&self) -> Result { + Ok(Self { + data: self.data.clone(), + num_values: self.num_values, + }) + } + + pub fn data_size(&self) -> u64 { + self.data.len() as u64 + } +} + +/// A data block for a single buffer of data where each element has a fixed number of bits +#[derive(Debug, PartialEq, Clone)] +pub struct FixedWidthDataBlock { + /// The data buffer + pub data: LanceBuffer, + /// The number of bits per value + pub bits_per_value: u64, + /// The number of values represented by this block + pub num_values: u64, + + pub block_info: BlockInfo, +} + +impl FixedWidthDataBlock { + fn do_into_arrow( + self, + data_type: DataType, + num_values: u64, + _validate: bool, + ) -> Result { + // Booleans expanded for full-zip (bits_per_value==8, one byte each) need re-packing to + // Arrow's bit-packed format. + let data_buffer = if matches!(data_type, DataType::Boolean) && self.bits_per_value == 8 { + let mut builder = BooleanBufferBuilder::new(num_values as usize); + for &byte in self.data.as_ref().iter().take(num_values as usize) { + builder.append(byte != 0); + } + builder.finish().into_inner() + } else { + self.data.into_buffer() + }; + let builder = ArrayDataBuilder::new(data_type) + .add_buffer(data_buffer) + .len(num_values as usize) + .null_count(0); + Ok(builder.build()?) + } + + /// Convert this block into Arrow data with full layout validation. + /// + /// The `validate` argument is retained for API compatibility. Conversion is + /// always validated because callers can construct this public type directly. + pub fn into_arrow(self, data_type: DataType, _validate: bool) -> Result { + let root_num_values = self.num_values; + self.do_into_arrow(data_type, root_num_values, true) + } + + pub fn into_buffers(self) -> Vec { + vec![self.data] + } + + pub fn try_clone(&self) -> Result { + Ok(Self { + data: self.data.clone(), + bits_per_value: self.bits_per_value, + num_values: self.num_values, + block_info: self.block_info.clone(), + }) + } + + pub fn data_size(&self) -> u64 { + self.data.len() as u64 + } +} + +#[derive(Debug)] +struct VariableWidthDataBlockBuilder { + offsets: Vec, + bytes: Vec, +} + +impl VariableWidthDataBlockBuilder { + fn new(estimated_size_bytes: u64) -> Self { + Self { + offsets: vec![T::from_usize(0).unwrap()], + bytes: Vec::with_capacity(estimated_size_bytes as usize), + } + } +} + +impl DataBlockBuilderImpl for VariableWidthDataBlockBuilder { + fn validate_append(&self, data_block: &DataBlock, selection: &Range) -> Result<()> { + let block = data_block.as_variable_width_ref().unwrap(); + block.validate_offsets_for_append::(selection) + } + + fn append_validated(&mut self, data_block: &DataBlock, selection: Range) -> Result<()> { + let block = data_block.as_variable_width_ref().unwrap(); + debug_assert_eq!(block.bits_per_offset, T::get_byte_width() as u8 * 8); + let offsets = block.offsets.borrow_to_typed_view::(); + + let start_offset = offsets[selection.start as usize]; + let end_offset = offsets[selection.end as usize]; + let selected_data_len = end_offset.as_usize() - start_offset.as_usize(); + let new_data_len = self + .bytes + .len() + .checked_add(selected_data_len) + .ok_or_else(|| { + Error::not_supported_source( + "appending variable-width data would overflow usize".into(), + ) + })?; + if T::from_usize(new_data_len).is_none() { + return Err(Error::not_supported_source( + format!( + "appending variable-width data would require {} bytes, which exceeds the \ + capacity of {}-bit offsets", + new_data_len, + T::get_byte_width() * 8 + ) + .into(), + )); + } + let previous_len = self.bytes.len(); + + self.bytes + .extend_from_slice(&block.data[start_offset.as_usize()..end_offset.as_usize()]); + + self.offsets.extend( + offsets[selection.start as usize + 1..=selection.end as usize] + .iter() + .map(|&offset| { + let rebased_offset = + previous_len + (offset.as_usize() - start_offset.as_usize()); + T::from_usize(rebased_offset).unwrap() + }), + ); + Ok(()) + } + + fn finish(self: Box) -> DataBlock { + let num_values = (self.offsets.len() - 1) as u64; + DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::from(self.bytes), + offsets: LanceBuffer::reinterpret_vec(self.offsets), + bits_per_offset: T::get_byte_width() as u8 * 8, + num_values, + block_info: BlockInfo::new(), + }) + } +} + +#[derive(Debug)] +struct BitmapDataBlockBuilder { + values: BooleanBufferBuilder, +} + +impl BitmapDataBlockBuilder { + fn new(estimated_size_bytes: u64) -> Self { + Self { + values: BooleanBufferBuilder::new(estimated_size_bytes as usize * 8), + } + } +} + +impl DataBlockBuilderImpl for BitmapDataBlockBuilder { + fn validate_append(&self, _data_block: &DataBlock, _selection: &Range) -> Result<()> { + Ok(()) + } + + fn append_validated(&mut self, data_block: &DataBlock, selection: Range) -> Result<()> { + let bitmap_blk = data_block.as_fixed_width_ref().unwrap(); + self.values.append_packed_range( + selection.start as usize..selection.end as usize, + &bitmap_blk.data, + ); + Ok(()) + } + + fn finish(mut self: Box) -> DataBlock { + let bool_buf = self.values.finish(); + let num_values = bool_buf.len() as u64; + let bits_buf = bool_buf.into_inner(); + DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(bits_buf), + bits_per_value: 1, + num_values, + block_info: BlockInfo::new(), + }) + } +} + +#[derive(Debug)] +struct FixedWidthDataBlockBuilder { + bits_per_value: u64, + bytes_per_value: u64, + values: Vec, +} + +impl FixedWidthDataBlockBuilder { + fn new(bits_per_value: u64, estimated_size_bytes: u64) -> Self { + assert!(bits_per_value.is_multiple_of(8)); + Self { + bits_per_value, + bytes_per_value: bits_per_value / 8, + values: Vec::with_capacity(estimated_size_bytes as usize), + } + } +} + +impl DataBlockBuilderImpl for FixedWidthDataBlockBuilder { + fn validate_append(&self, _data_block: &DataBlock, _selection: &Range) -> Result<()> { + Ok(()) + } + + fn append_validated(&mut self, data_block: &DataBlock, selection: Range) -> Result<()> { + let block = data_block.as_fixed_width_ref().unwrap(); + assert_eq!(self.bits_per_value, block.bits_per_value); + let start = selection.start as usize * self.bytes_per_value as usize; + let end = selection.end as usize * self.bytes_per_value as usize; + self.values.extend_from_slice(&block.data[start..end]); + Ok(()) + } + + fn finish(self: Box) -> DataBlock { + let num_values = (self.values.len() / self.bytes_per_value as usize) as u64; + DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(self.values), + bits_per_value: self.bits_per_value, + num_values, + block_info: BlockInfo::new(), + }) + } +} + +#[derive(Debug)] +struct StructDataBlockBuilder { + children: Vec>, +} + +impl StructDataBlockBuilder { + fn new(children: Vec>) -> Self { + Self { children } + } +} + +impl DataBlockBuilderImpl for StructDataBlockBuilder { + fn validate_append(&self, data_block: &DataBlock, selection: &Range) -> Result<()> { + let data_block = data_block.as_struct_ref().unwrap(); + for i in 0..self.children.len() { + self.children[i].validate_append(&data_block.children[i], selection)?; + } + Ok(()) + } + + fn append_validated(&mut self, data_block: &DataBlock, selection: Range) -> Result<()> { + let data_block = data_block.as_struct_ref().unwrap(); + for i in 0..self.children.len() { + self.children[i].append_validated(&data_block.children[i], selection.clone())?; + } + Ok(()) + } + + fn finish(self: Box) -> DataBlock { + let mut children_data_block = Vec::new(); + for child in self.children { + let child_data_block = child.finish(); + children_data_block.push(child_data_block); + } + DataBlock::Struct(StructDataBlock { + children: children_data_block, + block_info: BlockInfo::new(), + validity: None, + }) + } +} + +#[derive(Debug, Default)] +struct AllNullDataBlockBuilder { + num_values: u64, +} + +impl DataBlockBuilderImpl for AllNullDataBlockBuilder { + fn validate_append(&self, _data_block: &DataBlock, _selection: &Range) -> Result<()> { + Ok(()) + } + + fn append_validated(&mut self, _data_block: &DataBlock, selection: Range) -> Result<()> { + self.num_values += selection.end - selection.start; + Ok(()) + } + + fn finish(self: Box) -> DataBlock { + DataBlock::AllNull(AllNullDataBlock { + num_values: self.num_values, + }) + } +} + +/// A data block to represent a fixed size list +#[derive(Debug, Clone)] +pub struct FixedSizeListBlock { + /// The child data block + pub child: Box, + /// The number of items in each list + pub dimension: u64, +} + +impl FixedSizeListBlock { + pub fn num_values(&self) -> u64 { + self.child.num_values() / self.dimension + } + + /// Try to flatten a FixedSizeListBlock into a FixedWidthDataBlock + /// + /// Returns None if any children are nullable + pub fn try_into_flat(self) -> Option { + match *self.child { + // Cannot flatten a nullable child + DataBlock::Nullable(_) => None, + DataBlock::FixedSizeList(inner) => { + let mut flat = inner.try_into_flat()?; + flat.bits_per_value *= self.dimension; + flat.num_values /= self.dimension; + Some(flat) + } + DataBlock::FixedWidth(mut inner) => { + inner.bits_per_value *= self.dimension; + inner.num_values /= self.dimension; + Some(inner) + } + _ => panic!( + "Expected FixedSizeList or FixedWidth data block but found {:?}", + self + ), + } + } + + pub fn flatten_as_fixed(&mut self) -> FixedWidthDataBlock { + match self.child.as_mut() { + DataBlock::FixedSizeList(fsl) => fsl.flatten_as_fixed(), + DataBlock::FixedWidth(fw) => fw.clone(), + _ => panic!("Expected FixedSizeList or FixedWidth data block"), + } + } + + /// Convert a flattened values block into a FixedSizeListBlock + pub fn from_flat(data: FixedWidthDataBlock, data_type: &DataType) -> DataBlock { + match data_type { + DataType::FixedSizeList(child_field, dimension) => { + let mut data = data; + data.bits_per_value /= *dimension as u64; + data.num_values *= *dimension as u64; + let child_data = Self::from_flat(data, child_field.data_type()); + DataBlock::FixedSizeList(Self { + child: Box::new(child_data), + dimension: *dimension as u64, + }) + } + // Base case, we've hit a non-list type + _ => DataBlock::FixedWidth(data), + } + } + + fn into_arrow(self, data_type: DataType, _validate: bool) -> Result { + let num_values = self.num_values(); + let builder = match &data_type { + DataType::FixedSizeList(child_field, _) => { + let child_data = self + .child + .into_arrow_impl(child_field.data_type().clone(), true)?; + ArrayDataBuilder::new(data_type) + .add_child_data(child_data) + .len(num_values as usize) + .null_count(0) + } + _ => panic!("Expected FixedSizeList data type and got {:?}", data_type), + }; + Ok(builder.build()?) + } + + fn into_buffers(self) -> Vec { + self.child.into_buffers() + } + + fn data_size(&self) -> u64 { + self.child.data_size() + } +} + +#[derive(Debug)] +struct FixedSizeListBlockBuilder { + inner: Box, + dimension: u64, +} + +impl FixedSizeListBlockBuilder { + fn new(inner: Box, dimension: u64) -> Self { + Self { inner, dimension } + } +} + +impl DataBlockBuilderImpl for FixedSizeListBlockBuilder { + fn validate_append(&self, data_block: &DataBlock, selection: &Range) -> Result<()> { + let selection = selection.start * self.dimension..selection.end * self.dimension; + let fsl = data_block.as_fixed_size_list_ref().unwrap(); + self.inner.validate_append(fsl.child.as_ref(), &selection) + } + + fn append_validated(&mut self, data_block: &DataBlock, selection: Range) -> Result<()> { + let selection = selection.start * self.dimension..selection.end * self.dimension; + let fsl = data_block.as_fixed_size_list_ref().unwrap(); + self.inner.append_validated(fsl.child.as_ref(), selection) + } + + fn finish(self: Box) -> DataBlock { + let inner_block = self.inner.finish(); + DataBlock::FixedSizeList(FixedSizeListBlock { + child: Box::new(inner_block), + dimension: self.dimension, + }) + } +} + +#[derive(Debug)] +struct NullableDataBlockBuilder { + inner: Box, + validity: BooleanBufferBuilder, +} + +impl NullableDataBlockBuilder { + fn new(inner: Box, estimated_size_bytes: usize) -> Self { + Self { + inner, + validity: BooleanBufferBuilder::new(estimated_size_bytes * 8), + } + } +} + +impl DataBlockBuilderImpl for NullableDataBlockBuilder { + fn validate_append(&self, data_block: &DataBlock, selection: &Range) -> Result<()> { + let nullable = data_block.as_nullable_ref().unwrap(); + self.inner + .validate_append(nullable.data.as_ref(), selection) + } + + fn append_validated(&mut self, data_block: &DataBlock, selection: Range) -> Result<()> { + let nullable = data_block.as_nullable_ref().unwrap(); + self.inner + .append_validated(nullable.data.as_ref(), selection.clone())?; + let bool_buf = BooleanBuffer::new( + nullable.nulls.clone().into_buffer(), + selection.start as usize, + (selection.end - selection.start) as usize, + ); + self.validity.append_buffer(&bool_buf); + Ok(()) + } + + fn finish(mut self: Box) -> DataBlock { + let inner_block = self.inner.finish(); + DataBlock::Nullable(NullableDataBlock { + data: Box::new(inner_block), + nulls: LanceBuffer::from(self.validity.finish().into_inner()), + block_info: BlockInfo::new(), + }) + } +} + +/// A data block with no regular structure. There is no available spot to attach +/// validity / repdef information and it cannot be converted to Arrow without being +/// decoded +#[derive(Debug, Clone)] +pub struct OpaqueBlock { + pub buffers: Vec, + pub num_values: u64, + pub block_info: BlockInfo, +} + +impl OpaqueBlock { + pub fn data_size(&self) -> u64 { + self.buffers.iter().map(|b| b.len() as u64).sum() + } +} + +/// A data block for variable-width data (e.g. strings, packed rows, etc.) +#[derive(Debug, Clone)] +pub struct VariableWidthBlock { + /// The data buffer + pub data: LanceBuffer, + /// The offsets buffer (contains num_values + 1 offsets) + /// + /// Offsets MUST start at 0 + pub offsets: LanceBuffer, + /// The number of bits per offset + pub bits_per_offset: u8, + /// The number of values represented by this block + pub num_values: u64, + + pub block_info: BlockInfo, +} + +/// Proof that a [`VariableWidthBlock`] satisfies the Arrow layout contract for +/// its target data type (offsets buffer long enough, offsets monotonic and +/// within the data buffer, values valid UTF-8 where required). +/// +/// Only [`VariableWidthBlock::validate_layout`] can construct it, which ties the +/// unchecked Arrow build below to an actual validation pass instead of a +/// caller-controlled flag. +struct ValidVariableWidthLayout; + +impl VariableWidthBlock { + fn append_error(&self, selection: &Range, detail: impl std::fmt::Display) -> Error { + Error::corrupt_file_named( + "variable width data block", + format!( + "cannot append offsets for selection {}..{}: {} (num_values: {}, \ + bits_per_offset: {}, offsets buffer size: {} bytes, data buffer size: {} bytes)", + selection.start, + selection.end, + detail, + self.num_values, + self.bits_per_offset, + self.offsets.len(), + self.data.len(), + ), + ) + } + + fn validate_offsets_for_append(&self, selection: &Range) -> Result<()> + where + T: OffsetSizeTrait + bytemuck::Pod, + { + let expected_bits_per_offset = T::get_byte_width() as u8 * 8; + if self.bits_per_offset != expected_bits_per_offset { + return Err(self.append_error( + selection, + format!( + "expected {}-bit offsets but found {}-bit offsets", + expected_bits_per_offset, self.bits_per_offset + ), + )); + } + let offset_size = std::mem::size_of::(); + if !self.offsets.len().is_multiple_of(offset_size) { + return Err(self.append_error( + selection, + format!( + "offsets buffer length {} is not a multiple of the {}-byte offset width", + self.offsets.len(), + offset_size + ), + )); + } + if selection.start > selection.end || selection.end > self.num_values { + return Err( + self.append_error(selection, "selection is outside the block's value range") + ); + } + let selection_start = usize::try_from(selection.start) + .map_err(|_| self.append_error(selection, "selection start does not fit in usize"))?; + let selection_end = usize::try_from(selection.end) + .map_err(|_| self.append_error(selection, "selection end does not fit in usize"))?; + let offsets = self.offsets.borrow_to_typed_view::(); + if selection_end >= offsets.len() { + return Err(self.append_error( + selection, + format!( + "selection requires offset {} but the buffer holds {} offsets", + selection_end, + offsets.len() + ), + )); + } + let selected_offsets = &offsets[selection_start..=selection_end]; + if let Some(detail) = + Self::offset_violation_detail(selected_offsets, self.data.len(), selection_start) + { + return Err(self.append_error(selection, detail)); + } + Ok(()) + } + + // The offsets buffer comes straight from file bytes, so an unchecked build would + // let a corrupt file smuggle out-of-bounds offsets into an Arrow array whose + // consumers then read (or crash on) memory outside the data buffer. This + // boundary therefore always validates the layout, ignoring the optional + // `validate` flag. Lance validates the common layouts itself (a branchless + // scan, measurably cheaper than Arrow's element-wise checked build) and only + // falls back to Arrow's checked build for the cold cases. + fn into_arrow(self, data_type: DataType, _validate: bool) -> Result { + let Some(expected_bits_per_offset) = Self::expected_bits_per_offset(&data_type) else { + // Not an [offsets, bytes] layout we know how to prove; let Arrow + // check it. + return self.into_arrow_checked(data_type); + }; + if self.bits_per_offset != expected_bits_per_offset { + return Err(self.layout_error( + &data_type, + format!( + "expected {}-bit offsets but got {}-bit offsets", + expected_bits_per_offset, self.bits_per_offset + ), + )); + } + if self.num_values == 0 { + // Cold path; Arrow handles the empty-offsets special cases. + return self.into_arrow_checked(data_type); + } + let proof = self.validate_layout(&data_type)?; + Ok(self.into_arrow_unchecked(data_type, proof)) + } + + /// The offset width Arrow mandates for `data_type`, or `None` if the type + /// does not use the `[offsets, bytes]` layout this block represents. + fn expected_bits_per_offset(data_type: &DataType) -> Option { + match data_type { + DataType::Binary | DataType::Utf8 => Some(32), + DataType::LargeBinary | DataType::LargeUtf8 => Some(64), + _ => None, + } + } + + fn layout_error(&self, data_type: &DataType, detail: impl std::fmt::Display) -> Error { + Self::format_layout_error( + data_type, + detail, + self.num_values, + self.bits_per_offset, + self.offsets.len(), + self.data.len(), + ) + } + + fn format_layout_error( + data_type: &DataType, + detail: impl std::fmt::Display, + num_values: u64, + bits_per_offset: u8, + offsets_size: usize, + data_size: usize, + ) -> Error { + Error::corrupt_file_named( + "variable width data block", + format!( + "invalid variable-width layout for {}: {} (num_values: {}, bits_per_offset: {}, \ + offsets buffer size: {} bytes, data buffer size: {} bytes)", + data_type, detail, num_values, bits_per_offset, offsets_size, data_size, + ), + ) + } + + fn validate_layout(&self, data_type: &DataType) -> Result { + let bytes_per_offset = (self.bits_per_offset / 8) as u64; + let required_bytes = self + .num_values + .checked_add(1) + .and_then(|num_offsets| num_offsets.checked_mul(bytes_per_offset)) + .ok_or_else(|| self.layout_error(data_type, "offsets buffer size overflows"))?; + if (self.offsets.len() as u64) < required_bytes { + return Err(self.layout_error( + data_type, + format!( + "offsets buffer must hold at least {} offsets ({} bytes)", + self.num_values + 1, + required_bytes + ), + )); + } + let validate_utf8 = matches!(data_type, DataType::Utf8 | DataType::LargeUtf8); + match self.bits_per_offset { + 32 => self.validate_offsets_and_values::(data_type, validate_utf8), + 64 => self.validate_offsets_and_values::(data_type, validate_utf8), + other => Err(self.layout_error( + data_type, + format!("unsupported offset width: {} bits", other), + )), + } + } + + fn validate_offsets_and_values( + &self, + data_type: &DataType, + validate_utf8: bool, + ) -> Result { + let num_offsets = self.num_values as usize + 1; + // Slice before borrowing: the buffer may carry padding that is not a + // multiple of the offset width. + let offsets = self + .offsets + .slice_with_length(0, num_offsets * std::mem::size_of::()); + let offsets = offsets.borrow_to_typed_slice::(); + let offsets: &[T] = offsets.as_ref(); + let data = self.data.as_ref(); + + if let Some(detail) = Self::offset_violation_detail(offsets, data.len(), 0) { + return Err(self.layout_error(data_type, detail)); + } + + if validate_utf8 { + let (first, last) = (offsets[0].as_usize(), offsets[num_offsets - 1].as_usize()); + let values = std::str::from_utf8(&data[first..last]) + .map_err(|utf8_err| self.layout_error(data_type, utf8_err))?; + let mut on_char_boundaries = true; + for &offset in offsets { + on_char_boundaries &= values.is_char_boundary(offset.as_usize() - first); + } + if !on_char_boundaries { + // Cold path: rescan to pinpoint the offending offset. + let position = offsets + .iter() + .position(|offset| !values.is_char_boundary(offset.as_usize() - first)) + .expect("the fast scan found a non-boundary offset"); + return Err(self.layout_error( + data_type, + format!("offset at position {position} splits a UTF-8 character"), + )); + } + } + + Ok(ValidVariableWidthLayout) + } + + fn offset_violation_detail( + offsets: &[T], + data_size: usize, + position_base: usize, + ) -> Option { + // A monotonic sequence with a non-negative first offset and an + // in-bounds last offset is entirely within [0, data_size]. Keep this + // valid path branchless so it vectorizes, and only rescan on failure. + let mut is_monotonic = true; + for window in offsets.windows(2) { + is_monotonic &= window[0] <= window[1]; + } + let first = offsets[0]; + let last = offsets[offsets.len() - 1]; + let bounds_ok = + first >= T::usize_as(0) && last.to_usize().is_some_and(|last| last <= data_size); + if is_monotonic && bounds_ok { + return None; + } + + for (relative_position, window) in offsets.windows(2).enumerate() { + if window[0] > window[1] { + let position = position_base + relative_position + 1; + return Some(format!( + "non-monotonic offset at position {}: {:?} decreases from {:?}", + position, window[1], window[0] + )); + } + } + for (relative_position, offset) in offsets.iter().enumerate() { + let position = position_base + relative_position; + match offset.to_usize() { + None => { + return Some(format!( + "offset at position {} is negative: {:?}", + position, offset + )); + } + Some(offset) if offset > data_size => { + return Some(format!( + "offset at position {} is out of bounds: {} > {}", + position, offset, data_size + )); + } + Some(_) => {} + } + } + Some("offsets failed validation".to_string()) + } + + fn into_arrow_checked(self, data_type: DataType) -> Result { + let num_values = self.num_values; + let bits_per_offset = self.bits_per_offset; + let offsets_size = self.offsets.len(); + let data_size = self.data.len(); + let builder = self.into_arrow_builder(data_type.clone()); + builder.build().map_err(|arrow_err| { + Self::format_layout_error( + &data_type, + arrow_err, + num_values, + bits_per_offset, + offsets_size, + data_size, + ) + }) + } + + fn into_arrow_unchecked( + self, + data_type: DataType, + _proof: ValidVariableWidthLayout, + ) -> ArrayData { + let builder = self.into_arrow_builder(data_type); + // SAFETY: `_proof` witnesses that `validate_layout` proved this block + // satisfies the Arrow layout contract for `data_type`. + unsafe { builder.build_unchecked() } + } + + fn into_arrow_builder(self, data_type: DataType) -> ArrayDataBuilder { + let num_values = self.num_values; + let data_buffer = self.data.into_buffer(); + let offsets_buffer = self.offsets.into_buffer(); + ArrayDataBuilder::new(data_type) + .add_buffer(offsets_buffer) + .add_buffer(data_buffer) + .len(num_values as usize) + .null_count(0) + } + + fn into_buffers(self) -> Vec { + vec![self.offsets, self.data] + } + + pub fn offsets_as_block(&mut self) -> DataBlock { + let offsets = self.offsets.clone(); + DataBlock::FixedWidth(FixedWidthDataBlock { + data: offsets, + bits_per_value: self.bits_per_offset as u64, + num_values: self.num_values + 1, + block_info: BlockInfo::new(), + }) + } + + pub fn data_size(&self) -> u64 { + (self.data.len() + self.offsets.len()) as u64 + } +} + +/// A data block representing a struct +#[derive(Debug, Clone)] +pub struct StructDataBlock { + /// The child arrays + pub children: Vec, + pub block_info: BlockInfo, + /// The validity bitmap for the struct (None means all valid) + pub validity: Option, +} + +impl StructDataBlock { + fn into_arrow(self, data_type: DataType, _validate: bool) -> Result { + if let DataType::Struct(fields) = &data_type { + let mut builder = ArrayDataBuilder::new(DataType::Struct(fields.clone())); + let mut num_rows = 0; + for (field, child) in fields.iter().zip(self.children) { + let child_data = child.into_arrow_impl(field.data_type().clone(), true)?; + num_rows = child_data.len(); + builder = builder.add_child_data(child_data); + } + + // Apply validity if present + let builder = if let Some(validity) = self.validity { + let null_count = validity.null_count(); + builder + .null_bit_buffer(Some(validity.into_inner().into_inner())) + .null_count(null_count) + } else { + builder.null_count(0) + }; + + let builder = builder.len(num_rows); + Ok(builder.build()?) + } else { + Err(Error::internal(format!( + "Expected Struct, got {:?}", + data_type + ))) + } + } + + fn remove_outer_validity(self) -> Self { + Self { + children: self + .children + .into_iter() + .map(|c| c.remove_outer_validity()) + .collect(), + block_info: self.block_info, + validity: None, // Remove the validity + } + } + + fn into_buffers(self) -> Vec { + self.children + .into_iter() + .flat_map(|c| c.into_buffers()) + .collect() + } + + pub fn has_variable_width_child(&self) -> bool { + self.children + .iter() + .any(|child| !matches!(child, DataBlock::FixedWidth(_))) + } + + pub fn data_size(&self) -> u64 { + self.children + .iter() + .map(|data_block| data_block.data_size()) + .sum() + } +} + +/// A data block for dictionary encoded data +#[derive(Debug, Clone)] +pub struct DictionaryDataBlock { + /// The indices buffer + pub indices: FixedWidthDataBlock, + /// The dictionary itself + pub dictionary: Box, +} + +impl DictionaryDataBlock { + fn decode_helper(self) -> Result { + // Handle empty batch - this can happen when decoding a range that contains + // only empty/null lists, or when reading sparse data + if self.indices.num_values == 0 { + return Ok(DataBlock::AllNull(AllNullDataBlock { num_values: 0 })); + } + + // assume the indices are uniformly distributed. + let estimated_size_bytes = self.dictionary.data_size() + * (self.indices.num_values + self.dictionary.num_values() - 1) + / self.dictionary.num_values(); + let mut data_builder = DataBlockBuilder::with_capacity_estimate(estimated_size_bytes); + + let indices = self.indices.data.borrow_to_typed_slice::(); + let indices = indices.as_ref(); + + let selections = indices.iter().map(|idx| { + let idx = idx.to_usize().unwrap() as u64; + idx..idx + 1 + }); + data_builder.append_ranges(&self.dictionary, selections)?; + + Ok(data_builder.finish()) + } + + pub fn decode(self) -> Result { + match self.indices.bits_per_value { + 8 => self.decode_helper::(), + 16 => self.decode_helper::(), + 32 => self.decode_helper::(), + 64 => self.decode_helper::(), + _ => Err(lance_core::Error::internal(format!( + "Unsupported dictionary index bit width: {} bits", + self.indices.bits_per_value + ))), + } + } + + fn into_arrow_dict( + self, + key_type: Box, + value_type: Box, + _validate: bool, + ) -> Result { + let declared_key_bits = key_type.byte_width() as u64 * 8; + if self.indices.bits_per_value != declared_key_bits { + return Err(lance_core::Error::corrupt_file_named( + "dictionary", + format!( + "dictionary indices use {} bits but the declared {} key type uses {} bits", + self.indices.bits_per_value, key_type, declared_key_bits + ), + )); + } + let indices_num_values = self.indices.num_values; + let indices = self + .indices + .do_into_arrow((*key_type).clone(), indices_num_values, true)?; + let dictionary = self + .dictionary + .into_arrow_impl((*value_type).clone(), true)?; + + let builder = indices + .into_builder() + .add_child_data(dictionary) + .data_type(DataType::Dictionary(key_type, value_type)); + + Ok(builder.build()?) + } + + fn into_arrow(self, data_type: DataType, validate: bool) -> Result { + if let DataType::Dictionary(key_type, value_type) = data_type { + self.into_arrow_dict(key_type, value_type, validate) + } else { + self.decode()?.into_arrow_impl(data_type, validate) + } + } + + fn into_buffers(self) -> Vec { + let mut buffers = self.indices.into_buffers(); + buffers.extend(self.dictionary.into_buffers()); + buffers + } + + pub fn into_parts(self) -> (DataBlock, DataBlock) { + (DataBlock::FixedWidth(self.indices), *self.dictionary) + } + + pub fn from_parts(indices: FixedWidthDataBlock, dictionary: DataBlock) -> Self { + Self { + indices, + dictionary: Box::new(dictionary), + } + } +} + +/// A DataBlock is a collection of buffers that represents an "array" of data in very generic terms +/// +/// The output of each decoder is a DataBlock. Decoders can be chained together to transform +/// one DataBlock into a different kind of DataBlock. +/// +/// The DataBlock is somewhere in between Arrow's ArrayData and Array and represents a physical +/// layout of the data. +/// +/// A DataBlock can be converted into an Arrow ArrayData (and then Array) for a given array type. +/// For example, a FixedWidthDataBlock can be converted into any primitive type or a fixed size +/// list of a primitive type. This is a zero-copy operation. +/// +/// In addition, a DataBlock can be created from an Arrow array or arrays. This is not a zero-copy +/// operation as some normalization may be required. +#[derive(Debug, Clone)] +pub enum DataBlock { + Empty(), + Constant(ConstantDataBlock), + AllNull(AllNullDataBlock), + Nullable(NullableDataBlock), + FixedWidth(FixedWidthDataBlock), + FixedSizeList(FixedSizeListBlock), + VariableWidth(VariableWidthBlock), + Opaque(OpaqueBlock), + Struct(StructDataBlock), + Dictionary(DictionaryDataBlock), +} + +impl DataBlock { + /// Convert self into an Arrow ArrayData with full layout validation. + /// + /// The `validate` argument is retained for API compatibility. Conversion is + /// always validated because callers can construct data blocks directly. + pub fn into_arrow(self, data_type: DataType, _validate: bool) -> Result { + self.into_arrow_impl(data_type, true) + } + + fn into_arrow_impl(self, data_type: DataType, validate: bool) -> Result { + match self { + Self::Empty() => Ok(new_empty_array(&data_type).to_data()), + Self::Constant(inner) => inner.into_arrow(data_type, validate), + Self::AllNull(inner) => inner.into_arrow(data_type, validate), + Self::Nullable(inner) => inner.into_arrow(data_type, validate), + Self::FixedWidth(inner) => inner.into_arrow(data_type, validate), + Self::FixedSizeList(inner) => inner.into_arrow(data_type, validate), + Self::VariableWidth(inner) => inner.into_arrow(data_type, validate), + Self::Struct(inner) => inner.into_arrow(data_type, validate), + Self::Dictionary(inner) => inner.into_arrow(data_type, validate), + Self::Opaque(_) => Err(Error::internal( + "Cannot convert OpaqueBlock to Arrow".to_string(), + )), + } + } + + /// Convert the data block into a collection of buffers for serialization + /// + /// The order matters and will be used to reconstruct the data block at read time. + pub fn into_buffers(self) -> Vec { + match self { + Self::Empty() => Vec::default(), + Self::Constant(inner) => inner.into_buffers(), + Self::AllNull(inner) => inner.into_buffers(), + Self::Nullable(inner) => inner.into_buffers(), + Self::FixedWidth(inner) => inner.into_buffers(), + Self::FixedSizeList(inner) => inner.into_buffers(), + Self::VariableWidth(inner) => inner.into_buffers(), + Self::Struct(inner) => inner.into_buffers(), + Self::Dictionary(inner) => inner.into_buffers(), + Self::Opaque(inner) => inner.buffers, + } + } + + /// Converts the data buffers into borrowed mode and clones the block + /// + /// This is a zero-copy operation but requires a mutable reference to self and, afterwards, + /// all buffers will be in Borrowed mode. + /// Try and clone the block + /// + /// This will fail if any buffers are in owned mode. You can call borrow_and_clone() to + /// ensure that all buffers are in borrowed mode before calling this method. + pub fn try_clone(&self) -> Result { + match self { + Self::Empty() => Ok(Self::Empty()), + Self::Constant(inner) => Ok(Self::Constant(inner.clone())), + Self::AllNull(inner) => Ok(Self::AllNull(inner.clone())), + Self::Nullable(inner) => Ok(Self::Nullable(inner.clone())), + Self::FixedWidth(inner) => Ok(Self::FixedWidth(inner.clone())), + Self::FixedSizeList(inner) => Ok(Self::FixedSizeList(inner.clone())), + Self::VariableWidth(inner) => Ok(Self::VariableWidth(inner.clone())), + Self::Struct(inner) => Ok(Self::Struct(inner.clone())), + Self::Dictionary(inner) => Ok(Self::Dictionary(inner.clone())), + Self::Opaque(inner) => Ok(Self::Opaque(inner.clone())), + } + } + + pub fn name(&self) -> &'static str { + match self { + Self::Constant(_) => "Constant", + Self::Empty() => "Empty", + Self::AllNull(_) => "AllNull", + Self::Nullable(_) => "Nullable", + Self::FixedWidth(_) => "FixedWidth", + Self::FixedSizeList(_) => "FixedSizeList", + Self::VariableWidth(_) => "VariableWidth", + Self::Struct(_) => "Struct", + Self::Dictionary(_) => "Dictionary", + Self::Opaque(_) => "Opaque", + } + } + + pub fn is_variable(&self) -> bool { + match self { + Self::Constant(_) => false, + Self::Empty() => false, + Self::AllNull(_) => false, + Self::Nullable(nullable) => nullable.data.is_variable(), + Self::FixedWidth(_) => false, + Self::FixedSizeList(fsl) => fsl.child.is_variable(), + Self::VariableWidth(_) => true, + Self::Struct(strct) => strct.children.iter().any(|c| c.is_variable()), + Self::Dictionary(_) => { + todo!("is_variable for DictionaryDataBlock is not implemented yet") + } + Self::Opaque(_) => panic!("Does not make sense to ask if an Opaque block is variable"), + } + } + + pub fn is_nullable(&self) -> bool { + match self { + Self::AllNull(_) => true, + Self::Nullable(_) => true, + Self::FixedSizeList(fsl) => fsl.child.is_nullable(), + Self::Struct(strct) => strct.children.iter().any(|c| c.is_nullable()), + Self::Dictionary(_) => { + todo!("is_nullable for DictionaryDataBlock is not implemented yet") + } + Self::Opaque(_) => panic!("Does not make sense to ask if an Opaque block is nullable"), + _ => false, + } + } + + /// The number of values in the block + /// + /// This function does not recurse into child blocks. If this is a FSL then it will + /// be the number of lists and not the number of items. + pub fn num_values(&self) -> u64 { + match self { + Self::Empty() => 0, + Self::Constant(inner) => inner.num_values, + Self::AllNull(inner) => inner.num_values, + Self::Nullable(inner) => inner.data.num_values(), + Self::FixedWidth(inner) => inner.num_values, + Self::FixedSizeList(inner) => inner.num_values(), + Self::VariableWidth(inner) => inner.num_values, + Self::Struct(inner) => inner.children[0].num_values(), + Self::Dictionary(inner) => inner.indices.num_values, + Self::Opaque(inner) => inner.num_values, + } + } + + /// The number of items in a single row + /// + /// This is always 1 unless there are layers of FSL + pub fn items_per_row(&self) -> u64 { + match self { + Self::Empty() => todo!(), // Leave undefined until needed + Self::Constant(_) => todo!(), // Leave undefined until needed + Self::AllNull(_) => todo!(), // Leave undefined until needed + Self::Nullable(nullable) => nullable.data.items_per_row(), + Self::FixedWidth(_) => 1, + Self::FixedSizeList(fsl) => fsl.dimension * fsl.child.items_per_row(), + Self::VariableWidth(_) => 1, + Self::Struct(_) => todo!(), // Leave undefined until needed + Self::Dictionary(_) => 1, + Self::Opaque(_) => 1, + } + } + + /// The number of bytes in the data block (including any child blocks) + pub fn data_size(&self) -> u64 { + match self { + Self::Empty() => 0, + Self::Constant(inner) => inner.data_size(), + Self::AllNull(_) => 0, + Self::Nullable(inner) => inner.data_size(), + Self::FixedWidth(inner) => inner.data_size(), + Self::FixedSizeList(inner) => inner.data_size(), + Self::VariableWidth(inner) => inner.data_size(), + Self::Struct(inner) => inner.children.iter().map(|child| child.data_size()).sum(), + Self::Dictionary(inner) => inner.indices.data_size() + inner.dictionary.data_size(), + Self::Opaque(inner) => inner.data_size(), + } + } + + /// Removes any validity information from the block + /// + /// This does not filter the block (e.g. remove rows). It only removes + /// the validity bitmaps (if present). Any garbage masked by null bits + /// will now appear as proper values. + /// + /// If `recurse` is true, then this will also remove validity from any child blocks. + pub fn remove_outer_validity(self) -> Self { + match self { + Self::AllNull(_) => panic!("Cannot remove validity on all-null data"), + Self::Nullable(inner) => *inner.data, + Self::Struct(inner) => Self::Struct(inner.remove_outer_validity()), + other => other, + } + } + + fn make_builder(&self, estimated_size_bytes: u64) -> Box { + match self { + Self::FixedWidth(inner) => { + if inner.bits_per_value == 1 { + Box::new(BitmapDataBlockBuilder::new(estimated_size_bytes)) + } else { + Box::new(FixedWidthDataBlockBuilder::new( + inner.bits_per_value, + estimated_size_bytes, + )) + } + } + Self::VariableWidth(inner) => { + if inner.bits_per_offset == 32 { + Box::new(VariableWidthDataBlockBuilder::::new( + estimated_size_bytes, + )) + } else if inner.bits_per_offset == 64 { + Box::new(VariableWidthDataBlockBuilder::::new( + estimated_size_bytes, + )) + } else { + todo!() + } + } + Self::FixedSizeList(inner) => { + let inner_builder = inner.child.make_builder(estimated_size_bytes); + Box::new(FixedSizeListBlockBuilder::new( + inner_builder, + inner.dimension, + )) + } + Self::Nullable(nullable) => { + // There's no easy way to know what percentage of the data is in the valiidty buffer + // but 1/16th seems like a reasonable guess. + let estimated_validity_size_bytes = estimated_size_bytes / 16; + let inner_builder = nullable + .data + .make_builder(estimated_size_bytes - estimated_validity_size_bytes); + Box::new(NullableDataBlockBuilder::new( + inner_builder, + estimated_validity_size_bytes as usize, + )) + } + Self::Struct(struct_data_block) => { + let num_children = struct_data_block.children.len(); + let per_child_estimate = if num_children == 0 { + 0 + } else { + estimated_size_bytes / num_children as u64 + }; + let child_builders = struct_data_block + .children + .iter() + .map(|child| child.make_builder(per_child_estimate)) + .collect(); + Box::new(StructDataBlockBuilder::new(child_builders)) + } + Self::AllNull(_) => Box::new(AllNullDataBlockBuilder::default()), + _ => todo!("make_builder for {:?}", self), + } + } +} + +macro_rules! as_type { + ($fn_name:ident, $inner:tt, $inner_type:ident) => { + pub fn $fn_name(self) -> Option<$inner_type> { + match self { + Self::$inner(inner) => Some(inner), + _ => None, + } + } + }; +} + +macro_rules! as_type_ref { + ($fn_name:ident, $inner:tt, $inner_type:ident) => { + pub fn $fn_name(&self) -> Option<&$inner_type> { + match self { + Self::$inner(inner) => Some(inner), + _ => None, + } + } + }; +} + +macro_rules! as_type_ref_mut { + ($fn_name:ident, $inner:tt, $inner_type:ident) => { + pub fn $fn_name(&mut self) -> Option<&mut $inner_type> { + match self { + Self::$inner(inner) => Some(inner), + _ => None, + } + } + }; +} + +// Cast implementations +impl DataBlock { + as_type!(as_all_null, AllNull, AllNullDataBlock); + as_type!(as_nullable, Nullable, NullableDataBlock); + as_type!(as_fixed_width, FixedWidth, FixedWidthDataBlock); + as_type!(as_fixed_size_list, FixedSizeList, FixedSizeListBlock); + as_type!(as_variable_width, VariableWidth, VariableWidthBlock); + as_type!(as_struct, Struct, StructDataBlock); + as_type!(as_dictionary, Dictionary, DictionaryDataBlock); + as_type_ref!(as_all_null_ref, AllNull, AllNullDataBlock); + as_type_ref!(as_nullable_ref, Nullable, NullableDataBlock); + as_type_ref!(as_fixed_width_ref, FixedWidth, FixedWidthDataBlock); + as_type_ref!(as_fixed_size_list_ref, FixedSizeList, FixedSizeListBlock); + as_type_ref!(as_variable_width_ref, VariableWidth, VariableWidthBlock); + as_type_ref!(as_struct_ref, Struct, StructDataBlock); + as_type_ref!(as_dictionary_ref, Dictionary, DictionaryDataBlock); + as_type_ref_mut!(as_all_null_ref_mut, AllNull, AllNullDataBlock); + as_type_ref_mut!(as_nullable_ref_mut, Nullable, NullableDataBlock); + as_type_ref_mut!(as_fixed_width_ref_mut, FixedWidth, FixedWidthDataBlock); + as_type_ref_mut!( + as_fixed_size_list_ref_mut, + FixedSizeList, + FixedSizeListBlock + ); + as_type_ref_mut!(as_variable_width_ref_mut, VariableWidth, VariableWidthBlock); + as_type_ref_mut!(as_struct_ref_mut, Struct, StructDataBlock); + as_type_ref_mut!(as_dictionary_ref_mut, Dictionary, DictionaryDataBlock); +} + +// Methods to convert from Arrow -> DataBlock + +fn get_byte_range(offsets: &mut LanceBuffer) -> Range { + let offsets = offsets.borrow_to_typed_slice::(); + if offsets.as_ref().is_empty() { + 0..0 + } else { + offsets.as_ref().first().unwrap().as_usize()..offsets.as_ref().last().unwrap().as_usize() + } +} + +// Given multiple offsets arrays [0, 5, 10], [0, 3, 7], etc. stitch +// them together to get [0, 5, 10, 13, 20, ...] +// +// Also returns the data range referenced by each offset array (may +// not be 0..len if there is slicing involved) +fn stitch_offsets + std::ops::Sub>( + offsets: Vec, +) -> (LanceBuffer, Vec>) { + if offsets.is_empty() { + return (LanceBuffer::empty(), Vec::default()); + } + let len = offsets.iter().map(|b| b.len()).sum::(); + // Note: we are making a copy here, even if there is only one input, because we want to + // normalize that input if it doesn't start with zero. This could be micro-optimized out + // if needed. + let mut dest = Vec::with_capacity(len); + let mut byte_ranges = Vec::with_capacity(offsets.len()); + + // We insert one leading 0 before processing any of the inputs + dest.push(T::from_usize(0).unwrap()); + + for mut o in offsets.into_iter() { + if !o.is_empty() { + let last_offset = *dest.last().unwrap(); + let o = o.borrow_to_typed_slice::(); + let start = *o.as_ref().first().unwrap(); + // First, we skip the first offset + // Then, we subtract that first offset from each remaining offset + // + // This gives us a 0-based offset array (minus the leading 0) + // + // Then we add the last offset from the previous array to each offset + // which shifts our offset array to the correct position + // + // For example, let's assume the last offset from the previous array + // was 10 and we are given [13, 17, 22]. This means we have two values with + // length 4 (17 - 13) and 5 (22 - 17). The output from this step will be + // [14, 19]. Combined with our last offset of 10, this gives us [10, 14, 19] + // which is our same two values of length 4 and 5. + dest.extend(o.as_ref()[1..].iter().map(|&x| x + last_offset - start)); + } + byte_ranges.push(get_byte_range::(&mut o)); + } + (LanceBuffer::reinterpret_vec(dest), byte_ranges) +} + +fn arrow_binary_to_data_block( + arrays: &[ArrayRef], + num_values: u64, + bits_per_offset: u8, +) -> DataBlock { + let data_vec = arrays.iter().map(|arr| arr.to_data()).collect::>(); + arrow_binary_array_data_to_data_block(&data_vec, num_values, bits_per_offset) +} + +fn arrow_binary_array_data_to_data_block( + data_vec: &[ArrayData], + num_values: u64, + bits_per_offset: u8, +) -> DataBlock { + let bytes_per_offset = bits_per_offset as usize / 8; + let offsets = data_vec + .iter() + .map(|d| { + LanceBuffer::from(d.buffers()[0].slice_with_length( + d.offset() * bytes_per_offset, + (d.len() + 1) * bytes_per_offset, + )) + }) + .collect::>(); + let (offsets, data_ranges) = if bits_per_offset == 32 { + stitch_offsets::(offsets) + } else { + stitch_offsets::(offsets) + }; + let data = data_vec + .iter() + .zip(data_ranges) + .map(|(d, byte_range)| { + LanceBuffer::from( + d.buffers()[1] + .slice_with_length(byte_range.start, byte_range.end - byte_range.start), + ) + }) + .collect::>(); + let data = LanceBuffer::concat_into_one(data); + DataBlock::VariableWidth(VariableWidthBlock { + data, + offsets, + bits_per_offset, + num_values, + block_info: BlockInfo::new(), + }) +} + +fn encode_flat_data(arrays: &[ArrayRef], num_values: u64) -> LanceBuffer { + let bytes_per_value = arrays[0].data_type().byte_width(); + let mut buffer = Vec::with_capacity(num_values as usize * bytes_per_value); + for arr in arrays { + let data = arr.to_data(); + buffer.extend_from_slice(data.buffers()[0].as_slice()); + } + LanceBuffer::from(buffer) +} + +fn do_encode_bitmap_data(bitmaps: &[BooleanBuffer], num_values: u64) -> LanceBuffer { + let mut builder = BooleanBufferBuilder::new(num_values as usize); + + for buf in bitmaps { + builder.append_buffer(buf); + } + + let buffer = builder.finish().into_inner(); + LanceBuffer::from(buffer) +} + +fn encode_bitmap_data(arrays: &[ArrayRef], num_values: u64) -> LanceBuffer { + let bitmaps = arrays + .iter() + .map(|arr| arr.as_boolean().values().clone()) + .collect::>(); + do_encode_bitmap_data(&bitmaps, num_values) +} + +// Concatenate dictionary arrays. This is a bit tricky because we might overflow the +// index type. If we do, we need to upscale the indices to a larger type. +fn concat_dict_arrays(arrays: &[ArrayRef]) -> ArrayRef { + let value_type = arrays[0].as_any_dictionary().values().data_type(); + let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::>(); + match arrow_select::concat::concat(&array_refs) { + Ok(array) => array, + Err(arrow_schema::ArrowError::DictionaryKeyOverflowError) => { + // Slow, but hopefully a corner case. Optimize later + let upscaled = array_refs + .iter() + .map(|arr| { + match arrow_cast::cast( + *arr, + &DataType::Dictionary( + Box::new(DataType::UInt32), + Box::new(value_type.clone()), + ), + ) { + Ok(arr) => arr, + Err(arrow_schema::ArrowError::DictionaryKeyOverflowError) => { + // Technically I think this means the input type was u64 already + unimplemented!("Dictionary arrays with more than 2^32 unique values") + } + err => err.unwrap(), + } + }) + .collect::>(); + let array_refs = upscaled.iter().map(|arr| arr.as_ref()).collect::>(); + // Can still fail if concat pushes over u32 boundary + match arrow_select::concat::concat(&array_refs) { + Ok(array) => array, + Err(arrow_schema::ArrowError::DictionaryKeyOverflowError) => { + unimplemented!("Dictionary arrays with more than 2^32 unique values") + } + err => err.unwrap(), + } + } + // Shouldn't be any other possible errors in concat + err => err.unwrap(), + } +} + +fn max_index_val(index_type: &DataType) -> u64 { + match index_type { + DataType::Int8 => i8::MAX as u64, + DataType::Int16 => i16::MAX as u64, + DataType::Int32 => i32::MAX as u64, + DataType::Int64 => i64::MAX as u64, + DataType::UInt8 => u8::MAX as u64, + DataType::UInt16 => u16::MAX as u64, + DataType::UInt32 => u32::MAX as u64, + DataType::UInt64 => u64::MAX, + _ => panic!("Invalid dictionary index type"), + } +} + +// If we get multiple dictionary arrays and they don't all have the same dictionary +// then we need to normalize the indices. Otherwise we might have something like: +// +// First chunk ["hello", "foo"], [0, 0, 1, 1, 1] +// Second chunk ["bar", "world"], [0, 1, 0, 1, 1] +// +// If we simply encode as ["hello", "foo", "bar", "world"], [0, 0, 1, 1, 1, 0, 1, 0, 1, 1] +// then we will get the wrong answer because the dictionaries were not merged and the indices +// were not remapped. +// +// A simple way to do this today is to just concatenate all the arrays. This is because +// arrow's dictionary concatenation function already has the logic to merge dictionaries. +// +// TODO: We could be more efficient here by checking if the dictionaries are the same +// Also, if they aren't, we can possibly do something cheaper than concatenating +// +// In addition, we want to normalize the representation of nulls. The cheapest thing to +// do (space-wise) is to put the nulls in the dictionary. +fn arrow_dictionary_to_data_block(arrays: &[ArrayRef], validity: Option) -> DataBlock { + let array = concat_dict_arrays(arrays); + let array_dict = array.as_any_dictionary(); + let mut indices = array_dict.keys(); + let num_values = indices.len() as u64; + let mut values = array_dict.values().clone(); + + // TODO: Should we just always normalize indices to u32? That would make logic simpler + // and we're going to bitpack them soon anyways + + let indices_block = if let Some(validity) = validity { + // If there is validity then we find the first invalid index in the dictionary values, inserting + // a new value if we need to. Then we change all indices to point to that value. This way we + // never need to store nullability of the indices. + let mut first_invalid_index = None; + if let Some(values_validity) = values.nulls() { + first_invalid_index = (!values_validity.inner()).set_indices().next(); + } + let first_invalid_index = first_invalid_index.unwrap_or_else(|| { + let null_arr = new_null_array(values.data_type(), 1); + values = arrow_select::concat::concat(&[values.as_ref(), null_arr.as_ref()]).unwrap(); + values.len() - 1 + }); + let max_index_val = max_index_val(indices.data_type()); + let upcast = if first_invalid_index as u64 > max_index_val { + // Widen the index type when the null dictionary value cannot be addressed by the + // declared key type, whether the value already existed or was appended above. + if max_index_val >= u32::MAX as u64 { + unimplemented!("Dictionary arrays with 2^32 unique value (or more) and a null") + } + Some(arrow_cast::cast(indices, &DataType::UInt32).unwrap()) + } else { + None + }; + if let Some(upcast) = upcast.as_ref() { + indices = upcast; + } + // This can't fail since we already checked for fit + let null_index_arr = arrow_cast::cast( + &UInt64Array::from(vec![first_invalid_index as u64]), + indices.data_type(), + ) + .unwrap(); + + let bytes_per_index = indices.data_type().byte_width(); + let bits_per_index = bytes_per_index as u64 * 8; + + let null_index_arr = null_index_arr.into_data(); + let null_index_bytes = &null_index_arr.buffers()[0]; + // Need to make a copy here since indices isn't mutable, could be avoided in theory + let mut indices_bytes = indices.to_data().buffers()[0].to_vec(); + for invalid_idx in (!validity.inner()).set_indices() { + indices_bytes[invalid_idx * bytes_per_index..(invalid_idx + 1) * bytes_per_index] + .copy_from_slice(null_index_bytes.as_slice()); + } + FixedWidthDataBlock { + data: LanceBuffer::from(indices_bytes), + bits_per_value: bits_per_index, + num_values, + block_info: BlockInfo::new(), + } + } else { + FixedWidthDataBlock { + data: LanceBuffer::from(indices.to_data().buffers()[0].clone()), + bits_per_value: indices.data_type().byte_width() as u64 * 8, + num_values, + block_info: BlockInfo::new(), + } + }; + + let items = DataBlock::from(values); + DataBlock::Dictionary(DictionaryDataBlock { + indices: indices_block, + dictionary: Box::new(items), + }) +} + +enum Nullability { + None, + All, + Some(NullBuffer), +} + +impl Nullability { + fn to_option(&self) -> Option { + match self { + Self::Some(nulls) => Some(nulls.clone()), + _ => None, + } + } +} + +fn extract_nulls(arrays: &[ArrayRef], num_values: u64) -> Nullability { + let mut has_nulls = false; + let nulls_and_lens = arrays + .iter() + .map(|arr| { + let nulls = arr.logical_nulls(); + has_nulls |= nulls.is_some(); + (nulls, arr.len()) + }) + .collect::>(); + if !has_nulls { + return Nullability::None; + } + let mut builder = BooleanBufferBuilder::new(num_values as usize); + let mut num_nulls = 0; + for (null, len) in nulls_and_lens { + if let Some(null) = null { + num_nulls += null.null_count(); + builder.append_buffer(&null.into_inner()); + } else { + builder.append_n(len, true); + } + } + if num_nulls == num_values as usize { + Nullability::All + } else { + Nullability::Some(NullBuffer::new(builder.finish())) + } +} + +impl DataBlock { + fn validate_variable_width_offsets( + array_data: &ArrayData, + ) -> std::result::Result<(), String> { + if array_data.is_empty() && array_data.buffers()[0].is_empty() { + return Ok(()); + } + let offset_size = std::mem::size_of::(); + let offset_start = array_data.offset() * offset_size; + let offset_len = (array_data.len() + 1) * offset_size; + let offset_buffer = + LanceBuffer::from(array_data.buffers()[0].slice_with_length(offset_start, offset_len)); + let offsets = offset_buffer.borrow_to_typed_slice::(); + VariableWidthBlock::offset_violation_detail( + offsets.as_ref(), + array_data.buffers()[1].len(), + 0, + ) + .map_or(Ok(()), Err) + } + + // `validate_full` also rescans UTF-8 contents and character boundaries on every flush. + // Encoding only needs a complete monotonicity and bounds proof before slicing offsets. + fn validate_variable_width_layouts(array_data: &ArrayData) -> std::result::Result<(), String> { + match array_data.data_type() { + DataType::Binary | DataType::Utf8 => { + Self::validate_variable_width_offsets::(array_data)?; + } + DataType::LargeBinary | DataType::LargeUtf8 => { + Self::validate_variable_width_offsets::(array_data)?; + } + _ => {} + } + for child_data in array_data.child_data() { + Self::validate_variable_width_layouts(child_data)?; + } + Ok(()) + } + + fn validate_array_data( + array_data: &ArrayData, + field_name: &str, + array_index: usize, + ) -> Result<()> { + let validation = array_data + .validate() + .map_err(|error| error.to_string()) + .and_then(|_| Self::validate_variable_width_layouts(array_data)); + validation.map_err(|error| { + Error::invalid_input_source( + format!( + "Invalid Arrow array for field '{}' at buffered array {}: {}", + field_name, array_index, error + ) + .into(), + ) + }) + } + + pub(crate) fn validate_arrays(arrays: &[ArrayRef], field_name: &str) -> Result<()> { + for (array_index, array) in arrays.iter().enumerate() { + Self::validate_array_data(&array.to_data(), field_name, array_index)?; + } + Ok(()) + } + + pub fn from_arrays(arrays: &[ArrayRef], num_values: u64) -> Self { + if arrays.is_empty() || num_values == 0 { + return Self::AllNull(AllNullDataBlock { num_values: 0 }); + } + + let data_type = arrays[0].data_type(); + let nulls = extract_nulls(arrays, num_values); + + if let Nullability::All = nulls { + return Self::AllNull(AllNullDataBlock { num_values }); + } + + let mut encoded = match data_type { + DataType::Binary | DataType::Utf8 => arrow_binary_to_data_block(arrays, num_values, 32), + // View types have no Lance disk representation; cast to the classic offset layout. + DataType::Utf8View => { + let casted: Vec = arrays + .iter() + .map(|a| { + arrow_cast::cast(a.as_ref(), &DataType::Utf8) + .expect("Utf8View to Utf8 cast is always valid") + }) + .collect(); + arrow_binary_to_data_block(&casted, num_values, 32) + } + DataType::BinaryView => { + let casted: Vec = arrays + .iter() + .map(|a| { + arrow_cast::cast(a.as_ref(), &DataType::Binary) + .expect("BinaryView to Binary cast is always valid") + }) + .collect(); + arrow_binary_to_data_block(&casted, num_values, 32) + } + DataType::LargeBinary | DataType::LargeUtf8 => { + arrow_binary_to_data_block(arrays, num_values, 64) + } + DataType::Boolean => { + let data = encode_bitmap_data(arrays, num_values); + Self::FixedWidth(FixedWidthDataBlock { + data, + bits_per_value: 1, + num_values, + block_info: BlockInfo::new(), + }) + } + DataType::Date32 + | DataType::Date64 + | DataType::Decimal32(_, _) + | DataType::Decimal64(_, _) + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Duration(_) + | DataType::FixedSizeBinary(_) + | DataType::Float16 + | DataType::Float32 + | DataType::Float64 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Int8 + | DataType::Interval(_) + | DataType::Time32(_) + | DataType::Time64(_) + | DataType::Timestamp(_, _) + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::UInt8 => { + let data = encode_flat_data(arrays, num_values); + Self::FixedWidth(FixedWidthDataBlock { + data, + bits_per_value: data_type.byte_width() as u64 * 8, + num_values, + block_info: BlockInfo::new(), + }) + } + DataType::Null => Self::AllNull(AllNullDataBlock { num_values }), + DataType::Dictionary(_, _) => arrow_dictionary_to_data_block(arrays, nulls.to_option()), + DataType::Struct(fields) => { + let structs = arrays.iter().map(|arr| arr.as_struct()).collect::>(); + let mut children = Vec::with_capacity(fields.len()); + for child_idx in 0..fields.len() { + let child_vec = structs + .iter() + .map(|s| s.column(child_idx).clone()) + .collect::>(); + children.push(Self::from_arrays(&child_vec, num_values)); + } + + // Extract validity for the struct array + let validity = match &nulls { + Nullability::None => None, + Nullability::Some(null_buffer) => Some(null_buffer.clone()), + Nullability::All => unreachable!("Should have returned AllNull earlier"), + }; + + Self::Struct(StructDataBlock { + children, + block_info: BlockInfo::default(), + validity, + }) + } + DataType::FixedSizeList(_, dim) => { + let children = arrays + .iter() + .map(|arr| arr.as_fixed_size_list().values().clone()) + .collect::>(); + let child_block = Self::from_arrays(&children, num_values * *dim as u64); + Self::FixedSizeList(FixedSizeListBlock { + child: Box::new(child_block), + dimension: *dim as u64, + }) + } + DataType::LargeList(_) + | DataType::List(_) + | DataType::ListView(_) + | DataType::LargeListView(_) + | DataType::Map(_, _) + | DataType::RunEndEncoded(_, _) + | DataType::Union(_, _) => { + panic!( + "Field with data type {} cannot be converted to data block", + data_type + ) + } + }; + + // compute statistics + encoded.compute_stat(); + + if !matches!(data_type, DataType::Dictionary(_, _)) { + match nulls { + Nullability::None => encoded, + Nullability::Some(nulls) => Self::Nullable(NullableDataBlock { + data: Box::new(encoded), + nulls: LanceBuffer::from(nulls.into_inner().into_inner()), + block_info: BlockInfo::new(), + }), + _ => unreachable!(), + } + } else { + // Dictionaries already insert the nulls into the dictionary items + encoded + } + } + + pub fn from_array(array: T) -> Self { + let num_values = array.len(); + Self::from_arrays(&[Arc::new(array)], num_values as u64) + } +} + +impl From for DataBlock { + fn from(array: ArrayRef) -> Self { + let num_values = array.len() as u64; + Self::from_arrays(&[array], num_values) + } +} + +trait DataBlockBuilderImpl: std::fmt::Debug { + fn validate_append(&self, data_block: &DataBlock, selection: &Range) -> Result<()>; + + fn append_validated(&mut self, data_block: &DataBlock, selection: Range) -> Result<()>; + + fn append(&mut self, data_block: &DataBlock, selection: Range) -> Result<()> { + self.validate_append(data_block, &selection)?; + self.append_validated(data_block, selection) + } + + fn finish(self: Box) -> DataBlock; +} + +#[derive(Debug)] +pub struct DataBlockBuilder { + estimated_size_bytes: u64, + builder: Option>, +} + +impl DataBlockBuilder { + pub fn with_capacity_estimate(estimated_size_bytes: u64) -> Self { + Self { + estimated_size_bytes, + builder: None, + } + } + + fn get_builder(&mut self, block: &DataBlock) -> &mut dyn DataBlockBuilderImpl { + if self.builder.is_none() { + self.builder = Some(block.make_builder(self.estimated_size_bytes)); + } + self.builder.as_mut().unwrap().as_mut() + } + + pub fn append(&mut self, data_block: &DataBlock, selection: Range) -> Result<()> { + self.get_builder(data_block).append(data_block, selection) + } + + fn append_ranges( + &mut self, + data_block: &DataBlock, + selections: impl IntoIterator>, + ) -> Result<()> { + let full_selection = 0..data_block.num_values(); + let builder = self.get_builder(data_block); + builder.validate_append(data_block, &full_selection)?; + for selection in selections { + if selection.start > selection.end || selection.end > full_selection.end { + return Err(Error::corrupt_file_named( + "data block", + format!( + "cannot append selection {}..{} from a block with {} values", + selection.start, selection.end, full_selection.end + ), + )); + } + builder.append_validated(data_block, selection)?; + } + Ok(()) + } + + pub fn finish(self) -> DataBlock { + let builder = self.builder.expect("DataBlockBuilder didn't see any data"); + builder.finish() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{ + ArrayRef, BinaryArray, BinaryViewArray, DictionaryArray, Int8Array, LargeBinaryArray, + LargeStringArray, StringArray, StringViewArray, UInt8Array, UInt16Array, make_array, + new_null_array, + types::{Int8Type, Int32Type}, + }; + use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer}; + use arrow_data::ArrayData; + + use arrow_schema::{DataType, Field, Fields}; + use lance_core::Error; + use lance_datagen::{ArrayGeneratorExt, DEFAULT_SEED, RowCount, array}; + use rand::SeedableRng; + use rstest::rstest; + + use crate::buffer::LanceBuffer; + + use super::{ + AllNullDataBlock, BlockInfo, DataBlock, DataBlockBuilder, DictionaryDataBlock, + FixedWidthDataBlock, VariableWidthBlock, + }; + + use arrow_array::Array; + + #[test] + fn test_sliced_to_data_block() { + let ints = UInt16Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8]); + let ints = ints.slice(2, 4); + let data = DataBlock::from_array(ints); + + let fixed_data = data.as_fixed_width().unwrap(); + assert_eq!(fixed_data.num_values, 4); + assert_eq!(fixed_data.data.len(), 8); + + let nullable_ints = + UInt16Array::from(vec![Some(0), None, Some(2), None, Some(4), None, Some(6)]); + let nullable_ints = nullable_ints.slice(1, 3); + let data = DataBlock::from_array(nullable_ints); + + let nullable = data.as_nullable().unwrap(); + assert_eq!(nullable.nulls, LanceBuffer::from(vec![0b00000010])); + } + + #[test] + fn test_string_to_data_block() { + // Converting string arrays that contain nulls to DataBlock + let strings1 = StringArray::from(vec![Some("hello"), None, Some("world")]); + let strings2 = StringArray::from(vec![Some("a"), Some("b")]); + let strings3 = StringArray::from(vec![Option::<&'static str>::None, None]); + + let arrays = &[strings1, strings2, strings3] + .iter() + .map(|arr| Arc::new(arr.clone()) as ArrayRef) + .collect::>(); + + let block = DataBlock::from_arrays(arrays, 7); + + assert_eq!(block.num_values(), 7); + let block = block.as_nullable().unwrap(); + + assert_eq!(block.nulls, LanceBuffer::from(vec![0b00011101])); + + let data = block.data.as_variable_width().unwrap(); + assert_eq!( + data.offsets, + LanceBuffer::reinterpret_vec(vec![0, 5, 5, 10, 11, 12, 12, 12]) + ); + + assert_eq!(data.data, LanceBuffer::copy_slice(b"helloworldab")); + + // Converting string arrays that do not contain nulls to DataBlock + let strings1 = StringArray::from(vec![Some("a"), Some("bc")]); + let strings2 = StringArray::from(vec![Some("def")]); + + let arrays = &[strings1, strings2] + .iter() + .map(|arr| Arc::new(arr.clone()) as ArrayRef) + .collect::>(); + + let block = DataBlock::from_arrays(arrays, 3); + + assert_eq!(block.num_values(), 3); + // Should be no nullable wrapper + let data = block.as_variable_width().unwrap(); + assert_eq!(data.offsets, LanceBuffer::reinterpret_vec(vec![0, 1, 3, 6])); + assert_eq!(data.data, LanceBuffer::copy_slice(b"abcdef")); + } + + #[test] + fn test_string_view_to_data_block() { + let views1 = StringViewArray::from(vec![Some("hello"), None, Some("world")]); + let views2 = StringViewArray::from(vec![Some("a"), Some("b")]); + let views3 = StringViewArray::from(vec![Option::<&'static str>::None, None]); + + let arrays = &[views1, views2, views3] + .iter() + .map(|arr| Arc::new(arr.clone()) as ArrayRef) + .collect::>(); + + let block = DataBlock::from_arrays(arrays, 7); + + assert_eq!(block.num_values(), 7); + let block = block.as_nullable().unwrap(); + assert_eq!(block.nulls, LanceBuffer::from(vec![0b00011101])); + let data = block.data.as_variable_width().unwrap(); + assert_eq!( + data.offsets, + LanceBuffer::reinterpret_vec(vec![0, 5, 5, 10, 11, 12, 12, 12]) + ); + assert_eq!(data.data, LanceBuffer::copy_slice(b"helloworldab")); + + let views1 = StringViewArray::from(vec![Some("a"), Some("bc")]); + let views2 = StringViewArray::from(vec![Some("def")]); + + let arrays = &[views1, views2] + .iter() + .map(|arr| Arc::new(arr.clone()) as ArrayRef) + .collect::>(); + + let block = DataBlock::from_arrays(arrays, 3); + + assert_eq!(block.num_values(), 3); + let data = block.as_variable_width().unwrap(); + assert_eq!(data.offsets, LanceBuffer::reinterpret_vec(vec![0, 1, 3, 6])); + assert_eq!(data.data, LanceBuffer::copy_slice(b"abcdef")); + } + + #[test] + fn test_binary_view_to_data_block() { + let arr: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some(b"foo".as_slice()), + None, + Some(b"bar".as_slice()), + ])); + let block = DataBlock::from_arrays(&[arr], 3); + let block = block.as_nullable().unwrap(); + let data = block.data.as_variable_width().unwrap(); + assert_eq!(data.data, LanceBuffer::copy_slice(b"foobar")); + } + + #[test] + fn test_string_sliced() { + let check = |arr: Vec, expected_off: Vec, expected_data: &[u8]| { + let arrs = arr + .into_iter() + .map(|a| Arc::new(a) as ArrayRef) + .collect::>(); + let num_rows = arrs.iter().map(|a| a.len()).sum::() as u64; + let data = DataBlock::from_arrays(&arrs, num_rows); + + assert_eq!(data.num_values(), num_rows); + + let data = data.as_variable_width().unwrap(); + assert_eq!(data.offsets, LanceBuffer::reinterpret_vec(expected_off)); + assert_eq!(data.data, LanceBuffer::copy_slice(expected_data)); + }; + + let string = StringArray::from(vec![Some("hello"), Some("world")]); + check(vec![string.slice(1, 1)], vec![0, 5], b"world"); + check(vec![string.slice(0, 1)], vec![0, 5], b"hello"); + check( + vec![string.slice(0, 1), string.slice(1, 1)], + vec![0, 5, 10], + b"helloworld", + ); + + let string2 = StringArray::from(vec![Some("foo"), Some("bar")]); + check( + vec![string.slice(0, 1), string2.slice(0, 1)], + vec![0, 5, 8], + b"hellofoo", + ); + } + + #[rstest] + #[case::utf8( + DataType::Utf8, + Buffer::from_slice_ref([0_i32, 5, 10]), + 32, + LanceBuffer::reinterpret_vec(vec![0_i32, 5]) + )] + #[case::large_utf8( + DataType::LargeUtf8, + Buffer::from_slice_ref([0_i64, 5, 10]), + 64, + LanceBuffer::reinterpret_vec(vec![0_i64, 5]) + )] + fn test_variable_width_array_data_offset( + #[case] data_type: DataType, + #[case] offsets: Buffer, + #[case] bits_per_offset: u8, + #[case] expected_offsets: LanceBuffer, + ) { + let array_data = ArrayData::builder(data_type) + .len(1) + .offset(1) + .add_buffer(offsets) + .add_buffer(Buffer::from(b"helloworld")) + .build() + .unwrap(); + + DataBlock::validate_array_data(&array_data, "text", 0).unwrap(); + let data = super::arrow_binary_array_data_to_data_block(&[array_data], 1, bits_per_offset); + + let data = data.as_variable_width().unwrap(); + assert_eq!(data.offsets, expected_offsets); + assert_eq!(data.data, LanceBuffer::copy_slice(b"world")); + } + + #[rstest] + #[case::utf8(DataType::Utf8, Buffer::from_slice_ref([0_i32, -1]))] + #[case::large_utf8(DataType::LargeUtf8, Buffer::from_slice_ref([0_i64, -1]))] + fn test_invalid_string_offsets_rejected_before_encoding( + #[case] data_type: DataType, + #[case] offsets: Buffer, + ) { + let array_data = unsafe { + ArrayData::builder(data_type) + .len(1) + .add_buffer(offsets) + .add_buffer(Buffer::from(b"")) + .build_unchecked() + }; + + let error = DataBlock::validate_array_data(&array_data, "text", 0).unwrap_err(); + + assert!(matches!(error, Error::InvalidInput { .. })); + let message = error.to_string(); + assert!( + message.contains("field 'text'"), + "unexpected message: {message}" + ); + assert!( + message.contains("offset[1] (-1)"), + "unexpected message: {message}" + ); + } + + #[test] + fn test_large() { + let arr = LargeBinaryArray::from_vec(vec![b"hello", b"world"]); + let data = DataBlock::from_array(arr); + + assert_eq!(data.num_values(), 2); + let data = data.as_variable_width().unwrap(); + assert_eq!(data.bits_per_offset, 64); + assert_eq!(data.num_values, 2); + assert_eq!(data.data, LanceBuffer::copy_slice(b"helloworld")); + assert_eq!( + data.offsets, + LanceBuffer::reinterpret_vec(vec![0_u64, 5, 10]) + ); + } + + #[test] + fn test_dictionary_indices_normalized() { + let arr1 = DictionaryArray::::from_iter([Some("a"), Some("a"), Some("b")]); + let arr2 = DictionaryArray::::from_iter([Some("b"), Some("c")]); + + let data = DataBlock::from_arrays(&[Arc::new(arr1), Arc::new(arr2)], 5); + + assert_eq!(data.num_values(), 5); + let data = data.as_dictionary().unwrap(); + let indices = data.indices; + assert_eq!(indices.bits_per_value, 8); + assert_eq!(indices.num_values, 5); + assert_eq!( + indices.data, + // You might expect 0, 0, 1, 1, 2 but it seems that arrow's dictionary concat does + // not actually collapse dictionaries. This is an arrow problem however, and we don't + // need to fix it here. + LanceBuffer::reinterpret_vec::(vec![0, 0, 1, 2, 3]) + ); + + let items = data.dictionary.as_variable_width().unwrap(); + assert_eq!(items.bits_per_offset, 32); + assert_eq!(items.num_values, 4); + assert_eq!(items.data, LanceBuffer::copy_slice(b"abbc")); + assert_eq!( + items.offsets, + LanceBuffer::reinterpret_vec(vec![0, 1, 2, 3, 4],) + ); + } + + #[test] + fn test_dictionary_nulls() { + // Test both ways of encoding nulls + + // By default, nulls get encoded into the indices + let arr1 = DictionaryArray::::from_iter([None, Some("a"), Some("b")]); + let arr2 = DictionaryArray::::from_iter([Some("c"), None]); + + let data = DataBlock::from_arrays(&[Arc::new(arr1), Arc::new(arr2)], 5); + + let check_common = |data: DataBlock| { + assert_eq!(data.num_values(), 5); + let dict = data.as_dictionary().unwrap(); + + let nullable_items = dict.dictionary.as_nullable().unwrap(); + assert_eq!(nullable_items.nulls, LanceBuffer::from(vec![0b00000111])); + assert_eq!(nullable_items.data.num_values(), 4); + + let items = nullable_items.data.as_variable_width().unwrap(); + assert_eq!(items.bits_per_offset, 32); + assert_eq!(items.num_values, 4); + assert_eq!(items.data, LanceBuffer::copy_slice(b"abc")); + assert_eq!( + items.offsets, + LanceBuffer::reinterpret_vec(vec![0, 1, 2, 3, 3],) + ); + + let indices = dict.indices; + assert_eq!(indices.bits_per_value, 8); + assert_eq!(indices.num_values, 5); + assert_eq!( + indices.data, + LanceBuffer::reinterpret_vec::(vec![3, 0, 1, 2, 3]) + ); + }; + check_common(data); + + // However, we can manually create a dictionary where nulls are in the dictionary + let items = StringArray::from(vec![Some("a"), Some("b"), Some("c"), None]); + let indices = Int8Array::from(vec![Some(3), Some(0), Some(1), Some(2), Some(3)]); + let dict = DictionaryArray::new(indices, Arc::new(items)); + + let data = DataBlock::from_array(dict); + + check_common(data); + } + + #[test] + fn test_dictionary_cannot_add_null() { + // 256 unique strings + let items = StringArray::from( + (0..256) + .map(|i| Some(String::from_utf8(vec![0; i]).unwrap())) + .collect::>(), + ); + // 257 indices, covering the whole range, plus one null + let indices = UInt8Array::from( + (0..=256) + .map(|i| if i == 256 { None } else { Some(i as u8) }) + .collect::>(), + ); + // We want to normalize this by pushing nulls into the dictionary, but we cannot because + // the dictionary is too large for the index type + let dict = DictionaryArray::new(indices, Arc::new(items)); + let data = DataBlock::from_array(dict); + + assert_eq!(data.num_values(), 257); + + let dict = data.as_dictionary().unwrap(); + + assert_eq!(dict.indices.bits_per_value, 32); + assert_eq!( + dict.indices.data, + LanceBuffer::reinterpret_vec((0_u32..257).collect::>()) + ); + + let nullable_items = dict.dictionary.as_nullable().unwrap(); + let null_buffer = NullBuffer::new(BooleanBuffer::new( + nullable_items.nulls.into_buffer(), + 0, + 257, + )); + for i in 0..256 { + assert!(!null_buffer.is_null(i)); + } + assert!(null_buffer.is_null(256)); + + assert_eq!( + nullable_items.data.as_variable_width().unwrap().data.len(), + 32640 + ); + } + + #[test] + fn test_all_null() { + for data_type in [ + DataType::UInt32, + DataType::FixedSizeBinary(2), + DataType::List(Arc::new(Field::new("item", DataType::UInt32, true))), + DataType::Struct(Fields::from(vec![Field::new("a", DataType::UInt32, true)])), + ] { + let block = DataBlock::AllNull(AllNullDataBlock { num_values: 10 }); + let arr = block.into_arrow(data_type.clone(), true).unwrap(); + let arr = make_array(arr); + let expected = new_null_array(&data_type, 10); + assert_eq!(&arr, &expected); + } + } + + #[test] + fn test_dictionary_cannot_concatenate() { + // 256 unique strings + let items = StringArray::from( + (0..256) + .map(|i| Some(String::from_utf8(vec![0; i]).unwrap())) + .collect::>(), + ); + // 256 different unique strings + let other_items = StringArray::from( + (0..256) + .map(|i| Some(String::from_utf8(vec![1; i + 1]).unwrap())) + .collect::>(), + ); + let indices = UInt8Array::from_iter_values(0..=255); + let dict1 = DictionaryArray::new(indices.clone(), Arc::new(items)); + let dict2 = DictionaryArray::new(indices, Arc::new(other_items)); + let data = DataBlock::from_arrays(&[Arc::new(dict1), Arc::new(dict2)], 512); + assert_eq!(data.num_values(), 512); + + let dict = data.as_dictionary().unwrap(); + + assert_eq!(dict.indices.bits_per_value, 32); + assert_eq!( + dict.indices.data, + LanceBuffer::reinterpret_vec::((0..512).collect::>()) + ); + // What fun: 0 + 1 + .. + 255 + 1 + 2 + .. + 256 = 2^16 + assert_eq!( + dict.dictionary.as_variable_width().unwrap().data.len(), + 65536 + ); + } + + #[test] + fn test_data_size() { + let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0); + // test data_size() when input has no nulls + let mut genn = array::rand::().with_nulls(&[false, false, false]); + + let arr = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + assert!(block.data_size() == arr.get_buffer_memory_size() as u64); + + let arr = genn.generate(RowCount::from(400), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + assert!(block.data_size() == arr.get_buffer_memory_size() as u64); + + // test data_size() when input has nulls + let mut genn = array::rand::().with_nulls(&[false, true, false]); + let arr = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + + let array_data = arr.to_data(); + let total_buffer_size: usize = array_data.buffers().iter().map(|buffer| buffer.len()).sum(); + // the NullBuffer.len() returns the length in bits so we divide_round_up by 8 + let array_nulls_size_in_bytes = arr.nulls().unwrap().len().div_ceil(8); + assert!(block.data_size() == (total_buffer_size + array_nulls_size_in_bytes) as u64); + + let arr = genn.generate(RowCount::from(400), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + + let array_data = arr.to_data(); + let total_buffer_size: usize = array_data.buffers().iter().map(|buffer| buffer.len()).sum(); + let array_nulls_size_in_bytes = arr.nulls().unwrap().len().div_ceil(8); + assert!(block.data_size() == (total_buffer_size + array_nulls_size_in_bytes) as u64); + + let mut genn = array::rand::().with_nulls(&[true, true, false]); + let arr = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + + let array_data = arr.to_data(); + let total_buffer_size: usize = array_data.buffers().iter().map(|buffer| buffer.len()).sum(); + let array_nulls_size_in_bytes = arr.nulls().unwrap().len().div_ceil(8); + assert!(block.data_size() == (total_buffer_size + array_nulls_size_in_bytes) as u64); + + let arr = genn.generate(RowCount::from(400), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + + let array_data = arr.to_data(); + let total_buffer_size: usize = array_data.buffers().iter().map(|buffer| buffer.len()).sum(); + let array_nulls_size_in_bytes = arr.nulls().unwrap().len().div_ceil(8); + assert!(block.data_size() == (total_buffer_size + array_nulls_size_in_bytes) as u64); + + let mut genn = array::rand::().with_nulls(&[false, true, false]); + let arr1 = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let arr2 = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let arr3 = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_arrays(&[arr1.clone(), arr2.clone(), arr3.clone()], 9); + + let concatenated_array = arrow_select::concat::concat(&[ + &*Arc::new(arr1.clone()) as &dyn Array, + &*Arc::new(arr2.clone()) as &dyn Array, + &*Arc::new(arr3.clone()) as &dyn Array, + ]) + .unwrap(); + let total_buffer_size: usize = concatenated_array + .to_data() + .buffers() + .iter() + .map(|buffer| buffer.len()) + .sum(); + + let total_nulls_size_in_bytes = concatenated_array.nulls().unwrap().len().div_ceil(8); + assert!(block.data_size() == (total_buffer_size + total_nulls_size_in_bytes) as u64); + } + + #[test] + fn variable_width_rejects_out_of_bounds_offsets_without_optional_validation() { + let block = VariableWidthBlock { + data: LanceBuffer::copy_slice(b"alphabetagamma"), + offsets: LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + bits_per_offset: 32, + num_values: 3, + block_info: BlockInfo::new(), + }; + + let error = block + .into_arrow(DataType::Binary, false) + .expect_err("out-of-bounds offsets must be rejected"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile, got: {error}" + ); + let message = error.to_string(); + assert!( + message.contains("100000") && message.contains("data buffer size: 14 bytes"), + "error must report the offending offset and the data buffer size: {message}" + ); + } + + #[test] + fn public_fixed_width_conversion_always_validates_layout() { + let block = FixedWidthDataBlock { + data: LanceBuffer::from(vec![0_u8; 4]), + bits_per_value: 32, + num_values: 2, + block_info: BlockInfo::new(), + }; + + for validate in [false, true] { + block + .clone() + .into_arrow(DataType::Int32, validate) + .expect_err("a short values buffer must be rejected"); + DataBlock::FixedWidth(block.clone()) + .into_arrow(DataType::Int32, validate) + .expect_err("a short values buffer must be rejected"); + } + } + + #[rstest] + #[case::i32_decreasing( + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 2]), + 32, + 2, + "decreases" + )] + #[case::i64_decreasing( + LanceBuffer::reinterpret_vec(vec![0_i64, 5, 2]), + 64, + 2, + "decreases" + )] + #[case::i32_out_of_bounds( + LanceBuffer::reinterpret_vec(vec![0_i32, 6]), + 32, + 1, + "out of bounds" + )] + fn variable_width_builder_rejects_malformed_offsets( + #[case] offsets: LanceBuffer, + #[case] bits_per_offset: u8, + #[case] num_values: u64, + #[case] expected_message: &str, + ) { + let block = DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::copy_slice(b"abcde"), + offsets, + bits_per_offset, + num_values, + block_info: BlockInfo::new(), + }); + let mut builder = DataBlockBuilder::with_capacity_estimate(5); + + let error = builder + .append(&block, 0..num_values) + .expect_err("malformed offsets must fail concatenation"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile, got: {error}" + ); + assert!( + error.to_string().contains(expected_message), + "unexpected message: {error}" + ); + } + + #[rstest] + #[case::binary_i32_tail_out_of_bounds( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::utf8_i32_tail_out_of_bounds( + DataType::Utf8, + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::large_binary_i64_tail_out_of_bounds( + DataType::LargeBinary, + LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 100_000]), + 64, + 3, + b"alphabetagamma".as_slice() + )] + #[case::large_utf8_i64_tail_out_of_bounds( + DataType::LargeUtf8, + LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 100_000]), + 64, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_negative_offset( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, -1, 9, 14]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_non_monotonic_offsets( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 9, 5, 14]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_interior_offset_out_of_bounds( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 100_000, 100_000, 14]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_offsets_buffer_too_short( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::utf8_invalid_byte_sequence( + DataType::Utf8, + LanceBuffer::reinterpret_vec(vec![0_i32, 1, 2, 3]), + 32, + 3, + &[b'a', 0xFF, b'b'] + )] + #[case::utf8_offset_splits_multibyte_char( + DataType::Utf8, + LanceBuffer::reinterpret_vec(vec![0_i32, 1, 2]), + 32, + 2, + "é".as_bytes() + )] + #[case::large_utf8_invalid_byte_sequence( + DataType::LargeUtf8, + LanceBuffer::reinterpret_vec(vec![0_i64, 1, 2, 3]), + 64, + 3, + &[b'a', 0xFF, b'b'] + )] + fn variable_width_rejects_malformed_layout( + #[case] data_type: DataType, + #[case] offsets: LanceBuffer, + #[case] bits_per_offset: u8, + #[case] num_values: u64, + #[case] data: &[u8], + ) { + let block = VariableWidthBlock { + data: LanceBuffer::copy_slice(data), + offsets, + bits_per_offset, + num_values, + block_info: BlockInfo::new(), + }; + + // The malformed layout must be rejected regardless of the optional + // `validate` flag: the flag selects extra validation, not the memory + // safety proof required to construct an Arrow array. + for validate in [false, true] { + let error = DataBlock::VariableWidth(block.clone()) + .into_arrow(data_type.clone(), validate) + .expect_err("malformed variable-width layout must be rejected"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile with validate={validate}, got: {error}" + ); + } + } + + #[test] + fn dictionary_rejects_malformed_variable_width_values_without_optional_validation() { + let values = VariableWidthBlock { + data: LanceBuffer::copy_slice(b"alphabetagamma"), + offsets: LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + bits_per_offset: 32, + num_values: 3, + block_info: BlockInfo::new(), + }; + let dictionary = DataBlock::Dictionary(DictionaryDataBlock { + indices: FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(vec![0_i32, 1, 2]), + bits_per_value: 32, + num_values: 3, + block_info: BlockInfo::new(), + }, + dictionary: Box::new(DataBlock::VariableWidth(values)), + }); + + let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)); + let error = dictionary + .into_arrow(data_type, false) + .expect_err("dictionary with out-of-bounds value offsets must be rejected"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile, got: {error}" + ); + } + + #[test] + fn dictionary_rejects_indices_wider_than_declared_key_type() { + let dictionary = DataBlock::Dictionary(DictionaryDataBlock { + indices: FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(vec![0_u32, 1, 128]), + bits_per_value: 32, + num_values: 3, + block_info: BlockInfo::new(), + }, + dictionary: Box::new(DataBlock::from_array(StringArray::from(vec![ + Some("zero"), + Some("one"), + None, + ]))), + }); + + let data_type = DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)); + let error = dictionary + .into_arrow(data_type, false) + .expect_err("mismatched dictionary index widths must be rejected"); + + assert!(matches!(error, Error::CorruptFile { .. })); + assert!( + error.to_string().contains( + "dictionary indices use 32 bits but the declared Int8 key type uses 8 bits" + ) + ); + } + + #[rstest] + #[case::binary(Arc::new(BinaryArray::from_vec(vec![b"alpha", b"", b"gamma"])) as ArrayRef)] + #[case::large_binary( + Arc::new(LargeBinaryArray::from_vec(vec![b"alpha", b"", b"gamma"])) as ArrayRef + )] + #[case::utf8(Arc::new(StringArray::from(vec!["héllo", "", "world"])) as ArrayRef)] + #[case::large_utf8(Arc::new(LargeStringArray::from(vec!["héllo", "", "world"])) as ArrayRef)] + fn variable_width_valid_data_survives_mandatory_validation(#[case] array: ArrayRef) { + let block = DataBlock::from_array(array.clone()); + for validate in [false, true] { + let round_tripped = make_array( + block + .clone() + .into_arrow(array.data_type().clone(), validate) + .unwrap(), + ); + assert_eq!(&round_tripped, &array); + } + } +} diff --git a/vendor/lance-encoding/src/decoder.rs b/vendor/lance-encoding/src/decoder.rs new file mode 100644 index 000000000..8a4aebbd0 --- /dev/null +++ b/vendor/lance-encoding/src/decoder.rs @@ -0,0 +1,4203 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Utilities and traits for scheduling & decoding data +//! +//! Reading data involves two steps: scheduling and decoding. The +//! scheduling step is responsible for figuring out what data is needed +//! and issuing the appropriate I/O requests. The decoding step is +//! responsible for taking the loaded data and turning it into Arrow +//! arrays. +//! +//! # Scheduling +//! +//! Scheduling is split into `FieldScheduler` and `PageScheduler`. +//! There is one field scheduler for each output field, which may map to many +//! columns of actual data. A field scheduler is responsible for figuring out +//! the order in which pages should be scheduled. Field schedulers then delegate +//! to page schedulers to figure out the I/O requests that need to be made for +//! the page. +//! +//! Page schedulers also create the decoders that will be used to decode the +//! scheduled data. +//! +//! # Decoding +//! +//! Decoders are split into `PhysicalPageDecoder` and +//! [`LogicalPageDecoder`]. Note that both physical and logical decoding +//! happens on a per-page basis. There is no concept of a "field decoder" or +//! "column decoder". +//! +//! The physical decoders handle lower level encodings. They have a few advantages: +//! +//! * They do not need to decode into an Arrow array and so they don't need +//! to be enveloped into the Arrow filesystem (e.g. Arrow doesn't have a +//! bit-packed type. We can use variable-length binary but that is kind +//! of awkward) +//! * They can decode into an existing allocation. This can allow for "page +//! bridging". If we are trying to decode into a batch of 1024 rows and +//! the rows 0..1024 are spread across two pages then we can avoid a memory +//! copy by allocating once and decoding each page into the outer allocation. +//! (note: page bridging is not actually implemented yet) +//! +//! However, there are some limitations for physical decoders: +//! +//! * They are constrained to a single column +//! * The API is more complex +//! +//! The logical decoders are designed to map one or more columns of Lance +//! data into an Arrow array. +//! +//! Typically, a "logical encoding" will have both a logical decoder and a field scheduler. +//! Meanwhile, a "physical encoding" will have a physical decoder but no corresponding field +//! scheduler. +//! +//! +//! # General notes +//! +//! Encodings are typically nested into each other to form a tree. The top of the tree is +//! the user requested schema. Each field in that schema is assigned to one top-level logical +//! encoding. That encoding can then contain other logical encodings or physical encodings. +//! Physical encodings can also contain other physical encodings. +//! +//! So, for example, a single field in the Arrow schema might have the type `List` +//! +//! The encoding tree could then be: +//! +//! root: List (logical encoding) +//! - indices: Primitive (logical encoding) +//! - column: Basic (physical encoding) +//! - validity: Bitmap (physical encoding) +//! - values: RLE (physical encoding) +//! - runs: Value (physical encoding) +//! - values: Value (physical encoding) +//! - items: Primitive (logical encoding) +//! - column: Basic (physical encoding) +//! - values: Value (physical encoding) +//! +//! Note that, in this example, root.items.column does not have a validity because there were +//! no nulls in the page. +//! +//! ## Multiple buffers or multiple columns? +//! +//! Note that there are many different ways we can write encodings. For example, we might +//! store primitive fields in a single column with two buffers (one for validity and one for +//! values) +//! +//! On the other hand, we could also store a primitive field as two different columns. One +//! that yields a non-nullable boolean array and one that yields a non-nullable array of items. +//! Then we could combine these two arrays into a single array where the boolean array is the +//! bitmap. There are a few subtle differences between the approaches: +//! +//! * Storing things as multiple buffers within the same column is generally more efficient and +//! easier to schedule. For example, in-batch coalescing is very easy but can only be done +//! on data that is in the same page. +//! * When things are stored in multiple columns you have to worry about their pages not being +//! in sync. In our previous validity / values example this means we might have to do some +//! memory copies to get the validity array and values arrays to be the same length as +//! decode. +//! * When things are stored in a single column, projection is impossible. For example, if we +//! tried to store all the struct fields in a single column with lots of buffers then we wouldn't +//! be able to read back individual fields of the struct. +//! +//! The fixed size list decoding is an interesting example because it is actually both a physical +//! encoding and a logical encoding. A fixed size list of a physical encoding is, itself, a physical +//! encoding (e.g. a fixed size list of doubles). However, a fixed size list of a logical encoding +//! is a logical encoding (e.g. a fixed size list of structs). +//! +//! # The scheduling loop +//! +//! Reading a Lance file involves both scheduling and decoding. Its generally expected that these +//! will run as two separate threads. +//! +//! ```text +//! +//! I/O PARALLELISM +//! Issues +//! Requests ┌─────────────────┐ +//! │ │ Wait for +//! ┌──────────► I/O Service ├─────► Enough I/O ◄─┐ +//! │ │ │ For batch │ +//! │ └─────────────────┘ │3 │ +//! │ │ │ +//! │ │ │2 +//! ┌─────────────────────┴─┐ ┌─────────▼───────┴┐ +//! │ │ │ │Poll +//! │ Batch Decode │ Decode tasks sent via channel│ Batch Decode │1 +//! │ Scheduler ├─────────────────────────────►│ Stream ◄───── +//! │ │ │ │ +//! └─────▲─────────────┬───┘ └─────────┬────────┘ +//! │ │ │4 +//! │ │ │ +//! └─────────────┘ ┌────────┴────────┐ +//! Caller of schedule_range Buffer polling │ │ +//! will be scheduler thread to achieve CPU │ Decode Batch ├────► +//! and schedule one decode parallelism │ Task │ +//! task (and all needed I/O) (thread per │ │ +//! per logical page batch) └─────────────────┘ +//! ``` +//! +//! The scheduling thread will work through the file from the +//! start to the end as quickly as possible. Data is scheduled one page at a time in a row-major +//! fashion. For example, imagine we have a file with the following page structure: +//! +//! ```text +//! Score (Float32) | C0P0 | +//! Id (16-byte UUID) | C1P0 | C1P1 | C1P2 | C1P3 | +//! Vector (4096 bytes) | C2P0 | C2P1 | C2P2 | C2P3 | .. | C2P1024 | +//! ``` +//! +//! This would be quite common as each of these pages has the same number of bytes. Let's pretend +//! each page is 1MiB and so there are 256Ki rows of data. Each page of `Score` has 256Ki rows. +//! Each page of `Id` has 64Ki rows. Each page of `Vector` has 256 rows. The scheduler would then +//! schedule in the following order: +//! +//! C0 P0 +//! C1 P0 +//! C2 P0 +//! C2 P1 +//! ... (254 pages omitted) +//! C2 P255 +//! C1 P1 +//! C2 P256 +//! ... (254 pages omitted) +//! C2 P511 +//! C1 P2 +//! C2 P512 +//! ... (254 pages omitted) +//! C2 P767 +//! C1 P3 +//! C2 P768 +//! ... (254 pages omitted) +//! C2 P1024 +//! +//! This is the ideal scheduling order because it means we can decode complete rows as quickly as possible. +//! Note that the scheduler thread does not need to wait for I/O to happen at any point. As soon as it starts +//! it will start scheduling one page of I/O after another until it has scheduled the entire file's worth of +//! I/O. This is slightly different than other file readers which have "row group parallelism" and will +//! typically only schedule X row groups worth of reads at a time. +//! +//! In the near future there will be a backpressure mechanism and so it may need to stop/pause if the compute +//! falls behind. +//! +//! ## Indirect I/O +//! +//! Regrettably, there are times where we cannot know exactly what data we need until we have partially decoded +//! the file. This happens when we have variable sized list data. In that case the scheduling task for that +//! page will only schedule the first part of the read (loading the list offsets). It will then immediately +//! spawn a new tokio task to wait for that I/O and decode the list offsets. That follow-up task is not part +//! of the scheduling loop or the decode loop. It is a free task. Once the list offsets are decoded we submit +//! a follow-up I/O task. This task is scheduled at a high priority because the decoder is going to need it soon. +//! +//! # The decode loop +//! +//! As soon as the scheduler starts we can start decoding. Each time we schedule a page we +//! push a decoder for that page's data into a channel. The decode loop +//! ([`BatchDecodeStream`]) reads from that channel. Each time it receives a decoder it +//! waits until the decoder has all of its data. Then it grabs the next decoder. Once it has +//! enough loaded decoders to complete a batch worth of rows it will spawn a "decode batch task". +//! +//! These batch decode tasks perform the actual CPU work of decoding the loaded data into Arrow +//! arrays. This may involve signifciant CPU processing like decompression or arithmetic in order +//! to restore the data to its correct in-memory representation. +//! +//! ## Batch size +//! +//! The `BatchDecodeStream` is configured with a batch size. This does not need to have any +//! relation to the page size(s) used to write the data. This keeps our compute work completely +//! independent of our I/O work. We suggest using small batch sizes: +//! +//! * Batches should fit in CPU cache (at least L3) +//! * More batches means more opportunity for parallelism +//! * The "batch overhead" is very small in Lance compared to other formats because it has no +//! relation to the way the data is stored. + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{LazyLock, Mutex, Once, OnceLock}; +use std::{ops::Range, sync::Arc}; + +use arrow_array::cast::AsArray; +use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader}; +use arrow_schema::{ArrowError, DataType, Field as ArrowField, Fields, Schema as ArrowSchema}; +use bytes::Bytes; +use futures::channel::oneshot; +use futures::future::{BoxFuture, MaybeDone, Shared, maybe_done}; +use futures::stream::{self, BoxStream}; +use futures::{FutureExt, StreamExt}; +use lance_arrow::DataTypeExt; +use lance_core::cache::{Context, DeepSizeOf, LanceCache}; +use lance_core::datatypes::{ + BLOB_DESC_LANCE_FIELD, Field, Schema, validate_fixed_size_list_dimensions, +}; +use lance_core::utils::futures::{FinallyStreamExt, StreamOnDropExt}; +use lance_core::utils::parse::parse_env_as_bool; +use log::{debug, trace, warn}; +use prost::Message; +use tokio::sync::mpsc::error::SendError; +use tokio::sync::mpsc::{self, unbounded_channel}; + +use lance_core::error::{CloneableError, LanceOptionExt}; +use lance_core::{ArrowResult, Error, Result}; +use tracing::instrument; + +use crate::array_encoding::logical::list::OffsetPageInfo; +use crate::array_encoding::logical::r#struct::{SimpleStructDecoder, SimpleStructScheduler}; +use crate::array_encoding::logical::{ + binary::BinaryFieldScheduler, blob::BlobFieldScheduler, list::ListFieldScheduler, + primitive::PrimitiveFieldScheduler, +}; +use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy}; +use crate::data::DataBlock; +use crate::encoder::EncodedBatch; +use crate::encodings::logical::fixed_size_list::StructuralFixedSizeListScheduler; +use crate::encodings::logical::list::StructuralListScheduler; +use crate::encodings::logical::map::StructuralMapScheduler; +use crate::encodings::logical::primitive::StructuralPrimitiveFieldScheduler; +use crate::encodings::logical::r#struct::{StructuralStructDecoder, StructuralStructScheduler}; +use crate::format::pb::{self, column_encoding}; +use crate::format::pb21; +use crate::repdef::{CompositeRepDefUnraveler, RepDefUnraveler}; +use crate::{BufferScheduler, EncodingsIo}; + +/// Candidate batch sizes evaluated during byte-budget planning. +/// Powers of 4, covering 1–16Ki rows in 8 probes. +pub const CANDIDATE_BATCH_SIZES: [u32; 8] = [1, 4, 16, 64, 256, 1024, 4096, 16384]; + +pub trait SchedulingJob: std::fmt::Debug { + fn schedule_next( + &mut self, + context: &mut SchedulerContext, + priority: &dyn PriorityRange, + ) -> Result; + + fn num_rows(&self) -> u64; +} + +/// Schedules the I/O needed to decode one field. +pub trait FieldScheduler: Send + Sync + std::fmt::Debug { + fn initialize<'a>( + &'a self, + filter: &'a FilterExpression, + context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>>; + + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result>; + + fn num_rows(&self) -> u64; +} + +#[derive(Debug)] +pub struct DecoderReady { + pub decoder: Box, + pub path: VecDeque, +} + +/// Stateful decoder for one logical page. +pub trait LogicalPageDecoder: std::fmt::Debug + Send { + fn accept_child(&mut self, _child: DecoderReady) -> Result<()> { + Err(Error::internal(format!( + "The decoder {:?} does not expect children but received a child", + self + ))) + } + + fn wait_for_loaded(&'_ mut self, loaded_need: u64) -> BoxFuture<'_, Result<()>>; + + fn rows_loaded(&self) -> u64; + + fn rows_unloaded(&self) -> u64 { + self.num_rows() - self.rows_loaded() + } + + fn num_rows(&self) -> u64; + + fn rows_drained(&self) -> u64; + + fn rows_left(&self) -> u64 { + self.num_rows() - self.rows_drained() + } + + fn drain(&mut self, num_rows: u64) -> Result; + + /// Returns how many of the requested rows can be decoded into one Arrow array + /// while keeping every i32 offset buffer within `byte_budget` bytes. + /// + /// Most page decoders can always drain the full request and consume no budget. + /// Decoders that build arrays with i32 offsets (Utf8 / Binary / List) may return + /// fewer rows so that the concatenated array's offsets stay representable, and + /// report the budget those rows consume so callers can accumulate across pages. + fn max_rows_to_drain(&self, num_rows: u64, _byte_budget: u64) -> Result { + Ok(DrainLimit { + rows: num_rows, + bytes: 0, + }) + } + + fn data_type(&self) -> &DataType; +} + +/// Result of limiting a drain request to a variable-width byte budget. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DrainLimit { + /// How many of the requested rows can be decoded into one Arrow array + pub rows: u64, + /// The variable-width bytes (at the most constrained i32-offset nesting level) + /// those rows consume from the caller's budget + pub bytes: u64, +} + +/// The byte budget applied to each output batch of an i32-offset column. +/// +/// Concatenating decoded pages into one Arrow array fails when a variable-width +/// buffer's offsets exceed `i32::MAX`, so batches are split before that point. +pub(crate) const I32_OFFSET_BYTE_BUDGET: u64 = i32::MAX as u64; + +/// Clamp a batch's row count to the byte-budgeted limit. +/// +/// A batch always advances by at least one row: a first row whose (possibly +/// over-estimated) size exceeds the whole budget is emitted alone. A genuinely +/// oversized row then surfaces the same arrow-level offset error it always did, +/// while an over-estimated one decodes fine. +fn clamp_rows_to_i32_budget(limit: DrainLimit, to_take: u64) -> u64 { + if to_take == 0 { + return 0; + } + limit.rows.clamp(1, to_take) +} + +/// Returns whether concatenating arrays of this type can overflow an i32 offset buffer. +pub(crate) fn has_i32_offsets(data_type: &DataType) -> bool { + match data_type { + DataType::Binary | DataType::Utf8 | DataType::List(_) | DataType::ListView(_) => true, + DataType::Map(_, _) => true, + DataType::LargeList(field) + | DataType::LargeListView(field) + | DataType::FixedSizeList(field, _) => has_i32_offsets(field.data_type()), + DataType::Struct(fields) => fields + .iter() + .any(|field| has_i32_offsets(field.data_type())), + DataType::Dictionary(_, values) => has_i32_offsets(values), + DataType::RunEndEncoded(_, values) => has_i32_offsets(values.data_type()), + _ => false, + } +} + +// If users are getting batches over 10MiB large then it's time to reduce the batch size +const BATCH_SIZE_BYTES_WARNING: u64 = 10 * 1024 * 1024; +const ENV_LANCE_STRUCTURAL_BATCH_DECODE_SPAWN_MODE: &str = + "LANCE_STRUCTURAL_BATCH_DECODE_SPAWN_MODE"; +const ENV_LANCE_READ_CACHE_REPETITION_INDEX: &str = "LANCE_READ_CACHE_REPETITION_INDEX"; +const ENV_LANCE_INLINE_SCHEDULING_THRESHOLD: &str = "LANCE_INLINE_SCHEDULING_THRESHOLD"; + +// If a request is for at most this many rows we skip the scheduler-task spawn +// and run scheduling inline as part of the `schedule_and_decode` await. +const DEFAULT_INLINE_SCHEDULING_THRESHOLD: u64 = 16 * 1024; + +fn default_cache_repetition_index() -> bool { + static DEFAULT_CACHE_REPETITION_INDEX: OnceLock = OnceLock::new(); + *DEFAULT_CACHE_REPETITION_INDEX + .get_or_init(|| parse_env_as_bool(ENV_LANCE_READ_CACHE_REPETITION_INDEX, true)) +} + +fn inline_scheduling_threshold() -> u64 { + static THRESHOLD: OnceLock = OnceLock::new(); + *THRESHOLD.get_or_init(|| { + std::env::var(ENV_LANCE_INLINE_SCHEDULING_THRESHOLD) + .ok() + .and_then(|v| v.trim().parse::().ok()) + .unwrap_or(DEFAULT_INLINE_SCHEDULING_THRESHOLD) + }) +} + +/// Top-level encoding message for a page. Wraps both the v2.0 +/// [`pb::ArrayEncoding`] grammar and the structural [`pb21::PageLayout`] grammar. +/// +/// A file should only use one or the other and never both. +/// 2.0 decoders can always assume this is pb::ArrayEncoding +/// and 2.1+ decoders can always assume this is pb::PageLayout +#[derive(Debug, Clone)] +pub enum PageEncoding { + Legacy(pb::ArrayEncoding), + Structural(pb21::PageLayout), +} + +impl DeepSizeOf for PageEncoding { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + match self { + Self::Legacy(encoding) => encoding.encoded_len() * 4, + Self::Structural(encoding) => encoding.encoded_len() * 4, + } + } +} + +impl PageEncoding { + pub fn as_legacy(&self) -> &pb::ArrayEncoding { + match self { + Self::Legacy(enc) => enc, + Self::Structural(_) => panic!("Expected a legacy encoding"), + } + } + + pub fn as_structural(&self) -> &pb21::PageLayout { + match self { + Self::Structural(enc) => enc, + Self::Legacy(_) => panic!("Expected a structural encoding"), + } + } + + pub fn is_structural(&self) -> bool { + matches!(self, Self::Structural(_)) + } +} + +/// Metadata describing a page in a file +/// +/// This is typically created by reading the metadata section of a Lance file +#[derive(Debug)] +pub struct PageInfo { + /// The number of rows in the page + pub num_rows: u64, + /// The priority (top level row number) of the page + /// + /// This is only set in 2.1 files and will be 0 for 2.0 files + pub priority: u64, + /// The encoding that explains the buffers in the page + pub encoding: PageEncoding, + /// The offsets and sizes of the buffers in the file + pub buffer_offsets_and_sizes: Arc<[(u64, u64)]>, +} + +impl DeepSizeOf for PageInfo { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + self.encoding.deep_size_of_children(context) + + self.buffer_offsets_and_sizes.deep_size_of_children(context) + } +} + +/// Metadata describing a column in a file +/// +/// This is typically created by reading the metadata section of a Lance file +#[derive(Debug, Clone)] +pub struct ColumnInfo { + /// The index of the column in the file + pub index: u32, + /// The metadata for each page in the column + pub page_infos: Arc<[PageInfo]>, + /// File positions and their sizes of the column-level buffers + pub buffer_offsets_and_sizes: Arc<[(u64, u64)]>, + pub encoding: pb::ColumnEncoding, +} + +impl DeepSizeOf for ColumnInfo { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + self.page_infos.deep_size_of_children(context) + + self.buffer_offsets_and_sizes.deep_size_of_children(context) + + self.encoding.encoded_len() * 4 + } +} + +impl ColumnInfo { + /// Create a new instance + pub fn new( + index: u32, + page_infos: Arc<[PageInfo]>, + buffer_offsets_and_sizes: Vec<(u64, u64)>, + encoding: pb::ColumnEncoding, + ) -> Self { + Self { + index, + page_infos, + buffer_offsets_and_sizes: buffer_offsets_and_sizes.into_boxed_slice().into(), + encoding, + } + } + + pub fn is_structural(&self) -> bool { + self.page_infos + // Can just look at the first since all should be the same + .first() + .map(|page| page.encoding.is_structural()) + .unwrap_or(false) + } +} + +enum RootScheduler { + Structural(Box), + Array(Arc), +} + +impl RootScheduler { + fn as_array(&self) -> &Arc { + match self { + Self::Structural(_) => panic!("Expected an array scheduler"), + Self::Array(s) => s, + } + } + + fn as_structural(&self) -> &dyn StructuralFieldScheduler { + match self { + Self::Structural(s) => s.as_ref(), + Self::Array(_) => panic!("Expected a structural scheduler"), + } + } +} + +/// The scheduler for decoding batches +/// +/// Lance decoding is done in two steps, scheduling, and decoding. The +/// scheduling tends to be lightweight and should quickly figure what data +/// is needed from the disk issue the appropriate I/O requests. A decode task is +/// created to eventually decode the data (once it is loaded) and scheduling +/// moves on to scheduling the next page. +/// +/// Meanwhile, it's expected that a decode stream will be setup to run at the +/// same time. Decode tasks take the data that is loaded and turn it into +/// Arrow arrays. +/// +/// This approach allows us to keep our I/O parallelism and CPU parallelism +/// completely separate since those are often two very different values. +/// +/// Backpressure should be achieved via the I/O service. Requests that are +/// issued will pile up if the decode stream is not polling quickly enough. +/// The [`crate::EncodingsIo::submit_request`] function should return a pending +/// future once there are too many I/O requests in flight. +/// +/// TODO: Implement backpressure +pub struct DecodeBatchScheduler { + root_scheduler: RootScheduler, + pub root_fields: Fields, + cache: Arc, +} + +pub struct ColumnInfoIter<'a> { + column_infos: Vec>, + column_indices: &'a [u32], + column_info_pos: usize, + column_indices_pos: usize, +} + +impl<'a> ColumnInfoIter<'a> { + pub fn new(column_infos: Vec>, column_indices: &'a [u32]) -> Self { + let initial_pos = column_indices.first().copied().unwrap_or(0) as usize; + Self { + column_infos, + column_indices, + column_info_pos: initial_pos, + column_indices_pos: 0, + } + } + + pub fn peek(&self) -> &Arc { + &self.column_infos[self.column_info_pos] + } + + pub fn peek_transform(&mut self, transform: impl FnOnce(Arc) -> Arc) { + let column_info = self.column_infos[self.column_info_pos].clone(); + let transformed = transform(column_info); + self.column_infos[self.column_info_pos] = transformed; + } + + pub fn expect_next(&mut self) -> Result<&Arc> { + self.next().ok_or_else(|| { + Error::invalid_input( + "there were more fields in the schema than provided column indices / infos", + ) + }) + } + + fn next(&mut self) -> Option<&Arc> { + if self.column_info_pos < self.column_infos.len() { + let info = &self.column_infos[self.column_info_pos]; + self.column_info_pos += 1; + Some(info) + } else { + None + } + } + + pub(crate) fn next_top_level(&mut self) { + self.column_indices_pos += 1; + if self.column_indices_pos < self.column_indices.len() { + self.column_info_pos = self.column_indices[self.column_indices_pos] as usize; + } else { + self.column_info_pos = self.column_infos.len(); + } + } +} + +/// These contain the file buffers shared across the entire file +#[derive(Clone, Copy, Debug)] +pub struct FileBuffers<'a> { + pub positions_and_sizes: &'a [(u64, u64)], +} + +/// These contain the file buffers and also buffers specific to a column +#[derive(Clone, Copy, Debug)] +pub struct ColumnBuffers<'a, 'b> { + pub file_buffers: FileBuffers<'a>, + pub positions_and_sizes: &'b [(u64, u64)], +} + +/// These contain the file & column buffers and also buffers specific to a page +#[derive(Clone, Copy, Debug)] +pub struct PageBuffers<'a, 'b, 'c> { + pub column_buffers: ColumnBuffers<'a, 'b>, + pub positions_and_sizes: &'c [(u64, u64)], +} + +/// The core decoder strategy handles all the various Arrow types +#[derive(Debug)] +pub struct CoreFieldDecoderStrategy { + pub validate_data: bool, + pub decompressor_strategy: Arc, + pub cache_repetition_index: bool, +} + +impl Default for CoreFieldDecoderStrategy { + fn default() -> Self { + Self { + validate_data: false, + decompressor_strategy: Arc::new(DefaultDecompressionStrategy {}), + cache_repetition_index: false, + } + } +} + +impl CoreFieldDecoderStrategy { + /// Create a new strategy with cache_repetition_index enabled + pub fn with_cache_repetition_index(mut self, cache_repetition_index: bool) -> Self { + self.cache_repetition_index = cache_repetition_index; + self + } + + /// Create a new strategy from decoder config + pub fn from_decoder_config(config: &DecoderConfig) -> Self { + Self { + validate_data: config.validate_on_decode, + decompressor_strategy: Arc::new(DefaultDecompressionStrategy {}), + cache_repetition_index: config.cache_repetition_index, + } + } + + /// This is just a sanity check to ensure there is no "wrapped encodings" + /// that haven't been handled. + fn ensure_values_encoded(column_info: &ColumnInfo, field_name: &str) -> Result<()> { + let column_encoding = column_info + .encoding + .column_encoding + .as_ref() + .ok_or_else(|| { + Error::invalid_input(format!( + "the column at index {} was missing a ColumnEncoding", + column_info.index + )) + })?; + if matches!( + column_encoding, + pb::column_encoding::ColumnEncoding::Values(_) + ) { + Ok(()) + } else { + Err(Error::invalid_input(format!( + "the column at index {} mapping to the input field {} has column encoding {:?} and no decoder is registered to handle it", + column_info.index, field_name, column_encoding + ))) + } + } + + fn is_structural_primitive(data_type: &DataType) -> bool { + if data_type.is_primitive() { + true + } else { + match data_type { + // DataType::is_primitive doesn't consider these primitive but we do + DataType::Dictionary(_, value_type) => Self::is_structural_primitive(value_type), + DataType::Boolean + | DataType::Null + | DataType::FixedSizeBinary(_) + | DataType::Binary + | DataType::LargeBinary + | DataType::Utf8 + | DataType::LargeUtf8 => true, + DataType::FixedSizeList(inner, _) => { + Self::is_structural_primitive(inner.data_type()) + } + _ => false, + } + } + } + + fn is_array_primitive(data_type: &DataType) -> bool { + if data_type.is_primitive() { + true + } else { + match data_type { + // DataType::is_primitive doesn't consider these primitive but we do + DataType::Boolean | DataType::Null | DataType::FixedSizeBinary(_) => true, + DataType::FixedSizeList(inner, _) => Self::is_array_primitive(inner.data_type()), + _ => false, + } + } + } + + fn create_primitive_scheduler( + &self, + field: &Field, + column: &ColumnInfo, + buffers: FileBuffers, + ) -> Result> { + Self::ensure_values_encoded(column, &field.name)?; + // Primitive fields map to a single column + let column_buffers = ColumnBuffers { + file_buffers: buffers, + positions_and_sizes: &column.buffer_offsets_and_sizes, + }; + Ok(Box::new(PrimitiveFieldScheduler::new( + column.index, + field.data_type(), + column.page_infos.clone(), + column_buffers, + self.validate_data, + ))) + } + + /// Helper method to verify the page encoding of a struct header column + fn check_simple_struct(column_info: &ColumnInfo, field_name: &str) -> Result<()> { + Self::ensure_values_encoded(column_info, field_name)?; + if column_info.page_infos.len() != 1 { + return Err(Error::invalid_input_source(format!("Due to schema we expected a struct column but we received a column with {} pages and right now we only support struct columns with 1 page", column_info.page_infos.len()).into())); + } + let encoding = &column_info.page_infos[0].encoding; + match encoding.as_legacy().array_encoding.as_ref().unwrap() { + pb::array_encoding::ArrayEncoding::Struct(_) => Ok(()), + _ => Err(Error::invalid_input_source(format!("Expected a struct encoding because we have a struct field in the schema but got the encoding {:?}", encoding).into())), + } + } + + fn check_packed_struct(column_info: &ColumnInfo) -> bool { + let encoding = &column_info.page_infos[0].encoding; + matches!( + encoding.as_legacy().array_encoding.as_ref().unwrap(), + pb::array_encoding::ArrayEncoding::PackedStruct(_) + ) + } + + fn create_list_scheduler( + &self, + list_field: &Field, + column_infos: &mut ColumnInfoIter, + buffers: FileBuffers, + offsets_column: &ColumnInfo, + ) -> Result> { + Self::ensure_values_encoded(offsets_column, &list_field.name)?; + let offsets_column_buffers = ColumnBuffers { + file_buffers: buffers, + positions_and_sizes: &offsets_column.buffer_offsets_and_sizes, + }; + let items_scheduler = + self.create_array_field_scheduler(&list_field.children[0], column_infos, buffers)?; + + let mut inner_infos = Vec::with_capacity(offsets_column.page_infos.len()); + let mut null_offset_adjustments = Vec::with_capacity(offsets_column.page_infos.len()); + for (page_index, offsets_page) in offsets_column + .page_infos + .iter() + .enumerate() + .filter(|(_, offsets_page)| offsets_page.num_rows > 0) + { + let PageEncoding::Legacy(pb::ArrayEncoding { + array_encoding: Some(pb::array_encoding::ArrayEncoding::List(list_encoding)), + }) = &offsets_page.encoding + else { + return Err(Error::invalid_input(format!( + "expected list encoding for field '{}' in column {}, page {} but got {:?}", + list_field.name, offsets_column.index, page_index, offsets_page.encoding + ))); + }; + let offsets_encoding = list_encoding.offsets.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "list encoding for field '{}' in column {}, page {} is missing its offsets encoding", + list_field.name, offsets_column.index, page_index + )) + })?; + inner_infos.push(PageInfo { + buffer_offsets_and_sizes: offsets_page.buffer_offsets_and_sizes.clone(), + encoding: PageEncoding::Legacy(offsets_encoding.as_ref().clone()), + num_rows: offsets_page.num_rows, + priority: 0, + }); + null_offset_adjustments.push(OffsetPageInfo { + offsets_in_page: offsets_page.num_rows, + null_offset_adjustment: list_encoding.null_offset_adjustment, + num_items_referenced_by_page: list_encoding.num_items, + }); + } + let inner = Arc::new(PrimitiveFieldScheduler::new( + offsets_column.index, + DataType::UInt64, + Arc::from(inner_infos.into_boxed_slice()), + offsets_column_buffers, + self.validate_data, + )) as Arc; + let items_field = match list_field.data_type() { + DataType::List(inner) => inner, + DataType::LargeList(inner) => inner, + _ => unreachable!(), + }; + let offset_type = if matches!(list_field.data_type(), DataType::List(_)) { + DataType::Int32 + } else { + DataType::Int64 + }; + Ok(Box::new(ListFieldScheduler::new( + inner, + items_scheduler.into(), + items_field, + offset_type, + null_offset_adjustments, + ))) + } + + fn unwrap_blob(column_info: &ColumnInfo) -> Option { + if let column_encoding::ColumnEncoding::Blob(blob) = + column_info.encoding.column_encoding.as_ref().unwrap() + { + let mut column_info = column_info.clone(); + column_info.encoding = blob.inner.as_ref().unwrap().as_ref().clone(); + Some(column_info) + } else { + None + } + } + + fn create_structural_field_scheduler( + &self, + field: &Field, + column_infos: &mut ColumnInfoIter, + ) -> Result> { + let data_type = field.data_type(); + validate_fixed_size_list_dimensions(&field.name, &data_type)?; + if Self::is_structural_primitive(&data_type) { + let column_info = column_infos.expect_next()?; + let scheduler = Box::new(StructuralPrimitiveFieldScheduler::try_new_lazy( + column_info.clone(), + self.decompressor_strategy.clone(), + self.cache_repetition_index, + field, + )?); + + // advance to the next top level column + column_infos.next_top_level(); + + return Ok(scheduler); + } + match &data_type { + DataType::Struct(fields) => { + if field.is_packed_struct() { + // Packed struct + let column_info = column_infos.expect_next()?; + let scheduler = Box::new(StructuralPrimitiveFieldScheduler::try_new_lazy( + column_info.clone(), + self.decompressor_strategy.clone(), + self.cache_repetition_index, + field, + )?); + + // advance to the next top level column + column_infos.next_top_level(); + + return Ok(scheduler); + } + // Maybe a blob descriptions struct? + if field.is_blob() { + let column_info = column_infos.peek(); + if column_info.page_infos.iter().any(|page| { + matches!( + page.encoding, + PageEncoding::Structural(pb21::PageLayout { + layout: Some(pb21::page_layout::Layout::BlobLayout(_)) + }) + ) + }) { + let column_info = column_infos.expect_next()?; + let scheduler = Box::new(StructuralPrimitiveFieldScheduler::try_new_lazy( + column_info.clone(), + self.decompressor_strategy.clone(), + self.cache_repetition_index, + field, + )?); + column_infos.next_top_level(); + return Ok(scheduler); + } + } + + let mut child_schedulers = Vec::with_capacity(field.children.len()); + for field in field.children.iter() { + let field_scheduler = + self.create_structural_field_scheduler(field, column_infos)?; + child_schedulers.push(field_scheduler); + } + + let fields = fields.clone(); + Ok( + Box::new(StructuralStructScheduler::new(child_schedulers, fields)) + as Box, + ) + } + DataType::List(_) | DataType::LargeList(_) => { + let child = field.children.first().expect_ok()?; + let child_scheduler = + self.create_structural_field_scheduler(child, column_infos)?; + Ok(Box::new(StructuralListScheduler::new(child_scheduler)) + as Box) + } + DataType::FixedSizeList(inner, dimension) + if matches!(inner.data_type(), DataType::Struct(_)) => + { + let child = field.children.first().expect_ok()?; + let child_scheduler = + self.create_structural_field_scheduler(child, column_infos)?; + Ok(Box::new(StructuralFixedSizeListScheduler::new( + child_scheduler, + *dimension, + )) as Box) + } + DataType::Map(_, keys_sorted) => { + // TODO: We only support keys_sorted=false for now, + // because converting a rust arrow map field to the python arrow field will + // lose the keys_sorted property. + if *keys_sorted { + return Err(Error::not_supported_source(format!("Map data type is not supported with keys_sorted=true now, current value is {}", *keys_sorted).into())); + } + let entries_child = field.children.first().expect_ok()?; + let child_scheduler = + self.create_structural_field_scheduler(entries_child, column_infos)?; + Ok(Box::new(StructuralMapScheduler::new(child_scheduler)) + as Box) + } + _ => todo!("create_structural_field_scheduler for {}", data_type), + } + } + + fn create_array_field_scheduler( + &self, + field: &Field, + column_infos: &mut ColumnInfoIter, + buffers: FileBuffers, + ) -> Result> { + let data_type = field.data_type(); + validate_fixed_size_list_dimensions(&field.name, &data_type)?; + if Self::is_array_primitive(&data_type) { + let column_info = column_infos.expect_next()?; + let scheduler = self.create_primitive_scheduler(field, column_info, buffers)?; + return Ok(scheduler); + } else if data_type.is_binary_like() { + let column_info = column_infos.expect_next()?.clone(); + // Column is blob and user is asking for binary data + if let Some(blob_col) = Self::unwrap_blob(column_info.as_ref()) { + let desc_scheduler = + self.create_primitive_scheduler(&BLOB_DESC_LANCE_FIELD, &blob_col, buffers)?; + let blob_scheduler = Box::new(BlobFieldScheduler::new(desc_scheduler.into())); + return Ok(blob_scheduler); + } + if let Some(page_info) = column_info.page_infos.first() { + if matches!( + page_info.encoding.as_legacy(), + pb::ArrayEncoding { + array_encoding: Some(pb::array_encoding::ArrayEncoding::List(..)) + } + ) { + let list_type = if matches!(data_type, DataType::Utf8 | DataType::Binary) { + DataType::List(Arc::new(ArrowField::new("item", DataType::UInt8, false))) + } else { + DataType::LargeList(Arc::new(ArrowField::new( + "item", + DataType::UInt8, + false, + ))) + }; + let list_field = Field::try_from(ArrowField::new( + field.name.clone(), + list_type, + field.nullable, + )) + .unwrap(); + let list_scheduler = self.create_list_scheduler( + &list_field, + column_infos, + buffers, + &column_info, + )?; + let binary_scheduler = Box::new(BinaryFieldScheduler::new( + list_scheduler.into(), + field.data_type(), + )); + return Ok(binary_scheduler); + } else { + let scheduler = + self.create_primitive_scheduler(field, &column_info, buffers)?; + return Ok(scheduler); + } + } else { + return self.create_primitive_scheduler(field, &column_info, buffers); + } + } + match &data_type { + DataType::FixedSizeList(inner, _dimension) => { + // A fixed size list column could either be a physical or a logical decoder + // depending on the child data type. + if Self::is_array_primitive(inner.data_type()) { + let primitive_col = column_infos.expect_next()?; + let scheduler = + self.create_primitive_scheduler(field, primitive_col, buffers)?; + Ok(scheduler) + } else { + todo!() + } + } + DataType::Dictionary(_key_type, value_type) => { + if Self::is_array_primitive(value_type) || value_type.is_binary_like() { + let primitive_col = column_infos.expect_next()?; + let scheduler = + self.create_primitive_scheduler(field, primitive_col, buffers)?; + Ok(scheduler) + } else { + Err(Error::not_supported_source( + format!( + "No way to decode into a dictionary field of type {}", + value_type + ) + .into(), + )) + } + } + DataType::List(_) | DataType::LargeList(_) => { + let offsets_column = column_infos.expect_next()?.clone(); + column_infos.next_top_level(); + self.create_list_scheduler(field, column_infos, buffers, &offsets_column) + } + DataType::Struct(fields) => { + let column_info = column_infos.expect_next()?; + + // Column is blob and user is asking for descriptions + if let Some(blob_col) = Self::unwrap_blob(column_info.as_ref()) { + // Can use primitive scheduler here since descriptions are always packed struct + return self.create_primitive_scheduler(field, &blob_col, buffers); + } + + if Self::check_packed_struct(column_info) { + // use packed struct encoding + self.create_primitive_scheduler(field, column_info, buffers) + } else { + // use default struct encoding + Self::check_simple_struct(column_info, &field.name).unwrap(); + let num_rows = column_info + .page_infos + .iter() + .map(|page| page.num_rows) + .sum(); + let mut child_schedulers = Vec::with_capacity(field.children.len()); + for field in &field.children { + column_infos.next_top_level(); + let field_scheduler = + self.create_array_field_scheduler(field, column_infos, buffers)?; + child_schedulers.push(Arc::from(field_scheduler)); + } + + let fields = fields.clone(); + Ok(Box::new(SimpleStructScheduler::new( + child_schedulers, + fields, + num_rows, + ))) + } + } + // TODO: Still need support for RLE + _ => todo!(), + } + } +} + +/// Create's a dummy ColumnInfo for the root column +fn root_column(num_rows: u64) -> ColumnInfo { + let num_root_pages = num_rows.div_ceil(u32::MAX as u64); + let final_page_num_rows = num_rows % (u32::MAX as u64); + let root_pages = (0..num_root_pages) + .map(|i| PageInfo { + num_rows: if i == num_root_pages - 1 { + final_page_num_rows + } else { + u64::MAX + }, + encoding: PageEncoding::Legacy(pb::ArrayEncoding { + array_encoding: Some(pb::array_encoding::ArrayEncoding::Struct( + pb::SimpleStruct {}, + )), + }), + priority: 0, // not used by the array scheduler + buffer_offsets_and_sizes: Arc::new([]), + }) + .collect::>(); + ColumnInfo { + buffer_offsets_and_sizes: Arc::new([]), + encoding: pb::ColumnEncoding { + column_encoding: Some(pb::column_encoding::ColumnEncoding::Values(())), + }, + index: u32::MAX, + page_infos: Arc::from(root_pages), + } +} + +pub enum RootDecoder { + Structural(StructuralStructDecoder), + Array(SimpleStructDecoder), +} + +impl RootDecoder { + pub fn into_structural(self) -> StructuralStructDecoder { + match self { + Self::Structural(decoder) => decoder, + Self::Array(_) => panic!("Expected a structural decoder"), + } + } + + pub fn into_array(self) -> SimpleStructDecoder { + match self { + Self::Array(decoder) => decoder, + Self::Structural(_) => panic!("Expected an array decoder"), + } + } +} + +impl DecodeBatchScheduler { + /// Creates a new decode scheduler with the expected schema and the column + /// metadata of the file. + #[allow(clippy::too_many_arguments)] + pub async fn try_new<'a>( + schema: &'a Schema, + column_indices: &[u32], + column_infos: &[Arc], + file_buffer_positions_and_sizes: &'a Vec<(u64, u64)>, + num_rows: u64, + decoder_plugins: Arc, + io: Arc, + cache: Arc, + filter: &FilterExpression, + decoder_config: &DecoderConfig, + ) -> Result { + Self::try_new_with_ranges( + schema, + column_indices, + column_infos, + file_buffer_positions_and_sizes, + num_rows, + decoder_plugins, + io, + cache, + None, + filter, + decoder_config, + ) + .await + } + + /// Creates a decode scheduler and initializes only pages overlapping known ranges. + #[allow(clippy::too_many_arguments)] + pub(crate) async fn try_new_with_ranges<'a>( + schema: &'a Schema, + column_indices: &[u32], + column_infos: &[Arc], + file_buffer_positions_and_sizes: &'a Vec<(u64, u64)>, + num_rows: u64, + _decoder_plugins: Arc, + io: Arc, + cache: Arc, + // The top-level row ranges that will be scheduled, if known. This lets + // the structural path initialize only the pages those ranges touch. + // `None` preserves `try_new`'s eager initialization behavior. + requested_ranges: Option<&[Range]>, + filter: &FilterExpression, + decoder_config: &DecoderConfig, + ) -> Result { + assert!(num_rows > 0); + let buffers = FileBuffers { + positions_and_sizes: file_buffer_positions_and_sizes, + }; + let arrow_schema = ArrowSchema::from(schema); + let root_fields = arrow_schema.fields().clone(); + let root_type = DataType::Struct(root_fields.clone()); + let mut root_field = Field::try_from(&ArrowField::new("root", root_type, false))?; + // root_field.children and schema.fields should be identical at this point but the latter + // has field ids and the former does not. This line restores that. + // TODO: Is there another way to create the root field without forcing a trip through arrow? + root_field.children.clone_from(&schema.fields); + root_field + .metadata + .insert("__lance_decoder_root".to_string(), "true".to_string()); + + if column_infos.is_empty() || column_infos[0].is_structural() { + let mut column_iter = ColumnInfoIter::new(column_infos.to_vec(), column_indices); + + let strategy = CoreFieldDecoderStrategy::from_decoder_config(decoder_config); + let mut root_scheduler = + strategy.create_structural_field_scheduler(&root_field, &mut column_iter)?; + + let context = SchedulerContext::new(io, cache.clone()); + root_scheduler + .initialize(requested_ranges, filter, &context) + .await?; + + Ok(Self { + root_scheduler: RootScheduler::Structural(root_scheduler), + root_fields, + cache, + }) + } else { + // The old encoding style expected a header column for structs and so we + // need a header column for the top-level struct + let mut columns = Vec::with_capacity(column_infos.len() + 1); + columns.push(Arc::new(root_column(num_rows))); + columns.extend(column_infos.iter().cloned()); + + let adjusted_column_indices = [0_u32] + .into_iter() + .chain(column_indices.iter().map(|i| i.saturating_add(1))) + .collect::>(); + let mut column_iter = ColumnInfoIter::new(columns, &adjusted_column_indices); + let strategy = CoreFieldDecoderStrategy::from_decoder_config(decoder_config); + let root_scheduler = + strategy.create_array_field_scheduler(&root_field, &mut column_iter, buffers)?; + + let context = SchedulerContext::new(io, cache.clone()); + root_scheduler.initialize(filter, &context).await?; + + Ok(Self { + root_scheduler: RootScheduler::Array(root_scheduler.into()), + root_fields, + cache, + }) + } + } + + #[deprecated(since = "0.29.1", note = "This is for v2.0 array-encoding paths")] + pub fn from_scheduler( + root_scheduler: Arc, + root_fields: Fields, + cache: Arc, + ) -> Self { + Self { + root_scheduler: RootScheduler::Array(root_scheduler), + root_fields, + cache, + } + } + + fn do_schedule_ranges_structural( + &mut self, + ranges: &[Range], + filter: &FilterExpression, + io: Arc, + mut schedule_action: impl FnMut(Result) -> bool, + ) { + let root_scheduler = self.root_scheduler.as_structural(); + let mut context = SchedulerContext::new(io, self.cache.clone()); + let maybe_root_job = root_scheduler.schedule_ranges(ranges, filter); + if let Err(schedule_ranges_err) = maybe_root_job { + schedule_action(Err(schedule_ranges_err)); + return; + } + let mut root_job = maybe_root_job.unwrap(); + let mut num_rows_scheduled = 0; + loop { + let maybe_next_scan_lines = root_job.schedule_next(&mut context); + if let Err(err) = maybe_next_scan_lines { + schedule_action(Err(err)); + return; + } + let next_scan_lines = maybe_next_scan_lines.unwrap(); + context.flush_io(); + if next_scan_lines.is_empty() { + return; + } + for next_scan_line in next_scan_lines { + trace!( + "Scheduled scan line of {} rows and {} decoders", + next_scan_line.rows_scheduled, + next_scan_line.decoders.len() + ); + num_rows_scheduled += next_scan_line.rows_scheduled; + if !schedule_action(Ok(DecoderMessage { + scheduled_so_far: num_rows_scheduled, + decoders: next_scan_line.decoders, + })) { + // Decoder has disconnected + return; + } + } + } + } + + fn do_schedule_ranges_array( + &mut self, + ranges: &[Range], + filter: &FilterExpression, + io: Arc, + mut schedule_action: impl FnMut(Result) -> bool, + // If specified, this will be used as the top_level_row for all scheduling + // tasks. This is used by list scheduling to ensure all items scheduling + // tasks are scheduled at the same top level row. + priority: Option>, + ) { + let root_scheduler = self.root_scheduler.as_array(); + let rows_requested = ranges.iter().map(|r| r.end - r.start).sum::(); + trace!( + "Scheduling {} ranges across {}..{} ({} rows){}", + ranges.len(), + ranges.first().unwrap().start, + ranges.last().unwrap().end, + rows_requested, + priority + .as_ref() + .map(|p| format!(" (priority={:?})", p)) + .unwrap_or_default() + ); + + let mut context = SchedulerContext::new(io, self.cache.clone()); + let maybe_root_job = root_scheduler.schedule_ranges(ranges, filter); + if let Err(schedule_ranges_err) = maybe_root_job { + schedule_action(Err(schedule_ranges_err)); + return; + } + let mut root_job = maybe_root_job.unwrap(); + let mut num_rows_scheduled = 0; + let mut rows_to_schedule = root_job.num_rows(); + let mut priority = priority.unwrap_or(Box::new(SimplePriorityRange::new(0))); + trace!("Scheduled ranges refined to {} rows", rows_to_schedule); + while rows_to_schedule > 0 { + let maybe_next_scan_line = root_job.schedule_next(&mut context, priority.as_ref()); + if let Err(schedule_next_err) = maybe_next_scan_line { + schedule_action(Err(schedule_next_err)); + return; + } + let next_scan_line = maybe_next_scan_line.unwrap(); + context.flush_io(); + priority.advance(next_scan_line.rows_scheduled); + num_rows_scheduled += next_scan_line.rows_scheduled; + rows_to_schedule -= next_scan_line.rows_scheduled; + trace!( + "Scheduled scan line of {} rows and {} decoders", + next_scan_line.rows_scheduled, + next_scan_line.decoders.len() + ); + if !schedule_action(Ok(DecoderMessage { + scheduled_so_far: num_rows_scheduled, + decoders: next_scan_line.decoders, + })) { + // Decoder has disconnected + return; + } + + trace!("Finished scheduling {} ranges", ranges.len()); + } + } + + fn do_schedule_ranges( + &mut self, + ranges: &[Range], + filter: &FilterExpression, + io: Arc, + schedule_action: impl FnMut(Result) -> bool, + // If specified, this will be used as the top_level_row for all scheduling + // tasks. This is used by list scheduling to ensure all items scheduling + // tasks are scheduled at the same top level row. + priority: Option>, + ) { + match &self.root_scheduler { + RootScheduler::Array(_) => { + self.do_schedule_ranges_array(ranges, filter, io, schedule_action, priority) + } + RootScheduler::Structural(_) => { + self.do_schedule_ranges_structural(ranges, filter, io, schedule_action) + } + } + } + + // This method is similar to schedule_ranges but instead of + // sending the decoders to a channel it collects them all into a vector + pub fn schedule_ranges_to_vec( + &mut self, + ranges: &[Range], + filter: &FilterExpression, + io: Arc, + priority: Option>, + ) -> Result> { + let mut decode_messages = Vec::new(); + self.do_schedule_ranges( + ranges, + filter, + io, + |msg| { + decode_messages.push(msg); + true + }, + priority, + ); + decode_messages.into_iter().collect::>>() + } + + /// Schedules the load of multiple ranges of rows + /// + /// Ranges must be non-overlapping and in sorted order + /// + /// # Arguments + /// + /// * `ranges` - The ranges of rows to load + /// * `sink` - A channel to send the decode tasks + /// * `scheduler` An I/O scheduler to issue I/O requests + #[instrument(level = "debug", skip_all)] + pub fn schedule_ranges( + &mut self, + ranges: &[Range], + filter: &FilterExpression, + sink: mpsc::UnboundedSender>, + scheduler: Arc, + ) { + self.do_schedule_ranges( + ranges, + filter, + scheduler, + |msg| { + match sink.send(msg) { + Ok(_) => true, + Err(SendError { .. }) => { + // The receiver has gone away. We can't do anything about it + // so just ignore the error. + debug!( + "schedule_ranges aborting early since decoder appears to have been dropped" + ); + false + } + } + }, + None, + ) + } + + /// Schedules the load of a range of rows + /// + /// # Arguments + /// + /// * `range` - The range of rows to load + /// * `sink` - A channel to send the decode tasks + /// * `scheduler` An I/O scheduler to issue I/O requests + #[instrument(level = "debug", skip_all)] + pub fn schedule_range( + &mut self, + range: Range, + filter: &FilterExpression, + sink: mpsc::UnboundedSender>, + scheduler: Arc, + ) { + self.schedule_ranges(&[range], filter, sink, scheduler) + } + + /// Schedules the load of selected rows + /// + /// # Arguments + /// + /// * `indices` - The row indices to load (these must be in ascending order!) + /// * `sink` - A channel to send the decode tasks + /// * `scheduler` An I/O scheduler to issue I/O requests + pub fn schedule_take( + &mut self, + indices: &[u64], + filter: &FilterExpression, + sink: mpsc::UnboundedSender>, + scheduler: Arc, + ) { + if indices.is_empty() { + return; + } + trace!("Scheduling take of {} rows", indices.len()); + let ranges = match RequestedRows::indices_to_ranges(indices) { + Ok(ranges) => ranges, + Err(error) => { + if let Err(SendError { .. }) = sink.send(Err(error)) { + debug!( + "schedule_take could not report invalid indices because the decoder was dropped" + ); + } + return; + } + }; + self.schedule_ranges(&ranges, filter, sink, scheduler) + } +} + +pub struct ReadBatchTask { + pub task: BoxFuture<'static, Result>, + pub num_rows: u32, +} + +/// A stream that takes scheduled jobs and generates decode tasks from them. +pub struct BatchDecodeStream { + context: DecoderContext, + root_decoder: SimpleStructDecoder, + rows_remaining: u64, + rows_per_batch: u32, + rows_scheduled: u64, + rows_drained: u64, + scheduler_exhausted: bool, + emitted_batch_size_warning: Arc, + limit_i32_offset_batch_size: bool, +} + +impl BatchDecodeStream { + /// Create a new instance of a batch decode stream + /// + /// # Arguments + /// + /// * `scheduled` - an incoming stream of decode tasks from a `DecodeBatchScheduler` + /// * `schema` - the schema of the data to create + /// * `rows_per_batch` the number of rows to create before making a batch + /// * `num_rows` the total number of rows scheduled + /// * `num_columns` the total number of columns in the file + pub fn new( + scheduled: mpsc::UnboundedReceiver>, + rows_per_batch: u32, + num_rows: u64, + root_decoder: SimpleStructDecoder, + ) -> Self { + Self::new_with_i32_offset_limit(scheduled, rows_per_batch, num_rows, root_decoder, true) + } + + fn new_with_i32_offset_limit( + scheduled: mpsc::UnboundedReceiver>, + rows_per_batch: u32, + num_rows: u64, + root_decoder: SimpleStructDecoder, + limit_i32_offset_batch_size: bool, + ) -> Self { + Self { + context: DecoderContext::new(scheduled), + root_decoder, + rows_remaining: num_rows, + rows_per_batch, + rows_scheduled: 0, + rows_drained: 0, + scheduler_exhausted: false, + emitted_batch_size_warning: Arc::new(Once::new()), + limit_i32_offset_batch_size, + } + } + + fn accept_decoder(&mut self, decoder: DecoderReady) -> Result<()> { + if decoder.path.is_empty() { + // The root decoder we can ignore + Ok(()) + } else { + self.root_decoder.accept_child(decoder) + } + } + + #[instrument(level = "debug", skip_all)] + async fn wait_for_scheduled(&mut self, scheduled_need: u64) -> Result { + if self.scheduler_exhausted { + return Ok(self.rows_scheduled); + } + while self.rows_scheduled < scheduled_need { + let next_message = self.context.source.recv().await; + match next_message { + Some(scan_line) => { + let scan_line = scan_line?; + self.rows_scheduled = scan_line.scheduled_so_far; + for message in scan_line.decoders { + self.accept_decoder(message.into_array())?; + } + } + None => { + // Schedule ended before we got all the data we expected. This probably + // means some kind of pushdown filter was applied and we didn't load as + // much data as we thought we would. + self.scheduler_exhausted = true; + return Ok(self.rows_scheduled); + } + } + } + Ok(scheduled_need) + } + + #[instrument(level = "debug", skip_all)] + async fn next_batch_task(&mut self) -> Result> { + trace!( + "Draining batch task (rows_remaining={} rows_drained={} rows_scheduled={})", + self.rows_remaining, self.rows_drained, self.rows_scheduled, + ); + if self.rows_remaining == 0 { + return Ok(None); + } + + let mut to_take = self.rows_remaining.min(self.rows_per_batch as u64); + self.rows_remaining -= to_take; + + let scheduled_need = (self.rows_drained + to_take).saturating_sub(self.rows_scheduled); + trace!( + "scheduled_need = {} because rows_drained = {} and to_take = {} and rows_scheduled = {}", + scheduled_need, self.rows_drained, to_take, self.rows_scheduled + ); + if scheduled_need > 0 { + let desired_scheduled = scheduled_need + self.rows_scheduled; + trace!( + "Draining from scheduler (desire at least {} scheduled rows)", + desired_scheduled + ); + let actually_scheduled = self.wait_for_scheduled(desired_scheduled).await?; + if actually_scheduled < desired_scheduled { + let under_scheduled = desired_scheduled - actually_scheduled; + to_take -= under_scheduled; + } + } + + if to_take == 0 { + return Ok(None); + } + + // wait_for_loaded waits for *>* loaded_need (not >=) so we do a -1 here + let loaded_need = self.rows_drained + to_take - 1; + trace!( + "Waiting for I/O (desire at least {} fully loaded rows)", + loaded_need + ); + self.root_decoder.wait_for_loaded(loaded_need).await?; + + if self.limit_i32_offset_batch_size && has_i32_offsets(self.root_decoder.data_type()) { + let limit = LogicalPageDecoder::max_rows_to_drain( + &self.root_decoder, + to_take, + I32_OFFSET_BYTE_BUDGET, + )?; + let clamped = clamp_rows_to_i32_budget(limit, to_take); + self.rows_remaining += to_take - clamped; + to_take = clamped; + } + + let next_task = self.root_decoder.drain(to_take)?; + self.rows_drained += to_take; + Ok(Some(next_task)) + } + + pub fn into_stream(self) -> BoxStream<'static, ReadBatchTask> { + let stream = futures::stream::unfold(self, |mut slf| async move { + let next_task = match slf.next_batch_task().await { + Ok(Some(next_task)) => next_task, + Ok(None) => return None, + Err(err) => { + slf.rows_remaining = 0; + return Some(( + ReadBatchTask { + task: async move { Err(err) }.boxed(), + num_rows: 0, + }, + slf, + )); + } + }; + let num_rows = next_task.num_rows; + let emitted_batch_size_warning = slf.emitted_batch_size_warning.clone(); + let task = async move { + // Real decode work happens inside into_batch, which can block the current + // thread for a long time. By spawning it as a new task, we allow Tokio's + // worker threads to keep making progress. + let (batch, _data_size) = + tokio::spawn(async move { next_task.into_batch(emitted_batch_size_warning) }) + .await + .map_err(|err| Error::wrapped(err.into()))??; + Ok(batch) + }; + // This should be true since batch size is u32 + debug_assert!(num_rows <= u32::MAX as u64); + Some(( + ReadBatchTask { + task: task.boxed(), + num_rows: num_rows as u32, + }, + slf, + )) + }); + stream.boxed() + } +} + +// Utility types to smooth out the differences between the 2.0 and 2.1 decoders so that +// we can have a single implementation of the batch decode iterator +enum RootDecoderMessage { + LoadedPage(LoadedPageShard), + ArrayPage(DecoderReady), +} +trait RootDecoderType { + fn accept_message(&mut self, message: RootDecoderMessage) -> Result<()>; + fn drain_batch(&mut self, num_rows: u64) -> Result; + fn wait(&mut self, loaded_need: u64, runtime: &tokio::runtime::Runtime) -> Result<()>; + /// See [`LogicalPageDecoder::max_rows_to_drain`]. Decoders without i32-offset + /// accounting place no limit. + fn max_rows_to_drain(&self, num_rows: u64, _byte_budget: u64) -> Result { + Ok(DrainLimit { + rows: num_rows, + bytes: 0, + }) + } +} +impl RootDecoderType for StructuralStructDecoder { + fn accept_message(&mut self, message: RootDecoderMessage) -> Result<()> { + let RootDecoderMessage::LoadedPage(loaded_page) = message else { + unreachable!() + }; + self.accept_page(loaded_page) + } + fn drain_batch(&mut self, num_rows: u64) -> Result { + self.drain_batch_task(num_rows) + } + fn wait(&mut self, _: u64, _: &tokio::runtime::Runtime) -> Result<()> { + // Waiting happens elsewhere (not as part of the decoder) + Ok(()) + } +} +impl RootDecoderType for SimpleStructDecoder { + fn accept_message(&mut self, message: RootDecoderMessage) -> Result<()> { + let RootDecoderMessage::ArrayPage(array_page) = message else { + unreachable!() + }; + self.accept_child(array_page) + } + fn drain_batch(&mut self, num_rows: u64) -> Result { + self.drain(num_rows) + } + fn wait(&mut self, loaded_need: u64, runtime: &tokio::runtime::Runtime) -> Result<()> { + runtime.block_on(self.wait_for_loaded(loaded_need)) + } + fn max_rows_to_drain(&self, num_rows: u64, byte_budget: u64) -> Result { + if !has_i32_offsets(LogicalPageDecoder::data_type(self)) { + return Ok(DrainLimit { + rows: num_rows, + bytes: 0, + }); + } + LogicalPageDecoder::max_rows_to_drain(self, num_rows, byte_budget) + } +} + +/// A blocking batch decoder that performs synchronous decoding +struct BatchDecodeIterator { + messages: VecDeque>, + root_decoder: T, + rows_remaining: u64, + rows_per_batch: u32, + rows_scheduled: u64, + rows_drained: u64, + emitted_batch_size_warning: Arc, + // Note: this is not the runtime on which I/O happens. + // That's always in the scheduler. This is just a runtime we use to + // sleep the current thread if I/O is unready + wait_for_io_runtime: tokio::runtime::Runtime, + schema: Arc, +} + +impl BatchDecodeIterator { + /// Create a new instance of a batch decode iterator + pub fn new( + messages: VecDeque>, + rows_per_batch: u32, + num_rows: u64, + root_decoder: T, + schema: Arc, + ) -> Self { + Self { + messages, + root_decoder, + rows_remaining: num_rows, + rows_per_batch, + rows_scheduled: 0, + rows_drained: 0, + wait_for_io_runtime: tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(), + emitted_batch_size_warning: Arc::new(Once::new()), + schema, + } + } + + /// Wait for a single page of data to finish loading + /// + /// If the data is not available this will perform a *blocking* wait (put + /// the current thread to sleep) + fn wait_for_page(&self, unloaded_page: UnloadedPageShard) -> Result { + match maybe_done(unloaded_page.0) { + // Fast path, avoid all runtime shenanigans if the data is ready + MaybeDone::Done(loaded_page) => loaded_page, + // Slow path, we need to wait on I/O, enter the runtime + MaybeDone::Future(fut) => self.wait_for_io_runtime.block_on(fut), + MaybeDone::Gone => unreachable!(), + } + } + + /// Waits for I/O until `scheduled_need` rows have been loaded + /// + /// Note that `scheduled_need` is cumulative. E.g. this method + /// should be called with 5, 10, 15 and not 5, 5, 5 + #[instrument(level = "debug", skip_all)] + fn wait_for_io(&mut self, scheduled_need: u64, to_take: u64) -> Result { + while self.rows_scheduled < scheduled_need && !self.messages.is_empty() { + let message = self.messages.pop_front().unwrap()?; + self.rows_scheduled = message.scheduled_so_far; + for decoder_message in message.decoders { + match decoder_message { + MessageType::UnloadedPage(unloaded_page) => { + let loaded_page = self.wait_for_page(unloaded_page)?; + self.root_decoder + .accept_message(RootDecoderMessage::LoadedPage(loaded_page))?; + } + MessageType::DecoderReady(decoder_ready) => { + // The root decoder we can ignore + if !decoder_ready.path.is_empty() { + self.root_decoder + .accept_message(RootDecoderMessage::ArrayPage(decoder_ready))?; + } + } + } + } + } + + let loaded_need = self.rows_drained + to_take.min(self.rows_per_batch as u64) - 1; + + self.root_decoder + .wait(loaded_need, &self.wait_for_io_runtime)?; + Ok(self.rows_scheduled) + } + + #[instrument(level = "debug", skip_all)] + fn next_batch_task(&mut self) -> Result> { + trace!( + "Draining batch task (rows_remaining={} rows_drained={} rows_scheduled={})", + self.rows_remaining, self.rows_drained, self.rows_scheduled, + ); + if self.rows_remaining == 0 { + return Ok(None); + } + + let mut to_take = self.rows_remaining.min(self.rows_per_batch as u64); + self.rows_remaining -= to_take; + + let scheduled_need = (self.rows_drained + to_take).saturating_sub(self.rows_scheduled); + trace!( + "scheduled_need = {} because rows_drained = {} and to_take = {} and rows_scheduled = {}", + scheduled_need, self.rows_drained, to_take, self.rows_scheduled + ); + if scheduled_need > 0 { + let desired_scheduled = scheduled_need + self.rows_scheduled; + trace!( + "Draining from scheduler (desire at least {} scheduled rows)", + desired_scheduled + ); + let actually_scheduled = self.wait_for_io(desired_scheduled, to_take)?; + if actually_scheduled < desired_scheduled { + let under_scheduled = desired_scheduled - actually_scheduled; + to_take -= under_scheduled; + } + } + + if to_take == 0 { + return Ok(None); + } + + let limit = self + .root_decoder + .max_rows_to_drain(to_take, I32_OFFSET_BYTE_BUDGET)?; + let clamped = clamp_rows_to_i32_budget(limit, to_take); + self.rows_remaining += to_take - clamped; + to_take = clamped; + + let next_task = self.root_decoder.drain_batch(to_take)?; + + self.rows_drained += to_take; + + let (batch, _data_size) = next_task.into_batch(self.emitted_batch_size_warning.clone())?; + + Ok(Some(batch)) + } +} + +impl Iterator for BatchDecodeIterator { + type Item = ArrowResult; + + fn next(&mut self) -> Option { + let next = self.next_batch_task().transpose(); + if let Some(Err(_)) = &next { + // Fuse after an error instead of re-attempting every remaining batch + self.rows_remaining = 0; + } + next.map(|r| r.map_err(ArrowError::from)) + } +} + +impl RecordBatchReader for BatchDecodeIterator { + fn schema(&self) -> Arc { + self.schema.clone() + } +} + +/// Estimate the number of bytes per row for a given Arrow data type. +/// +/// For fixed-width types this is exact. For variable-width types (strings, +/// binary, lists) a rough default is used. The estimate is used as a +/// starting point when `batch_size_bytes` is set; a post-decode feedback +/// loop corrects it after the first batch. +/// +/// This estimate ignores validity bitmaps at the moment. We can't infer +/// their presence simply from the data_type and their impact is probably +/// fairly negligible. +/// Returns a schema-based estimate of the decoded bytes per row for `data_type`. +/// +/// Fixed-width types are exact. Variable-width types (strings, lists, etc.) use +/// heuristic constants. This estimate is used both in batch-size planning and as +/// a fallback for V1 files that lack structural decoders. +pub fn estimate_bytes_per_row(data_type: &DataType) -> f64 { + if let Some(w) = data_type.byte_width_opt() { + return w as f64; + } + match data_type { + DataType::Boolean => 1.0 / 8.0, + DataType::Utf8 | DataType::Binary | DataType::LargeUtf8 | DataType::LargeBinary => 64.0, + DataType::Struct(fields) => fields + .iter() + .map(|f| estimate_bytes_per_row(f.data_type())) + .sum(), + DataType::List(child) | DataType::LargeList(child) => { + 5.0 * estimate_bytes_per_row(child.data_type()) + } + DataType::FixedSizeList(child, dim) => { + *dim as f64 * estimate_bytes_per_row(child.data_type()) + } + DataType::Dictionary(_, value_type) => estimate_bytes_per_row(value_type), + DataType::Map(entries, _) => 5.0 * estimate_bytes_per_row(entries.data_type()), + _ => 64.0, + } +} + +/// A stream that takes scheduled jobs and generates decode tasks from them. +pub struct StructuralBatchDecodeStream { + context: DecoderContext, + root_decoder: StructuralStructDecoder, + rows_remaining: u64, + rows_per_batch: u32, + rows_scheduled: u64, + rows_drained: u64, + scheduler_exhausted: bool, + emitted_batch_size_warning: Arc, + // Decode scheduling policy selected at planning time. + // + // Performance tradeoff: + // - true: spawn `into_batch` onto Tokio, which improves scan throughput by allowing + // more decode parallelism. + // - false: run `into_batch` inline, which avoids Tokio scheduling overhead and is + // typically better for point lookups / small takes. + spawn_batch_decode_tasks: bool, + /// If set, target this many bytes per batch while retaining `rows_per_batch` + /// as an independent upper bound. + batch_size_bytes: Option, + /// Schema-based estimate of bytes per row, computed once at construction. + /// Only meaningful when `batch_size_bytes` is `Some`. + schema_bytes_per_row: f64, + /// Post-decode feedback: actual bytes-per-row measured from the most + /// recently decoded batch. Zero means no feedback yet (use schema estimate). + bytes_per_row_feedback: Arc, +} + +impl StructuralBatchDecodeStream { + /// Create a new instance of a batch decode stream + /// + /// # Arguments + /// + /// * `scheduled` - an incoming stream of decode tasks from a `DecodeBatchScheduler` + /// * `schema` - the schema of the data to create + /// * `rows_per_batch` the number of rows to create before making a batch + /// * `num_rows` the total number of rows scheduled + /// * `num_columns` the total number of columns in the file + pub fn new( + scheduled: mpsc::UnboundedReceiver>, + rows_per_batch: u32, + num_rows: u64, + root_decoder: StructuralStructDecoder, + spawn_batch_decode_tasks: bool, + batch_size_bytes: Option, + ) -> Self { + let schema_bytes_per_row = if batch_size_bytes.is_some() { + estimate_bytes_per_row(root_decoder.data_type()).max(1.0) + } else { + 0.0 + }; + Self { + context: DecoderContext::new(scheduled), + root_decoder, + rows_remaining: num_rows, + rows_per_batch, + rows_scheduled: 0, + rows_drained: 0, + scheduler_exhausted: false, + emitted_batch_size_warning: Arc::new(Once::new()), + spawn_batch_decode_tasks, + batch_size_bytes, + schema_bytes_per_row, + bytes_per_row_feedback: Arc::new(AtomicU64::new(0)), + } + } + + #[instrument(level = "debug", skip_all)] + async fn wait_for_scheduled(&mut self, scheduled_need: u64) -> Result { + if self.scheduler_exhausted { + return Ok(self.rows_scheduled); + } + while self.rows_scheduled < scheduled_need { + let next_message = self.context.source.recv().await; + match next_message { + Some(scan_line) => { + let scan_line = scan_line?; + self.rows_scheduled = scan_line.scheduled_so_far; + for message in scan_line.decoders { + let unloaded_page = message.into_structural(); + let loaded_page = unloaded_page.0.await?; + self.root_decoder.accept_page(loaded_page)?; + } + } + None => { + // Schedule ended before we got all the data we expected. This probably + // means some kind of pushdown filter was applied and we didn't load as + // much data as we thought we would. + self.scheduler_exhausted = true; + return Ok(self.rows_scheduled); + } + } + } + Ok(scheduled_need) + } + + #[instrument(level = "debug", skip_all)] + async fn next_batch_task(&mut self) -> Result> { + trace!( + "Draining batch task (rows_remaining={} rows_drained={} rows_scheduled={})", + self.rows_remaining, self.rows_drained, self.rows_scheduled, + ); + if self.rows_remaining == 0 { + return Ok(None); + } + + let row_limit = self.rows_remaining.min(self.rows_per_batch as u64); + let mut to_take = if let Some(batch_size_bytes) = self.batch_size_bytes { + let feedback = self.bytes_per_row_feedback.load(Ordering::Relaxed); + let bpr = if feedback > 0 { + feedback as f64 + } else { + self.schema_bytes_per_row + }; + let rows = (batch_size_bytes as f64 / bpr) as u64; + row_limit.min(rows.max(1)) + } else { + row_limit + }; + self.rows_remaining -= to_take; + + let scheduled_need = (self.rows_drained + to_take).saturating_sub(self.rows_scheduled); + trace!( + "scheduled_need = {} because rows_drained = {} and to_take = {} and rows_scheduled = {}", + scheduled_need, self.rows_drained, to_take, self.rows_scheduled + ); + if scheduled_need > 0 { + let desired_scheduled = scheduled_need + self.rows_scheduled; + trace!( + "Draining from scheduler (desire at least {} scheduled rows)", + desired_scheduled + ); + let actually_scheduled = self.wait_for_scheduled(desired_scheduled).await?; + if actually_scheduled < desired_scheduled { + let under_scheduled = desired_scheduled - actually_scheduled; + to_take -= under_scheduled; + } + } + + if to_take == 0 { + return Ok(None); + } + + let next_task = self.root_decoder.drain_batch_task(to_take)?; + self.rows_drained += to_take; + Ok(Some(next_task)) + } + + pub fn into_stream(self) -> BoxStream<'static, ReadBatchTask> { + let stream = futures::stream::unfold(self, |mut slf| async move { + let next_task = match slf.next_batch_task().await { + Ok(Some(next_task)) => next_task, + Ok(None) => return None, + Err(err) => { + slf.rows_remaining = 0; + return Some(( + ReadBatchTask { + task: async move { Err(err) }.boxed(), + num_rows: 0, + }, + slf, + )); + } + }; + let num_rows = next_task.num_rows; + let emitted_batch_size_warning = slf.emitted_batch_size_warning.clone(); + let bytes_per_row_feedback = slf.bytes_per_row_feedback.clone(); + // Capture the per-stream policy once so every emitted batch task follows the + // same throughput-vs-overhead choice made by the scheduler. + let spawn_batch_decode_tasks = slf.spawn_batch_decode_tasks; + let task = async move { + let (batch, data_size) = if spawn_batch_decode_tasks { + tokio::spawn(async move { next_task.into_batch(emitted_batch_size_warning) }) + .await + .map_err(|err| Error::wrapped(err.into()))?? + } else { + next_task.into_batch(emitted_batch_size_warning)? + }; + let num_rows = batch.num_rows() as u64; + if let Some(bpr) = data_size.checked_div(num_rows) { + let prev = bytes_per_row_feedback.load(Ordering::Relaxed); + let next = if prev == 0 || bpr >= prev { + // First batch or actual size is larger than estimate: + // adopt immediately to avoid OOM. + bpr + } else { + // Actual size is smaller: degrade gradually toward + // the true value to avoid over-correcting on a + // single anomalous batch. + (prev + bpr) / 2 + }; + bytes_per_row_feedback.store(next.max(1), Ordering::Relaxed); + } + Ok(batch) + }; + // This should be true since batch size is u32 + debug_assert!(num_rows <= u32::MAX as u64); + Some(( + ReadBatchTask { + task: task.boxed(), + num_rows: num_rows as u32, + }, + slf, + )) + }); + stream.boxed() + } +} + +#[derive(Debug)] +pub enum RequestedRows { + Ranges(Vec>), + Indices(Vec), +} + +impl RequestedRows { + pub fn num_rows(&self) -> u64 { + match self { + Self::Ranges(ranges) => ranges.iter().map(|r| r.end - r.start).sum(), + Self::Indices(indices) => indices.len() as u64, + } + } + + pub fn trim_empty_ranges(mut self) -> Self { + if let Self::Ranges(ranges) = &mut self { + ranges.retain(|r| !r.is_empty()); + } + self + } + + fn into_ranges(self) -> Result>> { + match self { + Self::Ranges(ranges) => Ok(ranges), + Self::Indices(indices) => Self::indices_to_ranges(&indices), + } + } + + // coalesce continuous indices if possible (the input indices must be sorted) + fn indices_to_ranges(indices: &[u64]) -> Result>> { + if indices.is_empty() { + return Ok(Vec::new()); + } + + let mut ranges = Vec::with_capacity(indices.len()); + let mut start = indices[0]; + + for (index, pair) in indices.windows(2).enumerate() { + if pair[0] >= pair[1] { + return Err(Error::invalid_input(format!( + "Requested row indices are not strictly increasing at index {index}: {} then {}", + pair[0], pair[1] + ))); + } + let end = pair[0] + .checked_add(1) + .ok_or_else(|| Error::invalid_input("Requested row index cannot equal u64::MAX"))?; + if pair[1] != end { + ranges.push(start..end); + start = pair[1]; + } + } + + let end = indices + .last() + .copied() + .and_then(|index| index.checked_add(1)) + .ok_or_else(|| Error::invalid_input("Requested row index cannot equal u64::MAX"))?; + ranges.push(start..end); + Ok(ranges) + } +} + +/// Configuration for decoder behavior +#[derive(Debug, Clone)] +pub struct DecoderConfig { + /// Whether to cache repetition indices for better performance. + /// + /// This defaults to the `LANCE_READ_CACHE_REPETITION_INDEX` environment variable + /// when present and is enabled by default. Set the env var to a non-truthy + /// value (for example `0` or `false`) to disable it. The env var is read + /// once per process. + pub cache_repetition_index: bool, + /// Whether to validate decoded data + pub validate_on_decode: bool, + /// Override the strategy used to dispatch the scheduling work in + /// [`schedule_and_decode`]. + /// + /// `schedule_and_decode` always awaits the scheduler's `initialize` (which + /// performs metadata I/O) before returning. This flag controls what + /// happens with the subsequent (synchronous) work of pushing decoder + /// messages into the channel that feeds the decode stream. + /// + /// * `None` - default behavior: the scheduling work runs inline (as part + /// of the `schedule_and_decode` await) when the request is small + /// (controlled by the `LANCE_INLINE_SCHEDULING_THRESHOLD` env var) and + /// is dispatched onto a spawned task otherwise. + /// * `Some(true)` - always run scheduling inline. The await of + /// `schedule_and_decode` does not return until every decoder message + /// has been queued. + /// * `Some(false)` - always spawn a task for scheduling so that it can + /// overlap with consumption of the decode stream. + pub inline_scheduling: Option, +} + +impl Default for DecoderConfig { + fn default() -> Self { + Self { + cache_repetition_index: default_cache_repetition_index(), + validate_on_decode: false, + inline_scheduling: None, + } + } +} + +#[derive(Debug, Clone)] +pub struct SchedulerDecoderConfig { + pub decoder_plugins: Arc, + pub batch_size: u32, + pub io: Arc, + pub cache: Arc, + /// Decoder configuration + pub decoder_config: DecoderConfig, + /// If set, target this many bytes per batch while retaining `batch_size` as + /// an independent row-count upper bound. + /// + /// Only supported for v2.1+ (structural) files. For v2.0 files this + /// option is ignored and a warning is logged. + pub batch_size_bytes: Option, +} + +fn check_scheduler_on_drop( + stream: BoxStream<'static, ReadBatchTask>, + scheduler_handle: tokio::task::JoinHandle<()>, +) -> BoxStream<'static, ReadBatchTask> { + // This is a bit weird but we create an "empty stream" that unwraps the scheduler handle (which + // will panic if the scheduler panicked). This let's us check if the scheduler panicked + // when the stream finishes. + let abort_handle = scheduler_handle.abort_handle(); + let mut scheduler_handle = Some(scheduler_handle); + let check_scheduler = stream::unfold((), move |_| { + let handle = scheduler_handle.take(); + async move { + if let Some(handle) = handle { + handle.await.unwrap(); + } + None + } + }); + stream + .chain(check_scheduler) + .on_drop(move || { + // Abort the scheduler task on early drop. The scheduler task holds + // a reference to the I/O scheduler (via config.io) which keeps the + // ScanScheduler alive. If the scheduler task is stuck waiting for + // initialization I/O (which is blocked on backpressure that will + // never drain because no one is consuming the stream), we need to + // abort it so it releases its I/O reference and allows the + // ScanScheduler to drop and cancel pending I/O. + abort_handle.abort(); + }) + .boxed() +} + +#[allow(clippy::too_many_arguments)] +pub fn create_decode_stream( + schema: &Schema, + num_rows: u64, + batch_size: u32, + is_structural: bool, + should_validate: bool, + spawn_structural_batch_decode_tasks: bool, + rx: mpsc::UnboundedReceiver>, + batch_size_bytes: Option, +) -> Result> { + create_decode_stream_with_i32_offset_limit( + schema, + num_rows, + batch_size, + is_structural, + should_validate, + spawn_structural_batch_decode_tasks, + rx, + batch_size_bytes, + true, + ) +} + +#[allow(clippy::too_many_arguments)] +fn create_decode_stream_with_i32_offset_limit( + schema: &Schema, + num_rows: u64, + batch_size: u32, + is_structural: bool, + should_validate: bool, + spawn_structural_batch_decode_tasks: bool, + rx: mpsc::UnboundedReceiver>, + batch_size_bytes: Option, + limit_i32_offset_batch_size: bool, +) -> Result> { + if is_structural { + let arrow_schema = ArrowSchema::from(schema); + let structural_decoder = StructuralStructDecoder::new( + arrow_schema.fields, + should_validate, + /*is_root=*/ true, + /*nullable=*/ false, + )?; + Ok(StructuralBatchDecodeStream::new( + rx, + batch_size, + num_rows, + structural_decoder, + spawn_structural_batch_decode_tasks, + batch_size_bytes, + ) + .into_stream()) + } else { + if batch_size_bytes.is_some() { + warn!("batch_size_bytes is not supported for v2.0 files and will be ignored"); + } + let arrow_schema = ArrowSchema::from(schema); + let root_fields = arrow_schema.fields; + + let simple_struct_decoder = SimpleStructDecoder::new(root_fields, num_rows); + Ok(BatchDecodeStream::new_with_i32_offset_limit( + rx, + batch_size, + num_rows, + simple_struct_decoder, + limit_i32_offset_batch_size, + ) + .into_stream()) + } +} + +/// Creates a iterator that decodes a set of messages in a blocking fashion +/// +/// See [`schedule_and_decode_blocking`] for more information. +pub fn create_decode_iterator( + schema: &Schema, + num_rows: u64, + batch_size: u32, + should_validate: bool, + is_structural: bool, + messages: VecDeque>, +) -> Result> { + let arrow_schema = Arc::new(ArrowSchema::from(schema)); + let root_fields = arrow_schema.fields.clone(); + if is_structural { + let simple_struct_decoder = StructuralStructDecoder::new( + root_fields, + should_validate, + /*is_root=*/ true, + /*nullable=*/ false, + )?; + Ok(Box::new(BatchDecodeIterator::new( + messages, + batch_size, + num_rows, + simple_struct_decoder, + arrow_schema, + ))) + } else { + let root_decoder = SimpleStructDecoder::new(root_fields, num_rows); + Ok(Box::new(BatchDecodeIterator::new( + messages, + batch_size, + num_rows, + root_decoder, + arrow_schema, + ))) + } +} + +async fn create_scheduler_decoder( + column_infos: Vec>, + requested_rows: RequestedRows, + filter: FilterExpression, + column_indices: Vec, + target_schema: Arc, + config: SchedulerDecoderConfig, +) -> Result> { + let num_rows = requested_rows.num_rows(); + let is_range_request = matches!(&requested_rows, RequestedRows::Ranges(_)); + + let is_structural = column_infos[0].is_structural(); + let mode = std::env::var(ENV_LANCE_STRUCTURAL_BATCH_DECODE_SPAWN_MODE); + let spawn_structural_batch_decode_tasks = match mode.ok().as_deref() { + Some("always") => true, + Some("never") => false, + _ => is_range_request, + }; + + let (tx, rx) = mpsc::unbounded_channel(); + + let decode_stream = create_decode_stream( + &target_schema, + num_rows, + config.batch_size, + is_structural, + config.decoder_config.validate_on_decode, + spawn_structural_batch_decode_tasks, + rx, + config.batch_size_bytes, + )?; + + // The scheduler's `initialize` may perform I/O to load column metadata + // unless that metadata is already in the cache. This metadata loading + // happens as part of this call and should be parallelized if reading + // multiple files. We pass the rows we are about to schedule so that the + // structural path only initializes the pages those rows touch. + let requested_ranges = requested_rows.into_ranges()?; + let mut decode_scheduler = DecodeBatchScheduler::try_new_with_ranges( + target_schema.as_ref(), + &column_indices, + &column_infos, + &vec![], + num_rows, + config.decoder_plugins, + config.io.clone(), + config.cache, + Some(&requested_ranges), + &filter, + &config.decoder_config, + ) + .await?; + + // For small requests the scheduling cost is dwarfed by the overhead of + // spawning a task, so we run scheduling inline (still as part of this + // await) before returning. The threshold is configurable via + // `LANCE_INLINE_SCHEDULING_THRESHOLD`, and callers can force either + // strategy via `DecoderConfig::inline_scheduling`. + let inline_scheduling = config + .decoder_config + .inline_scheduling + .unwrap_or_else(|| num_rows <= inline_scheduling_threshold()); + + if inline_scheduling { + decode_scheduler.schedule_ranges(&requested_ranges, &filter, tx, config.io); + Ok(decode_stream) + } else { + // Spawn the (still synchronous) scheduling work so that decoder + // messages can stream into the channel while the consumer is + // already pulling from the decode stream. + let scheduling = async move { + decode_scheduler.schedule_ranges(&requested_ranges, &filter, tx, config.io) + }; + let scheduler_handle = tokio::task::spawn(scheduling); + Ok(check_scheduler_on_drop(decode_stream, scheduler_handle)) + } +} + +/// Initializes the scheduler, schedules the requested rows, and returns a +/// stream of decode tasks for the resulting batches. +/// +/// This is a convenience function that creates both the scheduler and the +/// decoder, which can be a little tricky to get right. +/// +/// # Why is this async? +/// +/// Constructing the scheduler runs `initialize` which will perform I/O +/// unless the data required is already in the file metadata cache. +/// +/// When `DecoderConfig::inline_scheduling` resolves to `true`, the +/// subsequent (synchronous) scheduling work also runs before this function +/// returns, leaving a fully primed decode stream. +pub async fn schedule_and_decode( + column_infos: Vec>, + requested_rows: RequestedRows, + filter: FilterExpression, + column_indices: Vec, + target_schema: Arc, + config: SchedulerDecoderConfig, +) -> Result> { + if requested_rows.num_rows() == 0 { + return Ok(stream::empty().boxed()); + } + + // If the user requested any ranges that are empty, ignore them. They are pointless and + // trying to read them has caused bugs in the past. + let requested_rows = requested_rows.trim_empty_ranges(); + + let io = config.io.clone(); + + let stream = create_scheduler_decoder( + column_infos, + requested_rows, + filter, + column_indices, + target_schema, + config, + ) + .await?; + + // Keep the io alive until the stream is dropped or finishes. Otherwise the + // I/O drops as soon as the scheduling is finished and the I/O loop terminates. + Ok(stream.finally(move || drop(io)).boxed()) +} + +pub static WAITER_RT: LazyLock = LazyLock::new(|| { + tokio::runtime::Builder::new_current_thread() + .build() + .unwrap() +}); + +/// Schedules and decodes the requested data in a blocking fashion +/// +/// This function is a blocking version of [`schedule_and_decode`]. It schedules the requested data +/// and decodes it in the current thread. +/// +/// This can be useful when the disk is fast (or the data is in memory) and the amount +/// of data is relatively small. For example, when doing a take against NVMe or in-memory data. +/// +/// This should NOT be used for full scans. Even if the data is in memory this function will +/// not parallelize the decode and will be slower than the async version. Full scans typically +/// make relatively few IOPs and so the asynchronous overhead is much smaller. +/// +/// This method will first completely run the scheduling process. Then it will run the +/// decode process. +pub fn schedule_and_decode_blocking( + column_infos: Vec>, + requested_rows: RequestedRows, + filter: FilterExpression, + column_indices: Vec, + target_schema: Arc, + config: SchedulerDecoderConfig, +) -> Result> { + if requested_rows.num_rows() == 0 { + let arrow_schema = Arc::new(ArrowSchema::from(target_schema.as_ref())); + return Ok(Box::new(RecordBatchIterator::new(vec![], arrow_schema))); + } + + let requested_rows = requested_rows.trim_empty_ranges(); + let num_rows = requested_rows.num_rows(); + let is_structural = column_infos[0].is_structural(); + + let (tx, mut rx) = mpsc::unbounded_channel(); + + // Initialize the scheduler. This is still "asynchronous" but we run it with a current-thread + // runtime. Pass the rows we are about to schedule so the structural path + // only initializes the pages those rows touch. + let requested_ranges = requested_rows.into_ranges()?; + let mut decode_scheduler = WAITER_RT.block_on(DecodeBatchScheduler::try_new_with_ranges( + target_schema.as_ref(), + &column_indices, + &column_infos, + &vec![], + num_rows, + config.decoder_plugins, + config.io.clone(), + config.cache, + Some(&requested_ranges), + &filter, + &config.decoder_config, + ))?; + + // Schedule the requested rows + decode_scheduler.schedule_ranges(&requested_ranges, &filter, tx, config.io); + + // Drain the scheduler queue into a vec of decode messages + let mut messages = Vec::new(); + while rx + .recv_many(&mut messages, usize::MAX) + .now_or_never() + .unwrap() + != 0 + {} + + // Create a decoder to decode the messages + let decode_iterator = create_decode_iterator( + &target_schema, + num_rows, + config.batch_size, + config.decoder_config.validate_on_decode, + is_structural, + messages.into(), + )?; + + Ok(decode_iterator) +} + +/// A decoder for single-column encodings of primitive data (this includes fixed size +/// lists of primitive data) +/// +/// Physical decoders are able to decode into existing buffers for zero-copy operation. +/// +/// Instances should be stateless and `Send` / `Sync`. This is because multiple decode +/// tasks could reference the same page. For example, imagine a page covers rows 0-2000 +/// and the decoder stream has a batch size of 1024. The decoder will be needed by both +/// the decode task for batch 0 and the decode task for batch 1. +/// +/// See [`crate::decoder`] for more information +pub trait PrimitivePageDecoder: Send + Sync { + /// Decode data into buffers + /// + /// This may be a simple zero-copy from a disk buffer or could involve complex decoding + /// such as decompressing from some compressed representation. + /// + /// Capacity is stored as a tuple of (num_bytes: u64, is_needed: bool). The `is_needed` + /// portion only needs to be updated if the encoding has some concept of an "optional" + /// buffer. + /// + /// Encodings can have any number of input or output buffers. For example, a dictionary + /// decoding will convert two buffers (indices + dictionary) into a single buffer + /// + /// Binary decodings have two output buffers (one for values, one for offsets) + /// + /// Other decodings could even expand the # of output buffers. For example, we could decode + /// fixed size strings into variable length strings going from one input buffer to multiple output + /// buffers. + /// + /// Each Arrow data type typically has a fixed structure of buffers and the encoding chain will + /// generally end at one of these structures. However, intermediate structures may exist which + /// do not correspond to any Arrow type at all. For example, a bitpacking encoding will deal + /// with buffers that have bits-per-value that is not a multiple of 8. + /// + /// The `primitive_array_from_buffers` method has an expected buffer layout for each arrow + /// type (order matters) and encodings that aim to decode into arrow types should respect + /// this layout. + /// # Arguments + /// + /// * `rows_to_skip` - how many rows to skip (within the page) before decoding + /// * `num_rows` - how many rows to decode + /// * `all_null` - A mutable bool, set to true if a decoder determines all values are null + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result; + + /// Returns the variable-width **value** bytes (the bytes an offsets buffer + /// indexes into) that decoding `num_rows` rows starting at `rows_to_skip` + /// will produce, or `None` when the encoding cannot report sizes without + /// decoding. An upper bound is acceptable; callers use this to keep i32 + /// offset values within capacity, so over-estimating only splits batches + /// earlier while under-estimating would let them overflow. + fn variable_width_bytes(&self, _rows_to_skip: u64, _num_rows: u64) -> Result> { + Ok(None) + } +} + +/// A scheduler for single-column encodings of primitive data +/// +/// The scheduler is responsible for calculating what I/O is needed for the requested rows +/// +/// Instances should be stateless and `Send` and `Sync`. This is because instances can +/// be shared in follow-up I/O tasks. +/// +/// See [`crate::decoder`] for more information +pub trait PageScheduler: Send + Sync + std::fmt::Debug { + /// Schedules a batch of I/O to load the data needed for the requested ranges + /// + /// Returns a future that will yield a decoder once the data has been loaded + /// + /// # Arguments + /// + /// * `range` - the range of row offsets (relative to start of page) requested + /// these must be ordered and must not overlap + /// * `scheduler` - a scheduler to submit the I/O request to + /// * `top_level_row` - the row offset of the top level field currently being + /// scheduled. This can be used to assign priority to I/O requests + fn schedule_ranges( + &self, + ranges: &[Range], + scheduler: &Arc, + top_level_row: u64, + ) -> BoxFuture<'static, Result>>; +} + +/// A trait to control the priority of I/O +pub trait PriorityRange: std::fmt::Debug + Send + Sync { + fn advance(&mut self, num_rows: u64); + fn current_priority(&self) -> u64; + fn box_clone(&self) -> Box; +} + +/// A simple priority scheme for top-level fields with no parent +/// repetition +#[derive(Debug)] +pub struct SimplePriorityRange { + priority: u64, +} + +impl SimplePriorityRange { + fn new(priority: u64) -> Self { + Self { priority } + } +} + +impl PriorityRange for SimplePriorityRange { + fn advance(&mut self, num_rows: u64) { + self.priority += num_rows; + } + + fn current_priority(&self) -> u64 { + self.priority + } + + fn box_clone(&self) -> Box { + Box::new(Self { + priority: self.priority, + }) + } +} + +/// Determining the priority of a list request is tricky. We want +/// the priority to be the top-level row. So if we have a +/// `list>` and each outer list has 10 rows and each inner +/// list has 5 rows then the priority of the 100th item is 1 because +/// it is the 5th item in the 10th item of the *second* row. +/// +/// This structure allows us to keep track of this complicated priority +/// relationship. +/// +/// There's a fair amount of bookkeeping involved here. +/// +/// A better approach (using repetition levels) is coming in the future. +pub struct ListPriorityRange { + base: Box, + offsets: Arc<[u64]>, + cur_index_into_offsets: usize, + cur_position: u64, +} + +impl ListPriorityRange { + pub(crate) fn new(base: Box, offsets: Arc<[u64]>) -> Self { + Self { + base, + offsets, + cur_index_into_offsets: 0, + cur_position: 0, + } + } +} + +impl std::fmt::Debug for ListPriorityRange { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ListPriorityRange") + .field("base", &self.base) + .field("offsets.len()", &self.offsets.len()) + .field("cur_index_into_offsets", &self.cur_index_into_offsets) + .field("cur_position", &self.cur_position) + .finish() + } +} + +impl PriorityRange for ListPriorityRange { + fn advance(&mut self, num_rows: u64) { + // We've scheduled X items. Now walk through the offsets to + // determine how many rows we've scheduled. + self.cur_position += num_rows; + let mut idx_into_offsets = self.cur_index_into_offsets; + while idx_into_offsets + 1 < self.offsets.len() + && self.offsets[idx_into_offsets + 1] <= self.cur_position + { + idx_into_offsets += 1; + } + let base_rows_advanced = idx_into_offsets - self.cur_index_into_offsets; + self.cur_index_into_offsets = idx_into_offsets; + self.base.advance(base_rows_advanced as u64); + } + + fn current_priority(&self) -> u64 { + self.base.current_priority() + } + + fn box_clone(&self) -> Box { + Box::new(Self { + base: self.base.box_clone(), + offsets: self.offsets.clone(), + cur_index_into_offsets: self.cur_index_into_offsets, + cur_position: self.cur_position, + }) + } +} + +/// Contains the context for a scheduler +/// The writer accumulates about this many bytes of values per column before it +/// cuts pages (`EncodingOptions::cache_bytes_per_column`), and the rep/def +/// budget of a mini-block chunk then splits that block into many small pages +/// for nested columns. Batching reads back up to this size restores one +/// I/O request per accumulated block. +pub(crate) const IO_REQUEST_BATCH_BYTES: u64 = 8 * 1024 * 1024; + +type SharedRead = Shared>, CloneableError>>>; + +/// What a pending read receives once its batch is submitted. +enum SubmittedRead { + /// The read was alone in its batch and owns the request. + Alone(BoxFuture<'static, Result>>), + /// The batched request and the position of each of the read's ranges in it. + Shared(SharedRead, Vec), +} + +struct PendingRead { + ranges: Vec>, + priority: u64, + tx: oneshot::Sender, +} + +/// Collects the reads submitted while one scheduling step runs and sends them +/// to the I/O layer as a single request. +/// +/// The I/O scheduler only coalesces the ranges of one request, so the reads a +/// step issues for different pages, or for different columns, cost one request +/// each even when they sit next to each other in the file. A take of a few rows +/// reads a few bytes per column, and a file's small pages are packed together, +/// so batching them turns dozens of requests into a handful. +/// +/// A batch is submitted by [`Self::flush`] when the step ends, as soon as the +/// queued reads reach [`IO_REQUEST_BATCH_BYTES`] (a page that shards its own +/// reads to bound buffering keeps every shard in a bounded read), or when a +/// caller awaits one of its reads, so a read queued outside a step (indirect +/// I/O issued from a load future) never waits for a flush that does not come. +/// That last flush waits for one yield first: sibling futures polled in the +/// same pass (a struct's columns initializing their pages, a scan line's list +/// columns fetching their items) all queue their reads before the first of +/// them resumes and submits the lot. +pub(crate) struct RequestBatch { + inner: Arc, + pending: Arc>>, +} + +impl std::fmt::Debug for RequestBatch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RequestBatch").finish_non_exhaustive() + } +} + +impl RequestBatch { + pub(crate) fn new(inner: Arc) -> Self { + Self { + inner, + pending: Arc::new(Mutex::new(Vec::new())), + } + } + + fn queued_bytes(pending: &[PendingRead]) -> u64 { + pending + .iter() + .flat_map(|read| &read.ranges) + .map(|range| range.end - range.start) + .sum() + } + + /// Bytes requested by the reads collected so far. + pub(crate) fn pending_bytes(&self) -> u64 { + Self::queued_bytes(&self.pending.lock().unwrap()) + } + + /// Submits the collected reads as one request. + pub(crate) fn flush(&self) { + Self::flush_pending(&self.inner, &self.pending); + } + + fn flush_pending(inner: &Arc, pending: &Mutex>) { + let pending = std::mem::take(&mut *pending.lock().unwrap()); + Self::submit_batch(inner, pending); + } + + /// Submits `pending` as one request, sorted by file offset, and hands each + /// read its own slice of the result. + fn submit_batch(inner: &Arc, pending: Vec) { + if pending.is_empty() { + return; + } + if pending.len() == 1 { + let read = pending.into_iter().next().unwrap(); + let _ = read.tx.send(SubmittedRead::Alone( + inner.submit_request(read.ranges, read.priority), + )); + return; + } + // The lowest row number any of the batched reads delivers data for. + let priority = pending.iter().map(|read| read.priority).min().unwrap(); + let mut ordered = pending + .iter() + .enumerate() + .flat_map(|(read_idx, read)| { + read.ranges + .iter() + .enumerate() + .map(move |(range_idx, range)| (range.clone(), read_idx, range_idx)) + }) + .collect::>(); + ordered.sort_by_key(|(range, _, _)| (range.start, range.end)); + let ranges = ordered + .iter() + .map(|(range, _, _)| range.clone()) + .collect::>(); + // For every read, where each of its ranges landed in the sorted request. + let mut positions = pending + .iter() + .map(|read| vec![0; read.ranges.len()]) + .collect::>(); + for (position, (_, read_idx, range_idx)) in ordered.iter().enumerate() { + positions[*read_idx][*range_idx] = position; + } + let batched: SharedRead = inner + .submit_request(ranges, priority) + .map(|result| result.map(Arc::new).map_err(CloneableError)) + .boxed() + .shared(); + for (read, read_positions) in pending.into_iter().zip(positions) { + // The receiver is gone only when the read's future was dropped. + let _ = read + .tx + .send(SubmittedRead::Shared(batched.clone(), read_positions)); + } + } +} + +impl EncodingsIo for RequestBatch { + fn submit_request( + &self, + ranges: Vec>, + priority: u64, + ) -> BoxFuture<'static, Result>> { + let mut pending = self.pending.lock().unwrap(); + // Keep every batch within the budget: a page that shards its reads + // must not have all its shards collapse into one unbounded read. + let bytes = ranges + .iter() + .map(|range| range.end - range.start) + .sum::(); + if !pending.is_empty() && Self::queued_bytes(&pending) + bytes > IO_REQUEST_BATCH_BYTES { + Self::submit_batch(&self.inner, std::mem::take(&mut *pending)); + } + let (tx, mut rx) = oneshot::channel(); + pending.push(PendingRead { + ranges, + priority, + tx, + }); + drop(pending); + let inner = self.inner.clone(); + let pending = self.pending.clone(); + async move { + let submitted = match rx.try_recv() { + // The step that queued this read already flushed it. + Ok(Some(submitted)) => submitted, + _ => { + // Awaiting a read is the end of whatever step queued it. + // Yield once so the siblings polled in the same pass get + // to queue too, then submit them all. + tokio::task::yield_now().await; + Self::flush_pending(&inner, &pending); + rx.await.map_err(|_| { + Error::internal("a batched read was dropped before the batch was submitted") + })? + } + }; + match submitted { + SubmittedRead::Alone(request) => request.await, + SubmittedRead::Shared(batched, positions) => { + let bytes = batched.await.map_err(|err| err.0)?; + Ok(positions + .into_iter() + .map(|position| bytes[position].clone()) + .collect()) + } + } + } + .boxed() + } + + fn with_bypass_backpressure(&self) -> Option> { + self.inner.with_bypass_backpressure() + } + + fn with_io_stats( + &self, + stats: Arc, + ) -> Option> { + self.inner.with_io_stats(stats) + } +} + +pub struct SchedulerContext { + recv: Option>, + io: Arc, + /// `io`, as the batch that collects one scheduling step's reads. + batch: Arc, + cache: Arc, + name: String, + path: Vec, + path_names: Vec, +} + +pub struct ScopedSchedulerContext<'a> { + pub context: &'a mut SchedulerContext, +} + +impl<'a> ScopedSchedulerContext<'a> { + pub fn pop(self) -> &'a mut SchedulerContext { + self.context.pop(); + self.context + } +} + +impl SchedulerContext { + pub fn new(io: Arc, cache: Arc) -> Self { + let batch = Arc::new(RequestBatch::new(io)); + Self { + io: batch.clone(), + batch, + cache, + recv: None, + name: "".to_string(), + path: Vec::new(), + path_names: Vec::new(), + } + } + + pub fn io(&self) -> &Arc { + &self.io + } + + /// Submits the reads the schedulers queued since the last step as one request. + pub fn flush_io(&self) { + self.batch.flush(); + } + + pub fn cache(&self) -> &Arc { + &self.cache + } + + pub fn push(&'_ mut self, name: &str, index: u32) -> ScopedSchedulerContext<'_> { + self.path.push(index); + self.path_names.push(name.to_string()); + ScopedSchedulerContext { context: self } + } + + pub fn pop(&mut self) { + self.path.pop(); + self.path_names.pop(); + } + + pub fn path_name(&self) -> String { + let path = self.path_names.join("/"); + if self.recv.is_some() { + format!("TEMP({}){}", self.name, path) + } else { + format!("ROOT{}", path) + } + } + + pub fn current_path(&self) -> VecDeque { + VecDeque::from_iter(self.path.iter().copied()) + } + + #[deprecated(since = "0.29.1", note = "This is for v2.0 array-encoding paths")] + pub fn locate_decoder(&mut self, decoder: Box) -> DecoderReady { + trace!( + "Scheduling decoder of type {:?} for {:?}", + decoder.data_type(), + self.path, + ); + DecoderReady { + decoder, + path: self.current_path(), + } + } +} + +pub struct UnloadedPageShard(pub BoxFuture<'static, Result>); + +impl std::fmt::Debug for UnloadedPageShard { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UnloadedPage").finish() + } +} + +#[derive(Debug)] +pub struct ScheduledScanLine { + pub rows_scheduled: u64, + pub decoders: Vec, +} + +pub trait StructuralSchedulingJob: std::fmt::Debug { + /// Schedule the next batch of data + /// + /// Normally this equates to scheduling the next page of data into one task. Very large pages + /// might be split into multiple scan lines. Each scan line has one or more rows. + /// + /// If a scheduler ends early it may return an empty vector. + fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result>; +} + +/// A filter expression to apply to the data +/// +/// The core decoders do not currently take advantage of filtering in +/// any way. In order to maintain the abstraction we represent filters +/// as an arbitrary byte sequence. +/// +/// We recommend that encodings use Substrait for filters. +pub struct FilterExpression(pub Bytes); + +impl FilterExpression { + /// Create a filter expression that does not filter any data + /// + /// This is currently represented by an empty byte array. Encoders + /// that are "filter aware" should make sure they handle this case. + pub fn no_filter() -> Self { + Self(Bytes::new()) + } + + /// Returns true if the filter is the same as the [`Self::no_filter`] filter + pub fn is_noop(&self) -> bool { + self.0.is_empty() + } +} + +pub trait StructuralFieldScheduler: Send + std::fmt::Debug { + /// Loads the per-page metadata this column needs before scheduling. + /// + /// When `requested_ranges` is present, only pages overlapping those + /// top-level row ranges are initialized. `None` initializes the complete + /// field. + fn initialize<'a>( + &'a mut self, + requested_ranges: Option<&'a [Range]>, + filter: &'a FilterExpression, + context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>>; + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result>; +} + +/// A trait for tasks that decode data into an Arrow array +pub trait DecodeArrayTask: Send { + /// Decodes the data into an Arrow array and its data size in bytes + fn decode(self: Box) -> Result<(ArrayRef, u64)>; +} + +impl DecodeArrayTask for Box { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + let decoded_array = StructuralDecodeArrayTask::decode(*self)?; + decoded_array.repdef.ensure_exhausted()?; + Ok((decoded_array.array, decoded_array.data_size)) + } +} + +/// A task to decode data into an Arrow record batch +/// +/// It has a child `task` which decodes a struct array with no nulls. +/// This is then converted into a record batch. +pub struct NextDecodeTask { + /// The decode task itself + pub task: Box, + /// The number of rows that will be created + pub num_rows: u64, +} + +impl NextDecodeTask { + // Run the task and produce a record batch + // + // If the batch is very large this function will log a warning message + // suggesting the user try a smaller batch size. + #[instrument(name = "task_to_batch", level = "debug", skip_all)] + fn into_batch(self, emitted_batch_size_warning: Arc) -> Result<(RecordBatch, u64)> { + let (struct_arr, data_size) = self.task.decode()?; + let batch = RecordBatch::from(struct_arr.as_struct()); + if data_size > BATCH_SIZE_BYTES_WARNING { + emitted_batch_size_warning.call_once(|| { + let size_mb = data_size / 1024 / 1024; + debug!("Lance read in a single batch that contained more than {}MiB of data. You may want to consider reducing the batch size.", size_mb); + }); + } + Ok((batch, data_size)) + } +} + +// An envelope to wrap both 2.0 style messages and 2.1 style messages so we can +// share some code paths between the two. Decoders can safely unwrap into whatever +// style they expect since a file will be either all-2.0 or all-2.1 +#[derive(Debug)] +pub enum MessageType { + // The older v2.0 scheduler/decoder used a scheme where the message was the + // decoder itself. The messages were not sent in priority order and the decoder + // had to wait for I/O, figuring out the correct priority. This was a lot of + // complexity. + DecoderReady(DecoderReady), + // Starting in 2.1 we use a simpler scheme where the scheduling happens in priority + // order and the message is an unloaded decoder. These can be awaited, in order, and + // the decoder does not have to worry about waiting for I/O. + UnloadedPage(UnloadedPageShard), +} + +impl MessageType { + pub fn into_array(self) -> DecoderReady { + match self { + Self::DecoderReady(decoder) => decoder, + Self::UnloadedPage(_) => { + panic!("Expected DecoderReady but got UnloadedPage") + } + } + } + + pub fn into_structural(self) -> UnloadedPageShard { + match self { + Self::UnloadedPage(unloaded) => unloaded, + Self::DecoderReady(_) => { + panic!("Expected UnloadedPage but got DecoderReady") + } + } + } +} + +pub struct DecoderMessage { + pub scheduled_so_far: u64, + pub decoders: Vec, +} + +pub struct DecoderContext { + source: mpsc::UnboundedReceiver>, +} + +impl DecoderContext { + pub fn new(source: mpsc::UnboundedReceiver>) -> Self { + Self { source } + } +} + +pub struct DecodedPage { + pub data: DataBlock, + pub repdef: RepDefUnraveler, +} + +pub trait DecodePageTask: Send + std::fmt::Debug { + /// Decodes the data into an Arrow array + fn decode(self: Box) -> Result; +} + +pub trait StructuralPageDecoder: std::fmt::Debug + Send { + fn drain(&mut self, num_rows: u64) -> Result>; + fn num_rows(&self) -> u64; + /// Returns the exact decoded byte count for the next `num_rows` rows + /// from this decoder's current position, without consuming any rows. + fn decoded_bytes(&self, _num_rows: u64) -> Result { + Err(Error::not_supported( + "decoded_bytes is not implemented for this page decoder".to_string(), + )) + } +} + +#[derive(Debug)] +pub struct LoadedPageShard { + // The decoder that is ready to be decoded + pub decoder: Box, + // The path to the decoder, the first value is the column index + // following values, if present, are nested child indices + // + // For example, a path of [1, 1, 0] would mean to grab the second + // column, then the second child, and then the first child. + // + // It could represent x in the following schema: + // + // score: float64 + // points: struct + // color: string + // location: struct + // x: float64 + // + // Currently, only struct decoders have "children" although other + // decoders may at some point as well. List children are only + // handled through indirect I/O at the moment and so they don't + // need to be represented (yet) + pub path: VecDeque, +} + +pub struct DecodedArray { + pub array: ArrayRef, + pub repdef: CompositeRepDefUnraveler, + /// The number of bytes of data in this array (excluding Arrow overhead). + pub data_size: u64, +} + +pub trait StructuralDecodeArrayTask: std::fmt::Debug + Send { + fn decode(self: Box) -> Result; +} + +pub trait StructuralFieldDecoder: std::fmt::Debug + Send { + /// Add a newly scheduled child decoder + /// + /// The default implementation does not expect children and returns + /// an error. + fn accept_page(&mut self, _child: LoadedPageShard) -> Result<()>; + /// Creates a task to decode `num_rows` of data into an array + fn drain(&mut self, num_rows: u64) -> Result>; + /// The data type of the decoded data + fn data_type(&self) -> &DataType; + /// Returns the exact decoded byte count for each of [`CANDIDATE_BATCH_SIZES`] + /// row counts, clamped to `rows_remaining`. + /// + /// Implementations should do their best to estimate the exact size required for + /// the uncompressed data. In cases where this is not possible they should return + /// a worst-case estimate. + /// + /// The default implementation simply returns a "not supported" error though this + /// will hopefully be removed once implementation is complete. + fn plan_decoded_bytes(&self, _rows_remaining: u64) -> Result<[u64; 8]> { + Err(Error::not_supported( + "decoded_bytes is not implemented for this field decoder".to_string(), + )) + } +} + +#[derive(Debug, Default)] +pub struct DecoderPlugins {} + +/// The top-level column layout used by an in-memory encoded batch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EncodedBatchLayout { + /// Array pages include structural columns. + Array, + /// Structural pages include only leaf columns. + Structural, +} + +/// Decodes a batch of data from an in-memory structure created by [`crate::encoder::encode_batch`] +pub async fn decode_batch( + batch: &EncodedBatch, + filter: &FilterExpression, + decoder_plugins: Arc, + should_validate: bool, + layout: EncodedBatchLayout, + cache: Option>, +) -> Result { + // The io is synchronous so it shouldn't be possible for any async stuff to still be in progress + // Still, if we just use now_or_never we hit misfires because some futures (channels) need to be + // polled twice. + + let io_scheduler = Arc::new(BufferScheduler::new(batch.data.clone())) as Arc; + let cache = if let Some(cache) = cache { + cache + } else { + Arc::new(lance_core::cache::LanceCache::with_capacity( + 128 * 1024 * 1024, + )) + }; + let mut decode_scheduler = DecodeBatchScheduler::try_new( + batch.schema.as_ref(), + &batch.top_level_columns, + &batch.page_table, + &vec![], + batch.num_rows, + decoder_plugins, + io_scheduler.clone(), + cache, + filter, + &DecoderConfig::default(), + ) + .await?; + let (tx, rx) = unbounded_channel(); + decode_scheduler.schedule_range(0..batch.num_rows, filter, tx, io_scheduler); + let is_structural = layout == EncodedBatchLayout::Structural; + let mode = std::env::var(ENV_LANCE_STRUCTURAL_BATCH_DECODE_SPAWN_MODE); + let spawn_structural_batch_decode_tasks = !matches!(mode.ok().as_deref(), Some("never")); + // `EncodedBatch` is created from one RecordBatch, so every i32-offset array is + // already known to fit in one Arrow array. Preserve the single-batch contract here. + let mut decode_stream = create_decode_stream_with_i32_offset_limit( + &batch.schema, + batch.num_rows, + batch.num_rows as u32, + is_structural, + should_validate, + spawn_structural_batch_decode_tasks, + rx, + None, + false, + )?; + decode_stream.next().await.unwrap().task.await +} + +#[cfg(test)] +// test coalesce indices to ranges +mod tests { + use super::*; + use arrow_array::StringArray; + use std::collections::VecDeque; + + struct StaticArrayDecodeTask(ArrayRef); + + impl DecodeArrayTask for StaticArrayDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + Ok((self.0, 0)) + } + } + + #[derive(Debug)] + struct StaticArrayPageDecoder { + array: ArrayRef, + rows_drained: u64, + /// Pretend each row decodes to this many bytes instead of its real size, + /// so budget exhaustion can be tested without gigabyte allocations. + fake_bytes_per_row: Option, + } + + impl StaticArrayPageDecoder { + fn new(values: &[&str]) -> Self { + Self { + array: Arc::new(StringArray::from(values.to_vec())), + rows_drained: 0, + fake_bytes_per_row: None, + } + } + + fn with_fake_bytes_per_row(values: &[&str], fake_bytes_per_row: u64) -> Self { + Self { + fake_bytes_per_row: Some(fake_bytes_per_row), + ..Self::new(values) + } + } + + fn row_bytes(&self, row: usize) -> u64 { + self.fake_bytes_per_row.unwrap_or_else(|| { + self.array + .as_any() + .downcast_ref::() + .unwrap() + .value(row) + .len() as u64 + }) + } + } + + impl LogicalPageDecoder for StaticArrayPageDecoder { + fn wait_for_loaded(&mut self, _loaded_need: u64) -> BoxFuture<'_, Result<()>> { + std::future::ready(Ok(())).boxed() + } + + fn rows_loaded(&self) -> u64 { + self.array.len() as u64 + } + + fn num_rows(&self) -> u64 { + self.array.len() as u64 + } + + fn rows_drained(&self) -> u64 { + self.rows_drained + } + + fn drain(&mut self, num_rows: u64) -> Result { + let array = self + .array + .slice(self.rows_drained as usize, num_rows as usize); + self.rows_drained += num_rows; + Ok(NextDecodeTask { + task: Box::new(StaticArrayDecodeTask(array)), + num_rows, + }) + } + + fn max_rows_to_drain(&self, num_rows: u64, byte_budget: u64) -> Result { + let start = self.rows_drained as usize; + let mut rows = 0u64; + let mut bytes = 0u64; + for row in start..start + num_rows as usize { + let row_bytes = self.row_bytes(row); + if bytes + row_bytes > byte_budget { + break; + } + bytes += row_bytes; + rows += 1; + } + Ok(DrainLimit { rows, bytes }) + } + + fn data_type(&self) -> &DataType { + self.array.data_type() + } + } + + async fn array_stream_batch_sizes( + pages: Vec, + batch_size: u32, + ) -> Vec { + let num_rows = pages.iter().map(|page| page.num_rows()).sum::(); + let fields = Fields::from(vec![ArrowField::new("value", DataType::Utf8, false)]); + let root_decoder = SimpleStructDecoder::new(fields, num_rows); + let (tx, rx) = unbounded_channel(); + let pages = pages + .into_iter() + .map(|page| { + MessageType::DecoderReady(DecoderReady { + decoder: Box::new(page), + path: VecDeque::from([0]), + }) + }) + .collect(); + tx.send(Ok(DecoderMessage { + scheduled_so_far: num_rows, + decoders: pages, + })) + .unwrap(); + drop(tx); + + let mut stream = + BatchDecodeStream::new(rx, batch_size, num_rows, root_decoder).into_stream(); + let mut batch_sizes = Vec::new(); + while let Some(task) = stream.next().await { + batch_sizes.push(task.task.await.unwrap().num_rows()); + } + batch_sizes + } + + #[tokio::test] + async fn test_array_stream_spans_i32_offset_pages_within_budget() { + // Small variable-width pages must NOT split the batch at page boundaries. + let batch_sizes = array_stream_batch_sizes( + vec![ + StaticArrayPageDecoder::new(&["a", "b"]), + StaticArrayPageDecoder::new(&["c", "d"]), + ], + 4, + ) + .await; + assert_eq!(batch_sizes, vec![4]); + } + + #[tokio::test] + async fn test_array_stream_splits_when_i32_offset_budget_exhausted() { + // Three single-row pages pretending to hold 800MiB each: two fit the + // i32 offset budget (~2GiB), the third starts a new batch. + const FAKE_ROW_BYTES: u64 = 800 * 1024 * 1024; + let batch_sizes = array_stream_batch_sizes( + vec![ + StaticArrayPageDecoder::with_fake_bytes_per_row(&["a"], FAKE_ROW_BYTES), + StaticArrayPageDecoder::with_fake_bytes_per_row(&["b"], FAKE_ROW_BYTES), + StaticArrayPageDecoder::with_fake_bytes_per_row(&["c"], FAKE_ROW_BYTES), + ], + 4, + ) + .await; + assert_eq!(batch_sizes, vec![2, 1]); + } + + #[tokio::test] + async fn test_array_stream_emits_over_budget_rows_alone() { + // Rows whose (possibly over-estimated) size exceeds the whole budget are + // emitted as single-row batches rather than failing the scan. + const FAKE_ROW_BYTES: u64 = 3 * 1024 * 1024 * 1024; + let batch_sizes = array_stream_batch_sizes( + vec![ + StaticArrayPageDecoder::with_fake_bytes_per_row(&["a"], FAKE_ROW_BYTES), + StaticArrayPageDecoder::with_fake_bytes_per_row(&["b"], FAKE_ROW_BYTES), + ], + 4, + ) + .await; + assert_eq!(batch_sizes, vec![1, 1]); + } + + #[test] + fn requested_row_indices_convert_to_checked_ranges() { + let requested_rows = RequestedRows::Indices(vec![1, 2, 5, 8, 9]); + let ranges = requested_rows.into_ranges().unwrap(); + assert_eq!(ranges, [1..3, 5..6, 8..10]); + + let error = RequestedRows::Indices(vec![2, 2]) + .into_ranges() + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. })); + assert!(error.to_string().contains("not strictly increasing")); + + let error = RequestedRows::Indices(vec![u64::MAX]) + .into_ranges() + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. })); + assert!(error.to_string().contains("u64::MAX")); + } + + #[derive(Debug)] + struct FailingPageDecoder { + page_data_type: DataType, + total_rows: u64, + load_error_message: &'static str, + } + + impl FailingPageDecoder { + fn new( + page_data_type: DataType, + total_rows: u64, + load_error_message: &'static str, + ) -> Self { + Self { + page_data_type, + total_rows, + load_error_message, + } + } + } + + impl LogicalPageDecoder for FailingPageDecoder { + fn wait_for_loaded(&'_ mut self, _rows_needed: u64) -> BoxFuture<'_, Result<()>> { + let load_error_message = self.load_error_message; + async move { Err(Error::io(load_error_message)) }.boxed() + } + + fn rows_loaded(&self) -> u64 { + 0 + } + + fn num_rows(&self) -> u64 { + self.total_rows + } + + fn rows_drained(&self) -> u64 { + 0 + } + + fn drain(&mut self, requested_rows: u64) -> Result { + Err(Error::internal(format!( + "failing page decoder should not be drained after load error \ + (requested_rows={})", + requested_rows + ))) + } + + fn data_type(&self) -> &DataType { + &self.page_data_type + } + } + + struct InvalidInputDecodeTask; + + impl DecodeArrayTask for InvalidInputDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + Err(Error::invalid_input_source("malformed sparse page".into())) + } + } + + #[test] + fn next_decode_task_preserves_invalid_input_errors() { + let err = NextDecodeTask { + task: Box::new(InvalidInputDecodeTask), + num_rows: 0, + } + .into_batch(Arc::new(Once::new())) + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + } + + #[test] + fn test_read_zero_dimension_fsl_errors_instead_of_panicking() { + // Simulates reading a column whose stored schema declares a + // zero-dimension FixedSizeList, as old writers (before #5102) could + // persist. The read plan is built by the field-scheduler factories, + // which run the dimension guard before touching any column data, so + // an empty column iterator is sufficient to reach the guard. The read + // must surface a clean error rather than a divide-by-zero panic. + use arrow_schema::Field as ArrowField; + + let zero_dim = DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 0, + ); + let field = Field::try_from(&ArrowField::new("vec", zero_dim, true)).unwrap(); + let strategy = CoreFieldDecoderStrategy::default(); + + let mut structural_columns = ColumnInfoIter::new(vec![], &[]); + let err = strategy + .create_structural_field_scheduler(&field, &mut structural_columns) + .unwrap_err(); + assert!( + err.to_string() + .contains("dimension must be a positive integer"), + "unexpected error: {}", + err + ); + + let mut array_columns = ColumnInfoIter::new(vec![], &[]); + let err = strategy + .create_array_field_scheduler( + &field, + &mut array_columns, + FileBuffers { + positions_and_sizes: &[], + }, + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("dimension must be a positive integer"), + "unexpected error: {}", + err + ); + } + + #[test] + fn test_list_page_with_non_list_encoding_returns_error() { + let item = Arc::new(ArrowField::new("item", DataType::Int32, true)); + let list = DataType::List(item); + let field = Field::try_from(&ArrowField::new("values", list, true)).unwrap(); + let values_encoding = pb::ColumnEncoding { + column_encoding: Some(pb::column_encoding::ColumnEncoding::Values(())), + }; + let offsets_column = Arc::new(ColumnInfo::new( + 0, + Arc::new([PageInfo { + num_rows: 1, + priority: 0, + encoding: PageEncoding::Legacy(pb::ArrayEncoding { + array_encoding: Some(pb::array_encoding::ArrayEncoding::Flat( + pb::Flat::default(), + )), + }), + buffer_offsets_and_sizes: Arc::new([]), + }]), + vec![], + values_encoding.clone(), + )); + let items_column = Arc::new(ColumnInfo::new(1, Arc::new([]), vec![], values_encoding)); + let column_indices = [0, 1]; + let mut columns = ColumnInfoIter::new(vec![offsets_column, items_column], &column_indices); + + let err = CoreFieldDecoderStrategy::default() + .create_array_field_scheduler( + &field, + &mut columns, + FileBuffers { + positions_and_sizes: &[], + }, + ) + .unwrap_err(); + + assert!(matches!(err, Error::InvalidInput { .. })); + assert!( + err.to_string() + .contains("expected list encoding for field 'values' in column 0, page 0 but got"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn test_array_stream_stops_on_load_error() { + use arrow_schema::Field as ArrowField; + + let rows_per_batch = 1; + let total_rows = 2; + let scheduled_rows = 1; + let page_rows = 1; + let batch_readahead = 2; + let load_error_message = "simulated page load failure"; + let fields = Fields::from(vec![ArrowField::new("vector", DataType::Float32, true)]); + let root_decoder = SimpleStructDecoder::new(fields, total_rows); + let (tx, rx) = unbounded_channel(); + + tx.send(Ok(DecoderMessage { + scheduled_so_far: scheduled_rows, + decoders: vec![MessageType::DecoderReady(DecoderReady { + decoder: Box::new(FailingPageDecoder::new( + DataType::Float32, + page_rows, + load_error_message, + )), + path: VecDeque::from([0]), + })], + })) + .unwrap(); + drop(tx); + + let stream = + BatchDecodeStream::new(rx, rows_per_batch, total_rows, root_decoder).into_stream(); + let mut batches = stream.map(|task| task.task).buffered(batch_readahead); + + let err = batches + .next() + .await + .expect("stream should emit the array page-load error") + .unwrap_err(); + assert!( + err.to_string().contains(load_error_message), + "unexpected error: {}", + err + ); + assert!( + batches.next().await.is_none(), + "stream should stop after the array page-load error" + ); + } + + #[tokio::test] + async fn test_structural_stream_stops_on_load_error() { + let rows_per_batch = 1; + let total_rows = 2; + let scheduled_rows = 1; + let batch_readahead = 2; + let load_error_message = "simulated page load failure"; + let fields = Fields::from(vec![ArrowField::new("vector", DataType::Float32, true)]); + let root_decoder = StructuralStructDecoder::new( + fields, false, /*is_root=*/ true, /*nullable=*/ false, + ) + .unwrap(); + let (tx, rx) = unbounded_channel(); + let failed_page = async move { Err(Error::io(load_error_message)) }.boxed(); + + tx.send(Ok(DecoderMessage { + scheduled_so_far: scheduled_rows, + decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(failed_page))], + })) + .unwrap(); + drop(tx); + + let stream = StructuralBatchDecodeStream::new( + rx, + rows_per_batch, + total_rows, + root_decoder, + /*spawn_batch_decode_tasks=*/ true, + None, + ) + .into_stream(); + let mut batches = stream.map(|task| task.task).buffered(batch_readahead); + + let err = batches + .next() + .await + .expect("stream should emit the page-load error") + .unwrap_err(); + assert!( + err.to_string().contains(load_error_message), + "unexpected error: {}", + err + ); + assert!( + batches.next().await.is_none(), + "stream should stop after the page-load error" + ); + } + + #[test] + fn test_coalesce_indices_to_ranges_with_single_index() { + let indices = vec![1]; + let ranges = RequestedRows::indices_to_ranges(&indices).unwrap(); + assert_eq!(ranges, vec![1..2]); + } + + #[test] + fn test_coalesce_indices_to_ranges() { + let indices = vec![1, 2, 3, 4, 5, 6, 7, 8, 9]; + let ranges = RequestedRows::indices_to_ranges(&indices).unwrap(); + assert_eq!(ranges, vec![1..10]); + } + + #[test] + fn test_coalesce_indices_to_ranges_with_gaps() { + let indices = vec![1, 2, 3, 5, 6, 7, 9]; + let ranges = RequestedRows::indices_to_ranges(&indices).unwrap(); + assert_eq!(ranges, vec![1..4, 5..8, 9..10]); + } + + #[test] + fn test_estimate_bytes_per_row() { + assert_eq!(estimate_bytes_per_row(&DataType::Int32), 4.0); + assert_eq!(estimate_bytes_per_row(&DataType::Int64), 8.0); + assert_eq!(estimate_bytes_per_row(&DataType::Float32), 4.0); + assert_eq!(estimate_bytes_per_row(&DataType::Boolean), 1.0 / 8.0); + assert_eq!(estimate_bytes_per_row(&DataType::Utf8), 64.0); + assert_eq!(estimate_bytes_per_row(&DataType::Binary), 64.0); + // Struct of 4 x Int32 = 16 bytes + let struct_type = DataType::Struct(Fields::from(vec![ + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Int32, false), + ArrowField::new("c", DataType::Int32, false), + ArrowField::new("d", DataType::Int32, false), + ])); + assert_eq!(estimate_bytes_per_row(&struct_type), 16.0); + } + + /// Helper: encode a batch, then decode it as a stream with optional + /// `batch_size_bytes`, collecting all output batches. + async fn decode_batches_with_byte_limit( + batch: &RecordBatch, + batch_size: u32, + batch_size_bytes: Option, + ) -> Vec { + use crate::{ + encoder::{EncodingOptions, encode_batch}, + testing::{TestEncoding, test_encoding_strategy}, + }; + + let version = TestEncoding::StructuralU16; + let options = EncodingOptions::default(); + let strategy = test_encoding_strategy(version); + let schema = Schema::try_from(batch.schema().as_ref()).unwrap(); + let encoded = encode_batch(batch, Arc::new(schema.clone()), strategy.as_ref(), &options) + .await + .unwrap(); + + let io_scheduler = + Arc::new(BufferScheduler::new(encoded.data.clone())) as Arc; + let cache = Arc::new(lance_core::cache::LanceCache::with_capacity( + 128 * 1024 * 1024, + )); + let decoder_plugins = Arc::new(DecoderPlugins::default()); + + let mut decode_scheduler = DecodeBatchScheduler::try_new( + encoded.schema.as_ref(), + &encoded.top_level_columns, + &encoded.page_table, + &vec![], + encoded.num_rows, + decoder_plugins, + io_scheduler.clone(), + cache, + &FilterExpression::no_filter(), + &DecoderConfig::default(), + ) + .await + .unwrap(); + + let (tx, rx) = unbounded_channel(); + decode_scheduler.schedule_range( + 0..encoded.num_rows, + &FilterExpression::no_filter(), + tx, + io_scheduler, + ); + + let mut decode_stream = create_decode_stream( + &encoded.schema, + encoded.num_rows, + batch_size, + /*is_structural=*/ true, + /*should_validate=*/ true, + /*spawn_structural_batch_decode_tasks=*/ true, + rx, + batch_size_bytes, + ) + .unwrap(); + + let mut batches = Vec::new(); + while let Some(task) = decode_stream.next().await { + batches.push(task.task.await.unwrap()); + } + batches + } + + #[tokio::test] + async fn test_byte_sized_batches_fixed_width() { + use arrow_array::Int32Array; + + // 1000 rows x 4 Int32 columns = 16 bytes/row + let num_rows: i32 = 1000; + let arrays: Vec> = (0..4) + .map(|col| { + Arc::new(Int32Array::from_iter_values( + (0..num_rows).map(move |row| row * 10 + col), + )) as _ + }) + .collect(); + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Int32, false), + ArrowField::new("c", DataType::Int32, false), + ArrowField::new("d", DataType::Int32, false), + ])); + let input_batch = RecordBatch::try_new(schema, arrays).unwrap(); + + // 16 bytes/row, batch_size_bytes=1600 => 100 rows/batch + let batches = + decode_batches_with_byte_limit(&input_batch, /*batch_size=*/ 1024, Some(1600)).await; + + // Should produce 10 batches of 100 rows each + assert_eq!(batches.len(), 10); + for (i, batch) in batches.iter().enumerate() { + assert_eq!( + batch.num_rows(), + 100, + "batch {i} should have 100 rows, got {}", + batch.num_rows() + ); + } + + // Verify roundtrip: concatenate and compare + let all_batches: Vec<&RecordBatch> = batches.iter().collect(); + let concatenated = + arrow_select::concat::concat_batches(&batches[0].schema(), all_batches.iter().copied()) + .unwrap(); + assert_eq!(concatenated.num_rows(), num_rows as usize); + for col in 0..4 { + assert_eq!( + concatenated.column(col).as_ref(), + input_batch.column(col).as_ref(), + "column {col} roundtrip mismatch" + ); + } + } + + #[tokio::test] + async fn test_byte_sized_batches_none_unchanged() { + use arrow_array::Int32Array; + + // Without batch_size_bytes, rows_per_batch controls batching + let num_rows: i32 = 1000; + let arrays: Vec> = (0..2) + .map(|col| { + Arc::new(Int32Array::from_iter_values( + (0..num_rows).map(move |row| row * 10 + col), + )) as _ + }) + .collect(); + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("x", DataType::Int32, false), + ArrowField::new("y", DataType::Int32, false), + ])); + let input_batch = RecordBatch::try_new(schema, arrays).unwrap(); + + // batch_size=250, batch_size_bytes=None => 4 batches of 250 rows + let batches = decode_batches_with_byte_limit(&input_batch, /*batch_size=*/ 250, None).await; + assert_eq!(batches.len(), 4); + for (i, batch) in batches.iter().enumerate() { + assert_eq!( + batch.num_rows(), + 250, + "batch {i} should have 250 rows, got {}", + batch.num_rows() + ); + } + } + + #[tokio::test] + async fn test_byte_sized_batches_respect_row_limit() { + use arrow_array::Int32Array; + + let num_rows: i32 = 1000; + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "x", + DataType::Int32, + false, + )])); + let input_batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from_iter_values(0..num_rows))], + ) + .unwrap(); + + // The byte limit can hold every row, so the 100-row limit must win. + let batches = + decode_batches_with_byte_limit(&input_batch, /*batch_size=*/ 100, Some(10_000)).await; + assert_eq!(batches.len(), 10); + assert!(batches.iter().all(|batch| batch.num_rows() == 100)); + } + + #[tokio::test] + async fn test_byte_sized_batches_feedback_convergence() { + use arrow_array::StringArray; + + // Each row has a 100-byte string. Schema estimate = 64 bytes (default + // for Utf8), so the first batch will overshoot. The feedback loop + // should correct subsequent batches toward the target. + let num_rows = 500; + let value: String = "x".repeat(100); + let arrays: Vec> = vec![Arc::new(StringArray::from( + (0..num_rows).map(|_| value.as_str()).collect::>(), + ))]; + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Utf8, + false, + )])); + let input_batch = RecordBatch::try_new(schema, arrays).unwrap(); + + // Target 5000 bytes/batch. At 100 bytes/row the ideal is 50 rows/batch. + // Schema estimate is 64 bytes/row → first batch ~78 rows (overshoot). + // After feedback kicks in, batches should converge to ~50 rows. + let target_bytes: u64 = 5000; + let batches = decode_batches_with_byte_limit( + &input_batch, + /*batch_size=*/ 1024, + Some(target_bytes), + ) + .await; + + // Verify all data round-trips correctly + let all_batches: Vec<&RecordBatch> = batches.iter().collect(); + let concatenated = + arrow_select::concat::concat_batches(&batches[0].schema(), all_batches.iter().copied()) + .unwrap(); + assert_eq!(concatenated.num_rows(), num_rows as usize); + assert_eq!( + concatenated.column(0).as_ref(), + input_batch.column(0).as_ref() + ); + + // After the first batch, subsequent batches should be closer to the + // target. The ideal is 50 rows/batch. + assert!( + batches.len() >= 2, + "need at least 2 batches to test convergence" + ); + // The first batch uses the schema estimate (64 bytes/row) → + // ~78 rows. After feedback the rows should settle near 50. + if batches.len() >= 3 { + let second_batch_rows = batches[1].num_rows(); + let third_batch_rows = batches[2].num_rows(); + // Both should be within 20% of the ideal (50 rows) + assert!( + (40..=60).contains(&second_batch_rows), + "second batch should be near 50 rows, got {second_batch_rows}" + ); + assert!( + (40..=60).contains(&third_batch_rows), + "third batch should be near 50 rows, got {third_batch_rows}" + ); + } + } +} diff --git a/vendor/lance-encoding/src/encoder.rs b/vendor/lance-encoding/src/encoder.rs new file mode 100644 index 000000000..ee3c7d854 --- /dev/null +++ b/vendor/lance-encoding/src/encoder.rs @@ -0,0 +1,603 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! The top-level encoding module for Lance files. +//! +//! Lance files are encoded using a [`FieldEncodingStrategy`] which choose +//! what encoder to use for each field. +//! +//! Structural strategies build a tree of encoders for each field from the +//! version-free builders in [`structural`]. Struct and list encoders collect +//! validity and offsets; primitive leaf encoders accumulate values and emit +//! miniblock or full-zip pages. + +use std::{collections::HashMap, sync::Arc}; + +use arrow_array::{Array, ArrayRef, RecordBatch}; +use bytes::{Bytes, BytesMut}; +use futures::future::BoxFuture; +use lance_core::datatypes::{Field, Schema}; +use lance_core::utils::bit::{is_pwr_two, pad_bytes_to}; +use lance_core::{Error, Result}; + +use crate::buffer::LanceBuffer; +use crate::data::DataBlock; +use crate::decoder::PageEncoding; +use crate::repdef::RepDefBuilder; +use crate::{ + decoder::{ColumnInfo, PageInfo}, + format::pb, +}; + +pub use crate::array_encoding::ArrayFieldEncodingStrategy; + +pub mod structural; + +/// The minimum alignment for a page buffer. Writers must respect this. +pub const MIN_PAGE_BUFFER_ALIGNMENT: u64 = 8; + +/// An array encoded with the `pb::ArrayEncoding` grammar. +#[derive(Debug)] +pub struct EncodedArray { + pub data: DataBlock, + pub encoding: pb::ArrayEncoding, +} + +impl EncodedArray { + pub fn new(data: DataBlock, encoding: pb::ArrayEncoding) -> Self { + Self { data, encoding } + } + + pub fn into_buffers(self) -> (Vec, pb::ArrayEncoding) { + (self.data.into_buffers(), self.encoding) + } +} + +/// Encodes one data block and describes it with `pb::ArrayEncoding`. +pub trait ArrayEncoder: std::fmt::Debug + Send + Sync { + fn encode( + &self, + data: DataBlock, + data_type: &arrow_schema::DataType, + buffer_index: &mut u32, + ) -> Result; +} + +/// Selects an `ArrayEncoder` for one page. +pub trait ArrayEncodingStrategy: Send + Sync + std::fmt::Debug { + fn create_array_encoder( + &self, + arrays: &[ArrayRef], + field: &Field, + ) -> Result>; +} + +/// An encoded page of data +/// +/// Maps to a top-level array +/// +/// For example, `FixedSizeList` will have two EncodedArray instances and one EncodedPage +#[derive(Debug)] +pub struct EncodedPage { + // The encoded page buffers + pub data: Vec, + // A description of the encoding used to encode the page + pub description: PageEncoding, + /// The number of rows in the encoded page + pub num_rows: u64, + /// The top-level row number of the first row in the page + /// + /// Generally the number of "top-level" rows and the number of rows are the same. However, + /// when there is repetition (list/fixed-size-list) there will be more or less items than rows. + /// + /// A top-level row can never be split across a page boundary. + pub row_number: u64, + /// The index of the column + pub column_idx: u32, +} + +pub struct EncodedColumn { + pub column_buffers: Vec, + pub encoding: pb::ColumnEncoding, + pub final_pages: Vec, +} + +impl Default for EncodedColumn { + fn default() -> Self { + Self { + column_buffers: Default::default(), + encoding: pb::ColumnEncoding { + column_encoding: Some(pb::column_encoding::ColumnEncoding::Values(())), + }, + final_pages: Default::default(), + } + } +} + +/// A tool to reserve space for buffers that are not in-line with the data +/// +/// In most cases, buffers are stored in the page and referred to in the encoding +/// metadata by their index in the page. This keeps all buffers within a page together. +/// As a result, most encoders should not need to use this structure. +/// +/// In some cases (currently only the large binary encoding) there is a need to access +/// buffers that are not in the page (because storing the position / offset of every page +/// in the page metadata would be too expensive). +/// +/// To do this you can add a buffer with `add_buffer` and then use the returned position +/// in some way (in the large binary encoding the returned position is stored in the page +/// data as a position / size array). +pub struct OutOfLineBuffers { + position: u64, + buffer_alignment: u64, + buffers: Vec, +} + +impl OutOfLineBuffers { + pub fn new(base_position: u64, buffer_alignment: u64) -> Self { + Self { + position: base_position, + buffer_alignment, + buffers: Vec::new(), + } + } + + pub fn add_buffer(&mut self, buffer: LanceBuffer) -> u64 { + let position = self.position; + self.position += buffer.len() as u64; + self.position += pad_bytes_to(buffer.len(), self.buffer_alignment as usize) as u64; + self.buffers.push(buffer); + position + } + + pub fn take_buffers(self) -> Vec { + self.buffers + } + + pub fn reset_position(&mut self, position: u64) { + self.position = position; + } +} + +/// A task to create a page of data +pub type EncodeTask = BoxFuture<'static, Result>; + +/// Top level encoding trait to code any Arrow array type into one or more pages. +/// +/// The field encoder implements buffering and encoding of a single input column +/// but it may map to multiple output columns. For example, a list array or struct +/// array will be encoded into multiple columns. +/// +/// Also, fields may be encoded at different speeds. For example, given a struct +/// column with three fields (a boolean field, an int32 field, and a 4096-dimension +/// tensor field) the tensor field is likely to emit encoded pages much more frequently +/// than the boolean field. +pub trait FieldEncoder: Send { + /// Validate and prepare an array before any encoder state is mutated. + /// + /// Batch-oriented callers invoke this for every root field before calling + /// [`Self::maybe_encode`], so a later validation failure cannot leave an + /// earlier field encoder partially advanced. + fn prepare_array(&mut self, array: ArrayRef) -> Result { + Ok(array) + } + + /// Buffer the data and, if there is enough data in the buffer to form a page, return + /// an encoding task to encode the data. + /// + /// This may return more than one task because a single column may be mapped to multiple + /// output columns. For example, if encoding a struct column with three children then + /// up to three tasks may be returned from each call to maybe_encode. + /// + /// It may also return multiple tasks for a single column if the input array is larger + /// than a single disk page. + /// + /// It could also return an empty Vec if there is not enough data yet to encode any pages. + /// + /// The `row_number` must be passed which is the top-level row number currently being encoded + /// This is stored in any pages produced by this call so that we can know the priority of the + /// page. + /// + /// The `num_rows` is the number of top level rows. It is initially the same as `array.len()` + /// however it is passed seprately because array will become flattened over time (if there is + /// repetition) and we need to know the original number of rows for various purposes. + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result>; + /// Flush any remaining data from the buffers into encoding tasks + /// + /// Each encode task produces a single page. The order of these pages will be maintained + /// in the file (we do not worry about order between columns but all pages in the same + /// column should maintain order) + /// + /// This may be called intermittently throughout encoding but will always be called + /// once at the end of encoding just before calling finish + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result>; + /// Finish encoding and return column metadata + /// + /// This is called only once, after all encode tasks have completed + /// + /// This returns a Vec because a single field may have created multiple columns + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>>; + + /// The number of output columns this encoding will create + fn num_columns(&self) -> u32; +} + +/// Keeps track of the current column index and makes a mapping +/// from field id to column index +#[derive(Debug, Default)] +pub struct ColumnIndexSequence { + current_index: u32, + mapping: Vec<(u32, u32)>, +} + +impl ColumnIndexSequence { + pub fn next_column_index(&mut self, field_id: u32) -> u32 { + let idx = self.current_index; + self.current_index += 1; + self.mapping.push((field_id, idx)); + idx + } + + pub fn skip(&mut self) { + self.current_index += 1; + } +} + +/// Options that control the encoding process +pub struct EncodingOptions { + /// How much data (in bytes) to cache in-memory before writing a page + /// + /// This cache is applied on a per-column basis + pub cache_bytes_per_column: u64, + /// The maximum size of a page in bytes, if a single array would create + /// a page larger than this then it will be split into multiple pages + pub max_page_bytes: u64, + /// If false (the default) then arrays will be copied (deeply) before + /// being cached. This ensures any data kept alive by the array can + /// be discarded safely and helps avoid writer accumulation. However, + /// there is an associated cost. + pub keep_original_array: bool, + /// The alignment that the writer is applying to buffers + /// + /// The encoder needs to know this so it figures the position of out-of-line + /// buffers correctly + pub buffer_alignment: u64, +} + +impl Default for EncodingOptions { + fn default() -> Self { + Self { + cache_bytes_per_column: 8 * 1024 * 1024, + max_page_bytes: 32 * 1024 * 1024, + keep_original_array: true, + buffer_alignment: 64, + } + } +} + +/// A trait to pick which kind of field encoding to use for a field +/// +/// Unlike the ArrayEncodingStrategy, the field encoding strategy is +/// chosen before any data is generated and the same field encoder is +/// used for all data in the field. +pub trait FieldEncodingStrategy: Send + Sync + std::fmt::Debug { + /// Validate one top-level array before any field encoder state is mutated. + fn validate_array(&self, _array: &dyn Array, _field: &Field) -> Result<()> { + Ok(()) + } + + /// Choose and create an appropriate field encoder for the given + /// field. + /// + /// The field encoder can be chosen on the data type as well as + /// any metadata that is attached to the field. + /// + fn create_field_encoder( + &self, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result>; +} + +/// Context shared while one top-level field and all of its children are mapped +/// to concrete field encoders. +pub struct FieldEncodingContext<'a> { + /// The complete strategy composition used for recursive child fields. + pub strategy: &'a dyn FieldEncodingStrategy, + /// Runtime-only writer options. + pub options: &'a EncodingOptions, + /// Metadata inherited from the top-level field. + pub root_field_metadata: &'a HashMap, +} + +/// A batch encoder that encodes RecordBatch objects by delegating +/// to field encoders for each top-level field in the batch. +pub struct BatchEncoder { + pub field_encoders: Vec>, + pub field_id_to_column_index: Vec<(u32, u32)>, +} + +impl BatchEncoder { + pub fn try_new( + schema: &Schema, + strategy: &dyn FieldEncodingStrategy, + options: &EncodingOptions, + ) -> Result { + let mut col_idx = 0; + let mut col_idx_sequence = ColumnIndexSequence::default(); + let field_encoders = schema + .fields + .iter() + .map(|field| { + let context = FieldEncodingContext { + strategy, + options, + root_field_metadata: &field.metadata, + }; + let encoder = + strategy.create_field_encoder(field, &mut col_idx_sequence, &context)?; + col_idx += encoder.as_ref().num_columns(); + Ok(encoder) + }) + .collect::>>()?; + Ok(Self { + field_encoders, + field_id_to_column_index: col_idx_sequence.mapping, + }) + } + + pub fn num_columns(&self) -> u32 { + self.field_encoders + .iter() + .map(|field_encoder| field_encoder.num_columns()) + .sum::() + } +} + +/// An encoded batch of data and a page table describing it +/// +/// This is returned by [`crate::encoder::encode_batch`] +#[derive(Debug)] +pub struct EncodedBatch { + pub data: Bytes, + pub page_table: Vec>, + pub schema: Arc, + pub top_level_columns: Vec, + pub num_rows: u64, +} + +fn write_page_to_data_buffer(page: EncodedPage, data_buffer: &mut BytesMut) -> PageInfo { + let buffers = page.data; + let mut buffer_offsets_and_sizes = Vec::with_capacity(buffers.len()); + for buffer in buffers { + let buffer_offset = data_buffer.len() as u64; + data_buffer.extend_from_slice(&buffer); + let size = data_buffer.len() as u64 - buffer_offset; + buffer_offsets_and_sizes.push((buffer_offset, size)); + } + + PageInfo { + buffer_offsets_and_sizes: Arc::from(buffer_offsets_and_sizes), + encoding: page.description, + num_rows: page.num_rows, + priority: page.row_number, + } +} + +/// Helper method to encode a batch of data into memory +/// +/// This is primarily for testing and benchmarking but could be useful in other +/// niche situations like IPC. +pub async fn encode_batch( + batch: &RecordBatch, + schema: Arc, + encoding_strategy: &dyn FieldEncodingStrategy, + options: &EncodingOptions, +) -> Result { + if !is_pwr_two(options.buffer_alignment) || options.buffer_alignment < MIN_PAGE_BUFFER_ALIGNMENT + { + return Err(Error::invalid_input_source( + format!( + "buffer_alignment must be a power of two and at least {}", + MIN_PAGE_BUFFER_ALIGNMENT + ) + .into(), + )); + } + + let mut data_buffer = BytesMut::new(); + let lance_schema = Schema::try_from(batch.schema().as_ref())?; + let options = EncodingOptions { + keep_original_array: true, + ..*options + }; + let mut batch_encoder = BatchEncoder::try_new(&lance_schema, encoding_strategy, &options)?; + let arrays = batch + .columns() + .iter() + .cloned() + .zip(batch_encoder.field_encoders.iter_mut()) + .map(|(array, encoder)| encoder.prepare_array(array)) + .collect::>>()?; + let mut page_table = Vec::new(); + let mut col_idx_offset = 0; + for (arr, mut encoder) in arrays.into_iter().zip(batch_encoder.field_encoders) { + let mut external_buffers = + OutOfLineBuffers::new(data_buffer.len() as u64, options.buffer_alignment); + let repdef = RepDefBuilder::default(); + let encoder = encoder.as_mut(); + let num_rows = arr.len() as u64; + let mut tasks = encoder.maybe_encode(arr, &mut external_buffers, repdef, 0, num_rows)?; + tasks.extend(encoder.flush(&mut external_buffers)?); + for buffer in external_buffers.take_buffers() { + data_buffer.extend_from_slice(&buffer); + } + let mut pages = HashMap::>::new(); + for task in tasks { + let encoded_page = task.await?; + // Write external buffers first + pages + .entry(encoded_page.column_idx) + .or_default() + .push(write_page_to_data_buffer(encoded_page, &mut data_buffer)); + } + let mut external_buffers = + OutOfLineBuffers::new(data_buffer.len() as u64, options.buffer_alignment); + let encoded_columns = encoder.finish(&mut external_buffers).await?; + for buffer in external_buffers.take_buffers() { + data_buffer.extend_from_slice(&buffer); + } + let num_columns = encoded_columns.len(); + for (col_idx, encoded_column) in encoded_columns.into_iter().enumerate() { + let col_idx = col_idx + col_idx_offset; + let mut col_buffer_offsets_and_sizes = Vec::new(); + for buffer in encoded_column.column_buffers { + let buffer_offset = data_buffer.len() as u64; + data_buffer.extend_from_slice(&buffer); + let size = data_buffer.len() as u64 - buffer_offset; + col_buffer_offsets_and_sizes.push((buffer_offset, size)); + } + for page in encoded_column.final_pages { + pages + .entry(page.column_idx) + .or_default() + .push(write_page_to_data_buffer(page, &mut data_buffer)); + } + let col_pages = std::mem::take(pages.entry(col_idx as u32).or_default()); + page_table.push(Arc::new(ColumnInfo { + index: col_idx as u32, + buffer_offsets_and_sizes: Arc::from( + col_buffer_offsets_and_sizes.into_boxed_slice(), + ), + page_infos: Arc::from(col_pages.into_boxed_slice()), + encoding: encoded_column.encoding, + })) + } + col_idx_offset += num_columns; + } + let top_level_columns = batch_encoder + .field_id_to_column_index + .iter() + .map(|(_, idx)| *idx) + .collect(); + Ok(EncodedBatch { + data: data_buffer.freeze(), + top_level_columns, + page_table, + schema, + num_rows: batch.num_rows() as u64, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::testing::{TestEncoding, create_test_field_encoder, test_encoding_strategy}; + use arrow_array::make_array; + use arrow_buffer::Buffer; + use arrow_data::ArrayData; + use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Fields as ArrowFields}; + use rstest::rstest; + + #[rstest] + fn test_nested_variable_width_offsets_are_validated_before_dispatch( + #[values(TestEncoding::Array, TestEncoding::StructuralU32)] encoding: TestEncoding, + #[values(ArrowDataType::Utf8, ArrowDataType::LargeUtf8)] item_type: ArrowDataType, + ) { + let offsets = match &item_type { + ArrowDataType::Utf8 => Buffer::from_slice_ref([0_i32, 2, 1, 3]), + ArrowDataType::LargeUtf8 => Buffer::from_slice_ref([0_i64, 2, 1, 3]), + _ => unreachable!(), + }; + let child_data = unsafe { + ArrayData::builder(item_type.clone()) + .len(3) + .add_buffer(offsets) + .add_buffer(Buffer::from(b"abc")) + .build_unchecked() + }; + let item_field = Arc::new(ArrowField::new("item", item_type, false)); + let data_type = ArrowDataType::FixedSizeList(item_field, 1); + let array_data = unsafe { + ArrayData::builder(data_type.clone()) + .len(3) + .add_child_data(child_data) + .build_unchecked() + }; + let array = make_array(array_data); + let field = Field::try_from(&ArrowField::new("payload", data_type, false)).unwrap(); + let strategy = test_encoding_strategy(encoding); + let mut column_index = ColumnIndexSequence::default(); + let options = EncodingOptions { + cache_bytes_per_column: 0, + ..Default::default() + }; + let mut encoder = + create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options) + .unwrap(); + let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT); + + let error = encoder + .maybe_encode(array, &mut external_buffers, RepDefBuilder::default(), 0, 3) + .err() + .expect("malformed nested offsets should fail before task dispatch"); + + assert!(matches!(error, Error::InvalidInput { .. })); + let message = error.to_string(); + assert!( + message.contains("field 'payload'"), + "unexpected message: {message}" + ); + assert!( + message.contains("non-monotonic offset at position 2"), + "unexpected message: {message}" + ); + } + + #[test] + fn test_fixed_size_list_struct_requires_v2_2() { + let list_item = ArrowField::new( + "item", + ArrowDataType::Struct(ArrowFields::from(vec![ArrowField::new( + "x", + ArrowDataType::Int32, + true, + )])), + true, + ); + let arrow_field = ArrowField::new( + "list_struct", + ArrowDataType::FixedSizeList(Arc::new(list_item), 2), + true, + ); + let field = Field::try_from(&arrow_field).unwrap(); + + let strategy = test_encoding_strategy(TestEncoding::StructuralU16); + let mut column_index = ColumnIndexSequence::default(); + let options = EncodingOptions::default(); + + let result = + create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options); + assert!( + result.is_err(), + "FixedSizeList should be rejected for file version 2.1" + ); + let err = result.err().unwrap(); + + assert!( + err.to_string() + .contains("FixedSizeList is not enabled by the selected file format") + ); + } +} diff --git a/vendor/lance-encoding/src/encoder/structural.rs b/vendor/lance-encoding/src/encoder/structural.rs new file mode 100644 index 000000000..3d4e49f42 --- /dev/null +++ b/vendor/lance-encoding/src/encoder/structural.rs @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Version-free structural field encoder builders. + +use std::sync::Arc; + +use arrow_schema::DataType; +use lance_core::{Error, Result, datatypes::Field, error::LanceOptionExt}; + +pub use crate::encodings::logical::primitive::PrimitivePageEncoding; + +use crate::encodings::logical::{ + blob::{BlobStructuralEncoder, BlobV2StructuralEncoder}, + fixed_size_list::FixedSizeListStructuralEncoder, + list::ListStructuralEncoder, + map::MapStructuralEncoder, + primitive::PrimitiveStructuralEncoder, + r#struct::StructStructuralEncoder, +}; + +use super::{ColumnIndexSequence, FieldEncoder, FieldEncodingContext}; + +/// Encode primitive leaves, primitive fixed-size lists, dictionaries, and +/// packed or empty structs using one concrete primitive page grammar. +#[derive(Debug, Clone)] +pub struct PrimitiveFieldEncoding { + page_encodings: Arc<[PrimitivePageEncoding]>, +} + +impl PrimitiveFieldEncoding { + /// Create a primitive field mechanism from ordered executable page behaviors. + pub fn new(page_encodings: impl IntoIterator) -> Self { + Self { + page_encodings: page_encodings.into_iter().collect(), + } + } + + fn is_primitive_type(data_type: &DataType) -> bool { + match data_type { + DataType::FixedSizeList(inner, _) => Self::is_primitive_type(inner.data_type()), + _ => matches!( + data_type, + DataType::Boolean + | DataType::Date32 + | DataType::Date64 + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Duration(_) + | DataType::Float16 + | DataType::Float32 + | DataType::Float64 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Int8 + | DataType::Interval(_) + | DataType::Null + | DataType::Time32(_) + | DataType::Time64(_) + | DataType::Timestamp(_, _) + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::UInt8 + | DataType::FixedSizeBinary(_) + | DataType::Binary + | DataType::LargeBinary + | DataType::Utf8 + | DataType::LargeUtf8, + ), + } + } + + fn create_at( + &self, + field: Field, + column_index: u32, + context: &FieldEncodingContext<'_>, + ) -> Result> { + Ok(Box::new(PrimitiveStructuralEncoder::try_new( + context.options, + self.page_encodings.clone(), + column_index, + field, + Arc::new(context.root_field_metadata.clone()), + )?)) + } + + /// Create a primitive field encoder when this mechanism recognizes `field`. + pub fn try_create( + &self, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result>> { + if field.is_blob() { + return Ok(None); + } + + let data_type = field.data_type(); + let is_primitive = Self::is_primitive_type(&data_type); + let is_packed_or_empty_struct = matches!( + &data_type, + DataType::Struct(fields) if field.is_packed_struct() || fields.is_empty() + ); + let is_primitive_dictionary = matches!( + &data_type, + DataType::Dictionary(_, value_type) if Self::is_primitive_type(value_type) + ); + + if !is_primitive && !is_packed_or_empty_struct && !is_primitive_dictionary { + if let DataType::Dictionary(_, value_type) = data_type { + return Err(Error::not_supported_source( + format!( + "cannot encode a dictionary column whose value type is a logical type ({})", + value_type + ) + .into(), + )); + } + return Ok(None); + } + + Ok(Some(self.create_at( + field.clone(), + column_index.next_column_index(field.id as u32), + context, + )?)) + } +} + +/// Create the original binary blob descriptor when `field` matches. +pub fn try_create_binary_blob( + primitive: &PrimitiveFieldEncoding, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + if !field.is_blob() || !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) { + return Ok(None); + } + let descriptor_column_index = column_index.next_column_index(field.id as u32); + Ok(Some(Box::new(BlobStructuralEncoder::new( + field, + |descriptor_field| primitive.create_at(descriptor_field, descriptor_column_index, context), + )?))) +} + +/// Create the structural blob descriptor when `field` matches. +pub fn try_create_structural_blob( + primitive: &PrimitiveFieldEncoding, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + if !field.is_blob() || !matches!(field.data_type(), DataType::Struct(_)) { + return Ok(None); + } + let descriptor_column_index = column_index.next_column_index(field.id as u32); + Ok(Some(Box::new(BlobV2StructuralEncoder::new( + field, + |descriptor_field| primitive.create_at(descriptor_field, descriptor_column_index, context), + )?))) +} + +/// Create a variable-size list encoder when `field` matches. +pub fn try_create_list( + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + if !matches!( + field.data_type(), + DataType::List(_) | DataType::LargeList(_) + ) { + return Ok(None); + } + let child = field.children.first().expect_ok()?; + let child_encoder = context + .strategy + .create_field_encoder(child, column_index, context)?; + Ok(Some(Box::new(ListStructuralEncoder::new( + context.options.keep_original_array, + child_encoder, + )))) +} + +/// Create a fixed-size-list encoder whose child is a struct when applicable. +pub fn try_create_structural_fixed_size_list( + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + if !matches!( + field.data_type(), + DataType::FixedSizeList(inner, _) if matches!(inner.data_type(), DataType::Struct(_)) + ) { + return Ok(None); + } + let child = field.children.first().expect_ok()?; + let child_encoder = context + .strategy + .create_field_encoder(child, column_index, context)?; + Ok(Some(Box::new(FixedSizeListStructuralEncoder::new( + context.options.keep_original_array, + child_encoder, + )))) +} + +/// Create an Arrow map encoder when `field` matches. +pub fn try_create_map( + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + let DataType::Map(_, keys_sorted) = field.data_type() else { + return Ok(None); + }; + if keys_sorted { + return Err(Error::not_supported_source( + format!( + "Map data type is not supported with keys_sorted=true now, current value is {}", + keys_sorted + ) + .into(), + )); + } + let entries_child = field + .children + .first() + .ok_or_else(|| Error::schema("Map should have an entries child".to_string()))?; + let DataType::Struct(struct_fields) = entries_child.data_type() else { + return Err(Error::schema( + "Map entries field must be a Struct".to_string(), + )); + }; + if struct_fields.len() < 2 { + return Err(Error::schema( + "Map entries struct must contain both key and value fields".to_string(), + )); + } + let key_field = &struct_fields[0]; + if key_field.is_nullable() { + return Err(Error::schema(format!( + "Map key field '{}' must be non-nullable according to Arrow Map specification", + key_field.name() + ))); + } + let child_encoder = + context + .strategy + .create_field_encoder(entries_child, column_index, context)?; + Ok(Some(Box::new(MapStructuralEncoder::new( + context.options.keep_original_array, + child_encoder, + )))) +} + +/// Create a non-packed, non-empty struct encoder when `field` matches. +pub fn try_create_struct( + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + let DataType::Struct(fields) = field.data_type() else { + return Ok(None); + }; + if field.is_blob() || field.is_packed_struct() || fields.is_empty() { + return Ok(None); + } + let children_encoders = field + .children + .iter() + .map(|child| { + context + .strategy + .create_field_encoder(child, column_index, context) + }) + .collect::>>()?; + Ok(Some(Box::new(StructStructuralEncoder::new( + context.options.keep_original_array, + children_encoders, + )))) +} diff --git a/vendor/lance-encoding/src/encodings.rs b/vendor/lance-encoding/src/encodings.rs new file mode 100644 index 000000000..d3ae90286 --- /dev/null +++ b/vendor/lance-encoding/src/encodings.rs @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +pub mod logical; +pub mod physical; + +#[cfg(test)] +pub mod fuzz_tests; diff --git a/vendor/lance-encoding/src/encodings/fuzz_tests.rs b/vendor/lance-encoding/src/encodings/fuzz_tests.rs new file mode 100644 index 000000000..7fb2cb123 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/fuzz_tests.rs @@ -0,0 +1,559 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Comprehensive fuzz testing for Lance 2.1 encoding coverage +//! +//! This module implements property-based testing using proptest to ensure +//! correct behavior across 16 different encoding permutations as specified +//! in issue #3347. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::builder::{Int32Builder, ListBuilder}; +use arrow_array::*; +use arrow_schema::{DataType, Field}; +use proptest::prelude::*; +use proptest::test_runner::{Config, TestRunner}; + +use crate::testing::{TestCases, TestEncoding, check_round_trip_encoding_of_data}; +use lance_core::Result; +use lance_datagen::{ArrayGenerator, ByteCount, Dimension, RowCount, Seed, array, gen_batch}; + +/// Test configuration representing one of the 16 permutations +#[derive(Debug, Clone)] +struct EncodingTestConfig { + encoding_type: EncodingType, + data_structure: DataStructure, + data_width: DataWidth, + nullable: bool, +} + +#[derive(Debug, Clone, PartialEq)] +enum EncodingType { + Miniblock, + FullZip, +} + +#[derive(Debug, Clone, PartialEq)] +enum DataStructure { + Primitive, + List, + FixedSizeList(i32), +} + +#[derive(Debug, Clone)] +enum DataWidth { + Fixed(FixedWidthType), + Variable(VariableWidthType), +} + +#[derive(Debug, Clone)] +enum FixedWidthType { + Int8, + Int16, + Int32, + Int64, + UInt8, + UInt16, + UInt32, + UInt64, + Float16, + Float32, + Float64, + FixedBinary(i32), +} + +#[derive(Debug, Clone)] +enum VariableWidthType { + String, + Binary, + LargeString, + LargeBinary, +} + +impl EncodingTestConfig { + /// Create field based on configuration + fn to_field(&self, name: &str) -> Field { + let data_type = self.to_data_type(); + Field::new(name, data_type, self.nullable) + } + + /// Get the Arrow data type for this configuration + fn to_data_type(&self) -> DataType { + let base_type = match &self.data_width { + DataWidth::Fixed(fixed) => match fixed { + FixedWidthType::Int8 => DataType::Int8, + FixedWidthType::Int16 => DataType::Int16, + FixedWidthType::Int32 => DataType::Int32, + FixedWidthType::Int64 => DataType::Int64, + FixedWidthType::UInt8 => DataType::UInt8, + FixedWidthType::UInt16 => DataType::UInt16, + FixedWidthType::UInt32 => DataType::UInt32, + FixedWidthType::UInt64 => DataType::UInt64, + FixedWidthType::Float16 => DataType::Float16, + FixedWidthType::Float32 => DataType::Float32, + FixedWidthType::Float64 => DataType::Float64, + FixedWidthType::FixedBinary(size) => DataType::FixedSizeBinary(*size), + }, + DataWidth::Variable(var) => match var { + VariableWidthType::String => DataType::Utf8, + VariableWidthType::Binary => DataType::Binary, + VariableWidthType::LargeString => DataType::LargeUtf8, + VariableWidthType::LargeBinary => DataType::LargeBinary, + }, + }; + + match &self.data_structure { + DataStructure::Primitive => base_type, + DataStructure::List => DataType::List(Arc::new(Field::new("item", base_type, true))), + DataStructure::FixedSizeList(size) => { + DataType::FixedSizeList(Arc::new(Field::new("item", base_type, true)), *size) + } + } + } +} + +/// Generate all valid test configurations (excluding FSL+List combinations) +fn encoding_config_strategy() -> impl Strategy { + ( + prop_oneof![Just(EncodingType::Miniblock), Just(EncodingType::FullZip)], + prop_oneof![ + Just(DataStructure::Primitive), + Just(DataStructure::List), + (1..=100i32).prop_map(DataStructure::FixedSizeList) + ], + prop_oneof![ + fixed_width_strategy().prop_map(DataWidth::Fixed), + variable_width_strategy().prop_map(DataWidth::Variable) + ], + any::(), // nullable + ) + .prop_filter( + "Skip unsupported combinations", + |(_, structure, width, _)| { + // FSL with Variable width types is not supported + // This combination causes compute_stat to be called twice + !matches!( + (structure, width), + (DataStructure::FixedSizeList(_), DataWidth::Variable(_)) + ) + }, + ) + .prop_map( + |(encoding, structure, width, nullable)| EncodingTestConfig { + encoding_type: encoding, + data_structure: structure, + data_width: width, + nullable, + }, + ) +} + +fn fixed_width_strategy() -> impl Strategy { + prop_oneof![ + Just(FixedWidthType::Int8), + Just(FixedWidthType::Int16), + Just(FixedWidthType::Int32), + Just(FixedWidthType::Int64), + Just(FixedWidthType::UInt8), + Just(FixedWidthType::UInt16), + Just(FixedWidthType::UInt32), + Just(FixedWidthType::UInt64), + Just(FixedWidthType::Float16), + Just(FixedWidthType::Float32), + Just(FixedWidthType::Float64), + (1..=128i32).prop_map(FixedWidthType::FixedBinary) + ] +} + +fn variable_width_strategy() -> impl Strategy { + prop_oneof![ + Just(VariableWidthType::String), + Just(VariableWidthType::Binary), + Just(VariableWidthType::LargeString), + Just(VariableWidthType::LargeBinary) + ] +} + +/// Generate test data based on configuration +fn generate_test_data_for_config( + config: &EncodingTestConfig, + num_rows: usize, + seed: u64, +) -> Result> { + let mut batch_gen = gen_batch().with_seed(Seed::from(seed)); + + // Generate base values + let generator: Box = match &config.data_width { + DataWidth::Fixed(fixed) => match fixed { + FixedWidthType::Int8 => array::rand_type(&DataType::Int8), + FixedWidthType::Int16 => array::rand_type(&DataType::Int16), + FixedWidthType::Int32 => array::rand_type(&DataType::Int32), + FixedWidthType::Int64 => array::rand_type(&DataType::Int64), + FixedWidthType::UInt8 => array::rand_type(&DataType::UInt8), + FixedWidthType::UInt16 => array::rand_type(&DataType::UInt16), + FixedWidthType::UInt32 => array::rand_type(&DataType::UInt32), + FixedWidthType::UInt64 => array::rand_type(&DataType::UInt64), + FixedWidthType::Float16 => array::rand_type(&DataType::Float16), + FixedWidthType::Float32 => array::rand_type(&DataType::Float32), + FixedWidthType::Float64 => array::rand_type(&DataType::Float64), + FixedWidthType::FixedBinary(size) => { + array::rand_type(&DataType::FixedSizeBinary(*size)) + } + }, + DataWidth::Variable(var) => { + // Control string/binary size based on encoding type + let max_len = if config.encoding_type == EncodingType::Miniblock { + 10 // Small strings for miniblock + } else { + 1000 // Large strings for full-zip + }; + + match var { + VariableWidthType::String => { + array::rand_utf8(ByteCount::from(max_len as u64), false) + } + VariableWidthType::Binary => { + array::rand_varbin(ByteCount::from(1), ByteCount::from(max_len as u64)) + } + VariableWidthType::LargeString => { + array::rand_utf8(ByteCount::from(max_len as u64), true) + } + VariableWidthType::LargeBinary => { + // For large binary, just use regular binary with larger size + array::rand_varbin(ByteCount::from(1), ByteCount::from(max_len as u64)) + } + } + } + }; + + // Wrap in list structure if needed + let final_generator: Box = match &config.data_structure { + DataStructure::Primitive => generator, + DataStructure::List => { + // Use rand_list_any for wrapping any generator + array::rand_list_any(generator, false) + } + DataStructure::FixedSizeList(size) => { + // Use cycle_vec for fixed size lists + array::cycle_vec(generator, Dimension::from(*size as u32)) + } + }; + + // Add nulls if configured + if config.nullable { + batch_gen.with_random_nulls(0.1); // 10% null rate + } + + let batch = batch_gen + .anon_col(final_generator) + .into_batch_rows(RowCount::from(num_rows as u64))?; + + Ok(batch.column(0).clone()) +} + +#[rstest::rstest] +fn test_encoding_round_trip( + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values( + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, + 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, + 48, 49 + )] + _shard: usize, +) { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let strategy = (encoding_config_strategy(), 100..=1000usize, any::()); + let mut runner = TestRunner::new(Config { + cases: 1, + ..Config::default() + }); + runner + .run(&strategy, |(config, num_rows, seed)| { + rt.block_on(async { + let test_data = generate_test_data_for_config(&config, num_rows, seed) + .expect("Failed to generate test data"); + + let _field = config.to_field("test"); + + let mut metadata = HashMap::new(); + if config.encoding_type == EncodingType::Miniblock { + metadata.insert("encoding_hint".to_string(), "miniblock".to_string()); + } + + let test_cases = TestCases::default() + .with_encoding(encoding) + .with_batch_size(100) + .with_range(0..num_rows.min(500) as u64) + .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]); + + check_round_trip_encoding_of_data(vec![test_data], &test_cases, metadata).await; + }); + Ok(()) + }) + .unwrap(); +} + +#[tokio::test] +async fn test_edge_cases_single_value() { + // Test single value arrays + let single_int32 = Arc::new(Int32Array::from(vec![42])) as Arc; + let single_string = Arc::new(StringArray::from(vec!["test"])) as Arc; + + let test_cases = TestCases::default().with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![single_int32], &test_cases, HashMap::new()).await; + + check_round_trip_encoding_of_data(vec![single_string], &test_cases, HashMap::new()).await; +} + +#[tokio::test] +async fn test_edge_cases_all_nulls() { + // Test arrays with all null values + let all_nulls_int32 = + Arc::new(Int32Array::from(vec![None, None, None] as Vec>)) as Arc; + let all_nulls_string = Arc::new(StringArray::from( + vec![None, None, None] as Vec> + )) as Arc; + + let test_cases = TestCases::default().with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![all_nulls_int32], &test_cases, HashMap::new()).await; + + check_round_trip_encoding_of_data(vec![all_nulls_string], &test_cases, HashMap::new()).await; +} + +// Test list with repetition and definition levels +proptest! { + #[test] + fn test_list_repdef_handling( + list_sizes in prop::collection::vec(1..=20usize, 10..=100), + _seed in any::() + ) { + let rt = tokio::runtime::Runtime::new().unwrap(); + + rt.block_on(async { + // Create a list array with varying sizes + let mut list_builder = ListBuilder::new(Int32Builder::new()); + + for size in &list_sizes { + for i in 0..*size { + list_builder.values().append_value(i as i32); + } + list_builder.append(true); + } + + let list_array = Arc::new(list_builder.finish()) as Arc; + + let test_cases = TestCases::default() + .with_structural_encodings() + .with_range(0..list_sizes.len().min(50) as u64); + + check_round_trip_encoding_of_data( + vec![list_array], + &test_cases, + HashMap::new() + ).await; + }); + } +} + +#[rstest::rstest] +fn test_fixed_size_list_encoding( + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, +) { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let strategy = (1..=100i32, 10..=1000usize, any::()); + let mut runner = TestRunner::new(Config::default()); + runner + .run(&strategy, |(list_size, num_rows, seed)| { + rt.block_on(async { + let config = EncodingTestConfig { + encoding_type: EncodingType::Miniblock, + data_structure: DataStructure::FixedSizeList(list_size), + data_width: DataWidth::Fixed(FixedWidthType::Int32), + nullable: false, + }; + + let test_data = generate_test_data_for_config(&config, num_rows, seed) + .expect("Failed to generate test data"); + + let test_cases = TestCases::default().with_encoding(encoding); + + check_round_trip_encoding_of_data(vec![test_data], &test_cases, HashMap::new()) + .await; + }); + Ok(()) + }) + .unwrap(); +} + +#[tokio::test] +async fn test_list_dict_empty_batch() { + use arrow_array::builder::BinaryBuilder; + use arrow_array::builder::ListBuilder; + + // Create a list with some values followed by empty/null lists + let mut list_builder = ListBuilder::new(BinaryBuilder::new()); + + // First 50 lists have values with LOW CARDINALITY to trigger dictionary encoding + // Only 5 unique values repeated many times (150 total values, 5 unique) + let values = [b"aaaaa", b"bbbbb", b"ccccc", b"ddddd", b"eeeee"]; + for i in 0..50 { + // Each list has 3 values, cycling through the 5 unique values + list_builder.append_value([ + Some(values[i % 5]), + Some(values[(i + 1) % 5]), + Some(values[(i + 2) % 5]), + ]); + } + + // Next 50 lists are empty or null (no values) + for i in 0..50 { + if i % 2 == 0 { + list_builder.append_value(Vec::>::new()); // empty list + } else { + list_builder.append_null(); // null list + } + } + + let list_array = Arc::new(list_builder.finish()); + + let test_cases = TestCases::default() + .with_structural_encodings() + // Read only the empty/null lists (rows 50-99) + // This batch will have 0 underlying values + .with_range(50..100); + + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; +} + +// Test all valid combinations systematically (14 combinations) +// Excludes FSL with Variable width types which are not supported +#[tokio::test] +async fn test_all_valid_combinations() { + let combinations = vec![ + // Basic primitive types (4 combinations) + ( + EncodingType::Miniblock, + DataStructure::Primitive, + DataWidth::Fixed(FixedWidthType::Int32), + false, + ), + ( + EncodingType::Miniblock, + DataStructure::Primitive, + DataWidth::Variable(VariableWidthType::String), + false, + ), + ( + EncodingType::FullZip, + DataStructure::Primitive, + DataWidth::Fixed(FixedWidthType::Float64), + false, + ), + ( + EncodingType::FullZip, + DataStructure::Primitive, + DataWidth::Variable(VariableWidthType::Binary), + false, + ), + // List types (4 combinations) + ( + EncodingType::Miniblock, + DataStructure::List, + DataWidth::Fixed(FixedWidthType::Int32), + false, + ), + ( + EncodingType::Miniblock, + DataStructure::List, + DataWidth::Variable(VariableWidthType::String), + false, + ), + ( + EncodingType::FullZip, + DataStructure::List, + DataWidth::Fixed(FixedWidthType::Float64), + false, + ), + ( + EncodingType::FullZip, + DataStructure::List, + DataWidth::Variable(VariableWidthType::Binary), + false, + ), + // Fixed size list types (2 combinations - only with Fixed width types) + // Note: FSL with Variable width types is not supported + ( + EncodingType::Miniblock, + DataStructure::FixedSizeList(10), + DataWidth::Fixed(FixedWidthType::Int32), + false, + ), + ( + EncodingType::FullZip, + DataStructure::FixedSizeList(10), + DataWidth::Fixed(FixedWidthType::Float64), + false, + ), + // Nullable variants (4 combinations) + ( + EncodingType::Miniblock, + DataStructure::Primitive, + DataWidth::Fixed(FixedWidthType::Int32), + true, + ), + ( + EncodingType::Miniblock, + DataStructure::Primitive, + DataWidth::Variable(VariableWidthType::String), + true, + ), + ( + EncodingType::FullZip, + DataStructure::Primitive, + DataWidth::Fixed(FixedWidthType::Float64), + true, + ), + ( + EncodingType::FullZip, + DataStructure::Primitive, + DataWidth::Variable(VariableWidthType::Binary), + true, + ), + ]; + + for (encoding, structure, width, nullable) in combinations { + let config = EncodingTestConfig { + encoding_type: encoding, + data_structure: structure, + data_width: width, + nullable, + }; + + let test_data = + generate_test_data_for_config(&config, 100, 42).expect("Failed to generate test data"); + + let test_cases = TestCases::default().with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![test_data], &test_cases, HashMap::new()).await; + } +} diff --git a/vendor/lance-encoding/src/encodings/logical.rs b/vendor/lance-encoding/src/encodings/logical.rs new file mode 100644 index 000000000..199f470f5 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical.rs @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +pub mod blob; +pub mod fixed_size_list; +pub mod list; +pub mod map; +pub mod primitive; +pub mod r#struct; diff --git a/vendor/lance-encoding/src/encodings/logical/blob.rs b/vendor/lance-encoding/src/encodings/logical/blob.rs new file mode 100644 index 000000000..a5432fdc5 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/blob.rs @@ -0,0 +1,904 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::HashMap, sync::Arc}; + +use arrow_array::{ + Array, ArrayRef, StructArray, UInt64Array, + builder::{PrimitiveBuilder, StringBuilder}, + cast::AsArray, + types::{UInt8Type, UInt32Type, UInt64Type}, +}; +use arrow_buffer::Buffer; +use arrow_schema::{DataType, Field as ArrowField, Fields}; +use futures::{FutureExt, future::BoxFuture}; +use lance_core::{ + Error, Result, + datatypes::{BLOB_V2_DESC_FIELDS, BlobV2Layout, Field}, + error::LanceOptionExt, +}; + +use crate::{ + buffer::LanceBuffer, + constants::PACKED_STRUCT_META_KEY, + decoder::PageEncoding, + encoder::{EncodeTask, EncodedColumn, EncodedPage, FieldEncoder, OutOfLineBuffers}, + format::ProtobufUtils21, + repdef::{DefinitionInterpretation, RepDefBuilder}, +}; +use lance_core::datatypes::BlobKind; + +/// Blob structural encoder - stores large binary data in external buffers +/// +/// This encoder takes large binary arrays and stores them outside the normal +/// page structure. It creates a descriptor (position, size) for each blob +/// that is stored inline in the page. +pub struct BlobStructuralEncoder { + // Encoder for the descriptors (position/size struct) + descriptor_encoder: Box, + // Set when we first see data + def_meaning: Option>, +} + +impl BlobStructuralEncoder { + pub fn new( + field: &Field, + make_descriptor_encoder: impl FnOnce(Field) -> Result>, + ) -> Result { + // Create descriptor field: struct + // Preserve the original field's metadata for packed struct + let mut descriptor_metadata = HashMap::with_capacity(1); + descriptor_metadata.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string()); + + let descriptor_data_type = DataType::Struct(Fields::from(vec![ + ArrowField::new("position", DataType::UInt64, false), + ArrowField::new("size", DataType::UInt64, false), + ])); + + // Use the original field's name for the descriptor + let descriptor_field = Field::try_from( + ArrowField::new(&field.name, descriptor_data_type, field.nullable) + .with_metadata(descriptor_metadata), + )?; + + // Use PrimitiveStructuralEncoder to handle the descriptor + let descriptor_encoder = make_descriptor_encoder(descriptor_field)?; + + Ok(Self { + descriptor_encoder, + def_meaning: None, + }) + } + + fn wrap_tasks( + tasks: Vec, + def_meaning: Arc<[DefinitionInterpretation]>, + ) -> Vec { + tasks + .into_iter() + .map(|task| { + let def_meaning = def_meaning.clone(); + task.then(|encoded_page| async move { + let encoded_page = encoded_page?; + + let PageEncoding::Structural(inner_layout) = encoded_page.description else { + return Err(Error::internal( + "Expected inner encoding to return structural layout".to_string(), + )); + }; + + let wrapped = ProtobufUtils21::blob_layout(inner_layout, &def_meaning); + Ok(EncodedPage { + column_idx: encoded_page.column_idx, + data: encoded_page.data, + description: PageEncoding::Structural(wrapped), + num_rows: encoded_page.num_rows, + row_number: encoded_page.row_number, + }) + }) + .boxed() + }) + .collect::>() + } +} + +impl FieldEncoder for BlobStructuralEncoder { + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + mut repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + if let Some(validity) = array.nulls() { + repdef.add_validity_bitmap(validity.clone()); + } else { + repdef.add_no_null(array.len()); + } + + // Convert input array to LargeBinary + let binary_array = array.as_binary_opt::().ok_or_else(|| { + Error::invalid_input_source( + format!("Expected LargeBinary array, got {}", array.data_type()).into(), + ) + })?; + + let repdef = RepDefBuilder::serialize(vec![repdef]); + + let rep = repdef.repetition_levels.as_ref(); + let def = repdef.definition_levels.as_ref(); + let def_meaning: Arc<[DefinitionInterpretation]> = repdef.def_meaning.into(); + + // A blob page stores one definition interpretation for all of its rows. + // The descriptor encoder can buffer multiple input arrays, so finish the + // pending page before a later array changes from all-valid to nullable (or + // vice versa). + let mut encode_tasks = match self.def_meaning.as_ref() { + Some(existing) if existing != &def_meaning => { + let existing = existing.clone(); + Self::wrap_tasks(self.descriptor_encoder.flush(external_buffers)?, existing) + } + _ => Vec::new(), + }; + self.def_meaning = Some(def_meaning.clone()); + + // Collect positions and sizes + let mut positions = Vec::with_capacity(binary_array.len()); + let mut sizes = Vec::with_capacity(binary_array.len()); + + for i in 0..binary_array.len() { + if binary_array.is_null(i) { + // Null values are smuggled into the positions array + + // If we have null values we must have definition levels + let mut repdef = (def.expect_ok()?[i] as u64) << 16; + if let Some(rep) = rep { + repdef += rep[i] as u64; + } + + debug_assert_ne!(repdef, 0); + positions.push(repdef); + sizes.push(0); + } else { + let value = binary_array.value(i); + if value.is_empty() { + // Empty values + positions.push(0); + sizes.push(0); + } else { + // Add data to external buffers + let position = + external_buffers.add_buffer(LanceBuffer::from(Buffer::from(value))); + positions.push(position); + sizes.push(value.len() as u64); + } + } + } + + // Create descriptor array + let position_array = Arc::new(UInt64Array::from(positions)); + let size_array = Arc::new(UInt64Array::from(sizes)); + let descriptor_array = Arc::new(StructArray::new( + Fields::from(vec![ + ArrowField::new("position", DataType::UInt64, false), + ArrowField::new("size", DataType::UInt64, false), + ]), + vec![position_array as ArrayRef, size_array as ArrayRef], + None, // Descriptors are never null + )); + + // Delegate to descriptor encoder + let descriptor_tasks = self.descriptor_encoder.maybe_encode( + descriptor_array, + external_buffers, + RepDefBuilder::default(), + row_number, + num_rows, + )?; + encode_tasks.extend(Self::wrap_tasks(descriptor_tasks, def_meaning)); + + Ok(encode_tasks) + } + + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + let encode_tasks = self.descriptor_encoder.flush(external_buffers)?; + + // Use the cached def meaning. If we haven't seen any data yet then we can just use a dummy + // value (not clear there would be any encode tasks in that case) + let def_meaning = self + .def_meaning + .clone() + .unwrap_or_else(|| Arc::new([DefinitionInterpretation::AllValidItem])); + + Ok(Self::wrap_tasks(encode_tasks, def_meaning)) + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + self.descriptor_encoder.finish(external_buffers) + } + + fn num_columns(&self) -> u32 { + self.descriptor_encoder.num_columns() + } +} + +/// Blob v2 structural encoder +pub struct BlobV2StructuralEncoder { + descriptor_encoder: Box, +} + +impl BlobV2StructuralEncoder { + pub fn new( + field: &Field, + make_descriptor_encoder: impl FnOnce(Field) -> Result>, + ) -> Result { + let mut descriptor_metadata = HashMap::with_capacity(1); + descriptor_metadata.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string()); + + let descriptor_data_type = DataType::Struct(BLOB_V2_DESC_FIELDS.clone()); + + let descriptor_field = Field::try_from( + ArrowField::new(&field.name, descriptor_data_type, field.nullable) + .with_metadata(descriptor_metadata), + )?; + + let descriptor_encoder = make_descriptor_encoder(descriptor_field)?; + + Ok(Self { descriptor_encoder }) + } +} + +impl FieldEncoder for BlobV2StructuralEncoder { + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + let struct_arr = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Blob v2 encoder expected StructArray, got {}", + array.data_type() + ) + .into(), + ) + })?; + if BlobV2Layout::classify(struct_arr.fields()) != Some(BlobV2Layout::Prepared) { + let actual = BlobV2Layout::classify(struct_arr.fields()) + .map(|layout| layout.to_string()) + .unwrap_or_else(|| format!("unrecognized ({:?})", struct_arr.fields())); + return Err(Error::invalid_input_source( + format!("Blob v2 encoder expected prepared array layout, got {actual} layout") + .into(), + )); + } + + let kind_col = struct_arr + .column_by_name("kind") + .ok_or_else(|| { + Error::invalid_input_source("Blob v2 struct missing `kind` field".into()) + })? + .as_primitive::(); + let data_col = struct_arr + .column_by_name("data") + .ok_or_else(|| { + Error::invalid_input_source("Blob v2 struct missing `data` field".into()) + })? + .as_binary::(); + let uri_col = struct_arr + .column_by_name("uri") + .ok_or_else(|| { + Error::invalid_input_source("Blob v2 struct missing `uri` field".into()) + })? + .as_string::(); + let blob_id_col = struct_arr + .column_by_name("blob_id") + .ok_or_else(|| { + Error::invalid_input_source("Blob v2 struct missing `blob_id` field".into()) + })? + .as_primitive::(); + let blob_size_col = struct_arr + .column_by_name("blob_size") + .ok_or_else(|| { + Error::invalid_input_source("Blob v2 struct missing `blob_size` field".into()) + })? + .as_primitive::(); + let packed_position_col = struct_arr + .column_by_name("position") + .ok_or_else(|| { + Error::invalid_input_source("Blob v2 struct missing `position` field".into()) + })? + .as_primitive::(); + + let row_count = struct_arr.len(); + + let mut kind_builder = PrimitiveBuilder::::with_capacity(row_count); + let mut position_builder = PrimitiveBuilder::::with_capacity(row_count); + let mut size_builder = PrimitiveBuilder::::with_capacity(row_count); + let mut blob_id_builder = PrimitiveBuilder::::with_capacity(row_count); + let mut uri_builder = StringBuilder::with_capacity(row_count, row_count * 16); + + for i in 0..row_count { + let (kind_value, position_value, size_value, blob_id_value, uri_value) = + if struct_arr.is_null(i) || kind_col.is_null(i) { + (BlobKind::Inline as u8, 0, 0, 0, "".to_string()) + } else { + let kind_val = BlobKind::try_from(kind_col.value(i))?; + match kind_val { + BlobKind::Dedicated => ( + BlobKind::Dedicated as u8, + 0, + blob_size_col.value(i), + blob_id_col.value(i), + "".to_string(), + ), + BlobKind::External => { + let uri = uri_col.value(i).to_string(); + let position = if packed_position_col.is_null(i) { + 0 + } else { + packed_position_col.value(i) + }; + let size = if blob_size_col.is_null(i) { + 0 + } else { + blob_size_col.value(i) + }; + let external_base_id = if blob_id_col.is_null(i) { + 0 + } else { + blob_id_col.value(i) + }; + ( + BlobKind::External as u8, + position, + size, + external_base_id, + uri, + ) + } + BlobKind::Packed => ( + BlobKind::Packed as u8, + packed_position_col.value(i), + blob_size_col.value(i), + blob_id_col.value(i), + "".to_string(), + ), + BlobKind::Inline => { + let data_val = data_col.value(i); + let blob_len = data_val.len() as u64; + let position = external_buffers + .add_buffer(LanceBuffer::from(Buffer::from(data_val))); + + ( + BlobKind::Inline as u8, + position, + blob_len, + 0, + "".to_string(), + ) + } + } + }; + + kind_builder.append_value(kind_value); + position_builder.append_value(position_value); + size_builder.append_value(size_value); + blob_id_builder.append_value(blob_id_value); + uri_builder.append_value(uri_value); + } + let children: Vec = vec![ + Arc::new(kind_builder.finish()), + Arc::new(position_builder.finish()), + Arc::new(size_builder.finish()), + Arc::new(blob_id_builder.finish()), + Arc::new(uri_builder.finish()), + ]; + + let descriptor_array = Arc::new(StructArray::try_new( + BLOB_V2_DESC_FIELDS.clone(), + children, + struct_arr.nulls().cloned(), + )?) as ArrayRef; + + self.descriptor_encoder.maybe_encode( + descriptor_array, + external_buffers, + repdef, + row_number, + num_rows, + ) + } + + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + self.descriptor_encoder.flush(external_buffers) + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + self.descriptor_encoder.finish(external_buffers) + } + + fn num_columns(&self) -> u32 { + self.descriptor_encoder.num_columns() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + encoder::{ColumnIndexSequence, EncodingOptions}, + testing::{ + TestCases, TestEncoding, check_round_trip_encoding_of_data, + check_round_trip_encoding_of_data_with_expected, create_test_field_encoder, + test_encoding_strategy, + }, + }; + use arrow_array::{ + ArrayRef, LargeBinaryArray, StringArray, StructArray, UInt8Array, UInt32Array, UInt64Array, + }; + use arrow_schema::{DataType, Field as ArrowField}; + use lance_core::datatypes::BLOB_V2_LOGICAL_MINIMAL_FIELDS; + + #[test] + fn test_blob_encoder_creation() { + let field = Field::try_from( + ArrowField::new("blob_field", DataType::LargeBinary, true).with_metadata( + HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]), + ), + ) + .unwrap(); + let mut column_index = ColumnIndexSequence::default(); + let options = EncodingOptions::default(); + let strategy = test_encoding_strategy(TestEncoding::StructuralU16); + + let encoder = + create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options); + + assert!(encoder.is_ok()); + } + + #[test] + fn test_blob_v2_encoder_rejects_logical_array_layout() { + let field = Field::try_from( + ArrowField::new( + "blob_field", + DataType::Struct(BLOB_V2_LOGICAL_MINIMAL_FIELDS.clone()), + true, + ) + .with_metadata(HashMap::from([( + lance_arrow::ARROW_EXT_NAME_KEY.to_string(), + lance_arrow::BLOB_V2_EXT_NAME.to_string(), + )])), + ) + .unwrap(); + let mut column_index = ColumnIndexSequence::default(); + let options = EncodingOptions::default(); + let strategy = test_encoding_strategy(TestEncoding::StructuralU32); + let mut encoder = + create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options) + .unwrap(); + let array = Arc::new( + StructArray::try_new( + BLOB_V2_LOGICAL_MINIMAL_FIELDS.clone(), + vec![ + Arc::new(LargeBinaryArray::from(vec![Some(b"payload".as_ref())])) as ArrayRef, + Arc::new(StringArray::from(vec![None::<&str>])) as ArrayRef, + ], + None, + ) + .unwrap(), + ) as ArrayRef; + let mut external_buffers = OutOfLineBuffers::new(0, 8); + let Err(error) = + encoder.maybe_encode(array, &mut external_buffers, RepDefBuilder::default(), 0, 1) + else { + panic!("logical array layout unexpectedly reached the descriptor encoder"); + }; + assert!(matches!(error, Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("expected prepared array layout, got logical layout") + ); + } + + #[tokio::test] + async fn test_blob_encoding_simple() { + let field = Field::try_from( + ArrowField::new("blob_field", DataType::LargeBinary, true).with_metadata( + HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]), + ), + ) + .unwrap(); + let mut column_index = ColumnIndexSequence::default(); + let options = EncodingOptions::default(); + let strategy = test_encoding_strategy(TestEncoding::StructuralU16); + + let mut encoder = + create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options) + .unwrap(); + + // Create test data with larger blobs + let large_data = vec![0u8; 1024 * 100]; // 100KB blob + let data: Vec> = + vec![Some(b"hello world"), None, Some(&large_data), Some(b"")]; + let array = Arc::new(LargeBinaryArray::from(data)); + + // Test encoding + let mut external_buffers = OutOfLineBuffers::new(0, 8); + let repdef = RepDefBuilder::default(); + + let tasks = encoder + .maybe_encode(array, &mut external_buffers, repdef, 0, 4) + .unwrap(); + + // If no tasks yet, flush to force encoding + if tasks.is_empty() { + let _flush_tasks = encoder.flush(&mut external_buffers).unwrap(); + } + + // Should produce encode tasks for the descriptor (or we need more data) + // For now, just verify no errors occurred + assert!(encoder.num_columns() > 0); + + // Verify external buffers were used for large data + let buffers = external_buffers.take_buffers(); + assert!( + !buffers.is_empty(), + "Large blobs should be stored in external buffers" + ); + } + + #[tokio::test] + async fn test_blob_round_trip() { + // Test round-trip encoding with blob metadata + let blob_metadata = + HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]); + + // Create test data + let val1: &[u8] = &vec![1u8; 1024]; // 1KB + let val2: &[u8] = &vec![2u8; 10240]; // 10KB + let val3: &[u8] = &vec![3u8; 102400]; // 100KB + let array = Arc::new(LargeBinaryArray::from(vec![ + Some(val1), + None, + Some(val2), + Some(val3), + ])); + + // Use the standard test harness + check_round_trip_encoding_of_data( + vec![array], + &TestCases::default().with_array_and_u16_encodings(), + blob_metadata, + ) + .await; + } + + #[tokio::test] + async fn test_blob_round_trip_empty_values() { + // Empty values share size == 0 with nulls in the descriptor layout + // and schedule no read; each must decode to zero-length bytes without + // consuming the read result of a following non-empty blob. Empties + // are placed before payloads so a misassignment corrupts the output + // instead of only exhausting the read iterator. + let blob_metadata = + HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]); + + let val1: &[u8] = &vec![1u8; 1024]; + let val2: &[u8] = &vec![2u8; 10240]; + let empty: &[u8] = &[]; + let array = Arc::new(LargeBinaryArray::from(vec![ + Some(empty), + Some(val1), + None, + Some(empty), + Some(val2), + None, + Some(empty), + ])); + + check_round_trip_encoding_of_data(vec![array], &TestCases::default(), blob_metadata).await; + } + + #[tokio::test] + async fn test_blob_round_trip_varying_chunk_nullability() { + let blob_metadata = + HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]); + let all_valid = Arc::new(LargeBinaryArray::from(vec![Some(b"first".as_ref())])); + let with_null = Arc::new(LargeBinaryArray::from(vec![ + Some(b"second".as_ref()), + None, + Some(b"".as_ref()), + ])); + let all_valid_again = Arc::new(LargeBinaryArray::from(vec![Some(b"last".as_ref())])); + + check_round_trip_encoding_of_data( + vec![all_valid, with_null, all_valid_again], + &TestCases::default().with_encoding(TestEncoding::StructuralU16), + blob_metadata, + ) + .await; + } + + #[tokio::test] + async fn test_blob_v2_external_round_trip() { + let blob_metadata = HashMap::from([( + lance_arrow::ARROW_EXT_NAME_KEY.to_string(), + lance_arrow::BLOB_V2_EXT_NAME.to_string(), + )]); + + let kind_field = Arc::new(ArrowField::new("kind", DataType::UInt8, true)); + let data_field = Arc::new(ArrowField::new("data", DataType::LargeBinary, true)); + let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, true)); + let blob_id_field = Arc::new(ArrowField::new("blob_id", DataType::UInt32, true)); + let blob_size_field = Arc::new(ArrowField::new("blob_size", DataType::UInt64, true)); + let position_field = Arc::new(ArrowField::new("position", DataType::UInt64, true)); + + let kind_array = UInt8Array::from(vec![ + BlobKind::Inline as u8, + BlobKind::External as u8, + BlobKind::External as u8, + ]); + let data_array = LargeBinaryArray::from(vec![Some(b"inline".as_ref()), None, None]); + let uri_array = StringArray::from(vec![ + None, + Some("file:///tmp/external.bin"), + Some("s3://bucket/blob"), + ]); + let blob_id_array = UInt32Array::from(vec![0, 0, 0]); + let blob_size_array = UInt64Array::from(vec![0, 0, 0]); + let position_array = UInt64Array::from(vec![0, 0, 0]); + + let struct_array = StructArray::from(vec![ + (kind_field, Arc::new(kind_array) as ArrayRef), + (data_field, Arc::new(data_array) as ArrayRef), + (uri_field, Arc::new(uri_array) as ArrayRef), + (blob_id_field, Arc::new(blob_id_array) as ArrayRef), + (blob_size_field, Arc::new(blob_size_array) as ArrayRef), + (position_field, Arc::new(position_array) as ArrayRef), + ]); + + let expected_descriptor = StructArray::from(vec![ + ( + Arc::new(ArrowField::new("kind", DataType::UInt8, false)), + Arc::new(UInt8Array::from(vec![ + BlobKind::Inline as u8, + BlobKind::External as u8, + BlobKind::External as u8, + ])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("position", DataType::UInt64, false)), + Arc::new(UInt64Array::from(vec![0, 0, 0])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("size", DataType::UInt64, false)), + Arc::new(UInt64Array::from(vec![6, 0, 0])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("blob_id", DataType::UInt32, false)), + Arc::new(UInt32Array::from(vec![0, 0, 0])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("blob_uri", DataType::Utf8, false)), + Arc::new(StringArray::from(vec![ + "", + "file:///tmp/external.bin", + "s3://bucket/blob", + ])) as ArrayRef, + ), + ]); + + check_round_trip_encoding_of_data_with_expected( + vec![Arc::new(struct_array)], + Some(Arc::new(expected_descriptor)), + &TestCases::default().with_u32_structural_encodings(), + blob_metadata, + ) + .await; + } + + #[tokio::test] + async fn test_blob_v2_dedicated_round_trip() { + let blob_metadata = HashMap::from([( + lance_arrow::ARROW_EXT_NAME_KEY.to_string(), + lance_arrow::BLOB_V2_EXT_NAME.to_string(), + )]); + + let kind_field = Arc::new(ArrowField::new("kind", DataType::UInt8, true)); + let data_field = Arc::new(ArrowField::new("data", DataType::LargeBinary, true)); + let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, true)); + let blob_id_field = Arc::new(ArrowField::new("blob_id", DataType::UInt32, true)); + let blob_size_field = Arc::new(ArrowField::new("blob_size", DataType::UInt64, true)); + let position_field = Arc::new(ArrowField::new("position", DataType::UInt64, true)); + + let kind_array = UInt8Array::from(vec![BlobKind::Dedicated as u8, BlobKind::Inline as u8]); + let data_array = LargeBinaryArray::from(vec![None, Some(b"abc".as_ref())]); + let uri_array = StringArray::from(vec![Option::<&str>::None, None]); + let blob_id_array = UInt32Array::from(vec![42, 0]); + let blob_size_array = UInt64Array::from(vec![12, 0]); + let position_array = UInt64Array::from(vec![0, 0]); + + let struct_array = StructArray::from(vec![ + (kind_field, Arc::new(kind_array) as ArrayRef), + (data_field, Arc::new(data_array) as ArrayRef), + (uri_field, Arc::new(uri_array) as ArrayRef), + (blob_id_field, Arc::new(blob_id_array) as ArrayRef), + (blob_size_field, Arc::new(blob_size_array) as ArrayRef), + (position_field, Arc::new(position_array) as ArrayRef), + ]); + + let expected_descriptor = StructArray::from(vec![ + ( + Arc::new(ArrowField::new("kind", DataType::UInt8, false)), + Arc::new(UInt8Array::from(vec![ + BlobKind::Dedicated as u8, + BlobKind::Inline as u8, + ])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("position", DataType::UInt64, false)), + Arc::new(UInt64Array::from(vec![0, 0])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("size", DataType::UInt64, false)), + Arc::new(UInt64Array::from(vec![12, 3])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("blob_id", DataType::UInt32, false)), + Arc::new(UInt32Array::from(vec![42, 0])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("blob_uri", DataType::Utf8, false)), + Arc::new(StringArray::from(vec!["", ""])) as ArrayRef, + ), + ]); + + check_round_trip_encoding_of_data_with_expected( + vec![Arc::new(struct_array)], + Some(Arc::new(expected_descriptor)), + &TestCases::default().with_u32_structural_encodings(), + blob_metadata, + ) + .await; + } + + #[tokio::test] + async fn test_blob_v2_external_with_range_round_trip() { + let blob_metadata = HashMap::from([( + lance_arrow::ARROW_EXT_NAME_KEY.to_string(), + lance_arrow::BLOB_V2_EXT_NAME.to_string(), + )]); + + let kind_field = Arc::new(ArrowField::new("kind", DataType::UInt8, true)); + let data_field = Arc::new(ArrowField::new("data", DataType::LargeBinary, true)); + let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, true)); + let blob_id_field = Arc::new(ArrowField::new("blob_id", DataType::UInt32, true)); + let blob_size_field = Arc::new(ArrowField::new("blob_size", DataType::UInt64, true)); + let position_field = Arc::new(ArrowField::new("position", DataType::UInt64, true)); + + let kind_array = UInt8Array::from(vec![BlobKind::External as u8]); + let data_array = LargeBinaryArray::from(vec![None::<&[u8]>]); + let uri_array = StringArray::from(vec![Some("memory://container.pack")]); + let blob_id_array = UInt32Array::from(vec![0]); + let blob_size_array = UInt64Array::from(vec![42]); + let position_array = UInt64Array::from(vec![7]); + + let struct_array = StructArray::from(vec![ + (kind_field, Arc::new(kind_array) as ArrayRef), + (data_field, Arc::new(data_array) as ArrayRef), + (uri_field, Arc::new(uri_array) as ArrayRef), + (blob_id_field, Arc::new(blob_id_array) as ArrayRef), + (blob_size_field, Arc::new(blob_size_array) as ArrayRef), + (position_field, Arc::new(position_array) as ArrayRef), + ]); + + let expected_descriptor = StructArray::from(vec![ + ( + Arc::new(ArrowField::new("kind", DataType::UInt8, false)), + Arc::new(UInt8Array::from(vec![BlobKind::External as u8])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("position", DataType::UInt64, false)), + Arc::new(UInt64Array::from(vec![7])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("size", DataType::UInt64, false)), + Arc::new(UInt64Array::from(vec![42])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("blob_id", DataType::UInt32, false)), + Arc::new(UInt32Array::from(vec![0])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("blob_uri", DataType::Utf8, false)), + Arc::new(StringArray::from(vec!["memory://container.pack"])) as ArrayRef, + ), + ]); + + check_round_trip_encoding_of_data_with_expected( + vec![Arc::new(struct_array)], + Some(Arc::new(expected_descriptor)), + &TestCases::default().with_u32_structural_encodings(), + blob_metadata, + ) + .await; + } + + #[tokio::test] + async fn test_blob_v2_packed_round_trip() { + let blob_metadata = HashMap::from([( + lance_arrow::ARROW_EXT_NAME_KEY.to_string(), + lance_arrow::BLOB_V2_EXT_NAME.to_string(), + )]); + + let kind_field = Arc::new(ArrowField::new("kind", DataType::UInt8, true)); + let data_field = Arc::new(ArrowField::new("data", DataType::LargeBinary, true)); + let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, true)); + let blob_id_field = Arc::new(ArrowField::new("blob_id", DataType::UInt32, true)); + let blob_size_field = Arc::new(ArrowField::new("blob_size", DataType::UInt64, true)); + let position_field = Arc::new(ArrowField::new("position", DataType::UInt64, true)); + + let kind_array = UInt8Array::from(vec![BlobKind::Packed as u8]); + let data_array = LargeBinaryArray::from(vec![None::<&[u8]>]); + let uri_array = StringArray::from(vec![None::<&str>]); + let blob_id_array = UInt32Array::from(vec![7]); + let blob_size_array = UInt64Array::from(vec![5]); + let position_array = UInt64Array::from(vec![10]); + + let struct_array = StructArray::from(vec![ + (kind_field, Arc::new(kind_array) as ArrayRef), + (data_field, Arc::new(data_array) as ArrayRef), + (uri_field, Arc::new(uri_array) as ArrayRef), + (blob_id_field, Arc::new(blob_id_array) as ArrayRef), + (blob_size_field, Arc::new(blob_size_array) as ArrayRef), + (position_field, Arc::new(position_array) as ArrayRef), + ]); + + let expected_descriptor = StructArray::from(vec![ + ( + Arc::new(ArrowField::new("kind", DataType::UInt8, false)), + Arc::new(UInt8Array::from(vec![BlobKind::Packed as u8])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("position", DataType::UInt64, false)), + Arc::new(UInt64Array::from(vec![10])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("size", DataType::UInt64, false)), + Arc::new(UInt64Array::from(vec![5])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("blob_id", DataType::UInt32, false)), + Arc::new(UInt32Array::from(vec![7])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("blob_uri", DataType::Utf8, false)), + Arc::new(StringArray::from(vec![""])) as ArrayRef, + ), + ]); + + check_round_trip_encoding_of_data_with_expected( + vec![Arc::new(struct_array)], + Some(Arc::new(expected_descriptor)), + &TestCases::default().with_u32_structural_encodings(), + blob_metadata, + ) + .await; + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/fixed_size_list.rs b/vendor/lance-encoding/src/encodings/logical/fixed_size_list.rs new file mode 100644 index 000000000..745139bf7 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/fixed_size_list.rs @@ -0,0 +1,804 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Encoding support for complex FixedSizeList types (FSL with non-primitive children). +//! +//! Primitive FSL (e.g., `FixedSizeList`) is handled in the physical encoding layer. +//! This module handles FSL with complex children (Struct, Map, List) which require +//! structural encoding. + +use std::{ops::Range, sync::Arc}; + +use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait, StructArray, cast::AsArray}; +use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow_schema::DataType; +use futures::{FutureExt, future::BoxFuture}; +use lance_arrow::deepcopy::deep_copy_nulls; +use lance_core::{Error, Result}; + +use crate::{ + decoder::{ + DecodedArray, FilterExpression, ScheduledScanLine, SchedulerContext, + StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler, + StructuralSchedulingJob, + }, + encoder::{EncodeTask, FieldEncoder, OutOfLineBuffers}, + repdef::RepDefBuilder, +}; + +/// A structural encoder for complex fixed-size list fields +/// +/// The FSL's validity is added to the rep/def builder along with the dimension +/// and the FSL array's values are passed to the child encoder. +pub struct FixedSizeListStructuralEncoder { + keep_original_array: bool, + child: Box, +} + +impl FixedSizeListStructuralEncoder { + pub fn new(keep_original_array: bool, child: Box) -> Self { + Self { + keep_original_array, + child, + } + } +} + +impl FieldEncoder for FixedSizeListStructuralEncoder { + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + mut repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + let fsl_arr = array.as_fixed_size_list_opt().ok_or_else(|| { + Error::internal("FixedSizeList encoder used for non-fixed-size-list data".to_string()) + })?; + + let dimension = fsl_arr.value_length() as usize; + let values = fsl_arr.values().clone(); + + let validity = if self.keep_original_array { + array.nulls().cloned() + } else { + deep_copy_nulls(array.nulls()) + }; + repdef.add_fsl(validity.clone(), dimension, fsl_arr.len()); + + // FSL forces child elements to exist even under null rows. Normalize any + // nested lists under null FSL rows to null empty lists. + let values = if let Some(ref fsl_validity) = validity { + if needs_garbage_filtering(values.data_type()) { + let is_garbage = + expand_garbage_mask(&fsl_validity_to_garbage_mask(fsl_validity), dimension); + filter_fsl_child_garbage(values, &is_garbage) + } else { + values + } + } else { + values + }; + + self.child.maybe_encode( + values, + external_buffers, + repdef, + row_number, + num_rows * dimension as u64, + ) + } + + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + self.child.flush(external_buffers) + } + + fn num_columns(&self) -> u32 { + self.child.num_columns() + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + self.child.finish(external_buffers) + } +} + +/// A scheduler for complex fixed-size list fields +/// +/// Scales row ranges by the FSL dimension when scheduling child rows, +/// and scales scheduled rows back when reporting to the parent. +#[derive(Debug)] +pub struct StructuralFixedSizeListScheduler { + child: Box, + dimension: u64, +} + +impl StructuralFixedSizeListScheduler { + pub fn new(child: Box, dimension: i32) -> Self { + Self { + child, + dimension: dimension as u64, + } + } + + fn child_ranges(&self, ranges: &[Range]) -> Result>> { + ranges + .iter() + .map(|range| { + let start = range.start.checked_mul(self.dimension).ok_or_else(|| { + Error::invalid_input(format!( + "fixed-size-list child range start overflowed: {} * {}", + range.start, self.dimension + )) + })?; + let end = range.end.checked_mul(self.dimension).ok_or_else(|| { + Error::invalid_input(format!( + "fixed-size-list child range end overflowed: {} * {}", + range.end, self.dimension + )) + })?; + Ok(start..end) + }) + .collect() + } +} + +impl StructuralFieldScheduler for StructuralFixedSizeListScheduler { + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result> { + // Scale ranges by dimension for the child - each FSL row becomes `dimension` child rows + let child_ranges = self.child_ranges(ranges)?; + let child = self.child.schedule_ranges(&child_ranges, filter)?; + Ok(Box::new(StructuralFixedSizeListSchedulingJob::new( + child, + self.dimension, + ))) + } + + fn initialize<'a>( + &'a mut self, + requested_ranges: Option<&'a [Range]>, + filter: &'a FilterExpression, + context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + let child_ranges = match requested_ranges.map(|ranges| self.child_ranges(ranges)) { + Some(Ok(child_ranges)) => Some(child_ranges), + Some(Err(error)) => return std::future::ready(Err(error)).boxed(), + None => None, + }; + async move { + self.child + .initialize(child_ranges.as_deref(), filter, context) + .await + } + .boxed() + } +} + +#[derive(Debug)] +struct StructuralFixedSizeListSchedulingJob<'a> { + child: Box, + dimension: u64, +} + +impl<'a> StructuralFixedSizeListSchedulingJob<'a> { + fn new(child: Box, dimension: u64) -> Self { + Self { child, dimension } + } +} + +impl StructuralSchedulingJob for StructuralFixedSizeListSchedulingJob<'_> { + fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result> { + // Get the child's scan lines (scheduled in terms of child struct rows) + let child_scan_lines = self.child.schedule_next(context)?; + + // Scale down rows_scheduled by dimension to convert from child rows to FSL rows + Ok(child_scan_lines + .into_iter() + .map(|scan_line| ScheduledScanLine { + decoders: scan_line.decoders, + rows_scheduled: scan_line.rows_scheduled / self.dimension, + }) + .collect()) + } +} + +/// A decoder for complex fixed-size list fields +/// +/// Drains `num_rows * dimension` from the child decoder and reconstructs +/// the FSL array with validity from the rep/def information. +#[derive(Debug)] +pub struct StructuralFixedSizeListDecoder { + child: Box, + data_type: DataType, +} + +impl StructuralFixedSizeListDecoder { + pub fn new(child: Box, data_type: DataType) -> Self { + Self { child, data_type } + } +} + +impl StructuralFieldDecoder for StructuralFixedSizeListDecoder { + fn accept_page(&mut self, child: crate::decoder::LoadedPageShard) -> Result<()> { + self.child.accept_page(child) + } + + fn drain(&mut self, num_rows: u64) -> Result> { + // For FixedSizeList, we need to drain num_rows * dimension from the child + let dimension = match &self.data_type { + DataType::FixedSizeList(_, d) => *d as u64, + _ => { + return Err(Error::internal( + "FixedSizeListDecoder has non-FSL data type".to_string(), + )); + } + }; + let child_task = self.child.drain(num_rows * dimension)?; + Ok(Box::new(StructuralFixedSizeListDecodeTask::new( + child_task, + self.data_type.clone(), + num_rows, + ))) + } + + fn data_type(&self) -> &DataType { + &self.data_type + } +} + +#[derive(Debug)] +struct StructuralFixedSizeListDecodeTask { + child_task: Box, + data_type: DataType, + num_rows: u64, +} + +impl StructuralFixedSizeListDecodeTask { + fn new( + child_task: Box, + data_type: DataType, + num_rows: u64, + ) -> Self { + Self { + child_task, + data_type, + num_rows, + } + } +} + +impl StructuralDecodeArrayTask for StructuralFixedSizeListDecodeTask { + fn decode(self: Box) -> Result { + let DecodedArray { + array, + mut repdef, + data_size, + } = self.child_task.decode()?; + match &self.data_type { + DataType::FixedSizeList(child_field, dimension) => { + let num_rows = self.num_rows as usize; + let validity = repdef.unravel_fsl_validity(num_rows, *dimension as usize)?; + let fsl_array = arrow_array::FixedSizeListArray::try_new( + child_field.clone(), + *dimension, + array, + validity, + )?; + Ok(DecodedArray { + array: Arc::new(fsl_array), + repdef, + data_size, + }) + } + _ => Err(Error::internal( + "FixedSizeList decoder did not have a fixed-size list field".to_string(), + )), + } + } +} + +// ======================= +// Garbage filtering +// ======================= + +/// Returns true if the data type contains any variable-length list-like types +/// (List, LargeList, ListView, LargeListView, Map) that need garbage filtering. +fn needs_garbage_filtering(data_type: &DataType) -> bool { + match data_type { + DataType::List(_) + | DataType::LargeList(_) + | DataType::ListView(_) + | DataType::LargeListView(_) + | DataType::Map(_, _) => true, + DataType::Struct(fields) => fields + .iter() + .any(|f| needs_garbage_filtering(f.data_type())), + DataType::FixedSizeList(field, _) => needs_garbage_filtering(field.data_type()), + _ => false, + } +} + +/// Filters garbage (undefined data under null FSL rows) from nested list-like types. +/// Unlike variable-length lists which can remove null children entirely, FSL children +/// always exist, so we must clean any nested lists before encoding. +/// +/// NB: Nested FSL is currently precluded at a higher level in our system. However, this code +/// supports and tests it. +fn filter_fsl_child_garbage(array: ArrayRef, is_garbage: &[bool]) -> ArrayRef { + debug_assert_eq!(array.len(), is_garbage.len()); + + match array.data_type() { + DataType::List(_) => filter_list_garbage(array.as_list::(), is_garbage), + DataType::LargeList(_) => filter_list_garbage(array.as_list::(), is_garbage), + DataType::ListView(_) | DataType::LargeListView(_) => { + unimplemented!("ListView inside complex FSL is not yet supported") + } + DataType::Map(_, _) => filter_map_garbage(array.as_map(), is_garbage), + DataType::FixedSizeList(_, dim) => { + filter_nested_fsl_garbage(array.as_fixed_size_list(), is_garbage, *dim as usize) + } + DataType::Struct(_) => filter_struct_garbage(array.as_struct(), is_garbage), + _ => array, + } +} + +fn filter_struct_garbage(struct_arr: &StructArray, is_garbage: &[bool]) -> ArrayRef { + let needs_filtering = struct_arr + .fields() + .iter() + .any(|f| needs_garbage_filtering(f.data_type())); + + if !needs_filtering { + return Arc::new(struct_arr.clone()); + } + + let new_columns: Vec = struct_arr + .columns() + .iter() + .zip(struct_arr.fields().iter()) + .map(|(col, field)| { + if needs_garbage_filtering(field.data_type()) { + filter_fsl_child_garbage(col.clone(), is_garbage) + } else { + col.clone() + } + }) + .collect(); + + Arc::new(StructArray::new( + struct_arr.fields().clone(), + new_columns, + struct_arr.nulls().cloned(), + )) +} + +fn expand_garbage_mask(is_garbage: &[bool], dimension: usize) -> Vec { + let mut expanded = Vec::with_capacity(is_garbage.len() * dimension); + for &garbage in is_garbage { + for _ in 0..dimension { + expanded.push(garbage); + } + } + expanded +} + +fn fsl_validity_to_garbage_mask(fsl_validity: &NullBuffer) -> Vec { + fsl_validity.iter().map(|valid| !valid).collect() +} + +fn filter_list_garbage( + list_arr: &GenericListArray, + is_garbage: &[bool], +) -> ArrayRef { + debug_assert_eq!( + list_arr.len(), + is_garbage.len(), + "list length must match garbage mask length" + ); + + let old_offsets = list_arr.offsets(); + let value_field = match list_arr.data_type() { + DataType::List(f) | DataType::LargeList(f) => f.clone(), + _ => unreachable!(), + }; + + let mut new_offsets: Vec = Vec::with_capacity(list_arr.len() + 1); + let mut values_to_keep: Vec = Vec::new(); + let mut validity_builder = BooleanBufferBuilder::new(list_arr.len()); + let mut current_offset = O::usize_as(0); + new_offsets.push(current_offset); + let old_validity = list_arr.nulls(); + + for (i, &garbage) in is_garbage.iter().enumerate() { + if garbage { + new_offsets.push(current_offset); + validity_builder.append(false); + } else { + let start = old_offsets[i].as_usize(); + let end = old_offsets[i + 1].as_usize(); + values_to_keep.extend(start..end); + current_offset += O::usize_as(end - start); + new_offsets.push(current_offset); + validity_builder.append(old_validity.map(|v| v.is_valid(i)).unwrap_or(true)); + } + } + + let new_values = if values_to_keep.is_empty() { + list_arr.values().slice(0, 0) + } else { + let indices = + arrow_array::UInt64Array::from_iter_values(values_to_keep.iter().map(|&i| i as u64)); + arrow_select::take::take(list_arr.values().as_ref(), &indices, None) + .expect("take should succeed") + }; + + let new_values = if needs_garbage_filtering(value_field.data_type()) && !new_values.is_empty() { + let len = new_values.len(); + filter_fsl_child_garbage(new_values, &vec![false; len]) + } else { + new_values + }; + + let new_validity = NullBuffer::new(validity_builder.finish()); + Arc::new(GenericListArray::new( + value_field, + OffsetBuffer::new(ScalarBuffer::from(new_offsets)), + new_values, + Some(new_validity), + )) +} + +fn filter_map_garbage(map_arr: &arrow_array::MapArray, is_garbage: &[bool]) -> ArrayRef { + debug_assert_eq!(map_arr.len(), is_garbage.len()); + + let old_offsets = map_arr.offsets(); + let entries_field = match map_arr.data_type() { + DataType::Map(field, _) => field.clone(), + _ => unreachable!(), + }; + + let mut new_offsets: Vec = Vec::with_capacity(map_arr.len() + 1); + let mut values_to_keep: Vec = Vec::new(); + let mut validity_builder = BooleanBufferBuilder::new(map_arr.len()); + let mut current_offset: i32 = 0; + new_offsets.push(current_offset); + let old_validity = map_arr.nulls(); + + for (i, &garbage) in is_garbage.iter().enumerate() { + if garbage { + new_offsets.push(current_offset); + validity_builder.append(false); + } else { + let start = old_offsets[i] as usize; + let end = old_offsets[i + 1] as usize; + values_to_keep.extend(start..end); + current_offset += (end - start) as i32; + new_offsets.push(current_offset); + validity_builder.append(old_validity.map(|v| v.is_valid(i)).unwrap_or(true)); + } + } + + let new_entries: ArrayRef = if values_to_keep.is_empty() { + Arc::new(map_arr.entries().slice(0, 0)) + } else { + let indices = + arrow_array::UInt64Array::from_iter_values(values_to_keep.iter().map(|&i| i as u64)); + arrow_select::take::take(map_arr.entries(), &indices, None).expect("take should succeed") + }; + + let new_entries = + if needs_garbage_filtering(entries_field.data_type()) && !new_entries.is_empty() { + let len = new_entries.len(); + filter_fsl_child_garbage(new_entries, &vec![false; len]) + } else { + new_entries + }; + + let new_validity = NullBuffer::new(validity_builder.finish()); + let keys_sorted = matches!(map_arr.data_type(), DataType::Map(_, true)); + + Arc::new( + arrow_array::MapArray::try_new( + entries_field, + OffsetBuffer::new(ScalarBuffer::from(new_offsets)), + new_entries.as_struct().clone(), + Some(new_validity), + keys_sorted, + ) + .expect("MapArray construction should succeed"), + ) +} + +/// Filters garbage from nested FSL arrays that contain list-like children. +fn filter_nested_fsl_garbage( + fsl_arr: &arrow_array::FixedSizeListArray, + is_garbage: &[bool], + dimension: usize, +) -> ArrayRef { + debug_assert_eq!(fsl_arr.len(), is_garbage.len()); + + let child_field = match fsl_arr.data_type() { + DataType::FixedSizeList(field, _) => field.clone(), + _ => unreachable!(), + }; + + if !needs_garbage_filtering(child_field.data_type()) { + return Arc::new(fsl_arr.clone()); + } + + let child_garbage = expand_garbage_mask(is_garbage, dimension); + let new_values = filter_fsl_child_garbage(fsl_arr.values().clone(), &child_garbage); + + Arc::new(arrow_array::FixedSizeListArray::new( + child_field, + dimension as i32, + new_values, + fsl_arr.nulls().cloned(), + )) +} + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use arrow_array::{ + Array, FixedSizeListArray, + builder::{Int32Builder, ListBuilder}, + cast::AsArray, + }; + use arrow_schema::{DataType, Field, Fields}; + use rstest::rstest; + + use super::filter_nested_fsl_garbage; + use crate::{ + constants::{ + STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, + STRUCTURAL_ENCODING_MINIBLOCK, + }, + testing::{TestCases, TestEncoding, check_specific_random}, + }; + + fn make_fsl_struct_type(struct_fields: Fields, dimension: i32) -> DataType { + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Struct(struct_fields), true)), + dimension, + ) + } + + fn simple_struct_fields() -> Fields { + Fields::from(vec![ + Field::new("x", DataType::Float64, false), + Field::new("y", DataType::Float64, false), + ]) + } + + fn nested_struct_fields() -> Fields { + let inner = Fields::from(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + Fields::from(vec![ + Field::new("outer_val", DataType::Float64, false), + Field::new("inner", DataType::Struct(inner), true), + ]) + } + + fn nested_struct_with_list_fields() -> Fields { + let inner = Fields::from(vec![Field::new( + "values", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )]); + Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new("inner", DataType::Struct(inner), true), + ]) + } + + fn struct_with_list_fields() -> Fields { + Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "values", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + ), + ]) + } + + fn struct_with_large_list_fields() -> Fields { + Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "values", + DataType::LargeList(Arc::new(Field::new("item", DataType::Int64, true))), + true, + ), + ]) + } + + fn struct_with_nested_fsl_fields() -> Fields { + Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vectors", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + true, + ), + ]) + } + + fn struct_with_map_fields() -> Fields { + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", DataType::Int32, true), + ])), + false, + )); + Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new("props", DataType::Map(entries_field, false), true), + ]) + } + + fn make_fsl_of_list() -> DataType { + DataType::FixedSizeList( + Arc::new(Field::new( + "item", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )), + 2, + ) + } + + fn make_fsl_of_large_list() -> DataType { + DataType::FixedSizeList( + Arc::new(Field::new( + "item", + DataType::LargeList(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )), + 2, + ) + } + + fn make_fsl_of_map() -> DataType { + DataType::FixedSizeList( + Arc::new(Field::new( + "item", + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + ])), + false, + )), + false, + ), + true, + )), + 2, + ) + } + + fn make_fsl_of_nested_fsl_struct() -> DataType { + DataType::FixedSizeList( + Arc::new(Field::new( + "item", + DataType::FixedSizeList( + Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![Field::new( + "x", + DataType::Int32, + true, + )])), + true, + )), + 4, + ), + true, + )), + 2, + ) + } + + #[rstest] + #[case::simple(simple_struct_fields(), 2)] + #[case::nested_struct(nested_struct_fields(), 2)] + #[case::struct_with_list(struct_with_list_fields(), 2)] + #[case::struct_with_large_list(struct_with_large_list_fields(), 2)] + #[case::nested_struct_with_list(nested_struct_with_list_fields(), 2)] + #[case::struct_with_nested_fsl(struct_with_nested_fsl_fields(), 2)] + #[case::struct_with_map(struct_with_map_fields(), 2)] + #[test_log::test(tokio::test)] + async fn test_fsl_struct_random( + #[case] struct_fields: Fields, + #[case] dimension: i32, + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + #[values(TestEncoding::StructuralU32, TestEncoding::StructuralSparse)] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + #[values(1, 5, 10)] ingest_batch_count: u32, + ) { + let data_type = make_fsl_struct_type(struct_fields, dimension); + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + let field = Field::new("", data_type, true).with_metadata(field_metadata); + let test_cases = TestCases::basic() + .with_encoding(encoding) + .with_page_sizes(vec![page_size]) + .with_slicing_modes([use_slicing]) + .with_ingest_batch_counts([ingest_batch_count]); + check_specific_random(field, test_cases).await; + } + + #[rstest] + #[case::list(make_fsl_of_list())] + #[case::large_list(make_fsl_of_large_list())] + #[case::map(make_fsl_of_map())] + #[case::nested_fsl_struct(make_fsl_of_nested_fsl_struct())] + fn test_unsupported_fsl_child_types_return_error(#[case] data_type: DataType) { + let arrow_field = Field::new("test", data_type, true); + let err = lance_core::datatypes::Field::try_from(&arrow_field).unwrap_err(); + assert!(err.to_string().contains("Unsupported data type")); + } + + #[test] + fn test_filter_nested_fsl_garbage() { + // Create FSL> with dimension 2: [[[1], [2]], [[3], [4]], [[5], [6]]] + let mut list_builder = ListBuilder::new(Int32Builder::new()); + for i in 1..=6 { + list_builder.values().append_value(i); + list_builder.append(true); + } + let list_arr = list_builder.finish(); + + let fsl_field = Arc::new(Field::new( + "item", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )); + let fsl = FixedSizeListArray::new(fsl_field, 2, Arc::new(list_arr), None); + + // Mark second FSL row as garbage + let result = filter_nested_fsl_garbage(&fsl, &[false, true, false], 2); + let result = result.as_fixed_size_list(); + + // Child lists at positions 2,3 (garbage row 1) should be filtered to null + let child_list = result.values().as_list::(); + assert_eq!( + (0..6).map(|i| child_list.is_valid(i)).collect::>(), + vec![true, true, false, false, true, true] + ); + } + + #[test] + fn test_filter_nested_fsl_no_list_child() { + // FSL - no list child, should return unchanged + let fsl_field = Arc::new(Field::new("item", DataType::Int32, true)); + let values = arrow_array::Int32Array::from(vec![1, 2, 3, 4, 5, 6]); + let fsl = FixedSizeListArray::new(fsl_field, 2, Arc::new(values), None); + + let result = filter_nested_fsl_garbage(&fsl, &[false, true, false], 2); + // Should return the same array unchanged + assert_eq!(result.len(), 3); + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/list.rs b/vendor/lance-encoding/src/encodings/logical/list.rs new file mode 100644 index 000000000..a18608ace --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/list.rs @@ -0,0 +1,1662 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ops::Range, sync::Arc}; + +use arrow_array::{Array, ArrayRef, LargeListArray, ListArray, cast::AsArray, make_array}; +use arrow_schema::DataType; +use futures::future::BoxFuture; +use lance_arrow::deepcopy::deep_copy_nulls; +use lance_arrow::list::ListArrayExt; +use lance_core::Result; + +use crate::{ + decoder::{ + DecodedArray, FilterExpression, ScheduledScanLine, SchedulerContext, + StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler, + StructuralSchedulingJob, + }, + encoder::{EncodeTask, FieldEncoder, OutOfLineBuffers}, + repdef::RepDefBuilder, +}; + +/// A structural encoder for list fields +/// +/// The list's offsets are added to the rep/def builder +/// and the list array's values are passed to the child encoder +/// +/// The values will have any garbage values removed and will be trimmed +/// to only include the values that are actually used. +pub struct ListStructuralEncoder { + keep_original_array: bool, + child: Box, +} + +impl ListStructuralEncoder { + pub fn new(keep_original_array: bool, child: Box) -> Self { + Self { + keep_original_array, + child, + } + } +} + +impl FieldEncoder for ListStructuralEncoder { + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + mut repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + let values = if let Some(list_arr) = array.as_list_opt::() { + let has_garbage_values = if self.keep_original_array { + repdef.add_offsets(list_arr.offsets().clone(), array.nulls().cloned()) + } else { + // there is no need to deep copy offsets, because offset buffers will be cast to a common type (i64). + repdef.add_offsets(list_arr.offsets().clone(), deep_copy_nulls(array.nulls())) + }; + if has_garbage_values { + list_arr.filter_garbage_nulls().trimmed_values() + } else { + list_arr.trimmed_values() + } + } else if let Some(list_arr) = array.as_list_opt::() { + let has_garbage_values = if self.keep_original_array { + repdef.add_offsets(list_arr.offsets().clone(), array.nulls().cloned()) + } else { + repdef.add_offsets(list_arr.offsets().clone(), deep_copy_nulls(array.nulls())) + }; + if has_garbage_values { + list_arr.filter_garbage_nulls().trimmed_values() + } else { + list_arr.trimmed_values() + } + } else { + panic!("List encoder used for non-list data") + }; + self.child + .maybe_encode(values, external_buffers, repdef, row_number, num_rows) + } + + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + self.child.flush(external_buffers) + } + + fn num_columns(&self) -> u32 { + self.child.num_columns() + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + self.child.finish(external_buffers) + } +} + +#[derive(Debug)] +pub struct StructuralListScheduler { + child: Box, +} + +impl StructuralListScheduler { + pub fn new(child: Box) -> Self { + Self { child } + } +} + +impl StructuralFieldScheduler for StructuralListScheduler { + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result> { + let child = self.child.schedule_ranges(ranges, filter)?; + + Ok(Box::new(StructuralListSchedulingJob::new(child))) + } + + fn initialize<'a>( + &'a mut self, + requested_ranges: Option<&'a [Range]>, + filter: &'a FilterExpression, + context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + // The child (items) column is scheduled with the same top-level row + // ranges (see `schedule_ranges`); the page repetition index maps those + // rows to items, so forward the ranges unchanged. + self.child.initialize(requested_ranges, filter, context) + } +} + +/// Scheduling job for list data +/// +/// Scheduling is handled by the primitive encoder and nothing special +/// happens here. +#[derive(Debug)] +struct StructuralListSchedulingJob<'a> { + child: Box, +} + +impl<'a> StructuralListSchedulingJob<'a> { + fn new(child: Box) -> Self { + Self { child } + } +} + +impl StructuralSchedulingJob for StructuralListSchedulingJob<'_> { + fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result> { + self.child.schedule_next(context) + } +} + +#[derive(Debug)] +pub struct StructuralListDecoder { + child: Box, + data_type: DataType, +} + +impl StructuralListDecoder { + pub fn new(child: Box, data_type: DataType) -> Self { + Self { child, data_type } + } +} + +impl StructuralFieldDecoder for StructuralListDecoder { + fn accept_page(&mut self, child: crate::decoder::LoadedPageShard) -> Result<()> { + self.child.accept_page(child) + } + + fn drain(&mut self, num_rows: u64) -> Result> { + let child_task = self.child.drain(num_rows)?; + Ok(Box::new(StructuralListDecodeTask::new( + child_task, + self.data_type.clone(), + ))) + } + + fn data_type(&self) -> &DataType { + &self.data_type + } + + fn plan_decoded_bytes(&self, _rows_remaining: u64) -> lance_core::Result<[u64; 8]> { + Err(lance_core::Error::not_supported( + "plan_decoded_bytes is not yet supported for fields with repetition (list/repeated types)" + .to_string(), + )) + } +} + +#[derive(Debug)] +struct StructuralListDecodeTask { + child_task: Box, + data_type: DataType, +} + +impl StructuralListDecodeTask { + fn new(child_task: Box, data_type: DataType) -> Self { + Self { + child_task, + data_type, + } + } +} + +impl StructuralDecodeArrayTask for StructuralListDecodeTask { + fn decode(self: Box) -> Result { + let DecodedArray { + array, + mut repdef, + data_size, + } = self.child_task.decode()?; + match &self.data_type { + DataType::List(child_field) => { + let (offsets, validity) = repdef.unravel_offsets::()?; + let array = if !child_field.is_nullable() && array.null_count() == array.len() { + make_array(array.into_data().into_builder().nulls(None).build()?) + } else { + array + }; + let list_array = ListArray::try_new(child_field.clone(), offsets, array, validity)?; + + Ok(DecodedArray { + array: Arc::new(list_array), + repdef, + data_size, + }) + } + DataType::LargeList(child_field) => { + let (offsets, validity) = repdef.unravel_offsets::()?; + let list_array = + LargeListArray::try_new(child_field.clone(), offsets, array, validity)?; + Ok(DecodedArray { + array: Arc::new(list_array), + repdef, + data_size, + }) + } + _ => panic!("List decoder did not have a list field"), + } + } +} + +#[cfg(test)] +mod tests { + + use std::{collections::HashMap, sync::Arc}; + + use crate::constants::{ + STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, + }; + use arrow_array::{ + Array, ArrayRef, BooleanArray, DictionaryArray, LargeStringArray, ListArray, StringArray, + StructArray, UInt8Array, UInt64Array, + builder::{ + Int32Builder, Int64Builder, LargeListBuilder, ListBuilder, StringBuilder, UInt32Builder, + }, + }; + + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field, Fields}; + use lance_datagen::{RowCount, Seed, array, gen_batch}; + use rstest::rstest; + + use crate::testing::{ + TestCases, TestEncoding, check_round_trip_encoding_of_data, create_test_field_encoder, + test_encoding_strategy, + }; + + fn make_list_type(inner_type: DataType) -> DataType { + DataType::List(Arc::new(Field::new("item", inner_type, true))) + } + + fn make_large_list_type(inner_type: DataType) -> DataType { + DataType::LargeList(Arc::new(Field::new("item", inner_type, true))) + } + + #[derive(Clone, Copy)] + enum NullPattern { + None, + Mixed, + All, + } + + async fn check_nested_type( + data_type: DataType, + null_pattern: NullPattern, + encoding: TestEncoding, + ) { + check_nested_type_with_metadata(data_type, null_pattern, encoding, HashMap::new()).await; + } + + async fn check_nested_type_with_metadata( + data_type: DataType, + null_pattern: NullPattern, + encoding: TestEncoding, + field_metadata: HashMap, + ) { + let null_rate = match null_pattern { + NullPattern::None => None, + NullPattern::Mixed => Some(0.5), + NullPattern::All => Some(1.0), + }; + let make_batch = |seed, rows| { + let mut generator = gen_batch() + .with_seed(Seed::from(seed)) + .anon_col(array::rand_type(&data_type)); + if let Some(null_rate) = null_rate { + generator.with_random_nulls(null_rate); + } + generator + .into_batch_rows(RowCount::from(rows)) + .unwrap() + .column(0) + .clone() + }; + + // Combine a non-zero-offset slice with an independently generated batch. + // This covers both offset rebasing and rep/def accumulation at the ingest + // boundary without repeating the full generic random-test matrix. + let first = make_batch(0, 513).slice(1, 512); + let second = make_batch(1, 513); + let test_cases = TestCases::default() + .with_page_sizes(vec![4096]) + .with_encoding(encoding) + .with_batch_size(257) + .with_range(510..515) + .with_indices(vec![0, 511, 512, 1024]); + + check_round_trip_encoding_of_data(vec![first, second], &test_cases, field_metadata).await; + } + + async fn try_encode_v22_pages( + array: ArrayRef, + ) -> lance_core::Result> { + try_encode_v22_pages_with_metadata(array, HashMap::new()).await + } + + async fn try_encode_v22_pages_with_metadata( + array: ArrayRef, + field_metadata: HashMap, + ) -> lance_core::Result> { + let arrow_field = + Field::new("", array.data_type().clone(), true).with_metadata(field_metadata); + let lance_field = lance_core::datatypes::Field::try_from(&arrow_field).unwrap(); + let encoding_strategy = test_encoding_strategy(TestEncoding::StructuralU32); + let mut column_index_seq = crate::encoder::ColumnIndexSequence::default(); + let encoding_options = crate::encoder::EncodingOptions::default(); + let mut encoder = create_test_field_encoder( + encoding_strategy.as_ref(), + &lance_field, + &mut column_index_seq, + &encoding_options, + ) + .unwrap(); + let mut external_buffers = + crate::encoder::OutOfLineBuffers::new(0, crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT); + let num_rows = array.len() as u64; + let mut pages = Vec::new(); + for task in encoder + .maybe_encode( + array, + &mut external_buffers, + crate::repdef::RepDefBuilder::default(), + 0, + num_rows, + ) + .unwrap() + { + pages.push(task.await?); + } + for task in encoder.flush(&mut external_buffers).unwrap() { + pages.push(task.await?); + } + Ok(pages) + } + + async fn encode_v22_pages(array: ArrayRef) -> Vec { + try_encode_v22_pages(array).await.unwrap() + } + + fn assert_split_miniblock_layout( + pages: &[crate::encoder::EncodedPage], + min_miniblock_pages: usize, + expect_structural_only_page: bool, + ) { + let mut miniblock_pages = 0; + let mut fullzip_pages = 0; + let mut structural_only_pages = 0; + + for page in pages { + let crate::decoder::PageEncoding::Structural(layout) = &page.description else { + continue; + }; + match layout.layout.as_ref().unwrap() { + crate::format::pb21::page_layout::Layout::MiniBlockLayout(_) => { + miniblock_pages += 1; + } + crate::format::pb21::page_layout::Layout::FullZipLayout(_) => { + fullzip_pages += 1; + } + crate::format::pb21::page_layout::Layout::ConstantLayout(layout) => { + if layout.inline_value.is_none() + && (layout.num_rep_values > 0 || layout.num_def_values > 0) + { + structural_only_pages += 1; + } + } + crate::format::pb21::page_layout::Layout::BlobLayout(_) => {} + crate::format::pb21::page_layout::Layout::SparseLayout(_) => {} + } + } + + assert!( + miniblock_pages >= min_miniblock_pages, + "expected at least {min_miniblock_pages} mini-block pages, got {miniblock_pages}" + ); + assert_eq!( + fullzip_pages, 0, + "split list pages should not fall back to full-zip" + ); + assert_eq!( + structural_only_pages > 0, + expect_structural_only_page, + "structural-only page presence did not match expectation; got {structural_only_pages}" + ); + } + + fn assert_has_fullzip_layout(pages: &[crate::encoder::EncodedPage]) { + let has_fullzip = pages.iter().any(|page| { + let crate::decoder::PageEncoding::Structural(layout) = &page.description else { + return false; + }; + matches!( + layout.layout.as_ref().unwrap(), + crate::format::pb21::page_layout::Layout::FullZipLayout(_) + ) + }); + assert!(has_fullzip, "expected at least one full-zip page"); + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_list( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)] + null_pattern: NullPattern, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + check_nested_type_with_metadata( + make_list_type(DataType::Int32), + null_pattern, + encoding, + field_metadata, + ) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_deeply_nested_lists( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + #[values(1, 2, 3, 4, 5)] depth: usize, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + test_encoding: TestEncoding, + ) { + let mut data_type = DataType::Int32; + for _ in 0..depth { + data_type = make_list_type(data_type); + } + + let mut generator = gen_batch() + .with_seed(Seed::from(depth as u64)) + .anon_col(array::rand_type(&data_type)); + generator.with_random_nulls(0.2); + let source = generator + .into_batch_rows(RowCount::from(1026)) + .unwrap() + .column(0) + .clone(); + + // Two non-zero-offset slices cover nested offset rebasing across ingest + // batches. The selected range and indices straddle that batch boundary. + let data = vec![source.slice(1, 512), source.slice(513, 513)]; + let test_cases = TestCases::default() + .with_page_sizes(vec![4096]) + .with_encoding(test_encoding) + .with_batch_size(257) + .with_range(510..515) + .with_indices(vec![0, 511, 512, 1024]); + let field_metadata = HashMap::from([( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + )]); + + check_round_trip_encoding_of_data(data, &test_cases, field_metadata).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_large_list( + #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)] + null_pattern: NullPattern, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + check_nested_type( + make_large_list_type(DataType::Int32), + null_pattern, + encoding, + ) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_nested_strings( + #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)] + null_pattern: NullPattern, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + check_nested_type(make_list_type(DataType::Utf8), null_pattern, encoding).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_nested_list( + #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)] + null_pattern: NullPattern, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + check_nested_type( + make_list_type(make_list_type(DataType::Int32)), + null_pattern, + encoding, + ) + .await; + } + + /// Regression test: a `List>` column written as MULTIPLE + /// batches (chunks) whose flattened leaf values cross a value-page boundary + /// fails to decode with "Max offset N exceeds length of values M" (Arrow + /// error raised by `ListArray::try_new` in `StructuralListDecodeTask::decode`). + /// + /// The trigger (verified against the production file and pylance 7.0.0b12 / + /// 7.0.0 / 9.0.0-beta.10) requires ALL of: + /// 1. >= 2 list layers (`List>`), + /// 2. a leaf large enough to be chunked into multiple value pages, + /// 3. the column written as more than one batch. + /// A single batch of the identical data round-trips fine — which is why the + /// earlier single-chunk version of this test (and the small `test_nested_list` + /// cases) did not catch it. Found in production on the gaming TransNet + /// `dino_embedding_per_frame` column (rectangular 3 x 768 float per row). + /// + /// Each element of the `vec![..]` passed to `check_round_trip_encoding_of_data` + /// is encoded as a separate batch (its own `RepDefBuilder`), so we split the + /// rows into two chunks to exercise the multi-batch repdef accumulation path. + #[rstest] + #[test_log::test(tokio::test)] + async fn test_multipage_nested_float_list( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + test_encoding: TestEncoding, + ) { + use arrow_array::Float32Array; + + // Production shape: 3 inner lists per row, 768 floats each. + let inner_per_row: usize = 3; + let inner_len: usize = 768; + // Each chunk contains ~1.38 MiB of leaf values, so both cross the 1 MiB + // value-page limit. A read batch that spans the ingest boundary is where + // the multi-page outer-offset bug triggered. + let chunk_rows: &[usize] = &[150, 149]; + + let make_chunk = |start_row: usize, num_rows: usize| -> Arc { + let total_inner = num_rows * inner_per_row; + let total_values = total_inner * inner_len; + let values = Float32Array::from( + (0..total_values) + .map(|i| (start_row + i) as f32) + .collect::>(), + ); + let inner_offsets = ScalarBuffer::::from( + (0..=total_inner) + .map(|i| (i * inner_len) as i32) + .collect::>(), + ); + let inner_list = ListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + OffsetBuffer::new(inner_offsets), + Arc::new(values), + None, + ); + let outer_offsets = ScalarBuffer::::from( + (0..=num_rows) + .map(|i| (i * inner_per_row) as i32) + .collect::>(), + ); + Arc::new(ListArray::new( + Arc::new(Field::new( + "item", + DataType::List(Arc::new(Field::new("item", DataType::Float32, true))), + true, + )), + OffsetBuffer::new(outer_offsets), + Arc::new(inner_list), + None, + )) + }; + + let mut start = 0; + let chunks: Vec> = chunk_rows + .iter() + .map(|&n| { + let c = make_chunk(start, n); + start += n; + c + }) + .collect(); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default() + .with_page_sizes(vec![1024 * 1024]) + .with_encoding(test_encoding) + .with_batch_size(151) + .with_range(148..152) + .with_indices(vec![0, 149, 150, 298]); + check_round_trip_encoding_of_data(chunks, &test_cases, field_metadata).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_list_struct_list( + #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)] + null_pattern: NullPattern, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + let struct_type = DataType::Struct(Fields::from(vec![Field::new( + "inner_str", + DataType::Utf8, + false, + )])); + + check_nested_type(make_list_type(struct_type), null_pattern, encoding).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_list_struct_empty( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + let fields = Fields::from(vec![Field::new("inner", DataType::UInt64, true)]); + let items = UInt64Array::from(Vec::::new()); + let structs = StructArray::new(fields, vec![Arc::new(items)], None); + // Exceed two 1 MiB offset pages so flushing the empty struct child is + // still exercised multiple times (the original #2762 regression). + let num_rows = 2 * 1024 * 1024 / size_of::() + 1; + let offsets = OffsetBuffer::new(ScalarBuffer::::from(vec![0; num_rows + 1])); + let lists = ListArray::new( + Arc::new(Field::new("item", structs.data_type().clone(), true)), + offsets, + Arc::new(structs), + None, + ); + + check_round_trip_encoding_of_data( + vec![Arc::new(lists)], + &TestCases::default() + .with_page_sizes(vec![1024 * 1024]) + .with_encoding(encoding), + HashMap::new(), + ) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_simple_list( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + let items_builder = Int32Builder::new(); + let mut list_builder = ListBuilder::new(items_builder); + list_builder.append_value([Some(1), Some(2), Some(3)]); + list_builder.append_value([Some(4), Some(5)]); + list_builder.append_null(); + list_builder.append_value([Some(6), Some(7), Some(8)]); + let list_array = list_builder.finish(); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(1..3) + .with_indices(vec![1, 3]) + .with_indices(vec![2]); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_simple_nested_list_ends_with_null( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + use arrow_array::Int32Array; + + let values = Int32Array::from(vec![1, 2, 3, 4, 5]); + let inner_offsets = ScalarBuffer::::from(vec![0, 1, 2, 3, 4, 5, 5]); + let inner_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]); + let outer_offsets = ScalarBuffer::::from(vec![0, 1, 2, 3, 4, 5, 6, 6]); + let outer_validity = BooleanBuffer::from(vec![true, true, true, true, true, true, false]); + + let inner_list = ListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + OffsetBuffer::new(inner_offsets), + Arc::new(values), + Some(NullBuffer::new(inner_validity)), + ); + let outer_list = ListArray::new( + Arc::new(Field::new( + "item", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )), + OffsetBuffer::new(outer_offsets), + Arc::new(inner_list), + Some(NullBuffer::new(outer_validity)), + ); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(5..7) + .with_indices(vec![1, 6]) + .with_indices(vec![6]) + .with_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(outer_list)], &test_cases, field_metadata) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_simple_string_list( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + let items_builder = StringBuilder::new(); + let mut list_builder = ListBuilder::new(items_builder); + list_builder.append_value([Some("a"), Some("bc"), Some("def")]); + list_builder.append_value([Some("gh"), None]); + list_builder.append_null(); + list_builder.append_value([Some("ijk"), Some("lmnop"), Some("qrs")]); + let list_array = list_builder.finish(); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(1..3) + .with_indices(vec![1, 3]) + .with_indices(vec![2]) + .with_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_simple_string_list_no_null( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + let items_builder = StringBuilder::new(); + let mut list_builder = ListBuilder::new(items_builder); + list_builder.append_value([Some("a"), Some("bc"), Some("def")]); + list_builder.append_value([Some("gh"), Some("zxy")]); + list_builder.append_value([Some("gh"), Some("z")]); + list_builder.append_value([Some("ijk"), Some("lmnop"), Some("qrs")]); + let list_array = list_builder.finish(); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(1..3) + .with_indices(vec![1, 3]) + .with_indices(vec![2]) + .with_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_simple_sliced_list( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + let items_builder = Int32Builder::new(); + let mut list_builder = ListBuilder::new(items_builder); + list_builder.append_value([Some(1), Some(2), Some(3)]); + list_builder.append_value([Some(4), Some(5)]); + list_builder.append_null(); + list_builder.append_value([Some(6), Some(7), Some(8)]); + let list_array = list_builder.finish(); + + let list_array = list_array.slice(1, 2); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(1..2) + .with_indices(vec![0]) + .with_indices(vec![1]) + .with_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_list_dict() { + let values = LargeStringArray::from_iter_values(["a", "bb", "ccc"]); + let indices = UInt8Array::from(vec![0, 1, 2, 0, 1, 2, 0, 1, 2]); + let dict_array = DictionaryArray::new(indices, Arc::new(values)); + let offsets = OffsetBuffer::new(ScalarBuffer::::from(vec![0, 3, 5, 6, 9])); + let list_array = ListArray::new( + Arc::new(Field::new("item", dict_array.data_type().clone(), true)), + offsets, + Arc::new(dict_array), + None, + ); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(1..3) + .with_range(2..4) + .with_indices(vec![1]) + .with_indices(vec![2]); + check_round_trip_encoding_of_data( + vec![Arc::new(list_array)], + &test_cases, + HashMap::default(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_list_all_null() { + let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let offsets = ScalarBuffer::::from(vec![0, 5, 8, 10]); + let offsets = OffsetBuffer::new(offsets); + let list_validity = NullBuffer::new(BooleanBuffer::from(vec![false, false, false])); + + // The list array is nullable but the items are not. Then, all lists are null. + let list_arr = ListArray::new( + Arc::new(Field::new("item", DataType::UInt64, false)), + offsets, + Arc::new(items), + Some(list_validity), + ); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_range(1..2) + .with_indices(vec![1]) + .with_indices(vec![2]) + .with_structural_encodings(); + check_round_trip_encoding_of_data( + vec![Arc::new(list_arr)], + &test_cases, + HashMap::default(), + ) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_list_with_garbage_nulls( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + // In Arrow, list nulls are allowed to be non-empty, with masked garbage values + // Here we make a list with a null row in the middle with 3 garbage values + let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let offsets = ScalarBuffer::::from(vec![0, 5, 8, 10]); + let offsets = OffsetBuffer::new(offsets); + let list_validity = NullBuffer::new(BooleanBuffer::from(vec![true, false, true])); + let list_arr = ListArray::new( + Arc::new(Field::new("item", DataType::UInt64, true)), + offsets, + Arc::new(items), + Some(list_validity), + ); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_range(1..2) + .with_indices(vec![1]) + .with_indices(vec![2]) + .with_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(list_arr)], &test_cases, field_metadata) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_simple_two_page_list( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + // This is a simple pre-defined list that spans two pages. This test is useful for + // debugging the repetition index + + let items_builder = Int64Builder::new(); + let mut list_builder = ListBuilder::new(items_builder); + for i in 0..512 { + list_builder.append_value([Some(i), Some(i * 2)]); + } + let list_array_1 = list_builder.finish(); + + let items_builder = Int64Builder::new(); + let mut list_builder = ListBuilder::new(items_builder); + for i in 0..512 { + let i = i + 512; + list_builder.append_value([Some(i), Some(i * 2)]); + } + let list_array_2 = list_builder.finish(); + + let mut metadata = HashMap::new(); + metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default() + .with_structural_encodings() + .with_page_sizes(vec![100]) + .with_range(800..900); + check_round_trip_encoding_of_data( + vec![Arc::new(list_array_1), Arc::new(list_array_2)], + &test_cases, + metadata, + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_large_list() { + let items_builder = Int32Builder::new(); + let mut list_builder = LargeListBuilder::new(items_builder); + list_builder.append_value([Some(1), Some(2), Some(3)]); + list_builder.append_value([Some(4), Some(5)]); + list_builder.append_null(); + list_builder.append_value([Some(6), Some(7), Some(8)]); + let list_array = list_builder.finish(); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(1..3) + .with_indices(vec![1, 3]); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_empty_lists( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + // Scenario 1: Some lists are empty + + let values = [vec![Some(1), Some(2), Some(3)], vec![], vec![None]]; + // Test empty list at beginning, middle, and end + for order in [[0, 1, 2], [1, 0, 2], [2, 0, 1]] { + let items_builder = Int32Builder::new(); + let mut list_builder = ListBuilder::new(items_builder); + for idx in order { + list_builder.append_value(values[idx].clone()); + } + let list_array = Arc::new(list_builder.finish()); + let test_cases = TestCases::default() + .with_indices(vec![1]) + .with_indices(vec![0]) + .with_indices(vec![2]) + .with_indices(vec![0, 1]); + check_round_trip_encoding_of_data( + vec![list_array.clone()], + &test_cases, + field_metadata.clone(), + ) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data( + vec![list_array], + &test_cases, + field_metadata.clone(), + ) + .await; + } + + // Scenario 2: All lists are empty + + // When encoding a list of empty lists there are no items to encode + // which is strange and we want to ensure we handle it + let items_builder = Int32Builder::new(); + let mut list_builder = ListBuilder::new(items_builder); + list_builder.append(true); + list_builder.append_null(); + list_builder.append(true); + let list_array = Arc::new(list_builder.finish()); + + let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]); + check_round_trip_encoding_of_data( + vec![list_array.clone()], + &test_cases, + field_metadata.clone(), + ) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata.clone()) + .await; + + // Scenario 2B: All lists are empty (but now with strings) + + // When encoding a list of empty lists there are no items to encode + // which is strange and we want to ensure we handle it + let items_builder = StringBuilder::new(); + let mut list_builder = ListBuilder::new(items_builder); + list_builder.append(true); + list_builder.append_null(); + list_builder.append(true); + let list_array = Arc::new(list_builder.finish()); + + let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]); + check_round_trip_encoding_of_data( + vec![list_array.clone()], + &test_cases, + field_metadata.clone(), + ) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata.clone()) + .await; + + // Scenario 3: All lists are null + + let items_builder = Int32Builder::new(); + let mut list_builder = ListBuilder::new(items_builder); + list_builder.append_null(); + list_builder.append_null(); + list_builder.append_null(); + let list_array = Arc::new(list_builder.finish()); + + let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]); + check_round_trip_encoding_of_data( + vec![list_array.clone()], + &test_cases, + field_metadata.clone(), + ) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata.clone()) + .await; + + // Scenario 4: All lists are null and inside a struct (only valid for 2.1 since 2.0 doesn't + // support null structs) + let items_builder = Int32Builder::new(); + let mut list_builder = ListBuilder::new(items_builder); + list_builder.append_null(); + list_builder.append_null(); + list_builder.append_null(); + let list_array = Arc::new(list_builder.finish()); + + let struct_validity = NullBuffer::new(BooleanBuffer::from(vec![true, false, true])); + let struct_array = Arc::new(StructArray::new( + Fields::from(vec![Field::new( + "lists", + list_array.data_type().clone(), + true, + )]), + vec![list_array], + Some(struct_validity), + )); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_indices(vec![1]) + .with_structural_encodings(); + check_round_trip_encoding_of_data( + vec![struct_array.clone()], + &test_cases, + field_metadata.clone(), + ) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data(vec![struct_array], &test_cases, field_metadata.clone()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_empty_list_list() { + let items_builder = Int32Builder::new(); + let list_builder = ListBuilder::new(items_builder); + let mut outer_list_builder = ListBuilder::new(list_builder); + outer_list_builder.append_null(); + outer_list_builder.append_null(); + outer_list_builder.append_null(); + let list_array = Arc::new(outer_list_builder.finish()); + + let test_cases = TestCases::default().with_structural_encodings(); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + + #[test_log::test(tokio::test)] + #[ignore] // This test is quite slow in debug mode + async fn test_jumbo_list() { + // This is an overflow test. We have a list of lists where each list + // has 1Mi items. We encode 5000 of these lists and so we have over 4Gi in the + // offsets range + let items = BooleanArray::new_null(1024 * 1024); + let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 1024 * 1024])); + let list_arr = Arc::new(ListArray::new( + Arc::new(Field::new("item", DataType::Boolean, true)), + offsets, + Arc::new(items), + None, + )) as ArrayRef; + let arrs = vec![list_arr; 5000]; + + // We can't validate because our validation relies on concatenating all input arrays + let test_cases = TestCases::default().without_validation(); + check_round_trip_encoding_of_data(arrs, &test_cases, HashMap::new()).await; + } + + // Regression test for issue with ListArray encoding when crossing 1024 value boundary + // This test reproduces the bug where rows_avail assertion fails in schedule_instructions + // when encoding a ListArray with specific size patterns that cross the 1024 value boundary + #[tokio::test] + async fn test_fuzz_issue_4466() { + // This specific pattern of list sizes triggers the bug when total values cross 1024 + // 94 lists total 1009 values (passes), 95 lists total 1025 values (fails) + let list_sizes = vec![ + 13, 18, 12, 7, 14, 12, 6, 13, 18, 8, // 0-9: 119 values + 6, 11, 17, 12, 8, 19, 5, 6, 10, 13, // 10-19: 107 values + 8, 6, 10, 4, 8, 16, 14, 12, 18, 9, // 20-29: 105 values + 17, 8, 14, 18, 15, 3, 2, 4, 5, 1, // 30-39: 82 values + 3, 13, 1, 2, 10, 4, 10, 18, 7, 14, // 40-49: 75 values + 18, 13, 9, 17, 3, 13, 10, 14, 8, 19, // 50-59: 125 values + 17, 10, 5, 11, 6, 15, 10, 18, 18, 20, // 60-69: 130 values + 16, 11, 12, 15, 7, 9, 3, 10, 20, 5, // 70-79: 102 values + 2, 3, 17, 4, 8, 12, 15, 6, 3, 20, // 80-89: 90 values + 15, 20, 1, 19, 16, // 90-94: 71 values + ]; + + // Build the ListArray + let mut list_builder = ListBuilder::new(Int32Builder::new()); + let mut total_values = 0; + + for size in &list_sizes { + for i in 0..*size { + list_builder.values().append_value(i); + } + list_builder.append(true); + total_values += size; + } + + let list_array = Arc::new(list_builder.finish()); + + // Verify we have the expected number of values + assert_eq!(list_array.len(), 95); + assert_eq!(total_values, 1025); + + // This should trigger the assertion failure at primitive.rs:1362 + // debug_assert!(rows_avail > 0) + let test_cases = TestCases::default().with_structural_encodings(); + + // The bug manifests when encoding this specific pattern + // Expected: successful round-trip encoding + // Actual: panic at primitive.rs:1362 - assertion failed: rows_avail > 0 + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_sparse_large_string_list( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + // Three chunks' worth of rep/def levels (1 rep bit + 1 def bit each), so the + // planner must split the page. See #6184. + let levels_per_chunk = + crate::encodings::logical::primitive::miniblock::max_repdef_levels_per_chunk(2); + let num_rows = (levels_per_chunk * 3) as u32; + let num_non_empty = 100u32; + let strings_per_list = 10; + + let items_builder = StringBuilder::new(); + let mut list_builder = ListBuilder::new(items_builder); + + // Spread non-empty lists evenly across the range + let step = num_rows / num_non_empty; + let mut next_non_empty = step / 2; + + for i in 0..num_rows { + if i == next_non_empty { + let vals: Vec> = (0..strings_per_list) + .map(|j| match j % 4 { + 0 => Some("a"), + 1 => Some("bb"), + 2 => Some("ccc"), + _ => Some("d"), + }) + .collect(); + list_builder.append_value(vals); + next_non_empty = next_non_empty.saturating_add(step); + } else { + list_builder.append_value([] as [Option<&str>; 0]); + } + } + let list_array = list_builder.finish(); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let list_array = Arc::new(list_array) as ArrayRef; + let pages = try_encode_v22_pages_with_metadata(list_array.clone(), field_metadata.clone()) + .await + .unwrap(); + if structural_encoding == STRUCTURAL_ENCODING_MINIBLOCK { + assert_split_miniblock_layout(&pages, 2, true); + } + + let chunk_boundary = levels_per_chunk; + let test_cases = TestCases::default() + .with_range(chunk_boundary - 2..chunk_boundary + 2) + .with_indices(vec![ + 0, + (step / 2) as u64, + chunk_boundary - 1, + chunk_boundary, + num_rows as u64 - 1, + ]) + .with_batch_size(64 * 1024) + .with_page_sizes(vec![1024 * 1024]) + .with_encoding(TestEncoding::StructuralU32); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata).await; + } + + #[test_log::test(tokio::test)] + async fn test_sparse_boolean_list_uses_miniblock() { + // Redacted reproduction from a production schema shape containing ARRAY(BOOLEAN). + // The field names are not relevant; the failure requires sparse list structure + // with a 1-bit Boolean leaf value. + let levels_per_chunk = + crate::encodings::logical::primitive::miniblock::max_repdef_levels_per_chunk(2); + // One row past the chunk limit forces a split. Keeping values at both ends ensures + // both sides of that split remain mini-block pages instead of structural-only pages. + let num_rows = (levels_per_chunk + 1) as usize; + let booleans_per_list = 8usize; + + let mut offsets = Vec::with_capacity(num_rows + 1); + let mut values = Vec::with_capacity(2 * booleans_per_list); + offsets.push(0i32); + + for row in 0..num_rows { + if row == 0 || row == num_rows - 1 { + values.extend((0..booleans_per_list).map(|idx| idx % 2 == 0)); + } + offsets.push(values.len() as i32); + } + + let items = BooleanArray::from(values); + let list_array = ListArray::new( + Arc::new(Field::new("item", DataType::Boolean, true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + Arc::new(items), + None, + ); + + let test_cases = TestCases::default() + .with_range(0..1000) + .with_indices(vec![0, levels_per_chunk / 2, num_rows as u64 - 1]) + .with_batch_size(64 * 1024) + .with_page_sizes(vec![1024 * 1024]) + .with_encoding(TestEncoding::StructuralU32); + let list_array = Arc::new(list_array) as ArrayRef; + let pages = encode_v22_pages(list_array.clone()).await; + assert_split_miniblock_layout(&pages, 2, false); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + + #[test_log::test(tokio::test)] + async fn test_sparse_boolean_list_with_long_empty_prefix() { + let empty_prefix_rows = 70_000usize; + let trailing_empty_rows = 9usize; + let booleans_per_list = 8usize; + let num_rows = empty_prefix_rows + 1 + trailing_empty_rows; + + let mut offsets = Vec::with_capacity(num_rows + 1); + offsets.extend(std::iter::repeat_n(0i32, empty_prefix_rows + 1)); + let values = (0..booleans_per_list) + .map(|idx| idx % 2 == 0) + .collect::>(); + offsets.push(values.len() as i32); + offsets.extend(std::iter::repeat_n( + values.len() as i32, + trailing_empty_rows, + )); + + let items = BooleanArray::from(values); + let list_array = ListArray::new( + Arc::new(Field::new("item", DataType::Boolean, true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + Arc::new(items), + None, + ); + + let test_cases = TestCases::default() + .with_range(0..num_rows as u64) + .with_indices(vec![0, empty_prefix_rows as u64, num_rows as u64 - 1]) + .with_dense_encodings(); + let list_array = Arc::new(list_array) as ArrayRef; + let pages = encode_v22_pages(list_array.clone()).await; + assert_split_miniblock_layout(&pages, 1, true); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_sparse_boolean_list_with_long_null_prefix( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32 + )] + encoding: TestEncoding, + ) { + let null_prefix_rows = 70_000usize; + let trailing_empty_rows = 9usize; + let booleans_per_list = 8usize; + let num_rows = null_prefix_rows + 1 + trailing_empty_rows; + + let mut offsets = Vec::with_capacity(num_rows + 1); + offsets.extend(std::iter::repeat_n(0i32, null_prefix_rows + 1)); + let values = (0..booleans_per_list) + .map(|idx| idx % 2 == 0) + .collect::>(); + offsets.push(values.len() as i32); + offsets.extend(std::iter::repeat_n( + values.len() as i32, + trailing_empty_rows, + )); + let validity = BooleanBuffer::from_iter((0..num_rows).map(|row| row >= null_prefix_rows)); + + let items = BooleanArray::from(values); + let list_array = ListArray::new( + Arc::new(Field::new("item", DataType::Boolean, true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + Arc::new(items), + Some(NullBuffer::new(validity)), + ); + + let test_cases = TestCases::default() + .with_range(0..num_rows as u64) + .with_indices(vec![0, null_prefix_rows as u64, num_rows as u64 - 1]) + .with_encoding(encoding); + let list_array = Arc::new(list_array) as ArrayRef; + let pages = encode_v22_pages(list_array.clone()).await; + assert_split_miniblock_layout(&pages, 1, true); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + + #[test_log::test(tokio::test)] + async fn test_sparse_boolean_list_with_amortized_long_empty_prefix() { + let empty_prefix_rows = 62_000usize; + let booleans_per_list = 8_192usize; + let num_rows = empty_prefix_rows + 1; + + let mut offsets = Vec::with_capacity(num_rows + 1); + offsets.extend(std::iter::repeat_n(0i32, empty_prefix_rows + 1)); + let values = (0..booleans_per_list) + .map(|idx| idx % 2 == 0) + .collect::>(); + offsets.push(values.len() as i32); + + let items = BooleanArray::from(values); + let list_array = ListArray::new( + Arc::new(Field::new("item", DataType::Boolean, true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + Arc::new(items), + None, + ); + + let test_cases = TestCases::default() + .with_range(0..num_rows as u64) + .with_indices(vec![0, empty_prefix_rows as u64]) + .with_dense_encodings(); + let list_array = Arc::new(list_array) as ArrayRef; + let pages = encode_v22_pages(list_array.clone()).await; + assert_split_miniblock_layout(&pages, 1, true); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + + fn unsplittable_nested_list(items: ArrayRef, empty_inner_lists: usize) -> ArrayRef { + let mut inner_offsets = vec![0i32; empty_inner_lists + 1]; + inner_offsets.push(items.len() as i32); + let inner_list = ListArray::new( + Arc::new(Field::new("item", items.data_type().clone(), true)), + OffsetBuffer::new(ScalarBuffer::from(inner_offsets)), + items, + None, + ); + Arc::new(ListArray::new( + Arc::new(Field::new("item", inner_list.data_type().clone(), true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0i32, empty_inner_lists as i32 + 1])), + Arc::new(inner_list), + None, + )) + } + + #[rstest] + #[case::boolean(Arc::new(BooleanArray::from(vec![true, false, true, false, true, false, true, false])))] + #[case::string(Arc::new(StringArray::from(vec!["value", "other"])))] + #[test_log::test(tokio::test)] + async fn test_nested_sparse_single_row_falls_back_to_fullzip(#[case] items: ArrayRef) { + let list = unsplittable_nested_list(items, 70_000); + let pages = encode_v22_pages(list.clone()).await; + assert_has_fullzip_layout(&pages); + + let test_cases = TestCases::default() + .with_range(0..1) + .with_indices(vec![0]) + .with_dense_encodings(); + check_round_trip_encoding_of_data(vec![list], &test_cases, HashMap::new()).await; + } + + /// Builds the HNSW-flush repro shape: a dense prefix where every row has + /// `NEIGHBORS_PER_ROW` distinct values, followed by a long tail of empty + /// lists. Mirrors `HNSW::schema()` `__neighbors` / `__dists` columns: + /// dense level-0 lists, then ~6x as many mostly-empty higher-level rows. + fn make_hnsw_shaped_list_u32() -> ListArray { + const DENSE_ROWS: u32 = 5_000; + const NEIGHBORS_PER_ROW: u32 = 32; + const EMPTY_TAIL_ROWS: u32 = 70_000; + + let mut list_builder = ListBuilder::new(UInt32Builder::new()); + let mut next_val: u32 = 0; + for _ in 0..DENSE_ROWS { + for _ in 0..NEIGHBORS_PER_ROW { + list_builder.values().append_value(next_val); + next_val = next_val.wrapping_add(1); + } + list_builder.append(true); + } + for _ in 0..EMPTY_TAIL_ROWS { + list_builder.append(true); + } + list_builder.finish() + } + + /// Reproduces the HNSW-flush shape at v2.2 on the auto path (no + /// `STRUCTURAL_ENCODING` metadata): a dense level-0 prefix followed by a + /// long tail of empty lists. The global levels/values ratio looks dense, + /// so this used to encode as a single mini-block page whose final chunk + /// absorbed every trailing empty list and overflowed the per-chunk `u16` + /// `num_levels`, corrupting the read. The structural page planner now + /// splits on top-level row boundaries: the dense prefix stays on + /// mini-block pages and the empty tail becomes structural-only pages, so + /// the round-trip is lossless without falling back to full-zip. + #[test_log::test(tokio::test)] + async fn test_list_hnsw_shape_splits_to_miniblock_v2_2() { + let list_array = make_hnsw_shaped_list_u32(); + let dense_rows: u64 = 5_000; + let total_rows = list_array.len() as u64; + + let test_cases = TestCases::default() + .with_range(0..1000) + .with_range(dense_rows.saturating_sub(8)..(dense_rows + 8)) + .with_range(0..total_rows) + .with_indices(vec![0, dense_rows - 1, dense_rows, total_rows - 1]) + .with_encoding(TestEncoding::StructuralU32); + let list_array = Arc::new(list_array) as ArrayRef; + let pages = encode_v22_pages(list_array.clone()).await; + assert_split_miniblock_layout(&pages, 1, true); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + + /// Companion to the auto-path test: even when the user explicitly requests + /// `STRUCTURAL_ENCODING_MINIBLOCK`, the structural page planner splits the + /// HNSW shape so every emitted page fits the mini-block per-chunk budget. + /// The request is honored (the dense prefix stays on mini-block pages + /// rather than being forced to full-zip) and the round-trip is lossless. + #[test_log::test(tokio::test)] + async fn test_forced_miniblock_hnsw_shape_splits_to_miniblock_v2_2() { + let list_array = make_hnsw_shaped_list_u32(); + let total_rows = list_array.len() as u64; + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_MINIBLOCK.into(), + ); + + let test_cases = TestCases::default() + .with_range(0..total_rows) + .with_encoding(TestEncoding::StructuralU32); + let list_array = Arc::new(list_array) as ArrayRef; + let pages = try_encode_v22_pages_with_metadata(list_array.clone(), field_metadata.clone()) + .await + .unwrap(); + assert_split_miniblock_layout(&pages, 1, true); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata).await; + } + + #[test] + fn test_plan_decoded_bytes_list_i32_not_supported() { + // Lists encode items using repetition levels so the number of decoded + // bytes depends on the per-row list lengths, which are not known at + // planning time without inspecting page data. For this PR we + // explicitly reject the call rather than returning a wrong estimate. + use crate::decoder::StructuralFieldDecoder; + use crate::encodings::logical::primitive::StructuralPrimitiveFieldDecoder; + use arrow_array::builder::{Int32Builder, ListBuilder}; + use arrow_schema::Field; + + // Build a List array where every row has a random (variable) length + // to illustrate the case we cannot plan for. + let mut builder = ListBuilder::new(Int32Builder::new()); + let lengths = [0usize, 3, 1, 7, 2, 0, 5]; + let mut value = 0i32; + for &len in &lengths { + for _ in 0..len { + builder.values().append_value(value); + value += 1; + } + builder.append(true); + } + let _list_array = builder.finish(); // variable-length rows confirmed + + // Construct the decoder tree directly (no pages needed for this check). + let item_field = Arc::new(Field::new("item", DataType::Int32, true)); + let list_type = DataType::List(item_field.clone()); + + let child = Box::new(StructuralPrimitiveFieldDecoder::new(&item_field, false)); + let decoder = super::StructuralListDecoder::new(child, list_type); + + let err = decoder + .plan_decoded_bytes(lengths.len() as u64) + .unwrap_err(); + assert!( + matches!(err, lance_core::Error::NotSupported { .. }), + "expected NotSupported, got: {err:?}" + ); + assert!( + err.to_string().contains("repetition"), + "error should mention repetition, got: {err}" + ); + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/map.rs b/vendor/lance-encoding/src/encodings/logical/map.rs new file mode 100644 index 000000000..15a00471d --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/map.rs @@ -0,0 +1,801 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ops::Range, sync::Arc}; + +use arrow_array::{Array, ArrayRef, ListArray, MapArray}; +use arrow_schema::DataType; +use futures::future::BoxFuture; +use lance_arrow::deepcopy::deep_copy_nulls; +use lance_arrow::list::ListArrayExt; +use lance_core::{Error, Result}; + +use crate::{ + decoder::{ + DecodedArray, FilterExpression, ScheduledScanLine, SchedulerContext, + StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler, + StructuralSchedulingJob, + }, + encoder::{EncodeTask, FieldEncoder, OutOfLineBuffers}, + repdef::RepDefBuilder, +}; + +/// A structural encoder for map fields +/// +/// Map in Arrow is represented as List> +/// The map's offsets are added to the rep/def builder +/// and the map's entries (struct array) are passed to the child encoder +pub struct MapStructuralEncoder { + keep_original_array: bool, + child: Box, +} + +impl MapStructuralEncoder { + pub fn new(keep_original_array: bool, child: Box) -> Self { + Self { + keep_original_array, + child, + } + } +} + +impl FieldEncoder for MapStructuralEncoder { + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + mut repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + let map_array = array + .as_any() + .downcast_ref::() + .expect("MapEncoder used for non-map data"); + + // Add offsets to RepDefBuilder to handle nullability and list structure + let has_garbage_values = if self.keep_original_array { + repdef.add_offsets(map_array.offsets().clone(), array.nulls().cloned()) + } else { + repdef.add_offsets(map_array.offsets().clone(), deep_copy_nulls(array.nulls())) + }; + + // MapArray is physically a ListArray, so convert and use ListArrayExt + let list_array: ListArray = map_array.clone().into(); + let entries = if has_garbage_values { + list_array.filter_garbage_nulls().trimmed_values() + } else { + list_array.trimmed_values() + }; + + self.child + .maybe_encode(entries, external_buffers, repdef, row_number, num_rows) + } + + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + self.child.flush(external_buffers) + } + + fn num_columns(&self) -> u32 { + self.child.num_columns() + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + self.child.finish(external_buffers) + } +} + +#[derive(Debug)] +pub struct StructuralMapScheduler { + child: Box, +} + +impl StructuralMapScheduler { + pub fn new(child: Box) -> Self { + Self { child } + } +} + +impl StructuralFieldScheduler for StructuralMapScheduler { + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result> { + let child = self.child.schedule_ranges(ranges, filter)?; + + Ok(Box::new(StructuralMapSchedulingJob::new(child))) + } + + fn initialize<'a>( + &'a mut self, + requested_ranges: Option<&'a [Range]>, + filter: &'a FilterExpression, + context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + // Like lists, the child (entries) column is scheduled with the same + // top-level row ranges, so forward them unchanged. + self.child.initialize(requested_ranges, filter, context) + } +} + +/// Scheduling job for map data +/// +/// Scheduling is handled by the child encoder (struct) and nothing special +/// happens here, similar to list. +#[derive(Debug)] +struct StructuralMapSchedulingJob<'a> { + child: Box, +} + +impl<'a> StructuralMapSchedulingJob<'a> { + fn new(child: Box) -> Self { + Self { child } + } +} + +impl StructuralSchedulingJob for StructuralMapSchedulingJob<'_> { + fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result> { + self.child.schedule_next(context) + } +} + +#[derive(Debug)] +pub struct StructuralMapDecoder { + child: Box, + data_type: DataType, +} + +impl StructuralMapDecoder { + pub fn new(child: Box, data_type: DataType) -> Self { + Self { child, data_type } + } +} + +impl StructuralFieldDecoder for StructuralMapDecoder { + fn accept_page(&mut self, child: crate::decoder::LoadedPageShard) -> Result<()> { + self.child.accept_page(child) + } + + fn drain(&mut self, num_rows: u64) -> Result> { + let child_task = self.child.drain(num_rows)?; + Ok(Box::new(StructuralMapDecodeTask::new( + child_task, + self.data_type.clone(), + ))) + } + + fn data_type(&self) -> &DataType { + &self.data_type + } +} + +#[derive(Debug)] +struct StructuralMapDecodeTask { + child_task: Box, + data_type: DataType, +} + +impl StructuralMapDecodeTask { + fn new(child_task: Box, data_type: DataType) -> Self { + Self { + child_task, + data_type, + } + } +} + +impl StructuralDecodeArrayTask for StructuralMapDecodeTask { + fn decode(self: Box) -> Result { + let DecodedArray { + array, + mut repdef, + data_size, + } = self.child_task.decode()?; + + // Decode the offsets from RepDef + let (offsets, validity) = repdef.unravel_offsets::()?; + + // Extract the entries field and keys_sorted from the map data type + let (entries_field, keys_sorted) = match &self.data_type { + DataType::Map(field, keys_sorted) => { + if *keys_sorted { + return Err(Error::not_supported_source( + "Map type decoder does not support keys_sorted=true now" + .to_string() + .into(), + )); + } + (field.clone(), *keys_sorted) + } + _ => { + return Err(Error::schema( + "Map decoder did not have a map field".to_string(), + )); + } + }; + + // Convert the decoded array to StructArray + let entries = array + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::schema("Map entries should be a StructArray".to_string()))? + .clone(); + + // Build the MapArray from offsets, entries, validity, and keys_sorted + let map_array = MapArray::try_new(entries_field, offsets, entries, validity, keys_sorted) + .map_err(|error| Error::invalid_input_source(error.to_string().into()))?; + + Ok(DecodedArray { + array: Arc::new(map_array), + repdef, + data_size, + }) + } +} + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use arrow_array::{ + Array, Int32Array, MapArray, StringArray, StructArray, + builder::{Int32Builder, MapBuilder, StringBuilder}, + }; + use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field, Fields}; + + use crate::decoder::{DecodedArray, StructuralDecodeArrayTask}; + use crate::encoder::{ColumnIndexSequence, EncodingOptions}; + use crate::encodings::logical::primitive::sparse::{ + SparseCountSet, SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan, + SparseValidityMeaning, SparseValiditySet, + }; + use crate::repdef::{CompositeRepDefUnraveler, RepDefUnraveler}; + use crate::testing::{ + TestCases, TestEncoding, check_round_trip_encoding_of_data, test_encoding_strategy, + }; + use arrow_schema::Field as ArrowField; + use lance_core::datatypes::Field as LanceField; + + use super::StructuralMapDecodeTask; + + fn make_map_type(key_type: DataType, value_type: DataType) -> DataType { + // Note: Arrow MapBuilder uses "keys" and "values" as field names (plural) + let entries = Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", key_type, false), + Field::new("values", value_type, true), + ])), + false, + ); + DataType::Map(Arc::new(entries), false) + } + + #[derive(Debug)] + struct StaticMapEntriesTask { + entries: StructArray, + repdef: CompositeRepDefUnraveler, + } + + impl StructuralDecodeArrayTask for StaticMapEntriesTask { + fn decode(self: Box) -> lance_core::Result { + let Self { entries, repdef } = *self; + Ok(DecodedArray { + array: Arc::new(entries), + repdef, + data_size: 0, + }) + } + } + + #[test] + fn malformed_sparse_map_entries_return_invalid_input() { + let entry_fields = Fields::from(vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", DataType::Int32, true), + ]); + let entries = StructArray::try_new( + entry_fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![2])), + ], + Some(NullBuffer::from(vec![false])), + ) + .unwrap(); + let validity = SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions: SparsePositionSet::Empty, + }; + let plan = SparseStructuralPlan { + layers: vec![SparseStructuralLayerPlan::List { + num_slots: 1, + num_child_slots: 1, + non_empty_positions: SparsePositionSet::All { len: 1 }, + counts: SparseCountSet::Constant { value: 1, len: 1 }, + validity, + }], + num_items: 1, + num_visible_items: 1, + }; + let child_task = StaticMapEntriesTask { + entries, + repdef: CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new_sparse(plan)]), + }; + let map_type = DataType::Map( + Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)), + false, + ); + + let Err(err) = + Box::new(StructuralMapDecodeTask::new(Box::new(child_task), map_type)).decode() + else { + panic!("expected malformed map entries to be rejected"); + }; + assert!(matches!(err, lance_core::Error::InvalidInput { .. })); + } + + #[test_log::test(tokio::test)] + async fn test_simple_map() { + // Create a simple Map + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map_builder = MapBuilder::new(None, string_builder, int_builder); + + // Map 1: {"key1": 10, "key2": 20} + map_builder.keys().append_value("key1"); + map_builder.values().append_value(10); + map_builder.keys().append_value("key2"); + map_builder.values().append_value(20); + map_builder.append(true).unwrap(); + + // Map 2: {"key3": 30} + map_builder.keys().append_value("key3"); + map_builder.values().append_value(30); + map_builder.append(true).unwrap(); + + let map_array = map_builder.finish(); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_empty_maps() { + // Test maps with empty entries + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map_builder = MapBuilder::new(None, string_builder, int_builder); + + // Map 1: {"a": 1} + map_builder.keys().append_value("a"); + map_builder.values().append_value(1); + map_builder.append(true).unwrap(); + + // Map 2: {} (empty) + map_builder.append(true).unwrap(); + + // Map 3: null + map_builder.append(false).unwrap(); + + // Map 4: {} (empty) + map_builder.append(true).unwrap(); + + let map_array = map_builder.finish(); + + let test_cases = TestCases::default() + .with_range(0..4) + .with_indices(vec![1]) + .with_indices(vec![2]) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_map_with_null_values() { + // Test Map with null values + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map_builder = MapBuilder::new(None, string_builder, int_builder); + + // Map 1: {"key1": 10, "key2": null} + map_builder.keys().append_value("key1"); + map_builder.values().append_value(10); + map_builder.keys().append_value("key2"); + map_builder.values().append_null(); + map_builder.append(true).unwrap(); + + // Map 2: {"key3": null} + map_builder.keys().append_value("key3"); + map_builder.values().append_null(); + map_builder.append(true).unwrap(); + + let map_array = map_builder.finish(); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_indices(vec![0]) + .with_indices(vec![1]) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_map_in_struct() { + // Test Struct containing Map + // Struct> + + let string_key_builder = StringBuilder::new(); + let string_val_builder = StringBuilder::new(); + let mut map_builder = MapBuilder::new(None, string_key_builder, string_val_builder); + + // First struct: id=1, properties={"name": "Alice", "city": "NYC"} + map_builder.keys().append_value("name"); + map_builder.values().append_value("Alice"); + map_builder.keys().append_value("city"); + map_builder.values().append_value("NYC"); + map_builder.append(true).unwrap(); + + // Second struct: id=2, properties={"name": "Bob"} + map_builder.keys().append_value("name"); + map_builder.values().append_value("Bob"); + map_builder.append(true).unwrap(); + + // Third struct: id=3, properties=null + map_builder.append(false).unwrap(); + + let map_array = Arc::new(map_builder.finish()); + let id_array = Arc::new(Int32Array::from(vec![1, 2, 3])); + + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "properties", + make_map_type(DataType::Utf8, DataType::Utf8), + true, + ), + ]), + vec![id_array, map_array], + None, + ); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_indices(vec![0, 2]) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data( + vec![Arc::new(struct_array)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_map_in_nullable_struct() { + // Test Struct where null struct rows have garbage map entries. + // The encoder must filter these garbage entries before encoding. + let entries_fields = Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", DataType::Int32, true), + ]); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(entries_fields.clone()), + false, + )); + let map_entries = StructArray::new( + entries_fields, + vec![ + Arc::new(StringArray::from(vec!["a", "garbage", "b"])), + Arc::new(Int32Array::from(vec![1, 999, 2])), + ], + None, + ); + // map0: {"a": 1}, map1 (garbage): {"garbage": 999}, map2: {"b": 2} + let map_array: Arc = Arc::new(MapArray::new( + entries_field, + OffsetBuffer::new(ScalarBuffer::from(vec![0, 1, 2, 3])), + map_entries, + None, // No nulls at map level - nulls come from struct + false, + )); + + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("id", DataType::Int32, true), + Field::new("props", map_array.data_type().clone(), true), + ]), + vec![ + Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)])), + map_array, + ], + Some(NullBuffer::from(vec![true, false, true])), // Middle row is null + ); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data( + vec![Arc::new(struct_array)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_list_of_maps() { + // Test List> + use arrow_array::builder::ListBuilder; + + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let map_builder = MapBuilder::new(None, string_builder, int_builder); + let mut list_builder = ListBuilder::new(map_builder); + + // List 1: [{"a": 1}, {"b": 2}] + list_builder.values().keys().append_value("a"); + list_builder.values().values().append_value(1); + list_builder.values().append(true).unwrap(); + + list_builder.values().keys().append_value("b"); + list_builder.values().values().append_value(2); + list_builder.values().append(true).unwrap(); + + list_builder.append(true); + + // List 2: [{"c": 3}] + list_builder.values().keys().append_value("c"); + list_builder.values().values().append_value(3); + list_builder.values().append(true).unwrap(); + + list_builder.append(true); + + // List 3: [] (empty list) + list_builder.append(true); + + let list_array = list_builder.finish(); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_indices(vec![0, 2]) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_nested_map() { + // Test Map> + // This is more complex as we need to build nested maps manually + + // Build inner maps first + let inner_string_builder = StringBuilder::new(); + let inner_int_builder = Int32Builder::new(); + let mut inner_map_builder1 = MapBuilder::new(None, inner_string_builder, inner_int_builder); + + // Inner map 1: {"x": 10} + inner_map_builder1.keys().append_value("x"); + inner_map_builder1.values().append_value(10); + inner_map_builder1.append(true).unwrap(); + + // Inner map 2: {"y": 20, "z": 30} + inner_map_builder1.keys().append_value("y"); + inner_map_builder1.values().append_value(20); + inner_map_builder1.keys().append_value("z"); + inner_map_builder1.values().append_value(30); + inner_map_builder1.append(true).unwrap(); + + let inner_maps = Arc::new(inner_map_builder1.finish()); + + // Build outer map keys + let outer_keys = Arc::new(StringArray::from(vec!["key1", "key2"])); + + // Build outer map structure + let entries_struct = StructArray::new( + Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new( + "value", + make_map_type(DataType::Utf8, DataType::Int32), + true, + ), + ]), + vec![outer_keys, inner_maps], + None, + ); + + let offsets = OffsetBuffer::new(ScalarBuffer::::from(vec![0, 2])); + let entries_field = Field::new("entries", entries_struct.data_type().clone(), false); + + let outer_map = MapArray::new( + Arc::new(entries_field), + offsets, + entries_struct, + None, + false, + ); + + let test_cases = TestCases::default() + .with_range(0..1) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(outer_map)], &test_cases, HashMap::new()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_map_different_key_types() { + // Test Map (integer keys) + let int_builder = Int32Builder::new(); + let string_builder = StringBuilder::new(); + let mut map_builder = MapBuilder::new(None, int_builder, string_builder); + + // Map 1: {1: "one", 2: "two"} + map_builder.keys().append_value(1); + map_builder.values().append_value("one"); + map_builder.keys().append_value(2); + map_builder.values().append_value("two"); + map_builder.append(true).unwrap(); + + // Map 2: {3: "three"} + map_builder.keys().append_value(3); + map_builder.values().append_value("three"); + map_builder.append(true).unwrap(); + + let map_array = map_builder.finish(); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_indices(vec![0, 1]) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_map_with_extreme_sizes() { + // Test maps with large number of entries + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map_builder = MapBuilder::new(None, string_builder, int_builder); + + // Create a map with many entries + for i in 0..100 { + map_builder.keys().append_value(format!("key{}", i)); + map_builder.values().append_value(i); + } + map_builder.append(true).unwrap(); + + // Create a second map with no entries + map_builder.append(true).unwrap(); + + let map_array = map_builder.finish(); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_map_all_null() { + // Test map where all entries are null + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map_builder = MapBuilder::new(None, string_builder, int_builder); + + // All null maps + map_builder.append(false).unwrap(); // null + map_builder.append(false).unwrap(); // null + + let map_array = map_builder.finish(); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_u32_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_map_encoder_keep_original_array_scenarios() { + // Test scenarios that highlight the difference between keep_original_array=true/false + // This test focuses on round-trip behavior which should be equivalent in both cases + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map_builder = MapBuilder::new(None, string_builder, int_builder); + + // Create a map with mixed null and non-null values to test both scenarios + // Map 1: {"key1": 10, "key2": null} + map_builder.keys().append_value("key1"); + map_builder.values().append_value(10); + map_builder.keys().append_value("key2"); + map_builder.values().append_null(); + map_builder.append(true).unwrap(); + + // Map 2: null + map_builder.append(false).unwrap(); + + // Map 3: {"key3": 30} + map_builder.keys().append_value("key3"); + map_builder.values().append_value(30); + map_builder.append(true).unwrap(); + + let map_array = map_builder.finish(); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_indices(vec![0, 1, 2]) + .with_u32_structural_encodings(); + + // This test ensures that regardless of the internal keep_original_array setting, + // the end-to-end behavior produces equivalent results + check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) + .await; + } + + #[test] + fn test_map_not_supported_write_in_v2_1() { + // Create a map field using Arrow Field first, then convert to Lance Field + let map_arrow_field = ArrowField::new( + "map_field", + make_map_type(DataType::Utf8, DataType::Int32), + true, + ); + let map_field = LanceField::try_from(&map_arrow_field).unwrap(); + + // Test encoder: Try to create encoder with V2_1 version - should fail + let encoder_strategy = test_encoding_strategy(TestEncoding::StructuralU16); + let mut column_index = ColumnIndexSequence::default(); + let options = EncodingOptions::default(); + + let encoder_result = crate::testing::create_test_field_encoder( + encoder_strategy.as_ref(), + &map_field, + &mut column_index, + &options, + ); + + assert!( + encoder_result.is_err(), + "Map type should not be supported in V2_1 for encoder" + ); + let Err(encoder_err) = encoder_result else { + panic!("Expected error but got Ok") + }; + + let encoder_err_msg = format!("{}", encoder_err); + assert!( + encoder_err_msg.contains("not enabled by the selected file format"), + "unexpected encoder error: {encoder_err_msg}" + ); + assert!( + encoder_err_msg.contains("Map data type"), + "Encoder error message should mention Map data type, got: {}", + encoder_err_msg + ); + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive.rs b/vendor/lance-encoding/src/encodings/logical/primitive.rs new file mode 100644 index 000000000..77484be0b --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive.rs @@ -0,0 +1,12571 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ + any::Any, + collections::{HashMap, VecDeque}, + env, + fmt::Debug, + iter, + ops::Range, + sync::Arc, + vec, +}; + +use crate::{ + constants::{ + STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, + STRUCTURAL_ENCODING_SPARSE, + }, + data::DictionaryDataBlock, + encodings::logical::primitive::blob::{BlobDescriptionPageScheduler, BlobPageScheduler}, + format::{ + ProtobufUtils21, + pb21::{self, CompressiveEncoding, PageLayout, compressive_encoding::Compression}, + }, +}; +use arrow_array::{ + Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, new_null_array, types::UInt64Type, +}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer}; +use arrow_schema::{DataType, Field as ArrowField}; +use bytes::Bytes; +use futures::{FutureExt, TryStreamExt, future::BoxFuture, stream::FuturesOrdered}; +use itertools::Itertools; +use lance_arrow::DataTypeExt; +use lance_arrow::deepcopy::deep_copy_nulls; +use lance_core::{ + cache::{CacheKey, CacheKeySchema, Context, DeepSizeOf, KeyBuilder, LanceCache}, + error::{Error, LanceOptionExt}, + utils::bit::pad_bytes, +}; +use log::{debug, trace}; + +use crate::encodings::logical::primitive::miniblock::MiniBlockChunk; +use crate::encodings::physical::rle::{RleDecompressor, RleRuns}; +use crate::utils::bytepack::ByteUnpacker; +use crate::{ + compression::{ + BlockDecompressor, CompressionStrategy, DecompressionStrategy, MiniBlockDecompressor, + compress_required_block, create_rle_decompressor, + }, + data::{AllNullDataBlock, DataBlock, VariableWidthBlock}, + decoder::{IO_REQUEST_BATCH_BYTES, RequestBatch}, + utils::bytepack::BytepackedIntegerEncoder, +}; +use crate::{ + compression::{FixedPerValueDecompressor, VariablePerValueDecompressor}, + encodings::logical::primitive::fullzip::PerValueDataBlock, +}; +use crate::{ + encodings::logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressionContext}, + statistics::{ComputeStat, GetStat, Stat}, +}; +use crate::{ + repdef::{ + CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation, + MiniBlockRepDefBudget, NormalizedStructuralPlan, RepDefSlicer, SerializedRepDefs, + build_control_word_iterator, max_visible_level, + }, + utils::accumulation::AccumulationQueue, +}; +use lance_core::{Result, datatypes::Field, utils::tokio::spawn_cpu}; + +use crate::constants::{ + COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY, + DICT_SIZE_RATIO_META_KEY, DICT_VALUES_COMPRESSION_ENV_VAR, + DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR, DICT_VALUES_COMPRESSION_LEVEL_META_KEY, + DICT_VALUES_COMPRESSION_META_KEY, +}; +use crate::{ + EncodingsIo, + buffer::LanceBuffer, + data::{BlockInfo, DataBlockBuilder, FixedWidthDataBlock}, + decoder::{ + ColumnInfo, DecodePageTask, DecodedArray, DecodedPage, FilterExpression, LoadedPageShard, + MessageType, PageEncoding, PageInfo, ScheduledScanLine, SchedulerContext, + StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler, + StructuralPageDecoder, StructuralSchedulingJob, UnloadedPageShard, + }, + encoder::{ + EncodeTask, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, OutOfLineBuffers, + }, + repdef::{LevelBuffer, RepDefBuilder, RepDefUnraveler}, +}; + +pub mod blob; +mod chunk_index; +pub mod constant; +pub mod dict; +pub mod fullzip; +mod layout; +pub mod miniblock; +pub(crate) mod sparse; + +use chunk_index::{ItemCounts, MiniBlockChunkIndex, PrefixSums, RowMapping, parse_nested_rep}; + +const FILL_BYTE: u8 = 0xFE; +const DEFAULT_DICT_DIVISOR: u64 = 2; +const DEFAULT_DICT_MAX_CARDINALITY: u64 = 100_000; +const DEFAULT_DICT_SIZE_RATIO: f64 = 0.8; +const DEFAULT_DICT_VALUES_COMPRESSION: &str = "lz4"; + +/// Largest level count a direct legacy u16-value/U8-run RLE frame can prove. +/// +/// Mini-block level payload sizes are u16. After the eight-byte frame header, each run needs a +/// two-byte value and a one-byte length, and each U8 run can represent at most 255 levels. +const MAX_LEGACY_RLE_LEVELS: u64 = ((u16::MAX as u64 - std::mem::size_of::() as u64) + / (std::mem::size_of::() as u64 + std::mem::size_of::() as u64)) + * u8::MAX as u64; + +struct PageLoadTask { + decoder_fut: BoxFuture<'static, Result>>, + num_rows: u64, +} + +/// The ordered initialization-buffer slots for one page. +/// +/// Optional slots preserve the semantic position of absent buffers so range +/// declaration and buffer consumption share one representation. +#[derive(Debug)] +struct PageInitialization { + buffer_ranges: Vec>>, +} + +impl PageInitialization { + fn new(buffer_ranges: impl IntoIterator>>) -> Self { + Self { + buffer_ranges: buffer_ranges.into_iter().collect(), + } + } + + fn required_ranges(&self) -> impl Iterator> { + self.buffer_ranges.iter().flatten() + } + + fn num_required_ranges(&self) -> usize { + self.required_ranges().count() + } + + fn with_buffers(self, buffers: Vec) -> Result { + let expected = self.num_required_ranges(); + if buffers.len() != expected { + return Err(Error::invalid_input_source( + format!( + "Page initialization received {} buffers for {expected} declared ranges", + buffers.len() + ) + .into(), + )); + } + + let mut buffers = buffers.into_iter(); + let mut buffer_slots = Vec::with_capacity(self.buffer_ranges.len()); + for range in self.buffer_ranges { + if range.is_some() { + let buffer = buffers.next().ok_or_else(|| { + Error::internal("Page initialization buffer was not assigned") + })?; + buffer_slots.push(Some(buffer)); + } else { + buffer_slots.push(None); + } + } + debug_assert!( + buffers.next().is_none(), + "validated page initialization buffer count must be exhausted" + ); + Ok(PageInitializationBuffers { buffer_slots }) + } +} + +#[derive(Debug)] +struct PageInitializationBuffers { + buffer_slots: Vec>, +} + +impl PageInitializationBuffers { + fn try_into_array(self, layout: &str) -> Result<[Option; N]> { + let actual = self.buffer_slots.len(); + self.buffer_slots.try_into().map_err(|_| { + Error::internal(format!( + "{layout} page initialization declared {actual} buffer slots, expected {N}" + )) + }) + } + + fn into_required(self, layout: &str) -> Result> { + self.buffer_slots + .into_iter() + .enumerate() + .map(|(index, buffer)| { + buffer.ok_or_else(|| { + Error::internal(format!( + "{layout} page initialization buffer slot {index} is absent" + )) + }) + }) + .collect() + } +} + +/// A trait for figuring out how to schedule the data within +/// a single page. +trait StructuralPageScheduler: std::fmt::Debug + Send { + /// Whether this page requires initialized metadata before scheduling. + fn needs_initialization(&self) -> bool { + true + } + + /// Describes the file byte ranges this page must read to initialize its + /// metadata (chunk metadata, dictionary, repetition index, ...). + /// + /// The field scheduler sorts the combined I/O request by file offset and + /// restores buffers to the declared slots afterward. + fn init_layout(&self) -> Result; + + /// Returns initialized state from the bytes declared by [`Self::init_layout`]. + /// + /// The caller attaches the returned state with [`Self::try_load`]. + fn init_from_buffers<'a>( + &'a mut self, + buffers: PageInitializationBuffers, + io: &Arc, + ) -> BoxFuture<'a, Result>>; + + /// Tries to attach metadata from a previous initialization. + fn try_load(&mut self, data: &Arc) -> Result<()>; + + // Identifies the requested decode shape (e.g. blob descriptor struct vs + // raw bytes). Blob columns can produce multiple page scheduler variants + // for the same physical column depending on the target field's data type, + // and the cached page state types differ per variant. The stable view is + // mixed into the cache key so different variants do not collide. + fn cache_view(&self) -> PageCacheView { + PageCacheView::Native + } + + /// Schedules the read of the given ranges in the page + /// + /// The read may be split into multiple "shards" if the page is extremely large. + /// Each shard maps to one or more rows and can be decoded independently. + /// + /// Note: this sharding is for splitting up very large pages into smaller reads to + /// avoid buffering too much data in memory. It is not related to the batch size or + /// compute units in any way. + fn schedule_ranges( + &self, + ranges: &[Range], + io: &Arc, + ) -> Result>; +} + +/// Metadata describing the decoded size of a mini-block +#[derive(Debug)] +struct ChunkMeta { + num_values: u64, + chunk_size_bytes: u64, + offset_bytes: u64, +} + +/// A mini-block chunk that has been decoded and decompressed +#[derive(Debug, Clone)] +struct DecodedMiniBlockChunk { + rep: Option>, + def: Option>, + values: DataBlock, +} + +/// A task to decode a one or more mini-blocks of data into an output batch +/// +/// Note: Two batches might share the same mini-block of data. When this happens +/// then each batch gets a copy of the block and each batch decodes the block independently. +/// +/// This means we have duplicated work but it is necessary to avoid having to synchronize +/// the decoding of the block. (TODO: test this theory) +#[derive(Debug)] +struct DecodeMiniBlockTask { + rep_decompressor: Option>, + def_decompressor: Option>, + value_decompressor: Arc, + dictionary_data: Option>, + def_meaning: Arc<[DefinitionInterpretation]>, + num_buffers: u64, + max_visible_level: u16, + instructions: Vec<(ChunkDrainInstructions, LoadedChunk)>, + has_large_chunk: bool, +} + +impl DecodeMiniBlockTask { + fn decoded_size_bytes(&self) -> Option { + if self.rep_decompressor.is_some() || self.def_decompressor.is_some() { + return None; + } + let num_values = self + .instructions + .iter() + .try_fold(0_u64, |total, (instruction, _)| { + total.checked_add(instruction.rows_to_take) + })?; + self.value_decompressor.decoded_size_bytes(num_values) + } + + fn resolve_num_levels( + rep_decompressor: Option<&dyn BlockDecompressor>, + rep_levels: Option<&LanceBuffer>, + def_decompressor: Option<&dyn BlockDecompressor>, + def_levels: Option<&LanceBuffer>, + declared_num_levels: u16, + ) -> Result { + let rep_num_levels = match (rep_decompressor, rep_levels) { + (Some(decompressor), Some(levels)) => decompressor.infer_num_values(levels)?, + (None, None) => None, + _ => { + return Err(Error::invalid_input_source( + "miniblock repetition codec and payload presence disagree".into(), + )); + } + }; + let def_num_levels = match (def_decompressor, def_levels) { + (Some(decompressor), Some(levels)) => decompressor.infer_num_values(levels)?, + (None, None) => None, + _ => { + return Err(Error::invalid_input_source( + "miniblock definition codec and payload presence disagree".into(), + )); + } + }; + + let inferred_num_levels = match (rep_num_levels, def_num_levels) { + (Some(rep_num_levels), Some(def_num_levels)) if rep_num_levels != def_num_levels => { + return Err(Error::invalid_input_source( + format!( + "miniblock structural streams disagree on the level count: repetition inferred {rep_num_levels}, definition inferred {def_num_levels}" + ) + .into(), + )); + } + (Some(num_levels), _) | (_, Some(num_levels)) => num_levels, + (None, None) => return Ok(u64::from(declared_num_levels)), + }; + + let declared_num_levels = u64::from(declared_num_levels); + if inferred_num_levels == declared_num_levels { + return Ok(inferred_num_levels); + } + let u16_modulus = u64::from(u16::MAX) + 1; + if inferred_num_levels <= declared_num_levels + || inferred_num_levels % u16_modulus != declared_num_levels + { + return Err(Error::invalid_input_source( + format!( + "miniblock payload proves {inferred_num_levels} levels but the header declared {declared_num_levels}; the counts are not congruent modulo 65536" + ) + .into(), + )); + } + if inferred_num_levels > MAX_LEGACY_RLE_LEVELS { + return Err(Error::invalid_input_source( + format!( + "miniblock payload proves {inferred_num_levels} levels, exceeding the legacy RLE payload bound of {MAX_LEGACY_RLE_LEVELS}" + ) + .into(), + )); + } + Ok(inferred_num_levels) + } + + fn decode_levels( + decompressor: &dyn BlockDecompressor, + levels: LanceBuffer, + expected_num_levels: u64, + ) -> Result> { + let levels = decompressor.decompress(Some(levels), expected_num_levels)?; + let levels = levels.as_fixed_width().ok_or_else(|| { + Error::invalid_input_source( + "miniblock levels did not decode to fixed-width data".into(), + ) + })?; + if levels.num_values != expected_num_levels { + return Err(Error::invalid_input_source( + format!( + "miniblock levels decoded {} values, expected {expected_num_levels}", + levels.num_values + ) + .into(), + )); + } + if levels.bits_per_value != 16 { + return Err(Error::invalid_input_source( + format!( + "miniblock levels decoded {} bits per value, expected 16", + levels.bits_per_value + ) + .into(), + )); + } + let expected_bytes = expected_num_levels.checked_mul(2).ok_or_else(|| { + Error::invalid_input_source( + format!("miniblock level byte count overflowed for {expected_num_levels} levels") + .into(), + ) + })?; + if levels.data.len() as u64 != expected_bytes { + return Err(Error::invalid_input_source( + format!( + "miniblock levels decoded {} bytes, expected {expected_bytes}", + levels.data.len() + ) + .into(), + )); + } + Ok(levels.data.borrow_to_typed_slice::()) + } + + // We are building a LevelBuffer (levels) and want to copy into it `total_len` + // values from `level_buf` starting at `offset`. + // + // We need to handle both the case where `levels` is None (no nulls encountered + // yet) and the case where `level_buf` is None (the input we are copying from has + // no nulls) + fn extend_levels( + range: Range, + levels: &mut Option, + level_buf: &Option>, + dest_offset: usize, + ) { + if let Some(level_buf) = level_buf { + if levels.is_none() { + // This is the first non-empty def buf we've hit, fill in the past + // with 0 (valid) + let mut new_levels_vec = + LevelBuffer::with_capacity(dest_offset + (range.end - range.start) as usize); + new_levels_vec.extend(iter::repeat_n(0, dest_offset)); + *levels = Some(new_levels_vec); + } + levels.as_mut().unwrap().extend( + level_buf.as_ref()[range.start as usize..range.end as usize] + .iter() + .copied(), + ); + } else if let Some(levels) = levels { + let num_values = (range.end - range.start) as usize; + // This is an all-valid level_buf but we had nulls earlier and so we + // need to materialize it + levels.extend(iter::repeat_n(0, num_values)); + } + } + + /// Maps a range of rows to a range of items and a range of levels + /// + /// If there is no repetition information this just returns the range as-is. + /// + /// If there is repetition information then we need to do some work to figure out what + /// range of items corresponds to the requested range of rows. + /// + /// For example, if the data is [[1, 2, 3], [4, 5], [6, 7]] and the range is 1..2 (i.e. just row + /// 1) then the user actually wants items 3..5. In the above case the rep levels would be: + /// + /// Idx: 0 1 2 3 4 5 6 + /// Rep: 1 0 0 1 0 1 0 + /// + /// So the start (1) maps to the second 1 (idx=3) and the end (2) maps to the third 1 (idx=5) + /// + /// If there are invisible items then we don't count them when calculating the range of items we + /// are interested in but we do count them when calculating the range of levels we are interested + /// in. As a result we have to return both the item range (first return value) and the level range + /// (second return value). + /// + /// For example, if the data is [[1, 2, 3], [4, 5], NULL, [6, 7, 8]] and the range is 2..4 then the + /// user wants items 5..8 but they want levels 5..9. In the above case the rep/def levels would be: + /// + /// Idx: 0 1 2 3 4 5 6 7 8 + /// Rep: 1 0 0 1 0 1 1 0 0 + /// Def: 0 0 0 0 0 1 0 0 0 + /// Itm: 1 2 3 4 5 6 7 8 + /// + /// Finally, we have to contend with the fact that chunks may or may not start with a "preamble" of + /// trailing values that finish up a list from the previous chunk. In this case the first item does + /// not start at max_rep because it is a continuation of the previous chunk. For our purposes we do + /// not consider this a "row" and so the range 0..1 will refer to the first row AFTER the preamble. + /// + /// We have a separate parameter (`preamble_action`) to control whether we want the preamble or not. + /// + /// Note that the "trailer" is considered a "row" and if we want it we should include it in the range. + fn map_range( + range: Range, + rep: Option<&impl AsRef<[u16]>>, + def: Option<&impl AsRef<[u16]>>, + max_rep: u16, + max_visible_def: u16, + // The total number of items (not rows) in the chunk. This is not quite the same as + // rep.len() / def.len() because it doesn't count invisible items + total_items: u64, + preamble_action: PreambleAction, + ) -> (Range, Range) { + if let Some(rep) = rep { + let mut rep = rep.as_ref(); + // If there is a preamble and we need to skip it then do that first. The work is the same + // whether there is def information or not + let mut items_in_preamble = 0_u64; + let first_row_start = match preamble_action { + PreambleAction::Skip | PreambleAction::Take => { + let first_row_start = if let Some(def) = def.as_ref() { + let mut first_row_start = None; + for (idx, (rep, def)) in rep.iter().zip(def.as_ref()).enumerate() { + if *rep == max_rep { + first_row_start = Some(idx as u64); + break; + } + if *def <= max_visible_def { + items_in_preamble += 1; + } + } + first_row_start + } else { + let first_row_start = + rep.iter().position(|&r| r == max_rep).map(|r| r as u64); + items_in_preamble = first_row_start.unwrap_or(rep.len() as u64); + first_row_start + }; + // It is possible for a chunk to be entirely partial values but if it is then it + // should never show up as a preamble to skip + if first_row_start.is_none() { + assert!(preamble_action == PreambleAction::Take); + return (0..total_items, 0..rep.len() as u64); + } + let first_row_start = first_row_start.unwrap(); + rep = &rep[first_row_start as usize..]; + first_row_start + } + PreambleAction::Absent => { + debug_assert!(rep[0] == max_rep); + 0 + } + }; + + // We hit this case when all we needed was the preamble + if range.start == range.end { + debug_assert!(preamble_action == PreambleAction::Take); + debug_assert!(items_in_preamble <= total_items); + return (0..items_in_preamble, 0..first_row_start); + } + assert!(range.start < range.end); + + let mut rows_seen = 0; + let mut new_start = 0; + let mut new_levels_start = 0; + + if let Some(def) = def { + let def = &def.as_ref()[first_row_start as usize..]; + + // range.start == 0 always maps to 0 (even with invis items), otherwise we need to walk + let mut lead_invis_seen = 0; + + if range.start > 0 { + if def[0] > max_visible_def { + lead_invis_seen += 1; + } + for (idx, (rep, def)) in rep.iter().zip(def).skip(1).enumerate() { + if *rep == max_rep { + rows_seen += 1; + if rows_seen == range.start { + new_start = idx as u64 + 1 - lead_invis_seen; + new_levels_start = idx as u64 + 1; + break; + } + } + if *def > max_visible_def { + lead_invis_seen += 1; + } + } + } + + rows_seen += 1; + + let mut new_end = u64::MAX; + let mut new_levels_end = rep.len() as u64; + let new_start_is_visible = def[new_levels_start as usize] <= max_visible_def; + let mut tail_invis_seen = if new_start_is_visible { 0 } else { 1 }; + for (idx, (rep, def)) in rep[(new_levels_start + 1) as usize..] + .iter() + .zip(&def[(new_levels_start + 1) as usize..]) + .enumerate() + { + if *rep == max_rep { + rows_seen += 1; + if rows_seen == range.end + 1 { + new_end = idx as u64 + new_start + 1 - tail_invis_seen; + new_levels_end = idx as u64 + new_levels_start + 1; + break; + } + } + if *def > max_visible_def { + tail_invis_seen += 1; + } + } + + if new_end == u64::MAX { + new_levels_end = rep.len() as u64; + let total_invis_seen = lead_invis_seen + tail_invis_seen; + new_end = rep.len() as u64 - total_invis_seen; + } + + assert_ne!(new_end, u64::MAX); + + // Adjust for any skipped preamble + if preamble_action == PreambleAction::Skip { + new_start += items_in_preamble; + new_end += items_in_preamble; + new_levels_start += first_row_start; + new_levels_end += first_row_start; + } else if preamble_action == PreambleAction::Take { + debug_assert_eq!(new_start, 0); + debug_assert_eq!(new_levels_start, 0); + new_end += items_in_preamble; + new_levels_end += first_row_start; + } + + debug_assert!(new_end <= total_items); + (new_start..new_end, new_levels_start..new_levels_end) + } else { + // Easy case, there are no invisible items, so we don't need to check for them + // The items range and levels range will be the same. We do still need to walk + // the rep levels to find the row boundaries + + // range.start == 0 always maps to 0, otherwise we need to walk + if range.start > 0 { + for (idx, rep) in rep.iter().skip(1).enumerate() { + if *rep == max_rep { + rows_seen += 1; + if rows_seen == range.start { + new_start = idx as u64 + 1; + break; + } + } + } + } + let mut new_end = rep.len() as u64; + // range.end == max_items always maps to rep.len(), otherwise we need to walk + if range.end < total_items { + for (idx, rep) in rep[(new_start + 1) as usize..].iter().enumerate() { + if *rep == max_rep { + rows_seen += 1; + if rows_seen == range.end { + new_end = idx as u64 + new_start + 1; + break; + } + } + } + } + + // Adjust for any skipped preamble + if preamble_action == PreambleAction::Skip { + new_start += first_row_start; + new_end += first_row_start; + } else if preamble_action == PreambleAction::Take { + debug_assert_eq!(new_start, 0); + new_end += first_row_start; + } + + debug_assert!(new_end <= total_items); + (new_start..new_end, new_start..new_end) + } + } else { + // No repetition info, easy case, just use the range as-is and the item + // and level ranges are the same + (range.clone(), range) + } + } + + // read `num_buffers` buffer sizes from `buf` starting at `offset` + fn read_buffer_sizes( + buf: &[u8], + offset: &mut usize, + num_buffers: u64, + ) -> Vec { + let read_size = if LARGE { 4 } else { 2 }; + (0..num_buffers) + .map(|_| { + let bytes = &buf[*offset..*offset + read_size]; + let size = if LARGE { + u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) + } else { + // the buffer size is read from u16 but is stored as u32 after decoding for consistency + u16::from_le_bytes([bytes[0], bytes[1]]) as u32 + }; + *offset += read_size; + size + }) + .collect() + } + + // Unserialize a miniblock into a collection of vectors + fn decode_miniblock_chunk( + &self, + buf: &LanceBuffer, + items_in_chunk: u64, + ) -> Result { + let mut offset = 0; + let num_levels = u16::from_le_bytes([buf[offset], buf[offset + 1]]); + offset += 2; + + let rep_size = if self.rep_decompressor.is_some() { + let rep_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]); + offset += 2; + Some(rep_size) + } else { + None + }; + let def_size = if self.def_decompressor.is_some() { + let def_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]); + offset += 2; + Some(def_size) + } else { + None + }; + + let buffer_sizes = if self.has_large_chunk { + Self::read_buffer_sizes::(buf, &mut offset, self.num_buffers) + } else { + Self::read_buffer_sizes::(buf, &mut offset, self.num_buffers) + }; + + offset += pad_bytes::(offset); + + let rep = rep_size.map(|rep_size| { + let rep = buf.slice_with_length(offset, rep_size as usize); + offset += rep_size as usize; + offset += pad_bytes::(offset); + rep + }); + + let def = def_size.map(|def_size| { + let def = buf.slice_with_length(offset, def_size as usize); + offset += def_size as usize; + offset += pad_bytes::(offset); + def + }); + + let num_levels = Self::resolve_num_levels( + self.rep_decompressor.as_deref(), + rep.as_ref(), + self.def_decompressor.as_deref(), + def.as_ref(), + num_levels, + )?; + + let buffers = buffer_sizes + .into_iter() + .map(|buf_size| { + let buf = buf.slice_with_length(offset, buf_size as usize); + offset += buf_size as usize; + offset += pad_bytes::(offset); + buf + }) + .collect::>(); + + let values = self + .value_decompressor + .decompress(buffers, items_in_chunk)?; + + let rep = rep + .map(|rep| { + Self::decode_levels( + self.rep_decompressor.as_ref().unwrap().as_ref(), + rep, + num_levels, + ) + }) + .transpose()?; + let def = def + .map(|def| { + Self::decode_levels( + self.def_decompressor.as_ref().unwrap().as_ref(), + def, + num_levels, + ) + }) + .transpose()?; + + if let (Some(rep), Some(def)) = (&rep, &def) + && rep.len() != def.len() + { + return Err(Error::invalid_input_source( + format!( + "miniblock structural streams decoded different level counts: repetition {}, definition {}", + rep.len(), + def.len() + ) + .into(), + )); + } + + Ok(DecodedMiniBlockChunk { rep, def, values }) + } +} + +impl DecodePageTask for DecodeMiniBlockTask { + fn decode(self: Box) -> Result { + // First, we create output buffers for the rep and def and data + let mut repbuf: Option = None; + let mut defbuf: Option = None; + + let max_rep = self.def_meaning.iter().filter(|l| l.is_list()).count() as u16; + + let estimated_size_bytes = self.decoded_size_bytes().unwrap_or_else(|| { + // Variable-width and rep/def encoded output sizes are not known before decoding. + self.instructions + .iter() + .map(|(_, chunk)| chunk.data.len() as u64) + .sum::() + * 2 + }); + let mut data_builder = DataBlockBuilder::with_capacity_estimate(estimated_size_bytes); + + // We need to keep track of the offset into repbuf/defbuf that we are building up + let mut level_offset = 0; + + // Pre-compute caching needs for each chunk by checking if the next chunk is the same + let needs_caching: Vec = self + .instructions + .windows(2) + .map(|w| w[0].1.chunk_idx == w[1].1.chunk_idx) + .chain(std::iter::once(false)) // the last one never needs caching + .collect(); + + // Cache for storing decoded chunks when beneficial + let mut chunk_cache: Option<(usize, DecodedMiniBlockChunk)> = None; + + // Now we iterate through each instruction and process it + for (idx, (instructions, chunk)) in self.instructions.iter().enumerate() { + let should_cache_this_chunk = needs_caching[idx]; + + let decoded_chunk = match &chunk_cache { + Some((cached_chunk_idx, cached_chunk)) if *cached_chunk_idx == chunk.chunk_idx => { + // Clone only when we have a cache hit (much cheaper than decoding) + cached_chunk.clone() + } + _ => { + // Cache miss, need to decode + let decoded = self.decode_miniblock_chunk(&chunk.data, chunk.items_in_chunk)?; + + // Only update cache if this chunk will benefit the next access + if should_cache_this_chunk { + chunk_cache = Some((chunk.chunk_idx, decoded.clone())); + } + decoded + } + }; + + let DecodedMiniBlockChunk { rep, def, values } = decoded_chunk; + + // Our instructions tell us which rows we want to take from this chunk + let row_range_start = + instructions.rows_to_skip + instructions.chunk_instructions.rows_to_skip; + let row_range_end = row_range_start + instructions.rows_to_take; + + // We use the rep info to map the row range to an item range / levels range + let (item_range, level_range) = Self::map_range( + row_range_start..row_range_end, + rep.as_ref(), + def.as_ref(), + max_rep, + self.max_visible_level, + chunk.items_in_chunk, + instructions.preamble_action, + ); + if item_range.end - item_range.start > chunk.items_in_chunk { + return Err(lance_core::Error::internal(format!( + "Item range {:?} is greater than chunk items in chunk {:?}", + item_range, chunk.items_in_chunk + ))); + } + + // Now we append the data to the output buffers + Self::extend_levels(level_range.clone(), &mut repbuf, &rep, level_offset); + Self::extend_levels(level_range.clone(), &mut defbuf, &def, level_offset); + level_offset += (level_range.end - level_range.start) as usize; + data_builder.append(&values, item_range)?; + } + + let mut data = data_builder.finish(); + + let unraveler = + RepDefUnraveler::new(repbuf, defbuf, self.def_meaning.clone(), data.num_values()); + + if let Some(dictionary) = &self.dictionary_data { + // Don't decode here, that happens later (if needed) + let DataBlock::FixedWidth(indices) = data else { + return Err(lance_core::Error::internal(format!( + "Expected FixedWidth DataBlock for dictionary indices, got {:?}", + data + ))); + }; + data = DataBlock::Dictionary(DictionaryDataBlock::from_parts( + indices, + dictionary.as_ref().clone(), + )); + } + + Ok(DecodedPage { + data, + repdef: unraveler, + }) + } +} + +/// A chunk that has been loaded by the miniblock scheduler (but not +/// yet decoded) +#[derive(Debug)] +struct LoadedChunk { + data: LanceBuffer, + items_in_chunk: u64, + byte_range: Range, + chunk_idx: usize, +} + +impl Clone for LoadedChunk { + fn clone(&self) -> Self { + Self { + // Safe as we always create borrowed buffers here + data: self.data.clone(), + items_in_chunk: self.items_in_chunk, + byte_range: self.byte_range.clone(), + chunk_idx: self.chunk_idx, + } + } +} + +/// Decodes mini-block formatted data. See [`PrimitiveStructuralEncoder`] for more +/// details on the different layouts. +#[derive(Debug)] +struct MiniBlockDecoder { + rep_decompressor: Option>, + def_decompressor: Option>, + value_decompressor: Arc, + def_meaning: Arc<[DefinitionInterpretation]>, + loaded_chunks: VecDeque, + instructions: VecDeque, + offset_in_current_chunk: u64, + num_rows: u64, + num_buffers: u64, + dictionary: Option>, + has_large_chunk: bool, +} + +/// See [`MiniBlockScheduler`] for more details on the scheduling and decoding +/// process for miniblock encoded data. +impl StructuralPageDecoder for MiniBlockDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + let mut items_desired = num_rows; + let mut need_preamble = false; + let mut skip_in_chunk = self.offset_in_current_chunk; + let mut drain_instructions = Vec::new(); + while items_desired > 0 || need_preamble { + let (instructions, consumed) = self + .instructions + .front() + .unwrap() + .drain_from_instruction(&mut items_desired, &mut need_preamble, &mut skip_in_chunk); + + while self.loaded_chunks.front().unwrap().chunk_idx + != instructions.chunk_instructions.chunk_idx + { + self.loaded_chunks.pop_front(); + } + drain_instructions.push((instructions, self.loaded_chunks.front().unwrap().clone())); + if consumed { + self.instructions.pop_front(); + } + } + // We can throw away need_preamble here because it must be false. If it were true it would mean + // we were still in the middle of loading rows. We do need to latch skip_in_chunk though. + self.offset_in_current_chunk = skip_in_chunk; + + let max_visible_level = max_visible_level(&self.def_meaning); + + Ok(Box::new(DecodeMiniBlockTask { + instructions: drain_instructions, + def_decompressor: self.def_decompressor.clone(), + rep_decompressor: self.rep_decompressor.clone(), + value_decompressor: self.value_decompressor.clone(), + dictionary_data: self.dictionary.clone(), + def_meaning: self.def_meaning.clone(), + num_buffers: self.num_buffers, + max_visible_level, + has_large_chunk: self.has_large_chunk, + })) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +/// How a complex-all-null page's rep/def level buffer is compressed on disk. +/// Captured at scheduler construction so `initialize` can keep RLE levels in run +/// form instead of expanding them. +#[derive(Debug, Clone)] +pub(crate) enum LevelCodec { + /// Raw little-endian u16 levels (no block compression). + Uncompressed, + /// RLE-compressed levels; the validated physical runs select their cached representation. + Rle(Arc), + /// Any other block compression; decoded eagerly into [`LazyLevels::Dense`] + /// (these encodings don't expand, so laziness buys nothing). + Block(Arc), +} + +impl LevelCodec { + fn try_new( + encoding: Option<&CompressiveEncoding>, + decompression_strategy: &dyn DecompressionStrategy, + ) -> Result { + match encoding { + None => Ok(Self::Uncompressed), + Some(encoding) => match encoding.compression.as_ref() { + Some(Compression::Rle(rle)) => Ok(Self::Rle(Arc::new(create_rle_decompressor( + rle, + decompression_strategy, + )?))), + _ => Ok(Self::Block(Arc::from( + decompression_strategy.create_block_decompressor(encoding)?, + ))), + }, + } + } +} + +#[derive(Debug)] +enum RunEnds { + U16(Box<[u16]>), + U32(Box<[u32]>), + U64(Box<[u64]>), +} + +impl RunEnds { + fn width_for(num_values: usize) -> usize { + if u16::try_from(num_values).is_ok() { + std::mem::size_of::() + } else if u32::try_from(num_values).is_ok() { + std::mem::size_of::() + } else { + std::mem::size_of::() + } + } + + fn len(&self) -> usize { + match self { + Self::U16(ends) => ends.len(), + Self::U32(ends) => ends.len(), + Self::U64(ends) => ends.len(), + } + } + + fn get(&self, run: usize) -> usize { + match self { + Self::U16(ends) => ends[run] as usize, + Self::U32(ends) => ends[run] as usize, + Self::U64(ends) => ends[run] as usize, + } + } + + fn partition_point(&self, logical_index: usize) -> usize { + match self { + Self::U16(ends) => ends.partition_point(|&end| end as usize <= logical_index), + Self::U32(ends) => ends.partition_point(|&end| end as usize <= logical_index), + Self::U64(ends) => ends.partition_point(|&end| end as usize <= logical_index), + } + } + + fn deep_size(&self) -> usize { + match self { + Self::U16(ends) => std::mem::size_of_val(ends.as_ref()), + Self::U32(ends) => std::mem::size_of_val(ends.as_ref()), + Self::U64(ends) => std::mem::size_of_val(ends.as_ref()), + } + } +} + +enum RunEndsBuilder { + U16(Vec), + U32(Vec), + U64(Vec), +} + +impl RunEndsBuilder { + fn with_capacity(num_values: usize, capacity: usize) -> Self { + if u16::try_from(num_values).is_ok() { + Self::U16(Vec::with_capacity(capacity)) + } else if u32::try_from(num_values).is_ok() { + Self::U32(Vec::with_capacity(capacity)) + } else { + Self::U64(Vec::with_capacity(capacity)) + } + } + + fn push(&mut self, end: usize) -> Result<()> { + match self { + Self::U16(ends) => ends.push( + u16::try_from(end) + .map_err(|_| Error::internal(format!("Run end {end} does not fit in u16")))?, + ), + Self::U32(ends) => ends.push( + u32::try_from(end) + .map_err(|_| Error::internal(format!("Run end {end} does not fit in u32")))?, + ), + Self::U64(ends) => ends.push(end as u64), + } + Ok(()) + } + + fn set_last(&mut self, end: usize) -> Result<()> { + match self { + Self::U16(ends) => { + let last = ends.last_mut().ok_or_else(|| { + Error::internal("Cannot extend an empty coalesced run buffer") + })?; + *last = u16::try_from(end) + .map_err(|_| Error::internal(format!("Run end {end} does not fit in u16")))?; + } + Self::U32(ends) => { + let last = ends.last_mut().ok_or_else(|| { + Error::internal("Cannot extend an empty coalesced run buffer") + })?; + *last = u32::try_from(end) + .map_err(|_| Error::internal(format!("Run end {end} does not fit in u32")))?; + } + Self::U64(ends) => { + let last = ends.last_mut().ok_or_else(|| { + Error::internal("Cannot extend an empty coalesced run buffer") + })?; + *last = end as u64; + } + } + Ok(()) + } + + fn finish(self) -> RunEnds { + match self { + Self::U16(ends) => RunEnds::U16(ends.into_boxed_slice()), + Self::U32(ends) => RunEnds::U32(ends.into_boxed_slice()), + Self::U64(ends) => RunEnds::U64(ends.into_boxed_slice()), + } + } +} + +#[derive(Debug)] +enum RunStorage { + Physical(RleRuns), + Coalesced { values: Box<[u16]>, ends: RunEnds }, +} + +impl RunStorage { + fn len(&self) -> usize { + match self { + Self::Physical(runs) => runs.num_values(), + Self::Coalesced { ends, .. } => ends.get(ends.len() - 1), + } + } + + fn num_runs(&self) -> usize { + match self { + Self::Physical(runs) => runs.num_runs(), + Self::Coalesced { values, .. } => values.len(), + } + } + + fn value(&self, run: usize) -> u16 { + match self { + Self::Physical(runs) => runs.value(run), + Self::Coalesced { values, .. } => values[run], + } + } + + fn first_value_above(&self, max: u16) -> Option<(usize, u16)> { + (0..self.num_runs()).find_map(|run| { + let value = self.value(run); + (value > max).then_some((run, value)) + }) + } + + fn seek(&self, position: &mut RunPosition, logical_index: usize) { + if logical_index >= self.len() { + *position = RunPosition { + run: self.num_runs(), + start: self.len(), + end: self.len(), + }; + return; + } + + match self { + Self::Physical(runs) => { + if position.run >= runs.num_runs() + || position.end == 0 + || logical_index < position.start + { + *position = RunPosition { + run: 0, + start: 0, + end: runs.length(0), + }; + } + while position.end <= logical_index { + self.advance(position); + } + } + Self::Coalesced { ends, .. } => { + if logical_index < position.start || logical_index >= position.end { + let run = ends.partition_point(logical_index); + *position = RunPosition { + run, + start: if run == 0 { 0 } else { ends.get(run - 1) }, + end: ends.get(run), + }; + } + } + } + } + + fn advance(&self, position: &mut RunPosition) { + let next_run = position.run + 1; + if next_run >= self.num_runs() { + *position = RunPosition { + run: self.num_runs(), + start: self.len(), + end: self.len(), + }; + return; + } + + let start = position.end; + position.run = next_run; + position.start = start; + position.end = match self { + Self::Physical(runs) => start + runs.length(next_run), + Self::Coalesced { ends, .. } => ends.get(next_run), + }; + } + + fn deep_size(&self) -> usize { + match self { + Self::Physical(runs) => runs.deep_size(), + Self::Coalesced { values, ends } => { + std::mem::size_of_val(values.as_ref()) + ends.deep_size() + } + } + } +} + +/// Rep/def levels for a complex-all-null page. +/// +/// RLE pages retain the smallest of their validated physical runs, coalesced +/// runs, and dense values. The decoder materializes only the per-drain slices +/// it touches. +#[derive(Debug, Clone)] +enum LazyLevels { + Dense(ScalarBuffer), + Runs(Arc), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LevelPlan { + Physical, + Coalesced, + Dense, +} + +#[derive(Debug, Default, Clone, Copy)] +struct RunPosition { + run: usize, + start: usize, + end: usize, +} + +/// Monotonic forward cursor into a [`LazyLevels`] sequence. +/// +/// Drains seek to strictly increasing rows, so each [`LazyLevels::seek_row_start`] +/// resumes from the last position instead of rescanning — every run is visited at +/// most once per page while locating and counting level ranges. +#[derive(Debug, Default, Clone, Copy)] +struct LevelCursor { + /// Logical level index where the current row begins. + level: usize, + /// Row index at `level` (the number of `max_rep` occurrences before it). + row: u64, + /// Run containing `level`. Unused for [`LazyLevels::Dense`]. + run: RunPosition, +} + +impl LazyLevels { + fn from_rle_runs(runs: RleRuns) -> Result { + let plan = Self::select_plan(&runs); + match plan { + LevelPlan::Physical => Ok(Self::Runs(Arc::new(RunStorage::Physical( + runs.into_owned(), + )))), + LevelPlan::Coalesced => Self::build_coalesced(runs), + LevelPlan::Dense => Self::build_dense(runs), + } + } + + /// Minimize retained payload bytes first, then expected traversal work. + /// If both are equal, keep the physical runs and avoid another allocation. + fn select_plan(runs: &RleRuns) -> LevelPlan { + if runs.num_values() == 0 { + return LevelPlan::Dense; + } + + let run_storage_size = std::mem::size_of::() as u128; + let physical_size = run_storage_size + runs.owned_size() as u128; + let coalesced_size = run_storage_size + + (runs.coalesced_runs() as u128) + * (std::mem::size_of::() + RunEnds::width_for(runs.num_values())) as u128; + let dense_size = (runs.num_values() as u128) * std::mem::size_of::() as u128; + [ + (physical_size, runs.num_runs(), 0usize, LevelPlan::Physical), + ( + coalesced_size, + runs.coalesced_runs(), + 1usize, + LevelPlan::Coalesced, + ), + (dense_size, runs.num_values(), 2usize, LevelPlan::Dense), + ] + .into_iter() + .min_by_key(|(size, traversal, priority, _)| (*size, *traversal, *priority)) + .map(|(_, _, _, plan)| plan) + .unwrap_or(LevelPlan::Dense) + } + + fn build_coalesced(runs: RleRuns) -> Result { + let mut values = Vec::with_capacity(runs.coalesced_runs()); + let mut ends = RunEndsBuilder::with_capacity(runs.num_values(), runs.coalesced_runs()); + let mut logical_end = 0usize; + for (value, length) in runs.iter() { + logical_end = logical_end + .checked_add(length) + .ok_or_else(|| Error::internal("Validated RLE run length sum overflowed usize"))?; + if values.last().copied() == Some(value) { + ends.set_last(logical_end)?; + } else { + values.push(value); + ends.push(logical_end)?; + } + } + Ok(Self::Runs(Arc::new(RunStorage::Coalesced { + values: values.into_boxed_slice(), + ends: ends.finish(), + }))) + } + + fn build_dense(runs: RleRuns) -> Result { + let mut values = Vec::new(); + values.try_reserve_exact(runs.num_values()).map_err(|_| { + Error::internal(format!( + "Cannot allocate {} dense repetition/definition levels", + runs.num_values() + )) + })?; + for (value, length) in runs.iter() { + values.resize(values.len() + length, value); + } + Ok(Self::Dense(ScalarBuffer::from(values))) + } + + fn len(&self) -> usize { + match self { + Self::Dense(buf) => buf.len(), + Self::Runs(runs) => runs.len(), + } + } + + fn validate_max(&self, level_type: &str, max: u16) -> Result<()> { + let invalid = match self { + Self::Dense(levels) => levels + .iter() + .enumerate() + .find_map(|(index, &value)| (value > max).then_some(("index", index, value))), + Self::Runs(runs) => runs + .first_value_above(max) + .map(|(run, value)| ("run", run, value)), + }; + if let Some((position_type, position, value)) = invalid { + return Err(Error::invalid_input_source( + format!( + "Invalid {level_type} level {value} at {position_type} {position}: maximum is {max}" + ) + .into(), + )); + } + Ok(()) + } + + /// Advance `cursor` to the start of row `target_row`, returning that row's + /// starting level index. + /// + /// Rows begin at `max_rep` positions, so this finds the `target_row`-th one. + /// `target_row` must be `>= cursor.row`: the cursor only moves forward, which + /// is what keeps a full page decode O(runs) rather than O(rows). + fn seek_row_start( + &self, + cursor: &mut LevelCursor, + target_row: u64, + max_rep: u16, + ) -> Result { + let mut need = target_row.checked_sub(cursor.row).ok_or_else(|| { + Error::internal(format!( + "Complex all-null row ranges are not sorted: target row {target_row} follows {}", + cursor.row + )) + })?; + if need == 0 { + return Ok(cursor.level); + } + match self { + Self::Dense(buf) => { + let mut level = cursor.level; + while need > 0 { + if level >= buf.len() { + return Err(Error::internal( + "Invalid complex all-null layout: repetition buffer too short", + )); + } + if buf[level] != max_rep { + return Err(Error::internal( + "Invalid complex all-null layout: row did not start at max repetition level", + )); + } + level += 1; + while level < buf.len() && buf[level] != max_rep { + level += 1; + } + need -= 1; + } + cursor.level = level; + cursor.row = target_row; + Ok(level) + } + Self::Runs(runs) => { + let mut level = cursor.level; + let mut run = cursor.run; + runs.seek(&mut run, level); + while need > 0 { + if run.run >= runs.num_runs() { + return Err(Error::internal( + "Invalid complex all-null layout: repetition buffer too short", + )); + } + if runs.value(run.run) != max_rep { + return Err(Error::internal( + "Invalid complex all-null layout: row did not start at max repetition level", + )); + } + let avail = (run.end - level) as u64; + if need < avail { + // Target lands inside this max-rep run. + level += need as usize; + need = 0; + } else { + // Consume every row start in this run, then skip the + // trailing non-max-rep runs to reach the next row start. + need -= avail; + runs.advance(&mut run); + while run.run < runs.num_runs() && runs.value(run.run) != max_rep { + runs.advance(&mut run); + } + level = if run.run < runs.num_runs() { + run.start + } else { + self.len() + }; + } + } + cursor.level = level; + cursor.row = target_row; + cursor.run = run; + Ok(level) + } + } + } + + /// Count of levels in `range` that are `<= max`, resuming from `*run_cursor` + /// and leaving it on the last run that overlaps `range`. + /// + /// Successive calls must pass ascending, non-overlapping ranges (`range.start + /// >=` the previous `range.end`) so runs are swept at most once per page. + fn count_le_cursor( + &self, + run_cursor: &mut RunPosition, + range: Range, + max: u16, + ) -> (u64, RunPosition) { + if range.is_empty() { + return (0, *run_cursor); + } + match self { + Self::Dense(buf) => ( + buf[range].iter().filter(|&&d| d <= max).count() as u64, + RunPosition::default(), + ), + Self::Runs(runs) => { + // Advance to the first run overlapping the range. + runs.seek(run_cursor, range.start); + let start = *run_cursor; + let mut count = 0u64; + let mut current = *run_cursor; + while current.run < runs.num_runs() && current.start < range.end { + if runs.value(current.run) <= max { + let lo = current.start.max(range.start); + let hi = current.end.min(range.end); + count += (hi - lo) as u64; + } + if current.end >= range.end { + break; + } + runs.advance(&mut current); + } + // Resume the next (ascending) range from the last overlapping run; + // `current` remains valid because `range` is non-empty. + *run_cursor = current; + (count, start) + } + } + } + + fn extend_into(&self, range: Range, run: RunPosition, out: &mut Vec) { + if range.is_empty() { + return; + } + match self { + Self::Dense(buf) => out.extend_from_slice(&buf[range]), + Self::Runs(runs) => { + let mut current = run; + runs.seek(&mut current, range.start); + while current.run < runs.num_runs() && current.start < range.end { + let lo = current.start.max(range.start); + let hi = current.end.min(range.end); + if hi > lo { + out.resize(out.len() + (hi - lo), runs.value(current.run)); + } + runs.advance(&mut current); + } + } + } + } + + #[cfg(test)] + fn deep_size(&self) -> usize { + self.deep_size_of_children(&mut Context::new()) + } +} + +impl DeepSizeOf for LazyLevels { + fn deep_size_of_children(&self, ctx: &mut Context) -> usize { + match self { + Self::Dense(buf) => buf.deep_size_of_children(ctx), + Self::Runs(runs) => { + let pointer = Arc::as_ptr(runs) as *const () as usize; + if ctx.mark_seen(pointer) { + std::mem::size_of_val(runs.as_ref()) + runs.deep_size() + } else { + 0 + } + } + } + } +} + +fn validate_complex_all_null_levels( + rep: &Option, + def: &Option, + max_rep: u16, + max_def: u16, +) -> Result<()> { + if let Some(rep) = rep { + rep.validate_max("repetition", max_rep)?; + } + if let Some(def) = def { + def.validate_max("definition", max_def)?; + } + if let (Some(rep), Some(def)) = (rep, def) + && rep.len() != def.len() + { + return Err(Error::invalid_input_source( + format!( + "Mismatched complex all-null level counts: repetition has {}, definition has {}", + rep.len(), + def.len() + ) + .into(), + )); + } + Ok(()) +} + +fn expected_level_bytes(num_values: u64, level_type: &str) -> Result { + usize::try_from(num_values) + .ok() + .and_then(|num_values| num_values.checked_mul(std::mem::size_of::())) + .ok_or_else(|| { + Error::invalid_input_source( + format!("{level_type} level count {num_values} does not fit in memory").into(), + ) + }) +} + +fn dense_levels_from_block( + decompressed: DataBlock, + num_values: u64, + level_type: &str, +) -> Result { + let DataBlock::FixedWidth(block) = decompressed else { + return Err(Error::invalid_input_source( + format!("Expected fixed-width data block for {level_type} levels").into(), + )); + }; + if block.num_values != num_values { + return Err(Error::invalid_input_source( + format!( + "Unexpected {level_type} level count after decompression: expected {num_values}, got {}", + block.num_values + ) + .into(), + )); + } + if block.bits_per_value != 16 { + return Err(Error::invalid_input_source( + format!( + "Unexpected {level_type} level bit width after decompression: expected 16, got {}", + block.bits_per_value + ) + .into(), + )); + } + let expected_bytes = expected_level_bytes(num_values, level_type)?; + if block.data.len() != expected_bytes { + return Err(Error::invalid_input_source( + format!( + "Unexpected decompressed {level_type} level size: expected {expected_bytes} bytes for {num_values} values, got {}", + block.data.len() + ) + .into(), + )); + } + Ok(LazyLevels::Dense(block.data.borrow_to_typed_slice::())) +} + +#[derive(Debug)] +struct CachedComplexAllNullState { + rep: Option, + def: Option, +} + +impl DeepSizeOf for CachedComplexAllNullState { + fn deep_size_of_children(&self, ctx: &mut Context) -> usize { + self.rep.deep_size_of_children(ctx) + self.def.deep_size_of_children(ctx) + } +} + +impl CachedPageData for CachedComplexAllNullState { + fn as_arc_any(self: Arc) -> Arc { + self + } +} + +/// A scheduler for all-null data that has repetition and definition levels +/// +/// We still need to do some I/O in this case because we need to figure out what kind of null we +/// are dealing with (null list, null struct, what level null struct, etc.) +/// +/// TODO: Right now we just load the entire rep/def at initialization time and cache it. This is a touch +/// RAM aggressive and maybe we want something more lazy in the future. On the other hand, it's simple +/// and fast so...maybe not :) +#[derive(Debug)] +pub struct ComplexAllNullScheduler { + // Set from protobuf + buffer_offsets_and_sizes: Arc<[(u64, u64)]>, + def_meaning: Arc<[DefinitionInterpretation]>, + repdef: Option>, + max_rep: u16, + max_def: u16, + max_visible_level: u16, + rep_codec: LevelCodec, + def_codec: LevelCodec, + num_rep_values: u64, + num_def_values: u64, +} + +impl ComplexAllNullScheduler { + pub(crate) fn new( + buffer_offsets_and_sizes: Arc<[(u64, u64)]>, + def_meaning: Arc<[DefinitionInterpretation]>, + rep_codec: LevelCodec, + def_codec: LevelCodec, + num_rep_values: u64, + num_def_values: u64, + ) -> Self { + let max_rep = def_meaning.iter().filter(|l| l.is_list()).count() as u16; + let max_def = def_meaning + .iter() + .map(|meaning| meaning.num_def_levels()) + .sum::(); + let max_visible_level = max_visible_level(&def_meaning); + Self { + buffer_offsets_and_sizes, + def_meaning, + repdef: None, + max_rep, + max_def, + max_visible_level, + rep_codec, + def_codec, + num_rep_values, + num_def_values, + } + } +} + +impl StructuralPageScheduler for ComplexAllNullScheduler { + fn init_layout(&self) -> Result { + // Fully load the rep & def buffers, as needed + let (rep_pos, rep_size) = self.buffer_offsets_and_sizes[0]; + let (def_pos, def_size) = self.buffer_offsets_and_sizes[1]; + + Ok(PageInitialization::new([ + (rep_size > 0).then_some(rep_pos..rep_pos + rep_size), + (def_size > 0).then_some(def_pos..def_pos + def_size), + ])) + } + + fn init_from_buffers<'a>( + &'a mut self, + buffers: PageInitializationBuffers, + _io: &Arc, + ) -> BoxFuture<'a, Result>> { + let rep_codec = self.rep_codec.clone(); + let def_codec = self.def_codec.clone(); + let num_rep_values = self.num_rep_values; + let num_def_values = self.num_def_values; + let max_rep = self.max_rep; + let max_def = self.max_def; + + async move { + let [rep_bytes, def_bytes] = buffers.try_into_array("complex all-null")?; + + // RLE levels select the smallest validated cache representation; + // everything else expands eagerly to `LazyLevels::Dense`. + let build_levels = |compressed_bytes: Bytes, + codec: &LevelCodec, + num_values: u64, + level_type: &str| + -> Result { + match codec { + LevelCodec::Uncompressed => { + if num_values == 0 { + if !compressed_bytes + .len() + .is_multiple_of(std::mem::size_of::()) + { + return Err(Error::invalid_input_source( + format!( + "Unexpected uncompressed {level_type} level size: {} bytes is not divisible by {}", + compressed_bytes.len(), + std::mem::size_of::() + ) + .into(), + )); + } + } else { + let expected_bytes = expected_level_bytes(num_values, level_type)?; + if compressed_bytes.len() != expected_bytes { + return Err(Error::invalid_input_source( + format!( + "Unexpected uncompressed {level_type} level size: expected {expected_bytes} bytes for {num_values} values, got {}", + compressed_bytes.len() + ) + .into(), + )); + } + } + let buffer = LanceBuffer::from_bytes(compressed_bytes, 2); + Ok(LazyLevels::Dense(buffer.borrow_to_typed_slice::())) + } + LevelCodec::Rle(decompressor) => { + let frame = LanceBuffer::from_bytes(compressed_bytes, 1); + let runs = decompressor.decode_u16_runs(frame, num_values)?; + LazyLevels::from_rle_runs(runs) + } + LevelCodec::Block(decompressor) => { + let frame = LanceBuffer::from_bytes(compressed_bytes, 1); + let decompressed = decompressor.decompress(Some(frame), num_values)?; + dense_levels_from_block(decompressed, num_values, level_type) + } + } + }; + + let rep = if let Some(rep) = rep_bytes { + Some(build_levels(rep, &rep_codec, num_rep_values, "repetition")?) + } else { + None + }; + + let def = if let Some(def) = def_bytes { + Some(build_levels(def, &def_codec, num_def_values, "definition")?) + } else { + None + }; + + validate_complex_all_null_levels(&rep, &def, max_rep, max_def)?; + let repdef = Arc::new(CachedComplexAllNullState { rep, def }); + + Ok(repdef as Arc) + } + .boxed() + } + + fn try_load(&mut self, data: &Arc) -> Result<()> { + self.repdef = Some(data.clone().as_arc_any().downcast().map_err(|_| { + Error::invalid_input_source( + "Cached complex all-null page data has an unexpected type".into(), + ) + })?); + Ok(()) + } + + fn schedule_ranges( + &self, + ranges: &[Range], + _io: &Arc, + ) -> Result> { + let ranges = VecDeque::from_iter(ranges.iter().cloned()); + let num_rows = ranges.iter().map(|r| r.end - r.start).sum::(); + let decoder = Box::new(ComplexAllNullPageDecoder { + ranges, + rep: self.repdef.as_ref().unwrap().rep.clone(), + def: self.repdef.as_ref().unwrap().def.clone(), + num_rows, + def_meaning: self.def_meaning.clone(), + max_rep: self.max_rep, + max_visible_level: self.max_visible_level, + rep_cursor: LevelCursor::default(), + def_run_cursor: RunPosition::default(), + }) as Box; + let page_load_task = PageLoadTask { + decoder_fut: std::future::ready(Ok(decoder)).boxed(), + num_rows, + }; + Ok(vec![page_load_task]) + } +} + +#[derive(Debug)] +pub struct ComplexAllNullPageDecoder { + ranges: VecDeque>, + rep: Option, + def: Option, + num_rows: u64, + def_meaning: Arc<[DefinitionInterpretation]>, + max_rep: u16, + max_visible_level: u16, + /// Monotonic cursor into `rep` tracking the current row's level start. + rep_cursor: LevelCursor, + /// Monotonic run cursor into `def` for `count_le_cursor`. + def_run_cursor: RunPosition, +} + +impl ComplexAllNullPageDecoder { + fn drain_ranges(&mut self, num_rows: u64) -> Vec> { + let mut rows_desired = num_rows; + let mut ranges = Vec::with_capacity(self.ranges.len()); + while rows_desired > 0 { + let front = self.ranges.front_mut().unwrap(); + let avail = front.end - front.start; + if avail > rows_desired { + ranges.push(front.start..front.start + rows_desired); + front.start += rows_desired; + rows_desired = 0; + } else { + ranges.push(self.ranges.pop_front().unwrap()); + rows_desired -= avail; + } + } + ranges + } + + /// Level index at which row `target_row` starts, advancing the monotonic + /// repetition cursor. Callers must request non-decreasing `target_row`. + fn seek_row_start(&mut self, target_row: u64) -> Result { + match &self.rep { + Some(rep) => rep.seek_row_start(&mut self.rep_cursor, target_row, self.max_rep), + None => { + // Without repetition every level is its own row. + self.rep_cursor.row = target_row; + self.rep_cursor.level = target_row as usize; + Ok(target_row as usize) + } + } + } + + /// Number of visible items in the level range `levels` (definition levels + /// `<= max_visible_level`), advancing the monotonic definition cursor. + fn count_visible(&mut self, levels: Range) -> Result<(u64, RunPosition)> { + match &self.def { + Some(def) => { + if levels.end > def.len() { + return Err(Error::internal( + "Invalid complex all-null layout: definition buffer too short", + )); + } + Ok(def.count_le_cursor(&mut self.def_run_cursor, levels, self.max_visible_level)) + } + None => Ok(((levels.end - levels.start) as u64, RunPosition::default())), + } + } +} + +impl StructuralPageDecoder for ComplexAllNullPageDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + let drained_ranges = self.drain_ranges(num_rows); + let mut level_slices: Vec = Vec::with_capacity(drained_ranges.len()); + let mut visible_items_total = 0; + + // Each row range is one contiguous level slice `[start_row_level, + // end_row_level)`, so we seek both boundaries and count its visibility at + // once rather than per row. The cursors only move forward, so locating and + // counting all requested ranges visits each intervening run at most once. + for range in drained_ranges { + let level_start = self.seek_row_start(range.start)?; + let rep_run = self.rep_cursor.run; + let level_end = self.seek_row_start(range.end)?; + let (visible_items, def_run) = self.count_visible(level_start..level_end)?; + visible_items_total += visible_items; + if let Some(last) = level_slices.last_mut() + && last.range.end == level_start + { + last.range.end = level_end; + } else { + level_slices.push(LevelSlice { + range: level_start..level_end, + rep_run, + def_run, + }); + } + } + + Ok(Box::new(DecodeComplexAllNullTask { + level_slices, + visible_items_total, + rep: self.rep.clone(), + def: self.def.clone(), + def_meaning: self.def_meaning.clone(), + max_visible_level: self.max_visible_level, + })) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +/// We use `level_slices` to slice into `rep` and `def` and create rep/def buffers +/// for the null data. +#[derive(Debug, Clone)] +struct LevelSlice { + range: Range, + rep_run: RunPosition, + def_run: RunPosition, +} + +#[derive(Clone, Copy)] +enum LevelKind { + Repetition, + Definition, +} + +impl LevelSlice { + fn run(&self, kind: LevelKind) -> RunPosition { + match kind { + LevelKind::Repetition => self.rep_run, + LevelKind::Definition => self.def_run, + } + } +} + +#[derive(Debug)] +pub struct DecodeComplexAllNullTask { + level_slices: Vec, + visible_items_total: u64, + rep: Option, + def: Option, + def_meaning: Arc<[DefinitionInterpretation]>, + max_visible_level: u16, +} + +impl DecodeComplexAllNullTask { + fn decode_level(&self, levels: &Option, kind: LevelKind) -> Option> { + levels.as_ref().map(|levels| { + let num_levels = self + .level_slices + .iter() + .map(|slice| slice.range.end - slice.range.start) + .sum(); + let mut referenced_levels = Vec::with_capacity(num_levels); + for slice in &self.level_slices { + levels.extend_into(slice.range.clone(), slice.run(kind), &mut referenced_levels); + } + referenced_levels + }) + } +} + +impl DecodePageTask for DecodeComplexAllNullTask { + fn decode(self: Box) -> Result { + let rep = self.decode_level(&self.rep, LevelKind::Repetition); + let def = self.decode_level(&self.def, LevelKind::Definition); + + // If there are definition levels there may be empty / null lists which are not visible + // in the items array. We need to account for that here to figure out how many values + // should be in the items array. + let num_values = if let Some(def) = &def { + def.iter().filter(|&d| *d <= self.max_visible_level).count() as u64 + } else { + self.visible_items_total + }; + + let data = DataBlock::AllNull(AllNullDataBlock { num_values }); + let unraveler = RepDefUnraveler::new(rep, def, self.def_meaning, num_values); + Ok(DecodedPage { + data, + repdef: unraveler, + }) + } +} + +/// A scheduler for simple all-null data +/// +/// "simple" all-null data is data that is all null and only has a single level of definition and +/// no repetition. We don't need to read any data at all in this case. +#[derive(Debug, Default)] +pub struct SimpleAllNullScheduler {} + +impl StructuralPageScheduler for SimpleAllNullScheduler { + fn needs_initialization(&self) -> bool { + false + } + + fn init_layout(&self) -> Result { + Ok(PageInitialization::new([])) + } + + fn init_from_buffers<'a>( + &'a mut self, + _buffers: PageInitializationBuffers, + _io: &Arc, + ) -> BoxFuture<'a, Result>> { + std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc)).boxed() + } + + fn try_load(&mut self, _cache: &Arc) -> Result<()> { + Ok(()) + } + + fn schedule_ranges( + &self, + ranges: &[Range], + _io: &Arc, + ) -> Result> { + let num_rows = ranges.iter().map(|r| r.end - r.start).sum::(); + let decoder = + Box::new(SimpleAllNullPageDecoder { num_rows }) as Box; + let page_load_task = PageLoadTask { + decoder_fut: std::future::ready(Ok(decoder)).boxed(), + num_rows, + }; + Ok(vec![page_load_task]) + } +} + +/// A page decode task for all-null data without any +/// repetition and only a single level of definition +#[derive(Debug)] +struct SimpleAllNullDecodePageTask { + num_values: u64, +} +impl DecodePageTask for SimpleAllNullDecodePageTask { + fn decode(self: Box) -> Result { + let unraveler = RepDefUnraveler::new( + None, + Some(vec![1; self.num_values as usize]), + Arc::new([DefinitionInterpretation::NullableItem]), + self.num_values, + ); + Ok(DecodedPage { + data: DataBlock::AllNull(AllNullDataBlock { + num_values: self.num_values, + }), + repdef: unraveler, + }) + } +} + +#[derive(Debug)] +pub struct SimpleAllNullPageDecoder { + num_rows: u64, +} + +impl StructuralPageDecoder for SimpleAllNullPageDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + Ok(Box::new(SimpleAllNullDecodePageTask { + num_values: num_rows, + })) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +#[derive(Debug, Clone)] +struct MiniBlockSchedulerDictionary { + // These come from the protobuf + dictionary_decompressor: Arc, + dictionary_buf_position_and_size: (u64, u64), + dictionary_data_alignment: u64, + num_dictionary_items: u64, +} + +/// State that is loaded once and cached for future lookups +#[derive(Debug)] +struct MiniBlockCacheableState { + /// Compact per-chunk index (byte ranges + row/item mapping) for the page + chunk_index: MiniBlockChunkIndex, + /// The dictionary for the page, if any + dictionary: Option>, +} + +impl DeepSizeOf for MiniBlockCacheableState { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + self.chunk_index.deep_size_of_children(context) + + self + .dictionary + .as_ref() + .map(|dict| dict.data_size() as usize) + .unwrap_or(0) + } +} + +impl CachedPageData for MiniBlockCacheableState { + fn as_arc_any(self: Arc) -> Arc { + self + } +} + +/// A scheduler for a page that has been encoded with the mini-block layout +/// +/// Scheduling mini-block encoded data is simple in concept and somewhat complex +/// in practice. +/// +/// First, during initialization, we load the chunk metadata, the repetition index, +/// and the dictionary (these last two may not be present) +/// +/// Then, during scheduling, we use the user's requested row ranges and the repetition +/// index to determine which chunks we need and which rows we need from those chunks. +/// +/// For example, if the repetition index is: [50, 3], [50, 0], [10, 0] and the range +/// from the user is 40..60 then we need to: +/// +/// - Read the first chunk and skip the first 40 rows, then read 10 full rows, and +/// then read 3 items for the 11th row of our range. +/// - Read the second chunk and read the remaining items in our 11th row and then read +/// the remaining 9 full rows. +/// +/// Then, if we are going to decode that in batches of 5, we need to make decode tasks. +/// The first two decode tasks will just need the first chunk. The third decode task will +/// need the first chunk (for the trailer which has the 11th row in our range) and the second +/// chunk. The final decode task will just need the second chunk. +/// +/// The above prose descriptions are what are represented by `ChunkInstructions` and +/// `ChunkDrainInstructions`. +#[derive(Debug)] +pub struct MiniBlockScheduler { + // These come from the protobuf + buffer_offsets_and_sizes: Vec<(u64, u64)>, + priority: u64, + items_in_page: u64, + repetition_index_depth: u16, + num_buffers: u64, + rep_decompressor: Option>, + def_decompressor: Option>, + value_decompressor: Arc, + def_meaning: Arc<[DefinitionInterpretation]>, + dictionary: Option, + // This is set after initialization + page_meta: Option>, + has_large_chunk: bool, +} + +impl MiniBlockScheduler { + fn try_new( + buffer_offsets_and_sizes: &[(u64, u64)], + priority: u64, + items_in_page: u64, + layout: &pb21::MiniBlockLayout, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + let rep_decompressor = layout + .rep_compression + .as_ref() + .map(|rep_compression| { + decompressors + .create_block_decompressor(rep_compression) + .map(Arc::from) + }) + .transpose()?; + let def_decompressor = layout + .def_compression + .as_ref() + .map(|def_compression| { + decompressors + .create_block_decompressor(def_compression) + .map(Arc::from) + }) + .transpose()?; + let def_meaning = layout + .layers + .iter() + .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l)) + .collect::>(); + let value_decompressor = decompressors.create_miniblock_decompressor( + layout.value_compression.as_ref().unwrap(), + decompressors, + )?; + + let dictionary = if let Some(dictionary_encoding) = layout.dictionary.as_ref() { + let num_dictionary_items = layout.num_dictionary_items; + let dictionary_decompressor = decompressors + .create_block_decompressor(dictionary_encoding)? + .into(); + let dictionary_data_alignment = match dictionary_encoding.compression.as_ref().unwrap() + { + Compression::Variable(_) => 4, + Compression::Flat(_) => 16, + Compression::General(_) => 1, + Compression::InlineBitpacking(_) | Compression::OutOfLineBitpacking(_) => { + crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT + } + _ => { + return Err(Error::invalid_input_source( + format!( + "Unsupported mini-block dictionary encoding: {:?}", + dictionary_encoding.compression.as_ref().unwrap() + ) + .into(), + )); + } + }; + Some(MiniBlockSchedulerDictionary { + dictionary_decompressor, + dictionary_buf_position_and_size: buffer_offsets_and_sizes[2], + dictionary_data_alignment, + num_dictionary_items, + }) + } else { + None + }; + + Ok(Self { + buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(), + rep_decompressor, + def_decompressor, + value_decompressor: value_decompressor.into(), + repetition_index_depth: layout.repetition_index_depth as u16, + num_buffers: layout.num_buffers, + priority, + items_in_page, + dictionary, + def_meaning: def_meaning.into(), + page_meta: None, + has_large_chunk: layout.has_large_chunk, + }) + } + + fn lookup_chunks(&self, chunk_indices: &[usize]) -> Vec { + let chunk_index = &self.page_meta.as_ref().unwrap().chunk_index; + chunk_indices + .iter() + .map(|&chunk_idx| LoadedChunk { + byte_range: chunk_index.byte_range(chunk_idx), + items_in_chunk: chunk_index.items_in_chunk(chunk_idx), + chunk_idx, + data: LanceBuffer::empty(), + }) + .collect() + } +} + +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +enum PreambleAction { + Take, + Skip, + Absent, +} + +// When we schedule a chunk we use the repetition index (or, if none exists, just the # of items +// in each chunk) to map a user requested range into a set of ChunkInstruction objects which tell +// us how exactly to read from the chunk. +// +// Examples: +// +// | Chunk 0 | Chunk 1 | Chunk 2 | Chunk 3 | +// | xxxxyyyyzzz | zzzzzzzzz | zzzzzzzzz | aaabbcc | +// +// Full read (0..6) +// +// Chunk 0: (several rows, ends with trailer) +// preamble: absent +// rows_to_skip: 0 +// rows_to_take: 3 (x, y, z) +// take_trailer: true +// +// Chunk 1: (all preamble, ends with trailer) +// preamble: take +// rows_to_skip: 0 +// rows_to_take: 0 +// take_trailer: true +// +// Chunk 2: (all preamble, no trailer) +// preamble: take +// rows_to_skip: 0 +// rows_to_take: 0 +// take_trailer: false +// +// Chunk 3: (several rows, no trailer or preamble) +// preamble: absent +// rows_to_skip: 0 +// rows_to_take: 3 (a, b, c) +// take_trailer: false +#[derive(Clone, Debug, PartialEq, Eq)] +struct ChunkInstructions { + // The index of the chunk to read + chunk_idx: usize, + // A "preamble" is when a chunk begins with a continuation of the previous chunk's list. If there + // is no repetition index there is never a preamble. + // + // It's possible for a chunk to be entirely premable. For example, if there is a really large list + // that spans several chunks. + preamble: PreambleAction, + // How many complete rows (not including the preamble or trailer) to skip + // + // If this is non-zero then premable must not be Take + rows_to_skip: u64, + // How many rows to take. If a row splits across chunks then we will count the row in the first + // chunk that contains the row. + rows_to_take: u64, + // A "trailer" is when a chunk ends with a partial list. If there is no repetition index there is + // never a trailer. + // + // A chunk that is all preamble may or may not have a trailer. + // + // If this is true then we want to include the trailer + take_trailer: bool, +} + +// First, we schedule a bunch of [`ChunkInstructions`] based on the users ranges. Then we +// start decoding them, based on a batch size, which might not align with what we scheduled. +// +// This results in `ChunkDrainInstructions` which targets a contiguous slice of a `ChunkInstructions` +// +// So if `ChunkInstructions` is "skip preamble, skip 10, take 50, take trailer" and we are decoding in +// batches of size 10 we might have a `ChunkDrainInstructions` that targets that chunk and has its own +// skip of 17 and take of 10. This would mean we decode the chunk, skip the preamble and 27 rows, and +// then take 10 rows. +// +// One very confusing bit is that `rows_to_take` includes the trailer. So if we have two chunks: +// -no preamble, skip 5, take 10, take trailer +// -take preamble, skip 0, take 50, no trailer +// +// and we are draining 20 rows then the drain instructions for the first batch will be: +// - no preamble, skip 0 (from chunk 0), take 11 (from chunk 0) +// - take preamble (from chunk 1), skip 0 (from chunk 1), take 9 (from chunk 1) +#[derive(Debug, PartialEq, Eq)] +struct ChunkDrainInstructions { + chunk_instructions: ChunkInstructions, + rows_to_skip: u64, + rows_to_take: u64, + preamble_action: PreambleAction, +} + +impl ChunkInstructions { + // Given a repetition index and a set of user ranges we need to figure out how to read from the chunks + // + // We assume that `user_ranges` are in sorted order and non-overlapping + // + // The output will be a set of `ChunkInstructions` which tell us how to read from the chunks + fn schedule_instructions( + chunk_index: &MiniBlockChunkIndex, + user_ranges: &[Range], + ) -> Vec { + // Bind the per-page chunk count once; re-deriving it each iteration + // costs a width match plus a length read. + let num_chunks = chunk_index.num_chunks(); + // This is an in-exact capacity guess but pretty good. The actual capacity can be + // smaller if instructions are merged. It can be larger if there are multiple instructions + // per row which can happen with lists. + let mut chunk_instructions = Vec::with_capacity(user_ranges.len()); + + for user_range in user_ranges { + let mut rows_needed = user_range.end - user_range.start; + let mut need_preamble = false; + + // Need to find the first chunk with a first row >= user_range.start. If there are + // multiple chunks with the same first row we need to take the first one. + let mut block_index = chunk_index.find_chunk(user_range.start); + + let mut to_skip = user_range.start - chunk_index.first_row(block_index); + + while rows_needed > 0 || need_preamble { + // Check if we've gone past the last block (should not happen) + if block_index >= num_chunks { + log::warn!( + "schedule_instructions inconsistency: block_index >= num_chunks, exiting early" + ); + break; + } + + let starts_including_trailer = chunk_index.rows_in_chunk(block_index); + let has_preamble = chunk_index.has_preamble(block_index); + let has_trailer = chunk_index.has_trailer(block_index); + let rows_avail = starts_including_trailer.saturating_sub(to_skip); + + // Handle blocks that are entirely preamble (rows_avail = 0) + // These blocks have no rows to take but may have a preamble we need + // We only look for preamble if to_skip == 0 (we're not skipping rows) + if rows_avail == 0 && to_skip == 0 { + // Only process if this chunk has a preamble we need + if has_preamble && need_preamble { + chunk_instructions.push(Self { + chunk_idx: block_index, + preamble: PreambleAction::Take, + rows_to_skip: 0, + rows_to_take: 0, + // We still need to look at has_trailer to distinguish between "all preamble + // and row ends at end of chunk" and "all preamble and row bleeds into next + // chunk". Both cases will have 0 rows available. + take_trailer: has_trailer, + }); + // Only set need_preamble = false if the chunk has at least one row, + // Or we are reaching the last block, + // Otherwise, the chunk is entirely preamble and we need the next chunk's preamble too + if starts_including_trailer > 0 || block_index == num_chunks - 1 { + need_preamble = false; + } + } + // Move to next block + block_index += 1; + continue; + } + + // Edge case: if rows_avail == 0 but to_skip > 0 + // This theoretically shouldn't happen (binary search should avoid it) + // but handle it for safety + if rows_avail == 0 && to_skip > 0 { + // This block doesn't have enough rows to skip, move to next block + // Adjust to_skip by the number of rows in this block + to_skip -= starts_including_trailer; + block_index += 1; + continue; + } + + let rows_to_take = rows_avail.min(rows_needed); + rows_needed -= rows_to_take; + + let mut take_trailer = false; + let preamble = if has_preamble { + if need_preamble { + PreambleAction::Take + } else { + PreambleAction::Skip + } + } else { + PreambleAction::Absent + }; + + // Are we taking the trailer? If so, make sure we mark that we need the preamble + if rows_to_take == rows_avail && has_trailer { + take_trailer = true; + need_preamble = true; + } else { + need_preamble = false; + }; + + chunk_instructions.push(Self { + preamble, + chunk_idx: block_index, + rows_to_skip: to_skip, + rows_to_take, + take_trailer, + }); + + to_skip = 0; + block_index += 1; + } + } + + // If there were multiple ranges we may have multiple instructions for a single chunk. Merge them now if they + // are _adjacent_ (i.e. don't merge "take first row of chunk 0" and "take third row of chunk 0" into "take 2 + // rows of chunk 0 starting at 0") + if user_ranges.len() > 1 { + // Merge adjacent instructions in place. `write` indexes the last + // retained instruction; each following instruction is either folded + // into it (contiguous within the same chunk) or compacted forward. + let mut write = 0; + for read in 1..chunk_instructions.len() { + let merges = { + let last = &chunk_instructions[write]; + let candidate = &chunk_instructions[read]; + last.chunk_idx == candidate.chunk_idx + && last.rows_to_take + last.rows_to_skip == candidate.rows_to_skip + }; + if merges { + let rows_to_take = chunk_instructions[read].rows_to_take; + let take_trailer = chunk_instructions[read].take_trailer; + let last = &mut chunk_instructions[write]; + last.rows_to_take += rows_to_take; + last.take_trailer |= take_trailer; + } else { + write += 1; + if write != read { + chunk_instructions.swap(write, read); + } + } + } + chunk_instructions.truncate(write + 1); + } + chunk_instructions + } + + fn drain_from_instruction( + &self, + rows_desired: &mut u64, + need_preamble: &mut bool, + skip_in_chunk: &mut u64, + ) -> (ChunkDrainInstructions, bool) { + // If we need the premable then we shouldn't be skipping anything + debug_assert!(!*need_preamble || *skip_in_chunk == 0); + let rows_avail = self.rows_to_take - *skip_in_chunk; + let has_preamble = self.preamble != PreambleAction::Absent; + let preamble_action = match (*need_preamble, has_preamble) { + (true, true) => PreambleAction::Take, + (true, false) => panic!("Need preamble but there isn't one"), + (false, true) => PreambleAction::Skip, + (false, false) => PreambleAction::Absent, + }; + + // How many rows are we actually taking in this take step (including the preamble + // and trailer both as individual rows) + let rows_taking = if *rows_desired >= rows_avail { + // We want all the rows. If there is a trailer we are grabbing it and will need + // the preamble of the next chunk + // If there is a trailer and we are taking all the rows then we need the preamble + // of the next chunk. + // + // Also, if this chunk is entirely preamble (rows_avail == 0 && !take_trailer) then we + // need the preamble of the next chunk. + *need_preamble = self.take_trailer; + rows_avail + } else { + // We aren't taking all the rows. Even if there is a trailer we aren't taking + // it so we will not need the preamble + *need_preamble = false; + *rows_desired + }; + let rows_skipped = *skip_in_chunk; + + // Update the state for the next iteration + let consumed_chunk = if *rows_desired >= rows_avail { + *rows_desired -= rows_avail; + *skip_in_chunk = 0; + true + } else { + *skip_in_chunk += *rows_desired; + *rows_desired = 0; + false + }; + + ( + ChunkDrainInstructions { + chunk_instructions: self.clone(), + rows_to_skip: rows_skipped, + rows_to_take: rows_taking, + preamble_action, + }, + consumed_chunk, + ) + } +} + +enum Words { + U16(ScalarBuffer), + U32(ScalarBuffer), +} + +struct WordsIter<'a> { + iter: Box + 'a>, +} + +impl Words { + pub fn len(&self) -> usize { + match self { + Self::U16(b) => b.len(), + Self::U32(b) => b.len(), + } + } + + pub fn iter(&self) -> WordsIter<'_> { + match self { + Self::U16(buf) => WordsIter { + iter: Box::new(buf.iter().map(|&x| x as u32)), + }, + Self::U32(buf) => WordsIter { + iter: Box::new(buf.iter().copied()), + }, + } + } + + pub fn from_bytes(bytes: Bytes, has_large_chunk: bool) -> Result { + let bytes_per_value = if has_large_chunk { 4 } else { 2 }; + assert_eq!(bytes.len() % bytes_per_value, 0); + let buffer = LanceBuffer::from_bytes(bytes, bytes_per_value as u64); + if has_large_chunk { + Ok(Self::U32(buffer.borrow_to_typed_slice::())) + } else { + Ok(Self::U16(buffer.borrow_to_typed_slice::())) + } + } +} + +impl<'a> Iterator for WordsIter<'a> { + type Item = u32; + + fn next(&mut self) -> Option { + self.iter.next() + } +} + +/// Per-chunk leaf value-count analysis derived from the metadata words. +/// +/// `values_per_chunk` is the count shared by every non-last chunk (meaningful +/// when `uniform`), and `last_chunk_values` is the final chunk's count. +struct FlatValueCounts { + logs: Vec, + uniform: bool, + values_per_chunk: u64, + last_chunk_values: u64, +} + +fn analyze_value_counts(words: &Words, items_in_page: u64) -> Result { + let num_chunks = words.len(); + let logs = words.iter().map(|w| (w & 0x0F) as u8).collect::>(); + let mut counted = 0u64; + for (chunk_index, &log) in logs.iter().take(num_chunks.saturating_sub(1)).enumerate() { + if log == 0 { + return Err(Error::corrupt_file_named( + "miniblock_metadata", + format!( + "non-final chunk {chunk_index} of {num_chunks} has invalid log_num_values=0" + ), + )); + } + counted = counted.checked_add(1u64 << log).ok_or_else(|| { + Error::corrupt_file_named( + "miniblock_metadata", + format!( + "value count overflow at chunk {chunk_index}: counted_values={counted}, \ + log_num_values={log}, items_in_page={items_in_page}" + ), + ) + })?; + } + let last_chunk_values = items_in_page.checked_sub(counted).ok_or_else(|| { + Error::corrupt_file_named( + "miniblock_metadata", + format!( + "non-final chunks account for counted_values={counted}, exceeding \ + items_in_page={items_in_page}" + ), + ) + })?; + if let Some(&last_log) = logs.last() + && last_log != 0 + && (1u64 << last_log) != last_chunk_values + { + return Err(Error::corrupt_file_named( + "miniblock_metadata", + format!( + "final chunk log_num_values={last_log} does not match \ + last_chunk_values={last_chunk_values}: counted_values={counted}, \ + items_in_page={items_in_page}" + ), + )); + } + let uniform = num_chunks <= 1 || logs[..num_chunks - 1].iter().all(|&log| log == logs[0]); + // A single-chunk page has no "non-last" chunk to derive a stride from; use the + // page item count (min 1 so it stays a valid divisor in `find_chunk`). + let values_per_chunk = if num_chunks <= 1 { + items_in_page.max(1) + } else { + 1u64 << logs[0] + }; + Ok(FlatValueCounts { + logs, + uniform, + values_per_chunk, + last_chunk_values, + }) +} + +/// Iterator over per-chunk value counts for a non-uniform flat page. Non-last +/// chunks yield `1 << log`; the last yields the validated remaining item count. +fn flat_value_counts_iter(logs: &[u8], last_chunk_values: u64) -> impl Iterator + '_ { + let num_chunks = logs.len(); + (0..num_chunks).map(move |i| { + if i + 1 < num_chunks { + 1u64 << logs[i] + } else { + last_chunk_values + } + }) +} + +/// Builds the compact per-chunk index from the metadata words and, for nested +/// pages, the raw repetition-index bytes. The row axis is picked by page shape: +/// `UniformFlat` when all non-last chunks share a value count (fixed-width / +/// bitpacking), `Flat` for non-uniform flat pages (RLE / FSST), else `Nested`. +fn build_chunk_index( + words: &Words, + items_in_page: u64, + base: u64, + data_buf_size: u64, + rep_index_bytes: Option<&[u8]>, + repetition_index_depth: u16, +) -> Result { + let num_chunks = words.len(); + // Validate item counts before byte sizes because both share a metadata word, + // and an invalid count must not reach the final-chunk subtraction. + let value_counts = analyze_value_counts(words, items_in_page)?; + + // Each chunk stores `(divided_bytes + 1) * MINIBLOCK_ALIGNMENT` bytes, so the + // deltas are the chunk sizes and their grand total is the data buffer size. + let byte_starts = PrefixSums::from_deltas( + words + .iter() + .map(|word| ((word >> 4) as u64 + 1) * MINIBLOCK_ALIGNMENT as u64), + num_chunks, + data_buf_size, + ); + + // Nested pages track rows via the repetition index and keep leaf item counts + // separately; flat pages have row == value index, so value counts are rows. + let rows = if let Some(rep_index_data) = rep_index_bytes { + assert!(rep_index_data.len() % 8 == 0); + let stride = repetition_index_depth as usize + 1; + let (row_starts, has_trailer) = parse_nested_rep(rep_index_data, stride); + let item_counts = if value_counts.uniform { + ItemCounts::Uniform { + values_per_chunk: value_counts.values_per_chunk, + last_chunk_values: value_counts.last_chunk_values, + } + } else { + ItemCounts::PerChunkLog { + logs: value_counts.logs, + last_chunk_values: value_counts.last_chunk_values, + } + }; + RowMapping::Nested { + row_starts, + has_trailer, + item_counts, + } + } else { + if value_counts.uniform { + RowMapping::UniformFlat { + values_per_chunk: value_counts.values_per_chunk, + last_chunk_values: value_counts.last_chunk_values, + num_chunks, + } + } else { + let value_starts = PrefixSums::from_deltas( + flat_value_counts_iter(&value_counts.logs, value_counts.last_chunk_values), + num_chunks, + items_in_page, + ); + RowMapping::Flat { value_starts } + } + }; + + Ok(MiniBlockChunkIndex::new(base, byte_starts, rows)) +} + +impl StructuralPageScheduler for MiniBlockScheduler { + fn init_layout(&self) -> Result { + // We always need to fetch chunk metadata. We may also need to fetch a dictionary and + // we may also need to fetch the repetition index. Here, we gather what buffers we + // need. + let (meta_buf_position, meta_buf_size) = self.buffer_offsets_and_sizes[0]; + let dictionary_range = self.dictionary.as_ref().map(|dictionary| { + dictionary.dictionary_buf_position_and_size.0 + ..dictionary.dictionary_buf_position_and_size.0 + + dictionary.dictionary_buf_position_and_size.1 + }); + let repetition_index_range = if self.repetition_index_depth > 0 { + let (rep_index_pos, rep_index_size) = + self.buffer_offsets_and_sizes.last().ok_or_else(|| { + Error::invalid_input_source( + "Mini-block page has a repetition index but no buffers".into(), + ) + })?; + Some(*rep_index_pos..*rep_index_pos + *rep_index_size) + } else { + None + }; + Ok(PageInitialization::new([ + Some(meta_buf_position..meta_buf_position + meta_buf_size), + dictionary_range, + repetition_index_range, + ])) + } + + fn init_from_buffers<'a>( + &'a mut self, + buffers: PageInitializationBuffers, + _io: &Arc, + ) -> BoxFuture<'a, Result>> { + let base = self.buffer_offsets_and_sizes[1].0; + let data_buf_size = self.buffer_offsets_and_sizes[1].1; + + async move { + let [meta_bytes, dictionary_bytes, rep_index_bytes] = + buffers.try_into_array("mini-block")?; + let meta_bytes = meta_bytes.ok_or_else(|| { + Error::internal("Mini-block page initialization is missing chunk metadata") + })?; + + let words = Words::from_bytes(meta_bytes, self.has_large_chunk)?; + let chunk_index = build_chunk_index( + &words, + self.items_in_page, + base, + data_buf_size, + rep_index_bytes.as_deref(), + self.repetition_index_depth, + )?; + + // decode dictionary + let dictionary = match (self.dictionary.as_ref(), dictionary_bytes) { + (Some(dictionary), Some(dictionary_data)) => { + Some(Arc::new(dictionary.dictionary_decompressor.decompress( + Some(LanceBuffer::from_bytes( + dictionary_data, + dictionary.dictionary_data_alignment, + )), + dictionary.num_dictionary_items, + )?)) + } + (None, None) => None, + _ => { + return Err(Error::internal( + "Mini-block page initialization dictionary slot does not match its layout", + )); + } + }; + + let page_meta = Arc::new(MiniBlockCacheableState { + chunk_index, + dictionary, + }); + Ok(page_meta as Arc) + } + .boxed() + } + + fn try_load(&mut self, data: &Arc) -> Result<()> { + self.page_meta = Some(data.clone().as_arc_any().downcast().map_err(|_| { + Error::invalid_input_source("Cached mini-block page data has an unexpected type".into()) + })?); + Ok(()) + } + + fn schedule_ranges( + &self, + ranges: &[Range], + io: &Arc, + ) -> Result> { + let num_rows = ranges.iter().map(|r| r.end - r.start).sum(); + + let page_meta = self.page_meta.as_ref().unwrap(); + + let chunk_instructions = + ChunkInstructions::schedule_instructions(&page_meta.chunk_index, ranges); + + debug_assert_eq!( + num_rows, + chunk_instructions + .iter() + .map(|ci| ci.rows_to_take) + .sum::() + ); + + let chunks_needed = chunk_instructions + .iter() + .map(|ci| ci.chunk_idx) + .unique() + .collect::>(); + + let mut loaded_chunks = self.lookup_chunks(&chunks_needed); + let chunk_ranges = loaded_chunks + .iter() + .map(|c| c.byte_range.clone()) + .collect::>(); + let loaded_chunk_data = io.submit_request(chunk_ranges, self.priority); + + let rep_decompressor = self.rep_decompressor.clone(); + let def_decompressor = self.def_decompressor.clone(); + let value_decompressor = self.value_decompressor.clone(); + let num_buffers = self.num_buffers; + let has_large_chunk = self.has_large_chunk; + let dictionary = page_meta + .dictionary + .as_ref() + .map(|dictionary| dictionary.clone()); + let def_meaning = self.def_meaning.clone(); + + let res = async move { + let loaded_chunk_data = loaded_chunk_data.await?; + for (loaded_chunk, chunk_data) in loaded_chunks.iter_mut().zip(loaded_chunk_data) { + loaded_chunk.data = LanceBuffer::from_bytes(chunk_data, 1); + } + + Ok(Box::new(MiniBlockDecoder { + rep_decompressor, + def_decompressor, + value_decompressor, + def_meaning, + loaded_chunks: VecDeque::from_iter(loaded_chunks), + instructions: VecDeque::from(chunk_instructions), + offset_in_current_chunk: 0, + dictionary, + num_rows, + num_buffers, + has_large_chunk, + }) as Box) + } + .boxed(); + let page_load_task = PageLoadTask { + decoder_fut: res, + num_rows, + }; + Ok(vec![page_load_task]) + } +} + +#[derive(Debug, Clone, Copy)] +struct FullZipRepIndexDetails { + buf_position: u64, + bytes_per_value: u64, // Will be 1, 2, 4, or 8 +} + +#[derive(Debug)] +enum PerValueDecompressor { + Fixed(Arc), + Variable(Arc), +} + +#[derive(Debug)] +struct FullZipDecodeDetails { + value_decompressor: PerValueDecompressor, + def_meaning: Arc<[DefinitionInterpretation]>, + ctrl_word_parser: ControlWordParser, + max_rep: u16, + max_visible_def: u16, +} + +/// Describes where FullZip byte ranges should be read from. +/// +/// FullZip decoding always needs a list of byte ranges, but those bytes can come +/// from two different places: +/// - Remote I/O (normal path): ranges are fetched from the underlying `EncodingsIo`. +/// - A prefetched full page (full scan fast path): the entire page has already been +/// loaded once and ranges should be sliced from memory. +/// +/// This abstraction keeps scheduling code focused on "which ranges are needed" +/// instead of "how bytes are fetched", and it lets full-page scans avoid the +/// two-stage rep-index -> data I/O pipeline. +#[derive(Debug, Clone)] +enum FullZipReadSource { + /// Fetch ranges from the storage backend through the encoding I/O interface. + Remote(Arc), + /// Slice ranges from an already-loaded FullZip page buffer. + PrefetchedPage { base_offset: u64, data: LanceBuffer }, +} + +impl FullZipReadSource { + /// Materialize the requested ranges as decode-ready `LanceBuffer`s. + /// + /// The returned buffers preserve the input range order. + fn fetch( + &self, + ranges: &[Range], + priority: u64, + ) -> BoxFuture<'static, Result>> { + match self { + Self::Remote(io) => { + let io = io.clone(); + let ranges = ranges.to_vec(); + async move { + let data = io.submit_request(ranges, priority).await?; + Ok(data + .into_iter() + .map(|bytes| LanceBuffer::from_bytes(bytes, 1)) + .collect::>()) + } + .boxed() + } + Self::PrefetchedPage { base_offset, data } => { + let base_offset = *base_offset; + let data = data.clone(); + let page_end = base_offset + data.len() as u64; + std::future::ready( + ranges + .iter() + .map(|range| { + if range.start > range.end + || range.start < base_offset + || range.end > page_end + { + return Err(Error::internal(format!( + "Requested range {:?} is outside page range {}..{}", + range, base_offset, page_end + ))); + } + let start = (range.start - base_offset) as usize; + let len = (range.end - range.start) as usize; + Ok(data.slice_with_length(start, len)) + }) + .collect::>>(), + ) + .boxed() + } + } + } +} + +/// A scheduler for full-zip encoded data +/// +/// When the data type has a fixed-width then we simply need to map from +/// row ranges to byte ranges using the fixed-width of the data type. +/// +/// When the data type is variable-width or has any repetition then a +/// repetition index is required. +#[derive(Debug)] +pub struct FullZipScheduler { + data_buf_position: u64, + data_buf_size: u64, + rep_index: Option, + priority: u64, + rows_in_page: u64, + bits_per_offset: u8, + details: Arc, + /// Cached state containing the decoded repetition index + cached_state: Option>, + /// Whether repetition index metadata should be cached during initialize. + enable_cache: bool, +} + +impl FullZipScheduler { + fn try_new( + buffer_offsets_and_sizes: &[(u64, u64)], + priority: u64, + rows_in_page: u64, + layout: &pb21::FullZipLayout, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + let (data_buf_position, data_buf_size) = buffer_offsets_and_sizes[0]; + let rep_index = buffer_offsets_and_sizes.get(1).map(|(pos, len)| { + let num_reps = rows_in_page + 1; + let bytes_per_rep = len / num_reps; + debug_assert_eq!(len % num_reps, 0); + debug_assert!( + bytes_per_rep == 1 + || bytes_per_rep == 2 + || bytes_per_rep == 4 + || bytes_per_rep == 8 + ); + FullZipRepIndexDetails { + buf_position: *pos, + bytes_per_value: bytes_per_rep, + } + }); + + let value_decompressor = match layout.details { + Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => { + let decompressor = decompressors.create_fixed_per_value_decompressor( + layout.value_compression.as_ref().unwrap(), + )?; + PerValueDecompressor::Fixed(decompressor.into()) + } + Some(pb21::full_zip_layout::Details::BitsPerOffset(_)) => { + let decompressor = decompressors.create_variable_per_value_decompressor( + layout.value_compression.as_ref().unwrap(), + )?; + PerValueDecompressor::Variable(decompressor.into()) + } + None => { + panic!("Full-zip layout must have a `details` field"); + } + }; + let ctrl_word_parser = ControlWordParser::new( + layout.bits_rep.try_into().unwrap(), + layout.bits_def.try_into().unwrap(), + ); + let def_meaning = layout + .layers + .iter() + .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l)) + .collect::>(); + + let max_rep = def_meaning.iter().filter(|d| d.is_list()).count() as u16; + let max_visible_def = max_visible_level(&def_meaning); + + let bits_per_offset = match layout.details { + Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => 32, + Some(pb21::full_zip_layout::Details::BitsPerOffset(bits_per_offset)) => { + bits_per_offset as u8 + } + None => panic!("Full-zip layout must have a `details` field"), + }; + + let details = Arc::new(FullZipDecodeDetails { + value_decompressor, + def_meaning: def_meaning.into(), + ctrl_word_parser, + max_rep, + max_visible_def, + }); + Ok(Self { + data_buf_position, + data_buf_size, + rep_index, + details, + priority, + rows_in_page, + bits_per_offset, + cached_state: None, + enable_cache: false, + }) + } + + fn covers_entire_page(ranges: &[Range], rows_in_page: u64) -> bool { + if ranges.is_empty() { + return false; + } + let mut expected_start = 0; + for range in ranges { + if range.start != expected_start || range.end > rows_in_page || range.end < range.start + { + return false; + } + expected_start = range.end; + } + expected_start == rows_in_page + } + + fn create_page_load_task( + io_future: BoxFuture<'static, Result>>, + num_rows: u64, + details: Arc, + bits_per_offset: u8, + ) -> PageLoadTask { + let load_task = async move { + let buffers = io_future.await?; + let data = buffers + .into_iter() + .map(|bytes| LanceBuffer::from_bytes(bytes, 1)) + .collect::>(); + Self::create_decoder(details, data, num_rows, bits_per_offset) + } + .boxed(); + PageLoadTask { + decoder_fut: load_task, + num_rows, + } + } + + /// Creates a decoder from the loaded data + fn create_decoder( + details: Arc, + data: VecDeque, + num_rows: u64, + bits_per_offset: u8, + ) -> Result> { + match &details.value_decompressor { + PerValueDecompressor::Fixed(decompressor) => { + let bits_per_value = decompressor.bits_per_value(); + if bits_per_value % 8 != 0 { + return Err(lance_core::Error::not_supported_source("Bit-packed full-zip encoding (non-byte-aligned values) is not yet implemented".into())); + } + let bytes_per_value = bits_per_value / 8; + let total_bytes_per_value = + bytes_per_value as usize + details.ctrl_word_parser.bytes_per_word(); + // total_bytes_per_value == 0 is valid for constant-null FSL pages written by + // earlier encoders that produced bits_per_value=0 with no ctrl-word bytes. + // FixedFullZipDecoder::drain handles this case by producing AllNull output + // without touching the (empty) data buffer. + Ok(Box::new(FixedFullZipDecoder { + details, + data, + num_rows, + offset_in_current: 0, + bytes_per_value: bytes_per_value as usize, + total_bytes_per_value, + }) as Box) + } + PerValueDecompressor::Variable(_decompressor) => { + Ok(Box::new(VariableFullZipDecoder::new( + details, + data, + num_rows, + bits_per_offset, + bits_per_offset, + )?)) + } + } + } + + /// Extracts byte ranges from a repetition index buffer + /// The buffer contains pairs of (start, end) values for each range + fn extract_byte_ranges_from_pairs( + buffer: LanceBuffer, + bytes_per_value: u64, + data_buf_position: u64, + ) -> Vec> { + ByteUnpacker::new(buffer, bytes_per_value as usize) + .chunks(2) + .into_iter() + .map(|mut c| { + let start = c.next().unwrap() + data_buf_position; + let end = c.next().unwrap() + data_buf_position; + start..end + }) + .collect::>() + } + + /// Extracts byte ranges from a cached repetition index buffer + /// The buffer contains all values and we need to extract specific ranges + fn extract_byte_ranges_from_cached( + buffer: &LanceBuffer, + ranges: &[Range], + bytes_per_value: u64, + data_buf_position: u64, + ) -> Vec> { + ranges + .iter() + .map(|r| { + let start_offset = (r.start * bytes_per_value) as usize; + let end_offset = (r.end * bytes_per_value) as usize; + + let start_slice = &buffer[start_offset..start_offset + bytes_per_value as usize]; + let start_val = + ByteUnpacker::new(start_slice.iter().copied(), bytes_per_value as usize) + .next() + .unwrap(); + + let end_slice = &buffer[end_offset..end_offset + bytes_per_value as usize]; + let end_val = + ByteUnpacker::new(end_slice.iter().copied(), bytes_per_value as usize) + .next() + .unwrap(); + + (data_buf_position + start_val)..(data_buf_position + end_val) + }) + .collect() + } + + /// Computes the ranges in the repetition index that need to be loaded + fn compute_rep_index_ranges( + ranges: &[Range], + rep_index: &FullZipRepIndexDetails, + ) -> Vec> { + ranges + .iter() + .flat_map(|r| { + let first_val_start = + rep_index.buf_position + (r.start * rep_index.bytes_per_value); + let first_val_end = first_val_start + rep_index.bytes_per_value; + let last_val_start = rep_index.buf_position + (r.end * rep_index.bytes_per_value); + let last_val_end = last_val_start + rep_index.bytes_per_value; + [first_val_start..first_val_end, last_val_start..last_val_end] + }) + .collect() + } + + /// Schedules ranges in the presence of a repetition index + fn schedule_ranges_rep( + &self, + ranges: &[Range], + io: &Arc, + rep_index: FullZipRepIndexDetails, + ) -> Result> { + let num_rows = ranges.iter().map(|r| r.end - r.start).sum(); + let data_buf_position = self.data_buf_position; + let priority = self.priority; + let details = self.details.clone(); + let bits_per_offset = self.bits_per_offset; + + if Self::covers_entire_page(ranges, self.rows_in_page) { + let full_range = self.data_buf_position..(self.data_buf_position + self.data_buf_size); + let page_data = io.submit_single(full_range.clone(), priority); + let load_task = async move { + let page_data = page_data.await?; + let source = FullZipReadSource::PrefetchedPage { + base_offset: full_range.start, + data: LanceBuffer::from_bytes(page_data, 1), + }; + let read_ranges = vec![full_range]; + let data = source.fetch(&read_ranges, priority).await?; + Self::create_decoder(details, data, num_rows, bits_per_offset) + } + .boxed(); + let page_load_task = PageLoadTask { + decoder_fut: load_task, + num_rows, + }; + return Ok(vec![page_load_task]); + } + + if let Some(cached_state) = &self.cached_state { + let byte_ranges = Self::extract_byte_ranges_from_cached( + &cached_state.rep_index_buffer, + ranges, + rep_index.bytes_per_value, + data_buf_position, + ); + let io_future = io.submit_request(byte_ranges, priority); + let page_load_task = + Self::create_page_load_task(io_future, num_rows, details, bits_per_offset); + return Ok(vec![page_load_task]); + } + + let rep_ranges = Self::compute_rep_index_ranges(ranges, &rep_index); + let rep_data = io.submit_request(rep_ranges, priority); + let io_clone = io.clone(); + let load_task = async move { + let rep_data = rep_data.await?; + let rep_buffer = LanceBuffer::concat( + &rep_data + .into_iter() + .map(|d| LanceBuffer::from_bytes(d, 1)) + .collect::>(), + ); + let byte_ranges = Self::extract_byte_ranges_from_pairs( + rep_buffer, + rep_index.bytes_per_value, + data_buf_position, + ); + let source = FullZipReadSource::Remote(io_clone); + let data = source.fetch(&byte_ranges, priority).await?; + Self::create_decoder(details, data, num_rows, bits_per_offset) + } + .boxed(); + let page_load_task = PageLoadTask { + decoder_fut: load_task, + num_rows, + }; + Ok(vec![page_load_task]) + } + + // In the simple case there is no repetition and we just have large fixed-width + // rows of data. We can just map row ranges to byte ranges directly using the + // fixed-width of the data type. + fn schedule_ranges_simple( + &self, + ranges: &[Range], + io: &Arc, + ) -> Result> { + // Convert row ranges to item ranges (i.e. multiply by items per row) + let num_rows = ranges.iter().map(|r| r.end - r.start).sum(); + + let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else { + unreachable!() + }; + + // Convert item ranges to byte ranges (i.e. multiply by bytes per item) + let bits_per_value = decompressor.bits_per_value(); + if !bits_per_value.is_multiple_of(8) { + return Err(Error::invalid_input_source( + format!( + "Full-zip fixed-width values must be byte aligned, got {} bits per value", + bits_per_value + ) + .into(), + )); + } + let bytes_per_value = bits_per_value / 8; + let bytes_per_cw = self.details.ctrl_word_parser.bytes_per_word(); + let total_bytes_per_value = bytes_per_value + bytes_per_cw as u64; + let byte_ranges = ranges + .iter() + .map(|r| { + debug_assert!(r.end <= self.rows_in_page); + let start = self.data_buf_position + r.start * total_bytes_per_value; + let end = self.data_buf_position + r.end * total_bytes_per_value; + start..end + }) + .collect::>(); + + let io_future = io.submit_request(byte_ranges, self.priority); + let page_load_task = Self::create_page_load_task( + io_future, + num_rows, + self.details.clone(), + self.bits_per_offset, + ); + Ok(vec![page_load_task]) + } +} + +/// Cacheable state for FullZip encoding, storing the decoded repetition index +#[derive(Debug)] +struct FullZipCacheableState { + /// The raw repetition index buffer for future decoding + rep_index_buffer: LanceBuffer, +} + +impl DeepSizeOf for FullZipCacheableState { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + self.rep_index_buffer.len() + } +} + +impl CachedPageData for FullZipCacheableState { + fn as_arc_any(self: Arc) -> Arc { + self + } +} + +impl StructuralPageScheduler for FullZipScheduler { + fn needs_initialization(&self) -> bool { + self.enable_cache && self.rep_index.is_some() + } + + fn init_layout(&self) -> Result { + let repetition_index_range = if self.enable_cache + && let Some(rep_index) = self.rep_index + { + let total_size = (self.rows_in_page + 1) * rep_index.bytes_per_value; + Some(rep_index.buf_position..(rep_index.buf_position + total_size)) + } else { + None + }; + Ok(PageInitialization::new([repetition_index_range])) + } + + fn init_from_buffers<'a>( + &'a mut self, + buffers: PageInitializationBuffers, + _io: &Arc, + ) -> BoxFuture<'a, Result>> { + let [rep_index_buffer] = match buffers.try_into_array("full-zip") { + Ok(buffers) => buffers, + Err(error) => return std::future::ready(Err(error)).boxed(), + }; + if let Some(rep_index_buffer) = rep_index_buffer { + let state = Arc::new(FullZipCacheableState { + rep_index_buffer: LanceBuffer::from_bytes(rep_index_buffer, 1), + }); + std::future::ready(Ok(state as Arc)).boxed() + } else { + std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc)).boxed() + } + } + + /// Loads previously cached repetition index data from the cache system. + /// This method is called when a scheduler instance needs to use cached data + /// that was initialized by another instance or in a previous operation. + fn try_load(&mut self, cache: &Arc) -> Result<()> { + self.cached_state = Some(cache.clone().as_arc_any().downcast().map_err(|_| { + Error::invalid_input_source("Cached full-zip page data has an unexpected type".into()) + })?); + Ok(()) + } + + fn schedule_ranges( + &self, + ranges: &[Range], + io: &Arc, + ) -> Result> { + if let Some(rep_index) = self.rep_index { + self.schedule_ranges_rep(ranges, io, rep_index) + } else { + self.schedule_ranges_simple(ranges, io) + } + } +} + +/// A decoder for full-zip encoded data when the data has a fixed-width +/// +/// Here we need to unzip the control words from the values themselves and +/// then decompress the requested values. +/// +/// We use a PerValueDecompressor because we will only be decompressing the +/// requested data. This decoder / scheduler does not do any read amplification. +#[derive(Debug)] +struct FixedFullZipDecoder { + details: Arc, + data: VecDeque, + offset_in_current: usize, + bytes_per_value: usize, + total_bytes_per_value: usize, + num_rows: u64, +} + +impl FixedFullZipDecoder { + fn slice_next_task(&mut self, num_rows: u64) -> Result { + debug_assert!(num_rows > 0); + let cur_buf = self.data.front_mut().ok_or_else(|| { + Error::corrupt_file_named( + "fixed_full_zip", + format!("page data ran out with {num_rows} row(s) left to decode"), + ) + })?; + let start = self.offset_in_current; + if self.details.ctrl_word_parser.has_rep() { + // This is a slightly slower path. In order to figure out where to split we need to + // examine the rep index so we can convert num_lists to num_rows + let mut rows_started = 0; + // We always need at least one value. Now loop through until we have passed num_rows + // values + let mut num_items = 0; + let bytes_per_word = self.details.ctrl_word_parser.bytes_per_word(); + while self.offset_in_current < cur_buf.len() { + // The item walk is driven by the control words themselves, so a page that + // disagrees with its own layout can step past the end of the buffer. Check + // before every read and every advance so that a malformed page is reported + // instead of panicking in `slice_with_length`. + let bytes_avail = cur_buf.len() - self.offset_in_current; + if bytes_avail < bytes_per_word { + return Err(Error::corrupt_file_named( + "fixed_full_zip", + format!( + "truncated control word: {bytes_avail} byte(s) remain in the page \ + buffer but a control word requires {bytes_per_word}" + ), + )); + } + let control = self.details.ctrl_word_parser.parse_desc( + &cur_buf[self.offset_in_current..], + self.details.max_rep, + self.details.max_visible_def, + ); + if control.is_new_row { + if rows_started == num_rows { + break; + } + rows_started += 1; + } + num_items += 1; + let bytes_in_item = if control.is_visible { + self.total_bytes_per_value + } else { + bytes_per_word + }; + if bytes_in_item > bytes_avail { + return Err(Error::corrupt_file_named( + "fixed_full_zip", + format!( + "truncated item: {bytes_avail} byte(s) remain in the page buffer at \ + offset {} but the item requires {bytes_in_item}", + self.offset_in_current + ), + )); + } + self.offset_in_current += bytes_in_item; + } + + let task_slice = cur_buf.slice_with_length(start, self.offset_in_current - start); + if self.offset_in_current == cur_buf.len() { + self.data.pop_front(); + self.offset_in_current = 0; + } + + Ok(FullZipDecodeTaskItem { + data: PerValueDataBlock::Fixed(FixedWidthDataBlock { + data: task_slice, + bits_per_value: self.bytes_per_value as u64 * 8, + num_values: num_items, + block_info: BlockInfo::new(), + }), + rows_in_buf: rows_started, + }) + } else { + // If there's no repetition we can calculate the slicing point by just multiplying + // the number of rows by the total bytes per value + let cur_buf = self.data.front_mut().unwrap(); + let bytes_avail = cur_buf.len() - self.offset_in_current; + let offset_in_cur = self.offset_in_current; + + let bytes_needed = num_rows as usize * self.total_bytes_per_value; + let mut rows_taken = num_rows; + let task_slice = if bytes_needed >= bytes_avail { + self.offset_in_current = 0; + rows_taken = bytes_avail as u64 / self.total_bytes_per_value as u64; + self.data + .pop_front() + .unwrap() + .slice_with_length(offset_in_cur, bytes_avail) + } else { + self.offset_in_current += bytes_needed; + cur_buf.slice_with_length(offset_in_cur, bytes_needed) + }; + Ok(FullZipDecodeTaskItem { + data: PerValueDataBlock::Fixed(FixedWidthDataBlock { + data: task_slice, + bits_per_value: self.bytes_per_value as u64 * 8, + num_values: rows_taken, + block_info: BlockInfo::new(), + }), + rows_in_buf: rows_taken, + }) + } + } +} + +impl StructuralPageDecoder for FixedFullZipDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + if self.total_bytes_per_value == 0 { + // No bytes per row: constant-null page with no ctrl-word bytes. + // The decompressor (ConstantDecompressor) ignores its input and returns AllNull. + return Ok(Box::new(FixedFullZipDecodeTask { + details: self.details.clone(), + data: vec![FullZipDecodeTaskItem { + data: PerValueDataBlock::Fixed(FixedWidthDataBlock { + data: LanceBuffer::empty(), + bits_per_value: 0, + num_values: num_rows, + block_info: BlockInfo::new(), + }), + rows_in_buf: num_rows, + }], + bytes_per_value: 0, + num_rows: num_rows as usize, + })); + } + let mut task_data = Vec::with_capacity(self.data.len()); + let mut remaining = num_rows; + while remaining > 0 { + let task_item = self.slice_next_task(remaining)?; + remaining -= task_item.rows_in_buf; + task_data.push(task_item); + } + Ok(Box::new(FixedFullZipDecodeTask { + details: self.details.clone(), + data: task_data, + bytes_per_value: self.bytes_per_value, + num_rows: num_rows as usize, + })) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +/// A decoder for full-zip encoded data when the data has a variable-width +/// +/// Here we need to unzip the control words AND lengths from the values and +/// then decompress the requested values. +#[derive(Debug)] +struct VariableFullZipDecoder { + details: Arc, + decompressor: Arc, + data: LanceBuffer, + offsets: LanceBuffer, + rep: ScalarBuffer, + def: ScalarBuffer, + repdef_starts: Vec, + data_starts: Vec, + offset_starts: Vec, + visible_item_counts: Vec, + bits_per_offset: u8, + current_idx: usize, + num_rows: u64, +} + +impl VariableFullZipDecoder { + fn new( + details: Arc, + data: VecDeque, + num_rows: u64, + in_bits_per_length: u8, + out_bits_per_offset: u8, + ) -> Result { + let decompressor = match details.value_decompressor { + PerValueDecompressor::Variable(ref d) => d.clone(), + _ => unreachable!(), + }; + + assert_eq!(in_bits_per_length % 8, 0); + assert!(out_bits_per_offset == 32 || out_bits_per_offset == 64); + + let mut decoder = Self { + details, + decompressor, + data: LanceBuffer::empty(), + offsets: LanceBuffer::empty(), + rep: LanceBuffer::empty().borrow_to_typed_slice(), + def: LanceBuffer::empty().borrow_to_typed_slice(), + bits_per_offset: out_bits_per_offset, + repdef_starts: Vec::with_capacity(num_rows as usize + 1), + data_starts: Vec::with_capacity(num_rows as usize + 1), + offset_starts: Vec::with_capacity(num_rows as usize + 1), + visible_item_counts: Vec::with_capacity(num_rows as usize + 1), + current_idx: 0, + num_rows, + }; + + // There's no great time to do this and this is the least worst time. If we don't unzip then + // we can't slice the data during the decode phase. This is because we need the offsets to be + // unpacked to know where the values start and end. + // + // We don't want to unzip on the decode thread because that is a single-threaded path + // We don't want to unzip on the scheduling thread because that is a single-threaded path + // + // Fortunately, we know variable length data will always be read indirectly and so we can do it + // here, which should be on the indirect thread. The primary disadvantage to doing it here is that + // we load all the data into memory and then throw it away only to load it all into memory again during + // the decode. + // + // There are some alternatives to investigate: + // - Instead of just reading the beginning and end of the rep index we could read the entire + // range in between. This will give us the break points that we need for slicing and won't increase + // the number of IOPs but it will mean we are doing more total I/O and we need to load the rep index + // even when doing a full scan. + // - We could force each decode task to do a full unzip of all the data. Each decode task now + // has to do more work but the work is all fused. + // - We could just try doing this work on the decode thread and see if it is a problem. + decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?; + + Ok(decoder) + } + + fn slice_batch_data_and_rebase_offsets_typed( + data: &LanceBuffer, + offsets: &LanceBuffer, + ) -> Result<(LanceBuffer, LanceBuffer)> + where + T: arrow_buffer::ArrowNativeType + + Copy + + PartialOrd + + std::ops::Sub + + std::fmt::Display + + TryInto, + { + let offsets_slice = offsets.borrow_to_typed_slice::(); + let offsets_slice = offsets_slice.as_ref(); + if offsets_slice.is_empty() { + return Err(Error::internal( + "Variable offsets cannot be empty".to_string(), + )); + } + + let base = offsets_slice[0]; + let end = *offsets_slice.last().unwrap(); + if end < base { + return Err(Error::internal(format!( + "Invalid variable offsets: end ({end}) is less than base ({base})" + ))); + } + + let data_start = base.try_into().map_err(|_| { + Error::internal(format!("Variable offset ({base}) does not fit into usize")) + })?; + let data_end = end.try_into().map_err(|_| { + Error::internal(format!("Variable offset ({end}) does not fit into usize")) + })?; + if data_end > data.len() { + return Err(Error::internal(format!( + "Invalid variable offsets: end ({data_end}) exceeds data len ({})", + data.len() + ))); + } + + let mut rebased_offsets = Vec::with_capacity(offsets_slice.len()); + for &offset in offsets_slice { + if offset < base { + return Err(Error::internal(format!( + "Invalid variable offsets: offset ({offset}) is less than base ({base})" + ))); + } + rebased_offsets.push(offset - base); + } + + let sliced_data = data.slice_with_length(data_start, data_end - data_start); + // Copy into a compact buffer so each output batch owns only what it references. + let sliced_data = LanceBuffer::copy_slice(&sliced_data); + let rebased_offsets = LanceBuffer::reinterpret_vec(rebased_offsets); + Ok((sliced_data, rebased_offsets)) + } + + fn slice_batch_data_and_rebase_offsets( + data: &LanceBuffer, + offsets: &LanceBuffer, + bits_per_offset: u8, + ) -> Result<(LanceBuffer, LanceBuffer)> { + match bits_per_offset { + 32 => Self::slice_batch_data_and_rebase_offsets_typed::(data, offsets), + 64 => Self::slice_batch_data_and_rebase_offsets_typed::(data, offsets), + _ => Err(Error::internal(format!( + "Unsupported bits_per_offset={bits_per_offset}" + ))), + } + } + + /// Reads a single length prefix from the front of `data`. + /// + /// The bytes come from the file. A page whose item walk ends with a partial + /// trailing item leaves fewer than `bits_per_offset / 8` bytes here, so this + /// is bounds checked and reports a corrupt file rather than reading past the + /// end of the buffer. + fn parse_length(data: &[u8], bits_per_offset: u8) -> Result { + let width = bits_per_offset as usize / 8; + if data.len() < width { + return Err(Error::corrupt_file_named( + "variable_full_zip", + format!( + "truncated length prefix: {} byte(s) remain in the page buffer but a \ + {}-bit length prefix requires {}", + data.len(), + bits_per_offset, + width + ), + )); + } + Ok(match bits_per_offset { + 8 => data[0] as u64, + 16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64, + 32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64, + 64 => u64::from_le_bytes(data[..8].try_into().unwrap()), + _ => unreachable!(), + }) + } + + fn unzip( + &mut self, + data: VecDeque, + in_bits_per_length: u8, + out_bits_per_offset: u8, + num_rows: u64, + ) -> Result<()> { + // This undercounts if there are lists but, at this point, we don't really know how many items we have + let mut rep = Vec::with_capacity(num_rows as usize); + let mut def = Vec::with_capacity(num_rows as usize); + let bytes_cw = self.details.ctrl_word_parser.bytes_per_word() * num_rows as usize; + + // This undercounts if there are lists + // It can also overcount if there are invisible items + let bytes_per_offset = out_bits_per_offset as usize / 8; + let bytes_offsets = bytes_per_offset * (num_rows as usize + 1); + let mut offsets_data = Vec::with_capacity(bytes_offsets); + + let bytes_per_length = in_bits_per_length as usize / 8; + let bytes_lengths = bytes_per_length * num_rows as usize; + + let bytes_data = data.iter().map(|d| d.len()).sum::(); + // This overcounts since bytes_lengths and bytes_cw are undercounts + // It can also undercount if there are invisible items (hence the saturating_sub) + let mut unzipped_data = + Vec::with_capacity((bytes_data - bytes_cw).saturating_sub(bytes_lengths)); + + let mut current_offset = 0_u64; + let mut visible_item_count = 0_u64; + for databuf in data.into_iter() { + let mut databuf = databuf.as_ref(); + while !databuf.is_empty() { + let data_start = unzipped_data.len(); + let offset_start = offsets_data.len(); + // We might have only-rep or only-def, neither, or both. They move at the same + // speed though so we only need one index into it + let repdef_start = rep.len().max(def.len()); + // TODO: Kind of inefficient we parse the control word twice here + let ctrl_desc = self.details.ctrl_word_parser.parse_desc( + databuf, + self.details.max_rep, + self.details.max_visible_def, + ); + self.details + .ctrl_word_parser + .parse(databuf, &mut rep, &mut def); + databuf = &databuf[self.details.ctrl_word_parser.bytes_per_word()..]; + + if ctrl_desc.is_new_row { + self.repdef_starts.push(repdef_start); + self.data_starts.push(data_start); + self.offset_starts.push(offset_start); + self.visible_item_counts.push(visible_item_count); + } + if ctrl_desc.is_visible { + visible_item_count += 1; + if ctrl_desc.is_valid_item { + let length = Self::parse_length(databuf, in_bits_per_length)?; + match out_bits_per_offset { + 32 => offsets_data + .extend_from_slice(&(current_offset as u32).to_le_bytes()), + 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()), + _ => unreachable!(), + }; + databuf = &databuf[bytes_per_offset..]; + unzipped_data.extend_from_slice(&databuf[..length as usize]); + databuf = &databuf[length as usize..]; + current_offset += length; + } else { + // Null items still get an offset + match out_bits_per_offset { + 32 => offsets_data + .extend_from_slice(&(current_offset as u32).to_le_bytes()), + 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()), + _ => unreachable!(), + } + } + } + } + } + self.repdef_starts.push(rep.len().max(def.len())); + self.data_starts.push(unzipped_data.len()); + self.offset_starts.push(offsets_data.len()); + self.visible_item_counts.push(visible_item_count); + match out_bits_per_offset { + 32 => offsets_data.extend_from_slice(&(current_offset as u32).to_le_bytes()), + 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()), + _ => unreachable!(), + }; + self.rep = ScalarBuffer::from(rep); + self.def = ScalarBuffer::from(def); + self.data = LanceBuffer::from(unzipped_data); + self.offsets = LanceBuffer::from(offsets_data); + Ok(()) + } +} + +impl StructuralPageDecoder for VariableFullZipDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + let start = self.current_idx; + let end = start + num_rows as usize; + + let offset_start = self.offset_starts[start]; + let offset_end = self.offset_starts[end] + (self.bits_per_offset as usize / 8); + let offsets = self + .offsets + .slice_with_length(offset_start, offset_end - offset_start); + // Keep each batch's variable data buffer bounded to the selected rows. + let (data, offsets) = + Self::slice_batch_data_and_rebase_offsets(&self.data, &offsets, self.bits_per_offset)?; + + let repdef_start = self.repdef_starts[start]; + let repdef_end = self.repdef_starts[end]; + let rep = if self.rep.is_empty() { + self.rep.clone() + } else { + self.rep.slice(repdef_start, repdef_end - repdef_start) + }; + let def = if self.def.is_empty() { + self.def.clone() + } else { + self.def.slice(repdef_start, repdef_end - repdef_start) + }; + + let visible_item_counts_start = self.visible_item_counts[start]; + let visible_item_counts_end = self.visible_item_counts[end]; + let num_visible_items = visible_item_counts_end - visible_item_counts_start; + + self.current_idx += num_rows as usize; + + Ok(Box::new(VariableFullZipDecodeTask { + details: self.details.clone(), + decompressor: self.decompressor.clone(), + data, + offsets, + bits_per_offset: self.bits_per_offset, + num_visible_items, + rep, + def, + })) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +#[derive(Debug)] +struct VariableFullZipDecodeTask { + details: Arc, + decompressor: Arc, + data: LanceBuffer, + offsets: LanceBuffer, + bits_per_offset: u8, + num_visible_items: u64, + rep: ScalarBuffer, + def: ScalarBuffer, +} + +impl DecodePageTask for VariableFullZipDecodeTask { + fn decode(self: Box) -> Result { + let block = VariableWidthBlock { + data: self.data, + offsets: self.offsets, + bits_per_offset: self.bits_per_offset, + num_values: self.num_visible_items, + block_info: BlockInfo::new(), + }; + let decomopressed = self.decompressor.decompress(block)?; + let rep = if self.rep.is_empty() { + None + } else { + Some(self.rep.to_vec()) + }; + let def = if self.def.is_empty() { + None + } else { + Some(self.def.to_vec()) + }; + let unraveler = RepDefUnraveler::new( + rep, + def, + self.details.def_meaning.clone(), + self.num_visible_items, + ); + Ok(DecodedPage { + data: decomopressed, + repdef: unraveler, + }) + } +} + +#[derive(Debug)] +struct FullZipDecodeTaskItem { + data: PerValueDataBlock, + rows_in_buf: u64, +} + +/// A task to unzip and decompress full-zip encoded data when that data +/// has a fixed-width. +#[derive(Debug)] +struct FixedFullZipDecodeTask { + details: Arc, + data: Vec, + num_rows: usize, + bytes_per_value: usize, +} + +impl DecodePageTask for FixedFullZipDecodeTask { + fn decode(self: Box) -> Result { + let estimated_size_bytes = if self.details.ctrl_word_parser.bytes_per_word() == 0 { + let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else { + return Err(Error::internal( + "FixedFullZipDecodeTask requires a fixed-width decompressor", + )); + }; + decompressor + .decoded_size_bytes(self.num_rows as u64) + .unwrap_or_else(|| { + self.data + .iter() + .map(|task_item| task_item.data.data_size()) + .sum::() + * 2 + }) + } else { + // Rep/def levels can suppress values, so the exact output size is not known + // until they are decoded. Keep the existing conservative estimate. + self.data + .iter() + .map(|task_item| task_item.data.data_size()) + .sum::() + * 2 + }; + let mut data_builder = DataBlockBuilder::with_capacity_estimate(estimated_size_bytes); + + if self.details.ctrl_word_parser.bytes_per_word() == 0 { + // Fast path, no need to unzip because there is no rep/def + // + // We decompress each buffer and add it to our output buffer + for task_item in self.data.into_iter() { + let PerValueDataBlock::Fixed(fixed_data) = task_item.data else { + unreachable!() + }; + let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor + else { + unreachable!() + }; + debug_assert_eq!(fixed_data.num_values, task_item.rows_in_buf); + let decompressed = decompressor.decompress(fixed_data, task_item.rows_in_buf)?; + data_builder.append(&decompressed, 0..task_item.rows_in_buf)?; + } + + let unraveler = RepDefUnraveler::new( + None, + None, + self.details.def_meaning.clone(), + self.num_rows as u64, + ); + + Ok(DecodedPage { + data: data_builder.finish(), + repdef: unraveler, + }) + } else { + // Slow path, unzipping needed + let mut rep = Vec::with_capacity(self.num_rows); + let mut def = Vec::with_capacity(self.num_rows); + + for task_item in self.data.into_iter() { + let PerValueDataBlock::Fixed(fixed_data) = task_item.data else { + unreachable!() + }; + let mut buf_slice = fixed_data.data.as_ref(); + let num_values = fixed_data.num_values as usize; + // We will be unzipping repdef in to `rep` and `def` and the + // values into `values` (which contains the compressed values) + let mut values = Vec::with_capacity( + fixed_data.data.len() + - (self.details.ctrl_word_parser.bytes_per_word() * num_values), + ); + let mut visible_items = 0; + for _ in 0..num_values { + // Extract rep/def + self.details + .ctrl_word_parser + .parse(buf_slice, &mut rep, &mut def); + buf_slice = &buf_slice[self.details.ctrl_word_parser.bytes_per_word()..]; + + let is_visible = def + .last() + .map(|d| *d <= self.details.max_visible_def) + .unwrap_or(true); + if is_visible { + // Extract value + values.extend_from_slice(buf_slice[..self.bytes_per_value].as_ref()); + buf_slice = &buf_slice[self.bytes_per_value..]; + visible_items += 1; + } + } + + // Finally, we decompress the values and add them to our output buffer + let values_buf = LanceBuffer::from(values); + let fixed_data = FixedWidthDataBlock { + bits_per_value: self.bytes_per_value as u64 * 8, + block_info: BlockInfo::new(), + data: values_buf, + num_values: visible_items, + }; + let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor + else { + unreachable!() + }; + let decompressed = decompressor.decompress(fixed_data, visible_items)?; + data_builder.append(&decompressed, 0..visible_items)?; + } + + let repetition = if rep.is_empty() { None } else { Some(rep) }; + let definition = if def.is_empty() { None } else { Some(def) }; + + let unraveler = RepDefUnraveler::new( + repetition, + definition, + self.details.def_meaning.clone(), + self.num_rows as u64, + ); + let data = data_builder.finish(); + + Ok(DecodedPage { + data, + repdef: unraveler, + }) + } + } +} + +#[derive(Debug)] +struct StructuralPrimitiveFieldSchedulingJob<'a> { + scheduler: &'a StructuralPrimitiveFieldScheduler, + mappings: std::vec::IntoIter, +} + +impl<'a> StructuralPrimitiveFieldSchedulingJob<'a> { + pub fn try_new( + scheduler: &'a StructuralPrimitiveFieldScheduler, + ranges: Vec>, + ) -> Result { + let mappings = map_ranges_to_pages(&scheduler.page_row_ends, &ranges)?.into_iter(); + Ok(Self { + scheduler, + mappings, + }) + } +} + +impl StructuralSchedulingJob for StructuralPrimitiveFieldSchedulingJob<'_> { + fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result> { + let Some(first_mapping) = self.mappings.next() else { + return Ok(Vec::new()); + }; + // Pages are scheduled in row order until their reads add up to a + // batch, then submitted together so the I/O scheduler can coalesce + // neighbouring pages. + let batch = Arc::new(RequestBatch::new(context.io().clone())); + let batch_io: Arc = batch.clone(); + let cur_path = context.current_path(); + let mut scan_lines = Vec::new(); + let mut mapping = Some(first_mapping); + while let Some(cur_mapping) = mapping.take() { + let cur_page = self.scheduler.built_page(cur_mapping.page_idx)?; + + trace!( + "Scheduling {} rows across {} ranges from page with {} rows (column_index={}, page_index={})", + cur_mapping + .ranges_in_page + .iter() + .map(|r| r.end - r.start) + .sum::(), + cur_mapping.ranges_in_page.len(), + cur_page.row_range.end - cur_page.row_range.start, + self.scheduler.column_index, + cur_page.page_index, + ); + + let page_decoders = cur_page + .scheduler + .schedule_ranges(&cur_mapping.ranges_in_page, &batch_io)?; + + scan_lines.extend(page_decoders.into_iter().map(|page_load_task| { + let cur_path = cur_path.clone(); + let page_decoder = page_load_task.decoder_fut; + let unloaded_page = async move { + let page_decoder = page_decoder.await?; + Ok(LoadedPageShard { + decoder: page_decoder, + path: cur_path, + }) + } + .boxed(); + ScheduledScanLine { + decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(unloaded_page))], + rows_scheduled: page_load_task.num_rows, + } + })); + + if batch.pending_bytes() < IO_REQUEST_BATCH_BYTES { + mapping = self.mappings.next(); + } + } + batch.flush(); + Ok(scan_lines) + } +} + +#[derive(Debug)] +struct PageInfoAndScheduler { + page_index: usize, + row_range: Range, + scheduler: Box, +} + +/// Everything needed to build a page scheduler from its `PageInfo` on demand. +/// +/// Building a page scheduler clones buffer tables and instantiates +/// decompressors, so a column with hundreds of thousands of pages must not pay +/// for every page on every read. The field scheduler keeps this source and +/// builds only the pages `initialize` is asked to cover. +#[derive(Debug)] +struct PageSchedulerSource { + column_info: Arc, + decompressors: Arc, + cache_repetition_index: bool, + target_field: Field, +} + +/// A scheduler for a leaf node +/// +/// Here we look at the layout of the various pages and delegate scheduling to a scheduler +/// appropriate for the layout of the page. +#[derive(Debug)] +pub struct StructuralPrimitiveFieldScheduler { + /// Exclusive end row of every page in the column, in page order, so page + /// `i` covers `page_row_ends[i - 1]..page_row_ends[i]` (page 0 starts at 0). + page_row_ends: Vec, + /// Page schedulers built so far, ordered by `page_index`. After + /// `initialize` this holds exactly the pages the requested ranges touch, + /// or every page when no ranges were given. + page_schedulers: Vec, + /// Metadata to build pages that are not in `page_schedulers` yet. `None` + /// once every page has been built. + page_source: Option, + column_index: u32, + target_data_type_tag: Arc, +} + +impl StructuralPrimitiveFieldScheduler { + /// Creates the scheduler and builds every page scheduler up front. + /// + /// This pays for all pages of the column on every read. Prefer + /// [`Self::try_new_lazy`], which builds only the pages a read touches. + #[deprecated( + since = "13.0.0", + note = "builds every page scheduler eagerly; use try_new_lazy" + )] + pub fn try_new( + column_info: &ColumnInfo, + decompressors: &dyn DecompressionStrategy, + cache_repetition_index: bool, + target_field: &Field, + ) -> Result { + let page_row_ends = Self::page_row_ends(column_info)?; + let mut page_schedulers = Vec::with_capacity(page_row_ends.len()); + let mut row_start = 0_u64; + for (page_index, (page_info, &row_end)) in column_info + .page_infos + .iter() + .zip(&page_row_ends) + .enumerate() + { + page_schedulers.push(Self::page_info_to_scheduler( + page_info, + page_index, + row_start..row_end, + decompressors, + cache_repetition_index, + target_field, + )?); + row_start = row_end; + } + Ok(Self { + page_row_ends, + page_schedulers, + page_source: None, + column_index: column_info.index, + target_data_type_tag: format!("{:?}", target_field.data_type()).into(), + }) + } + + /// Creates the scheduler without building any page schedulers. + /// + /// Only the per-page row boundaries are computed here. Page schedulers + /// are built in [`StructuralFieldScheduler::initialize`] for the pages the + /// requested ranges touch (or for every page when no ranges are given), so + /// a small read from a column with many pages stays cheap. + pub fn try_new_lazy( + column_info: Arc, + decompressors: Arc, + cache_repetition_index: bool, + target_field: &Field, + ) -> Result { + Ok(Self { + page_row_ends: Self::page_row_ends(&column_info)?, + page_schedulers: Vec::new(), + column_index: column_info.index, + target_data_type_tag: format!("{:?}", target_field.data_type()).into(), + page_source: Some(PageSchedulerSource { + column_info, + decompressors, + cache_repetition_index, + target_field: target_field.clone(), + }), + }) + } + + /// Exclusive end row of every page of `column_info`, in page order. + fn page_row_ends(column_info: &ColumnInfo) -> Result> { + let mut page_row_ends = Vec::with_capacity(column_info.page_infos.len()); + let mut row_start = 0_u64; + for (page_index, page_info) in column_info.page_infos.iter().enumerate() { + let row_end = row_start.checked_add(page_info.num_rows).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Page row range overflowed for column {} at page {page_index}: start={row_start}, rows={}", + column_info.index, page_info.num_rows + ) + .into(), + ) + })?; + page_row_ends.push(row_end); + row_start = row_end; + } + Ok(page_row_ends) + } + + /// Wraps already-built page schedulers; nothing is built lazily. + #[cfg(test)] + fn from_page_schedulers( + page_schedulers: Vec, + column_index: u32, + target_data_type_tag: Arc, + ) -> Self { + Self { + page_row_ends: page_schedulers + .iter() + .map(|page| page.row_range.end) + .collect(), + page_schedulers, + page_source: None, + column_index, + target_data_type_tag, + } + } + + fn page_row_range(&self, page_index: usize) -> Range { + let start = match page_index.checked_sub(1) { + Some(previous) => self.page_row_ends[previous], + None => 0, + }; + start..self.page_row_ends[page_index] + } + + /// Position of `page_index` in `page_schedulers` if it has been built. + fn built_page_position(&self, page_index: usize) -> Option { + self.page_schedulers + .binary_search_by_key(&page_index, |page| page.page_index) + .ok() + } + + fn built_page(&self, page_index: usize) -> Result<&PageInfoAndScheduler> { + let position = self.built_page_position(page_index).ok_or_else(|| { + Error::internal(format!( + "Page {page_index} of column {} was scheduled but not initialized; the scheduled ranges must be covered by the ranges passed to initialize", + self.column_index + )) + })?; + Ok(&self.page_schedulers[position]) + } + + /// Builds the schedulers for `page_indices` (ascending) that do not exist + /// yet and returns the positions of all of them in `page_schedulers`. + fn build_pages(&mut self, page_indices: &[usize]) -> Result> { + if let Some(source) = &self.page_source { + let mut missing = Vec::new(); + for &page_index in page_indices { + if self.built_page_position(page_index).is_some() { + continue; + } + let page_info = source + .column_info + .page_infos + .get(page_index) + .ok_or_else(|| { + Error::internal(format!( + "Page {page_index} is out of bounds for column {} with {} pages", + self.column_index, + source.column_info.page_infos.len() + )) + })?; + missing.push(Self::page_info_to_scheduler( + page_info, + page_index, + self.page_row_range(page_index), + source.decompressors.as_ref(), + source.cache_repetition_index, + &source.target_field, + )?); + } + if !missing.is_empty() { + self.page_schedulers = std::mem::take(&mut self.page_schedulers) + .into_iter() + .merge_by(missing, |built, new| built.page_index < new.page_index) + .collect(); + } + if self.page_schedulers.len() == self.page_row_ends.len() { + self.page_source = None; + } + } + page_indices + .iter() + .map(|&page_index| { + self.built_page_position(page_index).ok_or_else(|| { + Error::internal(format!( + "Page {page_index} of column {} is missing after construction", + self.column_index + )) + }) + }) + .collect() + } + + fn page_layout_to_scheduler( + page_info: &PageInfo, + page_layout: &PageLayout, + decompressors: &dyn DecompressionStrategy, + cache_repetition_index: bool, + target_field: &Field, + ) -> Result> { + use pb21::page_layout::Layout; + Ok(match page_layout.layout.as_ref().expect_ok()? { + Layout::MiniBlockLayout(mini_block) => Box::new(MiniBlockScheduler::try_new( + &page_info.buffer_offsets_and_sizes, + page_info.priority, + mini_block.num_items, + mini_block, + decompressors, + )?), + Layout::SparseLayout(sparse_layout) => { + Box::new(sparse::SparseStructuralScheduler::try_new( + &page_info.buffer_offsets_and_sizes, + page_info.priority, + page_info.num_rows, + target_field.data_type(), + sparse_layout, + decompressors, + )?) + } + Layout::FullZipLayout(full_zip) => { + let mut scheduler = FullZipScheduler::try_new( + &page_info.buffer_offsets_and_sizes, + page_info.priority, + page_info.num_rows, + full_zip, + decompressors, + )?; + scheduler.enable_cache = cache_repetition_index; + Box::new(scheduler) + } + Layout::ConstantLayout(constant_layout) => { + let def_meaning = constant_layout + .layers + .iter() + .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l)) + .collect::>(); + let has_scalar_value = constant_layout.inline_value.is_some() + || page_info.buffer_offsets_and_sizes.len() == 1 + || page_info.buffer_offsets_and_sizes.len() == 3; + if has_scalar_value { + Box::new(constant::ConstantPageScheduler::try_new( + page_info.buffer_offsets_and_sizes.clone(), + constant_layout.inline_value.clone(), + target_field.data_type(), + def_meaning.into(), + )?) as Box + } else if def_meaning.len() == 1 + && def_meaning[0] == DefinitionInterpretation::NullableItem + { + Box::new(SimpleAllNullScheduler::default()) as Box + } else { + // RLE levels select a validated cache representation; other + // block compressions keep flowing through the eager decompressor. + let rep_codec = LevelCodec::try_new( + constant_layout.rep_compression.as_ref(), + decompressors, + )?; + let def_codec = LevelCodec::try_new( + constant_layout.def_compression.as_ref(), + decompressors, + )?; + + Box::new(ComplexAllNullScheduler::new( + page_info.buffer_offsets_and_sizes.clone(), + def_meaning.into(), + rep_codec, + def_codec, + constant_layout.num_rep_values, + constant_layout.num_def_values, + )) as Box + } + } + Layout::BlobLayout(blob) => { + let inner_scheduler = Self::page_layout_to_scheduler( + page_info, + blob.inner_layout.as_ref().expect_ok()?.as_ref(), + decompressors, + cache_repetition_index, + target_field, + )?; + let def_meaning = blob + .layers + .iter() + .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l)) + .collect::>(); + if matches!(target_field.data_type(), DataType::Struct(_)) { + // User wants to decode blob into struct + Box::new(BlobDescriptionPageScheduler::new( + inner_scheduler, + def_meaning.into(), + )) + } else { + // User wants to decode blob into binary data + Box::new(BlobPageScheduler::new( + inner_scheduler, + page_info.priority, + page_info.num_rows, + def_meaning.into(), + )) + } + } + }) + } + + fn page_info_to_scheduler( + page_info: &PageInfo, + page_index: usize, + row_range: Range, + decompressors: &dyn DecompressionStrategy, + cache_repetition_index: bool, + target_field: &Field, + ) -> Result { + let page_layout = page_info.encoding.as_structural(); + let scheduler = Self::page_layout_to_scheduler( + page_info, + page_layout, + decompressors, + cache_repetition_index, + target_field, + )?; + Ok(PageInfoAndScheduler { + page_index, + row_range, + scheduler, + }) + } +} + +pub trait CachedPageData: Any + Send + Sync + DeepSizeOf + 'static { + fn as_arc_any(self: Arc) -> Arc; +} + +pub struct NoCachedPageData; + +impl DeepSizeOf for NoCachedPageData { + fn deep_size_of_children(&self, _ctx: &mut Context) -> usize { + 0 + } +} +impl CachedPageData for NoCachedPageData { + fn as_arc_any(self: Arc) -> Arc { + self + } +} + +#[doc = "Deprecated compatibility wrapper; new code uses per-page PageDataCacheKey entries."] +pub struct CachedFieldData { + pages: Vec>, +} + +impl DeepSizeOf for CachedFieldData { + fn deep_size_of_children(&self, ctx: &mut Context) -> usize { + self.pages.deep_size_of_children(ctx) + } +} + +/// Compatibility key for the former whole-field page metadata cache. +#[derive(Debug, Clone)] +pub struct FieldDataCacheKey { + pub column_index: u32, + pub view_tag: String, +} + +impl CacheKey for FieldDataCacheKey { + type ValueType = CachedFieldData; + + fn key(&self) -> std::borrow::Cow<'_, str> { + format!("{}:{}", self.column_index, self.view_tag).into() + } + + fn type_name() -> &'static str { + "FieldData" + } + + fn schema() -> CacheKeySchema { + CacheKeySchema::new("lance.encoding.logical.primitive.field-data-key", 1) + } + + fn write_key(&self, builder: &mut KeyBuilder) { + builder.write_u32(self.column_index); + builder.write_str(&self.view_tag); + } +} + +/// Sized wrapper around a page's cached metadata so [`PageDataCacheKey`] can use +/// the sized [`CacheKey`] path; the concrete per-layout state still lives behind +/// `dyn CachedPageData`. +pub struct CachedPage { + data: Arc, +} + +impl DeepSizeOf for CachedPage { + fn deep_size_of_children(&self, ctx: &mut Context) -> usize { + self.data.deep_size_of_children(ctx) + } +} + +// Cache key for a single page's initialized metadata, keyed per page (not per +// column) so a read populates and reuses only the pages it touches. +// +// `view` distinguishes decode shapes of one physical column: a blob column +// materializes as a `Struct` descriptor in one variant and raw +// `LargeBinary` in another, each caching a different `CachedPageData` type. +// Without it, two readers at the same column/page but different shapes collide +// and crash on a downcast. +/// Decode shape used to isolate cached page metadata for one physical column. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum PageCacheView { + /// The column's native non-blob representation. + Native, + /// Blob position and size descriptors materialized as a struct. + BlobDescriptor, + /// Blob contents materialized as binary values. + BlobPayload, +} + +impl PageCacheView { + fn cache_tag(self) -> &'static str { + match self { + Self::Native => "n", + Self::BlobDescriptor => "bd", + Self::BlobPayload => "bp", + } + } +} + +#[derive(Debug, Clone)] +pub struct PageDataCacheKey { + pub column_index: u32, + pub page_index: u32, + pub view: PageCacheView, + pub target_data_type_tag: Arc, +} + +impl CacheKey for PageDataCacheKey { + type ValueType = CachedPage; + + fn key(&self) -> std::borrow::Cow<'_, str> { + format!( + "v2/{}/{}/{}/{}", + self.column_index, + self.page_index, + self.view.cache_tag(), + self.target_data_type_tag + ) + .into() + } + + fn type_name() -> &'static str { + "PageData" + } + + fn schema() -> CacheKeySchema { + CacheKeySchema::new("lance.encoding.logical.primitive.page-data-key", 2) + } + + fn write_key(&self, builder: &mut KeyBuilder) { + builder.write_u32(self.column_index); + builder.write_u32(self.page_index); + builder.write_variant(match self.view { + PageCacheView::Native => 0, + PageCacheView::BlobDescriptor => 1, + PageCacheView::BlobPayload => 2, + }); + builder.write_str(&self.target_data_type_tag); + } +} + +/// One page's slice of a scheduling request: the page's index and the +/// page-local row sub-ranges (relative to the page's first row) of the request +/// that fall in it. +#[derive(Debug)] +struct PageRangeMapping { + page_idx: usize, + ranges_in_page: Vec>, +} + +/// Visits every intersection between ordered row ranges and ordered pages. +/// +/// `page_row_ends` holds the exclusive end row of each page (see +/// [`StructuralPrimitiveFieldScheduler::page_row_ends`]). Each range +/// binary-searches its first overlapping page and then visits only the pages +/// it touches. For `R` ranges, `P` pages, and `T` intersections this takes +/// `O(R log P + T)` time without scanning sparse gaps. +fn for_each_page_range( + page_row_ends: &[u64], + ranges: &[Range], + mut visit: impl FnMut(usize, Range), +) -> Result<()> { + for (range_idx, range) in ranges.iter().enumerate() { + if range.start > range.end { + return Err(Error::invalid_input_source( + format!( + "Requested row range {range_idx} is reversed: {}..{}", + range.start, range.end + ) + .into(), + )); + } + if range_idx > 0 { + let previous = &ranges[range_idx - 1]; + if previous.start >= range.start || previous.end > range.start { + return Err(Error::invalid_input_source( + format!( + "Requested row ranges overlap or are not strictly ordered at index {range_idx}: {previous:?}, then {range:?}" + ) + .into(), + )); + } + } + } + + for range in ranges { + let mut page_idx = page_row_ends.partition_point(|page_end| *page_end <= range.start); + let mut page_start = match page_idx.checked_sub(1) { + Some(previous) => page_row_ends[previous], + None => 0, + }; + while let Some(&page_end) = page_row_ends.get(page_idx) { + if page_start >= range.end { + break; + } + let start_in_page = range.start.max(page_start) - page_start; + let end_in_page = range.end.min(page_end) - page_start; + if start_in_page < end_in_page { + visit(page_idx, start_in_page..end_in_page); + } + page_idx += 1; + page_start = page_end; + } + } + Ok(()) +} + +/// Maps global row `ranges` onto a column's pages, yielding the page-local +/// sub-ranges for every page any range touches. +/// +/// `ranges` must be ascending and non-overlapping (the precondition +/// `schedule_ranges` relies on). +fn map_ranges_to_pages( + page_row_ends: &[u64], + ranges: &[Range], +) -> Result> { + let mut result: Vec = + Vec::with_capacity(ranges.len().min(page_row_ends.len())); + for_each_page_range(page_row_ends, ranges, |page_idx, range_in_page| { + if let Some(mapping) = result.last_mut() + && mapping.page_idx == page_idx + { + mapping.ranges_in_page.push(range_in_page); + } else { + result.push(PageRangeMapping { + page_idx, + ranges_in_page: vec![range_in_page], + }); + } + })?; + Ok(result) +} + +fn pages_overlapping_ranges(page_row_ends: &[u64], ranges: &[Range]) -> Result> { + let mut result = Vec::with_capacity(ranges.len().min(page_row_ends.len())); + for_each_page_range(page_row_ends, ranges, |page_idx, _| { + if result.last().copied() != Some(page_idx) { + result.push(page_idx); + } + })?; + Ok(result) +} + +#[derive(Debug)] +struct PageCacheMiss { + /// Position of the page in `StructuralPrimitiveFieldScheduler::page_schedulers`, + /// which holds only the built pages and so may differ from `page_index`. + page_idx: usize, + cache_key: PageDataCacheKey, + initialization: PageInitialization, +} + +async fn read_page_initialization_buffers( + io: &Arc, + all_ranges: Vec>, +) -> Result> { + // `submit_request` coalesces (and later un-coalesces) in one forward + // pass, so submit ranges in file order. Page layouts are not required + // to declare their initialization ranges in that order, so retain each + // range's original position and restore it before initializing pages. + let mut sorted_ranges = all_ranges.into_iter().enumerate().collect::>(); + sorted_ranges.sort_unstable_by_key(|(_, range)| (range.start, range.end)); + let sorted_requests = sorted_ranges + .iter() + .map(|(_, range)| range.clone()) + .collect::>(); + let expected_buffer_count = sorted_requests.len(); + let sorted_buffers = if sorted_requests.is_empty() { + Vec::new() + } else { + io.submit_request(sorted_requests, 0).await? + }; + if sorted_buffers.len() != expected_buffer_count { + return Err(Error::invalid_input_source( + format!( + "Page initialization I/O returned {} buffers for {} ranges", + sorted_buffers.len(), + expected_buffer_count + ) + .into(), + )); + } + let mut buffers = vec![None; expected_buffer_count]; + for ((original_idx, _), buffer) in sorted_ranges.into_iter().zip(sorted_buffers) { + buffers[original_idx] = Some(buffer); + } + buffers + .into_iter() + .map(|buffer| { + buffer.ok_or_else(|| Error::internal("Page initialization buffer was not assigned")) + }) + .collect() +} + +async fn cache_initialized_page( + cache: Arc, + cache_key: PageDataCacheKey, + data: Arc, +) { + if data + .clone() + .as_arc_any() + .downcast::() + .is_err() + { + cache + .insert_with_key(&cache_key, Arc::new(CachedPage { data })) + .await; + } +} + +impl StructuralPrimitiveFieldScheduler { + /// Checks the built pages at `positions` (ascending indices into + /// `page_schedulers`) against the cache and returns those still needing + /// initialization I/O. + async fn find_cache_misses( + &mut self, + positions: impl ExactSizeIterator, + cache: &Arc, + ) -> Result> { + // Serve cache hits in place; the rest are misses we must read. + let mut misses = Vec::with_capacity(positions.len()); + for page_idx in positions { + let page = &self.page_schedulers[page_idx]; + if !page.scheduler.needs_initialization() { + continue; + } + let page_index = u32::try_from(page.page_index).map_err(|_| { + Error::invalid_input_source( + format!("Page index does not fit u32: {}", page.page_index).into(), + ) + })?; + let cache_key = PageDataCacheKey { + column_index: self.column_index, + page_index, + view: page.scheduler.cache_view(), + target_data_type_tag: self.target_data_type_tag.clone(), + }; + if let Some(cached_page) = cache.get_with_key(&cache_key).await { + match self.page_schedulers[page_idx] + .scheduler + .try_load(&cached_page.data) + { + Ok(()) => continue, + Err(error) => { + debug!("Replacing invalid cached page data for {cache_key:?}: {error}"); + } + } + } + let initialization = self.page_schedulers[page_idx].scheduler.init_layout()?; + misses.push(PageCacheMiss { + page_idx, + cache_key, + initialization, + }); + } + Ok(misses) + } + + fn initialization_ranges(misses: &[PageCacheMiss]) -> Vec> { + let num_ranges = misses + .iter() + .map(|miss| miss.initialization.num_required_ranges()) + .sum(); + let mut all_ranges = Vec::with_capacity(num_ranges); + for miss in misses { + all_ranges.extend(miss.initialization.required_ranges().cloned()); + } + all_ranges + } + + async fn initialize_cache_misses( + &mut self, + misses: Vec, + buffers: Vec, + cache: Arc, + io: &Arc, + ) -> Result<()> { + let mut buffers = buffers.into_iter(); + + // Initialize each missing page from its slice of the buffers, in file + // order. `init_from_buffers` can perform further I/O (e.g. loading blob + // descriptions), so run pages concurrently; peel disjoint `&mut` page + // borrows off the slice to satisfy the borrow checker without a lookup + // table. + let mut tail = self.page_schedulers.as_mut_slice(); + let mut next_page = 0; + let mut page_inits = FuturesOrdered::new(); + for miss in misses { + let PageCacheMiss { + page_idx, + cache_key, + initialization, + } = miss; + let relative_idx = page_idx + .checked_sub(next_page) + .ok_or_else(|| Error::internal("Page initialization indices are not ordered"))?; + if relative_idx >= tail.len() { + return Err(Error::internal(format!( + "Page initialization index {page_idx} is out of bounds" + ))); + } + let (_, rest) = tail.split_at_mut(relative_idx); + let (page, rest) = rest.split_first_mut().ok_or_else(|| { + Error::internal(format!( + "Page initialization index {page_idx} is out of bounds" + )) + })?; + tail = rest; + next_page = page_idx + 1; + + let count = initialization.num_required_ranges(); + let page_buffers = buffers.by_ref().take(count).collect::>(); + if page_buffers.len() != count { + return Err(Error::invalid_input_source( + format!( + "Page {page_idx} received {} initialization buffers, expected {count}", + page_buffers.len() + ) + .into(), + )); + } + let page_buffers = initialization.with_buffers(page_buffers)?; + let cache = cache.clone(); + let scheduler = &mut page.scheduler; + page_inits.push_back(async move { + let cached = scheduler.init_from_buffers(page_buffers, io).await?; + scheduler.try_load(&cached)?; + cache_initialized_page(cache, cache_key, cached).await; + Ok::<(), Error>(()) + }); + } + if buffers.next().is_some() { + return Err(Error::invalid_input_source( + "Page initialization returned unassigned buffers".into(), + )); + } + + page_inits.try_collect::>().await?; + Ok(()) + } + + async fn initialize_pages( + &mut self, + requested_ranges: Option<&[Range]>, + context: &SchedulerContext, + ) -> Result<()> { + let cache = context.cache().clone(); + let page_indices = match requested_ranges { + None => (0..self.page_row_ends.len()).collect::>(), + Some(ranges) => pages_overlapping_ranges(&self.page_row_ends, ranges)?, + }; + let positions = self.build_pages(&page_indices)?; + let misses = self + .find_cache_misses(positions.into_iter(), &cache) + .await?; + + if misses.is_empty() { + return Ok(()); + } + + // Concatenate the misses' metadata ranges into one request so adjacent + // ranges coalesce into shared GETs instead of one request per page. + let all_ranges = Self::initialization_ranges(&misses); + let buffers = read_page_initialization_buffers(context.io(), all_ranges).await?; + self.initialize_cache_misses(misses, buffers, cache, context.io()) + .await + } +} + +impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler { + fn initialize<'a>( + &'a mut self, + requested_ranges: Option<&'a [Range]>, + _filter: &'a FilterExpression, + context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + self.initialize_pages(requested_ranges, context).boxed() + } + + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + _filter: &FilterExpression, + ) -> Result> { + let ranges = ranges.to_vec(); + Ok(Box::new(StructuralPrimitiveFieldSchedulingJob::try_new( + self, ranges, + )?)) + } +} + +/// Takes the output from several pages decoders and +/// concatenates them. +#[derive(Debug)] +pub struct StructuralCompositeDecodeArrayTask { + tasks: Vec>, + should_validate: bool, + data_type: DataType, +} + +impl StructuralCompositeDecodeArrayTask { + fn restore_validity( + array: Arc, + unraveler: &mut CompositeRepDefUnraveler, + ) -> Result> { + let validity = unraveler.unravel_validity(array.len())?; + let Some(validity) = validity else { + return Ok(array); + }; + if array.data_type() == &DataType::Null { + // We unravel from a null array but we don't add the null buffer because arrow-rs doesn't like it + return Ok(array); + } + if validity.len() != array.len() { + return Err(Error::invalid_input_source( + format!( + "Structural validity has {} entries for an array with {} values", + validity.len(), + array.len() + ) + .into(), + )); + } + // SAFETY: The array buffers have already been validated and the null buffer length + // matches the array. We are only attaching the null buffer here. + Ok(make_array(unsafe { + array + .to_data() + .into_builder() + .nulls(Some(validity)) + .build_unchecked() + })) + } +} + +impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask { + fn decode(self: Box) -> Result { + let mut arrays = Vec::with_capacity(self.tasks.len()); + let mut unravelers = Vec::with_capacity(self.tasks.len()); + let mut data_size = 0u64; + for task in self.tasks { + let decoded = task.decode()?; + data_size += decoded.data.data_size(); + unravelers.push(decoded.repdef); + + let array = make_array( + decoded + .data + .into_arrow(self.data_type.clone(), self.should_validate)?, + ); + + arrays.push(array); + } + let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::>(); + let array = arrow_select::concat::concat(&array_refs)?; + let mut repdef = CompositeRepDefUnraveler::new(unravelers); + + let array = Self::restore_validity(array, &mut repdef)?; + + Ok(DecodedArray { + array, + repdef, + data_size, + }) + } +} + +/// Returns the total bytes consumed by validity bitmaps of nested child fields +/// for `rows` logical rows of `data_type`. +/// +/// Only `FixedSizeList` types have nested children that contribute additional +/// bitmaps; all other fixed-width leaf types return 0. +fn fsl_child_bitmap_bytes(data_type: &arrow_schema::DataType, rows: u64) -> u64 { + match data_type { + arrow_schema::DataType::FixedSizeList(child_field, dimension) => { + let child_rows = rows * *dimension as u64; + let own = if child_field.is_nullable() { + child_rows.div_ceil(8) + } else { + 0 + }; + own + fsl_child_bitmap_bytes(child_field.data_type(), child_rows) + } + _ => 0, + } +} + +#[derive(Debug)] +pub struct StructuralPrimitiveFieldDecoder { + field: Arc, + page_decoders: VecDeque>, + should_validate: bool, + rows_drained_in_current: u64, +} + +impl StructuralPrimitiveFieldDecoder { + pub fn new(field: &Arc, should_validate: bool) -> Self { + Self { + field: field.clone(), + page_decoders: VecDeque::new(), + should_validate, + rows_drained_in_current: 0, + } + } + + fn decoded_bytes_for_rows(&self, rows: u64) -> lance_core::Result { + let mut remaining = rows; + let mut total = 0u64; + for (page_num, page) in self.page_decoders.iter().enumerate() { + if remaining == 0 { + break; + } + let available = if page_num == 0 { + page.num_rows().saturating_sub(self.rows_drained_in_current) + } else { + page.num_rows() + }; + let take = available.min(remaining); + total += page.decoded_bytes(take)?; + remaining -= take; + } + if remaining > 0 { + return Err(lance_core::Error::not_supported(format!( + "plan_decoded_bytes: queued pages cover only {} of {} requested rows", + rows - remaining, + rows, + ))); + } + Ok(total) + } +} + +impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder { + fn accept_page(&mut self, child: LoadedPageShard) -> Result<()> { + assert!(child.path.is_empty()); + self.page_decoders.push_back(child.decoder); + Ok(()) + } + + fn drain(&mut self, num_rows: u64) -> Result> { + let mut remaining = num_rows; + let mut tasks = Vec::new(); + while remaining > 0 { + let queued_pages = self.page_decoders.len(); + let Some(cur_page) = self.page_decoders.front_mut() else { + return Err(Error::internal(format!( + "Primitive decoder missing page decoder while draining field '{}' (data_type={:?}, requested_rows={}, remaining_rows={}, rows_drained_in_current={}, queued_pages={})", + self.field.name(), + self.field.data_type(), + num_rows, + remaining, + self.rows_drained_in_current, + queued_pages + ))); + }; + let num_in_page = cur_page.num_rows() - self.rows_drained_in_current; + let to_take = num_in_page.min(remaining); + + let task = cur_page.drain(to_take)?; + tasks.push(task); + + if to_take == num_in_page { + self.page_decoders.pop_front(); + self.rows_drained_in_current = 0; + } else { + self.rows_drained_in_current += to_take; + } + + remaining -= to_take; + } + Ok(Box::new(StructuralCompositeDecodeArrayTask { + tasks, + should_validate: self.should_validate, + data_type: self.field.data_type().clone(), + })) + } + + fn data_type(&self) -> &DataType { + self.field.data_type() + } + + fn plan_decoded_bytes(&self, rows_remaining: u64) -> lance_core::Result<[u64; 8]> { + use crate::decoder::CANDIDATE_BATCH_SIZES; + let data_type = self.field.data_type(); + let is_nullable = self.field.is_nullable(); + + // Fixed-width: exact from type metadata, no page inspection needed. + if let Some(byte_width) = data_type.byte_width_opt() { + let mut out = [0u64; 8]; + for (i, &c) in CANDIDATE_BATCH_SIZES.iter().enumerate() { + let rows = (c as u64).min(rows_remaining); + out[i] = rows * byte_width as u64; + if is_nullable { + out[i] += rows.div_ceil(8); + } + out[i] += fsl_child_bitmap_bytes(data_type, rows); + } + return Ok(out); + } + + // Boolean: bit-packed, 1 bit per value. + if matches!(data_type, DataType::Boolean) { + let mut out = [0u64; 8]; + for (i, &c) in CANDIDATE_BATCH_SIZES.iter().enumerate() { + let rows = (c as u64).min(rows_remaining); + out[i] = rows.div_ceil(8); + if is_nullable { + out[i] += rows.div_ceil(8); + } + } + return Ok(out); + } + + // Null type: no data bytes. + if matches!(data_type, DataType::Null) { + return Ok([0u64; 8]); + } + + // Variable-width: walk page decoders for exact byte counts. + let mut out = [0u64; 8]; + for (i, &c) in CANDIDATE_BATCH_SIZES.iter().enumerate() { + let rows = (c as u64).min(rows_remaining); + out[i] = self.decoded_bytes_for_rows(rows)?; + if is_nullable { + out[i] += rows.div_ceil(8); + } + } + Ok(out) + } +} + +/// The serialized representation of full-zip data +struct SerializedFullZip { + /// The zipped values buffer + values: LanceBuffer, + /// The repetition index (only present if there is repetition) + repetition_index: Option, +} + +// We align and pad mini-blocks to 8 byte boundaries for two reasons. First, +// to allow us to store a chunk size in 12 bits. +// +// If we directly record the size in bytes with 12 bits we would be limited to +// 4KiB which is too small. Since we know each mini-block consists of 8 byte +// words we can store the # of words instead which gives us 32KiB. +// +// Second, each chunk in a mini-block is aligned to 8 bytes. This allows multi-byte +// values like offsets to be stored in a mini-block and safely read back out. It also +// helps ensure zero-copy reads in cases where zero-copy is possible (e.g. no decoding +// needed). +// +// Note: by "aligned to 8 bytes" we mean BOTH "aligned to 8 bytes from the start of +// the page" and "aligned to 8 bytes from the start of the file." +const MINIBLOCK_ALIGNMENT: usize = 8; + +/// An encoder for primitive (leaf) arrays +/// +/// This encoder is fairly complicated and follows a number of paths depending +/// on the data. +/// +/// First, we convert the validity & offsets information into repetition and +/// definition levels. Then we compress the data itself into a single buffer. +/// +/// If the data is narrow then we encode the data in small chunks (each chunk +/// should be a few disk sectors and contains a buffer of repetition, a buffer +/// of definition, and a buffer of value data). This approach is called +/// "mini-block". These mini-blocks are stored into a single data buffer. +/// +/// If the data is wide then we zip together the repetition and definition value +/// with the value data into a single buffer. This approach is called "zipped". +/// +/// If there is any repetition information then we create a repetition index +/// +/// In addition, the compression process may create zero or more metadata buffers. +/// For example, a dictionary compression will create dictionary metadata. Any +/// mini-block approach has a metadata buffer of block sizes. This metadata is +/// stored in a separate buffer on disk and read at initialization time. +/// +/// TODO: We should concatenate metadata buffers from all pages into a single buffer +/// at (roughly) the end of the file so there is, at most, one read per column of +/// metadata per file. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum MiniblockChunkSize { + U16, + U32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ComplexNullEncoding { + RawLevels, + CompressedLevels, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum FixedWidthDictionaryEncoding { + Exclude64Bit, + Include64Bit, +} + +trait PrimitivePageEncodingBehavior: Send + Sync + Debug { + fn validate_field(&self, _field: &Field, _metadata: &HashMap) -> Result<()> { + Ok(()) + } + + fn try_plan_pages( + &self, + _ctx: &PrimitivePlanContext<'_>, + _arrays: &[ArrayRef], + _normalized: &NormalizedStructuralPlan, + _row_number: u64, + _num_rows: u64, + _num_values: u64, + ) -> Result>> { + Ok(None) + } + + fn try_encode_page( + &self, + _ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + Ok(PrimitiveEncodeAttempt::Unhandled(page)) + } +} + +/// One executable primitive-page behavior selected by an exact file +/// composition. +#[derive(Debug, Clone)] +pub struct PrimitivePageEncoding { + behavior: Arc, +} + +impl PrimitivePageEncoding { + /// Reject an explicit request for sparse structural encoding. + pub fn reject_sparse() -> Self { + Self { + behavior: Arc::new(RejectSparsePrimitiveEncoding), + } + } + + /// Encode constant non-null values as a constant page when applicable. + pub fn constant() -> Self { + Self { + behavior: Arc::new(ConstantPrimitiveEncoding), + } + } + + /// Plan and encode sparse structural pages when applicable. + pub fn sparse(compression: Arc) -> Self { + Self { + behavior: Arc::new(SparsePrimitiveEncoding { compression }), + } + } + + /// Encode dense pages with the original u16 miniblock grammar. + pub fn dense_u16(compression: Arc) -> Self { + Self { + behavior: Arc::new(DenseU16PrimitiveEncoding { compression }), + } + } + + /// Encode dense pages with the u32 miniblock grammar. + pub fn dense_u32(compression: Arc) -> Self { + Self { + behavior: Arc::new(DenseU32PrimitiveEncoding { compression }), + } + } +} + +#[derive(Debug)] +struct RejectSparsePrimitiveEncoding; + +#[derive(Debug)] +struct ConstantPrimitiveEncoding; + +#[derive(Debug)] +struct SparsePrimitiveEncoding { + compression: Arc, +} + +#[derive(Debug)] +struct DenseU16PrimitiveEncoding { + compression: Arc, +} + +#[derive(Debug)] +struct DenseU32PrimitiveEncoding { + compression: Arc, +} + +pub struct PrimitiveStructuralEncoder { + // Accumulates arrays until we have enough data to justify a disk page + accumulation_queue: AccumulationQueue, + + keep_original_array: bool, + accumulated_repdefs: Vec, + page_encodings: Arc<[PrimitivePageEncoding]>, + column_index: u32, + field: Field, + encoding_metadata: Arc>, +} + +struct CompressedLevelsChunk { + data: LanceBuffer, + num_levels: u16, +} + +struct CompressedLevels { + data: Vec, + compression: CompressiveEncoding, + rep_index: Option, +} + +struct SerializedMiniBlockPage { + num_buffers: u64, + data: LanceBuffer, + metadata: LanceBuffer, +} + +#[derive(Debug, Clone, Copy)] +struct DictEncodingBudget { + max_dict_entries: u32, + max_encoded_size: usize, +} + +enum PrimitivePageStructure { + Dense { + repdef: SerializedRepDefs, + single_row_miniblock_repdef_levels: Option, + }, + Sparse { + plan: sparse::SparseStructuralPlan, + prepared_values: Option, + }, +} + +// A primitive page after structural encoding selection and optional dense splitting. +struct PrimitivePageData { + // Arrow leaf arrays that contain this page's visible values. + arrays: Vec, + // Structural representation aligned to this page. + structure: PrimitivePageStructure, + // Top-level row number of the first row in this page. + row_number: u64, + // Number of top-level rows in this page. + num_rows: u64, +} + +struct PrimitivePlanContext<'a> { + column_idx: u32, + field: &'a Field, + encoding_metadata: &'a HashMap, +} + +enum PrimitiveEncodeAttempt { + Encoded(EncodedPage), + Unhandled(PrimitivePageData), +} + +// Immutable encoder state shared by per-page encode tasks. +// +// Cloning this only clones Arc-backed configuration and field metadata. Page data +// stays in PrimitivePageData and is moved into exactly one task. +#[derive(Clone)] +struct PrimitiveEncodeContext { + // Column being encoded. + column_idx: u32, + field: Field, + encoding_metadata: Arc>, + is_simple_validity: bool, + has_repdef_info: bool, +} + +impl PrimitiveStructuralEncoder { + pub fn try_new( + options: &EncodingOptions, + page_encodings: Arc<[PrimitivePageEncoding]>, + column_index: u32, + field: Field, + encoding_metadata: Arc>, + ) -> Result { + for page_encoding in page_encodings.iter() { + page_encoding + .behavior + .validate_field(&field, &encoding_metadata)?; + } + Ok(Self { + accumulation_queue: AccumulationQueue::new( + options.cache_bytes_per_column, + column_index, + options.keep_original_array, + ), + keep_original_array: options.keep_original_array, + accumulated_repdefs: Vec::new(), + column_index, + page_encodings, + field, + encoding_metadata, + }) + } + + fn encode_page( + page_encodings: &[PrimitivePageEncoding], + ctx: &PrimitiveEncodeContext, + mut page: PrimitivePageData, + ) -> Result { + for page_encoding in page_encodings { + match page_encoding.behavior.try_encode_page(ctx, page)? { + PrimitiveEncodeAttempt::Encoded(page) => return Ok(page), + PrimitiveEncodeAttempt::Unhandled(unhandled) => page = unhandled, + } + } + Err(Error::invalid_input_source( + format!( + "No primitive page encoding atom supports field '{}'", + ctx.field.name + ) + .into(), + )) + } + + // TODO: This is a heuristic we may need to tune at some point + // + // As data gets narrow then the "zipping" process gets too expensive + // and we prefer mini-block + // As data gets wide then the # of values per block shrinks (very wide) + // data doesn't even fit in a mini-block and the block overhead gets + // too large and we prefer zipped. + fn is_narrow(data_block: &DataBlock) -> bool { + const MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE: u64 = 256; + + if let Some(max_len_array) = data_block.get_stat(Stat::MaxLength) { + let max_len_array = max_len_array + .as_any() + .downcast_ref::>() + .unwrap(); + if max_len_array.value(0) < MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE { + return true; + } + } + false + } + + fn prefers_miniblock( + data_block: &DataBlock, + encoding_metadata: &HashMap, + ) -> bool { + // If the user specifically requested miniblock then use it + if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) { + return user_requested.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK; + } + // Otherwise only use miniblock if it is narrow + Self::is_narrow(data_block) + } + + fn prefers_fullzip(encoding_metadata: &HashMap) -> bool { + // Fullzip is the backup option so the only reason we wouldn't use it is if the + // user specifically requested not to use it (in which case we're probably going + // to emit an error) + if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) { + return user_requested.to_lowercase() == STRUCTURAL_ENCODING_FULLZIP; + } + true + } + + // Converts value data, repetition levels, and definition levels into a single + // buffer of mini-blocks. In addition, creates a buffer of mini-block metadata + // which tells us the size of each block. Finally, if repetition is present then + // we also create a buffer for the repetition index. + // + // Each chunk is serialized as: + // | num_bufs (1 byte) | buf_lens (2 bytes per buffer) | P | buf0 | P | buf1 | ... | bufN | P | + // + // P - Padding inserted to ensure each buffer is 8-byte aligned and the buffer size is a multiple + // of 8 bytes (so that the next chunk is 8-byte aligned). + // + // Each block has a u16 word of metadata. The upper 12 bits contain the + // # of 8-byte words in the block (if the block does not fill the final word + // then up to 7 bytes of padding are added). The lower 4 bits describe the log_2 + // number of values (e.g. if there are 1024 then the lower 4 bits will be + // 0xA) All blocks except the last must have power-of-two number of values. + // This not only makes metadata smaller but it makes decoding easier since + // batch sizes are typically a power of 2. 4 bits would allow us to express + // up to 32Ki values. + // + // This means blocks can have 1 to 32Ki values and 8 - 32Ki bytes. + // + // All metadata words are serialized (as little endian) into a single buffer + // of metadata values. + // + // If there is repetition then we also create a repetition index. This is a + // single buffer of integer vectors (stored in row major order). There is one + // entry for each chunk. The size of the vector is based on the depth of random + // access we want to support. + // + // A vector of size 2 is the minimum and will support row-based random access (e.g. + // "take the 57th row"). A vector of size 3 will support 1 level of nested access + // (e.g. "take the 3rd item in the 57th row"). A vector of size 4 will support 2 + // levels of nested access and so on. + // + // The first number in the vector is the number of top-level rows that complete in + // the chunk. The second number is the number of second-level rows that complete + // after the final top-level row completed (or beginning of the chunk if no top-level + // row completes in the chunk). And so on. The final number in the vector is always + // the number of leftover items not covered by earlier entries in the vector. + // + // Currently we are limited to 0 levels of nested access but that will change in the + // future. + // + // The repetition index and the chunk metadata are read at initialization time and + // cached in memory. + fn serialize_miniblocks( + miniblocks: MiniBlockCompressed, + rep: Option>, + def: Option>, + miniblock_chunk_size: MiniblockChunkSize, + ) -> Result { + let bytes_rep = rep + .as_ref() + .map(|rep| rep.iter().map(|r| r.data.len()).sum::()) + .unwrap_or(0); + let bytes_def = def + .as_ref() + .map(|def| def.iter().map(|d| d.data.len()).sum::()) + .unwrap_or(0); + let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::(); + let mut num_buffers = miniblocks.data.len(); + if rep.is_some() { + num_buffers += 1; + } + if def.is_some() { + num_buffers += 1; + } + // 2 bytes for the length of each buffer and up to 7 bytes of padding per buffer + let max_extra = 9 * num_buffers; + let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra); + let chunk_size_bytes = match miniblock_chunk_size { + MiniblockChunkSize::U16 => 2, + MiniblockChunkSize::U32 => 4, + }; + let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes); + + let mut rep_iter = rep.map(|r| r.into_iter()); + let mut def_iter = def.map(|d| d.into_iter()); + + let mut buffer_offsets = vec![0; miniblocks.data.len()]; + for chunk in miniblocks.chunks { + let start_pos = data_buffer.len(); + // Start of chunk should be aligned + debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0); + + let rep = rep_iter.as_mut().map(|r| r.next().unwrap()); + let def = def_iter.as_mut().map(|d| d.next().unwrap()); + + // Write the number of levels, or 0 if there is no rep/def + let num_levels = rep + .as_ref() + .map(|r| r.num_levels) + .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0)); + data_buffer.extend_from_slice(&num_levels.to_le_bytes()); + + // Write the buffer lengths + if let Some(rep) = rep.as_ref() { + let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| { + Error::internal(format!( + "Repetition buffer size ({} bytes) too large", + rep.data.len() + )) + })?; + data_buffer.extend_from_slice(&bytes_rep.to_le_bytes()); + } + if let Some(def) = def.as_ref() { + let bytes_def = u16::try_from(def.data.len()).map_err(|_| { + Error::internal(format!( + "Definition buffer size ({} bytes) too large", + def.data.len() + )) + })?; + data_buffer.extend_from_slice(&bytes_def.to_le_bytes()); + } + + if miniblock_chunk_size == MiniblockChunkSize::U32 { + for &buffer_size in &chunk.buffer_sizes { + data_buffer.extend_from_slice(&buffer_size.to_le_bytes()); + } + } else { + for &buffer_size in &chunk.buffer_sizes { + let buffer_size = u16::try_from(buffer_size).map_err(|_| { + Error::internal(format!( + "Mini-block buffer size ({} bytes) too large for 16-bit metadata", + buffer_size + )) + })?; + data_buffer.extend_from_slice(&buffer_size.to_le_bytes()); + } + } + + // Pad + let add_padding = |data_buffer: &mut Vec| { + let pad = pad_bytes::(data_buffer.len()); + data_buffer.extend(iter::repeat_n(FILL_BYTE, pad)); + }; + add_padding(&mut data_buffer); + + // Write the buffers themselves + if let Some(rep) = rep.as_ref() { + data_buffer.extend_from_slice(&rep.data); + add_padding(&mut data_buffer); + } + if let Some(def) = def.as_ref() { + data_buffer.extend_from_slice(&def.data); + add_padding(&mut data_buffer); + } + for (buffer_size, (buffer, buffer_offset)) in chunk + .buffer_sizes + .iter() + .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut())) + { + let start = *buffer_offset; + let end = start + *buffer_size as usize; + *buffer_offset += *buffer_size as usize; + data_buffer.extend_from_slice(&buffer[start..end]); + add_padding(&mut data_buffer); + } + + let chunk_bytes = data_buffer.len() - start_pos; + let max_chunk_size = match miniblock_chunk_size { + MiniblockChunkSize::U16 => 32 * 1024, + MiniblockChunkSize::U32 => 1_u64 << 31, + }; + if chunk_bytes == 0 || chunk_bytes as u64 > max_chunk_size { + return Err(Error::internal(format!( + "Mini-block chunk size {} bytes exceeds the {} byte metadata limit", + chunk_bytes, max_chunk_size + ))); + } + if chunk_bytes % MINIBLOCK_ALIGNMENT != 0 { + return Err(Error::internal(format!( + "Mini-block chunk size {} bytes is not aligned to {} bytes", + chunk_bytes, MINIBLOCK_ALIGNMENT + ))); + } + if chunk.log_num_values > 15 { + return Err(Error::internal(format!( + "Mini-block log_num_values {} exceeds the 4-bit metadata limit", + chunk.log_num_values + ))); + } + // We subtract 1 here from chunk_bytes because we want to be able to express + // a size of 32KiB and not (32Ki - 8)B which is what we'd get otherwise with + // 0xFFF + let divided_bytes = chunk_bytes / MINIBLOCK_ALIGNMENT; + let divided_bytes_minus_one = (divided_bytes - 1) as u64; + + let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64; + if miniblock_chunk_size == MiniblockChunkSize::U32 { + meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes()); + } else { + meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes()); + } + } + + let data_buffer = LanceBuffer::from(data_buffer); + let metadata_buffer = LanceBuffer::from(meta_buffer); + + Ok(SerializedMiniBlockPage { + num_buffers: miniblocks.data.len() as u64, + data: data_buffer, + metadata: metadata_buffer, + }) + } + + /// Compresses a buffer of levels into chunks + /// + /// If these are repetition levels then we also calculate the repetition index here (that + /// is the third return value) + fn compress_levels( + mut levels: RepDefSlicer<'_>, + num_elements: u64, + compression_strategy: &dyn CompressionStrategy, + chunks: &[MiniBlockChunk], + // This will be 0 if we are compressing def levels + max_rep: u16, + ) -> Result { + let mut rep_index = if max_rep > 0 { + Vec::with_capacity(chunks.len()) + } else { + vec![] + }; + // Make the levels into a FixedWidth data block + let num_levels = levels.num_levels() as u64; + let levels_buf = levels.all_levels().clone(); + + let mut fixed_width_block = FixedWidthDataBlock { + data: levels_buf, + bits_per_value: 16, + num_values: num_levels, + block_info: BlockInfo::new(), + }; + // Compute statistics to enable optimal compression for rep/def levels + fixed_width_block.compute_stat(); + + let levels_block = DataBlock::FixedWidth(fixed_width_block); + let levels_field = Field::new_arrow("", DataType::UInt16, false)?; + // Pick a block compressor + let compressor = + compression_strategy.create_block_compressor(&levels_field, &levels_block)?; + let mut compressor_desc = None; + // Compress blocks of levels (sized according to the chunks) + let mut level_chunks = Vec::with_capacity(chunks.len()); + let mut values_counter = 0; + for (chunk_idx, chunk) in chunks.iter().enumerate() { + let chunk_num_values = chunk.num_values(values_counter, num_elements); + debug_assert!(chunk_num_values > 0); + values_counter += chunk_num_values; + let chunk_levels = if chunk_idx < chunks.len() - 1 { + levels.slice_next(chunk_num_values as usize) + } else { + levels.slice_rest() + }; + let num_chunk_levels = (chunk_levels.len() / 2) as u64; + if max_rep > 0 { + // If max_rep > 0 then we are working with rep levels and we need + // to calculate the repetition index. The repetition index for a + // chunk is currently 2 values (in the future it may be more). + // + // The first value is the number of rows that _finish_ in the + // chunk. + // + // The second value is the number of "leftovers" after the last + // finished row in the chunk. + let rep_values = chunk_levels.borrow_to_typed_slice::(); + let rep_values = rep_values.as_ref(); + + // We skip 1 here because a max_rep at spot 0 doesn't count as a finished list (we + // will count it in the previous chunk) + let mut num_rows = rep_values.iter().skip(1).filter(|v| **v == max_rep).count(); + let num_leftovers = if chunk_idx < chunks.len() - 1 { + rep_values + .iter() + .rev() + .position(|v| *v == max_rep) + // # of leftovers includes the max_rep spot + .map(|pos| pos + 1) + .unwrap_or(rep_values.len()) + } else { + // Last chunk can't have leftovers + 0 + }; + + if chunk_idx != 0 && rep_values.first() == Some(&max_rep) { + // This chunk starts with a new row and so, if we thought we had leftovers + // in the previous chunk, we were mistaken + // TODO: Can use unchecked here + let rep_len = rep_index.len(); + if rep_index[rep_len - 1] != 0 { + // We thought we had leftovers but that was actually a full row + rep_index[rep_len - 2] += 1; + rep_index[rep_len - 1] = 0; + } + } + + if chunk_idx == chunks.len() - 1 { + // The final list + num_rows += 1; + } + rep_index.push(num_rows as u64); + rep_index.push(num_leftovers as u64); + } + let mut chunk_fixed_width = FixedWidthDataBlock { + data: chunk_levels, + bits_per_value: 16, + num_values: num_chunk_levels, + block_info: BlockInfo::new(), + }; + chunk_fixed_width.compute_stat(); + let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width); + let (compressed_levels, chunk_compressor_desc) = + compressor.compress(chunk_levels_block)?; + if let Some(compressor_desc) = compressor_desc.as_ref() { + if compressor_desc != &chunk_compressor_desc { + return Err(Error::internal( + "Rep/def block compressor changed encoding between chunks".to_string(), + )); + } + } else { + compressor_desc = Some(chunk_compressor_desc); + } + let compressed_levels = compressed_levels.ok_or_else(|| { + Error::internal( + "Rep/def block compressor selected a metadata-only codec".to_string(), + ) + })?; + let num_levels = u16::try_from(num_chunk_levels).map_err(|_| { + Error::invalid_input_source( + format!( + "Mini-block cannot encode {} rep/def levels in one chunk. \ + This usually means a top-level row contains too much nested structure \ + for the current layout.", + num_chunk_levels + ) + .into(), + ) + })?; + level_chunks.push(CompressedLevelsChunk { + data: compressed_levels, + num_levels, + }); + } + debug_assert_eq!(levels.num_levels_remaining(), 0); + let rep_index = if rep_index.is_empty() { + None + } else { + Some(LanceBuffer::reinterpret_vec(rep_index)) + }; + Ok(CompressedLevels { + data: level_chunks, + compression: compressor_desc.ok_or_else(|| { + Error::internal("Rep/def compression produced no chunks".to_string()) + })?, + rep_index, + }) + } + + fn encode_simple_all_null( + column_idx: u32, + num_rows: u64, + row_number: u64, + ) -> Result { + let description = + ProtobufUtils21::constant_layout(&[DefinitionInterpretation::NullableItem], None); + Ok(EncodedPage { + column_idx, + data: vec![], + description: PageEncoding::Structural(description), + num_rows, + row_number, + }) + } + + fn encode_complex_all_null_vals( + data: &Arc<[u16]>, + compression_strategy: &dyn CompressionStrategy, + ) -> Result<(LanceBuffer, pb21::CompressiveEncoding)> { + let buffer = LanceBuffer::reinterpret_slice(data.clone()); + let mut fixed_width_block = FixedWidthDataBlock { + data: buffer, + bits_per_value: 16, + num_values: data.len() as u64, + block_info: BlockInfo::new(), + }; + fixed_width_block.compute_stat(); + + let levels_block = DataBlock::FixedWidth(fixed_width_block); + let levels_field = Field::new_arrow("", DataType::UInt16, false)?; + let (compressed_buffer, encoding) = + compress_required_block(compression_strategy, &levels_field, levels_block)?; + Ok((compressed_buffer, encoding)) + } + + // Encodes a page where all values are null but we have rep/def + // information that we need to store (e.g. to distinguish between + // different kinds of null) + fn encode_complex_all_null( + column_idx: u32, + repdef: crate::repdef::SerializedRepDefs, + row_number: u64, + num_rows: u64, + complex_null_encoding: ComplexNullEncoding, + compression_strategy: &dyn CompressionStrategy, + ) -> Result { + if complex_null_encoding == ComplexNullEncoding::RawLevels { + let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() { + LanceBuffer::reinterpret_slice(rep.clone()) + } else { + LanceBuffer::empty() + }; + + let def_bytes = if let Some(def) = repdef.definition_levels.as_ref() { + LanceBuffer::reinterpret_slice(def.clone()) + } else { + LanceBuffer::empty() + }; + + let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, None); + return Ok(EncodedPage { + column_idx, + data: vec![rep_bytes, def_bytes], + description: PageEncoding::Structural(description), + num_rows, + row_number, + }); + } + + let (rep_bytes, rep_encoding, num_rep_values) = if let Some(rep) = + repdef.repetition_levels.as_ref() + { + let num_values = rep.len() as u64; + let (buffer, encoding) = Self::encode_complex_all_null_vals(rep, compression_strategy)?; + (buffer, Some(encoding), num_values) + } else { + (LanceBuffer::empty(), None, 0) + }; + + let (def_bytes, def_encoding, num_def_values) = if let Some(def) = + repdef.definition_levels.as_ref() + { + let num_values = def.len() as u64; + let (buffer, encoding) = Self::encode_complex_all_null_vals(def, compression_strategy)?; + (buffer, Some(encoding), num_values) + } else { + (LanceBuffer::empty(), None, 0) + }; + + let description = ProtobufUtils21::compressed_all_null_constant_layout( + &repdef.def_meaning, + rep_encoding, + def_encoding, + num_rep_values, + num_def_values, + ); + Ok(EncodedPage { + column_idx, + data: vec![rep_bytes, def_bytes], + description: PageEncoding::Structural(description), + num_rows, + row_number, + }) + } + + fn leaf_validity( + repdef: &crate::repdef::SerializedRepDefs, + num_values: usize, + ) -> Result> { + let rep = repdef + .repetition_levels + .as_ref() + .map(|rep| rep.as_ref().to_vec()); + let def = repdef + .definition_levels + .as_ref() + .map(|def| def.as_ref().to_vec()); + let mut unraveler = RepDefUnraveler::new( + rep, + def, + repdef.def_meaning.clone().into(), + num_values as u64, + ); + if unraveler.is_all_valid() { + return Ok(None); + } + let mut validity = BooleanBufferBuilder::new(num_values); + unraveler.unravel_validity(&mut validity)?; + Ok(Some(validity.finish())) + } + + fn is_constant_values( + arrays: &[ArrayRef], + scalar: &ArrayRef, + validity: Option<&BooleanBuffer>, + ) -> Result { + debug_assert_eq!(scalar.len(), 1); + debug_assert_eq!(scalar.null_count(), 0); + + match scalar.data_type() { + DataType::Boolean => { + let mut global_idx = 0usize; + let scalar_val = scalar.as_boolean().value(0); + for arr in arrays { + let bool_arr = arr.as_boolean(); + for i in 0..arr.len() { + let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true); + global_idx += 1; + if !is_valid { + continue; + } + if bool_arr.value(i) != scalar_val { + return Ok(false); + } + } + } + Ok(true) + } + DataType::Utf8 => Self::is_constant_utf8::(arrays, scalar, validity), + DataType::LargeUtf8 => Self::is_constant_utf8::(arrays, scalar, validity), + DataType::Binary => Self::is_constant_binary::(arrays, scalar, validity), + DataType::LargeBinary => Self::is_constant_binary::(arrays, scalar, validity), + data_type => { + let mut global_idx = 0usize; + let Some(byte_width) = data_type.byte_width_opt() else { + return Ok(false); + }; + let scalar_data = scalar.to_data(); + if scalar_data.buffers().len() != 1 || !scalar_data.child_data().is_empty() { + return Ok(false); + } + let scalar_bytes = scalar_data.buffers()[0].as_slice(); + if scalar_bytes.len() != byte_width { + return Ok(false); + } + + for arr in arrays { + let data = arr.to_data(); + if data.buffers().is_empty() { + return Ok(false); + } + let buf = data.buffers()[0].as_slice(); + let base = data.offset(); + for i in 0..arr.len() { + let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true); + global_idx += 1; + if !is_valid { + continue; + } + let start = (base + i) * byte_width; + if buf[start..start + byte_width] != scalar_bytes[..] { + return Ok(false); + } + } + } + Ok(true) + } + } + } + + fn is_constant_utf8( + arrays: &[ArrayRef], + scalar: &ArrayRef, + validity: Option<&BooleanBuffer>, + ) -> Result { + debug_assert_eq!(scalar.len(), 1); + let scalar_val = scalar.as_string::().value(0).as_bytes(); + let mut global_idx = 0usize; + for arr in arrays { + let str_arr = arr.as_string::(); + for i in 0..arr.len() { + let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true); + global_idx += 1; + if !is_valid { + continue; + } + if str_arr.value(i).as_bytes() != scalar_val { + return Ok(false); + } + } + } + Ok(true) + } + + fn is_constant_binary( + arrays: &[ArrayRef], + scalar: &ArrayRef, + validity: Option<&BooleanBuffer>, + ) -> Result { + debug_assert_eq!(scalar.len(), 1); + let scalar_val = scalar.as_binary::().value(0); + let mut global_idx = 0usize; + for arr in arrays { + let bin_arr = arr.as_binary::(); + for i in 0..arr.len() { + let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true); + global_idx += 1; + if !is_valid { + continue; + } + if bin_arr.value(i) != scalar_val { + return Ok(false); + } + } + } + Ok(true) + } + + fn find_constant_scalar( + arrays: &[ArrayRef], + validity: Option<&BooleanBuffer>, + ) -> Result> { + if arrays.is_empty() { + return Ok(None); + } + + let global_scalar_idx = if let Some(validity) = validity { + let Some(idx) = (0..validity.len()).find(|&i| validity.value(i)) else { + return Ok(None); + }; + idx + } else { + 0 + }; + + let mut idx_remaining = global_scalar_idx; + let mut scalar_arr_idx = 0usize; + while scalar_arr_idx < arrays.len() { + let len = arrays[scalar_arr_idx].len(); + if idx_remaining < len { + break; + } + idx_remaining -= len; + scalar_arr_idx += 1; + } + + if scalar_arr_idx >= arrays.len() { + return Ok(None); + } + + let scalar = + lance_arrow::scalar::extract_scalar_value(&arrays[scalar_arr_idx], idx_remaining)?; + if scalar.null_count() != 0 { + return Ok(None); + } + if !Self::is_constant_values(arrays, &scalar, validity)? { + return Ok(None); + } + Ok(Some(scalar)) + } + + fn resolve_dict_values_compression_metadata( + field_metadata: &HashMap, + env_compression: Option, + env_compression_level: Option, + ) -> HashMap { + let mut metadata = HashMap::new(); + + let compression = field_metadata + .get(DICT_VALUES_COMPRESSION_META_KEY) + .cloned() + .or(env_compression) + .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string()); + metadata.insert(COMPRESSION_META_KEY.to_string(), compression); + + if let Some(compression_level) = field_metadata + .get(DICT_VALUES_COMPRESSION_LEVEL_META_KEY) + .cloned() + .or(env_compression_level) + { + metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), compression_level); + } + + metadata + } + + fn build_dict_values_compressor_field(field: &Field) -> Result { + // This is an internal synthetic field used only to feed metadata into + // `create_block_compressor` for dictionary values. The concrete type/name here + // are not semantically meaningful; we rely on explicit metadata below to control + // general compression selection for dictionary values. + let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?; + dict_values_field.metadata = Self::resolve_dict_values_compression_metadata( + &field.metadata, + env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(), + env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(), + ); + Ok(dict_values_field) + } + + #[allow(clippy::too_many_arguments)] + fn encode_miniblock( + column_idx: u32, + field: &Field, + compression_strategy: &dyn CompressionStrategy, + data: DataBlock, + repdef: crate::repdef::SerializedRepDefs, + row_number: u64, + dictionary_data: Option, + num_rows: u64, + miniblock_chunk_size: MiniblockChunkSize, + ) -> Result { + if let DataBlock::AllNull(_null_block) = data { + // We should not be using mini-block for all-null. There are other structural + // encodings for that. + unreachable!() + } + + let num_items = data.num_values(); + + let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; + let common_chunk_buffers = + u64::from(repdef.rep_slicer().is_some()) + u64::from(repdef.def_slicer().is_some()); + let support_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32; + let compression_context = + MiniBlockCompressionContext::new(common_chunk_buffers, support_large_chunk, true); + let (compressed_data, value_encoding) = compressor.compress(compression_context, data)?; + + let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16; + + let mut compressed_rep = repdef + .rep_slicer() + .map(|rep_slicer| { + Self::compress_levels( + rep_slicer, + num_items, + compression_strategy, + &compressed_data.chunks, + max_rep, + ) + }) + .transpose()?; + + let (rep_index, rep_index_depth) = + match compressed_rep.as_mut().and_then(|cr| cr.rep_index.as_mut()) { + Some(rep_index) => (Some(rep_index.clone()), 1), + None => (None, 0), + }; + + let mut compressed_def = repdef + .def_slicer() + .map(|def_slicer| { + Self::compress_levels( + def_slicer, + num_items, + compression_strategy, + &compressed_data.chunks, + /*max_rep=*/ 0, + ) + }) + .transpose()?; + + // TODO: Parquet sparsely encodes values here. We could do the same but + // then we won't have log2 values per chunk. This means more metadata + // and potentially more decoder asymmetry. However, it may be worth + // investigating at some point + + let rep_data = compressed_rep + .as_mut() + .map(|cr| std::mem::take(&mut cr.data)); + let def_data = compressed_def + .as_mut() + .map(|cd| std::mem::take(&mut cd.data)); + + let serialized = + Self::serialize_miniblocks(compressed_data, rep_data, def_data, miniblock_chunk_size)?; + let has_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32; + + // Metadata, Data, Dictionary, (maybe) Repetition Index + let mut data = Vec::with_capacity(4); + data.push(serialized.metadata); + data.push(serialized.data); + + if let Some(dictionary_data) = dictionary_data { + let num_dictionary_items = dictionary_data.num_values(); + let dict_values_field = Self::build_dict_values_compressor_field(field)?; + + let (dictionary_buffer, dictionary_encoding) = + compress_required_block(compression_strategy, &dict_values_field, dictionary_data)?; + + data.push(dictionary_buffer); + if let Some(rep_index) = rep_index { + data.push(rep_index); + } + + let description = ProtobufUtils21::miniblock_layout( + compressed_rep.map(|cr| cr.compression), + compressed_def.map(|cd| cd.compression), + value_encoding, + rep_index_depth, + serialized.num_buffers, + Some((dictionary_encoding, num_dictionary_items)), + &repdef.def_meaning, + num_items, + has_large_chunk, + ); + Ok(EncodedPage { + num_rows, + column_idx, + data, + description: PageEncoding::Structural(description), + row_number, + }) + } else { + let description = ProtobufUtils21::miniblock_layout( + compressed_rep.map(|cr| cr.compression), + compressed_def.map(|cd| cd.compression), + value_encoding, + rep_index_depth, + serialized.num_buffers, + None, + &repdef.def_meaning, + num_items, + has_large_chunk, + ); + + if let Some(rep_index) = rep_index { + let view = rep_index.borrow_to_typed_slice::(); + let total = view.chunks_exact(2).map(|c| c[0]).sum::(); + debug_assert_eq!(total, num_rows); + + data.push(rep_index); + } + + Ok(EncodedPage { + num_rows, + column_idx, + data, + description: PageEncoding::Structural(description), + row_number, + }) + } + } + + // For fixed-size data we encode < control word | data > for each value + fn serialize_full_zip_fixed( + fixed: FixedWidthDataBlock, + mut repdef: ControlWordIterator, + num_values: u64, + ) -> Result { + if !fixed.bits_per_value.is_multiple_of(8) { + return Err(Error::invalid_input_source( + format!( + "Full-zip fixed-width values must be byte aligned, got {} bits per value", + fixed.bits_per_value + ) + .into(), + )); + } + + let len = fixed.data.len() + repdef.bytes_per_word() * num_values as usize; + let mut zipped_data = Vec::with_capacity(len); + + let max_rep_index_val = if repdef.has_repetition() { + len as u64 + } else { + // Setting this to 0 means we won't write a repetition index + 0 + }; + let mut rep_index_builder = + BytepackedIntegerEncoder::with_capacity(num_values as usize + 1, max_rep_index_val); + + let bytes_per_value = fixed.bits_per_value as usize / 8; + let mut offset = 0; + + if bytes_per_value == 0 { + // No data, just dump the repdef into the buffer + while let Some(control) = repdef.append_next(&mut zipped_data) { + if control.is_new_row { + // We have finished a row + debug_assert!(offset <= len); + rep_index_builder.append_trusted(offset as u64); + } + offset = zipped_data.len(); + } + } else { + // We have data, zip it with the repdef + let mut data_iter = fixed.data.chunks_exact(bytes_per_value); + while let Some(control) = repdef.append_next(&mut zipped_data) { + if control.is_new_row { + // We have finished a row + debug_assert!(offset <= len); + rep_index_builder.append_trusted(offset as u64); + } + if control.is_visible { + let value = data_iter.next().unwrap(); + zipped_data.extend_from_slice(value); + } + offset = zipped_data.len(); + } + } + + debug_assert_eq!(zipped_data.len(), len); + // Put the final value in the rep index + rep_index_builder.append_trusted(zipped_data.len() as u64); + + let zipped_data = LanceBuffer::from(zipped_data); + let rep_index = rep_index_builder.into_data(); + let rep_index = if rep_index.is_empty() { + None + } else { + Some(LanceBuffer::from(rep_index)) + }; + Ok(SerializedFullZip { + values: zipped_data, + repetition_index: rep_index, + }) + } + + // For variable-size data we encode < control word | length | data > for each value + // + // In addition, we create a second buffer, the repetition index + fn serialize_full_zip_variable( + variable: VariableWidthBlock, + mut repdef: ControlWordIterator, + num_items: u64, + ) -> Result { + let bytes_per_offset = variable.bits_per_offset as usize / 8; + if !variable.bits_per_offset.is_multiple_of(8) { + return Err(Error::invalid_input_source( + format!( + "Full-zip variable-width offsets must be byte aligned, got {} bits per offset", + variable.bits_per_offset + ) + .into(), + )); + } + let len = variable.data.len() + + repdef.bytes_per_word() * num_items as usize + + bytes_per_offset * variable.num_values as usize; + let mut buf = Vec::with_capacity(len); + + let max_rep_index_val = len as u64; + let mut rep_index_builder = + BytepackedIntegerEncoder::with_capacity(num_items as usize + 1, max_rep_index_val); + + // TODO: byte pack the item lengths with varint encoding + match bytes_per_offset { + 4 => { + let offs = variable.offsets.borrow_to_typed_slice::(); + let mut rep_offset = 0; + let mut windows_iter = offs.as_ref().windows(2); + while let Some(control) = repdef.append_next(&mut buf) { + if control.is_new_row { + // We have finished a row + debug_assert!(rep_offset <= len); + rep_index_builder.append_trusted(rep_offset as u64); + } + if control.is_visible { + let window = windows_iter.next().unwrap(); + if control.is_valid_item { + buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes()); + buf.extend_from_slice( + &variable.data[window[0] as usize..window[1] as usize], + ); + } + } + rep_offset = buf.len(); + } + } + 8 => { + let offs = variable.offsets.borrow_to_typed_slice::(); + let mut rep_offset = 0; + let mut windows_iter = offs.as_ref().windows(2); + while let Some(control) = repdef.append_next(&mut buf) { + if control.is_new_row { + // We have finished a row + debug_assert!(rep_offset <= len); + rep_index_builder.append_trusted(rep_offset as u64); + } + if control.is_visible { + let window = windows_iter.next().unwrap(); + if control.is_valid_item { + buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes()); + buf.extend_from_slice( + &variable.data[window[0] as usize..window[1] as usize], + ); + } + } + rep_offset = buf.len(); + } + } + _ => { + return Err(Error::invalid_input_source( + format!( + "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits", + variable.bits_per_offset + ) + .into(), + )); + } + } + + // We might have saved a few bytes by not copying lengths when the length was zero. However, + // if we are over `len` then we have a bug. + debug_assert!(buf.len() <= len); + // Put the final value in the rep index + rep_index_builder.append_trusted(buf.len() as u64); + + let zipped_data = LanceBuffer::from(buf); + let rep_index = rep_index_builder.into_data(); + debug_assert!(!rep_index.is_empty()); + let rep_index = Some(LanceBuffer::from(rep_index)); + Ok(SerializedFullZip { + values: zipped_data, + repetition_index: rep_index, + }) + } + + /// Serializes data into a single buffer according to the full-zip format which zips + /// together the repetition, definition, and value data into a single buffer. + fn serialize_full_zip( + compressed_data: PerValueDataBlock, + repdef: ControlWordIterator, + num_items: u64, + ) -> Result { + match compressed_data { + PerValueDataBlock::Fixed(fixed) => { + Self::serialize_full_zip_fixed(fixed, repdef, num_items) + } + PerValueDataBlock::Variable(var) => { + Self::serialize_full_zip_variable(var, repdef, num_items) + } + } + } + + fn expand_boolean_to_bytes(fixed: FixedWidthDataBlock) -> FixedWidthDataBlock { + debug_assert_eq!(fixed.bits_per_value, 1); + let num_values = fixed.num_values as usize; + let bool_buf = BooleanBuffer::new(fixed.data.into_buffer(), 0, num_values); + let expanded: Vec = (0..num_values).map(|i| bool_buf.value(i) as u8).collect(); + FixedWidthDataBlock { + data: LanceBuffer::from(expanded), + bits_per_value: 8, + num_values: fixed.num_values, + block_info: BlockInfo::new(), + } + } + + fn encode_full_zip( + column_idx: u32, + field: &Field, + compression_strategy: &dyn CompressionStrategy, + data: DataBlock, + repdef: crate::repdef::SerializedRepDefs, + row_number: u64, + num_lists: u64, + ) -> Result { + let max_rep = repdef + .repetition_levels + .as_ref() + .map_or(0, |r| r.iter().max().copied().unwrap_or(0)); + let max_def = repdef + .definition_levels + .as_ref() + .map_or(0, |d| d.iter().max().copied().unwrap_or(0)); + + // To handle FSL we just flatten + // let data = data.flatten(); + + let (num_items, num_visible_items) = + if let Some(rep_levels) = repdef.repetition_levels.as_ref() { + // If there are rep levels there may be "invisible" items and we need to encode + // rep_levels.len() things which might be larger than data.num_values() + (rep_levels.len() as u64, data.num_values()) + } else { + // If there are no rep levels then we encode data.num_values() things + (data.num_values(), data.num_values()) + }; + + let max_visible_def = repdef.max_visible_level.unwrap_or(u16::MAX); + + let repdef_iter = build_control_word_iterator( + repdef.repetition_levels.as_deref(), + max_rep, + repdef.definition_levels.as_deref(), + max_def, + max_visible_def, + num_items as usize, + ); + let bits_rep = repdef_iter.bits_rep(); + let bits_def = repdef_iter.bits_def(); + + // Full-zip requires byte-aligned values; expand 1-bit booleans to 1 byte each. + let data = match data { + DataBlock::FixedWidth(fixed) if fixed.bits_per_value == 1 => { + DataBlock::FixedWidth(Self::expand_boolean_to_bytes(fixed)) + } + other => other, + }; + + let compressor = compression_strategy.create_per_value(field, &data)?; + let (compressed_data, value_encoding) = compressor.compress(data)?; + + let description = match &compressed_data { + PerValueDataBlock::Fixed(fixed) => ProtobufUtils21::fixed_full_zip_layout( + bits_rep, + bits_def, + fixed.bits_per_value as u32, + value_encoding, + &repdef.def_meaning, + num_items as u32, + num_visible_items as u32, + ), + PerValueDataBlock::Variable(variable) => ProtobufUtils21::variable_full_zip_layout( + bits_rep, + bits_def, + variable.bits_per_offset as u32, + value_encoding, + &repdef.def_meaning, + num_items as u32, + num_visible_items as u32, + ), + }; + + let zipped = Self::serialize_full_zip(compressed_data, repdef_iter, num_items)?; + + let data = if let Some(repindex) = zipped.repetition_index { + vec![zipped.values, repindex] + } else { + vec![zipped.values] + }; + + Ok(EncodedPage { + num_rows: num_lists, + column_idx, + data, + description: PageEncoding::Structural(description), + row_number, + }) + } + + fn should_dictionary_encode( + data_block: &DataBlock, + field: &Field, + fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding, + ) -> Option { + const DEFAULT_SAMPLE_SIZE: usize = 4096; + const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98; + + // Since we only dictionary encode FixedWidth and VariableWidth blocks for now, we skip + // estimating the size for other types. + match data_block { + DataBlock::FixedWidth(fixed) => { + if fixed.bits_per_value == 64 + && fixed_width_dictionary_encoding == FixedWidthDictionaryEncoding::Exclude64Bit + { + return None; + } + if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 { + return None; + } + if fixed.bits_per_value % 8 != 0 { + return None; + } + } + DataBlock::VariableWidth(var) => { + if var.bits_per_offset != 32 && var.bits_per_offset != 64 { + return None; + } + } + _ => return None, + } + + // Don't dictionary encode tiny arrays. + let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL") + .ok() + .and_then(|val| val.parse().ok()) + .unwrap_or(100); + if data_block.num_values() < too_small { + return None; + } + + let num_values = data_block.num_values(); + + // Apply divisor threshold and cap. This is intentionally conservative: the goal is to + // avoid spending too much CPU trying to estimate very high cardinalities. + let divisor: u64 = field + .metadata + .get(DICT_DIVISOR_META_KEY) + .and_then(|val| val.parse().ok()) + .or_else(|| { + env::var("LANCE_ENCODING_DICT_DIVISOR") + .ok() + .and_then(|val| val.parse().ok()) + }) + .unwrap_or(DEFAULT_DICT_DIVISOR); + + let max_cardinality: u64 = env::var("LANCE_ENCODING_DICT_MAX_CARDINALITY") + .ok() + .and_then(|val| val.parse().ok()) + .unwrap_or(DEFAULT_DICT_MAX_CARDINALITY); + + let threshold_cardinality = num_values + .checked_div(divisor.max(1)) + .unwrap_or(0) + .min(max_cardinality); + if threshold_cardinality == 0 { + return None; + } + + // Get size ratio from metadata or env var. + let threshold_ratio = field + .metadata + .get(DICT_SIZE_RATIO_META_KEY) + .and_then(|val| val.parse::().ok()) + .or_else(|| { + env::var("LANCE_ENCODING_DICT_SIZE_RATIO") + .ok() + .and_then(|val| val.parse().ok()) + }) + .unwrap_or(DEFAULT_DICT_SIZE_RATIO); + + if threshold_ratio <= 0.0 || threshold_ratio > 1.0 { + panic!( + "Invalid parameter: dict-size-ratio is {} which is not in the range (0, 1].", + threshold_ratio + ); + } + + let data_size = data_block.data_size(); + if data_size == 0 { + return None; + } + + let max_encoded_size = (data_size as f64 * threshold_ratio) as u64; + let max_encoded_size = usize::try_from(max_encoded_size).ok()?; + + // Avoid probing dictionary encoding on data that appears to be near-unique + // or likely to exceed the dictionary budget. + if let Some(sample_unique_ratio) = + Self::sample_unique_ratio(data_block, DEFAULT_SAMPLE_SIZE)? + { + if sample_unique_ratio >= DEFAULT_SAMPLE_UNIQUE_RATIO { + return None; + } + + let projected_cardinality = (sample_unique_ratio * num_values as f64).ceil() as u64; + if projected_cardinality > threshold_cardinality { + return None; + } + } + + let max_dict_entries = u32::try_from(threshold_cardinality.min(i32::MAX as u64)).ok()?; + Some(DictEncodingBudget { + max_dict_entries, + max_encoded_size, + }) + } + + /// Samples whether a page looks near-unique before attempting dictionary encoding. + /// + /// The probe uses deterministic block sampling (not RNG sampling), which keeps + /// the check cheap and reproducible across runs. The result is only a gate for + /// whether we try dictionary encoding, not a cardinality statistic. + /// Returns `Some(None)` when there are too few reliable samples or the block type does not + /// support dictionary encoding. Returns `None` for malformed data. + fn sample_unique_ratio(data_block: &DataBlock, max_samples: usize) -> Option> { + use std::collections::HashSet; + + const NUM_SAMPLE_BLOCKS: usize = 32; + const MIN_RELIABLE_SAMPLES: usize = 1024; + + let num_values = usize::try_from(data_block.num_values()).ok()?; + if num_values == 0 { + return Some(None); + } + + let sample_count = num_values.min(max_samples).max(1); + if sample_count < MIN_RELIABLE_SAMPLES { + return Some(None); + } + + let block_count = NUM_SAMPLE_BLOCKS.min(sample_count).min(num_values).max(1); + let samples_per_block = (sample_count / block_count).max(1); + let mut indices = Vec::with_capacity(sample_count); + for block_idx in 0..block_count { + let block_start = block_idx * num_values / block_count; + let next_block_start = ((block_idx + 1) * num_values / block_count).min(num_values); + let block_len = next_block_start.saturating_sub(block_start); + let samples_in_block = samples_per_block.min(block_len); + indices.extend((0..samples_in_block).map(|offset| block_start + offset)); + } + + if indices.len() < MIN_RELIABLE_SAMPLES { + return Some(None); + } + + let ratio = match data_block { + DataBlock::FixedWidth(fixed) => match fixed.bits_per_value { + 64 => { + let values = fixed.data.borrow_to_typed_slice::(); + let values = values.as_ref(); + let mut unique: HashSet = + HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES)); + for idx in indices.iter().copied() { + unique.insert(values.get(idx).copied()?); + } + unique.len() as f64 / indices.len() as f64 + } + 128 => { + let values = fixed.data.borrow_to_typed_slice::(); + let values = values.as_ref(); + let mut unique: HashSet = + HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES)); + for idx in indices.iter().copied() { + unique.insert(values.get(idx).copied()?); + } + unique.len() as f64 / indices.len() as f64 + } + _ => return Some(None), + }, + DataBlock::VariableWidth(var) => { + use xxhash_rust::xxh3::xxh3_64; + + // Hash variable-width slices instead of storing borrowed slice keys. + let mut unique: HashSet = + HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES)); + match var.bits_per_offset { + 32 => { + let offsets_ref = var.offsets.borrow_to_typed_slice::(); + let offsets: &[u32] = offsets_ref.as_ref(); + for i in indices.iter().copied() { + let start = usize::try_from(*offsets.get(i)?).ok()?; + let end = usize::try_from(*offsets.get(i + 1)?).ok()?; + if start > end || end > var.data.len() { + return None; + } + unique.insert(xxh3_64(&var.data[start..end])); + } + } + 64 => { + let offsets_ref = var.offsets.borrow_to_typed_slice::(); + let offsets: &[u64] = offsets_ref.as_ref(); + for i in indices.iter().copied() { + let start = usize::try_from(*offsets.get(i)?).ok()?; + let end = usize::try_from(*offsets.get(i + 1)?).ok()?; + if start > end || end > var.data.len() { + return None; + } + unique.insert(xxh3_64(&var.data[start..end])); + } + } + _ => return Some(None), + } + unique.len() as f64 / indices.len() as f64 + } + _ => return Some(None), + }; + + Some(Some(ratio)) + } + + fn slice_repdef(repdef: &SerializedRepDefs, range: Range) -> SerializedRepDefs { + let repetition_levels = repdef + .repetition_levels + .as_ref() + .map(|levels| levels[range.clone()].to_vec()); + let definition_levels = repdef + .definition_levels + .as_ref() + .map(|levels| levels[range].to_vec()); + SerializedRepDefs::new_with_fixed_size_list_levels( + repetition_levels, + definition_levels, + repdef.def_meaning.clone(), + repdef.has_fixed_size_list_levels(), + ) + } + + fn slice_arrays( + arrays: &[ArrayRef], + value_start: u64, + num_values: u64, + ) -> Result> { + if num_values == 0 { + return Ok(Vec::new()); + } + + let mut values_to_skip = usize::try_from(value_start).map_err(|_| { + Error::invalid_input(format!("Value start {} is too large", value_start)) + })?; + let mut values_remaining = usize::try_from(num_values).map_err(|_| { + Error::invalid_input(format!("Value count {} is too large", num_values)) + })?; + let mut sliced = Vec::new(); + + for array in arrays { + if values_to_skip >= array.len() { + values_to_skip -= array.len(); + continue; + } + + let offset = values_to_skip; + let len = (array.len() - offset).min(values_remaining); + sliced.push(array.slice(offset, len)); + values_remaining -= len; + values_to_skip = 0; + + if values_remaining == 0 { + break; + } + } + + if values_remaining != 0 { + return Err(Error::internal(format!( + "Page split requested {} values starting at {}, but the page did not contain enough values", + num_values, value_start + ))); + } + + Ok(sliced) + } + + fn split_pages_for_miniblock_repdef_budget( + arrays: Vec, + repdef: SerializedRepDefs, + budget: MiniBlockRepDefBudget, + row_number: u64, + num_rows: u64, + ) -> Result> { + if budget == MiniBlockRepDefBudget::WithinBudget { + return Ok(vec![PrimitivePageData { + arrays, + structure: PrimitivePageStructure::Dense { + repdef, + single_row_miniblock_repdef_levels: None, + }, + row_number, + num_rows, + }]); + } + if let MiniBlockRepDefBudget::SingleRowOverBudget(num_levels) = budget { + return Ok(vec![PrimitivePageData { + arrays, + structure: PrimitivePageStructure::Dense { + repdef, + single_row_miniblock_repdef_levels: Some(num_levels), + }, + row_number, + num_rows, + }]); + } + + let MiniBlockRepDefBudget::RequiresPageSplit(splits) = budget else { + unreachable!(); + }; + + let mut pages = Vec::with_capacity(splits.len()); + for split in splits { + let arrays = Self::slice_arrays(&arrays, split.value_start, split.num_values)?; + let repdef = Self::slice_repdef(&repdef, split.level_range); + pages.push(PrimitivePageData { + arrays, + structure: PrimitivePageStructure::Dense { + repdef, + single_row_miniblock_repdef_levels: None, + }, + row_number: row_number + split.row_start, + num_rows: split.num_rows, + }); + } + Ok(pages) + } + + fn encode_dense_page( + ctx: PrimitiveEncodeContext, + page: PrimitivePageData, + compression_strategy: Arc, + miniblock_chunk_size: MiniblockChunkSize, + complex_null_encoding: ComplexNullEncoding, + fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding, + ) -> Result { + let PrimitiveEncodeContext { + column_idx, + field, + encoding_metadata, + is_simple_validity, + has_repdef_info, + } = ctx; + let PrimitivePageData { + arrays, + structure, + row_number, + num_rows, + } = page; + let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); + + let (repdef, single_row_miniblock_repdef_levels) = match structure { + PrimitivePageStructure::Dense { + repdef, + single_row_miniblock_repdef_levels, + } => (repdef, single_row_miniblock_repdef_levels), + PrimitivePageStructure::Sparse { .. } => { + unreachable!("dense atom received sparse page") + } + }; + + if num_values == 0 { + // This page contains only structural events, such as empty/null list rows. + // The existing complex-null layout stores the rep/def stream without value buffers. + log::debug!( + "Encoding column {} with {} items ({} rows) using complex-null layout", + column_idx, + num_values, + num_rows + ); + return Self::encode_complex_all_null( + column_idx, + repdef, + row_number, + num_rows, + complex_null_encoding, + compression_strategy.as_ref(), + ); + } + + let leaf_validity = Self::leaf_validity(&repdef, num_values as usize)?; + let all_null = leaf_validity + .as_ref() + .map(|validity| validity.count_set_bits() == 0) + .unwrap_or(false); + + if all_null { + return if is_simple_validity { + log::debug!( + "Encoding column {} with {} items ({} rows) using simple-null layout", + column_idx, + num_values, + num_rows + ); + Self::encode_simple_all_null(column_idx, num_values, row_number) + } else { + log::debug!( + "Encoding column {} with {} items ({} rows) using complex-null layout", + column_idx, + num_values, + num_rows + ); + Self::encode_complex_all_null( + column_idx, + repdef, + row_number, + num_rows, + complex_null_encoding, + compression_strategy.as_ref(), + ) + }; + } + + if let DataType::Struct(fields) = &field.data_type() + && fields.is_empty() + { + if has_repdef_info { + return Err(Error::invalid_input_source(format!("Empty structs with rep/def information are not yet supported. The field {} is an empty struct that either has nulls or is in a list.", field.name).into())); + } + // This is maybe a little confusing but the reader should never look at this anyways and it + // seems like overkill to invent a new layout just for "empty structs". + return Self::encode_simple_all_null(column_idx, num_values, row_number); + } + + let data_block = DataBlock::from_arrays(&arrays, num_values); + + if let Some(num_levels) = single_row_miniblock_repdef_levels { + let requested_encoding = encoding_metadata + .get(STRUCTURAL_ENCODING_META_KEY) + .map(|requested| requested.to_lowercase()); + let fullzip_error = match &data_block { + // 1-bit booleans are widened to bytes inside `encode_full_zip`. + DataBlock::FixedWidth(fixed) + if fixed.bits_per_value != 1 && !fixed.bits_per_value.is_multiple_of(8) => + { + Some(format!( + "Full-zip fixed-width values must be byte aligned, got {} bits per value", + fixed.bits_per_value + )) + } + DataBlock::VariableWidth(variable) + if !variable.bits_per_offset.is_multiple_of(8) => + { + Some(format!( + "Full-zip variable-width offsets must be byte aligned, got {} bits per offset", + variable.bits_per_offset + )) + } + DataBlock::VariableWidth(variable) + if variable.bits_per_offset != 32 && variable.bits_per_offset != 64 => + { + Some(format!( + "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits", + variable.bits_per_offset + )) + } + DataBlock::Dictionary(_) => { + Some("Full-zip does not encode dictionary data blocks directly".to_string()) + } + DataBlock::FixedSizeList(fsl) => match fsl.clone().try_into_flat() { + Some(flat) if flat.bits_per_value.is_multiple_of(8) => None, + Some(flat) => Some(format!( + "Full-zip fixed-size-list values must be byte aligned after flattening, got {} bits per value", + flat.bits_per_value + )), + None => Some( + "Full-zip fixed-size-list capability requires a flat fixed-width child" + .to_string(), + ), + }, + DataBlock::FixedWidth(_) | DataBlock::VariableWidth(_) | DataBlock::Struct(_) => { + None + } + other => Some(format!( + "Full-zip does not support value block type {}", + other.name() + )), + }; + match requested_encoding.as_deref() { + Some(STRUCTURAL_ENCODING_FULLZIP) => { + if let Some(reason) = fullzip_error { + return Err(Error::invalid_input_source(reason.into())); + } + return Self::encode_full_zip( + column_idx, + &field, + compression_strategy.as_ref(), + data_block, + repdef, + row_number, + num_rows, + ); + } + Some(STRUCTURAL_ENCODING_MINIBLOCK) | None => { + if requested_encoding.is_none() && fullzip_error.is_none() { + log::debug!( + "Encoding column {} with {} items using full-zip layout because mini-block cannot split the structural page", + column_idx, + num_values + ); + return Self::encode_full_zip( + column_idx, + &field, + compression_strategy.as_ref(), + data_block, + repdef, + row_number, + num_rows, + ); + } + return Err(Error::invalid_input_source( + format!( + "Mini-block cannot encode {} rep/def levels in one top-level row. \ + This usually means the row contains too much nested structure \ + for the current layout.", + num_levels + ) + .into(), + )); + } + _ => {} + } + } + + let requires_full_zip_packed_struct = + if let DataBlock::Struct(ref struct_data_block) = data_block { + struct_data_block.has_variable_width_child() + } else { + false + }; + + if requires_full_zip_packed_struct { + log::debug!( + "Encoding column {} with {} items using full-zip packed struct layout", + column_idx, + num_values + ); + return Self::encode_full_zip( + column_idx, + &field, + compression_strategy.as_ref(), + data_block, + repdef, + row_number, + num_rows, + ); + } + + if let DataBlock::Dictionary(dict) = data_block { + log::debug!( + "Encoding column {} with {} items using dictionary encoding (already dictionary encoded)", + column_idx, + num_values + ); + let (mut indices_data_block, dictionary_data_block) = dict.into_parts(); + // TODO: https://github.com/lancedb/lance/issues/4809 + // If we compute stats on dictionary_data_block => panic. + // If we don't compute stats on indices_data_block => panic. + // This is messy. Don't make me call compute_stat ever. + indices_data_block.compute_stat(); + return Self::encode_miniblock( + column_idx, + &field, + compression_strategy.as_ref(), + indices_data_block, + repdef, + row_number, + Some(dictionary_data_block), + num_rows, + miniblock_chunk_size, + ); + } + + // Try dictionary encoding first if applicable. If encoding aborts, fall back to the + // preferred structural encoding. + let dict_result = Self::should_dictionary_encode( + &data_block, + &field, + fixed_width_dictionary_encoding, + ) + .and_then(|budget| { + log::debug!( + "Encoding column {} with {} items using dictionary encoding (mini-block layout)", + column_idx, + num_values + ); + dict::dictionary_encode( + &data_block, + budget.max_dict_entries, + budget.max_encoded_size, + ) + }); + + if let Some((indices_data_block, dictionary_data_block)) = dict_result { + Self::encode_miniblock( + column_idx, + &field, + compression_strategy.as_ref(), + indices_data_block, + repdef, + row_number, + Some(dictionary_data_block), + num_rows, + miniblock_chunk_size, + ) + } else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) { + log::debug!( + "Encoding column {} with {} items using mini-block layout", + column_idx, + num_values + ); + Self::encode_miniblock( + column_idx, + &field, + compression_strategy.as_ref(), + data_block, + repdef, + row_number, + None, + num_rows, + miniblock_chunk_size, + ) + } else if Self::prefers_fullzip(encoding_metadata.as_ref()) { + log::debug!( + "Encoding column {} with {} items using full-zip layout", + column_idx, + num_values + ); + Self::encode_full_zip( + column_idx, + &field, + compression_strategy.as_ref(), + data_block, + repdef, + row_number, + num_rows, + ) + } else { + Err(Error::invalid_input_source(format!("Cannot determine structural encoding for field {}. This typically indicates an invalid value of the field metadata key {}", field.name, STRUCTURAL_ENCODING_META_KEY).into())) + } + } + + /// Rebuilds dictionary chunks whose values array is empty. + /// + /// Such chunks are entirely null and retain their key validity until rep/def has been + /// recorded. At flush time their keys are zeroed and attached to a neighboring non-empty + /// dictionary. If every chunk is empty, one non-null placeholder value makes key zero valid. + /// Rep/def retains the logical nullness, so these replacement keys are never exposed. + fn rebuild_empty_dictionary_chunks(arrays: Vec) -> Result> { + if !arrays.iter().any(|array| { + array + .as_any_dictionary_opt() + .is_some_and(|dictionary| dictionary.values().is_empty()) + }) { + return Ok(arrays); + } + + let zeroed = |data_type: &DataType, len: usize| { + new_null_array(data_type, len) + .to_data() + .into_builder() + .nulls(None) + .build() + .map(make_array) + }; + let rebuild = + |keys: Vec, data_type: &DataType, values: ArrayRef| -> Result { + let keys = keys.iter().map(|keys| keys.as_ref()).collect::>(); + let keys = arrow_select::concat::concat(&keys)?; + let data = keys + .to_data() + .into_builder() + .data_type(data_type.clone()) + .child_data(vec![values.to_data()]) + .build()?; + Ok(make_array(data)) + }; + + let mut rebuilt = Vec::with_capacity(arrays.len()); + let mut pending_keys = Vec::with_capacity(arrays.len()); + let mut empty_data_type = None; + for array in arrays { + let Some(dictionary) = array.as_any_dictionary_opt() else { + return Err(Error::invalid_input_source( + "Cannot mix dictionary and non-dictionary chunks".into(), + )); + }; + if dictionary.values().is_empty() { + pending_keys.push(zeroed(dictionary.keys().data_type(), array.len())?); + empty_data_type.get_or_insert_with(|| array.data_type().clone()); + } else if pending_keys.is_empty() { + rebuilt.push(array); + } else { + pending_keys.push(make_array(dictionary.keys().to_data())); + rebuilt.push(rebuild( + std::mem::take(&mut pending_keys), + array.data_type(), + dictionary.values().clone(), + )?); + } + } + + if pending_keys.is_empty() { + return Ok(rebuilt); + } + if let Some(array) = rebuilt.pop() { + let Some(dictionary) = array.as_any_dictionary_opt() else { + return Err(Error::invalid_input_source( + "Cannot mix dictionary and non-dictionary chunks".into(), + )); + }; + let mut keys = Vec::with_capacity(pending_keys.len() + 1); + keys.push(make_array(dictionary.keys().to_data())); + keys.append(&mut pending_keys); + rebuilt.push(rebuild( + keys, + array.data_type(), + dictionary.values().clone(), + )?); + } else { + let data_type = empty_data_type.ok_or_else(|| { + Error::internal("Missing data type for an empty dictionary chunk") + })?; + let DataType::Dictionary(_, value_type) = &data_type else { + return Err(Error::internal(format!( + "Expected dictionary data type, got {data_type}" + ))); + }; + rebuilt.push(rebuild(pending_keys, &data_type, zeroed(value_type, 1)?)?); + } + Ok(rebuilt) + } + + // Creates encode tasks, consuming all buffered data + fn do_flush( + &mut self, + arrays: Vec, + repdefs: Vec, + row_number: u64, + num_rows: u64, + ) -> Result> { + let arrays = Self::rebuild_empty_dictionary_chunks(arrays)?; + DataBlock::validate_arrays(&arrays, &self.field.name)?; + let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); + let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity()); + let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty()); + let normalized = RepDefBuilder::normalize(repdefs); + let plan_ctx = PrimitivePlanContext { + column_idx: self.column_index, + field: &self.field, + encoding_metadata: &self.encoding_metadata, + }; + let mut pages = None; + for page_encoding in self.page_encodings.iter() { + if let Some(planned) = page_encoding.behavior.try_plan_pages( + &plan_ctx, + &arrays, + &normalized, + row_number, + num_rows, + num_values, + )? { + pages = Some(planned); + break; + } + } + let pages = pages.ok_or_else(|| { + Error::invalid_input_source( + format!( + "No primitive page planner supports field '{}'", + self.field.name + ) + .into(), + ) + })?; + + let mut tasks = Vec::with_capacity(pages.len()); + let ctx = PrimitiveEncodeContext { + column_idx: self.column_index, + field: self.field.clone(), + encoding_metadata: self.encoding_metadata.clone(), + is_simple_validity, + has_repdef_info, + }; + for page in pages { + let ctx = ctx.clone(); + let page_encodings = self.page_encodings.clone(); + let task = + spawn_cpu(move || Self::encode_page(page_encodings.as_ref(), &ctx, page)).boxed(); + tasks.push(task); + } + Ok(tasks) + } + + fn extract_validity_buf( + array: Arc, + repdef: &mut RepDefBuilder, + keep_original_array: bool, + ) -> Result> { + if let Some(validity) = array.nulls() { + if keep_original_array { + repdef.add_validity_bitmap(validity.clone()); + } else { + repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap()); + } + // Empty dictionaries have no valid key payload. Keep the validity until rep/def is + // grouped with the buffered page; `do_flush` then rebuilds the keys against either a + // neighboring values array or an all-empty-page placeholder. + if let Some(dictionary) = array.as_any_dictionary_opt() + && dictionary.values().is_empty() + { + return Ok(array); + } + let data_no_nulls = array.to_data().into_builder().nulls(None).build()?; + Ok(make_array(data_no_nulls)) + } else { + repdef.add_no_null(array.len()); + Ok(array) + } + } + + fn extract_validity( + mut array: Arc, + repdef: &mut RepDefBuilder, + keep_original_array: bool, + ) -> Result> { + match array.data_type() { + DataType::Null => { + repdef.add_validity_bitmap(NullBuffer::new(BooleanBuffer::new_unset(array.len()))); + Ok(array) + } + DataType::Dictionary(_, _) => { + array = dict::normalize_dict_nulls(array)?; + array = dict::clear_out_of_range_null_keys(array)?; + Self::extract_validity_buf(array, repdef, keep_original_array) + } + // Extract our validity buf but NOT any child validity bufs. (they will be encoded in + // as part of the values). Note: for FSL we do not use repdef.add_fsl because we do + // NOT want to increase the repdef depth. + // + // This would be quite catasrophic for something like vector embeddings. Imagine we + // had thousands of vectors and some were null but no vector contained null items. If + // we treated the vectors (primitive FSL) like we treat structural FSL we would end up + // with a rep/def value for every single item in the vector. + _ => Self::extract_validity_buf(array, repdef, keep_original_array), + } + } +} + +impl PrimitivePageEncodingBehavior for RejectSparsePrimitiveEncoding { + fn validate_field(&self, field: &Field, metadata: &HashMap) -> Result<()> { + if metadata + .get(STRUCTURAL_ENCODING_META_KEY) + .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE)) + { + return Err(Error::invalid_input_source( + format!( + "Field '{}' requests sparse structural encoding, which is not enabled by the selected file format", + field.name + ) + .into(), + )); + } + Ok(()) + } +} + +fn plan_dense_primitive_pages( + arrays: &[ArrayRef], + normalized: &NormalizedStructuralPlan, + row_number: u64, + num_rows: u64, + num_values: u64, +) -> Result> { + let (repdef, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget( + miniblock::max_repdef_levels_per_chunk, + num_rows, + num_values, + )?; + PrimitiveStructuralEncoder::split_pages_for_miniblock_repdef_budget( + arrays.to_vec(), + repdef, + miniblock_repdef_budget, + row_number, + num_rows, + ) +} + +impl PrimitivePageEncodingBehavior for DenseU16PrimitiveEncoding { + fn try_plan_pages( + &self, + _ctx: &PrimitivePlanContext<'_>, + arrays: &[ArrayRef], + normalized: &NormalizedStructuralPlan, + row_number: u64, + num_rows: u64, + num_values: u64, + ) -> Result>> { + Ok(Some(plan_dense_primitive_pages( + arrays, normalized, row_number, num_rows, num_values, + )?)) + } + + fn try_encode_page( + &self, + ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + Ok(PrimitiveEncodeAttempt::Encoded( + PrimitiveStructuralEncoder::encode_dense_page( + ctx.clone(), + page, + self.compression.clone(), + MiniblockChunkSize::U16, + ComplexNullEncoding::RawLevels, + FixedWidthDictionaryEncoding::Exclude64Bit, + )?, + )) + } +} + +impl PrimitivePageEncodingBehavior for DenseU32PrimitiveEncoding { + fn try_plan_pages( + &self, + _ctx: &PrimitivePlanContext<'_>, + arrays: &[ArrayRef], + normalized: &NormalizedStructuralPlan, + row_number: u64, + num_rows: u64, + num_values: u64, + ) -> Result>> { + Ok(Some(plan_dense_primitive_pages( + arrays, normalized, row_number, num_rows, num_values, + )?)) + } + + fn try_encode_page( + &self, + ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + Ok(PrimitiveEncodeAttempt::Encoded( + PrimitiveStructuralEncoder::encode_dense_page( + ctx.clone(), + page, + self.compression.clone(), + MiniblockChunkSize::U32, + ComplexNullEncoding::CompressedLevels, + FixedWidthDictionaryEncoding::Include64Bit, + )?, + )) + } +} + +impl PrimitivePageEncodingBehavior for SparsePrimitiveEncoding { + fn try_plan_pages( + &self, + ctx: &PrimitivePlanContext<'_>, + arrays: &[ArrayRef], + normalized: &NormalizedStructuralPlan, + row_number: u64, + num_rows: u64, + num_values: u64, + ) -> Result>> { + let requested_encoding = ctx.encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY); + let requests_sparse = requested_encoding + .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE)); + if requests_sparse { + let plan = sparse::writer::plan(normalized, num_values)?; + if sparse::writer::uses_constant_layout(&plan, ctx.field) { + return Ok(None); + } + return Ok(Some(vec![PrimitivePageData { + arrays: arrays.to_vec(), + structure: PrimitivePageStructure::Sparse { + plan, + prepared_values: None, + }, + row_number, + num_rows, + }])); + } + + let (_, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget( + miniblock::max_repdef_levels_per_chunk, + num_rows, + num_values, + )?; + let automatic_sparse = layout::select_automatic_sparse( + requested_encoding.map(String::as_str), + &miniblock_repdef_budget, + || { + let data = DataBlock::from_arrays(arrays, num_values); + if !sparse::writer::supports_value_block(&data) { + return Ok(None); + } + let prepared_values = match sparse::writer::prepare_values( + ctx.field, + self.compression.as_ref(), + data, + MiniblockChunkSize::U32, + ) { + Ok(prepared_values) => prepared_values, + Err(error) => { + debug!( + "Keeping column {} on its dense structural path because sparse value preparation is unavailable: {}", + ctx.column_idx, error + ); + return Ok(None); + } + }; + let plan = sparse::writer::plan(normalized, num_values)?; + if sparse::writer::uses_constant_layout(&plan, ctx.field) { + return Ok(None); + } + Ok(Some((plan, prepared_values))) + }, + )?; + Ok(automatic_sparse.map(|(plan, prepared_values)| { + vec![PrimitivePageData { + arrays: arrays.to_vec(), + structure: PrimitivePageStructure::Sparse { + plan, + prepared_values: Some(prepared_values), + }, + row_number, + num_rows, + }] + })) + } + + fn try_encode_page( + &self, + ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + if !matches!(&page.structure, PrimitivePageStructure::Sparse { .. }) { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + let PrimitivePageData { + arrays, + structure: + PrimitivePageStructure::Sparse { + plan, + prepared_values, + }, + row_number, + num_rows, + } = page + else { + unreachable!() + }; + let num_values = arrays.iter().map(|array| array.len() as u64).sum(); + log::debug!( + "Encoding column {} with {} visible items ({} rows) using sparse layout", + ctx.column_idx, + num_values, + num_rows + ); + Ok(PrimitiveEncodeAttempt::Encoded( + sparse::writer::encode_page( + ctx.column_idx, + &ctx.field, + self.compression.as_ref(), + prepared_values.map_or_else( + || { + sparse::writer::SparseValueInput::Unprepared(DataBlock::from_arrays( + &arrays, num_values, + )) + }, + sparse::writer::SparseValueInput::Prepared, + ), + plan, + row_number, + num_rows, + MiniblockChunkSize::U32, + )?, + )) + } +} + +impl PrimitivePageEncodingBehavior for ConstantPrimitiveEncoding { + fn try_encode_page( + &self, + ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + let PrimitivePageStructure::Dense { repdef, .. } = &page.structure else { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + }; + let num_values: u64 = page.arrays.iter().map(|array| array.len() as u64).sum(); + if num_values == 0 { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + let leaf_validity = PrimitiveStructuralEncoder::leaf_validity(repdef, num_values as usize)?; + if leaf_validity + .as_ref() + .is_some_and(|validity| validity.count_set_bits() == 0) + || matches!(ctx.field.data_type(), DataType::Struct(fields) if fields.is_empty()) + { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + let Some(scalar) = + PrimitiveStructuralEncoder::find_constant_scalar(&page.arrays, leaf_validity.as_ref())? + else { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + }; + let PrimitivePageData { + structure: PrimitivePageStructure::Dense { repdef, .. }, + row_number, + num_rows, + .. + } = page + else { + unreachable!() + }; + log::debug!( + "Encoding column {} with {} items ({} rows) using constant layout", + ctx.column_idx, + num_values, + num_rows + ); + Ok(PrimitiveEncodeAttempt::Encoded( + constant::encode_constant_page(ctx.column_idx, scalar, repdef, row_number, num_rows)?, + )) + } +} + +impl FieldEncoder for PrimitiveStructuralEncoder { + // Buffers data, if there is enough to write a page then we create an encode task + fn maybe_encode( + &mut self, + array: ArrayRef, + _external_buffers: &mut OutOfLineBuffers, + mut repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + let array = Self::extract_validity(array, &mut repdef, self.keep_original_array)?; + self.accumulated_repdefs.push(repdef); + + if let Some((arrays, row_number, num_rows)) = + self.accumulation_queue.insert(array, row_number, num_rows) + { + let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs); + Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?) + } else { + Ok(vec![]) + } + } + + // If there is any data left in the buffer then create an encode task from it + fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result> { + if let Some((arrays, row_number, num_rows)) = self.accumulation_queue.flush() { + let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs); + Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?) + } else { + Ok(vec![]) + } + } + + fn num_columns(&self) -> u32 { + 1 + } + + fn finish( + &mut self, + _external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + std::future::ready(Ok(vec![EncodedColumn::default()])).boxed() + } +} + +#[cfg(test)] +#[allow(clippy::single_range_in_vec_init)] +mod tests { + use super::{ + CachedPage, CachedPageData, ChunkInstructions, DataBlock, DecodeMiniBlockTask, + DecodePageTask, FixedFullZipDecodeTask, FixedPerValueDecompressor, FixedWidthDataBlock, + FixedWidthDictionaryEncoding, FullZipCacheableState, FullZipDecodeDetails, + FullZipDecodeTaskItem, FullZipReadSource, FullZipRepIndexDetails, FullZipScheduler, + LazyLevels, LevelCodec, LevelCursor, LevelPlan, MiniBlockChunk, MiniBlockChunkIndex, + MiniBlockCompressed, MiniblockChunkSize, PageCacheView, PageDataCacheKey, + PageInfoAndScheduler, PageInitialization, PerValueDataBlock, PerValueDecompressor, + PreambleAction, RunEndsBuilder, RunPosition, RunStorage, SimpleAllNullScheduler, + StructuralPageScheduler, StructuralPrimitiveFieldScheduler, VariableFullZipDecoder, + dense_levels_from_block, map_ranges_to_pages, pages_overlapping_ranges, + validate_complex_all_null_levels, + }; + use crate::EncodingsIo; + use crate::buffer::LanceBuffer; + use crate::compression::{ + BlockCompressor, DefaultDecompressionStrategy, MiniBlockDecompressor, + }; + use crate::constants::{ + COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY, + DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY, + STRUCTURAL_ENCODING_MINIBLOCK, + }; + use crate::data::BlockInfo; + use crate::decoder::{ + FilterExpression, PageEncoding, RequestBatch, SchedulerContext, StructuralFieldDecoder, + StructuralFieldScheduler, + }; + use crate::encodings::logical::primitive::fullzip::PerValueCompressor; + use crate::encodings::logical::primitive::{ + ChunkDrainInstructions, LoadedChunk, PrimitiveStructuralEncoder, + StructuralPrimitiveFieldDecoder, + }; + use crate::encodings::physical::rle::{RleDecompressor, RleEncoder, RleRuns, RunLengthWidth}; + use crate::encodings::physical::value::{ValueDecompressor, ValueEncoder}; + use crate::format::ProtobufUtils21; + use crate::format::pb21; + use crate::format::pb21::compressive_encoding::Compression; + use crate::repdef::build_control_word_iterator; + use crate::testing::TestEncoding; + use crate::testing::{SimulatedScheduler, TestCases, check_round_trip_encoding_of_data}; + use arrow_array::{ + Array, ArrayRef, DictionaryArray, FixedSizeListArray, Float32Array, Int8Array, + PrimitiveArray, StringArray, UInt8Array, make_array, new_null_array, types::Int32Type, + }; + use arrow_buffer::{BooleanBuffer, NullBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field as ArrowField}; + use std::collections::HashMap; + use std::{collections::VecDeque, ops::Range, sync::Arc}; + + #[derive(Debug)] + struct CacheMarker(u8); + + impl lance_core::cache::DeepSizeOf for CacheMarker { + fn deep_size_of_children(&self, _context: &mut lance_core::cache::Context) -> usize { + 0 + } + } + + impl CachedPageData for CacheMarker { + fn as_arc_any(self: Arc) -> Arc { + self + } + } + + #[test] + fn test_is_narrow() { + let int8_array = Int8Array::from(vec![1, 2, 3]); + let array_ref: ArrayRef = Arc::new(int8_array); + let block = DataBlock::from_array(array_ref); + + assert!(PrimitiveStructuralEncoder::is_narrow(&block)); + + let string_array = StringArray::from(vec![Some("hello"), Some("world")]); + let block = DataBlock::from_array(string_array); + assert!(PrimitiveStructuralEncoder::is_narrow(&block)); + + let string_array = StringArray::from(vec![ + Some("hello world".repeat(100)), + Some("world".to_string()), + ]); + let block = DataBlock::from_array(string_array); + assert!((!PrimitiveStructuralEncoder::is_narrow(&block))); + } + + fn valued_dictionary() -> ArrayRef { + Arc::new( + DictionaryArray::::try_new( + PrimitiveArray::::from(vec![0]), + Arc::new(StringArray::from(vec!["a"])), + ) + .unwrap(), + ) + } + + fn empty_dictionary() -> ArrayRef { + new_null_array( + &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + 1, + ) + } + + fn null_valued_dictionary() -> ArrayRef { + Arc::new( + DictionaryArray::::try_new( + PrimitiveArray::::from(vec![0]), + Arc::new(StringArray::from(vec![None::<&str>])), + ) + .unwrap(), + ) + } + + fn sliced_empty_dictionary_with_hidden_key_payload() -> ArrayRef { + let keys = PrimitiveArray::::new( + ScalarBuffer::from(vec![5, 7, 9]), + Some(NullBuffer::new(BooleanBuffer::new_unset(3))), + ); + let values = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef; + let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); + Arc::new(dictionary.slice(1, 1)) + } + + #[rstest::rstest] + #[case::empty_after_value(vec![valued_dictionary(), empty_dictionary()])] + #[case::empty_before_value(vec![empty_dictionary(), valued_dictionary()])] + #[case::null_value_after_value(vec![valued_dictionary(), null_valued_dictionary()])] + #[case::hidden_sliced_after_value(vec![ + valued_dictionary(), + sliced_empty_dictionary_with_hidden_key_payload(), + ])] + #[tokio::test] + async fn test_mixed_valued_and_all_null_dictionary_chunks(#[case] dictionaries: Vec) { + check_round_trip_encoding_of_data( + dictionaries, + &TestCases::default() + .with_structural_encodings() + .with_page_sizes(vec![1]) + .with_range(0..1) + .with_range(1..2) + .with_indices(vec![0]) + .with_indices(vec![1]), + HashMap::new(), + ) + .await; + } + + #[tokio::test] + async fn test_sliced_empty_dictionary_with_hidden_key_payload() { + check_round_trip_encoding_of_data( + vec![sliced_empty_dictionary_with_hidden_key_payload()], + &TestCases::default() + .with_structural_encodings() + .with_page_sizes(vec![4096]), + HashMap::new(), + ) + .await; + } + + #[test] + fn test_primitive_decoder_empty_page_queue_returns_error() { + let field = Arc::new(ArrowField::new("vector", DataType::Float32, true)); + let mut decoder = StructuralPrimitiveFieldDecoder::new(&field, false); + + let err = decoder.drain(1).unwrap_err(); + assert!( + matches!(&err, lance_core::Error::Internal { .. }), + "expected internal error, got: {err:?}" + ); + let message = err.to_string(); + for expected in [ + "Primitive decoder missing page decoder", + "field 'vector'", + "data_type=Float32", + "requested_rows=1", + "remaining_rows=1", + "rows_drained_in_current=0", + "queued_pages=0", + ] { + assert!( + message.contains(expected), + "expected error to contain {expected:?}, got: {message}" + ); + } + } + + #[test] + fn test_plan_decoded_bytes_i32() { + use crate::decoder::CANDIDATE_BATCH_SIZES; + + let field = Arc::new(ArrowField::new("x", DataType::Int32, false)); + let decoder = StructuralPrimitiveFieldDecoder::new(&field, false); + + let rows_remaining = 100u64; + let bytes = decoder.plan_decoded_bytes(rows_remaining).unwrap(); + + for (i, &candidate) in CANDIDATE_BATCH_SIZES.iter().enumerate() { + let rows = (candidate as u64).min(rows_remaining); + assert_eq!(bytes[i], rows * 4, "candidate batch size {candidate}"); + } + } + + #[test] + fn test_plan_decoded_bytes_i32_with_nulls() { + use arrow_array::Int32Array; + + // Use N == CANDIDATE_BATCH_SIZES[5] == 1024. At this size both the + // values buffer (1024*4 = 4096 bytes) and the validity bitmap + // (ceil(1024/8) = 128 bytes) are naturally aligned to Arrow's 64-byte + // allocation boundary, so plan_decoded_bytes matches + // get_buffer_memory_size exactly without needing alignment arithmetic. + const N: i32 = 1024; + let array = Int32Array::from_iter((0..N).map(|i| if i % 2 == 0 { Some(i) } else { None })); + let expected = array.get_buffer_memory_size() as u64; + + let field = Arc::new(ArrowField::new("x", DataType::Int32, true)); + let decoder = StructuralPrimitiveFieldDecoder::new(&field, false); + + let bytes = decoder.plan_decoded_bytes(N as u64).unwrap(); + // bytes[5] corresponds to CANDIDATE_BATCH_SIZES[5] == 1024 == N + assert_eq!( + bytes[5], expected, + "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}", + bytes[5], + ); + } + + #[test] + fn test_plan_decoded_bytes_nullable_fsl_with_nullable_items() { + // FixedSizeList where the FSL itself is also nullable. + // Arrow allocates three buffers: + // 1. FSL null bitmap: ceil(rows / 8) + // 2. Child i32 null bitmap: ceil(rows * DIM / 8) + // 3. Child i32 values: rows * DIM * 4 + // + // N=1024, DIM=4 keeps all three naturally 64-byte aligned so the + // comparison against get_buffer_memory_size() is exact. + use arrow_array::{FixedSizeListArray, Int32Array}; + use arrow_buffer::NullBuffer; + + const N: usize = 1024; + const DIM: i32 = 4; + + let child_values = Int32Array::from_iter( + (0..(N as i32 * DIM)).map(|i| if i % 2 == 0 { Some(i) } else { None }), + ); + let item_field = Arc::new(ArrowField::new("item", DataType::Int32, true)); + let fsl_nulls = NullBuffer::from((0..N).map(|i| i % 3 != 0).collect::>()); + let fsl_array = FixedSizeListArray::new( + item_field.clone(), + DIM, + Arc::new(child_values), + Some(fsl_nulls), + ); + let expected = fsl_array.get_buffer_memory_size() as u64; + + let fsl_type = DataType::FixedSizeList(item_field, DIM); + let field = Arc::new(ArrowField::new("v", fsl_type, true)); + let decoder = StructuralPrimitiveFieldDecoder::new(&field, false); + + let bytes = decoder.plan_decoded_bytes(N as u64).unwrap(); + // bytes[5] corresponds to CANDIDATE_BATCH_SIZES[5] == 1024 == N + assert_eq!( + bytes[5], expected, + "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}", + bytes[5], + ); + } + + #[test] + fn test_plan_decoded_bytes_fixed_size_list_of_f32() { + use crate::decoder::CANDIDATE_BATCH_SIZES; + + const DIMENSION: i32 = 8; + // Items are non-nullable, which is the typical case for vector embeddings. + let item_field = Arc::new(ArrowField::new("item", DataType::Float32, false)); + let fsl_type = DataType::FixedSizeList(item_field, DIMENSION); + let field = Arc::new(ArrowField::new("vector", fsl_type, true)); + let decoder = StructuralPrimitiveFieldDecoder::new(&field, false); + + let rows_remaining = 100u64; + let bytes = decoder.plan_decoded_bytes(rows_remaining).unwrap(); + + // Each FSL row is DIMENSION f32 values = DIMENSION * 4 bytes, plus + // 1 bit of validity bitmap (the field is nullable). + let bytes_per_row = DIMENSION as u64 * 4; + for (i, &candidate) in CANDIDATE_BATCH_SIZES.iter().enumerate() { + let rows = (candidate as u64).min(rows_remaining); + assert_eq!( + bytes[i], + rows * bytes_per_row + rows.div_ceil(8), + "candidate batch size {candidate}" + ); + } + } + + #[test] + fn test_fullzip_fixed_rejects_non_byte_aligned_values() { + let fixed = FixedWidthDataBlock { + data: LanceBuffer::from(vec![0_u8]), + bits_per_value: 1, + num_values: 8, + block_info: BlockInfo::new(), + }; + let repdef = build_control_word_iterator(None, 0, None, 0, u16::MAX, 8); + + let Err(err) = PrimitiveStructuralEncoder::serialize_full_zip_fixed(fixed, repdef, 8) + else { + panic!("expected full-zip to reject 1-bit fixed-width values"); + }; + assert!( + err.to_string().contains("byte aligned"), + "unexpected error: {err}" + ); + } + + fn decode_fixed_fullzip_no_levels( + decompressor: Arc, + data: Vec, + num_rows: usize, + bytes_per_value: usize, + ) -> DataBlock { + Box::new(FixedFullZipDecodeTask { + details: Arc::new(FullZipDecodeDetails { + value_decompressor: PerValueDecompressor::Fixed(decompressor), + def_meaning: Arc::from([]), + ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 0), + max_rep: 0, + max_visible_def: u16::MAX, + }), + data, + num_rows, + bytes_per_value, + }) + .decode() + .unwrap() + .data + } + + #[test] + fn test_fixed_fullzip_decode_preallocates_exact_output_size() { + #[derive(Debug)] + struct IdentityFixedDecompressor; + + impl FixedPerValueDecompressor for IdentityFixedDecompressor { + fn decompress( + &self, + data: FixedWidthDataBlock, + num_rows: u64, + ) -> crate::Result { + assert_eq!(data.num_values, num_rows); + Ok(DataBlock::FixedWidth(data)) + } + + fn bits_per_value(&self) -> u64 { + 32 + } + + fn decoded_size_bytes(&self, num_values: u64) -> Option { + num_values.checked_mul(4) + } + } + + let make_item = |num_rows: u64| FullZipDecodeTaskItem { + data: PerValueDataBlock::Fixed(FixedWidthDataBlock { + data: LanceBuffer::from(vec![7_u8; num_rows as usize * 4]), + bits_per_value: 32, + num_values: num_rows, + block_info: BlockInfo::new(), + }), + rows_in_buf: num_rows, + }; + + let num_rows = 512; + let decoded = decode_fixed_fullzip_no_levels( + Arc::new(IdentityFixedDecompressor), + vec![make_item(128), make_item(384)], + num_rows, + 4, + ); + let values = decoded.as_fixed_width_ref().unwrap(); + let expected_size = num_rows * 4; + assert_eq!(values.data.len(), expected_size); + assert_eq!(values.data.clone().into_buffer().capacity(), expected_size); + } + + #[test] + fn test_fixed_fullzip_decode_falls_back_when_output_size_is_not_exact() { + #[derive(Debug)] + struct FallbackFixedDecompressor; + + impl FixedPerValueDecompressor for FallbackFixedDecompressor { + fn decompress( + &self, + data: FixedWidthDataBlock, + num_rows: u64, + ) -> crate::Result { + assert_eq!(data.num_values, num_rows); + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(vec![7_u8; num_rows as usize * 4]), + bits_per_value: 32, + num_values: num_rows, + block_info: BlockInfo::new(), + })) + } + + fn bits_per_value(&self) -> u64 { + // This deliberately cannot be multiplied by num_rows. If FullZip treats + // bits_per_value as an exact decoded-size estimate, decoding will fail. + u64::MAX - 7 + } + } + + let num_rows = 2; + let decoded = decode_fixed_fullzip_no_levels( + Arc::new(FallbackFixedDecompressor), + vec![FullZipDecodeTaskItem { + data: PerValueDataBlock::Fixed(FixedWidthDataBlock { + data: LanceBuffer::from(vec![0_u8; num_rows * 4]), + bits_per_value: 32, + num_values: num_rows as u64, + block_info: BlockInfo::new(), + }), + rows_in_buf: num_rows as u64, + }], + num_rows, + 4, + ); + let values = decoded.as_fixed_width_ref().unwrap(); + assert_eq!(values.num_values, num_rows as u64); + assert_eq!(values.data.len(), num_rows * 4); + } + + #[test] + fn test_fixed_fullzip_real_fsl_preallocates_exact_output_size() { + let num_rows = 64; + let dimension = 32; + let items = Arc::new(Float32Array::from_iter_values( + (0..num_rows * dimension).map(|value| value as f32), + )); + let item_field = Arc::new(ArrowField::new("item", DataType::Float32, false)); + let sample = FixedSizeListArray::new(item_field, dimension as i32, items, None); + + let (data, compression) = PerValueCompressor::compress( + &ValueEncoder::default(), + DataBlock::from_array(sample.clone()), + ) + .unwrap(); + let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { + panic!("expected fixed-size-list compression"); + }; + let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap(); + let expected_size = num_rows * dimension * size_of::(); + assert_eq!( + FixedPerValueDecompressor::decoded_size_bytes(&decompressor, num_rows as u64), + Some(expected_size as u64) + ); + + let decoded = decode_fixed_fullzip_no_levels( + Arc::new(decompressor), + vec![FullZipDecodeTaskItem { + data, + rows_in_buf: num_rows as u64, + }], + num_rows, + dimension * size_of::(), + ); + let fsl = decoded.as_fixed_size_list_ref().unwrap(); + let values = fsl.child.as_fixed_width_ref().unwrap(); + assert_eq!(values.data.len(), expected_size); + assert_eq!(values.data.clone().into_buffer().capacity(), expected_size); + + let decoded_array = make_array( + decoded + .into_arrow(sample.data_type().clone(), true) + .unwrap(), + ); + assert_eq!(decoded_array.as_ref(), &sample); + } + + #[test] + fn test_fixed_fullzip_nullable_fsl_uses_fallback_end_to_end() { + #[derive(Debug)] + struct NullableFslDecompressor { + inner: ValueDecompressor, + } + + impl FixedPerValueDecompressor for NullableFslDecompressor { + fn decompress( + &self, + data: FixedWidthDataBlock, + num_rows: u64, + ) -> crate::Result { + FixedPerValueDecompressor::decompress(&self.inner, data, num_rows) + } + + fn bits_per_value(&self) -> u64 { + // FullZip must not use this physical row width as an exact decoded-size + // estimate for the nullable, multi-buffer Arrow output. + u64::MAX - 7 + } + + fn decoded_size_bytes(&self, num_values: u64) -> Option { + FixedPerValueDecompressor::decoded_size_bytes(&self.inner, num_values) + } + } + + let num_rows = 64; + let items = Arc::new(UInt8Array::from_iter( + (0..num_rows).map(|value| (value % 3 != 0).then_some(value as u8)), + )); + let item_field = Arc::new(ArrowField::new("item", DataType::UInt8, true)); + let sample = FixedSizeListArray::new(item_field, 1, items, None); + + let (data, compression) = PerValueCompressor::compress( + &ValueEncoder::default(), + DataBlock::from_array(sample.clone()), + ) + .unwrap(); + let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { + panic!("expected fixed-size-list compression"); + }; + let decompressor = NullableFslDecompressor { + inner: ValueDecompressor::from_fsl(fsl.as_ref()).unwrap(), + }; + assert_eq!( + FixedPerValueDecompressor::decoded_size_bytes(&decompressor, num_rows as u64), + None + ); + + let decoded = decode_fixed_fullzip_no_levels( + Arc::new(decompressor), + vec![FullZipDecodeTaskItem { + data, + rows_in_buf: num_rows as u64, + }], + num_rows, + 2, + ); + let decoded_array = make_array( + decoded + .into_arrow(sample.data_type().clone(), true) + .unwrap(), + ); + assert_eq!(decoded_array.as_ref(), &sample); + } + + #[test] + fn test_miniblock_decode_uses_exact_fixed_width_output_size() { + #[derive(Debug)] + struct FixedWidthMiniBlockDecompressor; + + impl MiniBlockDecompressor for FixedWidthMiniBlockDecompressor { + fn decompress( + &self, + data: Vec, + num_values: u64, + ) -> crate::Result { + assert_eq!(data.len(), 1); + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: data.into_iter().next().unwrap(), + bits_per_value: 32, + num_values, + block_info: BlockInfo::new(), + })) + } + + fn decoded_size_bytes(&self, num_values: u64) -> Option { + num_values.checked_mul(4) + } + } + + let num_rows = 512; + let expected_size = num_rows * 4; + let mut chunk_data = Vec::new(); + chunk_data.extend_from_slice(&0_u16.to_le_bytes()); + chunk_data.extend_from_slice(&(expected_size as u16).to_le_bytes()); + let header_padding = + lance_core::utils::bit::pad_bytes::<{ super::MINIBLOCK_ALIGNMENT }>(chunk_data.len()); + chunk_data.resize(chunk_data.len() + header_padding, 0); + chunk_data.resize(chunk_data.len() + expected_size as usize, 7); + + let task = DecodeMiniBlockTask { + rep_decompressor: None, + def_decompressor: None, + value_decompressor: Arc::new(FixedWidthMiniBlockDecompressor), + dictionary_data: None, + def_meaning: Arc::from([]), + num_buffers: 1, + max_visible_level: 0, + instructions: vec![( + ChunkDrainInstructions { + chunk_instructions: ChunkInstructions { + chunk_idx: 0, + preamble: PreambleAction::Absent, + rows_to_skip: 0, + rows_to_take: num_rows, + take_trailer: false, + }, + rows_to_skip: 0, + rows_to_take: num_rows, + preamble_action: PreambleAction::Absent, + }, + LoadedChunk { + byte_range: 0..chunk_data.len() as u64, + data: LanceBuffer::from(chunk_data), + items_in_chunk: num_rows, + chunk_idx: 0, + }, + )], + has_large_chunk: false, + }; + + let decoded = Box::new(task).decode().unwrap(); + let values = decoded.data.as_fixed_width_ref().unwrap(); + assert_eq!(values.data.len(), expected_size as usize); + assert_eq!( + values.data.clone().into_buffer().capacity(), + expected_size as usize + ); + } + + #[test] + fn test_map_range() { + // Null in the middle + // [[A, B, C], [D, E], NULL, [F, G, H]] + let rep = Some(vec![1, 0, 0, 1, 0, 1, 1, 0, 0]); + let def = Some(vec![0, 0, 0, 0, 0, 1, 0, 0, 0]); + let max_visible_def = 0; + let total_items = 8; + let max_rep = 1; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Absent, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(0..1, 0..3, 0..3); + check(1..2, 3..5, 3..5); + check(2..3, 5..5, 5..6); + check(3..4, 5..8, 6..9); + check(0..2, 0..5, 0..5); + check(1..3, 3..5, 3..6); + check(2..4, 5..8, 5..9); + check(0..3, 0..5, 0..6); + check(1..4, 3..8, 3..9); + check(0..4, 0..8, 0..9); + + // Null at start + // [NULL, [A, B], [C]] + let rep = Some(vec![1, 1, 0, 1]); + let def = Some(vec![1, 0, 0, 0]); + let max_visible_def = 0; + let total_items = 3; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Absent, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(0..1, 0..0, 0..1); + check(1..2, 0..2, 1..3); + check(2..3, 2..3, 3..4); + check(0..2, 0..2, 0..3); + check(1..3, 0..3, 1..4); + check(0..3, 0..3, 0..4); + + // Null at end + // [[A], [B, C], NULL] + let rep = Some(vec![1, 1, 0, 1]); + let def = Some(vec![0, 0, 0, 1]); + let max_visible_def = 0; + let total_items = 3; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Absent, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(0..1, 0..1, 0..1); + check(1..2, 1..3, 1..3); + check(2..3, 3..3, 3..4); + check(0..2, 0..3, 0..3); + check(1..3, 1..3, 1..4); + check(0..3, 0..3, 0..4); + + // No nulls, with repetition + // [[A, B], [C, D], [E, F]] + let rep = Some(vec![1, 0, 1, 0, 1, 0]); + let def: Option<&[u16]> = None; + let max_visible_def = 0; + let total_items = 6; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Absent, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(0..1, 0..2, 0..2); + check(1..2, 2..4, 2..4); + check(2..3, 4..6, 4..6); + check(0..2, 0..4, 0..4); + check(1..3, 2..6, 2..6); + check(0..3, 0..6, 0..6); + + // No repetition, with nulls (this case is trivial) + // [A, B, NULL, C] + let rep: Option<&[u16]> = None; + let def = Some(vec![0, 0, 1, 0]); + let max_visible_def = 1; + let total_items = 4; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Absent, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(0..1, 0..1, 0..1); + check(1..2, 1..2, 1..2); + check(2..3, 2..3, 2..3); + check(0..2, 0..2, 0..2); + check(1..3, 1..3, 1..3); + check(0..3, 0..3, 0..3); + + // Tricky case, this chunk is a continuation and starts with a rep-index = 0 + // [[..., A] [B, C], NULL] + // + // What we do will depend on the preamble action + let rep = Some(vec![0, 1, 0, 1]); + let def = Some(vec![0, 0, 0, 1]); + let max_visible_def = 0; + let total_items = 3; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Take, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + // If we are taking the preamble then the range must start at 0 + check(0..1, 0..3, 0..3); + check(0..2, 0..3, 0..4); + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Skip, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(0..1, 1..3, 1..3); + check(1..2, 3..3, 3..4); + check(0..2, 1..3, 1..4); + + // Another preamble case but now it doesn't end with a new list + // [[..., A], NULL, [D, E]] + // + // What we do will depend on the preamble action + let rep = Some(vec![0, 1, 1, 0]); + let def = Some(vec![0, 1, 0, 0]); + let max_visible_def = 0; + let total_items = 4; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Take, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + // If we are taking the preamble then the range must start at 0 + check(0..1, 0..1, 0..2); + check(0..2, 0..3, 0..4); + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Skip, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + // If we are taking the preamble then the range must start at 0 + check(0..1, 1..1, 1..2); + check(1..2, 1..3, 2..4); + check(0..2, 1..3, 1..4); + + // Now a preamble case without any definition levels + // [[..., A] [B, C], [D]] + let rep = Some(vec![0, 1, 0, 1]); + let def: Option> = None; + let max_visible_def = 0; + let total_items = 4; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Take, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + // If we are taking the preamble then the range must start at 0 + check(0..1, 0..3, 0..3); + check(0..2, 0..4, 0..4); + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Skip, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(0..1, 1..3, 1..3); + check(1..2, 3..4, 3..4); + check(0..2, 1..4, 1..4); + + // If we have nested lists then non-top level lists may be empty/null + // and we need to make sure we still handle them as invisible items (we + // failed to do this previously) + let rep = Some(vec![2, 1, 2, 0, 1, 2]); + let def = Some(vec![0, 1, 2, 0, 0, 0]); + let max_rep = 2; + let max_visible_def = 0; + let total_items = 4; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Absent, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(0..3, 0..4, 0..6); + check(0..1, 0..1, 0..2); + check(1..2, 1..3, 2..5); + check(2..3, 3..4, 5..6); + + // Invisible items in a preamble that we are taking (regressing a previous failure) + let rep = Some(vec![0, 0, 1, 0, 1, 1]); + let def = Some(vec![0, 1, 0, 0, 0, 0]); + let max_rep = 1; + let max_visible_def = 0; + let total_items = 5; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Take, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(0..0, 0..1, 0..2); + check(0..1, 0..3, 0..4); + check(0..2, 0..4, 0..5); + + // Skip preamble (with invis items) and skip a few rows (with invis items) + // and then take a few rows but not all the rows + let rep = Some(vec![0, 1, 0, 1, 0, 1, 0, 1]); + let def = Some(vec![1, 0, 1, 1, 0, 0, 0, 0]); + let max_rep = 1; + let max_visible_def = 0; + let total_items = 5; + + let check = |range, expected_item_range, expected_level_range| { + let (item_range, level_range) = DecodeMiniBlockTask::map_range( + range, + rep.as_ref(), + def.as_ref(), + max_rep, + max_visible_def, + total_items, + PreambleAction::Skip, + ); + assert_eq!(item_range, expected_item_range); + assert_eq!(level_range, expected_level_range); + }; + + check(2..3, 2..4, 5..7); + } + + #[test] + fn test_slice_batch_data_and_rebase_offsets_u32() { + let data = LanceBuffer::copy_slice(b"0123456789abcdefghij"); + let offsets = LanceBuffer::reinterpret_vec(vec![6_u32, 8_u32, 8_u32, 12_u32]); + + let (sliced_data, normalized_offsets) = + VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32) + .unwrap(); + + assert_eq!(sliced_data.as_ref(), b"6789ab"); + let normalized = normalized_offsets.borrow_to_typed_slice::(); + assert_eq!(normalized.as_ref(), &[0, 2, 2, 6]); + } + + #[test] + fn test_slice_batch_data_and_rebase_offsets_u64() { + let data = LanceBuffer::copy_slice(b"abcdefghijklmnopqrstuvwxyz"); + let offsets = LanceBuffer::reinterpret_vec(vec![10_u64, 12_u64, 16_u64, 20_u64]); + + let (sliced_data, normalized_offsets) = + VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 64) + .unwrap(); + + assert_eq!(sliced_data.as_ref(), b"klmnopqrst"); + let normalized = normalized_offsets.borrow_to_typed_slice::(); + assert_eq!(normalized.as_ref(), &[0, 2, 6, 10]); + } + + #[test] + fn test_slice_batch_data_and_rebase_offsets_rejects_invalid_offsets() { + let data = LanceBuffer::copy_slice(b"abcd"); + let offsets = LanceBuffer::reinterpret_vec(vec![3_u32, 2_u32]); + + let err = VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32) + .expect_err("offset end before start should error"); + assert!(err.to_string().contains("less than base")); + } + + #[test] + fn test_schedule_instructions() { + // Convert repetition index to bytes for testing + let rep_data: Vec = vec![5, 2, 3, 0, 4, 7, 2, 0]; + let rep_bytes: Vec = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect(); + let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2); + + let check = |user_ranges, expected_instructions| { + let instructions = ChunkInstructions::schedule_instructions(&chunk_index, user_ranges); + assert_eq!(instructions, expected_instructions); + }; + + // The instructions we expect if we're grabbing the whole range + let expected_take_all = vec![ + ChunkInstructions { + chunk_idx: 0, + preamble: PreambleAction::Absent, + rows_to_skip: 0, + rows_to_take: 6, + take_trailer: true, + }, + ChunkInstructions { + chunk_idx: 1, + preamble: PreambleAction::Take, + rows_to_skip: 0, + rows_to_take: 2, + take_trailer: false, + }, + ChunkInstructions { + chunk_idx: 2, + preamble: PreambleAction::Absent, + rows_to_skip: 0, + rows_to_take: 5, + take_trailer: true, + }, + ChunkInstructions { + chunk_idx: 3, + preamble: PreambleAction::Take, + rows_to_skip: 0, + rows_to_take: 1, + take_trailer: false, + }, + ]; + + // Take all as 1 range + check(&[0..14], expected_take_all.clone()); + + // Take all a individual rows + check( + &[ + 0..1, + 1..2, + 2..3, + 3..4, + 4..5, + 5..6, + 6..7, + 7..8, + 8..9, + 9..10, + 10..11, + 11..12, + 12..13, + 13..14, + ], + expected_take_all, + ); + + // Test some partial takes + + // 2 rows in the same chunk but not contiguous + check( + &[0..1, 3..4], + vec![ + ChunkInstructions { + chunk_idx: 0, + preamble: PreambleAction::Absent, + rows_to_skip: 0, + rows_to_take: 1, + take_trailer: false, + }, + ChunkInstructions { + chunk_idx: 0, + preamble: PreambleAction::Absent, + rows_to_skip: 3, + rows_to_take: 1, + take_trailer: false, + }, + ], + ); + + // Taking just a trailer/preamble + check( + &[5..6], + vec![ + ChunkInstructions { + chunk_idx: 0, + preamble: PreambleAction::Absent, + rows_to_skip: 5, + rows_to_take: 1, + take_trailer: true, + }, + ChunkInstructions { + chunk_idx: 1, + preamble: PreambleAction::Take, + rows_to_skip: 0, + rows_to_take: 0, + take_trailer: false, + }, + ], + ); + + // Skipping an entire chunk + check( + &[7..10], + vec![ + ChunkInstructions { + chunk_idx: 1, + preamble: PreambleAction::Skip, + rows_to_skip: 1, + rows_to_take: 1, + take_trailer: false, + }, + ChunkInstructions { + chunk_idx: 2, + preamble: PreambleAction::Absent, + rows_to_skip: 0, + rows_to_take: 2, + take_trailer: false, + }, + ], + ); + } + + #[test] + fn test_drain_instructions() { + fn drain_from_instructions( + instructions: &mut VecDeque, + mut rows_desired: u64, + need_preamble: &mut bool, + skip_in_chunk: &mut u64, + ) -> Vec { + // Note: instructions.len() is an upper bound, we typically take much fewer + let mut drain_instructions = Vec::with_capacity(instructions.len()); + while rows_desired > 0 || *need_preamble { + let (next_instructions, consumed_chunk) = instructions + .front() + .unwrap() + .drain_from_instruction(&mut rows_desired, need_preamble, skip_in_chunk); + if consumed_chunk { + instructions.pop_front(); + } + drain_instructions.push(next_instructions); + } + drain_instructions + } + + // Convert repetition index to bytes for testing + let rep_data: Vec = vec![5, 2, 3, 0, 4, 7, 2, 0]; + let rep_bytes: Vec = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect(); + let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2); + let user_ranges = vec![1..7, 10..14]; + + // First, schedule the ranges + let scheduled = ChunkInstructions::schedule_instructions(&chunk_index, &user_ranges); + + let mut to_drain = VecDeque::from(scheduled.clone()); + + // Now we drain in batches of 4 + + let mut need_preamble = false; + let mut skip_in_chunk = 0; + + let next_batch = + drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk); + + assert!(!need_preamble); + assert_eq!(skip_in_chunk, 4); + assert_eq!( + next_batch, + vec![ChunkDrainInstructions { + chunk_instructions: scheduled[0].clone(), + rows_to_take: 4, + rows_to_skip: 0, + preamble_action: PreambleAction::Absent, + }] + ); + + let next_batch = + drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk); + + assert!(!need_preamble); + assert_eq!(skip_in_chunk, 2); + + assert_eq!( + next_batch, + vec![ + ChunkDrainInstructions { + chunk_instructions: scheduled[0].clone(), + rows_to_take: 1, + rows_to_skip: 4, + preamble_action: PreambleAction::Absent, + }, + ChunkDrainInstructions { + chunk_instructions: scheduled[1].clone(), + rows_to_take: 1, + rows_to_skip: 0, + preamble_action: PreambleAction::Take, + }, + ChunkDrainInstructions { + chunk_instructions: scheduled[2].clone(), + rows_to_take: 2, + rows_to_skip: 0, + preamble_action: PreambleAction::Absent, + } + ] + ); + + let next_batch = + drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk); + + assert!(!need_preamble); + assert_eq!(skip_in_chunk, 0); + + assert_eq!( + next_batch, + vec![ + ChunkDrainInstructions { + chunk_instructions: scheduled[2].clone(), + rows_to_take: 1, + rows_to_skip: 2, + preamble_action: PreambleAction::Absent, + }, + ChunkDrainInstructions { + chunk_instructions: scheduled[3].clone(), + rows_to_take: 1, + rows_to_skip: 0, + preamble_action: PreambleAction::Take, + }, + ] + ); + + // Regression case. Need a chunk with preamble, rows, and trailer (the middle chunk here) + let rep_data: Vec = vec![5, 2, 3, 3, 20, 0]; + let rep_bytes: Vec = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect(); + let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2); + let user_ranges = vec![0..28]; + + // First, schedule the ranges + let scheduled = ChunkInstructions::schedule_instructions(&chunk_index, &user_ranges); + + let mut to_drain = VecDeque::from(scheduled.clone()); + + // Drain first chunk and some of second chunk + + let mut need_preamble = false; + let mut skip_in_chunk = 0; + + let next_batch = + drain_from_instructions(&mut to_drain, 7, &mut need_preamble, &mut skip_in_chunk); + + assert_eq!( + next_batch, + vec![ + ChunkDrainInstructions { + chunk_instructions: scheduled[0].clone(), + rows_to_take: 6, + rows_to_skip: 0, + preamble_action: PreambleAction::Absent, + }, + ChunkDrainInstructions { + chunk_instructions: scheduled[1].clone(), + rows_to_take: 1, + rows_to_skip: 0, + preamble_action: PreambleAction::Take, + }, + ] + ); + + assert!(!need_preamble); + assert_eq!(skip_in_chunk, 1); + + // Now, the tricky part. We drain the second chunk, including the trailer, and need to make sure + // we get a drain task to take the preamble of the third chunk (and nothing else) + let next_batch = + drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk); + + assert_eq!( + next_batch, + vec![ + ChunkDrainInstructions { + chunk_instructions: scheduled[1].clone(), + rows_to_take: 2, + rows_to_skip: 1, + preamble_action: PreambleAction::Skip, + }, + ChunkDrainInstructions { + chunk_instructions: scheduled[2].clone(), + rows_to_take: 0, + rows_to_skip: 0, + preamble_action: PreambleAction::Take, + }, + ] + ); + + assert!(!need_preamble); + assert_eq!(skip_in_chunk, 0); + } + + use super::chunk_index::{PrefixSums, RowMapping}; + use super::{MINIBLOCK_ALIGNMENT, Words, build_chunk_index}; + use bytes::Bytes; + use lance_core::cache::{Context, DeepSizeOf}; + use rstest::rstest; + + /// Builds a `Words` metadata buffer (u16 words) from `(log_num_values, num_bytes)` + /// pairs, returning the words and the total data-buffer size. + fn words_from(entries: &[(u32, u32)]) -> (Words, u64) { + let mut raw = Vec::with_capacity(entries.len() * 2); + let mut total = 0u64; + for &(log, num_bytes) in entries { + assert!(num_bytes > 0 && num_bytes % MINIBLOCK_ALIGNMENT as u32 == 0); + let divided = num_bytes / MINIBLOCK_ALIGNMENT as u32 - 1; + let word = (divided << 4) | log; + assert!(word <= u16::MAX as u32, "test word {word} exceeds u16"); + raw.extend_from_slice(&(word as u16).to_le_bytes()); + total += num_bytes as u64; + } + (Words::from_bytes(Bytes::from(raw), false).unwrap(), total) + } + + fn rep_bytes_from(values: &[u64]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() + } + + #[rstest] + // Two full chunks of 8 values (log 3) plus a partial last chunk; byte sizes vary + // independently of value counts. + #[case::uniform_partial_last(&[(3, 16), (3, 24), (0, 8)], 19, "uniform_flat", 8, 3)] + // Single chunk covers the whole page. + #[case::single_chunk(&[(0, 24)], 5, "uniform_flat", 5, 5)] + // Last chunk is also full (exact multiple). + #[case::exact_multiple(&[(3, 16), (3, 16)], 16, "uniform_flat", 8, 8)] + // Non-last chunks differ in size, so this is a non-uniform flat page. + #[case::non_uniform(&[(4, 16), (2, 16), (0, 8)], 21, "flat", 16, 1)] + fn test_flat_detection( + #[case] entries: &[(u32, u32)], + #[case] items_in_page: u64, + #[case] expected_kind: &str, + #[case] expected_first_items: u64, + #[case] expected_last_items: u64, + ) { + let base = 100u64; + let (words, data_buf_size) = words_from(entries); + let index = build_chunk_index(&words, items_in_page, base, data_buf_size, None, 0).unwrap(); + + assert_eq!(index.row_mapping_debug(), expected_kind); + assert_eq!(index.num_chunks(), entries.len()); + assert_eq!(index.items_in_chunk(0), expected_first_items); + assert_eq!(index.items_in_chunk(entries.len() - 1), expected_last_items); + + // Byte ranges are absolute, contiguous, and exactly cover the data buffer. + let mut expected_start = base; + for (i, &(_, num_bytes)) in entries.iter().enumerate() { + let range = index.byte_range(i); + assert_eq!(range.start, expected_start); + assert_eq!(range.end - range.start, num_bytes as u64); + expected_start = range.end; + } + assert_eq!(expected_start, base + data_buf_size); + + // For flat pages rows == items, so the per-chunk items sum to the page total. + let total_items: u64 = (0..index.num_chunks()) + .map(|i| index.items_in_chunk(i)) + .sum(); + assert_eq!(total_items, items_in_page); + } + + #[test] + fn test_nested_detection_and_axes() { + // Repetition index (stride 2): three chunks holding 5, 4, 3 rows, no trailers. + let rep = rep_bytes_from(&[5, 0, 4, 0, 3, 0]); + + // Uniform leaf chunking: value counts 4, 4, 2. + let (words, data_buf_size) = words_from(&[(2, 8), (2, 8), (0, 8)]); + let index = build_chunk_index(&words, 10, 0, data_buf_size, Some(&rep), 1).unwrap(); + assert_eq!(index.row_mapping_debug(), "nested"); + assert_eq!(index.num_chunks(), 3); + // Rows come from the repetition index, not the value counts. + assert_eq!(index.first_row(0), 0); + assert_eq!(index.rows_in_chunk(0), 5); + assert_eq!(index.first_row(1), 5); + assert_eq!(index.rows_in_chunk(1), 4); + assert_eq!(index.first_row(2), 9); + assert_eq!(index.rows_in_chunk(2), 3); + // Items come from the value words. + assert_eq!(index.items_in_chunk(0), 4); + assert_eq!(index.items_in_chunk(1), 4); + assert_eq!(index.items_in_chunk(2), 2); + + // Non-uniform leaf chunking: value counts 8, 2, 5. + let (words_nu, dbs_nu) = words_from(&[(3, 8), (1, 8), (0, 8)]); + let index_nu = build_chunk_index(&words_nu, 15, 0, dbs_nu, Some(&rep), 1).unwrap(); + assert_eq!(index_nu.row_mapping_debug(), "nested"); + assert_eq!(index_nu.items_in_chunk(0), 8); + assert_eq!(index_nu.items_in_chunk(1), 2); + assert_eq!(index_nu.items_in_chunk(2), 5); + // The row axis is unchanged by the leaf chunking. + assert_eq!(index_nu.rows_in_chunk(0), 5); + } + + #[test] + fn test_uniform_flat_matches_prefix_sum_flat() { + // Distribution: 4 chunks of 4 values, last chunk 3 (15 items total). + let (words, data_buf_size) = words_from(&[(2, 8), (2, 8), (2, 8), (0, 8)]); + let uniform = build_chunk_index(&words, 15, 0, data_buf_size, None, 0).unwrap(); + assert_eq!(uniform.row_mapping_debug(), "uniform_flat"); + + // The same distribution expressed as a non-uniform Flat prefix-sum index. + let byte_starts = PrefixSums::from_deltas([8u64, 8, 8, 8].into_iter(), 4, 32); + let value_starts = PrefixSums::from_deltas([4u64, 4, 4, 3].into_iter(), 4, 15); + let flat = MiniBlockChunkIndex::new(0, byte_starts, RowMapping::Flat { value_starts }); + assert_eq!(flat.row_mapping_debug(), "flat"); + + // Lookup parity: identical byte ranges and item counts. + for i in 0..4 { + assert_eq!(uniform.byte_range(i), flat.byte_range(i)); + assert_eq!(uniform.items_in_chunk(i), flat.items_in_chunk(i)); + } + + // Scheduler parity across scan / single-row / partial / scattered multi-range. + let range_sets: Vec>> = vec![ + vec![0..15], + vec![0..1], + vec![7..8], + vec![14..15], + vec![3..10], + vec![0..2, 5..6, 12..15], + ]; + for ranges in &range_sets { + let from_uniform = ChunkInstructions::schedule_instructions(&uniform, ranges); + let from_flat = ChunkInstructions::schedule_instructions(&flat, ranges); + assert_eq!(from_uniform, from_flat, "mismatch for ranges {ranges:?}"); + } + + // A full scan yields one Absent, no-trailer instruction per chunk. + let full = ChunkInstructions::schedule_instructions(&uniform, &[0..15]); + assert_eq!(full.len(), 4); + for (i, inst) in full.iter().enumerate() { + assert_eq!(inst.chunk_idx, i); + assert_eq!(inst.preamble, PreambleAction::Absent); + assert_eq!(inst.rows_to_skip, 0); + assert!(!inst.take_trailer); + } + assert_eq!(full.iter().map(|i| i.rows_to_take).sum::(), 15); + } + + #[test] + fn test_deep_size_per_variant_below_legacy() { + // The previous representation cached 48 bytes per chunk (24 for ChunkMeta plus + // 24 for a rep-index block); every variant's heap must be well below that. + const LEGACY_PER_CHUNK: usize = 48; + let num_chunks = 3; + let heap = |index: &MiniBlockChunkIndex| index.deep_size_of_children(&mut Context::new()); + + let (uniform_words, uniform_dbs) = words_from(&[(2, 8), (2, 8), (0, 8)]); + let uniform = build_chunk_index(&uniform_words, 10, 0, uniform_dbs, None, 0).unwrap(); + assert_eq!(uniform.row_mapping_debug(), "uniform_flat"); + assert!(heap(&uniform) < LEGACY_PER_CHUNK * num_chunks); + + let (flat_words, flat_dbs) = words_from(&[(3, 8), (1, 8), (0, 8)]); + let flat = build_chunk_index(&flat_words, 11, 0, flat_dbs, None, 0).unwrap(); + assert_eq!(flat.row_mapping_debug(), "flat"); + assert!(heap(&flat) < LEGACY_PER_CHUNK * num_chunks); + // Flat carries a value-starts array that UniformFlat derives arithmetically. + assert!(heap(&flat) > heap(&uniform)); + + let rep = rep_bytes_from(&[4, 0, 3, 0, 3, 0]); + let (nested_words, nested_dbs) = words_from(&[(2, 8), (2, 8), (0, 8)]); + let nested = build_chunk_index(&nested_words, 10, 0, nested_dbs, Some(&rep), 1).unwrap(); + assert_eq!(nested.row_mapping_debug(), "nested"); + assert!(heap(&nested) < LEGACY_PER_CHUNK * num_chunks); + } + + #[tokio::test] + async fn test_fullzip_initialize_is_lazy() { + use futures::{FutureExt, future::BoxFuture}; + use std::ops::Range; + use std::sync::Mutex; + + #[derive(Debug, Clone)] + struct RecordingScheduler { + data: bytes::Bytes, + requests: Arc>>>>, + } + + impl RecordingScheduler { + fn new(data: bytes::Bytes) -> Self { + Self { + data, + requests: Arc::new(Mutex::new(Vec::new())), + } + } + + fn requests(&self) -> Vec>> { + self.requests.lock().unwrap().clone() + } + } + + impl crate::EncodingsIo for RecordingScheduler { + fn submit_request( + &self, + ranges: Vec>, + _priority: u64, + ) -> BoxFuture<'static, crate::Result>> { + self.requests.lock().unwrap().push(ranges.clone()); + let data = ranges + .into_iter() + .map(|range| self.data.slice(range.start as usize..range.end as usize)) + .collect::>(); + std::future::ready(Ok(data)).boxed() + } + } + + #[derive(Debug)] + struct TestFixedDecompressor; + + impl FixedPerValueDecompressor for TestFixedDecompressor { + fn decompress( + &self, + _data: FixedWidthDataBlock, + _num_rows: u64, + ) -> crate::Result { + unimplemented!("Test decompressor") + } + + fn bits_per_value(&self) -> u64 { + 32 + } + } + + let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(vec![ + 0; + 16 * 1024 + ]))); + let mut scheduler = FullZipScheduler { + data_buf_position: 0, + data_buf_size: 4096, + rep_index: Some(FullZipRepIndexDetails { + buf_position: 1000, + bytes_per_value: 4, + }), + priority: 0, + rows_in_page: 100, + bits_per_offset: 32, + details: Arc::new(FullZipDecodeDetails { + value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)), + def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]), + ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1), + max_rep: 0, + max_visible_def: 0, + }), + cached_state: None, + enable_cache: false, + }; + + let io_dyn: Arc = io.clone(); + let initialization = scheduler.init_layout().unwrap(); + assert!( + initialization.required_ranges().next().is_none(), + "FullZip should not request any ranges when caching is disabled" + ); + let buffers = initialization.with_buffers(vec![]).unwrap(); + let cached_data = scheduler.init_from_buffers(buffers, &io_dyn).await.unwrap(); + + assert!( + cached_data + .as_arc_any() + .downcast_ref::() + .is_some(), + "FullZip initialize should not eagerly load repetition index data" + ); + assert!(scheduler.cached_state.is_none()); + assert!( + io.requests().is_empty(), + "FullZip initialize should not issue any I/O" + ); + } + + /// Prove that an N-page cache miss issues one `submit_request` carrying every + /// page's metadata range, rather than one submission per page. The production + /// I/O scheduler can then coalesce adjacent ranges within that batch. + #[tokio::test] + async fn test_initialize_coalesces_missed_page_metadata() { + use std::ops::Range; + + use futures::FutureExt; + use futures::future::BoxFuture; + + use super::{ + CachedPage, CachedPageData, Error, PageCacheView, PageDataCacheKey, + PageInfoAndScheduler, PageInitialization, PageInitializationBuffers, PageLoadTask, + StructuralPrimitiveFieldScheduler, + }; + use crate::decoder::{FilterExpression, SchedulerContext, StructuralFieldScheduler}; + + // Give each fake page one fixed-size metadata buffer. The ranges are far + // apart because this test observes batching, not the I/O layer's distance- + // based range coalescing. + #[derive(Debug)] + struct FakePageScheduler { + meta_range: Range, + } + impl StructuralPageScheduler for FakePageScheduler { + fn init_layout(&self) -> crate::Result { + Ok(PageInitialization::new([Some(self.meta_range.clone())])) + } + fn init_from_buffers<'a>( + &'a mut self, + buffers: PageInitializationBuffers, + _io: &Arc, + ) -> BoxFuture<'a, crate::Result>> { + let [buffer] = buffers.try_into_array("fake").unwrap(); + let buffer = buffer.unwrap(); + assert_eq!( + &buffer[..8], + &self.meta_range.start.to_le_bytes(), + "each page gets the buffer for its own range" + ); + std::future::ready(Ok(Arc::new(CacheMarker(42)) as Arc)).boxed() + } + fn try_load(&mut self, data: &Arc) -> crate::Result<()> { + data.clone() + .as_arc_any() + .downcast::() + .map(|_| ()) + .map_err(|_| { + Error::invalid_input_source( + "Fake page cache data has an unexpected type".into(), + ) + }) + } + fn schedule_ranges( + &self, + _ranges: &[Range], + _io: &Arc, + ) -> crate::Result> { + unreachable!("initialize never schedules data reads"); + } + } + + const NUM_PAGES: u64 = 64; + const META_LEN: u64 = 64; + const STRIDE: u64 = 1024 * 1024; + let page_schedulers = (0..NUM_PAGES) + .map(|i| PageInfoAndScheduler { + page_index: i as usize, + row_range: (i * 100)..((i + 1) * 100), + scheduler: Box::new(FakePageScheduler { + meta_range: (i * STRIDE)..(i * STRIDE + META_LEN), + }), + }) + .collect(); + let mut scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers( + page_schedulers, + 0, + Arc::from("test"), + ); + + let io = Arc::new(RecordingScheduler::default()); + // A no-op cache always misses, so every page is a miss. + let cache = Arc::new(lance_core::cache::LanceCache::no_cache()); + let context = SchedulerContext::new(io.clone(), cache.clone()); + + // `None` requests every page; against an always-miss cache that's an + // N-page miss. + let filter = FilterExpression::no_filter(); + scheduler.initialize(None, &filter, &context).await.unwrap(); + + { + let requests = io.requests.lock().unwrap(); + assert_eq!( + requests.len(), + 1, + "an N-page miss must issue exactly one request, not one per page" + ); + assert_eq!( + requests[0].len() as u64, + NUM_PAGES, + "the single request must carry every missed page's metadata range" + ); + } + + io.requests.lock().unwrap().clear(); + let cache_misses_before = cache.stats().await.misses; + let requested_ranges = Arc::<[Range]>::from([150..151, 6_350..6_351]); + scheduler + .initialize(Some(&requested_ranges), &filter, &context) + .await + .unwrap(); + assert_eq!( + cache.stats().await.misses - cache_misses_before, + 2, + "range-scoped initialization must only look up touched pages" + ); + + { + let requests = io.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0], + vec![ + STRIDE..(STRIDE + META_LEN), + 63 * STRIDE..(63 * STRIDE + META_LEN), + ], + "range-scoped initialization must omit untouched pages" + ); + } + + io.requests.lock().unwrap().clear(); + let cache = Arc::new(lance_core::cache::LanceCache::with_capacity(4 * 1024)); + let cache_key = PageDataCacheKey { + column_index: 9, + page_index: 0, + view: PageCacheView::Native, + target_data_type_tag: Arc::from("test"), + }; + cache + .insert_with_key( + &cache_key, + Arc::new(CachedPage { + data: Arc::new(super::NoCachedPageData), + }), + ) + .await; + let mut recovering_scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers( + vec![PageInfoAndScheduler { + page_index: 0, + row_range: 0..100, + scheduler: Box::new(FakePageScheduler { + meta_range: 0..META_LEN, + }), + }], + 9, + Arc::from("test"), + ); + let recovering_context = SchedulerContext::new(io.clone(), cache.clone()); + recovering_scheduler + .initialize(None, &filter, &recovering_context) + .await + .unwrap(); + assert_eq!(io.requests.lock().unwrap().as_slice(), &[vec![0..META_LEN]]); + + io.requests.lock().unwrap().clear(); + let mut warm_scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers( + vec![PageInfoAndScheduler { + page_index: 0, + row_range: 0..100, + scheduler: Box::new(FakePageScheduler { + meta_range: 0..META_LEN, + }), + }], + 9, + Arc::from("test"), + ); + warm_scheduler + .initialize(None, &filter, &recovering_context) + .await + .unwrap(); + assert!( + io.requests.lock().unwrap().is_empty(), + "the rebuilt cache entry must be reusable" + ); + + let mut out_of_order_scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers( + vec![ + PageInfoAndScheduler { + page_index: 0, + row_range: 0..100, + scheduler: Box::new(FakePageScheduler { + meta_range: STRIDE..(STRIDE + META_LEN), + }), + }, + PageInfoAndScheduler { + page_index: 1, + row_range: 100..200, + scheduler: Box::new(FakePageScheduler { + meta_range: 0..META_LEN, + }), + }, + ], + 0, + Arc::from("test"), + ); + out_of_order_scheduler + .initialize(None, &filter, &context) + .await + .unwrap(); + assert_eq!( + io.requests.lock().unwrap().as_slice(), + &[vec![0..META_LEN, STRIDE..(STRIDE + META_LEN)]], + "metadata reads must be sorted without changing page buffer assignment" + ); + } + + #[tokio::test] + async fn page_cache_key_distinguishes_page_view_and_target_type() { + let cache = lance_core::cache::LanceCache::with_capacity(4 * 1024); + let cases = [ + ( + PageDataCacheKey { + column_index: 7, + page_index: 0, + view: PageCacheView::Native, + target_data_type_tag: Arc::from("Boolean"), + }, + 10, + ), + ( + PageDataCacheKey { + column_index: 7, + page_index: 1, + view: PageCacheView::Native, + target_data_type_tag: Arc::from("Boolean"), + }, + 11, + ), + ( + PageDataCacheKey { + column_index: 7, + page_index: 0, + view: PageCacheView::Native, + target_data_type_tag: Arc::from("UInt8"), + }, + 12, + ), + ( + PageDataCacheKey { + column_index: 7, + page_index: 0, + view: PageCacheView::BlobDescriptor, + target_data_type_tag: Arc::from("Struct"), + }, + 20, + ), + ( + PageDataCacheKey { + column_index: 7, + page_index: 0, + view: PageCacheView::BlobPayload, + target_data_type_tag: Arc::from("LargeBinary"), + }, + 21, + ), + ]; + + for (key, marker) in &cases { + cache + .insert_with_key( + key, + Arc::new(CachedPage { + data: Arc::new(CacheMarker(*marker)), + }), + ) + .await; + } + + assert_eq!(cache.size().await, cases.len()); + for (key, expected) in &cases { + let cached = cache.get_with_key(key).await.unwrap(); + let marker = cached + .data + .clone() + .as_arc_any() + .downcast::() + .unwrap(); + assert_eq!(marker.0, *expected); + } + } + + #[tokio::test] + async fn no_op_pages_do_not_create_cache_entries() { + let mut scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers( + vec![PageInfoAndScheduler { + page_index: 0, + row_range: 0..10, + scheduler: Box::new(SimpleAllNullScheduler::default()), + }], + 3, + Arc::from("test"), + ); + let io: Arc = + Arc::new(SimulatedScheduler::new(bytes::Bytes::new())); + let cache = Arc::new(lance_core::cache::LanceCache::with_capacity(4 * 1024)); + let context = SchedulerContext::new(io, cache.clone()); + + scheduler + .initialize(None, &FilterExpression::no_filter(), &context) + .await + .unwrap(); + + assert_eq!(cache.size().await, 0); + } + + #[test] + fn page_initialization_preserves_optional_buffer_slots() { + let initialization = PageInitialization::new([None, Some(20..22), None, Some(10..11)]); + assert_eq!( + initialization + .required_ranges() + .cloned() + .collect::>(), + [20..22, 10..11] + ); + + let buffers = initialization + .with_buffers(vec![ + bytes::Bytes::from_static(b"ab"), + bytes::Bytes::from_static(b"c"), + ]) + .unwrap(); + let [first, second, third, fourth] = buffers.try_into_array("test").unwrap(); + assert!(first.is_none()); + assert_eq!(second.unwrap(), bytes::Bytes::from_static(b"ab")); + assert!(third.is_none()); + assert_eq!(fourth.unwrap(), bytes::Bytes::from_static(b"c")); + + let error = PageInitialization::new([Some(0..1)]) + .with_buffers(Vec::new()) + .unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!(error.to_string().contains("1 declared ranges")); + } + + #[tokio::test] + async fn initialize_builds_only_requested_page_schedulers() { + use crate::compression::DefaultDecompressionStrategy; + use crate::decoder::{ + ColumnInfo, FilterExpression, PageEncoding, PageInfo, SchedulerContext, + StructuralFieldScheduler, + }; + use crate::format::pb; + use lance_core::datatypes::Field; + + const NUM_PAGES: u64 = 64; + const ROWS_PER_PAGE: u64 = 100; + // All-null constant pages need no initialization I/O, so this test + // observes construction alone. + let page_infos = (0..NUM_PAGES) + .map(|page_index| PageInfo { + num_rows: ROWS_PER_PAGE, + priority: page_index * ROWS_PER_PAGE, + encoding: PageEncoding::Structural(pb21::PageLayout { + layout: Some(pb21::page_layout::Layout::ConstantLayout( + pb21::ConstantLayout { + layers: vec![pb21::RepDefLayer::RepdefNullableItem as i32], + ..Default::default() + }, + )), + }), + buffer_offsets_and_sizes: Arc::new([]), + }) + .collect::>(); + let column_info = Arc::new(ColumnInfo::new( + 3, + Arc::from(page_infos), + Vec::new(), + pb::ColumnEncoding::default(), + )); + let field = Field::try_from(&ArrowField::new("x", DataType::Int32, true)).unwrap(); + #[allow(deprecated)] + let eager = StructuralPrimitiveFieldScheduler::try_new( + &column_info, + &DefaultDecompressionStrategy {}, + false, + &field, + ) + .unwrap(); + assert_eq!(eager.page_schedulers.len(), NUM_PAGES as usize); + assert!(eager.page_source.is_none()); + + let mut scheduler = StructuralPrimitiveFieldScheduler::try_new_lazy( + column_info, + Arc::new(DefaultDecompressionStrategy {}), + false, + &field, + ) + .unwrap(); + assert!( + scheduler.page_schedulers.is_empty(), + "construction must not build page schedulers" + ); + assert_eq!(scheduler.page_row_ends.len(), NUM_PAGES as usize); + + let io: Arc = Arc::new(SimulatedScheduler::new(bytes::Bytes::new())); + let cache = Arc::new(lance_core::cache::LanceCache::no_cache()); + let mut context = SchedulerContext::new(io, cache); + let filter = FilterExpression::no_filter(); + + let requested_ranges = [150..151, 6_350..6_351]; + scheduler + .initialize(Some(&requested_ranges), &filter, &context) + .await + .unwrap(); + let built_pages = |scheduler: &StructuralPrimitiveFieldScheduler| { + scheduler + .page_schedulers + .iter() + .map(|page| page.page_index) + .collect::>() + }; + assert_eq!(built_pages(&scheduler), [1, 63]); + assert_eq!(scheduler.page_schedulers[1].row_range, 6_300..6_400); + assert!(scheduler.page_source.is_some()); + + { + let mut job = scheduler + .schedule_ranges(&requested_ranges, &filter) + .unwrap(); + // Both pages fit one read batch, so one call schedules them together. + assert_eq!(job.schedule_next(&mut context).unwrap().len(), 2); + assert!(job.schedule_next(&mut context).unwrap().is_empty()); + + // Rows outside the initialized ranges must fail cleanly, not panic + // on an uninitialized page. + let mut job = scheduler.schedule_ranges(&[0..1], &filter).unwrap(); + let error = job.schedule_next(&mut context).unwrap_err(); + assert!(matches!(error, lance_core::Error::Internal { .. })); + assert!(error.to_string().contains("not initialized")); + } + + // Widening the request builds only the new pages and keeps page order. + scheduler + .initialize(Some(&[0..250]), &filter, &context) + .await + .unwrap(); + assert_eq!(built_pages(&scheduler), [0, 1, 2, 63]); + + // The eager path builds every page and releases the page source. + scheduler.initialize(None, &filter, &context).await.unwrap(); + assert_eq!(scheduler.page_schedulers.len(), NUM_PAGES as usize); + assert!(scheduler.page_source.is_none()); + let mut job = scheduler.schedule_ranges(&[0..1], &filter).unwrap(); + assert_eq!(job.schedule_next(&mut context).unwrap().len(), 1); + } + + /// Records every `submit_request` (ranges and priority) and answers each + /// range with a zero buffer that starts with the range's offset, so a test + /// can count requests and check that every reader got its own ranges back. + #[derive(Debug, Default)] + struct RecordingScheduler { + requests: std::sync::Mutex>>>, + priorities: std::sync::Mutex>, + } + + impl crate::EncodingsIo for RecordingScheduler { + fn submit_request( + &self, + ranges: Vec>, + priority: u64, + ) -> futures::future::BoxFuture<'static, crate::Result>> { + use futures::FutureExt; + + let buffers = ranges + .iter() + .map(|range| { + let mut buffer = vec![0u8; (range.end - range.start) as usize]; + let marker = range.start.to_le_bytes(); + let marker_len = marker.len().min(buffer.len()); + buffer[..marker_len].copy_from_slice(&marker[..marker_len]); + Bytes::from(buffer) + }) + .collect(); + self.requests.lock().unwrap().push(ranges); + self.priorities.lock().unwrap().push(priority); + std::future::ready(Ok(buffers)).boxed() + } + } + + #[tokio::test] + async fn page_read_batch_submits_pending_reads_as_one_request() { + let io = Arc::new(RecordingScheduler::default()); + let batch = RequestBatch::new(io.clone()); + let first = batch.submit_request(vec![100..108, 300..308], 7); + let second = batch.submit_request(vec![200..208], 3); + assert_eq!(batch.pending_bytes(), 24); + assert!( + io.requests.lock().unwrap().is_empty(), + "nothing is read before the flush" + ); + + batch.flush(); + { + let requests = io.requests.lock().unwrap(); + assert_eq!(requests.len(), 1, "the batch is one request"); + assert_eq!( + requests[0], + vec![100..108, 200..208, 300..308], + "the batch is sorted by file offset" + ); + } + assert_eq!( + io.priorities.lock().unwrap()[0], + 3, + "the batch takes the lowest priority" + ); + let first = first.await.unwrap(); + assert_eq!(first.len(), 2); + assert_eq!(first[0].as_ref(), &100_u64.to_le_bytes()); + assert_eq!(first[1].as_ref(), &300_u64.to_le_bytes()); + let second = second.await.unwrap(); + assert_eq!(second.len(), 1); + assert_eq!(second[0].as_ref(), &200_u64.to_le_bytes()); + + // A read submitted after the flush is batched until it is awaited. + let late = batch.submit_request(vec![400..408], 9).await.unwrap(); + assert_eq!(late[0].as_ref(), &400_u64.to_le_bytes()); + assert_eq!(io.requests.lock().unwrap().len(), 2); + + // Queued reads that would exceed the budget are submitted first, so + // no batch grows past it. + let io = Arc::new(RecordingScheduler::default()); + let batch = RequestBatch::new(io.clone()); + let budget = crate::decoder::IO_REQUEST_BATCH_BYTES; + let first = batch.submit_request(vec![0..budget], 1); + assert!(io.requests.lock().unwrap().is_empty()); + let second = batch.submit_request(vec![budget..(budget + 8)], 2); + assert_eq!( + io.requests.lock().unwrap().as_slice(), + &[vec![0..budget]], + "the full batch is submitted before the read that would overflow it" + ); + batch.flush(); + assert_eq!(io.requests.lock().unwrap().len(), 2); + assert_eq!(first.await.unwrap()[0].len(), budget as usize); + assert_eq!(second.await.unwrap()[0].as_ref(), &budget.to_le_bytes()); + } + + #[tokio::test] + async fn request_batch_flush_keeps_collecting_and_awaiting_flushes() { + let io = Arc::new(RecordingScheduler::default()); + let batch = RequestBatch::new(io.clone()); + let first = batch.submit_request(vec![100..108], 1); + let second = batch.submit_request(vec![300..308], 2); + batch.flush(); + assert_eq!( + io.requests.lock().unwrap().as_slice(), + &[vec![100..108, 300..308]], + "a flush submits the collected reads as one request" + ); + // The batch keeps collecting after a flush. + let third = batch.submit_request(vec![500..508], 3); + let fourth = batch.submit_request(vec![700..708], 4); + assert_eq!(batch.pending_bytes(), 16); + assert_eq!(io.requests.lock().unwrap().len(), 1); + // Awaiting a read submits whatever is queued, so a read queued outside + // a scheduling step (indirect I/O) is never stuck behind a missing flush. + assert_eq!(third.await.unwrap()[0].as_ref(), &500_u64.to_le_bytes()); + assert_eq!( + io.requests.lock().unwrap().as_slice(), + &[vec![100..108, 300..308], vec![500..508, 700..708]] + ); + assert_eq!(fourth.await.unwrap()[0].as_ref(), &700_u64.to_le_bytes()); + assert_eq!(first.await.unwrap()[0].as_ref(), &100_u64.to_le_bytes()); + assert_eq!(second.await.unwrap()[0].as_ref(), &300_u64.to_le_bytes()); + assert_eq!(io.requests.lock().unwrap().len(), 2); + + // Futures that submit when first polled and are polled in the same + // pass (struct children initializing, list columns fetching items) + // end up in one request: the flush waits for one yield. + let (fifth, sixth) = futures::join!( + async { batch.submit_request(vec![900..908], 5).await.unwrap() }, + async { batch.submit_request(vec![1100..1108], 6).await.unwrap() }, + ); + assert_eq!(fifth[0].as_ref(), &900_u64.to_le_bytes()); + assert_eq!(sixth[0].as_ref(), &1100_u64.to_le_bytes()); + assert_eq!( + io.requests.lock().unwrap().last().unwrap(), + &vec![900..908, 1100..1108], + "reads submitted by sibling futures share one request" + ); + assert_eq!(io.requests.lock().unwrap().len(), 3); + } + + #[tokio::test] + async fn schedule_next_batches_page_reads_up_to_the_budget() { + use super::{ + CachedPageData, IO_REQUEST_BATCH_BYTES, PageInfoAndScheduler, PageInitialization, + PageInitializationBuffers, PageLoadTask, StructuralPageScheduler, + StructuralPrimitiveFieldScheduler, + }; + use crate::decoder::{FilterExpression, SchedulerContext, StructuralFieldScheduler}; + + /// A page that reads `data_ranges`, one shard (and one load task) per + /// range. The bytes it gets back are covered by the batch test, so its + /// load futures just fail and the test never has to build a decoder. + #[derive(Debug)] + struct FakeDataPage { + data_ranges: Vec>, + } + + impl StructuralPageScheduler for FakeDataPage { + fn needs_initialization(&self) -> bool { + false + } + + fn init_layout(&self) -> crate::Result { + unreachable!("fake data pages need no initialization") + } + + fn init_from_buffers<'a>( + &'a mut self, + _buffers: PageInitializationBuffers, + _io: &Arc, + ) -> futures::future::BoxFuture<'a, crate::Result>> + { + unreachable!("fake data pages need no initialization") + } + + fn try_load(&mut self, _data: &Arc) -> crate::Result<()> { + Ok(()) + } + + fn schedule_ranges( + &self, + ranges: &[Range], + io: &Arc, + ) -> crate::Result> { + use futures::FutureExt; + + let num_rows = ranges.iter().map(|r| r.end - r.start).sum(); + Ok(self + .data_ranges + .iter() + .enumerate() + .map(|(shard, range)| { + drop(io.submit_request(vec![range.clone()], range.start)); + PageLoadTask { + decoder_fut: std::future::ready(Err(lance_core::Error::internal( + "fake page", + ))) + .boxed(), + // The page's rows are reported once, on the first shard. + num_rows: if shard == 0 { num_rows } else { 0 }, + } + }) + .collect()) + } + } + + fn scheduler_with_pages( + page_bytes: u64, + num_pages: u64, + ) -> StructuralPrimitiveFieldScheduler { + // Leave a gap between pages so the test can tell batching from + // the I/O layer's coalescing of adjacent ranges. + let stride = page_bytes * 2; + let pages = (0..num_pages) + .map(|page| PageInfoAndScheduler { + page_index: page as usize, + row_range: (page * 10)..((page + 1) * 10), + scheduler: Box::new(FakeDataPage { + data_ranges: vec![(page * stride)..(page * stride + page_bytes)], + }), + }) + .collect(); + StructuralPrimitiveFieldScheduler::from_page_schedulers(pages, 0, Arc::from("test")) + } + + let filter = FilterExpression::no_filter(); + let cache = Arc::new(lance_core::cache::LanceCache::no_cache()); + + // Small pages: every page of the request is read in one batch. + let io = Arc::new(RecordingScheduler::default()); + let scheduler = scheduler_with_pages(64, 6); + let mut context = SchedulerContext::new(io.clone(), cache.clone()); + let mut job = scheduler.schedule_ranges(&[5..45], &filter).unwrap(); + let scan_lines = job.schedule_next(&mut context).unwrap(); + context.flush_io(); + assert_eq!(scan_lines.len(), 5, "one scan line per touched page"); + assert_eq!( + scan_lines + .iter() + .map(|line| line.rows_scheduled) + .sum::(), + 40 + ); + assert!(job.schedule_next(&mut context).unwrap().is_empty()); + { + let requests = io.requests.lock().unwrap(); + assert_eq!(requests.len(), 1, "five pages cost one request"); + assert_eq!( + requests[0], + vec![0..64, 128..192, 256..320, 384..448, 512..576], + "the untouched last page is not read" + ); + } + assert_eq!( + io.priorities.lock().unwrap()[0], + 0, + "the batch carries the first page's priority" + ); + + // Pages as large as the batch budget are read one request at a time. + let io = Arc::new(RecordingScheduler::default()); + let scheduler = scheduler_with_pages(IO_REQUEST_BATCH_BYTES, 2); + let mut context = SchedulerContext::new(io.clone(), cache); + let mut job = scheduler.schedule_ranges(&[0..20], &filter).unwrap(); + let first = job.schedule_next(&mut context).unwrap(); + context.flush_io(); + assert_eq!(first.len(), 1); + assert_eq!(io.requests.lock().unwrap().len(), 1); + let second = job.schedule_next(&mut context).unwrap(); + context.flush_io(); + assert_eq!(second.len(), 1); + assert_eq!(io.requests.lock().unwrap().len(), 2); + assert!(job.schedule_next(&mut context).unwrap().is_empty()); + + // A page that shards its own reads (like a blob page) keeps every + // shard in a bounded read instead of one request for the whole page. + let io = Arc::new(RecordingScheduler::default()); + let shard = IO_REQUEST_BATCH_BYTES; + let scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers( + vec![PageInfoAndScheduler { + page_index: 0, + row_range: 0..10, + scheduler: Box::new(FakeDataPage { + data_ranges: vec![0..shard, (2 * shard)..(3 * shard)], + }), + }], + 0, + Arc::from("test"), + ); + let mut context = SchedulerContext::new( + io.clone(), + Arc::new(lance_core::cache::LanceCache::no_cache()), + ); + let mut job = scheduler.schedule_ranges(&[0..10], &filter).unwrap(); + let scan_lines = job.schedule_next(&mut context).unwrap(); + context.flush_io(); + assert_eq!(scan_lines.len(), 2, "one scan line per shard"); + assert_eq!( + io.requests.lock().unwrap().as_slice(), + &[vec![0..shard], vec![(2 * shard)..(3 * shard)]], + "each shard stays its own bounded request" + ); + assert!(job.schedule_next(&mut context).unwrap().is_empty()); + } + + #[test] + fn page_range_mapping_validates_and_splits_ranges() { + // Exclusive end rows of three ten-row pages. + let pages = [10_u64, 20, 30]; + + let mappings = map_ranges_to_pages(&pages, &[8..22]).unwrap(); + assert_eq!(mappings.len(), 3); + assert_eq!(mappings[0].page_idx, 0); + assert_eq!(mappings[0].ranges_in_page, [8..10]); + assert_eq!(mappings[1].page_idx, 1); + assert_eq!(mappings[1].ranges_in_page, [0..10]); + assert_eq!(mappings[2].page_idx, 2); + assert_eq!(mappings[2].ranges_in_page, [0..2]); + + let sparse_mappings = map_ranges_to_pages(&pages, &[1..2, 3..4, 28..29]).unwrap(); + assert_eq!(sparse_mappings.len(), 2); + assert_eq!(sparse_mappings[0].page_idx, 0); + assert_eq!(sparse_mappings[0].ranges_in_page, [1..2, 3..4]); + assert_eq!(sparse_mappings[1].page_idx, 2); + assert_eq!(sparse_mappings[1].ranges_in_page, [8..9]); + assert_eq!( + pages_overlapping_ranges(&pages, &[1..2, 3..4, 28..29]).unwrap(), + [0, 2] + ); + + for (ranges, expected_message) in [ + (vec![Range { start: 2, end: 1 }], "is reversed"), + (vec![15..16, 5..6], "not strictly ordered"), + (vec![5..15, 10..20], "overlap"), + (vec![5..5, 5..6], "not strictly ordered"), + ] { + let error = map_ranges_to_pages(&pages, &ranges).unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!(error.to_string().contains(expected_message)); + } + } + + #[tokio::test] + async fn test_fullzip_read_source_slices_prefetched_page() { + let page_start = 200_u64; + let page_data = LanceBuffer::copy_slice(&[0, 1, 2, 3, 4, 5, 6, 7]); + let source = FullZipReadSource::PrefetchedPage { + base_offset: page_start, + data: page_data, + }; + let ranges = vec![ + page_start..(page_start + 3), + (page_start + 4)..(page_start + 8), + ]; + let mut data = source.fetch(&ranges, 0).await.unwrap(); + assert_eq!(data.pop_front().unwrap().as_ref(), &[0, 1, 2]); + assert_eq!(data.pop_front().unwrap().as_ref(), &[4, 5, 6, 7]); + } + + #[tokio::test] + async fn test_fullzip_initialize_caches_rep_index_when_enabled() { + use futures::{FutureExt, future::BoxFuture}; + use std::ops::Range; + use std::sync::Mutex; + + #[derive(Debug, Clone)] + struct RecordingScheduler { + data: bytes::Bytes, + requests: Arc>>>>, + } + + impl RecordingScheduler { + fn new(data: bytes::Bytes) -> Self { + Self { + data, + requests: Arc::new(Mutex::new(Vec::new())), + } + } + + fn requests(&self) -> Vec>> { + self.requests.lock().unwrap().clone() + } + } + + impl crate::EncodingsIo for RecordingScheduler { + fn submit_request( + &self, + ranges: Vec>, + _priority: u64, + ) -> BoxFuture<'static, crate::Result>> { + self.requests.lock().unwrap().push(ranges.clone()); + let data = ranges + .into_iter() + .map(|range| self.data.slice(range.start as usize..range.end as usize)) + .collect::>(); + std::future::ready(Ok(data)).boxed() + } + } + + #[derive(Debug)] + struct TestFixedDecompressor; + + impl FixedPerValueDecompressor for TestFixedDecompressor { + fn decompress( + &self, + _data: FixedWidthDataBlock, + _num_rows: u64, + ) -> crate::Result { + unimplemented!("Test decompressor") + } + + fn bits_per_value(&self) -> u64 { + 32 + } + } + + let rows_in_page = 100_u64; + let bytes_per_value = 4_u64; + let rep_start = 1000_u64; + let rep_size = ((rows_in_page + 1) * bytes_per_value) as usize; + let mut data = vec![0_u8; 16 * 1024]; + data[rep_start as usize..rep_start as usize + rep_size].fill(7); + let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(data))); + + let mut scheduler = FullZipScheduler { + data_buf_position: 0, + data_buf_size: 4096, + rep_index: Some(FullZipRepIndexDetails { + buf_position: rep_start, + bytes_per_value, + }), + priority: 0, + rows_in_page, + bits_per_offset: 32, + details: Arc::new(FullZipDecodeDetails { + value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)), + def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]), + ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1), + max_rep: 0, + max_visible_def: 0, + }), + cached_state: None, + enable_cache: true, + }; + + let io_dyn: Arc = io.clone(); + let initialization = scheduler.init_layout().unwrap(); + let ranges = initialization + .required_ranges() + .cloned() + .collect::>(); + assert_eq!( + ranges, + vec![rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)] + ); + // The field scheduler fetches the declared ranges (coalesced across all + // pages) and hands the bytes back to the page to finish initialization. + let buffers = io_dyn.submit_request(ranges, 0).await.unwrap(); + let buffers = initialization.with_buffers(buffers).unwrap(); + let cached_data = scheduler.init_from_buffers(buffers, &io_dyn).await.unwrap(); + assert!( + cached_data + .clone() + .as_arc_any() + .downcast_ref::() + .is_some() + ); + assert!(scheduler.cached_state.is_none()); + scheduler.try_load(&cached_data).unwrap(); + assert!(scheduler.cached_state.is_some()); + assert_eq!( + io.requests(), + vec![vec![ + rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value) + ]] + ); + } + + #[tokio::test] + async fn test_fullzip_full_page_bypasses_rep_index_io() { + use futures::{FutureExt, future::BoxFuture}; + use std::ops::Range; + use std::sync::Mutex; + + #[derive(Debug, Clone)] + struct RecordingScheduler { + data: bytes::Bytes, + requests: Arc>>>>, + } + + impl RecordingScheduler { + fn new(data: bytes::Bytes) -> Self { + Self { + data, + requests: Arc::new(Mutex::new(Vec::new())), + } + } + + fn requests(&self) -> Vec>> { + self.requests.lock().unwrap().clone() + } + } + + impl crate::EncodingsIo for RecordingScheduler { + fn submit_request( + &self, + ranges: Vec>, + _priority: u64, + ) -> BoxFuture<'static, crate::Result>> { + self.requests.lock().unwrap().push(ranges.clone()); + let data = ranges + .into_iter() + .map(|range| self.data.slice(range.start as usize..range.end as usize)) + .collect::>(); + std::future::ready(Ok(data)).boxed() + } + } + + #[derive(Debug)] + struct TestFixedDecompressor; + + impl FixedPerValueDecompressor for TestFixedDecompressor { + fn decompress( + &self, + _data: FixedWidthDataBlock, + _num_rows: u64, + ) -> crate::Result { + unimplemented!("Test decompressor") + } + + fn bits_per_value(&self) -> u64 { + 32 + } + } + + let rows_in_page = 100_u64; + let data_start = 256_u64; + let data_size = 500_u64; + let rep_start = 4096_u64; + let bytes_per_value = 4_u64; + + let mut bytes = vec![0_u8; 16 * 1024]; + for i in 0..=rows_in_page { + let offset = (i * 5) as u32; + let pos = rep_start as usize + (i * bytes_per_value) as usize; + bytes[pos..pos + 4].copy_from_slice(&offset.to_le_bytes()); + } + let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(bytes))); + + let scheduler = FullZipScheduler { + data_buf_position: data_start, + data_buf_size: data_size, + rep_index: Some(FullZipRepIndexDetails { + buf_position: rep_start, + bytes_per_value, + }), + priority: 0, + rows_in_page, + bits_per_offset: 32, + details: Arc::new(FullZipDecodeDetails { + value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)), + def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]), + ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1), + max_rep: 0, + max_visible_def: 0, + }), + cached_state: None, + enable_cache: false, + }; + + let io_dyn: Arc = io.clone(); + let tasks = scheduler + .schedule_ranges_rep( + &[0..rows_in_page], + &io_dyn, + FullZipRepIndexDetails { + buf_position: rep_start, + bytes_per_value, + }, + ) + .unwrap(); + + let requests = io.requests(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0], vec![data_start..(data_start + data_size)]); + + let _ = tasks.into_iter().next().unwrap().decoder_fut.await.unwrap(); + let requests_after_await = io.requests(); + assert_eq!( + requests_after_await.len(), + 1, + "full page path should not issue rep-index I/O" + ); + } + + /// This test is used to reproduce fuzz test https://github.com/lancedb/lance/issues/4492 + #[tokio::test] + async fn test_fuzz_issue_4492_empty_rep_values() { + use lance_datagen::{RowCount, Seed, array, gen_batch}; + + let seed = 1823859942947654717u64; + let num_rows = 2741usize; + + // Generate the exact same data that caused the failure + let batch_gen = gen_batch().with_seed(Seed::from(seed)); + let base_generator = array::rand_type(&DataType::FixedSizeBinary(32)); + let list_generator = array::rand_list_any(base_generator, false); + + let batch = batch_gen + .anon_col(list_generator) + .into_batch_rows(RowCount::from(num_rows as u64)) + .unwrap(); + + let list_array = batch.column(0).clone(); + + // Force miniblock encoding + let mut metadata = HashMap::new(); + metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_MINIBLOCK.to_string(), + ); + + let test_cases = TestCases::default() + .with_structural_encodings() + .with_batch_size(100) + .with_range(0..num_rows.min(500) as u64) + .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]); + + check_round_trip_encoding_of_data(vec![list_array], &test_cases, metadata).await + } + + async fn test_minichunk_size_helper( + string_data: Vec>, + minichunk_size: u64, + encodings: &[TestEncoding], + ) { + use crate::constants::MINICHUNK_SIZE_META_KEY; + use crate::testing::{TestCases, check_round_trip_encoding_of_data}; + use arrow_array::{ArrayRef, StringArray}; + use std::sync::Arc; + + let string_array: ArrayRef = Arc::new(StringArray::from(string_data)); + + let mut metadata = HashMap::new(); + metadata.insert( + MINICHUNK_SIZE_META_KEY.to_string(), + minichunk_size.to_string(), + ); + metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_MINIBLOCK.to_string(), + ); + + let test_cases = TestCases::default() + .with_encodings(encodings.iter().copied()) + .with_batch_size(1000); + + check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await; + } + + #[tokio::test] + async fn test_minichunk_size_roundtrip() { + // Test that minichunk size can be configured and works correctly in round-trip encoding + let mut string_data = Vec::new(); + for i in 0..100 { + string_data.push(Some(format!("test_string_{}", i).repeat(50))); + } + // configure minichunk size to 64 bytes (smaller than the default 4kb) for Lance 2.1 + test_minichunk_size_helper( + string_data, + 64, + &[ + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse, + ], + ) + .await; + } + + #[tokio::test] + async fn test_minichunk_size_128kb_v2_2() { + // Test that minichunk size can be configured to 128KB and works correctly with Lance 2.2 + let mut string_data = Vec::new(); + // create a 500kb string array + for i in 0..10000 { + string_data.push(Some(format!("test_string_{}", i).repeat(50))); + } + test_minichunk_size_helper( + string_data, + 128 * 1024, + &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse], + ) + .await; + } + + #[tokio::test] + async fn test_binary_large_minichunk_size_over_max_miniblock_values() { + let mut string_data = Vec::new(); + // 128kb/chunk / 6 bytes (t_9999) = 21845 items per chunk + for i in 0..10000 { + string_data.push(Some(format!("t_{}", i))); + } + test_minichunk_size_helper( + string_data, + 128 * 1024, + &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse], + ) + .await; + } + + #[tokio::test] + async fn test_large_dictionary_general_compression() { + use arrow_array::{ArrayRef, StringArray}; + use std::collections::HashMap; + use std::sync::Arc; + + // Create large string dictionary data (>32KiB) with low cardinality + // Use 100 unique strings, each 500 bytes long = 50KB dictionary + let unique_values: Vec = (0..100) + .map(|i| format!("value_{:04}_{}", i, "x".repeat(500))) + .collect(); + + // Repeat these strings many times to create a large array + let repeated_strings: Vec<_> = unique_values + .iter() + .cycle() + .take(10_000) + .map(|s| Some(s.as_str())) + .collect(); + + let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef; + + // Configure test to use V2_2 and verify encoding + let test_cases = TestCases::default() + .with_u32_structural_encodings() + .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| { + assert_eq!(cols.len(), 1); + let col = &cols[0]; + + // Navigate to the dictionary encoding in the page layout + if let Some(PageEncoding::Structural(page_layout)) = + &col.final_pages.first().map(|p| &p.description) + && let Some(pb21::page_layout::Layout::MiniBlockLayout(mini_block)) = + &page_layout.layout + && let Some(dictionary_encoding) = &mini_block.dictionary + { + match dictionary_encoding.compression.as_ref() { + Some(Compression::General(general)) => { + // Verify it's using LZ4 or Zstd + let compression = general.compression.as_ref().unwrap(); + assert!( + compression.scheme() + == pb21::CompressionScheme::CompressionAlgorithmLz4 + || compression.scheme() + == pb21::CompressionScheme::CompressionAlgorithmZstd, + "Expected LZ4 or Zstd compression for large dictionary" + ); + } + _ => panic!("Expected General compression for large dictionary"), + } + } + })); + + check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await; + } + + fn dictionary_encoding_from_page( + page: &crate::encoder::EncodedPage, + ) -> &crate::format::pb21::CompressiveEncoding { + let PageEncoding::Structural(layout) = &page.description else { + panic!("Expected structural page encoding"); + }; + let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap() + else { + panic!("Expected mini-block layout"); + }; + layout + .dictionary + .as_ref() + .unwrap_or_else(|| panic!("Expected dictionary encoding")) + } + + async fn encode_variable_dict_page( + metadata: HashMap, + ) -> crate::encoder::EncodedPage { + use arrow_array::types::Int32Type; + use arrow_array::{ArrayRef, DictionaryArray, Int32Array, StringArray}; + + let values = Arc::new(StringArray::from( + (0..128) + .map(|i| format!("value_{i:04}_{}", "x".repeat(256))) + .collect::>(), + )) as ArrayRef; + let keys = Int32Array::from_iter_values((0..20_000).map(|i| i % 128)); + let dict_array = + Arc::new(DictionaryArray::::try_new(keys, values).unwrap()) as ArrayRef; + + let field = arrow_schema::Field::new( + "dict_col", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + ) + .with_metadata(metadata); + + encode_first_page(field, dict_array, TestEncoding::StructuralU32).await + } + + async fn encode_auto_fixed_dict_page( + metadata: HashMap, + ) -> crate::encoder::EncodedPage { + use arrow_array::{ArrayRef, Decimal128Array}; + + // 128-bit fixed-width values with low cardinality to trigger dictionary encoding. + let values = (0..20_000) + .map(|i| match i % 3 { + 0 => 10_i128, + 1 => 20_i128, + _ => 30_i128, + }) + .collect::>(); + let decimal = Decimal128Array::from_iter_values(values) + .with_precision_and_scale(38, 0) + .unwrap(); + let decimal = Arc::new(decimal) as ArrayRef; + + let mut field_metadata = metadata; + // Strongly encourage dictionary encoding for this synthetic test data. + field_metadata.insert( + "lance-encoding:dict-size-ratio".to_string(), + "0.99".to_string(), + ); + let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false) + .with_metadata(field_metadata); + + encode_first_page(field, decimal, TestEncoding::StructuralU32).await + } + + #[tokio::test] + async fn test_dict_values_general_compression_default_lz4_for_variable_dict_values() { + let page = encode_variable_dict_page(HashMap::new()).await; + let dictionary_encoding = dictionary_encoding_from_page(&page); + let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else { + panic!("Expected General compression for dictionary values"); + }; + let compression = general.compression.as_ref().unwrap(); + assert_eq!( + compression.scheme(), + pb21::CompressionScheme::CompressionAlgorithmLz4 + ); + } + + #[tokio::test] + async fn test_dict_values_general_compression_default_lz4_for_fixed_dict_values() { + let page = encode_auto_fixed_dict_page(HashMap::new()).await; + let dictionary_encoding = dictionary_encoding_from_page(&page); + let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else { + panic!("Expected General compression for dictionary values"); + }; + let compression = general.compression.as_ref().unwrap(); + assert_eq!( + compression.scheme(), + pb21::CompressionScheme::CompressionAlgorithmLz4 + ); + } + + #[tokio::test] + async fn test_dict_values_general_compression_zstd() { + let mut metadata = HashMap::new(); + metadata.insert( + DICT_VALUES_COMPRESSION_META_KEY.to_string(), + "zstd".to_string(), + ); + let page = encode_variable_dict_page(metadata).await; + let dictionary_encoding = dictionary_encoding_from_page(&page); + let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else { + panic!("Expected General compression for dictionary values"); + }; + let compression = general.compression.as_ref().unwrap(); + assert_eq!( + compression.scheme(), + pb21::CompressionScheme::CompressionAlgorithmZstd + ); + } + + #[tokio::test] + async fn test_dict_values_general_compression_none() { + let mut metadata = HashMap::new(); + metadata.insert( + DICT_VALUES_COMPRESSION_META_KEY.to_string(), + "none".to_string(), + ); + let page = encode_variable_dict_page(metadata).await; + let dictionary_encoding = dictionary_encoding_from_page(&page); + assert!( + !matches!( + dictionary_encoding.compression.as_ref(), + Some(Compression::General(_)) + ), + "Expected dictionary values to avoid General compression" + ); + } + + #[test] + fn test_resolve_dict_values_compression_metadata_defaults_to_lz4() { + let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata( + &HashMap::new(), + None, + None, + ); + assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),); + assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY)); + } + + #[test] + fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() { + let field_metadata = HashMap::from([ + ( + DICT_VALUES_COMPRESSION_META_KEY.to_string(), + "none".to_string(), + ), + ( + DICT_VALUES_COMPRESSION_LEVEL_META_KEY.to_string(), + "7".to_string(), + ), + ]); + let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata( + &field_metadata, + Some("zstd".to_string()), + Some("3".to_string()), + ); + assert_eq!( + metadata.get(COMPRESSION_META_KEY), + Some(&"none".to_string()), + ); + assert_eq!( + metadata.get(COMPRESSION_LEVEL_META_KEY), + Some(&"7".to_string()), + ); + } + + #[test] + fn test_resolve_dict_values_compression_metadata_env_fallback() { + let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata( + &HashMap::new(), + Some("zstd".to_string()), + Some("9".to_string()), + ); + assert_eq!( + metadata.get(COMPRESSION_META_KEY), + Some(&"zstd".to_string()), + ); + assert_eq!( + metadata.get(COMPRESSION_LEVEL_META_KEY), + Some(&"9".to_string()), + ); + } + + #[tokio::test] + async fn test_dictionary_encode_int64() { + use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY}; + use crate::testing::{TestCases, check_round_trip_encoding_of_data}; + use arrow_array::{ArrayRef, Int64Array}; + use std::collections::HashMap; + use std::sync::Arc; + + // Low cardinality with poor RLE opportunity. + let values = (0..1000) + .map(|i| match i % 3 { + 0 => 10i64, + 1 => 20i64, + _ => 30i64, + }) + .collect::>(); + let array = Arc::new(Int64Array::from(values)) as ArrayRef; + + let mut metadata = HashMap::new(); + metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_MINIBLOCK.to_string(), + ); + metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string()); + + let test_cases = TestCases::default() + .with_u32_structural_encodings() + .with_batch_size(1000) + .with_range(0..1000) + .with_indices(vec![0, 1, 10, 999]) + .with_expected_encoding("dictionary"); + + check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await; + } + + #[tokio::test] + async fn test_dictionary_encode_float64() { + use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY}; + use crate::testing::{TestCases, check_round_trip_encoding_of_data}; + use arrow_array::{ArrayRef, Float64Array}; + use std::collections::HashMap; + use std::sync::Arc; + + // Low cardinality with poor RLE opportunity. + let values = (0..1000) + .map(|i| match i % 3 { + 0 => 0.1f64, + 1 => 0.2f64, + _ => 0.3f64, + }) + .collect::>(); + let array = Arc::new(Float64Array::from(values)) as ArrayRef; + + let mut metadata = HashMap::new(); + metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_MINIBLOCK.to_string(), + ); + metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string()); + + let test_cases = TestCases::default() + .with_u32_structural_encodings() + .with_batch_size(1000) + .with_range(0..1000) + .with_indices(vec![0, 1, 10, 999]) + .with_expected_encoding("dictionary"); + + check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await; + } + + #[test] + fn test_miniblock_dictionary_out_of_line_bitpacking_decode() { + let rows = 10_000; + let unique_values = 2_000; + + let dictionary_encoding = + ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(11, None)); + let layout = pb21::MiniBlockLayout { + rep_compression: None, + def_compression: None, + value_compression: Some(ProtobufUtils21::flat(64, None)), + dictionary: Some(dictionary_encoding), + num_dictionary_items: unique_values, + layers: vec![pb21::RepDefLayer::RepdefAllValidItem as i32], + num_buffers: 1, + repetition_index_depth: 0, + num_items: rows, + has_large_chunk: false, + }; + + let buffer_offsets_and_sizes = vec![(0, 0), (0, 0), (0, 0)]; + let scheduler = super::MiniBlockScheduler::try_new( + &buffer_offsets_and_sizes, + /*priority=*/ 0, + /*items_in_page=*/ rows, + &layout, + &DefaultDecompressionStrategy::default(), + ) + .unwrap(); + + let dictionary = scheduler.dictionary.unwrap(); + assert_eq!(dictionary.num_dictionary_items, unique_values); + assert_eq!( + dictionary.dictionary_data_alignment, + crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT + ); + } + + // Dictionary encoding decision tests + fn create_test_fixed_data_block( + num_values: u64, + cardinality: u64, + bits_per_value: u64, + ) -> DataBlock { + assert!(cardinality > 0); + assert!(cardinality <= num_values); + let block_info = BlockInfo::default(); + + assert_eq!(bits_per_value % 8, 0); + let data = match bits_per_value { + 32 => { + let values = (0..num_values) + .map(|i| (i % cardinality) as u32) + .collect::>(); + crate::buffer::LanceBuffer::reinterpret_vec(values) + } + 64 => { + let values = (0..num_values).map(|i| i % cardinality).collect::>(); + crate::buffer::LanceBuffer::reinterpret_vec(values) + } + 128 => { + let values = (0..num_values) + .map(|i| (i % cardinality) as u128) + .collect::>(); + crate::buffer::LanceBuffer::reinterpret_vec(values) + } + _ => unreachable!(), + }; + DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value, + data, + num_values, + block_info, + }) + } + + /// Helper to create VariableWidth (string) test data block with exact cardinality + fn create_test_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock { + use arrow_array::StringArray; + + assert!(cardinality <= num_values && cardinality > 0); + + let mut values = Vec::with_capacity(num_values as usize); + for i in 0..num_values { + values.push(format!("value_{:016}", i % cardinality)); + } + + let array = StringArray::from(values); + DataBlock::from_array(Arc::new(array) as ArrayRef) + } + + fn create_sorted_string_array(num_values: u64, cardinality: u64) -> ArrayRef { + use arrow_array::StringArray; + + assert!(cardinality <= num_values && cardinality > 0); + + let mut values = Vec::with_capacity(num_values as usize); + for i in 0..num_values { + let value_idx = i * cardinality / num_values; + values.push(format!("value_{:016}", value_idx)); + } + + Arc::new(StringArray::from(values)) as ArrayRef + } + + fn create_sorted_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock { + DataBlock::from_array(create_sorted_string_array(num_values, cardinality)) + } + + #[test] + fn test_should_dictionary_encode() { + use crate::constants::DICT_SIZE_RATIO_META_KEY; + use lance_core::datatypes::Field as LanceField; + + // Create data where dict encoding saves space + let block = create_test_variable_width_block(1000, 10); + + let mut metadata = HashMap::new(); + metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string()); + let arrow_field = + arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata); + let field = LanceField::try_from(&arrow_field).unwrap(); + + let result = PrimitiveStructuralEncoder::should_dictionary_encode( + &block, + &field, + FixedWidthDictionaryEncoding::Exclude64Bit, + ); + + assert!( + result.is_some(), + "Should use dictionary encode based on size" + ); + } + + #[test] + fn test_block_sampling_detects_low_cardinality_in_short_sorted_runs() { + let sample_count: usize = 4096; + let num_values: u64 = 200_000; + let cardinality: u64 = 8_000; + let run_length = num_values / cardinality; + let stride = num_values as usize / sample_count; + assert!( + stride > run_length as usize, + "test must construct the stride > run_length case" + ); + + let block = create_sorted_variable_width_block(num_values, cardinality); + let sample_unique_ratio = + PrimitiveStructuralEncoder::sample_unique_ratio(&block, sample_count).unwrap(); + + assert!( + sample_unique_ratio.is_some_and(|ratio| ratio < 0.98), + "sorted low-cardinality data must not be classified as near-unique" + ); + } + + #[test] + fn test_should_dictionary_encode_sorted_low_cardinality() { + use crate::constants::DICT_SIZE_RATIO_META_KEY; + use lance_core::datatypes::Field as LanceField; + + let block = create_sorted_variable_width_block(200_000, 8_000); + + let mut metadata = HashMap::new(); + metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string()); + let arrow_field = + arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata); + let field = LanceField::try_from(&arrow_field).unwrap(); + + let result = PrimitiveStructuralEncoder::should_dictionary_encode( + &block, + &field, + FixedWidthDictionaryEncoding::Include64Bit, + ); + + assert!( + result.is_some(), + "sorted low-cardinality data should reach dictionary encoding" + ); + } + + #[test] + fn test_should_not_dictionary_encode_sorted_high_cardinality_short_runs() { + use crate::constants::DICT_SIZE_RATIO_META_KEY; + use lance_core::datatypes::Field as LanceField; + + let num_values = 200_002; + let cardinality = 100_001; + let block = create_sorted_variable_width_block(num_values, cardinality); + + let mut metadata = HashMap::new(); + metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string()); + let arrow_field = + arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata); + let field = LanceField::try_from(&arrow_field).unwrap(); + + let result = PrimitiveStructuralEncoder::should_dictionary_encode( + &block, + &field, + FixedWidthDictionaryEncoding::Include64Bit, + ); + + assert!( + result.is_none(), + "sorted high-cardinality short runs should not trigger a full dictionary probe" + ); + } + + #[tokio::test] + async fn test_encode_sorted_low_cardinality_uses_dictionary_layout() { + use crate::constants::DICT_SIZE_RATIO_META_KEY; + + let mut metadata = HashMap::new(); + metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string()); + let field = arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata); + let array = create_sorted_string_array(200_000, 8_000); + + let page = encode_first_page(field, array, TestEncoding::StructuralU32).await; + let _ = dictionary_encoding_from_page(&page); + } + + #[test] + fn test_should_not_dictionary_encode_unsupported_bits() { + use crate::constants::DICT_SIZE_RATIO_META_KEY; + use lance_core::datatypes::Field as LanceField; + + let block = create_test_fixed_data_block(1000, 1000, 32); + + let mut metadata = HashMap::new(); + metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string()); + let arrow_field = + arrow_schema::Field::new("test", DataType::Int32, false).with_metadata(metadata); + let field = LanceField::try_from(&arrow_field).unwrap(); + + let result = PrimitiveStructuralEncoder::should_dictionary_encode( + &block, + &field, + FixedWidthDictionaryEncoding::Exclude64Bit, + ); + + assert!( + result.is_none(), + "Should not use dictionary encode for unsupported bit width" + ); + } + + #[test] + fn test_should_not_dictionary_encode_near_unique_sample() { + use crate::constants::DICT_SIZE_RATIO_META_KEY; + use lance_core::datatypes::Field as LanceField; + + let num_values = 5000; + let block = create_test_variable_width_block(num_values, num_values); + + let mut metadata = HashMap::new(); + metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "1.0".to_string()); + let arrow_field = + arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata); + let field = LanceField::try_from(&arrow_field).unwrap(); + + let result = PrimitiveStructuralEncoder::should_dictionary_encode( + &block, + &field, + FixedWidthDictionaryEncoding::Exclude64Bit, + ); + + assert!( + result.is_none(), + "Should not probe dictionary encoding for near-unique data" + ); + } + + #[test] + fn test_v2_1_miniblock_serializes_log_num_values_15() { + let miniblocks = MiniBlockCompressed { + data: vec![LanceBuffer::from(vec![1_u8; 16])], + chunks: vec![ + MiniBlockChunk { + buffer_sizes: vec![8], + log_num_values: 15, + }, + MiniBlockChunk { + buffer_sizes: vec![8], + log_num_values: 0, + }, + ], + num_values: 32_769, + }; + + let serialized = PrimitiveStructuralEncoder::serialize_miniblocks( + miniblocks, + None, + None, + MiniblockChunkSize::U16, + ) + .unwrap(); + + let chunk_metadata = serialized.metadata.borrow_to_typed_slice::(); + assert_eq!(chunk_metadata.len(), 2); + assert_eq!( + chunk_metadata[0] & 0x0F, + 15, + "V2.1 metadata should use all 4 bits for log_num_values" + ); + } + + async fn encode_first_page( + field: arrow_schema::Field, + array: ArrayRef, + version: TestEncoding, + ) -> crate::encoder::EncodedPage { + use crate::repdef::RepDefBuilder; + use crate::{ + encoder::{ + ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers, + }, + testing::{create_test_field_encoder, test_encoding_strategy}, + }; + + let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap(); + let encoding_strategy = test_encoding_strategy(version); + let mut column_index_seq = ColumnIndexSequence::default(); + let encoding_options = EncodingOptions { + cache_bytes_per_column: 1, + max_page_bytes: 32 * 1024 * 1024, + keep_original_array: true, + buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT, + }; + + let mut encoder = create_test_field_encoder( + encoding_strategy.as_ref(), + &lance_field, + &mut column_index_seq, + &encoding_options, + ) + .unwrap(); + + let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT); + let repdef = RepDefBuilder::default(); + let num_rows = array.len() as u64; + let mut pages = Vec::new(); + for task in encoder + .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows) + .unwrap() + { + pages.push(task.await.unwrap()); + } + for task in encoder.flush(&mut external_buffers).unwrap() { + pages.push(task.await.unwrap()); + } + pages.into_iter().next().unwrap() + } + + #[tokio::test] + async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() { + use crate::format::pb21::page_layout::Layout; + + let val = vec![0xABu8; 33]; + let arr: ArrayRef = Arc::new( + arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size( + std::iter::repeat_n(Some(val.as_slice()), 256), + 33, + ) + .unwrap(), + ); + let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true); + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; + + let PageEncoding::Structural(layout) = &page.description else { + panic!("Expected structural encoding"); + }; + let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else { + panic!("Expected constant layout in slot 2"); + }; + assert!(layout.inline_value.is_none()); + assert_eq!(page.data.len(), 1); + + let test_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralU32) + .with_page_sizes(vec![4096]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_constant_layout_out_of_line_utf8_v2_2() { + use crate::format::pb21::page_layout::Layout; + + let arr: ArrayRef = Arc::new(arrow_array::StringArray::from_iter_values( + std::iter::repeat_n("hello", 512), + )); + let field = arrow_schema::Field::new("c", DataType::Utf8, true); + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; + + let PageEncoding::Structural(layout) = &page.description else { + panic!("Expected structural encoding"); + }; + let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else { + panic!("Expected constant layout in slot 2"); + }; + assert!(layout.inline_value.is_none()); + assert_eq!(page.data.len(), 1); + + let test_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralU32) + .with_page_sizes(vec![4096]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_constant_layout_nullable_item_v2_2() { + use crate::format::pb21::page_layout::Layout; + + let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![ + Some(7), + None, + Some(7), + None, + Some(7), + ])); + let field = arrow_schema::Field::new("c", DataType::Int32, true); + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; + + let PageEncoding::Structural(layout) = &page.description else { + panic!("Expected structural encoding"); + }; + let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else { + panic!("Expected constant layout in slot 2"); + }; + assert!(layout.inline_value.is_some()); + assert_eq!(page.data.len(), 2); + + let test_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralU32) + .with_page_sizes(vec![4096]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_constant_layout_list_repdef_v2_2() { + use crate::format::pb21::page_layout::Layout; + use arrow_array::builder::{Int32Builder, ListBuilder}; + + let mut builder = ListBuilder::new(Int32Builder::new()); + builder.values().append_value(7); + builder.values().append_null(); + builder.values().append_value(7); + builder.append(true); + + builder.append(true); + + builder.values().append_value(7); + builder.append(true); + + builder.append_null(); + + let arr: ArrayRef = Arc::new(builder.finish()); + let field = arrow_schema::Field::new( + "c", + DataType::List(Arc::new(arrow_schema::Field::new( + "item", + DataType::Int32, + true, + ))), + true, + ); + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; + + let PageEncoding::Structural(layout) = &page.description else { + panic!("Expected structural encoding"); + }; + let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else { + panic!("Expected constant layout in slot 2"); + }; + assert!(layout.inline_value.is_some()); + assert_eq!(page.data.len(), 2); + + let test_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralU32) + .with_page_sizes(vec![4096]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_constant_layout_fixed_size_list_not_used_v2_2() { + use crate::format::pb21::page_layout::Layout; + use arrow_array::builder::{FixedSizeListBuilder, Int32Builder}; + + let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3); + for _ in 0..64 { + builder.values().append_value(1); + builder.values().append_null(); + builder.values().append_value(3); + builder.append(true); + } + let arr: ArrayRef = Arc::new(builder.finish()); + let field = arrow_schema::Field::new( + "c", + DataType::FixedSizeList( + Arc::new(arrow_schema::Field::new("item", DataType::Int32, true)), + 3, + ), + true, + ); + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; + + if let PageEncoding::Structural(layout) = &page.description { + assert!( + !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)), + "FixedSizeList should not use constant layout yet" + ); + } + + let test_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralU32) + .with_page_sizes(vec![4096]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_constant_layout_not_written_before_v2_2() { + use crate::format::pb21::page_layout::Layout; + + let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024])); + let field = arrow_schema::Field::new("c", DataType::Int32, true); + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU16).await; + + let PageEncoding::Structural(layout) = &page.description else { + return; + }; + assert!( + !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)), + "Should not emit constant layout before v2.2" + ); + + let test_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralU16) + .with_page_sizes(vec![4096]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_all_null_constant_layout_still_works_v2_2() { + use crate::format::pb21::page_layout::Layout; + + let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None])); + let field = arrow_schema::Field::new("c", DataType::Int32, true); + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; + + let PageEncoding::Structural(layout) = &page.description else { + panic!("Expected structural encoding"); + }; + let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else { + panic!("Expected layout in slot 2"); + }; + assert!(layout.inline_value.is_none()); + assert_eq!(page.data.len(), 0); + + let test_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralU32) + .with_page_sizes(vec![4096]); + check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; + } + + fn hand_built_dictionary_with_out_of_range_null_keys() -> ArrayRef { + use arrow_array::{DictionaryArray, Int32Array, types::Int32Type}; + use arrow_buffer::NullBuffer; + + let keys = Int32Array::new( + vec![0, 7, 7].into(), + Some(NullBuffer::from(vec![true, false, false])), + ); + let values = Arc::new(StringArray::from(vec!["a"])); + Arc::new(DictionaryArray::::try_new(keys, values).unwrap()) as ArrayRef + } + + fn concatenated_dictionary_with_out_of_range_null_keys() -> ArrayRef { + use arrow_array::{builder::StringDictionaryBuilder, new_null_array, types::Int32Type}; + + let mut builder = StringDictionaryBuilder::::new(); + builder.append_value("a"); + for _ in 0..7 { + builder.append_null(); + } + let valued = Arc::new(builder.finish()) as ArrayRef; + let all_null = new_null_array(valued.data_type(), 8); + arrow_select::concat::concat(&[valued.as_ref(), all_null.as_ref()]).unwrap() + } + + #[rstest::rstest] + #[case::hand_built(hand_built_dictionary_with_out_of_range_null_keys(), 1)] + #[case::concatenated(concatenated_dictionary_with_out_of_range_null_keys(), 8)] + #[tokio::test] + async fn test_dictionary_out_of_range_null_keys_round_trip( + #[case] dictionary: ArrayRef, + #[case] split_at: usize, + ) { + let num_rows = dictionary.len(); + let dictionaries = vec![ + dictionary.slice(0, split_at), + dictionary.slice(split_at, num_rows - split_at), + ]; + let test_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralU32) + .with_page_sizes(vec![1]) + .with_range(split_at as u64..num_rows as u64) + .with_indices(vec![split_at as u64]); + + check_round_trip_encoding_of_data(dictionaries, &test_cases, HashMap::new()).await; + } + + #[test] + fn test_encode_decode_complex_all_null_vals_roundtrip() { + use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy}; + + let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::>()); + + let compression_strategy = crate::testing::test_compression_strategy( + TestEncoding::StructuralU16, + crate::compression_config::CompressionParams::default(), + ); + let decompression_strategy = DefaultDecompressionStrategy::default(); + + let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals( + &values, + compression_strategy.as_ref(), + ) + .unwrap(); + + let decompressor = decompression_strategy + .create_block_decompressor(&encoding) + .unwrap(); + let decompressed = decompressor + .decompress(Some(compressed_buf), values.len() as u64) + .unwrap(); + let decompressed_fixed_width = decompressed.as_fixed_width().unwrap(); + assert_eq!(decompressed_fixed_width.num_values, values.len() as u64); + assert_eq!(decompressed_fixed_width.bits_per_value, 16); + let rep_result = decompressed_fixed_width.data.borrow_to_typed_slice::(); + assert_eq!(rep_result.as_ref(), values.as_ref()); + } + + #[tokio::test] + async fn test_complex_all_null_compression_gated_by_version() { + use crate::format::pb21::page_layout::Layout; + use arrow_array::ListArray; + + let list_array = ListArray::from_iter_primitive::( + (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }), + ); + let arr: ArrayRef = Arc::new(list_array); + let field = arrow_schema::Field::new( + "c", + DataType::List(Arc::new(arrow_schema::Field::new( + "item", + DataType::Int32, + true, + ))), + true, + ); + + let page_v21 = + encode_first_page(field.clone(), arr.clone(), TestEncoding::StructuralU16).await; + let PageEncoding::Structural(layout_v21) = &page_v21.description else { + panic!("Expected structural encoding"); + }; + let Layout::ConstantLayout(layout_v21) = layout_v21.layout.as_ref().unwrap() else { + panic!("Expected constant layout"); + }; + assert!(layout_v21.rep_compression.is_none()); + assert!(layout_v21.def_compression.is_none()); + assert_eq!(layout_v21.num_rep_values, 0); + assert_eq!(layout_v21.num_def_values, 0); + + let page_v22 = encode_first_page(field, arr, TestEncoding::StructuralU32).await; + let PageEncoding::Structural(layout_v22) = &page_v22.description else { + panic!("Expected structural encoding"); + }; + let Layout::ConstantLayout(layout_v22) = layout_v22.layout.as_ref().unwrap() else { + panic!("Expected constant layout"); + }; + assert!(layout_v22.def_compression.is_some()); + assert!(layout_v22.num_def_values > 0); + } + + #[tokio::test] + async fn test_complex_all_null_round_trip() { + use arrow_array::ListArray; + + let list_array = ListArray::from_iter_primitive::( + (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }), + ); + + let test_cases = TestCases::default().with_u32_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) + .await; + } + + #[tokio::test] + async fn test_complex_all_null_constant_def_round_trip() { + use arrow_array::ListArray; + + // Every row is a null list => constant def levels => a single RLE run, + // exercising the lazy run-form decode end to end. + let list_array = ListArray::from_iter_primitive::( + (0..5000).map(|_| None::>>), + ); + + let test_cases = TestCases::default().with_u32_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) + .await; + } + + fn encoded_u16_frame(levels: &[u16], run_length_width: RunLengthWidth) -> LanceBuffer { + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_slice(Arc::from(levels)), + bits_per_value: 16, + num_values: levels.len() as u64, + block_info: BlockInfo::new(), + }); + BlockCompressor::compress(&RleEncoder::with_run_length_width(run_length_width), block) + .unwrap() + .0 + .unwrap() + } + + fn encoded_u16_runs(levels: &[u16], run_length_width: RunLengthWidth) -> RleRuns { + let frame = encoded_u16_frame(levels, run_length_width); + RleDecompressor::with_run_length_width(16, run_length_width) + .decode_u16_runs(frame, levels.len() as u64) + .unwrap() + } + + #[test] + fn miniblock_levels_use_one_count_for_mixed_codecs() { + let actual_num_levels = usize::from(u16::MAX) + 8; + let levels = vec![1_u16; actual_num_levels]; + let rep_frame = encoded_u16_frame(&levels, RunLengthWidth::U8); + let rep_decompressor = RleDecompressor::new(16); + let def_frame = LanceBuffer::reinterpret_slice(Arc::from(levels.clone())); + let def_decompressor = ValueDecompressor::from_flat(&pb21::Flat { + bits_per_value: 16, + data: None, + }); + + let num_levels = DecodeMiniBlockTask::resolve_num_levels( + Some(&rep_decompressor), + Some(&rep_frame), + Some(&def_decompressor), + Some(&def_frame), + 7, + ) + .unwrap(); + assert_eq!(num_levels, actual_num_levels as u64); + + let rep = + DecodeMiniBlockTask::decode_levels(&rep_decompressor, rep_frame, num_levels).unwrap(); + let def = + DecodeMiniBlockTask::decode_levels(&def_decompressor, def_frame, num_levels).unwrap(); + assert_eq!(rep.as_ref(), levels); + assert_eq!(def, rep); + } + + #[test] + fn miniblock_levels_reject_non_wrapped_count_mismatch() { + let frame = encoded_u16_frame(&[1_u16; 8], RunLengthWidth::U8); + let decompressor = RleDecompressor::new(16); + + let error = DecodeMiniBlockTask::resolve_num_levels( + Some(&decompressor), + Some(&frame), + None, + None, + 7, + ) + .unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!(error.to_string().contains("not congruent modulo 65536")); + } + + #[test] + fn miniblock_levels_reject_cross_stream_count_disagreement() { + let rep_levels = vec![1_u16; usize::from(u16::MAX) + 8]; + let def_levels = vec![1_u16; rep_levels.len() + usize::from(u16::MAX) + 1]; + let rep_frame = encoded_u16_frame(&rep_levels, RunLengthWidth::U8); + let def_frame = encoded_u16_frame(&def_levels, RunLengthWidth::U8); + let decompressor = RleDecompressor::new(16); + + let error = DecodeMiniBlockTask::resolve_num_levels( + Some(&decompressor), + Some(&rep_frame), + Some(&decompressor), + Some(&def_frame), + 7, + ) + .unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("structural streams disagree on the level count") + ); + } + + fn physical_levels(levels: &[u16]) -> LazyLevels { + LazyLevels::Runs(Arc::new(RunStorage::Physical( + encoded_u16_runs(levels, RunLengthWidth::U8).into_owned(), + ))) + } + + fn coalesced_levels(levels: &[u16]) -> LazyLevels { + let mut values = Vec::new(); + let mut ends = RunEndsBuilder::with_capacity(levels.len(), levels.len()); + for (index, &value) in levels.iter().enumerate() { + if values.last() == Some(&value) { + ends.set_last(index + 1).unwrap(); + } else { + values.push(value); + ends.push(index + 1).unwrap(); + } + } + LazyLevels::Runs(Arc::new(RunStorage::Coalesced { + values: values.into_boxed_slice(), + ends: ends.finish(), + })) + } + + #[test] + fn lazy_levels_runs_match_dense() { + // Runs: 3x2, 1x1, 3x3, 0x2 => [3,3,1,3,3,3,0,0] + let expanded: Vec = vec![3, 3, 1, 3, 3, 3, 0, 0]; + let coalesced = coalesced_levels(&expanded); + let physical = physical_levels(&expanded); + let dense = LazyLevels::Dense(ScalarBuffer::::from(expanded.clone())); + let n = expanded.len(); + + assert_eq!(coalesced.len(), n); + assert_eq!(physical.len(), n); + assert_eq!(dense.len(), n); + + // Rows begin at each `max_rep` (3) position; row `num_rows` maps to `len`. + let max_rep = 3u16; + let row_starts: Vec = (0..n).filter(|&i| expanded[i] == max_rep).collect(); + for target in 0..=row_starts.len() as u64 { + let want = row_starts.get(target as usize).copied().unwrap_or(n); + for runs in [&coalesced, &physical] { + let mut cursor = LevelCursor::default(); + assert_eq!( + runs.seek_row_start(&mut cursor, target, max_rep).unwrap(), + want, + "seek_row_start({target})" + ); + } + let mut c_dense = LevelCursor::default(); + assert_eq!( + dense.seek_row_start(&mut c_dense, target, max_rep).unwrap(), + want + ); + } + + // `count_le_cursor` (fresh cursor per range) and `extend_into` agree with + // the dense reference on every sub-range. + for start in 0..=n { + for end in start..=n { + for max in [0u16, 1, 2, 3] { + let want = expanded[start..end].iter().filter(|&&d| d <= max).count() as u64; + for runs in [&coalesced, &physical] { + let mut cursor = RunPosition::default(); + assert_eq!( + runs.count_le_cursor(&mut cursor, start..end, max).0, + want, + "count_le_cursor({start}..{end}, {max})" + ); + } + let mut d_cur = RunPosition::default(); + assert_eq!(dense.count_le_cursor(&mut d_cur, start..end, max).0, want); + } + for runs in [&coalesced, &physical] { + let mut got = Vec::new(); + runs.extend_into(start..end, RunPosition::default(), &mut got); + assert_eq!( + got, + expanded[start..end].to_vec(), + "extend_into({start}..{end})" + ); + } + let mut got_dense = Vec::new(); + dense.extend_into(start..end, RunPosition::default(), &mut got_dense); + assert_eq!(got_dense, expanded[start..end].to_vec()); + } + } + } + + #[test] + fn physical_run_hints_support_deferred_materialization() { + let expanded: Vec = vec![3, 3, 1, 1, 2, 2, 0, 0]; + let physical = physical_levels(&expanded); + let LazyLevels::Runs(runs) = &physical else { + panic!("expected physical runs"); + }; + let mut first_hint = RunPosition::default(); + runs.seek(&mut first_hint, 2); + let mut second_hint = RunPosition::default(); + runs.seek(&mut second_hint, 6); + + let mut second = Vec::new(); + physical.extend_into(6..8, second_hint, &mut second); + let mut first = Vec::new(); + physical.extend_into(2..4, first_hint, &mut first); + assert_eq!(second, expanded[6..8]); + assert_eq!(first, expanded[2..4]); + } + + /// Fuzz parity for the run-oriented complex-all-null drain: the cursor walk + /// over `LazyLevels` must yield the exact level slices and visible + /// count that a brute-force reference over the fully expanded levels does, for + /// dense, physical-run, and coalesced-run forms and arbitrarily shaped range requests. + mod complex_all_null_drain_parity { + use std::ops::Range; + + use arrow_buffer::ScalarBuffer; + use proptest::prelude::*; + + use super::super::{LazyLevels, LevelCursor, RunPosition}; + use super::{coalesced_levels, physical_levels}; + use crate::Result; + + #[derive(Debug, Clone)] + struct DrainInput { + max_rep: u16, + max_visible: u16, + rep: Option>, + def: Option>, + ranges: Vec>, + } + + fn dense_levels(levels: &[u16]) -> LazyLevels { + LazyLevels::Dense(ScalarBuffer::from(levels.to_vec())) + } + + fn rle_levels(levels: &[u16]) -> LazyLevels { + coalesced_levels(levels) + } + + fn seek( + rep: Option<&LazyLevels>, + cursor: &mut LevelCursor, + row: u64, + max_rep: u16, + ) -> Result { + match rep { + Some(rep) => rep.seek_row_start(cursor, row, max_rep), + None => { + cursor.row = row; + cursor.level = row as usize; + Ok(row as usize) + } + } + } + + /// Mirror of `ComplexAllNullPageDecoder::drain`, driving the real + /// `seek_row_start` / `count_le_cursor` with monotonic cursors. + fn simulate_drain( + rep: Option<&LazyLevels>, + def: Option<&LazyLevels>, + max_rep: u16, + max_visible: u16, + ranges: &[Range], + ) -> Result<(Vec>, u64)> { + let mut rep_cursor = LevelCursor::default(); + let mut def_run_cursor = RunPosition::default(); + let mut slices: Vec> = Vec::new(); + let mut visible = 0u64; + for range in ranges { + let level_start = seek(rep, &mut rep_cursor, range.start, max_rep)?; + let level_end = seek(rep, &mut rep_cursor, range.end, max_rep)?; + visible += match def { + Some(def) => { + def.count_le_cursor( + &mut def_run_cursor, + level_start..level_end, + max_visible, + ) + .0 + } + None => (level_end - level_start) as u64, + }; + match slices.last_mut() { + Some(last) if last.end == level_start => last.end = level_end, + _ => slices.push(level_start..level_end), + } + } + Ok((slices, visible)) + } + + /// Independent brute-force reference over fully expanded levels. + fn reference_drain( + rep: Option<&[u16]>, + def: Option<&[u16]>, + max_rep: u16, + max_visible: u16, + ranges: &[Range], + ) -> (Vec>, u64) { + let total_levels = rep + .map(|r| r.len()) + .or_else(|| def.map(|d| d.len())) + .unwrap_or(0); + // Level index where each row starts (or `total_levels` for the end row). + let row_starts: Vec = match rep { + Some(rep) => (0..rep.len()).filter(|&i| rep[i] == max_rep).collect(), + None => (0..total_levels).collect(), + }; + let level_of_row = |row: u64| { + row_starts + .get(row as usize) + .copied() + .unwrap_or(total_levels) + }; + + let mut slices: Vec> = Vec::new(); + let mut visible = 0u64; + for range in ranges { + let ls = level_of_row(range.start); + let le = level_of_row(range.end); + visible += match def { + Some(def) => def[ls..le].iter().filter(|&&d| d <= max_visible).count() as u64, + None => (le - ls) as u64, + }; + match slices.last_mut() { + Some(last) if last.end == ls => last.end = le, + _ => slices.push(ls..le), + } + } + (slices, visible) + } + + fn ranges_strategy(num_rows: u64) -> BoxedStrategy>> { + if num_rows == 0 { + return Just(Vec::new()).boxed(); + } + // (gap, len) pairs; a zero gap yields ranges adjacent in row space, + // which exercises the level-slice coalescing path. + proptest::collection::vec((0u64..=3, 1u64..=4), 0..=8) + .prop_map(move |pairs| { + let mut ranges = Vec::new(); + let mut pos = 0u64; + for (gap, len) in pairs { + pos = pos.saturating_add(gap); + if pos >= num_rows { + break; + } + let end = (pos + len).min(num_rows); + ranges.push(pos..end); + pos = end; + } + ranges + }) + .boxed() + } + + fn drain_input() -> impl Strategy { + ( + 1u16..=3, + 0u16..=3, + any::(), + any::(), + 1usize..=48, + ) + .prop_flat_map(|(max_rep, max_visible, has_rep, has_def, len)| { + // Complex-all-null always has definition levels when there is + // no repetition, so force `def` present in that case. + let has_def = has_def || !has_rep; + let rep = if has_rep { + proptest::collection::vec(0u16..=max_rep, len) + .prop_map(move |mut v| { + // Row 0 must start at a max-rep boundary. + v[0] = max_rep; + Some(v) + }) + .boxed() + } else { + Just(None).boxed() + }; + let def = if has_def { + proptest::collection::vec(0u16..=(max_visible + 2), len) + .prop_map(Some) + .boxed() + } else { + Just(None).boxed() + }; + (Just(max_rep), Just(max_visible), rep, def) + }) + .prop_flat_map(|(max_rep, max_visible, rep, def)| { + let num_rows = match &rep { + Some(rep) => rep.iter().filter(|&&v| v == max_rep).count() as u64, + None => def.as_ref().map(|d| d.len() as u64).unwrap_or(0), + }; + ranges_strategy(num_rows).prop_map(move |ranges| DrainInput { + max_rep, + max_visible, + rep: rep.clone(), + def: def.clone(), + ranges, + }) + }) + } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(512))] + + #[test] + fn drain_matches_reference(input in drain_input()) { + let DrainInput { max_rep, max_visible, rep, def, ranges } = input; + + let reference = + reference_drain(rep.as_deref(), def.as_deref(), max_rep, max_visible, &ranges); + + let rep_dense = rep.as_deref().map(dense_levels); + let def_dense = def.as_deref().map(dense_levels); + let got_dense = + simulate_drain(rep_dense.as_ref(), def_dense.as_ref(), max_rep, max_visible, &ranges) + .unwrap(); + prop_assert_eq!(&got_dense, &reference, "dense form diverged from reference"); + + let rep_rle = rep.as_deref().map(rle_levels); + let def_rle = def.as_deref().map(rle_levels); + let got_rle = + simulate_drain(rep_rle.as_ref(), def_rle.as_ref(), max_rep, max_visible, &ranges) + .unwrap(); + prop_assert_eq!(&got_rle, &reference, "rle form diverged from reference"); + + let rep_physical = rep.as_deref().map(physical_levels); + let def_physical = def.as_deref().map(physical_levels); + let got_physical = + simulate_drain(rep_physical.as_ref(), def_physical.as_ref(), max_rep, max_visible, &ranges) + .unwrap(); + prop_assert_eq!(&got_physical, &reference, "physical form diverged from reference"); + } + } + } + + #[test] + fn lazy_levels_runs_are_compact() { + let single_run = |n: usize| { + let mut ends = RunEndsBuilder::with_capacity(n, 1); + ends.push(n).unwrap(); + LazyLevels::Runs(Arc::new(RunStorage::Coalesced { + values: vec![1u16].into_boxed_slice(), + ends: ends.finish(), + })) + }; + // Run-form footprint is independent of the logical length within an end width... + assert_eq!(single_run(100).deep_size(), single_run(10_000).deep_size()); + assert!(single_run(10_000_000).deep_size() < 100); + assert_eq!(single_run(10_000_000).len(), 10_000_000); + // ...while Dense pays 2 bytes per value. + assert_eq!( + LazyLevels::Dense(ScalarBuffer::::from(vec![1u16; 1000])).deep_size(), + 2000 + ); + } + + #[test] + fn lazy_levels_selects_smallest_representation() { + let runs = encoded_u16_runs(&[7u16; 10], RunLengthWidth::U8); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense); + + let equal_size: Vec = std::iter::repeat_n(0, 256) + .chain(std::iter::repeat_n(1, 100)) + .chain(std::iter::repeat_n(2, 100)) + .collect(); + let runs = encoded_u16_runs(&equal_size, RunLengthWidth::U8); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced); + + let moderate_runs: Vec = (0..250) + .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4)) + .collect(); + let runs = encoded_u16_runs(&moderate_runs, RunLengthWidth::U8); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical); + + let split_constant = vec![7u16; 5000]; + let runs = encoded_u16_runs(&split_constant, RunLengthWidth::U8); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced); + + let high_density: Vec = (0..70_000).map(|index| (index % 2) as u16).collect(); + let runs = encoded_u16_runs(&high_density, RunLengthWidth::U32); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense); + } + + #[test] + fn physical_runs_detach_from_large_encoded_frame() { + let levels: Vec = (0..250) + .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4)) + .collect(); + let frame = encoded_u16_frame(&levels, RunLengthWidth::U8); + let frame_offset = 4096; + let mut allocation = vec![0; frame_offset + frame.len() + 1_000_000]; + allocation[frame_offset..frame_offset + frame.len()].copy_from_slice(frame.as_ref()); + let frame = LanceBuffer::from(allocation).slice_with_length(frame_offset, frame.len()); + let runs = RleDecompressor::with_run_length_width(16, RunLengthWidth::U8) + .decode_u16_runs(frame, levels.len() as u64) + .unwrap(); + + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical); + let cached = LazyLevels::from_rle_runs(runs).unwrap(); + assert!( + matches!(cached, LazyLevels::Runs(ref runs) if matches!(runs.as_ref(), RunStorage::Physical(_))) + ); + assert_eq!(cached.len(), levels.len()); + assert!(cached.deep_size() < 4096); + } + + #[test] + fn complex_all_null_levels_reject_invalid_values_and_lengths() { + let invalid_levels = vec![0u16, 3]; + for levels in [ + LazyLevels::Dense(ScalarBuffer::from(invalid_levels.clone())), + physical_levels(&invalid_levels), + coalesced_levels(&invalid_levels), + ] { + let error = validate_complex_all_null_levels(&None, &Some(levels), 0, 2).unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!(error.to_string().contains("Invalid definition level 3")); + } + + let rep = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16; 2]))); + let def = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16]))); + let error = validate_complex_all_null_levels(&rep, &def, 0, 0).unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("repetition has 2, definition has 1") + ); + } + + #[test] + fn block_levels_reject_malformed_payload_size() { + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(vec![0]), + bits_per_value: 16, + num_values: 1, + block_info: BlockInfo::new(), + }); + let error = dense_levels_from_block(block, 1, "definition").unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("expected 2 bytes for 1 values, got 1") + ); + } + + #[test] + fn complex_all_null_level_codec_validates_rle_metadata() { + let encoding = pb21::CompressiveEncoding { + compression: Some(Compression::Rle(Box::new(pb21::Rle { + values: None, + run_lengths: Some(Box::new(ProtobufUtils21::flat(8, None))), + }))), + }; + + let error = LevelCodec::try_new(Some(&encoding), &DefaultDecompressionStrategy::default()) + .unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("RLE compression missing values encoding") + ); + } + + // https://github.com/lance-format/lance/issues/6681 + #[tokio::test] + async fn test_sparse_boolean_list_roundtrip() { + use arrow_array::builder::{BooleanBuilder, ListBuilder}; + + let mut list_builder = ListBuilder::new(BooleanBuilder::new()); + for i in 0..1000i32 { + if i % 64 == 0 { + // Alternate true/false so the array is not constant (constant path avoids the bug). + list_builder.values().append_value(i % 128 == 0); + list_builder.append(true); + } else { + list_builder.append(false); + } + } + let list_array = Arc::new(list_builder.finish()); + + let test_cases = TestCases::default().with_structural_encodings(); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + + fn truncated_tail_details() -> std::sync::Arc { + use crate::compression::VariablePerValueDecompressor; + use crate::encodings::physical::binary::VariableDecoder; + use crate::repdef::{ControlWordParser, DefinitionInterpretation}; + use std::sync::Arc; + Arc::new(super::FullZipDecodeDetails { + value_decompressor: super::PerValueDecompressor::Variable(Arc::new( + VariableDecoder::default(), + ) + as Arc), + def_meaning: vec![DefinitionInterpretation::NullableItem].into(), + ctrl_word_parser: ControlWordParser::new(0, 0), + max_rep: 0, + max_visible_def: 0, + }) + } + + fn decode_variable_full_zip( + buf: Vec, + bits_per_offset: u8, + ) -> lance_core::Result { + use std::collections::VecDeque; + let mut data = VecDeque::new(); + data.push_back(crate::buffer::LanceBuffer::from(buf)); + super::VariableFullZipDecoder::new( + truncated_tail_details(), + data, + 1, + bits_per_offset, + bits_per_offset, + ) + } + + /// A well-formed length prefix decodes without incident, for both widths. + #[test] + fn variable_full_zip_wellformed_length_prefix() { + assert!(decode_variable_full_zip(0u32.to_le_bytes().to_vec(), 32).is_ok()); + assert!(decode_variable_full_zip(0u64.to_le_bytes().to_vec(), 64).is_ok()); + } + + /// A page whose item walk ends with a partial length prefix must surface a + /// corrupt-file error rather than read past the end of the buffer. + /// + /// This asserts the error variant and message rather than merely expecting a + /// panic: before the length prefix was bounds checked, the read was + /// `get_unchecked` behind a `debug_assert!`, so a debug build panicked here + /// (which a `#[should_panic]` test would have accepted as a pass) while a + /// release build read up to 8 bytes out of a 4 byte allocation. + #[test] + fn variable_full_zip_truncated_length_prefix_is_corrupt_file() { + use lance_core::Error; + + for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] { + let err = decode_variable_full_zip(vec![0xAA; buf_len], bits) + .expect_err("a truncated length prefix must not decode"); + assert!( + matches!(err, Error::CorruptFile { .. }), + "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}", + bits, + buf_len, + err + ); + let msg = err.to_string(); + assert!( + msg.contains("truncated length prefix"), + "error should say what is wrong, got: {msg}" + ); + } + } + + /// Drains `num_rows` from a single-buffer fixed full-zip page built from `buf`. + /// + /// `bits_rep` picks the control word width (1 -> one byte, 9 -> two bytes) and each + /// visible item is a one-word control plus a 4 byte value. + fn drain_fixed_full_zip(buf: Vec, bits_rep: u8, num_rows: u64) -> lance_core::Result<()> { + use crate::compression::FixedPerValueDecompressor; + use crate::decoder::StructuralPageDecoder; + use crate::repdef::{ControlWordParser, DefinitionInterpretation}; + use std::collections::VecDeque; + use std::sync::Arc; + + #[derive(Debug)] + struct UnusedFixedDecompressor; + + impl FixedPerValueDecompressor for UnusedFixedDecompressor { + fn decompress(&self, _: FixedWidthDataBlock, _: u64) -> crate::Result { + unreachable!("these cases fail while slicing, before any value is decoded") + } + + fn bits_per_value(&self) -> u64 { + 32 + } + } + + let details = Arc::new(super::FullZipDecodeDetails { + value_decompressor: super::PerValueDecompressor::Fixed(Arc::new( + UnusedFixedDecompressor, + )), + def_meaning: vec![DefinitionInterpretation::AllValidList].into(), + ctrl_word_parser: ControlWordParser::new(bits_rep, 0), + max_rep: 1, + max_visible_def: 0, + }); + let bytes_per_word = bits_rep.div_ceil(8) as usize; + + let mut data = VecDeque::new(); + data.push_back(crate::buffer::LanceBuffer::from(buf)); + let mut decoder = super::FixedFullZipDecoder { + details, + data, + offset_in_current: 0, + bytes_per_value: 4, + total_bytes_per_value: 4 + bytes_per_word, + num_rows, + }; + decoder.drain(num_rows).map(|_| ()) + } + + /// Two one-item rows, each a control word plus a 4 byte value, slice cleanly. + #[test] + fn fixed_full_zip_wellformed_page_slices() { + assert!(drain_fixed_full_zip(vec![1, 0, 0, 0, 0, 1, 0, 0, 0, 0], 1, 2).is_ok()); + } + + /// The item walk is driven by the page's own control words, so a page that disagrees + /// with its layout can run off the end of the buffer. Each way of doing so must + /// surface a corrupt-file error. + /// + /// These assert the error variant and message rather than merely expecting a panic: + /// before the walk was bounds checked these aborted the process inside + /// `LanceBuffer::slice_with_length`, which a `#[should_panic]` test would have + /// accepted as a pass. + #[rstest::rstest] + // Second value is one byte short of a full 4 byte value + #[case::truncated_item(vec![1, 0, 0, 0, 0, 1, 0, 0, 0], 1, 2, "truncated item")] + // Two byte control words, with a lone trailing byte after the first item + #[case::truncated_control_word( + vec![1, 0, 0, 0, 0, 0, 1], + 9, + 2, + "truncated control word" + )] + // A well-formed two row page cannot satisfy a request for three rows + #[case::page_exhausted(vec![1, 0, 0, 0, 0, 1, 0, 0, 0, 0], 1, 3, "page data ran out")] + fn fixed_full_zip_malformed_page_is_corrupt_file( + #[case] buf: Vec, + #[case] bits_rep: u8, + #[case] num_rows: u64, + #[case] expected_msg: &str, + ) { + use lance_core::Error; + + let err = drain_fixed_full_zip(buf, bits_rep, num_rows) + .expect_err("a malformed page must not slice"); + assert!( + matches!(err, Error::CorruptFile { .. }), + "expected CorruptFile, got: {err:?}" + ); + let msg = err.to_string(); + assert!( + msg.contains(expected_msg), + "error should say what is wrong, got: {msg}" + ); + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive/blob.rs b/vendor/lance-encoding/src/encodings/logical/primitive/blob.rs new file mode 100644 index 000000000..d9ed1b390 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive/blob.rs @@ -0,0 +1,552 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Routines for decoding blob data +//! +//! The blob structural encoding is a structural encoding where the values (blobs) are stored +//! out-of-line in the file. The page contains the descriptions, encoded using some other layout. + +use std::{collections::VecDeque, ops::Range, sync::Arc}; + +use arrow_array::{Array, UInt64Array, cast::AsArray, make_array}; +use bytes::Bytes; +use futures::{FutureExt, future::BoxFuture}; + +use lance_core::{ + Error, Result, cache::DeepSizeOf, datatypes::BLOB_DESC_TYPE, error::LanceOptionExt, +}; + +use crate::{ + EncodingsIo, + buffer::LanceBuffer, + data::{BlockInfo, DataBlock, VariableWidthBlock}, + decoder::{DecodePageTask, DecodedPage, StructuralPageDecoder}, + encodings::logical::primitive::{ + CachedPageData, PageCacheView, PageInitialization, PageInitializationBuffers, PageLoadTask, + StructuralPageScheduler, + }, + repdef::{DefinitionInterpretation, RepDefUnraveler}, +}; + +/// How many bytes to target in each unloaded / loaded shard. A larger value means +/// we buffer more data in memory / make bigger requests to the I/O scheduler while +/// a smaller value means more requests to the I/O scheduler. +/// +/// This is probably a reasonable default for most cases. +pub const TARGET_SHARD_SIZE: u64 = 32 * 1024 * 1024; + +#[derive(Debug)] +pub(super) struct BlobDescriptionPageScheduler { + inner_scheduler: Box, + def_meaning: Arc<[DefinitionInterpretation]>, +} + +impl BlobDescriptionPageScheduler { + pub fn new( + inner_scheduler: Box, + def_meaning: Arc<[DefinitionInterpretation]>, + ) -> Self { + Self { + inner_scheduler, + def_meaning, + } + } + + fn wrap_decoder_fut( + decoder_fut: BoxFuture<'static, Result>>, + def_meaning: Arc<[DefinitionInterpretation]>, + ) -> BoxFuture<'static, Result>> { + async move { + let decoder = decoder_fut.await?; + Ok( + Box::new(BlobDescriptionPageDecoder::new(decoder, def_meaning)) + as Box, + ) + } + .boxed() + } +} + +impl StructuralPageScheduler for BlobDescriptionPageScheduler { + fn needs_initialization(&self) -> bool { + self.inner_scheduler.needs_initialization() + } + + fn init_layout(&self) -> Result { + self.inner_scheduler.init_layout() + } + + fn init_from_buffers<'a>( + &'a mut self, + buffers: PageInitializationBuffers, + io: &Arc, + ) -> BoxFuture<'a, Result>> { + self.inner_scheduler.init_from_buffers(buffers, io) + } + + fn try_load(&mut self, data: &Arc) -> Result<()> { + self.inner_scheduler.try_load(data) + } + + fn cache_view(&self) -> PageCacheView { + PageCacheView::BlobDescriptor + } + + fn schedule_ranges( + &self, + ranges: &[Range], + io: &Arc, + ) -> Result> { + let tasks = self.inner_scheduler.schedule_ranges(ranges, io)?; + Ok(tasks + .into_iter() + .map(|task| PageLoadTask { + decoder_fut: Self::wrap_decoder_fut(task.decoder_fut, self.def_meaning.clone()), + num_rows: task.num_rows, + }) + .collect()) + } +} + +#[derive(Debug)] +struct BlobDescriptionPageDecoder { + inner: Box, + def_meaning: Arc<[DefinitionInterpretation]>, +} + +impl BlobDescriptionPageDecoder { + fn new( + inner: Box, + def_meaning: Arc<[DefinitionInterpretation]>, + ) -> Self { + Self { inner, def_meaning } + } +} + +impl StructuralPageDecoder for BlobDescriptionPageDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + Ok(Box::new(BlobDescriptionDecodePageTask::new( + self.inner.drain(num_rows)?, + self.def_meaning.clone(), + ))) + } + + fn num_rows(&self) -> u64 { + self.inner.num_rows() + } +} + +#[derive(Debug)] +struct BlobDescriptionDecodePageTask { + inner: Box, + def_meaning: Arc<[DefinitionInterpretation]>, +} + +impl BlobDescriptionDecodePageTask { + fn new(inner: Box, def_meaning: Arc<[DefinitionInterpretation]>) -> Self { + Self { inner, def_meaning } + } +} + +impl DecodePageTask for BlobDescriptionDecodePageTask { + fn decode(self: Box) -> Result { + let decoded = self.inner.decode()?; + let num_values = decoded.data.num_values(); + + // Need to extract out the repdef information + let DataBlock::Struct(descriptions) = &decoded.data else { + return Err(Error::internal( + "Expected struct data block for descriptions", + )); + }; + let mut description_children = descriptions.children.iter(); + let DataBlock::FixedWidth(positions) = description_children.next().expect_ok()? else { + return Err(Error::internal( + "Expected fixed width data block for positions", + )); + }; + let DataBlock::FixedWidth(sizes) = description_children.next().expect_ok()? else { + return Err(Error::internal("Expected fixed width data block for sizes")); + }; + let positions = positions.data.borrow_to_typed_slice::(); + let sizes = sizes.data.borrow_to_typed_slice::(); + + let mut rep = Vec::with_capacity(num_values as usize); + let mut def = Vec::with_capacity(num_values as usize); + + for (position, size) in positions.iter().copied().zip(sizes.iter().copied()) { + if size == 0 { + if position == 0 { + rep.push(0); + def.push(0); + } else { + let repval = (position & 0xFFFF) as u16; + let defval = ((position >> 16) & 0xFFFF) as u16; + rep.push(repval); + def.push(defval); + } + } else { + rep.push(0); + def.push(0); + } + } + + let rep = if rep.iter().any(|r| *r != 0) { + Some(rep) + } else { + None + }; + let def = if self.def_meaning.len() > 1 + || self.def_meaning[0] != DefinitionInterpretation::AllValidItem + { + Some(def) + } else { + None + }; + + let repdef = + RepDefUnraveler::new(rep, def, self.def_meaning.clone(), positions.len() as u64); + + Ok(DecodedPage { + data: decoded.data, + repdef, + }) + } +} + +struct BlobCacheableState { + positions: Arc, + sizes: Arc, +} + +impl DeepSizeOf for BlobCacheableState { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + (self.positions.as_ref() as &dyn arrow_array::Array).deep_size_of_children(context) + + (self.sizes.as_ref() as &dyn arrow_array::Array).deep_size_of_children(context) + } +} + +impl CachedPageData for BlobCacheableState { + fn as_arc_any(self: Arc) -> Arc { + self + } +} + +#[derive(Debug)] +pub(super) struct BlobPageScheduler { + inner_scheduler: Box, + row_number: u64, + num_rows: u64, + def_meaning: Arc<[DefinitionInterpretation]>, + positions: Option>, + sizes: Option>, +} + +impl BlobPageScheduler { + pub fn new( + inner_scheduler: Box, + row_number: u64, + num_rows: u64, + def_meaning: Arc<[DefinitionInterpretation]>, + ) -> Self { + Self { + inner_scheduler, + row_number, + num_rows, + def_meaning, + positions: None, + sizes: None, + } + } + + fn create_page_load_task( + ranges_to_read: Vec>, + mut loaded_blobs: Vec, + first_row_number: u64, + io: &dyn EncodingsIo, + def_meaning: Arc<[DefinitionInterpretation]>, + ) -> Result { + let num_rows = loaded_blobs.len() as u64; + let read_fut = io.submit_request(ranges_to_read, first_row_number); + let decoder_fut = async move { + let bytes = read_fut.await?; + let mut bytes_iter = bytes.into_iter(); + for blob in loaded_blobs.iter_mut() { + // Empty values have def == 0 too but scheduled no read; their + // bytes were set at scheduling time. + if blob.def == 0 && blob.bytes.is_none() { + blob.set_bytes(bytes_iter.next().expect_ok()?); + } + } + debug_assert!(bytes_iter.next().is_none()); + Ok(Box::new(BlobPageDecoder::new(loaded_blobs, def_meaning)) + as Box) + } + .boxed(); + Ok(PageLoadTask { + decoder_fut, + num_rows, + }) + } +} + +impl StructuralPageScheduler for BlobPageScheduler { + fn init_layout(&self) -> Result { + self.inner_scheduler.init_layout() + } + + fn init_from_buffers<'a>( + &'a mut self, + buffers: PageInitializationBuffers, + io: &Arc, + ) -> BoxFuture<'a, Result>> { + let io = io.clone(); + let num_rows = self.num_rows; + async move { + let cached = self.inner_scheduler.init_from_buffers(buffers, &io).await?; + self.inner_scheduler.try_load(&cached)?; + let mut desc_decoders = self.inner_scheduler.schedule_ranges(&[0..num_rows], &io)?; + if desc_decoders.len() != 1 { + // This can't happen yet today so being a little lazy but if it did happen we just + // need to concatenate the descriptions. I'm guessing by then we might be doing something + // different than "load all descriptors in initialize" anyways. + return Err(Error::not_supported_source( + "Expected exactly one descriptor decoder".into(), + )); + } + let desc_decoder_task = desc_decoders.pop().unwrap(); + let mut desc_decoder = desc_decoder_task.decoder_fut.await?; + + let descs = desc_decoder.drain(desc_decoder_task.num_rows)?; + let descs = descs.decode()?; + let descs = make_array(descs.data.into_arrow(BLOB_DESC_TYPE.clone(), true)?); + let descs = descs.as_struct(); + let positions = Arc::new( + descs + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(), + ); + let sizes = Arc::new( + descs + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .clone(), + ); + let state = Arc::new(BlobCacheableState { positions, sizes }); + Ok(state as Arc) + } + .boxed() + } + + fn try_load(&mut self, data: &Arc) -> Result<()> { + let blob_state = data + .clone() + .as_arc_any() + .downcast::() + .map_err(|_| { + Error::invalid_input_source("Cached blob page data has an unexpected type".into()) + })?; + self.positions = Some(blob_state.positions.clone()); + self.sizes = Some(blob_state.sizes.clone()); + Ok(()) + } + + fn cache_view(&self) -> PageCacheView { + PageCacheView::BlobPayload + } + + fn schedule_ranges( + &self, + ranges: &[Range], + io: &Arc, + ) -> Result> { + let num_rows: u64 = ranges.iter().map(|r| r.end - r.start).sum(); + + let positions = self.positions.as_ref().expect_ok()?; + let sizes = self.sizes.as_ref().expect_ok()?; + + let mut page_load_tasks = Vec::new(); + let mut bytes_so_far = 0; + let mut ranges_to_read = Vec::with_capacity(num_rows as usize); + let mut loaded_blobs = Vec::with_capacity(num_rows as usize); + let mut first_row_number = None; + for range in ranges { + for row in range.start..range.end { + if first_row_number.is_none() { + first_row_number = Some(row + self.row_number); + } + let position = positions.value(row as usize); + let size = sizes.value(row as usize); + + if size == 0 { + let rep = (position & 0xFFFF) as u16; + let def = ((position >> 16) & 0xFFFF) as u16; + let mut blob = LoadedBlob::new(rep, def); + if def == 0 { + // A size-0 descriptor with definition level 0 is a + // valid, empty value (nulls carry their non-zero + // packed rep/def levels in `position`). No read is + // scheduled for it, so it gets its zero-length bytes + // here rather than consuming another blob's read + // result in the load task. + blob.set_bytes(Bytes::new()); + } + loaded_blobs.push(blob); + } else { + loaded_blobs.push(LoadedBlob::new(0, 0)); + ranges_to_read.push(position..(position + size)); + bytes_so_far += size; + } + + if bytes_so_far >= TARGET_SHARD_SIZE { + let page_load_task = Self::create_page_load_task( + std::mem::take(&mut ranges_to_read), + std::mem::take(&mut loaded_blobs), + first_row_number.unwrap(), + io.as_ref(), + self.def_meaning.clone(), + )?; + page_load_tasks.push(page_load_task); + bytes_so_far = 0; + first_row_number = None; + } + } + } + if !loaded_blobs.is_empty() { + let page_load_task = Self::create_page_load_task( + std::mem::take(&mut ranges_to_read), + std::mem::take(&mut loaded_blobs), + first_row_number.unwrap(), + io.as_ref(), + self.def_meaning.clone(), + )?; + page_load_tasks.push(page_load_task); + } + + Ok(page_load_tasks) + } +} + +#[derive(Debug)] +struct LoadedBlob { + bytes: Option, + rep: u16, + def: u16, +} + +impl LoadedBlob { + fn new(rep: u16, def: u16) -> Self { + Self { + bytes: None, + rep, + def, + } + } + + fn set_bytes(&mut self, bytes: Bytes) { + self.bytes = Some(bytes); + } +} + +#[derive(Debug)] +struct BlobPageDecoder { + blobs: VecDeque, + def_meaning: Arc<[DefinitionInterpretation]>, + num_rows: u64, +} + +impl BlobPageDecoder { + fn new(blobs: Vec, def_meaning: Arc<[DefinitionInterpretation]>) -> Self { + Self { + num_rows: blobs.len() as u64, + blobs: blobs.into_iter().collect(), + def_meaning, + } + } +} + +impl StructuralPageDecoder for BlobPageDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + let blobs = self.blobs.drain(0..num_rows as usize).collect::>(); + Ok(Box::new(BlobDecodePageTask::new( + blobs, + self.def_meaning.clone(), + ))) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +#[derive(Debug)] +struct BlobDecodePageTask { + blobs: Vec, + def_meaning: Arc<[DefinitionInterpretation]>, +} + +impl BlobDecodePageTask { + fn new(blobs: Vec, def_meaning: Arc<[DefinitionInterpretation]>) -> Self { + Self { blobs, def_meaning } + } +} + +impl DecodePageTask for BlobDecodePageTask { + fn decode(self: Box) -> Result { + let num_values = self.blobs.len() as u64; + let num_bytes = self + .blobs + .iter() + .filter_map(|b| b.bytes.as_ref()) + .map(|b| b.len()) + .sum::(); + let mut buffer = Vec::with_capacity(num_bytes); + let mut offsets = Vec::with_capacity(num_values as usize + 1); + let mut rep = Vec::with_capacity(num_values as usize); + let mut def = Vec::with_capacity(num_values as usize); + offsets.push(0_u64); + for blob in self.blobs { + rep.push(blob.rep); + def.push(blob.def); + if let Some(bytes) = blob.bytes { + offsets.push(offsets.last().unwrap() + bytes.len() as u64); + buffer.extend_from_slice(&bytes); + } else { + // Null / emptyvalue + offsets.push(*offsets.last().unwrap()); + } + } + let offsets = LanceBuffer::reinterpret_vec(offsets); + let data = LanceBuffer::from(buffer); + let data_block = DataBlock::VariableWidth(VariableWidthBlock { + data, + offsets, + bits_per_offset: 64, + num_values, + block_info: BlockInfo::new(), + }); + + let rep = if rep.iter().any(|r| *r != 0) { + Some(rep) + } else { + None + }; + let def = if self.def_meaning.len() > 1 + || self.def_meaning[0] != DefinitionInterpretation::AllValidItem + { + Some(def) + } else { + None + }; + + Ok(DecodedPage { + data: data_block, + repdef: RepDefUnraveler::new(rep, def, self.def_meaning, num_values), + }) + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive/chunk_index.rs b/vendor/lance-encoding/src/encodings/logical/primitive/chunk_index.rs new file mode 100644 index 000000000..19fb0fdb0 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive/chunk_index.rs @@ -0,0 +1,536 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Compact per-page chunk index for the mini-block structural encoding. +//! +//! The chunk index is stored on disk in an extremely compressed form that +//! requires a lot of CPU to work with. However, extracting it out to its full +//! width can be RAM-intensive. As a compromise we extract into a prefix-sum +//! array that we fit into `u32` if possible and we avoid storing per-block row +//! counts when those are redundant. The scheduler looks chunks up by index +//! (byte range, leaf value count) and by row (which chunk holds a row). +//! +//! ```text +//! MiniBlockChunkIndex +//! |- base: u64 absolute file position of the value buffer +//! |- byte_starts: PrefixSums cumulative chunk byte sizes (all pages) +//! `- rows: RowMapping +//! |- UniformFlat { .. } flat page, uniform leaf chunking (arithmetic) +//! |- Flat { value_starts } flat page, non-uniform leaf chunking +//! `- Nested { row_starts, .. } repetition present; rows tracked as prefix sums +//! ``` + +use std::ops::Range; + +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; +use lance_core::cache::{Context, DeepSizeOf}; + +/// Cumulative (prefix-sum) array of length `num_chunks + 1` (entry `0` is `0`, +/// the last entry is the grand total). Stored as `u32` when the total fits, +/// else `u64`. +#[derive(Debug, DeepSizeOf)] +pub enum PrefixSums { + U32(Vec), + U64(Vec), +} + +impl PrefixSums { + /// Builds the cumulative array from per-chunk `deltas`. `total` selects the + /// storage width (callers must pass the true sum); `num_chunks` only pre-sizes. + pub fn from_deltas(deltas: impl Iterator, num_chunks: usize, total: u64) -> Self { + if total <= u32::MAX as u64 { + let mut values = Vec::with_capacity(num_chunks + 1); + let mut acc = 0u32; + values.push(0); + for delta in deltas { + acc += delta as u32; + values.push(acc); + } + debug_assert_eq!(values.len(), num_chunks + 1); + debug_assert_eq!(acc as u64, total); + Self::U32(values) + } else { + let mut values = Vec::with_capacity(num_chunks + 1); + let mut acc = 0u64; + values.push(0); + for delta in deltas { + acc += delta; + values.push(acc); + } + debug_assert_eq!(values.len(), num_chunks + 1); + debug_assert_eq!(acc, total); + Self::U64(values) + } + } + + /// Builds a `PrefixSums` from an already-cumulative array (`[0, .., total]`), + /// narrowing to `u32` when the total fits. Avoids the deltas buffer + /// [`Self::from_deltas`] would need. + fn from_prefix(prefix: Vec) -> Self { + debug_assert!(!prefix.is_empty()); + debug_assert_eq!(prefix[0], 0); + let total = prefix.last().copied().unwrap_or(0); + if total <= u32::MAX as u64 { + Self::U32(prefix.into_iter().map(|v| v as u32).collect()) + } else { + Self::U64(prefix) + } + } + + /// Cumulative value at position `i` (i.e. the start of chunk `i`). + pub fn get(&self, i: usize) -> u64 { + match self { + Self::U32(values) => values[i] as u64, + Self::U64(values) => values[i], + } + } + + /// Start and end of chunk `i` (positions `i`, `i + 1`) behind one width + /// match -- halves the branching of two `get` calls on the hot per-chunk path. + pub fn get_pair(&self, i: usize) -> (u64, u64) { + match self { + Self::U32(values) => (values[i] as u64, values[i + 1] as u64), + Self::U64(values) => (values[i], values[i + 1]), + } + } + + /// Number of chunks (array length minus the trailing total). + pub fn num_chunks(&self) -> usize { + match self { + Self::U32(values) => values.len() - 1, + Self::U64(values) => values.len() - 1, + } + } + + /// Size of chunk `i` (the delta between consecutive cumulative values). + pub fn delta(&self, i: usize) -> u64 { + let (start, end) = self.get_pair(i); + end - start + } + + /// Index of the chunk whose half-open span `[get(i), get(i+1))` contains + /// `value`. On an exact hit against a chunk start, returns the *first* chunk + /// with that start (chunks can share a start row). + pub fn find(&self, value: u64) -> usize { + // Match the width once, then binary-search only the starts (not the + // trailing total). `partition_point` already yields the first of any + // duplicated starts; the `idx - 1` fallback is safe since `get(0) == 0`. + match self { + Self::U32(values) => { + let starts = &values[..values.len() - 1]; + let idx = starts.partition_point(|&start| (start as u64) < value); + if idx < starts.len() && starts[idx] as u64 == value { + idx + } else { + idx - 1 + } + } + Self::U64(values) => { + let starts = &values[..values.len() - 1]; + let idx = starts.partition_point(|&start| start < value); + if idx < starts.len() && starts[idx] == value { + idx + } else { + idx - 1 + } + } + } + } +} + +/// Leaf value counts per chunk, needed to decode. Tracked only for nested +/// pages; flat pages read items off the row mapping (rows == items). +#[derive(Debug, DeepSizeOf)] +pub enum ItemCounts { + /// Every non-last chunk holds the same number of values. + Uniform { + values_per_chunk: u64, + last_chunk_values: u64, + }, + /// `log2` of each chunk's value count, stored as one byte per chunk rather + /// than the full count because this index stays cached in RAM; the last + /// chunk is handled via `last_chunk_values`. + PerChunkLog { + logs: Vec, + last_chunk_values: u64, + }, +} + +impl ItemCounts { + fn get(&self, i: usize, num_chunks: usize) -> u64 { + match self { + Self::Uniform { + values_per_chunk, + last_chunk_values, + } => { + if i == num_chunks - 1 { + *last_chunk_values + } else { + *values_per_chunk + } + } + Self::PerChunkLog { + logs, + last_chunk_values, + } => { + if i == num_chunks - 1 { + *last_chunk_values + } else { + 1u64 << logs[i] + } + } + } + } +} + +/// How row ranges map onto chunks. +/// +/// Flat pages have row == value index and no preamble/trailer. Nested pages +/// track rows separately from leaf items and store a trailer bit per chunk; a +/// chunk's preamble is the previous chunk's trailer. +#[derive(Debug)] +pub enum RowMapping { + /// Flat page whose non-last chunks all hold `values_per_chunk` values, so + /// row->chunk is pure arithmetic. + UniformFlat { + values_per_chunk: u64, + last_chunk_values: u64, + num_chunks: usize, + }, + /// Flat page with non-uniform chunk sizes; `value_starts` are the cumulative + /// value counts (final entry == number of items in the page). + Flat { value_starts: PrefixSums }, + /// Nested page. `row_starts` are cumulative row counts (final entry == + /// number of rows in the page); `has_trailer[i]` is set when chunk `i` ends + /// with a partial list. + Nested { + row_starts: PrefixSums, + has_trailer: BooleanBuffer, + item_counts: ItemCounts, + }, +} + +impl DeepSizeOf for RowMapping { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + match self { + Self::UniformFlat { .. } => 0, + Self::Flat { value_starts } => value_starts.deep_size_of_children(context), + Self::Nested { + row_starts, + has_trailer, + item_counts, + } => { + row_starts.deep_size_of_children(context) + + has_trailer.len().div_ceil(8) + + item_counts.deep_size_of_children(context) + } + } + } +} + +/// Compact per-page chunk index that avoids fully materializing the repetition +/// index into u64's to save RAM. See the module docs for the layout. +#[derive(Debug)] +pub struct MiniBlockChunkIndex { + base: u64, + byte_starts: PrefixSums, + rows: RowMapping, +} + +impl MiniBlockChunkIndex { + pub fn new(base: u64, byte_starts: PrefixSums, rows: RowMapping) -> Self { + Self { + base, + byte_starts, + rows, + } + } + + /// Number of chunks in the page. + pub fn num_chunks(&self) -> usize { + self.byte_starts.num_chunks() + } + + /// Absolute byte range of chunk `i` within the file. + pub fn byte_range(&self, i: usize) -> Range { + let (start, end) = self.byte_starts.get_pair(i); + (self.base + start)..(self.base + end) + } + + /// Number of leaf values in chunk `i` (passed to the value decompressor). + pub fn items_in_chunk(&self, i: usize) -> u64 { + let num_chunks = self.num_chunks(); + match &self.rows { + RowMapping::UniformFlat { + values_per_chunk, + last_chunk_values, + .. + } => { + if i == num_chunks - 1 { + *last_chunk_values + } else { + *values_per_chunk + } + } + RowMapping::Flat { value_starts } => value_starts.delta(i), + RowMapping::Nested { item_counts, .. } => item_counts.get(i, num_chunks), + } + } + + /// Index of the chunk that contains `row`. + pub fn find_chunk(&self, row: u64) -> usize { + match &self.rows { + RowMapping::UniformFlat { + values_per_chunk, + num_chunks, + .. + } => ((row / values_per_chunk) as usize).min(num_chunks - 1), + RowMapping::Flat { value_starts } => value_starts.find(row), + RowMapping::Nested { row_starts, .. } => row_starts.find(row), + } + } + + /// First row (relative to the page) that begins in chunk `i`. + pub fn first_row(&self, i: usize) -> u64 { + match &self.rows { + RowMapping::UniformFlat { + values_per_chunk, .. + } => i as u64 * values_per_chunk, + RowMapping::Flat { value_starts } => value_starts.get(i), + RowMapping::Nested { row_starts, .. } => row_starts.get(i), + } + } + + /// Number of rows that start in chunk `i`, including a trailer but not a + /// preamble (the previous `starts_including_trailer`). + pub fn rows_in_chunk(&self, i: usize) -> u64 { + let num_chunks = self.num_chunks(); + match &self.rows { + RowMapping::UniformFlat { + values_per_chunk, + last_chunk_values, + .. + } => { + if i == num_chunks - 1 { + *last_chunk_values + } else { + *values_per_chunk + } + } + RowMapping::Flat { value_starts } => value_starts.delta(i), + RowMapping::Nested { row_starts, .. } => row_starts.delta(i), + } + } + + /// Whether chunk `i` begins with a preamble (a continuation of the previous + /// chunk's list). Always false for flat pages; for nested pages this is the + /// previous chunk's trailer. + pub fn has_preamble(&self, i: usize) -> bool { + match &self.rows { + RowMapping::Nested { has_trailer, .. } => i > 0 && has_trailer.value(i - 1), + _ => false, + } + } + + /// Whether chunk `i` ends with a trailer (a partial list continued in the + /// next chunk). Always false for flat pages. + pub fn has_trailer(&self, i: usize) -> bool { + match &self.rows { + RowMapping::Nested { has_trailer, .. } => has_trailer.value(i), + _ => false, + } + } + + /// Name of the active row-mapping variant, used to assert detection in tests. + #[cfg(test)] + pub fn row_mapping_debug(&self) -> &'static str { + match &self.rows { + RowMapping::UniformFlat { .. } => "uniform_flat", + RowMapping::Flat { .. } => "flat", + RowMapping::Nested { .. } => "nested", + } + } + + /// Builds a nested index from raw repetition-index bytes, using placeholder + /// byte offsets and item counts. Only the row axis is populated, which is + /// all the scheduler exercises. + #[cfg(test)] + pub fn new_nested_for_test(rep_bytes: &[u8], stride: usize) -> Self { + let (row_starts, has_trailer) = parse_nested_rep(rep_bytes, stride); + let num_chunks = row_starts.num_chunks(); + let byte_starts = PrefixSums::from_deltas( + std::iter::repeat_n(8u64, num_chunks), + num_chunks, + 8 * num_chunks as u64, + ); + Self { + base: 0, + byte_starts, + rows: RowMapping::Nested { + row_starts, + has_trailer, + item_counts: ItemCounts::Uniform { + values_per_chunk: 1, + last_chunk_values: 1, + }, + }, + } + } +} + +impl DeepSizeOf for MiniBlockChunkIndex { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + self.byte_starts.deep_size_of_children(context) + self.rows.deep_size_of_children(context) + } +} + +/// Parses a mini-block repetition index into the compact nested row mapping. +/// +/// Bytes are `u64`s in groups of `stride`; the first two are `ends` (lists +/// finishing in the chunk) and `partial` (leftover items). Only cumulative row +/// starts and a trailer bit are kept: `has_preamble[i] = has_trailer[i-1]` and +/// `starts_including_trailer = ends + has_trailer - has_preamble`. +pub fn parse_nested_rep(rep_bytes: &[u8], stride: usize) -> (PrefixSums, BooleanBuffer) { + // Read the two `u64`s per group straight from the little-endian bytes rather + // than copying the buffer to reinterpret it. The caller guarantees + // `rep_bytes.len() % 8 == 0`, so the 8-byte windows stay in bounds. + const WORD: usize = std::mem::size_of::(); + let read_word = |word_idx: usize| -> u64 { + let byte = word_idx * WORD; + u64::from_le_bytes(rep_bytes[byte..byte + WORD].try_into().unwrap()) + }; + let num_chunks = (rep_bytes.len() / WORD) / stride; + + let mut has_trailer_builder = BooleanBufferBuilder::new(num_chunks); + // Accumulate the cumulative row starts in a single pass (entry 0 is 0, the + // trailing entry is the total) so there is no separate deltas buffer. + let mut row_starts = Vec::with_capacity(num_chunks + 1); + row_starts.push(0u64); + let mut acc = 0u64; + let mut chunk_has_preamble = false; + + for i in 0..num_chunks { + let base_idx = i * stride; + let ends = read_word(base_idx); + let partial = read_word(base_idx + 1); + + let has_trailer = partial > 0; + let starts_including_trailer = ends + (has_trailer as u64) - (chunk_has_preamble as u64); + + has_trailer_builder.append(has_trailer); + acc += starts_including_trailer; + row_starts.push(acc); + + chunk_has_preamble = has_trailer; + } + + ( + PrefixSums::from_prefix(row_starts), + has_trailer_builder.finish(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::buffer::LanceBuffer; + + /// Reference decode: the previous per-block repetition index, used as an + /// oracle for the compact `parse_nested_rep`. + struct RefBlock { + first_row: u64, + starts_including_trailer: u64, + has_preamble: bool, + has_trailer: bool, + } + + fn reference_decode(rep_bytes: &[u8], stride: usize) -> Vec { + let buffer = LanceBuffer::from(rep_bytes.to_vec()); + let u64_slice = buffer.borrow_to_typed_slice::(); + let n = u64_slice.len() / stride; + let mut blocks = Vec::with_capacity(n); + let mut chunk_has_preamble = false; + let mut offset = 0u64; + for i in 0..n { + let base_idx = i * stride; + let ends = u64_slice[base_idx]; + let partial = u64_slice[base_idx + 1]; + let has_trailer = partial > 0; + let starts_including_trailer = + ends + (has_trailer as u64) - (chunk_has_preamble as u64); + blocks.push(RefBlock { + first_row: offset, + starts_including_trailer, + has_preamble: chunk_has_preamble, + has_trailer, + }); + chunk_has_preamble = has_trailer; + offset += starts_including_trailer; + } + blocks + } + + fn rep_bytes(values: &[u64]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() + } + + #[test] + fn test_prefix_sums_u32() { + let sums = PrefixSums::from_deltas([2u64, 3, 5].into_iter(), 3, 10); + assert!(matches!(sums, PrefixSums::U32(_))); + assert_eq!(sums.num_chunks(), 3); + assert_eq!(sums.get(0), 0); + assert_eq!(sums.get(1), 2); + assert_eq!(sums.get(3), 10); + assert_eq!(sums.delta(1), 3); + } + + #[test] + fn test_prefix_sums_u64_selected_by_total() { + let big = u32::MAX as u64 + 1; + let sums = PrefixSums::from_deltas([big].into_iter(), 1, big); + assert!(matches!(sums, PrefixSums::U64(_))); + assert_eq!(sums.get(1), big); + } + + #[test] + fn test_prefix_sums_find() { + // Chunk starts: 0, 5, 5, 12 (a zero-width chunk creates a duplicate start) + let sums = PrefixSums::from_deltas([5u64, 0, 7].into_iter(), 3, 12); + // Inside the first chunk + assert_eq!(sums.find(3), 0); + // Exact match on a duplicated start returns the first such chunk + assert_eq!(sums.find(5), 1); + // Inside the last chunk + assert_eq!(sums.find(11), 2); + // Start of the first chunk + assert_eq!(sums.find(0), 0); + } + + #[test] + fn test_parse_nested_rep_matches_reference() { + let cases: Vec> = vec![ + vec![5, 2, 3, 0, 4, 7, 2, 0], + vec![5, 2, 3, 3, 20, 0], + vec![0, 5, 0, 3, 10, 0], + vec![1, 0], + ]; + for values in cases { + let bytes = rep_bytes(&values); + let reference = reference_decode(&bytes, 2); + let (row_starts, has_trailer) = parse_nested_rep(&bytes, 2); + assert_eq!(row_starts.num_chunks(), reference.len()); + for (i, block) in reference.iter().enumerate() { + assert_eq!(row_starts.get(i), block.first_row, "first_row[{i}]"); + assert_eq!( + row_starts.delta(i), + block.starts_including_trailer, + "starts_including_trailer[{i}]" + ); + assert_eq!(has_trailer.value(i), block.has_trailer, "has_trailer[{i}]"); + let derived_preamble = i > 0 && has_trailer.value(i - 1); + assert_eq!(derived_preamble, block.has_preamble, "has_preamble[{i}]"); + } + } + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive/constant.rs b/vendor/lance-encoding/src/encodings/logical/primitive/constant.rs new file mode 100644 index 000000000..3ac011ceb --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive/constant.rs @@ -0,0 +1,522 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{any::Any, collections::VecDeque, ops::Range, sync::Arc}; + +use arrow_array::{Array, ArrayRef, new_empty_array}; +use arrow_buffer::ScalarBuffer; +use arrow_schema::DataType; +use bytes::Bytes; +use futures::FutureExt; +use futures::future::BoxFuture; + +use lance_core::{ + Error, Result, + cache::{Context, DeepSizeOf}, +}; + +use crate::{ + EncodingsIo, + buffer::LanceBuffer, + decoder::PageEncoding, + encoder::EncodedPage, + encodings::logical::primitive::{ + CachedPageData, PageInitialization, PageInitializationBuffers, PageLoadTask, + }, + format::ProtobufUtils21, + repdef::{DefinitionInterpretation, RepDefUnraveler, max_visible_level}, +}; + +pub(crate) fn encode_constant_page( + column_idx: u32, + scalar: ArrayRef, + repdef: crate::repdef::SerializedRepDefs, + row_number: u64, + num_rows: u64, +) -> Result { + let inline_value = lance_arrow::scalar::try_inline_value(&scalar); + let value_buffer = if inline_value.is_some() { + None + } else { + Some(LanceBuffer::from( + lance_arrow::scalar::encode_scalar_value_buffer(&scalar)?, + )) + }; + + let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, inline_value); + + let has_repdef = repdef.repetition_levels.is_some() || repdef.definition_levels.is_some(); + + let data = if !has_repdef { + value_buffer.into_iter().collect::>() + } else { + let rep_bytes = repdef + .repetition_levels + .as_ref() + .map(|rep| LanceBuffer::reinterpret_slice(rep.clone())) + .unwrap_or_else(LanceBuffer::empty); + let def_bytes = repdef + .definition_levels + .as_ref() + .map(|def| LanceBuffer::reinterpret_slice(def.clone())) + .unwrap_or_else(LanceBuffer::empty); + + match value_buffer { + Some(value_buffer) => vec![value_buffer, rep_bytes, def_bytes], + None => vec![rep_bytes, def_bytes], + } + }; + + Ok(EncodedPage { + column_idx, + data, + description: PageEncoding::Structural(description), + num_rows, + row_number, + }) +} + +#[derive(Debug)] +struct CachedConstantState { + scalar: ArrayRef, + rep: Option>, + def: Option>, +} + +impl DeepSizeOf for CachedConstantState { + fn deep_size_of_children(&self, _ctx: &mut Context) -> usize { + self.scalar.get_buffer_memory_size() + + self.rep.as_ref().map(|buf| buf.len() * 2).unwrap_or(0) + + self.def.as_ref().map(|buf| buf.len() * 2).unwrap_or(0) + } +} + +impl CachedPageData for CachedConstantState { + fn as_arc_any(self: Arc) -> Arc { + self + } +} + +#[derive(Debug, Clone)] +enum ScalarSource { + Inline(Vec), + ValueBuffer(usize), +} + +/// The (scalar, rep, def) file ranges a constant page reads, in the order +/// `init_layout` declares and `init_from_buffers` consumes them. Any field is +/// `None` when that buffer is absent. Computed once at construction so the two +/// halves share one source of truth and their buffer order cannot drift. +#[derive(Debug)] +struct ConstantBufferLayout { + scalar: Option>, + rep: Option>, + def: Option>, +} + +#[derive(Debug)] +pub struct ConstantPageScheduler { + scalar_source: ScalarSource, + data_type: DataType, + def_meaning: Arc<[DefinitionInterpretation]>, + max_rep: u16, + max_visible_def: u16, + layout: ConstantBufferLayout, + repdef: Option>, +} + +impl ConstantPageScheduler { + pub fn try_new( + buffer_offsets_and_sizes: Arc<[(u64, u64)]>, + inline_value: Option, + data_type: DataType, + def_meaning: Arc<[DefinitionInterpretation]>, + ) -> Result { + let max_rep = def_meaning.iter().filter(|d| d.is_list()).count() as u16; + let max_visible_def = max_visible_level(&def_meaning); + + let (scalar_source, rep_buf_idx, def_buf_idx) = + match (inline_value, buffer_offsets_and_sizes.len()) { + (Some(inline), 0) => (ScalarSource::Inline(inline.to_vec()), None, None), + (Some(inline), 2) => (ScalarSource::Inline(inline.to_vec()), Some(0), Some(1)), + (None, 1) => (ScalarSource::ValueBuffer(0), None, None), + (None, 3) => (ScalarSource::ValueBuffer(0), Some(1), Some(2)), + (Some(_inline), 1) => { + return Err(Error::invalid_input(format!( + "Invalid constant layout: inline_value present with {} buffers", + 1 + ))); + } + (Some(_inline), 3) => { + return Err(Error::invalid_input( + "Invalid constant layout: inline_value present with 3 buffers", + )); + } + (None, 0) => { + return Err(Error::invalid_input( + "Invalid constant layout: missing scalar source", + )); + } + (None, 2) => { + return Err(Error::invalid_input( + "Invalid constant layout: ambiguous (2 buffers and no inline_value)", + )); + } + (Some(_), n) => { + return Err(Error::invalid_input(format!( + "Invalid constant layout: inline_value present with {} buffers", + n + ))); + } + (None, n) => { + return Err(Error::invalid_input(format!( + "Invalid constant layout: unexpected buffer count {}", + n + ))); + } + }; + + let range_of = |idx: usize| { + let (pos, len) = buffer_offsets_and_sizes[idx]; + pos..pos + len + }; + let layout = ConstantBufferLayout { + scalar: match &scalar_source { + ScalarSource::ValueBuffer(idx) => Some(range_of(*idx)), + ScalarSource::Inline(_) => None, + }, + rep: rep_buf_idx + .filter(|&idx| buffer_offsets_and_sizes[idx].1 > 0) + .map(range_of), + def: def_buf_idx + .filter(|&idx| buffer_offsets_and_sizes[idx].1 > 0) + .map(range_of), + }; + + Ok(Self { + scalar_source, + data_type, + def_meaning, + max_rep, + max_visible_def, + layout, + repdef: None, + }) + } +} + +impl crate::encodings::logical::primitive::StructuralPageScheduler for ConstantPageScheduler { + fn init_layout(&self) -> Result { + // Order must match `init_from_buffers`' consumption: scalar, rep, def. + Ok(PageInitialization::new([ + self.layout.scalar.clone(), + self.layout.rep.clone(), + self.layout.def.clone(), + ])) + } + + fn init_from_buffers<'a>( + &'a mut self, + buffers: PageInitializationBuffers, + _io: &Arc, + ) -> BoxFuture<'a, Result>> { + // Consume `buffers` in the same scalar, rep, def order `init_layout` + // declared them, using the shared layout's presence flags. + let scalar_source = self.scalar_source.clone(); + let data_type = self.data_type.clone(); + async move { + let [scalar_buffer, rep_buffer, def_buffer] = buffers.try_into_array("constant")?; + + let scalar = match (scalar_source, scalar_buffer) { + (ScalarSource::Inline(inline), None) => { + lance_arrow::scalar::decode_scalar_from_inline_value(&data_type, &inline)? + } + (ScalarSource::ValueBuffer(_), Some(bytes)) => { + let buf = LanceBuffer::from_bytes(bytes, 1); + lance_arrow::scalar::decode_scalar_from_value_buffer(&data_type, buf.as_ref())? + } + _ => { + return Err(Error::internal( + "Constant page scalar buffer does not match its initialization layout", + )); + } + }; + + let rep = rep_buffer.map(|rep| { + let rep = LanceBuffer::from_bytes(rep, 2); + rep.borrow_to_typed_slice::() + }); + + let def = def_buffer.map(|def| { + let def = LanceBuffer::from_bytes(def, 2); + def.borrow_to_typed_slice::() + }); + + let cached = Arc::new(CachedConstantState { scalar, rep, def }); + Ok(cached as Arc) + } + .boxed() + } + + fn try_load(&mut self, data: &Arc) -> Result<()> { + let cached: Arc = + data.clone().as_arc_any().downcast().map_err(|_| { + Error::invalid_input_source( + "Cached constant page data has an unexpected type".into(), + ) + })?; + if cached.scalar.data_type() != &self.data_type { + return Err(Error::invalid_input_source( + format!( + "Cached constant page scalar has data type {:?}, expected {:?}", + cached.scalar.data_type(), + self.data_type + ) + .into(), + )); + } + self.repdef = Some(cached); + Ok(()) + } + + fn schedule_ranges( + &self, + ranges: &[Range], + _io: &Arc, + ) -> Result> { + let num_rows = ranges.iter().map(|r| r.end - r.start).sum::(); + let decoder = Box::new(ConstantPageDecoder { + ranges: VecDeque::from_iter(ranges.iter().cloned()), + scalar: self.repdef.as_ref().unwrap().scalar.clone(), + rep: self.repdef.as_ref().unwrap().rep.clone(), + def: self.repdef.as_ref().unwrap().def.clone(), + def_meaning: self.def_meaning.clone(), + max_rep: self.max_rep, + max_visible_def: self.max_visible_def, + cursor_row: 0, + cursor_level: 0, + num_rows, + }) + as Box; + Ok(vec![PageLoadTask { + decoder_fut: std::future::ready(Ok(decoder)).boxed(), + num_rows, + }]) + } +} + +#[derive(Debug)] +struct ConstantPageDecoder { + ranges: VecDeque>, + scalar: ArrayRef, + rep: Option>, + def: Option>, + def_meaning: Arc<[DefinitionInterpretation]>, + max_rep: u16, + max_visible_def: u16, + cursor_row: u64, + cursor_level: usize, + num_rows: u64, +} + +impl ConstantPageDecoder { + fn drain_ranges(&mut self, num_rows: u64) -> Vec> { + let mut rows_desired = num_rows; + let mut ranges = Vec::with_capacity(self.ranges.len()); + while rows_desired > 0 { + let front = self.ranges.front_mut().unwrap(); + let avail = front.end - front.start; + if avail > rows_desired { + ranges.push(front.start..front.start + rows_desired); + front.start += rows_desired; + rows_desired = 0; + } else { + ranges.push(self.ranges.pop_front().unwrap()); + rows_desired -= avail; + } + } + ranges + } + + fn take_row(&mut self) -> Result<(Range, u64)> { + let start = self.cursor_level; + let end = if let Some(rep) = &self.rep { + if start >= rep.len() { + return Err(Error::internal( + "Invalid constant layout: repetition buffer too short", + )); + } + if rep[start] != self.max_rep { + return Err(Error::internal( + "Invalid constant layout: row did not start at max_rep", + )); + } + let mut end = start + 1; + while end < rep.len() && rep[end] != self.max_rep { + end += 1; + } + end + } else { + start + 1 + }; + + let visible = if let Some(def) = &self.def { + def[start..end] + .iter() + .filter(|d| **d <= self.max_visible_def) + .count() as u64 + } else { + (end - start) as u64 + }; + + self.cursor_level = end; + self.cursor_row += 1; + Ok((start..end, visible)) + } + + fn skip_to_row(&mut self, target_row: u64) -> Result<()> { + while self.cursor_row < target_row { + self.take_row()?; + } + Ok(()) + } +} + +impl crate::encodings::logical::primitive::StructuralPageDecoder for ConstantPageDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + let drained_ranges = self.drain_ranges(num_rows); + + let mut level_slices: Vec> = Vec::new(); + let mut visible_items_total: u64 = 0; + + for range in drained_ranges { + self.skip_to_row(range.start)?; + for _ in range.start..range.end { + let (level_range, visible) = self.take_row()?; + visible_items_total += visible; + if let Some(last) = level_slices.last_mut() + && last.end == level_range.start + { + last.end = level_range.end; + continue; + } + level_slices.push(level_range); + } + } + + Ok(Box::new(DecodeConstantTask { + scalar: self.scalar.clone(), + rep: self.rep.clone(), + def: self.def.clone(), + level_slices, + visible_items_total, + def_meaning: self.def_meaning.clone(), + max_visible_def: self.max_visible_def, + })) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +#[derive(Debug)] +struct DecodeConstantTask { + scalar: ArrayRef, + rep: Option>, + def: Option>, + level_slices: Vec>, + visible_items_total: u64, + def_meaning: Arc<[DefinitionInterpretation]>, + max_visible_def: u16, +} + +impl DecodeConstantTask { + fn slice_levels( + levels: &Option>, + slices: &[Range], + ) -> Option> { + levels.as_ref().map(|levels| { + let total = slices.iter().map(|r| r.end - r.start).sum(); + let mut out = Vec::with_capacity(total); + for r in slices { + out.extend(levels[r.start..r.end].iter().copied()); + } + out + }) + } + + fn materialize_values(&self, num_values: u64) -> Result { + if num_values == 0 { + return Ok(new_empty_array(self.scalar.data_type())); + } + + if let DataType::Struct(fields) = self.scalar.data_type() + && fields.is_empty() + { + return Ok(Arc::new(arrow_array::StructArray::new_empty_fields( + num_values as usize, + None, + )) as ArrayRef); + } + + let indices = arrow_array::UInt64Array::from(vec![0u64; num_values as usize]); + Ok(arrow_select::take::take( + self.scalar.as_ref(), + &indices, + None, + )?) + } +} + +impl crate::decoder::DecodePageTask for DecodeConstantTask { + fn decode(self: Box) -> Result { + let rep = Self::slice_levels(&self.rep, &self.level_slices); + let def = Self::slice_levels(&self.def, &self.level_slices); + + let visible_items_total = if let Some(def) = &def { + def.iter().filter(|d| **d <= self.max_visible_def).count() as u64 + } else { + self.visible_items_total + }; + + let values = self.materialize_values(visible_items_total)?; + let data = crate::data::DataBlock::from_array(values); + let unraveler = + RepDefUnraveler::new(rep, def, self.def_meaning.clone(), visible_items_total); + + Ok(crate::decoder::DecodedPage { + data, + repdef: unraveler, + }) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::BooleanArray; + + use super::*; + use crate::encodings::logical::primitive::StructuralPageScheduler; + + #[test] + fn rejects_cached_scalar_with_different_data_type() { + let mut scheduler = ConstantPageScheduler::try_new( + Arc::from([]), + Some(Bytes::new()), + DataType::UInt8, + Arc::from([DefinitionInterpretation::AllValidItem]), + ) + .unwrap(); + let cached: Arc = Arc::new(CachedConstantState { + scalar: Arc::new(BooleanArray::from(vec![true])), + rep: None, + def: None, + }); + + let error = scheduler.try_load(&cached).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. })); + assert!(error.to_string().contains("Boolean")); + assert!(error.to_string().contains("UInt8")); + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive/dict.rs b/vendor/lance-encoding/src/encodings/logical/primitive/dict.rs new file mode 100644 index 000000000..19582c167 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive/dict.rs @@ -0,0 +1,632 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::HashMap, sync::Arc}; + +/// Bits per value for FixedWidth dictionary values (legacy default for 128-bit values) +pub const DICT_FIXED_WIDTH_BITS_PER_VALUE: u64 = 128; +/// Bits per index for dictionary indices (always i32) +pub const DICT_INDICES_BITS_PER_VALUE: u64 = 32; + +use arrow_array::{ + Array, DictionaryArray, PrimitiveArray, UInt64Array, + cast::AsArray, + types::{ + ArrowDictionaryKeyType, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, + UInt32Type, UInt64Type, + }, +}; +use arrow_buffer::ArrowNativeType; +use arrow_schema::DataType; +use arrow_select::take::TakeOptions; +use lance_core::{Error, Result, error::LanceOptionExt, utils::hash::U8SliceKey}; + +use crate::{ + buffer::LanceBuffer, + data::{BlockInfo, DataBlock, FixedWidthDataBlock, VariableWidthBlock}, + statistics::{ComputeStat, GetStat, Stat}, +}; + +// Helper function for normalize_dict_nulls +fn normalize_dict_nulls_impl( + array: Arc, +) -> Result> { + // TODO: Fast path when there is only one null index? (common case) + + let dict_array = array.as_dictionary_opt::().expect_ok()?; + + if dict_array.values().null_count() == 0 { + return Ok(array); + } + + let mut mapping = vec![None; dict_array.values().len()]; + let mut skipped = 0; + let mut valid_indices = Vec::with_capacity(dict_array.values().len()); + for (old_idx, is_valid) in dict_array.values().nulls().expect_ok()?.iter().enumerate() { + if is_valid { + // Should be safe since we are only decreasing K values (e.g. won't overflow u8 keys into u16) + mapping[old_idx] = Some(K::Native::from_usize(old_idx - skipped).expect_ok()?); + valid_indices.push(old_idx as u64); + } else { + skipped += 1; + mapping[old_idx] = None; + } + } + + let mut keys_builder = PrimitiveArray::::builder(dict_array.keys().len()); + for key in dict_array.keys().iter() { + if let Some(key) = key { + if let Some(mapped) = mapping[key.to_usize().expect_ok()?] { + // Valid item + keys_builder.append_value(mapped); + } else { + // Null via values + keys_builder.append_null(); + } + } else { + // Null via keys + keys_builder.append_null(); + } + } + let keys = keys_builder.finish(); + + let valid_indices = UInt64Array::from(valid_indices); + let values = arrow_select::take::take( + dict_array.values(), + &valid_indices, + Some(TakeOptions { + check_bounds: false, + }), + )?; + + Ok(Arc::new(DictionaryArray::new(keys, values)) as Arc) +} + +/// In Arrow a dictionary array can have nulls in two different places: +/// 1. The keys can be null +/// 2. The values can be null +/// +/// We want to normalize this so that all nulls are in the keys. This way we can store +/// the nulls with the keys as rep-def values the same as any other array. +pub fn normalize_dict_nulls(array: Arc) -> Result> { + match array.data_type() { + DataType::Dictionary(key_type, _) => match key_type.as_ref() { + DataType::UInt8 => normalize_dict_nulls_impl::(array), + DataType::UInt16 => normalize_dict_nulls_impl::(array), + DataType::UInt32 => normalize_dict_nulls_impl::(array), + DataType::UInt64 => normalize_dict_nulls_impl::(array), + DataType::Int8 => normalize_dict_nulls_impl::(array), + DataType::Int16 => normalize_dict_nulls_impl::(array), + DataType::Int32 => normalize_dict_nulls_impl::(array), + DataType::Int64 => normalize_dict_nulls_impl::(array), + _ => Err(Error::not_supported_source( + format!("Unsupported dictionary key type: {}", key_type).into(), + )), + }, + _ => Err(Error::internal(format!( + "Data type is not a dictionary: {}", + array.data_type() + ))), + } +} + +fn clear_out_of_range_null_keys_impl( + array: Arc, +) -> Result> { + let dict_array = array.as_dictionary_opt::().expect_ok()?; + let num_values = dict_array.values().len(); + let Some(nulls) = dict_array.keys().nulls() else { + return Ok(array); + }; + + // There is no valid replacement key for an empty dictionary, so that case + // requires separate handling and must remain unchanged here. + if num_values == 0 { + return Ok(array); + } + + let has_out_of_range_null_key = dict_array + .keys() + .values() + .iter() + .zip(nulls.iter()) + .any(|(key, is_valid)| !is_valid && key.to_usize().is_none_or(|key| key >= num_values)); + if !has_out_of_range_null_key { + return Ok(array); + } + + // Building from the logical iterator writes the default physical key into + // every null slot while preserving the original validity bitmap. + let keys = PrimitiveArray::::from_iter(dict_array.keys().iter()); + let values = dict_array.values().clone(); + Ok(Arc::new(DictionaryArray::::try_new(keys, values)?) as Arc) +} + +/// Replaces out-of-range physical keys in null dictionary slots with a valid key. +/// +/// Arrow permits arbitrary keys in null slots, but the structural encoder removes +/// key validity after recording it as rep-def. The replacement keeps the array +/// valid when that null buffer is removed without changing its logical values. +pub(super) fn clear_out_of_range_null_keys(array: Arc) -> Result> { + match array.data_type() { + DataType::Dictionary(key_type, _) => match key_type.as_ref() { + DataType::UInt8 => clear_out_of_range_null_keys_impl::(array), + DataType::UInt16 => clear_out_of_range_null_keys_impl::(array), + DataType::UInt32 => clear_out_of_range_null_keys_impl::(array), + DataType::UInt64 => clear_out_of_range_null_keys_impl::(array), + DataType::Int8 => clear_out_of_range_null_keys_impl::(array), + DataType::Int16 => clear_out_of_range_null_keys_impl::(array), + DataType::Int32 => clear_out_of_range_null_keys_impl::(array), + DataType::Int64 => clear_out_of_range_null_keys_impl::(array), + _ => Err(Error::not_supported_source( + format!("Unsupported dictionary key type: {}", key_type).into(), + )), + }, + _ => Err(Error::internal(format!( + "Data type is not a dictionary: {}", + array.data_type() + ))), + } +} + +fn dict_encode_variable_width( + variable_width_data_block: &VariableWidthBlock, + bits_per_offset: u8, + max_dict_entries: u32, + max_encoded_size: usize, +) -> Option<(DataBlock, DataBlock)> +where + T: ArrowNativeType, + usize: TryFrom, +{ + use std::collections::hash_map::Entry; + let mut map = HashMap::new(); + let offsets = variable_width_data_block + .offsets + .borrow_to_typed_slice::(); + let offsets = offsets.as_ref(); + + let max_len = variable_width_data_block + .get_stat(Stat::MaxLength) + .expect("VariableWidth DataBlock should have valid `Stat::MaxLength` statistics"); + let max_len = max_len.as_primitive::().value(0); + + let max_dict_data_len = variable_width_data_block.data.len(); + let max_len: usize = max_len.try_into().unwrap_or(usize::MAX); + let dict_data_capacity = max_len + .saturating_mul(32) + .max(1024) + .min(max_dict_data_len) + .min(max_encoded_size); + + let mut dictionary_buffer: Vec = Vec::with_capacity(dict_data_capacity); + let mut dictionary_offsets_buffer = vec![T::default()]; + let mut curr_idx = 0; + let mut indices_buffer = Vec::with_capacity(variable_width_data_block.num_values as usize); + let bytes_per_offset = (bits_per_offset / 8) as usize; + + for window in offsets.windows(2) { + let start = usize::try_from(window[0]).ok()?; + let end = usize::try_from(window[1]).ok()?; + if start > end || end > variable_width_data_block.data.len() { + return None; + } + + let key = &variable_width_data_block.data[start..end]; + + let idx = match map.entry(U8SliceKey(key)) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + if max_dict_entries == 0 || curr_idx as u32 >= max_dict_entries { + return None; + } + if curr_idx == i32::MAX { + return None; + } + dictionary_buffer.extend_from_slice(key); + let dict_offset = T::from_usize(dictionary_buffer.len())?; + dictionary_offsets_buffer.push(dict_offset); + let idx = curr_idx; + entry.insert(idx); + curr_idx += 1; + idx + } + }; + + indices_buffer.push(idx); + + let indices_bytes = indices_buffer + .len() + .saturating_mul(DICT_INDICES_BITS_PER_VALUE as usize / 8); + let offsets_bytes = dictionary_offsets_buffer + .len() + .saturating_mul(bytes_per_offset); + let encoded_size = dictionary_buffer + .len() + .saturating_add(indices_bytes) + .saturating_add(offsets_bytes); + if encoded_size > max_encoded_size { + return None; + } + } + + let mut dictionary_data_block = DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::reinterpret_vec(dictionary_buffer), + offsets: LanceBuffer::reinterpret_vec(dictionary_offsets_buffer), + bits_per_offset, + num_values: curr_idx as u64, + block_info: BlockInfo::default(), + }); + dictionary_data_block.compute_stat(); + + let mut indices_data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(indices_buffer), + bits_per_value: DICT_INDICES_BITS_PER_VALUE, + num_values: variable_width_data_block.num_values, + block_info: BlockInfo::default(), + }); + indices_data_block.compute_stat(); + + Some((indices_data_block, dictionary_data_block)) +} + +/// Dictionary encodes a data block +/// +/// Currently only supported for some common cases (string / binary / 64-bit / 128-bit) +/// +/// Returns a block of indices (will always be a fixed width data block) and a block of dictionary +pub fn dictionary_encode( + data_block: &DataBlock, + max_dict_entries: u32, + max_encoded_size: usize, +) -> Option<(DataBlock, DataBlock)> { + match data_block { + DataBlock::FixedWidth(fixed_width_data_block) => { + use std::collections::hash_map::Entry; + + let bytes_per_value = match fixed_width_data_block.bits_per_value { + 64 => 8usize, + 128 => 16usize, + _ => return None, + }; + + match fixed_width_data_block.bits_per_value { + 64 => { + let mut map = HashMap::new(); + let u64_slice = fixed_width_data_block.data.borrow_to_typed_slice::(); + let u64_slice = u64_slice.as_ref(); + let mut dictionary_buffer = + Vec::with_capacity((fixed_width_data_block.num_values as usize).min(1024)); + let mut indices_buffer = + Vec::with_capacity(fixed_width_data_block.num_values as usize); + let mut curr_idx: i32 = 0; + + for &value in u64_slice.iter() { + let idx = match map.entry(value) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + if max_dict_entries == 0 || curr_idx as u32 >= max_dict_entries { + return None; + } + if curr_idx == i32::MAX { + return None; + } + dictionary_buffer.push(value); + let idx = curr_idx; + entry.insert(idx); + curr_idx += 1; + idx + } + }; + indices_buffer.push(idx); + let dict_bytes = dictionary_buffer.len().saturating_mul(bytes_per_value); + let indices_bytes = indices_buffer + .len() + .saturating_mul(DICT_INDICES_BITS_PER_VALUE as usize / 8); + let encoded_size = dict_bytes.saturating_add(indices_bytes); + if encoded_size > max_encoded_size { + return None; + } + } + + let mut dictionary_data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(dictionary_buffer), + bits_per_value: 64, + num_values: curr_idx as u64, + block_info: BlockInfo::default(), + }); + dictionary_data_block.compute_stat(); + let mut indices_data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(indices_buffer), + bits_per_value: DICT_INDICES_BITS_PER_VALUE, + num_values: fixed_width_data_block.num_values, + block_info: BlockInfo::default(), + }); + indices_data_block.compute_stat(); + + Some((indices_data_block, dictionary_data_block)) + } + 128 => { + // TODO: a follow up PR to support `FixedWidth DataBlock with bits_per_value == 256`. + let mut map = HashMap::new(); + let u128_slice = fixed_width_data_block.data.borrow_to_typed_slice::(); + let u128_slice = u128_slice.as_ref(); + let mut dictionary_buffer = + Vec::with_capacity((fixed_width_data_block.num_values as usize).min(1024)); + let mut indices_buffer = + Vec::with_capacity(fixed_width_data_block.num_values as usize); + let mut curr_idx: i32 = 0; + + for &value in u128_slice.iter() { + let idx = match map.entry(value) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + if max_dict_entries == 0 || curr_idx as u32 >= max_dict_entries { + return None; + } + if curr_idx == i32::MAX { + return None; + } + dictionary_buffer.push(value); + let idx = curr_idx; + entry.insert(idx); + curr_idx += 1; + idx + } + }; + indices_buffer.push(idx); + let dict_bytes = dictionary_buffer.len().saturating_mul(bytes_per_value); + let indices_bytes = indices_buffer + .len() + .saturating_mul(DICT_INDICES_BITS_PER_VALUE as usize / 8); + let encoded_size = dict_bytes.saturating_add(indices_bytes); + if encoded_size > max_encoded_size { + return None; + } + } + + let mut dictionary_data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(dictionary_buffer), + bits_per_value: DICT_FIXED_WIDTH_BITS_PER_VALUE, + num_values: curr_idx as u64, + block_info: BlockInfo::default(), + }); + dictionary_data_block.compute_stat(); + let mut indices_data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(indices_buffer), + bits_per_value: DICT_INDICES_BITS_PER_VALUE, + num_values: fixed_width_data_block.num_values, + block_info: BlockInfo::default(), + }); + indices_data_block.compute_stat(); + + Some((indices_data_block, dictionary_data_block)) + } + _ => None, + } + } + DataBlock::VariableWidth(variable_width_data_block) => { + match variable_width_data_block.bits_per_offset { + 32 => dict_encode_variable_width::( + variable_width_data_block, + 32, + max_dict_entries, + max_encoded_size, + ), + 64 => dict_encode_variable_width::( + variable_width_data_block, + 64, + max_dict_entries, + max_encoded_size, + ), + _ => None, + } + } + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + buffer::LanceBuffer, + data::{BlockInfo, FixedWidthDataBlock}, + }; + use arrow_array::{Array, StringArray}; + use std::sync::Arc; + + #[test] + fn test_dictionary_encode_abort_fixed_width() { + // Create a u128 block with very high cardinality where dict encoding + // would result in larger data (dictionary overhead + indices > original) + let num_values = 120u64; + + // Create actual data: each value is unique u128 so dictionary encode will not be helpful + let mut data = Vec::with_capacity(num_values as usize); + for i in 0..num_values { + data.push(i as u128); + } + + let mut data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: DICT_FIXED_WIDTH_BITS_PER_VALUE, + data: LanceBuffer::reinterpret_vec(data), + num_values, + block_info: BlockInfo::default(), + }); + + // Compute stats naturally + data_block.compute_stat(); + + // Dictionary encoding should abort and return None + let max_encoded_size = usize::try_from(data_block.data_size()).unwrap_or(usize::MAX); + let result = dictionary_encode(&data_block, 1000, max_encoded_size); + assert!( + result.is_none(), + "Dictionary encoding should abort for high cardinality u128 data" + ); + } + + #[test] + fn test_dictionary_encode_success_fixed_width() { + // Create a u128 block with low cardinality where dict encoding helps + let num_values = 120u64; + let cardinality = 3u64; + + // Create data with few unique u128 values + let mut data = Vec::with_capacity(num_values as usize); + for i in 0..num_values { + data.push((i % cardinality) as u128); + } + + let mut data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: DICT_FIXED_WIDTH_BITS_PER_VALUE, + data: LanceBuffer::reinterpret_vec(data), + num_values, + block_info: BlockInfo::default(), + }); + + // Compute stats naturally + data_block.compute_stat(); + + // Dictionary encoding should succeed and return Some + let max_encoded_size = usize::try_from(data_block.data_size()).unwrap_or(usize::MAX); + let result = dictionary_encode(&data_block, 1000, max_encoded_size); + assert!( + result.is_some(), + "Dictionary encoding should succeed for low cardinality u128 data" + ); + + if let Some((indices, dictionary)) = result { + // Verify indices block + if let DataBlock::FixedWidth(indices_block) = indices { + assert_eq!(indices_block.num_values, num_values); + assert_eq!(indices_block.bits_per_value, DICT_INDICES_BITS_PER_VALUE); + } else { + panic!("Expected FixedWidth indices block"); + } + + // Verify dictionary block + if let DataBlock::FixedWidth(dict_block) = dictionary { + assert_eq!(dict_block.num_values, cardinality); + assert_eq!(dict_block.bits_per_value, DICT_FIXED_WIDTH_BITS_PER_VALUE); + } else { + panic!("Expected FixedWidth dictionary block"); + } + } + } + + #[test] + fn test_dictionary_encode_abort_variable_width() { + // Create a variable-width block with high cardinality where dict encoding + // won't provide sufficient benefit + let num_values = 120u64; + let mut values = Vec::with_capacity(num_values as usize); + for i in 0..num_values { + values.push(format!("unique_value_{:04}", i)); + } + let array = StringArray::from(values); + // from_array already computes stats + let data_block = DataBlock::from_array(Arc::new(array) as Arc); + + // Dictionary encoding should abort and return None + let max_encoded_size = usize::try_from(data_block.data_size()).unwrap_or(usize::MAX); + let result = dictionary_encode(&data_block, 10, max_encoded_size); + assert!( + result.is_none(), + "Dictionary encoding should abort for high cardinality string data" + ); + } + + #[test] + fn test_dictionary_encode_success_low_cardinality() { + // Create a variable-width block with low cardinality where dict encoding helps + let num_values = 120u64; + let cardinality = 3u64; + + let mut values = Vec::with_capacity(num_values as usize); + for i in 0..num_values { + values.push(format!("value_{}", i % cardinality)); + } + + let array = StringArray::from(values); + let data_block = DataBlock::from_array(Arc::new(array) as Arc); + + // Dictionary encoding should succeed and return Some + let max_encoded_size = usize::try_from(data_block.data_size()).unwrap_or(usize::MAX); + let result = dictionary_encode(&data_block, 100, max_encoded_size); + assert!( + result.is_some(), + "Dictionary encoding should succeed for low cardinality data" + ); + + if let Some((indices, dictionary)) = result { + // Verify indices block + if let DataBlock::FixedWidth(indices_block) = indices { + assert_eq!(indices_block.num_values, num_values); + assert_eq!(indices_block.bits_per_value, DICT_INDICES_BITS_PER_VALUE); + } else { + panic!("Expected FixedWidth indices block"); + } + + // Verify dictionary block + if let DataBlock::VariableWidth(dict_block) = dictionary { + assert_eq!(dict_block.num_values, cardinality); + } else { + panic!("Expected VariableWidth dictionary block"); + } + } + } + + #[test] + fn test_dictionary_encode_invalid_offset_width_returns_none() { + let array = StringArray::from(vec!["a", "b", "c", "a"]); + let data_block = DataBlock::from_array(Arc::new(array) as Arc); + let invalid_block = match data_block { + DataBlock::VariableWidth(mut var) => { + var.bits_per_offset = 16; + DataBlock::VariableWidth(var) + } + other => panic!("Expected VariableWidth data block, got {:?}", other), + }; + let max_encoded_size = usize::try_from(invalid_block.data_size()).unwrap_or(usize::MAX); + assert!(dictionary_encode(&invalid_block, 100, max_encoded_size).is_none()); + } + + #[test] + fn test_dictionary_encode_respects_size_limit() { + let num_values = 10_000u64; + let cardinality = 50u64; + + let mut values = Vec::with_capacity(num_values as usize); + for i in 0..num_values { + values.push(format!("value_{:08}", i % cardinality)); + } + + let array = StringArray::from(values); + let data_block = DataBlock::from_array(Arc::new(array) as Arc); + + let full_size = usize::try_from(data_block.data_size()).unwrap_or(usize::MAX); + let too_small_limit = full_size / 10; + assert!(dictionary_encode(&data_block, 1000, too_small_limit).is_none()); + assert!(dictionary_encode(&data_block, 1000, full_size).is_some()); + } + + #[test] + fn test_dictionary_encode_respects_entry_limit() { + let num_values = 10_000u64; + let cardinality = 200u64; + + let mut values = Vec::with_capacity(num_values as usize); + for i in 0..num_values { + values.push(format!("value_{:08}", i % cardinality)); + } + + let array = StringArray::from(values); + let data_block = DataBlock::from_array(Arc::new(array) as Arc); + + let max_encoded_size = usize::try_from(data_block.data_size()).unwrap_or(usize::MAX); + assert!(dictionary_encode(&data_block, 10, max_encoded_size).is_none()); + assert!(dictionary_encode(&data_block, 500, max_encoded_size).is_some()); + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive/fullzip.rs b/vendor/lance-encoding/src/encodings/logical/primitive/fullzip.rs new file mode 100644 index 000000000..aaafa1887 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive/fullzip.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Routines for encoding and decoding full-zip data +//! +//! Full-zip is one of the two structural encodings in Lance 2.1. +//! In this approach the various compressed buffers are zipped +//! together so that all parts of a value are stored contiguously in memory. +//! +//! This requires transparent compression and is most suitable for +//! large data types. + +use crate::{ + data::{DataBlock, FixedWidthDataBlock, VariableWidthBlock}, + format::pb21::CompressiveEncoding, +}; + +use lance_core::Result; + +/// Per-value compression must either: +/// +/// A single buffer of fixed-width values +/// A single buffer of value data and a buffer of offsets +/// +/// TODO: In the future we may allow metadata buffers +#[derive(Debug)] +pub enum PerValueDataBlock { + Fixed(FixedWidthDataBlock), + Variable(VariableWidthBlock), +} + +impl PerValueDataBlock { + pub fn data_size(&self) -> u64 { + match self { + Self::Fixed(fixed) => fixed.data_size(), + Self::Variable(variable) => variable.data_size(), + } + } +} + +/// Trait for compression algorithms that are suitable for use in the zipped structural encoding +/// +/// This compression must return either a FixedWidthDataBlock or a VariableWidthBlock. This is because +/// we need to zip the data and those are the only two blocks we know how to zip today. +/// +/// In addition, the compressed data must be able to be decompressed in a random-access fashion. +/// This means that the decompression algorithm must be able to decompress any value without +/// decompressing all values before it. +pub trait PerValueCompressor: std::fmt::Debug + Send + Sync { + /// Compress the data into a single buffer + /// + /// Also returns a description of the compression that can be used to decompress when reading the data back + fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)>; +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive/layout.rs b/vendor/lance-encoding/src/encodings/logical/primitive/layout.rs new file mode 100644 index 000000000..95df12ede --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive/layout.rs @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use lance_core::Result; + +use crate::repdef::MiniBlockRepDefBudget; + +/// Runs automatic sparse planning only after the dense mini-block budget makes it useful. +pub(super) fn select_automatic_sparse( + requested_encoding: Option<&str>, + dense_budget: &MiniBlockRepDefBudget, + candidate: impl FnOnce() -> Result>, +) -> Result> { + if requested_encoding.is_some() + || !matches!( + dense_budget, + MiniBlockRepDefBudget::RequiresPageSplit(_) + | MiniBlockRepDefBudget::SingleRowOverBudget(_) + ) + { + return Ok(None); + } + candidate() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::constants::{ + STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_SPARSE, + }; + + #[test] + fn within_budget_does_not_construct_sparse_candidate() { + let selected = + select_automatic_sparse::<()>(None, &MiniBlockRepDefBudget::WithinBudget, || { + panic!("within-budget pages must not construct sparse candidates") + }) + .unwrap(); + assert!(selected.is_none()); + } + + #[test] + fn over_budget_selects_only_eligible_candidates() { + let split = MiniBlockRepDefBudget::RequiresPageSplit(Vec::new()); + let selected = select_automatic_sparse(None, &split, || Ok(Some(42))).unwrap(); + assert_eq!(selected, Some(42)); + + let ineligible = select_automatic_sparse::<()>(None, &split, || Ok(None)).unwrap(); + assert!(ineligible.is_none()); + + let unsplittable = MiniBlockRepDefBudget::SingleRowOverBudget(70_000); + let selected = select_automatic_sparse(None, &unsplittable, || Ok(Some(7))).unwrap(); + assert_eq!(selected, Some(7)); + } + + #[test] + fn explicit_modes_and_lance_2_2_do_not_auto_select() { + let split = MiniBlockRepDefBudget::RequiresPageSplit(Vec::new()); + for requested in [ + STRUCTURAL_ENCODING_MINIBLOCK, + STRUCTURAL_ENCODING_FULLZIP, + STRUCTURAL_ENCODING_SPARSE, + ] { + let selected = select_automatic_sparse::<()>(Some(requested), &split, || { + panic!("explicit modes must not invoke automatic sparse planning") + }) + .unwrap(); + assert!(selected.is_none()); + } + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive/miniblock.rs b/vendor/lance-encoding/src/encodings/logical/primitive/miniblock.rs new file mode 100644 index 000000000..0ee408a85 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive/miniblock.rs @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Routines for encoding and decoding miniblock data +//! +//! Miniblock encoding is one of the two structural encodings in Lance 2.1. +//! In this approach the data is compressed into a series of chunks put into +//! a single buffer. +//! +//! A chunk must be encoded or decoded as a unit. There is a small amount of +//! chunk metadata such as the number and size of each buffer in the chunk. +//! +//! Any form of compression can be used since we are compressing and decompressing +//! entire chunks. +use crate::{buffer::LanceBuffer, data::DataBlock, format::pb21::CompressiveEncoding}; + +use lance_core::Result; + +pub const MAX_MINIBLOCK_BYTES: u64 = 8 * 1024 - 6; + +const DEFAULT_MAX_MINIBLOCK_VALUES: u64 = 4096; +/// Maximum number of values that any mini-block decoder accepts from page metadata. +pub(crate) const MAX_CONFIGURABLE_MINIBLOCK_VALUES: u64 = 32768; + +fn parse_max_miniblock_values() -> u64 { + let val = std::env::var("LANCE_MINIBLOCK_MAX_VALUES") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(DEFAULT_MAX_MINIBLOCK_VALUES); + val.clamp(1, MAX_CONFIGURABLE_MINIBLOCK_VALUES) +} + +pub static MAX_MINIBLOCK_VALUES: std::sync::LazyLock = + std::sync::LazyLock::new(parse_max_miniblock_values); + +/// Maximum number of rep/def levels the structural planner should place into +/// a single mini-block chunk. +pub fn max_repdef_levels_per_chunk(bits_per_level: u64) -> u64 { + debug_assert!(bits_per_level > 0); + const REPDEF_BUDGET_BITS: u64 = 16 * 1024 * 8; + let budgeted_levels = REPDEF_BUDGET_BITS / bits_per_level; + budgeted_levels.min(u16::MAX as u64) +} + +/// Page data that has been compressed into a series of chunks put into +/// a single buffer. +#[derive(Debug)] +pub struct MiniBlockCompressed { + /// The buffers of compressed data + pub data: Vec, + /// Describes the size of each chunk + pub chunks: Vec, + /// The number of values in the entire page + pub num_values: u64, +} + +/// Per-page framing details that can affect a mini-block compressor's choice. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MiniBlockCompressionContext { + common_chunk_buffers: u64, + support_large_chunk: bool, + allow_generic_offsets: bool, +} + +impl MiniBlockCompressionContext { + /// Creates the framing context supplied by the owning mini-block page. + pub fn new( + common_chunk_buffers: u64, + support_large_chunk: bool, + allow_generic_offsets: bool, + ) -> Self { + Self { + common_chunk_buffers, + support_large_chunk, + allow_generic_offsets, + } + } +} + +/// Describes the size of a mini-block chunk of data +/// +/// Mini-block chunks are designed to be small (just a few disk sectors) +/// and contain a power-of-two number of values (except for the last chunk) +/// +/// By default we limit a chunk to 4Ki values and slightly less than +/// 8KiB of compressed value data. The byte budget remains the primary +/// constraint, so only encodings that compress many values into that +/// budget can use larger value counts when explicitly configured. +/// +/// The maximum number of values per chunk can be configured via the +/// `LANCE_MINIBLOCK_MAX_VALUES` environment variable. This is only +/// useful in extremely bandwidth-limited environments; the default is +/// appropriate for local disks and same-region cloud object storage. +#[derive(Debug)] +pub struct MiniBlockChunk { + // The size in bytes of each buffer in the chunk. + // + // In Lance 2.1, the chunk size is limited to 32KiB, so only 16-bits are used. + // Since Lance 2.2, the chunk size uses u32 to support larger chunk size + pub buffer_sizes: Vec, + // The log (base 2) of the number of values in the chunk. If this is the final chunk + // then this should be 0 (the number of values will be calculated by subtracting the + // size of all other chunks from the total size of the page) + // + // For example, 1 would mean there are 2 values in the chunk and 15 would mean there + // are 32Ki values in the chunk. + // + // This must be <= log2(MAX_MINIBLOCK_VALUES) (i.e. <= 12 at the default of 4096) + pub log_num_values: u8, +} + +impl MiniBlockChunk { + /// Gets the number of values in this block + /// + /// This requires `vals_in_prev_blocks` and `total_num_values` because the + /// last block in a page is a special case which stores 0 for log_num_values + /// and, in that case, the number of values is determined by subtracting + /// `vals_in_prev_blocks` from `total_num_values` + pub fn num_values(&self, vals_in_prev_blocks: u64, total_num_values: u64) -> u64 { + if self.log_num_values == 0 { + total_num_values - vals_in_prev_blocks + } else { + 1 << self.log_num_values + } + } +} + +/// Trait for compression algorithms that are suitable for use in the miniblock structural encoding +/// +/// These compression algorithms should be capable of encoding the data into small chunks +/// where each chunk (except the last) has 2^N values (N can vary between chunks) +pub trait MiniBlockCompressor: std::fmt::Debug + Send + Sync { + /// Compress a `page` of data into multiple chunks + /// + /// See [`MiniBlockCompressed`] for details on how chunks should be sized. + /// + /// This method also returns a description of the encoding applied that will be + /// used at decode time to read the data. + fn compress( + &self, + context: MiniBlockCompressionContext, + page: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)>; +} + +#[cfg(test)] +mod tests { + use serial_test::serial; + + use super::*; + + #[test] + #[serial] + fn test_parse_default() { + unsafe { std::env::remove_var("LANCE_MINIBLOCK_MAX_VALUES") }; + assert_eq!(parse_max_miniblock_values(), 4096); + } + + #[test] + #[serial] + fn test_parse_custom_value() { + unsafe { std::env::set_var("LANCE_MINIBLOCK_MAX_VALUES", "256") }; + assert_eq!(parse_max_miniblock_values(), 256); + unsafe { std::env::remove_var("LANCE_MINIBLOCK_MAX_VALUES") }; + } + + #[test] + #[serial] + fn test_parse_can_raise_to_32k() { + unsafe { std::env::set_var("LANCE_MINIBLOCK_MAX_VALUES", "32768") }; + assert_eq!(parse_max_miniblock_values(), 32768); + unsafe { std::env::remove_var("LANCE_MINIBLOCK_MAX_VALUES") }; + } + + #[test] + #[serial] + fn test_parse_clamps_zero_to_one() { + unsafe { std::env::set_var("LANCE_MINIBLOCK_MAX_VALUES", "0") }; + assert_eq!(parse_max_miniblock_values(), 1); + unsafe { std::env::remove_var("LANCE_MINIBLOCK_MAX_VALUES") }; + } + + #[test] + #[serial] + fn test_parse_clamps_above_max() { + unsafe { std::env::set_var("LANCE_MINIBLOCK_MAX_VALUES", "99999") }; + assert_eq!( + parse_max_miniblock_values(), + MAX_CONFIGURABLE_MINIBLOCK_VALUES + ); + unsafe { std::env::remove_var("LANCE_MINIBLOCK_MAX_VALUES") }; + } + + #[test] + #[serial] + fn test_parse_invalid_falls_back_to_default() { + unsafe { std::env::set_var("LANCE_MINIBLOCK_MAX_VALUES", "not_a_number") }; + assert_eq!(parse_max_miniblock_values(), DEFAULT_MAX_MINIBLOCK_VALUES); + unsafe { std::env::remove_var("LANCE_MINIBLOCK_MAX_VALUES") }; + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive/sparse.rs b/vendor/lance-encoding/src/encodings/logical/primitive/sparse.rs new file mode 100644 index 000000000..ec8e0b418 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive/sparse.rs @@ -0,0 +1,5628 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use super::*; +use arrow_array::new_empty_array; +use arrow_buffer::ArrowNativeType; + +fn invalid_enum( + value: std::result::Result, + label: &str, +) -> Result { + value.map_err(|error| { + Error::invalid_input_source( + format!("Sparse structural {label} has an invalid enum value: {error}").into(), + ) + }) +} + +pub(super) mod writer; + +fn usize_from_u64(value: u64, label: &str) -> Result { + usize::try_from(value).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} {value} exceeds usize::MAX").into(), + ) + }) +} + +/// Native sparse structural representation used by the 2.3 sparse layout. +/// +/// Layers are stored from outer-most to inner-most, matching the order Arrow structural +/// encoders record offsets and validity. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SparseStructuralPlan { + pub(crate) layers: Vec, + pub(crate) num_items: u64, + pub(crate) num_visible_items: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SparsePositionSet { + Empty, + All { len: u64 }, + Range { start: u64, len: u64 }, + Explicit(Vec), +} + +impl SparsePositionSet { + pub(crate) fn from_positions( + positions: Vec, + domain_len: u64, + label: &str, + ) -> Result { + if positions.is_empty() { + return Ok(Self::Empty); + } + for window in positions.windows(2) { + let [previous, current] = window else { + continue; + }; + if previous >= current { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} positions must be strictly increasing") + .into(), + )); + } + } + if let Some(position) = positions.iter().find(|position| **position >= domain_len) { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} position {} is outside layer with {} slots", + position, domain_len + ) + .into(), + )); + } + + let len = u64::try_from(positions.len()).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} position count exceeds u64::MAX").into(), + ) + })?; + let first = positions.first().copied().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} positions are unexpectedly empty").into(), + ) + })?; + let last = positions.last().copied().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} positions are unexpectedly empty").into(), + ) + })?; + if first == 0 && len == domain_len && domain_len > 0 && last == domain_len - 1 { + return Ok(Self::All { len: domain_len }); + } + if last - first + 1 == len { + return Ok(Self::Range { start: first, len }); + } + Ok(Self::Explicit(positions)) + } + + pub(crate) fn empty() -> Self { + Self::Empty + } + + pub(crate) fn all(len: u64) -> Self { + if len == 0 { + Self::Empty + } else { + Self::All { len } + } + } + + pub(crate) fn range(start: u64, len: u64) -> Self { + if len == 0 { + Self::Empty + } else { + Self::Range { start, len } + } + } + + pub(crate) fn len(&self) -> u64 { + match self { + Self::Empty => 0, + Self::All { len } | Self::Range { len, .. } => *len, + Self::Explicit(positions) => positions.len() as u64, + } + } + + pub(crate) fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub(crate) fn deep_size(&self) -> usize { + match self { + Self::Explicit(positions) => positions.len() * std::mem::size_of::(), + Self::Empty | Self::All { .. } | Self::Range { .. } => 0, + } + } + + pub(crate) fn materialize(&self) -> Result> { + match self { + Self::Empty => Ok(Vec::new()), + Self::All { len } => Self::materialize_range(0, *len), + Self::Range { start, len } => Self::materialize_range(*start, *len), + Self::Explicit(positions) => Ok(positions.clone()), + } + } + + fn materialize_range(start: u64, len: u64) -> Result> { + let len_usize = usize::try_from(len).map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural position range length {} exceeds usize::MAX", + len + ) + .into(), + ) + })?; + let end = start.checked_add(len).ok_or_else(|| { + Error::invalid_input_source("Sparse structural position range overflows".into()) + })?; + let mut positions = Vec::with_capacity(len_usize); + positions.extend(start..end); + Ok(positions) + } + + fn contains(&self, position: u64) -> bool { + match self { + Self::Empty => false, + Self::All { len } => position < *len, + Self::Range { start, len } => { + position >= *start && position < start.saturating_add(*len) + } + Self::Explicit(positions) => positions.binary_search(&position).is_ok(), + } + } + + fn is_subset_of(&self, other: &Self, domain_len: u64) -> Result { + self.validate_domain(domain_len, "subset")?; + other.validate_domain(domain_len, "superset")?; + Ok(match self { + Self::Empty => true, + Self::All { .. } => other.len() == domain_len, + Self::Range { start, len } => { + let end = start.checked_add(*len).ok_or_else(|| { + Error::invalid_input_source("Sparse structural subset range overflows".into()) + })?; + match other { + Self::All { .. } => true, + Self::Range { + start: other_start, + len: other_len, + } => { + let other_end = other_start.saturating_add(*other_len); + *start >= *other_start && end <= other_end + } + Self::Explicit(positions) => { + let first = positions.partition_point(|position| *position < *start); + let last = positions.partition_point(|position| *position < end); + u64::try_from(last.saturating_sub(first)).ok() == Some(*len) + } + Self::Empty => false, + } + } + Self::Explicit(positions) => positions.iter().all(|position| other.contains(*position)), + }) + } + + fn is_disjoint(&self, other: &Self, domain_len: u64) -> Result { + self.validate_domain(domain_len, "first disjoint set")?; + other.validate_domain(domain_len, "second disjoint set")?; + let (smaller, larger) = if self.len() <= other.len() { + (self, other) + } else { + (other, self) + }; + Ok(match smaller { + Self::Empty => true, + Self::All { .. } => larger.is_empty(), + Self::Range { start, len } => { + let end = start.saturating_add(*len); + match larger { + Self::Empty => true, + Self::All { .. } => false, + Self::Range { + start: other_start, + len: other_len, + } => end <= *other_start || other_start.saturating_add(*other_len) <= *start, + Self::Explicit(positions) => { + let index = positions.partition_point(|position| *position < *start); + positions.get(index).is_none_or(|position| *position >= end) + } + } + } + Self::Explicit(positions) => { + positions.iter().all(|position| !larger.contains(*position)) + } + }) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SparseValidityMeaning { + NullPositions, + ValidPositions, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SparseValiditySet { + pub(crate) meaning: SparseValidityMeaning, + pub(crate) positions: SparsePositionSet, +} + +impl SparseValiditySet { + pub(crate) fn deep_size(&self) -> usize { + self.positions.deep_size() + } + + fn contains_only_valid_positions( + &self, + positions: &SparsePositionSet, + num_slots: u64, + ) -> Result { + match self.meaning { + SparseValidityMeaning::NullPositions => { + positions.is_disjoint(&self.positions, num_slots) + } + SparseValidityMeaning::ValidPositions => { + positions.is_subset_of(&self.positions, num_slots) + } + } + } + + fn append_to(&self, validity: &mut BooleanBufferBuilder, num_slots: u64) -> Result<()> { + self.positions.validate_domain(num_slots, "validity")?; + let num_slots_usize = usize_from_u64(num_slots, "validity slot count")?; + match (self.meaning, &self.positions) { + (SparseValidityMeaning::NullPositions, SparsePositionSet::Empty) => { + validity.append_n(num_slots_usize, true); + } + (SparseValidityMeaning::ValidPositions, SparsePositionSet::Empty) => { + validity.append_n(num_slots_usize, false); + } + (SparseValidityMeaning::NullPositions, SparsePositionSet::All { .. }) => { + validity.append_n(num_slots_usize, false); + } + (SparseValidityMeaning::ValidPositions, SparsePositionSet::All { .. }) => { + validity.append_n(num_slots_usize, true); + } + (meaning, SparsePositionSet::Range { start, len }) => { + let range_end = start.checked_add(*len).ok_or_else(|| { + Error::invalid_input_source("Sparse structural validity range overflows".into()) + })?; + let default_valid = matches!(meaning, SparseValidityMeaning::NullPositions); + let range_valid = !default_valid; + validity.append_n( + usize_from_u64(*start, "validity range start")?, + default_valid, + ); + validity.append_n(usize_from_u64(*len, "validity range length")?, range_valid); + validity.append_n( + usize_from_u64(num_slots - range_end, "validity range tail")?, + default_valid, + ); + } + (_, SparsePositionSet::Explicit(_)) => { + let mut cursor = SparseValidityCursor::new(self, num_slots, "validity")?; + for slot in 0..num_slots { + validity.append(cursor.is_valid(slot)?); + } + cursor.finish()?; + } + } + Ok(()) + } +} + +impl SparsePositionSet { + fn validate_domain(&self, domain_len: u64, label: &str) -> Result<()> { + match self { + Self::Empty => {} + Self::All { len } => { + if *len != domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} all set length {} does not match domain {}", + len, domain_len + ) + .into(), + )); + } + } + Self::Range { start, len } => { + let end = start.checked_add(*len).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} range overflows").into(), + ) + })?; + if end > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} range {}..{} is outside domain {}", + start, end, domain_len + ) + .into(), + )); + } + } + Self::Explicit(positions) => { + for window in positions.windows(2) { + let [previous, current] = window else { + continue; + }; + if previous >= current { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} positions must be strictly increasing" + ) + .into(), + )); + } + } + if let Some(position) = positions.last() + && *position >= domain_len + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} position {} is outside layer with {} slots", + position, domain_len + ) + .into(), + )); + } + } + } + Ok(()) + } +} + +struct SparsePositionSetCursor<'a> { + set: &'a SparsePositionSet, + explicit: Option>>, +} + +impl<'a> SparsePositionSetCursor<'a> { + fn new(set: &'a SparsePositionSet, domain_len: u64, label: &str) -> Result { + set.validate_domain(domain_len, label)?; + let explicit = match set { + SparsePositionSet::Explicit(positions) => Some(positions.iter().peekable()), + _ => None, + }; + Ok(Self { set, explicit }) + } + + fn contains(&mut self, slot: u64) -> Result { + Ok(match self.set { + SparsePositionSet::Empty => false, + SparsePositionSet::All { .. } => true, + SparsePositionSet::Range { start, len } => { + slot >= *start && slot < start.saturating_add(*len) + } + SparsePositionSet::Explicit(_) => { + let iter = self.explicit.as_mut().ok_or_else(|| { + Error::internal("Sparse structural explicit cursor is missing".to_string()) + })?; + if let Some(position) = iter.peek() + && **position < slot + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural explicit position {} was skipped before slot {}", + **position, slot + ) + .into(), + )); + } + if iter.peek().is_some_and(|position| **position == slot) { + iter.next(); + true + } else { + false + } + } + }) + } + + fn finish(&mut self) -> Result<()> { + if let Some(iter) = self.explicit.as_mut() + && let Some(position) = iter.next() + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural explicit position {} was not consumed", + position + ) + .into(), + )); + } + Ok(()) + } +} + +struct SparseValidityCursor<'a> { + meaning: SparseValidityMeaning, + positions: SparsePositionSetCursor<'a>, +} + +impl<'a> SparseValidityCursor<'a> { + fn new(validity: &'a SparseValiditySet, domain_len: u64, label: &str) -> Result { + Ok(Self { + meaning: validity.meaning, + positions: SparsePositionSetCursor::new(&validity.positions, domain_len, label)?, + }) + } + + fn is_valid(&mut self, slot: u64) -> Result { + let is_stored = self.positions.contains(slot)?; + Ok(match self.meaning { + SparseValidityMeaning::NullPositions => !is_stored, + SparseValidityMeaning::ValidPositions => is_stored, + }) + } + + fn finish(&mut self) -> Result<()> { + self.positions.finish() + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SparseCountSet { + Empty, + Constant { + value: u64, + len: u64, + }, + Explicit { + counts: Arc<[u64]>, + offsets: Arc<[u64]>, + }, +} + +impl SparseCountSet { + pub(crate) fn from_counts(counts: Vec) -> Result { + if counts.is_empty() { + return Ok(Self::Empty); + } + if let Some(first) = counts.first().copied() + && counts.iter().all(|count| *count == first) + { + return Ok(Self::Constant { + value: first, + len: counts.len() as u64, + }); + } + let offsets = offsets_from_counts(&counts)?; + Ok(Self::Explicit { + counts: counts.into(), + offsets: offsets.into(), + }) + } + + pub(crate) fn constant(value: u64, len: u64) -> Self { + if len == 0 { + Self::Empty + } else { + Self::Constant { value, len } + } + } + + pub(crate) fn len(&self) -> u64 { + match self { + Self::Empty => 0, + Self::Constant { len, .. } => *len, + Self::Explicit { counts, .. } => counts.len() as u64, + } + } + + pub(crate) fn deep_size(&self) -> usize { + match self { + Self::Explicit { counts, offsets } => { + counts.len() * std::mem::size_of::() + + offsets.len() * std::mem::size_of::() + } + Self::Empty | Self::Constant { .. } => 0, + } + } + + pub(crate) fn materialize(&self) -> Result> { + match self { + Self::Empty => Ok(Vec::new()), + Self::Constant { value, len } => { + let len = usize::try_from(*len).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count set length exceeds usize::MAX".into(), + ) + })?; + Ok(vec![*value; len]) + } + Self::Explicit { counts, .. } => Ok(counts.to_vec()), + } + } + + pub(crate) fn sum(&self) -> Result { + match self { + Self::Empty => Ok(0), + Self::Constant { value, len } => value.checked_mul(*len).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural constant count sum overflows: value={}, len={}", + value, len + ) + .into(), + ) + }), + Self::Explicit { offsets, .. } => offsets.last().copied().ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural explicit count offsets are empty".into(), + ) + }), + } + } + + fn validate_positive(&self) -> Result<()> { + let has_zero = match self { + Self::Empty => false, + Self::Constant { value, .. } => *value == 0, + Self::Explicit { counts, .. } => counts.contains(&0), + }; + if has_zero { + return Err(Error::invalid_input_source( + "Sparse structural non-empty list count is zero".into(), + )); + } + Ok(()) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SparseStructuralLayerPlan { + Validity { + num_slots: u64, + validity: SparseValiditySet, + }, + List { + num_slots: u64, + num_child_slots: u64, + non_empty_positions: SparsePositionSet, + counts: SparseCountSet, + validity: SparseValiditySet, + }, + FixedSizeList { + num_slots: u64, + dimension: u64, + validity: SparseValiditySet, + }, +} + +impl SparseStructuralPlan { + fn expected_num_items( + layers: &[SparseStructuralLayerPlan], + num_visible_items: u64, + ) -> Result { + layers.iter().try_fold(num_visible_items, |items, layer| { + let additional = match layer { + SparseStructuralLayerPlan::List { + num_slots, + non_empty_positions, + .. + } => num_slots + .checked_sub(non_empty_positions.len()) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list has more non-empty positions than slots".into(), + ) + })?, + SparseStructuralLayerPlan::Validity { .. } + | SparseStructuralLayerPlan::FixedSizeList { .. } => 0, + }; + items.checked_add(additional).ok_or_else(|| { + Error::invalid_input_source("Sparse structural item count overflows".into()) + }) + }) + } + + fn validate(&self, row_domain: u64) -> Result<()> { + usize_from_u64(self.num_visible_items, "visible item count")?; + let expected_num_items = Self::expected_num_items(&self.layers, self.num_visible_items)?; + if self.num_items != expected_num_items { + return Err(Error::invalid_input_source( + format!( + "Sparse structural item count {} does not match the {} items implied by its layers", + self.num_items, expected_num_items + ) + .into(), + )); + } + let mut expected_slots = row_domain; + for (layer_index, layer) in self.layers.iter().enumerate() { + let (num_slots, num_child_slots, validity) = match layer { + SparseStructuralLayerPlan::Validity { + num_slots, + validity, + } => (*num_slots, *num_slots, validity), + SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + } => { + counts.validate_positive()?; + if non_empty_positions.len() != counts.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list layer {} has {} non-empty positions but {} counts", + layer_index, + non_empty_positions.len(), + counts.len() + ) + .into(), + )); + } + if counts.sum()? != *num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list layer {} count sum does not match {} child slots", + layer_index, num_child_slots + ) + .into(), + )); + } + if !validity.contains_only_valid_positions(non_empty_positions, *num_slots)? { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list layer {} contains a non-empty null slot", + layer_index + ) + .into(), + )); + } + (*num_slots, *num_child_slots, validity) + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots, + dimension, + validity, + } => { + if *dimension == 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list layer {} has dimension zero", + layer_index + ) + .into(), + )); + } + let num_child_slots = num_slots.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural fixed-size-list child domain overflows".into(), + ) + })?; + (*num_slots, num_child_slots, validity) + } + }; + usize_from_u64(num_slots, "layer slot count")?; + usize_from_u64(num_child_slots, "layer child slot count")?; + if num_slots != expected_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural layer {} has {} slots, expected {}", + layer_index, num_slots, expected_slots + ) + .into(), + )); + } + validity.positions.validate_domain(num_slots, "validity")?; + expected_slots = num_child_slots; + } + if expected_slots != self.num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse structural terminal domain has {} slots, expected {} visible items", + expected_slots, self.num_visible_items + ) + .into(), + )); + } + Ok(()) + } +} + +#[derive(Debug)] +pub struct SparseStructuralUnraveler { + layers: Vec, + next_layer: usize, + pending_fixed_size_list: bool, +} + +impl SparseStructuralUnraveler { + pub(crate) fn new(plan: SparseStructuralPlan) -> Self { + let next_layer = plan.layers.len(); + Self { + layers: plan.layers, + next_layer, + pending_fixed_size_list: false, + } + } + + fn current_layer(&self) -> Option<&SparseStructuralLayerPlan> { + self.next_layer + .checked_sub(1) + .and_then(|idx| self.layers.get(idx)) + } + + fn consume_current_layer(&mut self) -> Result<()> { + self.next_layer = self.next_layer.checked_sub(1).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural metadata has fewer layers than the Arrow schema".into(), + ) + })?; + Ok(()) + } + + pub(crate) fn ensure_exhausted(&self) -> Result<()> { + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse structural metadata has an unconsumed fixed-size-list layer".into(), + )); + } + if self.next_layer != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural metadata has {} unconsumed layer(s)", + self.next_layer + ) + .into(), + )); + } + Ok(()) + } + + pub(crate) fn is_all_valid(&self) -> bool { + match self.current_layer() { + Some(SparseStructuralLayerPlan::Validity { + num_slots, + validity, + }) + | Some(SparseStructuralLayerPlan::FixedSizeList { + num_slots, + validity, + .. + }) + | Some(SparseStructuralLayerPlan::List { + num_slots, + validity, + .. + }) => match validity.meaning { + SparseValidityMeaning::NullPositions => validity.positions.is_empty(), + SparseValidityMeaning::ValidPositions => validity.positions.len() == *num_slots, + }, + None => true, + } + } + + pub(crate) fn max_lists(&self) -> Result { + match self.current_layer() { + Some(SparseStructuralLayerPlan::List { num_slots, .. }) => { + usize_from_u64(*num_slots, "list slot count") + } + _ => Ok(0), + } + } + + pub(crate) fn skip_validity(&mut self) -> Result<()> { + match self.current_layer() { + Some(SparseStructuralLayerPlan::Validity { .. }) => { + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list schema does not match a validity layer".into(), + )); + } + self.consume_current_layer()?; + } + Some(SparseStructuralLayerPlan::FixedSizeList { .. }) => { + if !self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list layer does not match the Arrow schema".into(), + )); + } + self.pending_fixed_size_list = false; + self.consume_current_layer()?; + } + None => { + return Err(Error::invalid_input_source( + "Sparse structural metadata has fewer layers than the Arrow schema".into(), + )); + } + Some(SparseStructuralLayerPlan::List { .. }) => { + return Err(Error::invalid_input_source( + "Sparse structural list layer does not match an Arrow validity layer".into(), + )); + } + } + Ok(()) + } + + pub(crate) fn unravel_validity(&mut self, validity: &mut BooleanBufferBuilder) -> Result<()> { + match self.current_layer() { + Some(SparseStructuralLayerPlan::Validity { + num_slots, + validity: layer_validity, + }) => { + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list schema does not match a validity layer".into(), + )); + } + layer_validity.append_to(validity, *num_slots)?; + self.consume_current_layer()?; + } + Some(SparseStructuralLayerPlan::FixedSizeList { + num_slots, + validity: layer_validity, + .. + }) => { + if !self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list layer does not match the Arrow schema".into(), + )); + } + layer_validity.append_to(validity, *num_slots)?; + self.pending_fixed_size_list = false; + self.consume_current_layer()?; + } + None => { + return Err(Error::invalid_input_source( + "Sparse structural metadata has fewer layers than the Arrow schema".into(), + )); + } + Some(SparseStructuralLayerPlan::List { .. }) => { + return Err(Error::invalid_input_source( + "Sparse structural list layer does not match an Arrow validity layer".into(), + )); + } + } + Ok(()) + } + + pub(crate) fn decimate(&mut self, dimension: usize) -> Result<()> { + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list layer was decimated more than once".into(), + )); + } + let Some(SparseStructuralLayerPlan::FixedSizeList { + dimension: actual_dimension, + .. + }) = self.current_layer() + else { + return Err(Error::invalid_input_source( + "Sparse structural layer does not match an Arrow fixed-size-list layer".into(), + )); + }; + if usize_from_u64(*actual_dimension, "fixed-size-list dimension")? != dimension { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list dimension {} does not match Arrow dimension {}", + actual_dimension, dimension + ) + .into(), + )); + } + self.pending_fixed_size_list = true; + Ok(()) + } + + fn to_offset(value: u64) -> Result { + let value = usize::try_from(value).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural offset {} exceeds usize::MAX", value).into(), + ) + })?; + T::from_usize(value).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural offset does not fit the Arrow offset type".into(), + ) + }) + } + + pub(crate) fn unravel_offsets( + &mut self, + offsets: &mut Vec, + validity: Option<&mut BooleanBufferBuilder>, + ) -> Result<()> { + let Some(SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity: layer_validity, + .. + }) = self.current_layer() + else { + return Err(Error::invalid_input_source( + "Sparse structural layer does not match an Arrow list layer".into(), + )); + }; + if self.pending_fixed_size_list { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list schema does not match an Arrow list layer".into(), + )); + } + + if non_empty_positions.len() != counts.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list has {} non-empty positions but {} counts", + non_empty_positions.len(), + counts.len() + ) + .into(), + )); + } + let actual_child_slots = counts.sum()?; + if actual_child_slots != *num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match child slots {}", + actual_child_slots, num_child_slots + ) + .into(), + )); + } + + let mut current_offset = offsets + .last() + .map(|offset| offset.as_usize() as u64) + .unwrap_or(0); + if offsets.is_empty() { + offsets.push(Self::to_offset(current_offset)?); + } + + if non_empty_positions.is_empty() { + if let Some(validity) = validity { + layer_validity.append_to(validity, *num_slots)?; + } + let offset = Self::to_offset(current_offset)?; + let new_len = offsets + .len() + .checked_add(usize_from_u64(*num_slots, "list slot count")?) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list offset length overflows usize".into(), + ) + })?; + offsets.resize(new_len, offset); + self.consume_current_layer()?; + return Ok(()); + } + + let non_empty_positions = non_empty_positions.materialize()?; + let counts = counts.materialize()?; + let mut non_empty_iter = non_empty_positions + .iter() + .copied() + .zip(counts.iter().copied()) + .peekable(); + let mut validity_cursor = + SparseValidityCursor::new(layer_validity, *num_slots, "list validity")?; + let mut validity = validity; + for slot in 0..*num_slots { + let is_valid = validity_cursor.is_valid(slot)?; + if let Some(validity) = validity.as_mut() { + validity.append(is_valid); + } + + if non_empty_iter + .peek() + .is_some_and(|(non_empty_pos, _)| *non_empty_pos == slot) + { + let (_, count) = non_empty_iter.next().ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list position is missing its child count".into(), + ) + })?; + if !is_valid { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list slot {} is both invalid and non-empty", + slot + ) + .into(), + )); + } + current_offset = current_offset.checked_add(count).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list offsets overflow u64".into(), + ) + })?; + } + offsets.push(Self::to_offset(current_offset)?); + } + validity_cursor.finish()?; + if let Some((extra_pos, _)) = non_empty_iter.next() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural non-empty position {} is outside layer with {} slots", + extra_pos, num_slots + ) + .into(), + )); + } + + self.consume_current_layer()?; + Ok(()) + } +} + +#[derive(Debug, Clone)] +enum SparsePositionSetDecoder { + Empty, + All { + len: u64, + }, + Range { + start: u64, + len: u64, + }, + Explicit { + decompressor: Arc, + encoding: CompressiveEncoding, + count: u64, + domain_len: u64, + }, +} + +#[derive(Debug, Clone)] +enum SparseCountSetDecoder { + Empty, + Constant { + value: u64, + len: u64, + }, + Explicit { + decompressor: Arc, + encoding: CompressiveEncoding, + count: u64, + }, +} + +#[derive(Debug, Clone)] +enum SparseLayerDecompressors { + Validity { + num_slots: u64, + validity: SparseValiditySetDecoder, + }, + List { + num_slots: u64, + num_child_slots: u64, + non_empty_positions: SparsePositionSetDecoder, + counts: SparseCountSetDecoder, + validity: SparseValiditySetDecoder, + }, + FixedSizeList { + num_slots: u64, + num_child_slots: u64, + dimension: u64, + validity: SparseValiditySetDecoder, + }, +} + +#[derive(Debug, Clone)] +struct SparseValiditySetDecoder { + meaning: SparseValidityMeaning, + positions: SparsePositionSetDecoder, +} + +#[derive(Debug)] +struct SparseStructuralCacheableState { + chunk_meta: Vec, + chunk_value_offsets: Arc<[u64]>, + plan: SparseStructuralPlan, + row_domain: u64, +} + +impl DeepSizeOf for SparseStructuralCacheableState { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + let structural_size = self + .plan + .layers + .iter() + .map(|layer| match layer { + SparseStructuralLayerPlan::Validity { validity, .. } => validity.deep_size(), + SparseStructuralLayerPlan::List { + non_empty_positions, + counts, + validity, + .. + } => non_empty_positions.deep_size() + counts.deep_size() + validity.deep_size(), + SparseStructuralLayerPlan::FixedSizeList { validity, .. } => validity.deep_size(), + }) + .sum::(); + self.chunk_meta.len() * std::mem::size_of::() + + self.chunk_value_offsets.len() * std::mem::size_of::() + + structural_size + } +} + +impl CachedPageData for SparseStructuralCacheableState { + fn as_arc_any(self: Arc) -> Arc { + self + } +} + +#[derive(Debug)] +pub(super) struct SparseStructuralScheduler { + buffer_offsets_and_sizes: Vec<(u64, u64)>, + priority: u64, + row_domain: u64, + row_scale: u64, + num_items: u64, + num_visible_items: u64, + num_buffers: u64, + value_encoding: CompressiveEncoding, + value_decompressor: Arc, + layer_decompressors: Vec, + data_type: DataType, + page_meta: Option>, + has_large_chunk: bool, +} + +impl SparseStructuralScheduler { + fn require_layer( + layer: &pb21::SparseStructuralLayer, + ) -> Result<&pb21::sparse_structural_layer::Layer> { + layer.layer.as_ref().ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural layer is missing its layer variant".into(), + ) + }) + } + + fn layer_num_slots(layer: &pb21::SparseStructuralLayer) -> Result { + Ok(match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::Validity(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::List(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => layer.num_slots, + }) + } + + fn layer_num_child_slots(layer: &pb21::SparseStructuralLayer) -> Result { + Ok(match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::Validity(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::List(layer) => layer.num_child_slots, + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => layer + .num_slots + .checked_mul(layer.dimension) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + layer.num_slots, layer.dimension + ) + .into(), + ) + })?, + }) + } + + pub(super) fn try_new( + buffer_offsets_and_sizes: &[(u64, u64)], + priority: u64, + encoded_row_domain: u64, + data_type: DataType, + layout: &pb21::SparseLayout, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + let value_compression = layout.value_compression.as_ref().ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing value compression".into()) + })?; + let value_buffer_count = Self::validate_value_encoding(value_compression)?; + if layout.num_buffers != value_buffer_count { + return Err(Error::invalid_input_source( + format!( + "Sparse layout declares {} value buffers, but its compression descriptor requires {}", + layout.num_buffers, value_buffer_count + ) + .into(), + )); + } + let row_domain = match layout.structural_layers.first() { + Some(layer) => Self::layer_num_slots(layer)?, + None => encoded_row_domain, + }; + Self::validate_domain_chain( + &layout.structural_layers, + row_domain, + layout.num_items, + layout.num_visible_items, + )?; + let expected_buffers = 2 + Self::structural_buffer_count(&layout.structural_layers)?; + if buffer_offsets_and_sizes.len() != expected_buffers { + return Err(Error::invalid_input_source( + format!( + "Sparse layout has {} buffers, expected {}", + buffer_offsets_and_sizes.len(), + expected_buffers + ) + .into(), + )); + } + Self::validate_page_buffers(buffer_offsets_and_sizes, layout.num_visible_items)?; + let row_scale = + layout + .structural_layers + .iter() + .try_fold(1_u64, |scale, layer| -> Result { + match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => { + scale.checked_mul(layer.dimension).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural fixed-size-list row scale overflows".into(), + ) + }) + } + pb21::sparse_structural_layer::Layer::Validity(_) + | pb21::sparse_structural_layer::Layer::List(_) => Ok(scale), + } + })?; + let expected_encoded_row_domain = row_domain.checked_mul(row_scale).ok_or_else(|| { + Error::invalid_input_source("Sparse structural encoded row domain overflows".into()) + })?; + if encoded_row_domain != expected_encoded_row_domain { + return Err(Error::invalid_input_source( + format!( + "Sparse structural encoded row domain {} does not match outer domain {} * fixed-size-list scale {}", + encoded_row_domain, row_domain, row_scale + ) + .into(), + )); + } + let value_decompressor = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + decompressors.create_miniblock_decompressor(value_compression, decompressors) + })) + .map_err(|_| { + Error::invalid_input_source( + "Sparse value compression descriptor caused decompressor construction to panic" + .into(), + ) + })? + .map_err(|error| { + Error::invalid_input_source( + format!("Sparse value decompressor construction failed: {error}").into(), + ) + })?; + let layer_decompressors = layout + .structural_layers + .iter() + .map(|layer| Self::layer_decompressors(layer, decompressors)) + .collect::>>()?; + + Ok(Self { + buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(), + priority, + row_domain, + row_scale, + num_items: layout.num_items, + num_visible_items: layout.num_visible_items, + num_buffers: layout.num_buffers, + value_encoding: value_compression.clone(), + value_decompressor: value_decompressor.into(), + layer_decompressors, + data_type, + page_meta: None, + has_large_chunk: layout.has_large_chunk, + }) + } + + fn validate_compression<'a>( + compression: &'a CompressiveEncoding, + label: &str, + ) -> Result<&'a CompressiveEncoding> { + compression.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} is missing compression details").into(), + ) + })?; + Ok(compression) + } + + fn validate_buffer_compression( + compression: Option<&pb21::BufferCompression>, + label: &str, + ) -> Result<()> { + let Some(compression) = compression else { + return Ok(()); + }; + match invalid_enum( + pb21::CompressionScheme::try_from(compression.scheme), + "compression scheme", + )? { + pb21::CompressionScheme::CompressionAlgorithmUnspecified => { + Err(Error::invalid_input_source( + format!("Sparse structural {label} buffer compression is unspecified").into(), + )) + } + pb21::CompressionScheme::CompressionAlgorithmLz4 + | pb21::CompressionScheme::CompressionAlgorithmZstd => Ok(()), + } + } + + fn encoding_contains_general(root: &CompressiveEncoding) -> bool { + use pb21::compressive_encoding::Compression; + + let mut stack = vec![root]; + while let Some(encoding) = stack.pop() { + let Some(compression) = encoding.compression.as_ref() else { + continue; + }; + match compression { + Compression::General(_) => return true, + Compression::Variable(variable) => { + stack.extend(variable.offsets.as_deref()); + } + Compression::OutOfLineBitpacking(bitpacking) => { + stack.extend(bitpacking.values.as_deref()); + } + Compression::Fsst(fsst) => stack.extend(fsst.values.as_deref()), + Compression::Dictionary(dictionary) => { + stack.extend(dictionary.indices.as_deref()); + stack.extend(dictionary.items.as_deref()); + } + Compression::Rle(rle) => { + stack.extend(rle.values.as_deref()); + stack.extend(rle.run_lengths.as_deref()); + } + Compression::ByteStreamSplit(split) => { + stack.extend(split.values.as_deref()); + } + Compression::FixedSizeList(fsl) => stack.extend(fsl.values.as_deref()), + Compression::PackedStruct(packed) => stack.extend(packed.values.as_deref()), + Compression::VariablePackedStruct(packed) => { + stack.extend( + packed + .fields + .iter() + .filter_map(|field| field.value.as_ref()), + ); + } + Compression::Flat(_) + | Compression::Constant(_) + | Compression::InlineBitpacking(_) => {} + } + } + false + } + + fn require_encoding<'a>( + encoding: &'a Option>, + label: &str, + ) -> Result<&'a CompressiveEncoding> { + encoding.as_deref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} encoding is required").into(), + ) + }) + } + + fn validate_flat(flat: &pb21::Flat, label: &str) -> Result<()> { + if flat.bits_per_value == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} flat bit width is zero").into(), + )); + } + if flat.data.is_some() { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} uses unsupported leaf buffer compression") + .into(), + )); + } + Ok(()) + } + + fn validate_value_encoding(encoding: &CompressiveEncoding) -> Result { + use pb21::compressive_encoding::Compression; + + let compression = encoding.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + "Sparse value compression is missing compression details".into(), + ) + })?; + match compression { + Compression::Flat(flat) => { + Self::validate_flat(flat, "value")?; + Ok(1) + } + Compression::InlineBitpacking(bitpacking) => { + if !matches!(bitpacking.uncompressed_bits_per_value, 8 | 16 | 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse inline bitpacking width {} is not supported", + bitpacking.uncompressed_bits_per_value + ) + .into(), + )); + } + if bitpacking.values.is_some() { + return Err(Error::invalid_input_source( + "Sparse inline bitpacking uses unsupported leaf buffer compression".into(), + )); + } + Ok(1) + } + Compression::Variable(variable) => { + let offsets = variable.offsets.as_deref().ok_or_else(|| { + Error::invalid_input_source( + "Sparse variable compression is missing offsets".into(), + ) + })?; + let Some(Compression::Flat(offsets)) = offsets.compression.as_ref() else { + return Err(Error::invalid_input_source( + "Sparse variable offsets must use flat compression".into(), + )); + }; + Self::validate_flat(offsets, "variable offsets")?; + if !matches!(offsets.bits_per_value, 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse variable offset width {} is not supported", + offsets.bits_per_value + ) + .into(), + )); + } + if variable.values.is_some() { + return Err(Error::invalid_input_source( + "Sparse variable values use unsupported leaf buffer compression".into(), + )); + } + Ok(1) + } + Compression::Fsst(fsst) => { + if fsst.symbol_table.is_empty() { + return Err(Error::invalid_input_source( + "Sparse FSST compression has an empty symbol table".into(), + )); + } + let values = Self::require_encoding(&fsst.values, "FSST values")?; + if !matches!(values.compression.as_ref(), Some(Compression::Variable(_))) { + return Err(Error::invalid_input_source( + "Sparse FSST values must use variable compression".into(), + )); + } + Self::validate_value_encoding(values) + } + Compression::ByteStreamSplit(split) => { + let values = Self::require_encoding(&split.values, "byte-stream-split values")?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input_source( + "Sparse byte-stream-split values must use flat compression".into(), + )); + }; + Self::validate_flat(flat, "byte-stream-split values")?; + if !matches!(flat.bits_per_value, 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse byte-stream-split width {} is not supported", + flat.bits_per_value + ) + .into(), + )); + } + Ok(1) + } + Compression::FixedSizeList(fsl) => Self::validate_fsl_value_encoding(fsl), + Compression::PackedStruct(packed) => Self::validate_packed_value_encoding(packed), + Compression::Rle(rle) => { + let values = Self::require_encoding(&rle.values, "RLE values")?; + let lengths = Self::require_encoding(&rle.run_lengths, "RLE run lengths")?; + Self::validate_block_encoding(values, "RLE values")?; + Self::validate_block_encoding(lengths, "RLE run lengths")?; + Ok(2) + } + Compression::General(general) => { + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + "Sparse general compression is missing its buffer compression".into(), + ) + })?; + Self::validate_buffer_compression(Some(compression), "general")?; + Self::validate_value_encoding(Self::require_encoding( + &general.values, + "general values", + )?) + } + Compression::Constant(_) + | Compression::OutOfLineBitpacking(_) + | Compression::Dictionary(_) + | Compression::VariablePackedStruct(_) => Err(Error::invalid_input_source( + "Sparse value compression uses an unsupported mini-block encoding".into(), + )), + } + } + + fn validate_fsl_value_encoding(fsl: &pb21::FixedSizeList) -> Result { + use pb21::compressive_encoding::Compression; + + if fsl.items_per_value == 0 { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list value compression has dimension zero".into(), + )); + } + let values = Self::require_encoding(&fsl.values, "fixed-size-list values")?; + let child_buffers = match values.compression.as_ref() { + Some(Compression::Flat(flat)) => { + Self::validate_flat(flat, "fixed-size-list values")?; + 1_u64 + } + Some(Compression::FixedSizeList(inner)) => Self::validate_fsl_value_encoding(inner)?, + _ => { + return Err(Error::invalid_input_source( + "Sparse fixed-size-list values must use fixed-size-list or flat compression" + .into(), + )); + } + }; + child_buffers + .checked_add(u64::from(fsl.has_validity)) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse fixed-size-list value buffer count overflows".into(), + ) + }) + } + + fn validate_packed_value_encoding(packed: &pb21::PackedStruct) -> Result { + use pb21::compressive_encoding::Compression; + + if packed.bits_per_value.is_empty() + || packed + .bits_per_value + .iter() + .any(|bits| *bits == 0 || !bits.is_multiple_of(8)) + { + return Err(Error::invalid_input_source( + "Sparse packed-struct widths must be non-empty positive byte widths".into(), + )); + } + let values = Self::require_encoding(&packed.values, "packed-struct values")?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input_source( + "Sparse packed-struct values must use flat compression".into(), + )); + }; + Self::validate_flat(flat, "packed-struct values")?; + let total_bits = packed.bits_per_value.iter().try_fold(0_u64, |sum, bits| { + sum.checked_add(*bits).ok_or_else(|| { + Error::invalid_input_source("Sparse packed-struct bit width sum overflows".into()) + }) + })?; + if total_bits != flat.bits_per_value { + return Err(Error::invalid_input_source( + format!( + "Sparse packed-struct child widths sum to {}, but values use {} bits", + total_bits, flat.bits_per_value + ) + .into(), + )); + } + Ok(1) + } + + fn validate_block_encoding(encoding: &CompressiveEncoding, label: &str) -> Result<()> { + use pb21::compressive_encoding::Compression; + + match encoding.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} is missing compression details").into(), + ) + })? { + Compression::Flat(flat) => Self::validate_flat(flat, label), + Compression::InlineBitpacking(bitpacking) => { + if !matches!(bitpacking.uncompressed_bits_per_value, 8 | 16 | 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} inline bitpacking width {} is unsupported", + bitpacking.uncompressed_bits_per_value + ) + .into(), + )); + } + if bitpacking.values.is_some() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} uses unsupported leaf buffer compression" + ) + .into(), + )); + } + Ok(()) + } + Compression::OutOfLineBitpacking(bitpacking) => { + if !matches!(bitpacking.uncompressed_bits_per_value, 8 | 16 | 32 | 64) { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} out-of-line bitpacking width {} is unsupported", + bitpacking.uncompressed_bits_per_value + ) + .into(), + )); + } + let values = Self::require_encoding(&bitpacking.values, label)?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} bitpacked values must be flat").into(), + )); + }; + Self::validate_flat(flat, label) + } + Compression::Constant(constant) => { + if constant + .value + .as_ref() + .is_some_and(|value| value.len() != 8) + { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} constant must be 64 bits").into(), + )); + } + Ok(()) + } + Compression::General(general) => { + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} general compression is missing config") + .into(), + ) + })?; + Self::validate_buffer_compression(Some(compression), label)?; + Self::validate_block_encoding( + Self::require_encoding(&general.values, label)?, + label, + ) + } + Compression::Rle(rle) => { + Self::validate_block_encoding( + Self::require_encoding(&rle.values, "RLE values")?, + "RLE values", + )?; + Self::validate_block_encoding( + Self::require_encoding(&rle.run_lengths, "RLE run lengths")?, + "RLE run lengths", + ) + } + _ => Err(Error::invalid_input_source( + format!("Sparse structural {label} uses an unsupported block encoding").into(), + )), + } + } + + fn validate_domain_chain( + layers: &[pb21::SparseStructuralLayer], + row_domain: u64, + num_items: u64, + num_visible_items: u64, + ) -> Result<()> { + let expected_num_items = + layers + .iter() + .try_fold(num_visible_items, |items, layer| -> Result { + let additional = match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::List(layer) => { + let positions = Self::require_position_set( + &layer.non_empty_positions, + "list non-empty", + )?; + let num_non_empty = Self::position_cardinality( + positions, + layer.num_slots, + "list non-empty positions", + )?; + layer.num_slots.checked_sub(num_non_empty).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list has more non-empty positions than slots" + .into(), + ) + })? + } + pb21::sparse_structural_layer::Layer::Validity(_) + | pb21::sparse_structural_layer::Layer::FixedSizeList(_) => 0, + }; + items.checked_add(additional).ok_or_else(|| { + Error::invalid_input_source("Sparse structural item count overflows".into()) + }) + })?; + if num_items != expected_num_items { + return Err(Error::invalid_input_source( + format!( + "Sparse layout has {} structural items, but its layers imply {}", + num_items, expected_num_items + ) + .into(), + )); + } + let mut expected_slots = row_domain; + for (layer_index, layer) in layers.iter().enumerate() { + let num_slots = Self::layer_num_slots(layer)?; + let num_child_slots = Self::layer_num_child_slots(layer)?; + usize_from_u64(num_slots, "layer slot count")?; + usize_from_u64(num_child_slots, "layer child slot count")?; + if num_slots != expected_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural layer {} has {} slots, expected {} from the outer domain", + layer_index, num_slots, expected_slots + ) + .into(), + )); + } + expected_slots = num_child_slots; + } + if expected_slots != num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse structural terminal domain has {} slots, expected {} visible items", + expected_slots, num_visible_items + ) + .into(), + )); + } + Ok(()) + } + + fn metadata_buffer(&self) -> Result<(u64, u64)> { + self.buffer_offsets_and_sizes + .first() + .copied() + .ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing metadata buffer".into()) + }) + } + + fn value_buffer(&self) -> Result<(u64, u64)> { + self.buffer_offsets_and_sizes + .get(1) + .copied() + .ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing value buffer".into()) + }) + } + + fn checked_buffer_range(position: u64, size: u64, label: &str) -> Result> { + let end = position.checked_add(size).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} buffer range overflows").into(), + ) + })?; + Ok(position..end) + } + + fn validate_page_buffers( + buffer_offsets_and_sizes: &[(u64, u64)], + num_visible_items: u64, + ) -> Result<()> { + let (_, metadata_size) = buffer_offsets_and_sizes.first().copied().ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing metadata buffer".into()) + })?; + if !metadata_size.is_multiple_of(8) { + return Err(Error::invalid_input_source( + format!( + "Sparse layout metadata buffer has {metadata_size} bytes, which is not a multiple of 8" + ) + .into(), + )); + } + let chunk_count = metadata_size / 8; + if num_visible_items == 0 { + if chunk_count != 0 { + return Err(Error::invalid_input_source( + "Sparse layout with no visible items must have empty chunk metadata".into(), + )); + } + } else { + let min_chunks = + num_visible_items.div_ceil(miniblock::MAX_CONFIGURABLE_MINIBLOCK_VALUES); + if chunk_count < min_chunks || chunk_count > num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse layout metadata declares {chunk_count} chunks for {num_visible_items} visible items, expected {min_chunks}..={num_visible_items}" + ) + .into(), + )); + } + } + + let (_, value_size) = buffer_offsets_and_sizes.get(1).copied().ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing value buffer".into()) + })?; + if (num_visible_items == 0) != (value_size == 0) { + return Err(Error::invalid_input_source( + format!( + "Sparse layout value buffer has {value_size} bytes for {num_visible_items} visible items" + ) + .into(), + )); + } + + for (position, size) in buffer_offsets_and_sizes.iter().skip(2) { + Self::checked_buffer_range(*position, *size, "structural")?; + } + + for (position, size) in buffer_offsets_and_sizes.iter().take(2) { + Self::checked_buffer_range(*position, *size, "page")?; + } + Ok(()) + } + + fn require_position_set<'a>( + set: &'a Option, + label: &str, + ) -> Result<&'a pb21::SparsePositionSet> { + set.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} position set is required").into(), + ) + }) + } + + fn require_count_set<'a>( + set: &'a Option, + label: &str, + ) -> Result<&'a pb21::SparseCountSet> { + set.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} count set is required").into(), + ) + }) + } + + fn require_validity_set<'a>( + set: &'a Option, + label: &str, + ) -> Result<&'a pb21::SparseValiditySet> { + set.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} validity set is required").into(), + ) + }) + } + + fn validity_meaning( + validity_set: &pb21::SparseValiditySet, + label: &str, + ) -> Result { + match invalid_enum( + pb21::sparse_validity_set::Meaning::try_from(validity_set.meaning), + "validity meaning", + )? { + pb21::sparse_validity_set::Meaning::SparseValidityUnspecified => { + Err(Error::invalid_input_source( + format!("Sparse structural {label} meaning is unspecified").into(), + )) + } + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions => { + Ok(SparseValidityMeaning::NullPositions) + } + pb21::sparse_validity_set::Meaning::SparseValidityValidPositions => { + Ok(SparseValidityMeaning::ValidPositions) + } + } + } + + fn validity_buffer_count( + validity_set: &pb21::SparseValiditySet, + domain_len: u64, + label: &str, + ) -> Result<(usize, SparseValidityMeaning, u64)> { + let meaning = Self::validity_meaning(validity_set, label)?; + let position_set = validity_set.positions.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} positions are required").into(), + ) + })?; + let cardinality = Self::position_cardinality(position_set, domain_len, label)?; + let buffer_count = Self::position_buffer_count(position_set, domain_len, label)?; + Ok((buffer_count, meaning, cardinality)) + } + + fn position_cardinality( + position_set: &pb21::SparsePositionSet, + domain_len: u64, + label: &str, + ) -> Result { + let positions = position_set.positions.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} position set is missing its variant").into(), + ) + })?; + let cardinality = position_set.num_positions; + if cardinality > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} cardinality {} exceeds domain {}", + cardinality, domain_len + ) + .into(), + )); + } + match positions { + pb21::sparse_position_set::Positions::Empty(_) => { + if cardinality != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} empty set has cardinality {}", + cardinality + ) + .into(), + )); + } + } + pb21::sparse_position_set::Positions::All(_) => { + if domain_len == 0 || cardinality != domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} all set has cardinality {}, expected {}", + cardinality, domain_len + ) + .into(), + )); + } + } + pb21::sparse_position_set::Positions::Range(range) => { + let end = range.start.checked_add(range.length).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} range overflows").into(), + ) + })?; + if range.length == 0 || range.length != cardinality || end > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} range {}..{} does not match cardinality {} in domain {}", + range.start, end, cardinality, domain_len + ) + .into(), + )); + } + } + pb21::sparse_position_set::Positions::Explicit(compression) => { + if cardinality == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} has compression but no values").into(), + )); + } + Self::validate_compression(compression, label)?; + } + } + Ok(cardinality) + } + + fn position_buffer_count( + position_set: &pb21::SparsePositionSet, + domain_len: u64, + label: &str, + ) -> Result { + Self::position_cardinality(position_set, domain_len, label)?; + let positions = position_set.positions.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} position set is missing its variant").into(), + ) + })?; + Ok(usize::from(matches!( + positions, + pb21::sparse_position_set::Positions::Explicit(_) + ))) + } + + fn count_buffer_count( + count_set: &pb21::SparseCountSet, + cardinality: u64, + label: &str, + ) -> Result { + let counts = count_set.counts.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} count set is missing its variant").into(), + ) + })?; + match counts { + pb21::sparse_count_set::Counts::Empty(_) => { + if cardinality != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} empty count set has cardinality {}", + cardinality + ) + .into(), + )); + } + Ok(0) + } + pb21::sparse_count_set::Counts::Constant(constant) => { + if cardinality == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} constant count has no values").into(), + )); + } + if constant.value == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} constant count is zero").into(), + )); + } + Ok(0) + } + pb21::sparse_count_set::Counts::Explicit(compression) => { + if cardinality == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} has compression but no values").into(), + )); + } + Self::validate_compression(compression, label)?; + Ok(1) + } + } + } + + fn count_set_child_slots( + count_set: &pb21::SparseCountSet, + cardinality: u64, + label: &str, + ) -> Result> { + let counts = count_set.counts.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} count set is missing its variant").into(), + ) + })?; + match counts { + pb21::sparse_count_set::Counts::Empty(_) => Ok(Some(0)), + pb21::sparse_count_set::Counts::Constant(constant) => constant + .value + .checked_mul(cardinality) + .map(Some) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} constant count sum overflows: value={}, len={}", + constant.value, cardinality + ) + .into(), + ) + }), + pb21::sparse_count_set::Counts::Explicit(_) => Ok(None), + } + } + + fn create_position_decompressor( + compression: &CompressiveEncoding, + label: &str, + decompressors: &dyn DecompressionStrategy, + ) -> Result> { + let compression = Self::validate_compression(compression, label)?; + Self::validate_block_encoding(compression, label)?; + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + decompressors.create_block_decompressor(compression) + })) + .map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural {label} descriptor caused decompressor construction to panic" + ) + .into(), + ) + })? + .map(Arc::from) + .map_err(|error| { + Error::invalid_input_source( + format!("Sparse structural {label} decompressor construction failed: {error}") + .into(), + ) + }) + } + + fn position_set_decoder( + position_set: &pb21::SparsePositionSet, + domain_len: u64, + label: &str, + decompressors: &dyn DecompressionStrategy, + ) -> Result<(SparsePositionSetDecoder, u64)> { + let cardinality = Self::position_cardinality(position_set, domain_len, label)?; + let positions = position_set.positions.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} position set is missing its variant").into(), + ) + })?; + Ok(( + match positions { + pb21::sparse_position_set::Positions::Empty(_) => SparsePositionSetDecoder::Empty, + pb21::sparse_position_set::Positions::All(_) => { + SparsePositionSetDecoder::All { len: domain_len } + } + pb21::sparse_position_set::Positions::Range(range) => { + SparsePositionSetDecoder::Range { + start: range.start, + len: range.length, + } + } + pb21::sparse_position_set::Positions::Explicit(compression) => { + SparsePositionSetDecoder::Explicit { + decompressor: Self::create_position_decompressor( + compression, + label, + decompressors, + )?, + encoding: compression.clone(), + count: cardinality, + domain_len, + } + } + }, + cardinality, + )) + } + + fn validity_set_decoder( + validity_set: &pb21::SparseValiditySet, + domain_len: u64, + label: &str, + decompressors: &dyn DecompressionStrategy, + ) -> Result<(SparseValiditySetDecoder, u64)> { + let meaning = Self::validity_meaning(validity_set, label)?; + let position_set = validity_set.positions.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} positions are required").into(), + ) + })?; + let (positions, cardinality) = + Self::position_set_decoder(position_set, domain_len, label, decompressors)?; + Ok((SparseValiditySetDecoder { meaning, positions }, cardinality)) + } + + fn num_valid_slots( + meaning: SparseValidityMeaning, + cardinality: u64, + num_slots: u64, + label: &str, + ) -> Result { + match meaning { + SparseValidityMeaning::NullPositions => { + num_slots.checked_sub(cardinality).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} null cardinality {} exceeds slots {}", + cardinality, num_slots + ) + .into(), + ) + }) + } + SparseValidityMeaning::ValidPositions => Ok(cardinality), + } + } + + fn count_set_decoder( + count_set: &pb21::SparseCountSet, + cardinality: u64, + label: &str, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + Self::count_buffer_count(count_set, cardinality, label)?; + let counts = count_set.counts.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} count set is missing its variant").into(), + ) + })?; + Ok(match counts { + pb21::sparse_count_set::Counts::Empty(_) => SparseCountSetDecoder::Empty, + pb21::sparse_count_set::Counts::Constant(constant) => SparseCountSetDecoder::Constant { + value: constant.value, + len: cardinality, + }, + pb21::sparse_count_set::Counts::Explicit(compression) => { + SparseCountSetDecoder::Explicit { + decompressor: Self::create_position_decompressor( + compression, + label, + decompressors, + )?, + encoding: compression.clone(), + count: cardinality, + } + } + }) + } + + fn add_buffer_count(count: &mut usize, additional: usize) -> Result<()> { + *count = count.checked_add(additional).ok_or_else(|| { + Error::invalid_input_source("Sparse structural buffer count overflows".into()) + })?; + Ok(()) + } + + fn structural_buffer_count(layers: &[pb21::SparseStructuralLayer]) -> Result { + layers + .iter() + .try_fold(0_usize, |mut count, layer| -> Result { + match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::Validity(layer) => { + let (validity_buffers, _, _) = Self::validity_buffer_count( + Self::require_validity_set(&layer.validity, "validity")?, + layer.num_slots, + "validity positions", + )?; + Self::add_buffer_count(&mut count, validity_buffers)?; + } + pb21::sparse_structural_layer::Layer::List(layer) => { + let non_empty_positions = Self::require_position_set( + &layer.non_empty_positions, + "list non-empty", + )?; + let num_non_empty = Self::position_cardinality( + non_empty_positions, + layer.num_slots, + "list non-empty positions", + )?; + let non_empty_buffers = Self::position_buffer_count( + non_empty_positions, + layer.num_slots, + "list non-empty positions", + )?; + Self::add_buffer_count(&mut count, non_empty_buffers)?; + let (validity_buffers, validity_meaning, validity_cardinality) = + Self::validity_buffer_count( + Self::require_validity_set(&layer.validity, "list")?, + layer.num_slots, + "list validity positions", + )?; + Self::add_buffer_count(&mut count, validity_buffers)?; + let num_valid_slots = Self::num_valid_slots( + validity_meaning, + validity_cardinality, + layer.num_slots, + "list validity", + )?; + if num_non_empty > num_valid_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list has {} non-empty slots but only {} valid slots", + num_non_empty, num_valid_slots + ) + .into(), + )); + } + let counts = Self::require_count_set(&layer.counts, "list counts")?; + let count_buffers = + Self::count_buffer_count(counts, num_non_empty, "list counts")?; + Self::add_buffer_count(&mut count, count_buffers)?; + if let Some(child_slots) = + Self::count_set_child_slots(counts, num_non_empty, "list counts")? + && child_slots != layer.num_child_slots + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match child slots {}", + child_slots, layer.num_child_slots + ) + .into(), + )); + } + } + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => { + if layer.dimension == 0 { + return Err(Error::invalid_input_source( + "Sparse structural fixed-size-list dimension is zero".into(), + )); + } + layer.num_slots.checked_mul(layer.dimension).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + layer.num_slots, layer.dimension + ) + .into(), + ) + })?; + let (validity_buffers, _, _) = Self::validity_buffer_count( + Self::require_validity_set(&layer.validity, "fixed-size-list")?, + layer.num_slots, + "fixed-size-list validity positions", + )?; + Self::add_buffer_count(&mut count, validity_buffers)?; + } + } + Ok(count) + }) + } + + fn layer_decompressors( + layer: &pb21::SparseStructuralLayer, + decompressors: &dyn DecompressionStrategy, + ) -> Result { + Ok(match Self::require_layer(layer)? { + pb21::sparse_structural_layer::Layer::Validity(layer) => { + let (validity, _) = Self::validity_set_decoder( + Self::require_validity_set(&layer.validity, "validity")?, + layer.num_slots, + "validity positions", + decompressors, + )?; + SparseLayerDecompressors::Validity { + num_slots: layer.num_slots, + validity, + } + } + pb21::sparse_structural_layer::Layer::List(layer) => { + let non_empty_positions = + Self::require_position_set(&layer.non_empty_positions, "list non-empty")?; + let (non_empty_positions, num_non_empty) = Self::position_set_decoder( + non_empty_positions, + layer.num_slots, + "list non-empty positions", + decompressors, + )?; + let counts = Self::count_set_decoder( + Self::require_count_set(&layer.counts, "list counts")?, + num_non_empty, + "list counts", + decompressors, + )?; + let (validity, _) = Self::validity_set_decoder( + Self::require_validity_set(&layer.validity, "list")?, + layer.num_slots, + "list validity positions", + decompressors, + )?; + SparseLayerDecompressors::List { + num_slots: layer.num_slots, + num_child_slots: layer.num_child_slots, + non_empty_positions, + counts, + validity, + } + } + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => { + let (validity, _) = Self::validity_set_decoder( + Self::require_validity_set(&layer.validity, "fixed-size-list")?, + layer.num_slots, + "fixed-size-list validity positions", + decompressors, + )?; + SparseLayerDecompressors::FixedSizeList { + num_slots: layer.num_slots, + num_child_slots: layer.num_slots.checked_mul(layer.dimension).ok_or_else( + || { + Error::invalid_input_source( + "Sparse structural fixed-size-list child slot count overflows" + .into(), + ) + }, + )?, + dimension: layer.dimension, + validity, + } + } + }) + } + + fn parse_chunk_meta(&self, meta_bytes: Bytes) -> Result> { + if !meta_bytes.len().is_multiple_of(8) { + return Err(Error::invalid_input_source( + format!( + "Sparse layout metadata length {} is not a multiple of 8", + meta_bytes.len() + ) + .into(), + )); + } + + let (value_buf_position, value_buf_size) = self.value_buffer()?; + let value_buf_end = value_buf_position + .checked_add(value_buf_size) + .ok_or_else(|| { + Error::invalid_input_source("Sparse layout value buffer range overflows".into()) + })?; + let mut rows_counter = 0_u64; + let mut offset_bytes = value_buf_position; + let mut chunk_meta = Vec::with_capacity(meta_bytes.len() / 8); + for chunk in meta_bytes.chunks_exact(8) { + let entry: [u8; 8] = chunk.try_into().map_err(|_| { + Error::invalid_input_source( + "Sparse layout chunk metadata entry is not 8 bytes".into(), + ) + })?; + let divided_bytes_minus_one = u32::from_le_bytes( + entry + .get(..4) + .and_then(|bytes| bytes.try_into().ok()) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse layout chunk byte-size field is malformed".into(), + ) + })?, + ); + let num_values = u64::from(u32::from_le_bytes( + entry + .get(4..) + .and_then(|bytes| bytes.try_into().ok()) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse layout chunk value-count field is malformed".into(), + ) + })?, + )); + if num_values == 0 { + return Err(Error::invalid_input_source( + "Sparse layout contains an empty value chunk".into(), + )); + } + if num_values > miniblock::MAX_CONFIGURABLE_MINIBLOCK_VALUES { + return Err(Error::invalid_input_source( + format!( + "Sparse layout value chunk has {} values, exceeding the mini-block limit {}", + num_values, + miniblock::MAX_CONFIGURABLE_MINIBLOCK_VALUES + ) + .into(), + )); + } + let num_bytes = u64::from(divided_bytes_minus_one) + .checked_add(1) + .and_then(|units| units.checked_mul(MINIBLOCK_ALIGNMENT as u64)) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse layout value chunk byte size overflows".into(), + ) + })?; + rows_counter = rows_counter.checked_add(num_values).ok_or_else(|| { + Error::invalid_input_source("Sparse layout visible item count overflows".into()) + })?; + chunk_meta.push(ChunkMeta { + num_values, + chunk_size_bytes: num_bytes, + offset_bytes, + }); + offset_bytes = offset_bytes.checked_add(num_bytes).ok_or_else(|| { + Error::invalid_input_source("Sparse layout value chunk byte range overflows".into()) + })?; + } + if rows_counter != self.num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse layout visible item count mismatch: metadata has {}, layout has {}", + rows_counter, self.num_visible_items + ) + .into(), + )); + } + if offset_bytes != value_buf_end { + return Err(Error::invalid_input_source( + format!( + "Sparse layout chunk metadata describes {} value bytes, but value buffer has {} bytes", + offset_bytes - value_buf_position, + value_buf_size + ) + .into(), + )); + } + Ok(chunk_meta) + } + + fn validate_general_buffer_header( + general: &pb21::General, + data: &[u8], + label: &str, + ) -> Result<()> { + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} general compression is missing config").into(), + ) + })?; + Self::validate_buffer_compression(Some(compression), label)?; + let values = Self::require_encoding(&general.values, label)?; + if Self::encoding_contains_general(values) { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} contains nested general compression, which is unsupported" + ) + .into(), + )); + } + + let scheme = invalid_enum( + pb21::CompressionScheme::try_from(compression.scheme), + "compression scheme", + )?; + match scheme { + pb21::CompressionScheme::CompressionAlgorithmLz4 => { + data.get(..4).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} LZ4 buffer is missing its length prefix" + ) + .into(), + ) + })?; + } + pb21::CompressionScheme::CompressionAlgorithmZstd => { + data.get(..8).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} Zstd buffer is missing its length prefix" + ) + .into(), + ) + })?; + } + pb21::CompressionScheme::CompressionAlgorithmUnspecified => { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} general compression scheme is unspecified") + .into(), + )); + } + } + Ok(()) + } + + fn validate_general_child_buffer( + encoding: &CompressiveEncoding, + data: &[u8], + label: &str, + ) -> Result<()> { + if let Some(pb21::compressive_encoding::Compression::General(general)) = + encoding.compression.as_ref() + { + Self::validate_general_buffer_header(general, data, label)?; + } + Ok(()) + } + + fn validate_structural_buffer_headers( + encoding: &CompressiveEncoding, + data: &[u8], + label: &str, + ) -> Result<()> { + use pb21::compressive_encoding::Compression; + + match encoding.compression.as_ref() { + Some(Compression::General(general)) => { + Self::validate_general_buffer_header(general, data, label) + } + Some(Compression::Rle(rle)) => { + let values_size = u64::from_le_bytes( + data.get(..8) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} RLE buffer is missing its header" + ) + .into(), + ) + })? + .try_into() + .map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} RLE header is malformed").into(), + ) + })?, + ); + let values_size = usize_from_u64(values_size, "RLE values buffer size")?; + let values_end = 8_usize.checked_add(values_size).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} RLE values range overflows").into(), + ) + })?; + let values_data = data.get(8..values_end).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} RLE values buffer is truncated").into(), + ) + })?; + let lengths_data = data.get(values_end..).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} RLE run-length buffer is missing") + .into(), + ) + })?; + Self::validate_general_child_buffer( + Self::require_encoding(&rle.values, "RLE values")?, + values_data, + "RLE values", + )?; + Self::validate_general_child_buffer( + Self::require_encoding(&rle.run_lengths, "RLE run lengths")?, + lengths_data, + "RLE run lengths", + ) + } + _ => Ok(()), + } + } + + fn decode_u64_values( + decompressor: &dyn BlockDecompressor, + encoding: &CompressiveEncoding, + data: Bytes, + num_values: u64, + label: &str, + ) -> Result> { + Self::validate_structural_buffer_headers(encoding, &data, label)?; + let decoded = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + decompressor.decompress(Some(LanceBuffer::from_bytes(data, 1)), num_values) + })) + .map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} decompression panicked").into(), + ) + })? + .map_err(|error| { + Error::invalid_input_source( + format!("Sparse structural {label} decompression failed: {error}").into(), + ) + })?; + let fixed = decoded.as_fixed_width().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} did not decode to fixed width data").into(), + ) + })?; + if fixed.bits_per_value != 64 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} decoded to {} bits per value, expected 64", + fixed.bits_per_value + ) + .into(), + )); + } + if fixed.num_values != num_values { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} decoded {} values, expected {}", + fixed.num_values, num_values + ) + .into(), + )); + } + let num_values_usize = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} value count exceeds usize::MAX").into(), + ) + })?; + let expected_len = num_values_usize + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} decoded byte length overflows").into(), + ) + })?; + if fixed.data.len() != expected_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} decoded {} bytes, expected {}", + fixed.data.len(), + expected_len + ) + .into(), + )); + } + let values = fixed.data.borrow_to_typed_slice::(); + if values.len() != num_values_usize { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} decoded {} u64 values, expected {}", + values.len(), + num_values + ) + .into(), + )); + } + Ok(values.to_vec()) + } + + fn decode_explicit_positions( + decompressor: &Arc, + encoding: &CompressiveEncoding, + data: Bytes, + num_positions: u64, + num_slots: u64, + label: &str, + ) -> Result { + let deltas = + Self::decode_u64_values(decompressor.as_ref(), encoding, data, num_positions, label)?; + let mut positions = Vec::with_capacity(deltas.len()); + let mut current = 0_u64; + for (idx, delta) in deltas.into_iter().enumerate() { + if idx == 0 { + current = delta; + } else { + if delta == 0 { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} positions must be strictly increasing") + .into(), + )); + } + current = current.checked_add(delta).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} position overflow").into(), + ) + })?; + } + if current >= num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} position {} is outside layer with {} slots", + current, num_slots + ) + .into(), + )); + } + positions.push(current); + } + SparsePositionSet::from_positions(positions, num_slots, label) + } + + fn decode_position_set( + decoder: &SparsePositionSetDecoder, + buffers: &mut impl Iterator, + label: &str, + ) -> Result { + match decoder { + SparsePositionSetDecoder::Empty => Ok(SparsePositionSet::empty()), + SparsePositionSetDecoder::All { len } => Ok(SparsePositionSet::all(*len)), + SparsePositionSetDecoder::Range { start, len } => { + Ok(SparsePositionSet::range(*start, *len)) + } + SparsePositionSetDecoder::Explicit { + decompressor, + encoding, + count, + domain_len, + } => Self::decode_explicit_positions( + decompressor, + encoding, + Self::next_structural_buffer(buffers, label)?, + *count, + *domain_len, + label, + ), + } + } + + fn decode_validity_set( + decoder: &SparseValiditySetDecoder, + buffers: &mut impl Iterator, + label: &str, + ) -> Result { + Ok(SparseValiditySet { + meaning: decoder.meaning, + positions: Self::decode_position_set(&decoder.positions, buffers, label)?, + }) + } + + fn decode_count_set( + decoder: &SparseCountSetDecoder, + buffers: &mut impl Iterator, + label: &str, + ) -> Result { + match decoder { + SparseCountSetDecoder::Empty => Ok(SparseCountSet::Empty), + SparseCountSetDecoder::Constant { value, len } => { + Ok(SparseCountSet::constant(*value, *len)) + } + SparseCountSetDecoder::Explicit { + decompressor, + encoding, + count, + } => { + let counts = Self::decode_u64_values( + decompressor.as_ref(), + encoding, + Self::next_structural_buffer(buffers, label)?, + *count, + label, + )?; + SparseCountSet::from_counts(counts) + } + } + } + + fn next_structural_buffer( + buffers: &mut impl Iterator, + label: &str, + ) -> Result { + buffers.next().ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} is missing its buffer").into(), + ) + }) + } + + fn decode_layer( + layer: &SparseLayerDecompressors, + buffers: &mut impl Iterator, + ) -> Result { + Ok(match layer { + SparseLayerDecompressors::Validity { + num_slots, + validity, + } => SparseStructuralLayerPlan::Validity { + num_slots: *num_slots, + validity: Self::decode_validity_set(validity, buffers, "validity positions")?, + }, + SparseLayerDecompressors::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + } => { + let non_empty_positions = Self::decode_position_set( + non_empty_positions, + buffers, + "list non-empty positions", + )?; + let counts = Self::decode_count_set(counts, buffers, "list counts")?; + let validity = + Self::decode_validity_set(validity, buffers, "list validity positions")?; + let actual_child_slots = counts.sum()?; + if actual_child_slots != *num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match declared child slots {}", + actual_child_slots, num_child_slots + ) + .into(), + )); + } + SparseStructuralLayerPlan::List { + num_slots: *num_slots, + num_child_slots: *num_child_slots, + non_empty_positions, + counts, + validity, + } + } + SparseLayerDecompressors::FixedSizeList { + num_slots, + num_child_slots, + dimension, + validity, + } => { + let expected_child_slots = num_slots.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + num_slots, dimension + ) + .into(), + ) + })?; + if expected_child_slots != *num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count {} does not match slots {} * dimension {}", + num_child_slots, num_slots, dimension + ) + .into(), + )); + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots: *num_slots, + dimension: *dimension, + validity: Self::decode_validity_set( + validity, + buffers, + "fixed-size-list validity positions", + )?, + } + } + }) + } + + fn lookup_value_chunks(&self, chunk_indices: &[usize]) -> Result> { + let page_meta = self.page_meta.as_ref().ok_or_else(|| { + Error::internal("Sparse page scheduler has not been initialized".to_string()) + })?; + chunk_indices + .iter() + .map(|&chunk_idx| { + let chunk_meta = page_meta.chunk_meta.get(chunk_idx).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse layout missing value chunk metadata for chunk {chunk_idx}") + .into(), + ) + })?; + let bytes_start = chunk_meta.offset_bytes; + let bytes_end = bytes_start + .checked_add(chunk_meta.chunk_size_bytes) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse layout value chunk {} byte range overflows", + chunk_idx + ) + .into(), + ) + })?; + Ok(LoadedChunk { + byte_range: bytes_start..bytes_end, + items_in_chunk: chunk_meta.num_values, + chunk_idx, + data: LanceBuffer::empty(), + }) + }) + .collect() + } + + fn value_chunk_index(chunk_value_offsets: &[u64], value: u64) -> Result { + let total_values = chunk_value_offsets.last().copied().ok_or_else(|| { + Error::invalid_input_source("Sparse layout has no value chunk offsets".into()) + })?; + if chunk_value_offsets.len() < 2 || value >= total_values { + return Err(Error::invalid_input_source( + format!( + "Sparse layout value index {} is outside {} visible items", + value, total_values + ) + .into(), + )); + } + chunk_value_offsets + .partition_point(|&offset| offset <= value) + .checked_sub(1) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse layout value index {value} is before the first chunk").into(), + ) + }) + } + + fn value_chunk_range( + chunk_value_offsets: &[u64], + value_range: Range, + ) -> Result> { + if value_range.is_empty() { + return Ok(0..0); + } + let total_values = chunk_value_offsets.last().copied().ok_or_else(|| { + Error::invalid_input_source("Sparse layout has no value chunk offsets".into()) + })?; + if value_range.start > value_range.end || value_range.end > total_values { + return Err(Error::invalid_input_source( + format!( + "Sparse layout value range {}..{} is outside {} visible items", + value_range.start, value_range.end, total_values + ) + .into(), + )); + } + let start = Self::value_chunk_index(chunk_value_offsets, value_range.start)?; + let end = chunk_value_offsets + .partition_point(|&offset| offset < value_range.end) + .max(start + 1); + Ok(start..end) + } +} + +impl StructuralPageScheduler for SparseStructuralScheduler { + fn init_layout(&self) -> Result { + let (meta_buf_position, meta_buf_size) = self.metadata_buffer()?; + let mut buffer_ranges = Vec::with_capacity(self.buffer_offsets_and_sizes.len() - 1); + buffer_ranges.push(Some(Self::checked_buffer_range( + meta_buf_position, + meta_buf_size, + "metadata", + )?)); + for (position, size) in self.buffer_offsets_and_sizes.iter().skip(2) { + buffer_ranges.push(Some(Self::checked_buffer_range( + *position, + *size, + "structural", + )?)); + } + Ok(PageInitialization::new(buffer_ranges)) + } + + fn init_from_buffers<'a>( + &'a mut self, + buffers: PageInitializationBuffers, + _io: &Arc, + ) -> BoxFuture<'a, Result>> { + async move { + let mut buffers = buffers.into_required("sparse")?.into_iter(); + let meta_bytes = buffers.next().ok_or_else(|| { + Error::invalid_input_source("Sparse layout is missing chunk metadata buffer".into()) + })?; + + let chunk_meta = self.parse_chunk_meta(meta_bytes)?; + let mut chunk_value_offsets = Vec::with_capacity(chunk_meta.len() + 1); + let mut value_offset = 0_u64; + chunk_value_offsets.push(value_offset); + for chunk in &chunk_meta { + value_offset = value_offset.checked_add(chunk.num_values).ok_or_else(|| { + Error::invalid_input_source( + "Sparse layout visible item offset overflows".into(), + ) + })?; + chunk_value_offsets.push(value_offset); + } + + let layers = self + .layer_decompressors + .iter() + .map(|layer| Self::decode_layer(layer, &mut buffers)) + .collect::>>()?; + let plan = SparseStructuralPlan { + layers, + num_items: self.num_items, + num_visible_items: self.num_visible_items, + }; + plan.validate(self.row_domain)?; + if buffers.next().is_some() { + return Err(Error::invalid_input_source( + "Sparse layout has unused structural buffers".into(), + )); + } + + let page_meta = Arc::new(SparseStructuralCacheableState { + chunk_meta, + chunk_value_offsets: chunk_value_offsets.into(), + plan, + row_domain: self.row_domain, + }); + Ok(page_meta as Arc) + } + .boxed() + } + + fn try_load(&mut self, data: &Arc) -> Result<()> { + self.page_meta = Some(data.clone().as_arc_any().downcast().map_err(|_| { + Error::invalid_input_source("Cached sparse page data has an unexpected type".into()) + })?); + Ok(()) + } + + fn schedule_ranges( + &self, + ranges: &[Range], + io: &Arc, + ) -> Result> { + let page_meta = self.page_meta.as_ref().ok_or_else(|| { + Error::internal("Sparse page scheduler has not been initialized".to_string()) + })?; + let encoded_row_domain = page_meta + .row_domain + .checked_mul(self.row_scale) + .ok_or_else(|| { + Error::invalid_input_source("Sparse structural encoded row domain overflows".into()) + })?; + let num_rows = validate_slice_ranges(ranges, encoded_row_domain, "row")?; + let ranges = ranges + .iter() + .map(|range| { + if !range.start.is_multiple_of(self.row_scale) + || !range.end.is_multiple_of(self.row_scale) + { + return Err(Error::invalid_input_source( + format!( + "Sparse structural encoded row range {}..{} is not aligned to fixed-size-list scale {}", + range.start, range.end, self.row_scale + ) + .into(), + )); + } + Ok((range.start / self.row_scale)..(range.end / self.row_scale)) + }) + .collect::>>()?; + + let mut chunks_needed = Vec::new(); + let selection = slice_sparse_plan(&page_meta.plan, &ranges, page_meta.row_domain)?; + for value_range in &selection.leaf_ranges { + chunks_needed.extend(Self::value_chunk_range( + &page_meta.chunk_value_offsets, + value_range.clone(), + )?); + } + chunks_needed.sort_unstable(); + chunks_needed.dedup(); + + let mut loaded_chunks = self.lookup_value_chunks(&chunks_needed)?; + let chunk_ranges = loaded_chunks + .iter() + .map(|chunk| chunk.byte_range.clone()) + .collect::>(); + let loaded_chunk_data = io.submit_request(chunk_ranges, self.priority); + let ranges = VecDeque::from(ranges); + let value_decompressor = self.value_decompressor.clone(); + let value_encoding = self.value_encoding.clone(); + let data_type = self.data_type.clone(); + let page_meta = page_meta.clone(); + let num_buffers = self.num_buffers; + let has_large_chunk = self.has_large_chunk; + let row_scale = self.row_scale; + + let res = async move { + let loaded_chunk_data = loaded_chunk_data.await?; + for (loaded_chunk, chunk_data) in loaded_chunks.iter_mut().zip(loaded_chunk_data) { + loaded_chunk.data = LanceBuffer::from_bytes(chunk_data, 1); + } + + Ok(Box::new(SparseStructuralDecoder { + value_decompressor, + value_encoding, + data_type, + page_meta, + loaded_chunks: Arc::new(loaded_chunks), + ranges, + offset_in_current_range: 0, + num_rows, + row_scale, + num_buffers, + has_large_chunk, + }) as Box) + } + .boxed(); + Ok(vec![PageLoadTask { + decoder_fut: res, + num_rows, + }]) + } +} + +#[derive(Debug)] +struct SparseStructuralDecoder { + value_decompressor: Arc, + value_encoding: CompressiveEncoding, + data_type: DataType, + page_meta: Arc, + loaded_chunks: Arc>, + ranges: VecDeque>, + offset_in_current_range: u64, + num_rows: u64, + row_scale: u64, + num_buffers: u64, + has_large_chunk: bool, +} + +impl SparseStructuralDecoder { + fn drain_ranges(&mut self, mut rows_desired: u64) -> Result>> { + if !rows_desired.is_multiple_of(self.row_scale) { + return Err(Error::invalid_input_source( + format!( + "Sparse page decoder drain of {} encoded rows is not aligned to fixed-size-list scale {}", + rows_desired, self.row_scale + ) + .into(), + )); + } + rows_desired /= self.row_scale; + let mut ranges = Vec::new(); + while rows_desired > 0 { + let range = self.ranges.front().ok_or_else(|| { + Error::invalid_input_source( + "Sparse page decoder was asked to drain more rows than were scheduled".into(), + ) + })?; + let start = range + .start + .checked_add(self.offset_in_current_range) + .ok_or_else(|| { + Error::invalid_input_source("Sparse page decoder row offset overflows".into()) + })?; + let rows_available = range.end.checked_sub(start).ok_or_else(|| { + Error::invalid_input_source( + "Sparse page decoder range offset exceeds the scheduled range".into(), + ) + })?; + let rows_to_take = rows_available.min(rows_desired); + let end = start.checked_add(rows_to_take).ok_or_else(|| { + Error::invalid_input_source("Sparse page decoder row range overflows".into()) + })?; + ranges.push(start..end); + rows_desired -= rows_to_take; + self.offset_in_current_range += rows_to_take; + if self.offset_in_current_range == range.end - range.start { + self.offset_in_current_range = 0; + self.ranges.pop_front(); + } + } + Ok(ranges) + } +} + +impl StructuralPageDecoder for SparseStructuralDecoder { + fn drain(&mut self, num_rows: u64) -> Result> { + Ok(Box::new(DecodeSparseStructuralTask { + row_ranges: self.drain_ranges(num_rows)?, + value_decompressor: self.value_decompressor.clone(), + value_encoding: self.value_encoding.clone(), + data_type: self.data_type.clone(), + page_meta: self.page_meta.clone(), + loaded_chunks: self.loaded_chunks.clone(), + num_buffers: self.num_buffers, + has_large_chunk: self.has_large_chunk, + })) + } + + fn num_rows(&self) -> u64 { + self.num_rows + } +} + +#[derive(Debug)] +struct DecodeSparseStructuralTask { + row_ranges: Vec>, + value_decompressor: Arc, + value_encoding: CompressiveEncoding, + data_type: DataType, + page_meta: Arc, + loaded_chunks: Arc>, + num_buffers: u64, + has_large_chunk: bool, +} + +impl DecodeSparseStructuralTask { + fn read_chunk_size( + buf: &[u8], + offset: &mut usize, + width: usize, + chunk_idx: usize, + ) -> Result { + let end = offset.checked_add(width).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural value chunk {chunk_idx} size header overflows").into(), + ) + })?; + let bytes = buf.get(*offset..end).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural value chunk {chunk_idx} has a truncated size header") + .into(), + ) + })?; + let size = match width { + 2 => u32::from(u16::from_le_bytes(bytes.try_into().map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural value chunk {chunk_idx} has a malformed u16 size") + .into(), + ) + })?)), + 4 => u32::from_le_bytes(bytes.try_into().map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural value chunk {chunk_idx} has a malformed u32 size") + .into(), + ) + })?), + _ => { + return Err(Error::internal(format!( + "Unsupported sparse value chunk size width {width}" + ))); + } + }; + *offset = end; + Ok(size) + } + + fn expected_fixed_bytes(num_values: u64, bits_per_value: u64, label: &str) -> Result { + let bits = num_values.checked_mul(bits_per_value).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} decoded bit length overflows").into(), + ) + })?; + usize_from_u64(bits.div_ceil(8), label) + } + + fn validate_fixed_buffer( + buffer: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + label: &str, + ) -> Result<()> { + let expected = Self::expected_fixed_bytes(num_values, bits_per_value, label)?; + if buffer.len() != expected { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} buffer has {} bytes, expected {}", + buffer.len(), + expected + ) + .into(), + )); + } + Ok(()) + } + + fn validate_variable_buffer( + buffer: &LanceBuffer, + num_values: u64, + bits_per_offset: u64, + ) -> Result<()> { + let width = usize_from_u64(bits_per_offset / 8, "variable offset width")?; + let offset_count = num_values.checked_add(1).ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset count overflows".into()) + })?; + let table_len = usize_from_u64(offset_count, "variable offset count")? + .checked_mul(width) + .ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset table size overflows".into()) + })?; + if buffer.len() < table_len { + return Err(Error::invalid_input_source( + format!( + "Sparse variable buffer has {} bytes, smaller than its {}-byte offset table", + buffer.len(), + table_len + ) + .into(), + )); + } + let mut previous = None; + for index in 0..usize_from_u64(offset_count, "variable offset count")? { + let start = index.checked_mul(width).ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset index overflows".into()) + })?; + let end = start.checked_add(width).ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset range overflows".into()) + })?; + let bytes = buffer.as_ref().get(start..end).ok_or_else(|| { + Error::invalid_input_source("Sparse variable offset table is truncated".into()) + })?; + let offset = match width { + 4 => u64::from(u32::from_le_bytes(bytes.try_into().map_err(|_| { + Error::invalid_input_source("Sparse variable u32 offset is malformed".into()) + })?)), + 8 => u64::from_le_bytes(bytes.try_into().map_err(|_| { + Error::invalid_input_source("Sparse variable u64 offset is malformed".into()) + })?), + _ => { + return Err(Error::invalid_input_source( + format!( + "Sparse variable offset width {} is unsupported", + bits_per_offset + ) + .into(), + )); + } + }; + if offset < table_len as u64 || offset > buffer.len() as u64 { + return Err(Error::invalid_input_source( + format!( + "Sparse variable offset {} is outside payload range {}..{}", + offset, + table_len, + buffer.len() + ) + .into(), + )); + } + if previous.is_some_and(|previous| offset < previous) { + return Err(Error::invalid_input_source( + "Sparse variable offsets are not monotonically increasing".into(), + )); + } + previous = Some(offset); + } + Ok(()) + } + + fn validate_fsl_buffers( + fsl: &pb21::FixedSizeList, + buffers: &[LanceBuffer], + num_values: u64, + buffer_index: &mut usize, + ) -> Result<()> { + use pb21::compressive_encoding::Compression; + + let child_values = num_values.checked_mul(fsl.items_per_value).ok_or_else(|| { + Error::invalid_input_source("Sparse fixed-size-list value count overflows".into()) + })?; + if fsl.has_validity { + let validity = buffers.get(*buffer_index).ok_or_else(|| { + Error::invalid_input_source( + "Sparse fixed-size-list value validity buffer is missing".into(), + ) + })?; + Self::validate_fixed_buffer(validity, child_values, 1, "fixed-size-list validity")?; + *buffer_index = buffer_index.checked_add(1).ok_or_else(|| { + Error::invalid_input_source("Sparse fixed-size-list buffer index overflows".into()) + })?; + } + let values = fsl.values.as_deref().ok_or_else(|| { + Error::invalid_input_source("Sparse fixed-size-list value encoding is missing".into()) + })?; + match values.compression.as_ref() { + Some(Compression::FixedSizeList(inner)) => { + Self::validate_fsl_buffers(inner, buffers, child_values, buffer_index) + } + Some(Compression::Flat(flat)) => { + let values = buffers.get(*buffer_index).ok_or_else(|| { + Error::invalid_input_source( + "Sparse fixed-size-list leaf value buffer is missing".into(), + ) + })?; + Self::validate_fixed_buffer( + values, + child_values, + flat.bits_per_value, + "fixed-size-list leaf values", + )?; + *buffer_index = buffer_index.checked_add(1).ok_or_else(|| { + Error::invalid_input_source( + "Sparse fixed-size-list buffer index overflows".into(), + ) + })?; + Ok(()) + } + _ => Err(Error::invalid_input_source( + "Sparse fixed-size-list value encoding is malformed".into(), + )), + } + } + + fn validate_value_buffers(&self, buffers: &[LanceBuffer], num_values: u64) -> Result<()> { + use pb21::compressive_encoding::Compression; + + let compression = self.value_encoding.compression.as_ref().ok_or_else(|| { + Error::invalid_input_source("Sparse value compression is missing".into()) + })?; + match compression { + Compression::Flat(flat) => Self::validate_fixed_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source("Sparse flat value buffer is missing".into()) + })?, + num_values, + flat.bits_per_value, + "flat values", + ), + Compression::InlineBitpacking(bitpacking) => { + let buffer = buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline-bitpacked value buffer is missing".into(), + ) + })?; + if num_values > 1024 { + return Err(Error::invalid_input_source( + format!( + "Sparse inline-bitpacked chunk has {} values, exceeding 1024", + num_values + ) + .into(), + )); + } + let word_bytes = usize_from_u64( + bitpacking.uncompressed_bits_per_value / 8, + "inline bitpacking word width", + )?; + let header = buffer.as_ref().get(..word_bytes).ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline-bitpacked buffer is missing its bit-width header".into(), + ) + })?; + let bit_width = + header + .iter() + .enumerate() + .try_fold(0_u64, |value, (idx, byte)| { + let shift = u32::try_from(idx.checked_mul(8).ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline bit-width shift overflows".into(), + ) + })?) + .map_err(|_| { + Error::invalid_input_source( + "Sparse inline bit-width shift exceeds u32".into(), + ) + })?; + Ok::<_, Error>(value | (u64::from(*byte) << shift)) + })?; + if bit_width > bitpacking.uncompressed_bits_per_value { + return Err(Error::invalid_input_source( + format!( + "Sparse inline bit width {} exceeds uncompressed width {}", + bit_width, bitpacking.uncompressed_bits_per_value + ) + .into(), + )); + } + let payload_bytes = usize_from_u64( + bit_width.checked_mul(1024).ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline-bitpacked payload size overflows".into(), + ) + })? / 8, + "inline-bitpacked payload size", + )?; + let expected = word_bytes.checked_add(payload_bytes).ok_or_else(|| { + Error::invalid_input_source( + "Sparse inline-bitpacked buffer size overflows".into(), + ) + })?; + if buffer.len() != expected { + return Err(Error::invalid_input_source( + format!( + "Sparse inline-bitpacked buffer has {} bytes, expected {}", + buffer.len(), + expected + ) + .into(), + )); + } + Ok(()) + } + Compression::Variable(variable) => { + let offsets = variable + .offsets + .as_deref() + .and_then(|encoding| encoding.compression.as_ref()) + .and_then(|compression| match compression { + Compression::Flat(flat) => Some(flat), + _ => None, + }) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse variable offset encoding is malformed".into(), + ) + })?; + Self::validate_variable_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse variable value buffer is missing".into(), + ) + })?, + num_values, + offsets.bits_per_value, + ) + } + Compression::Fsst(fsst) => { + let variable = fsst + .values + .as_deref() + .and_then(|encoding| encoding.compression.as_ref()) + .and_then(|compression| match compression { + Compression::Variable(variable) => Some(variable), + _ => None, + }) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse FSST value encoding is malformed".into(), + ) + })?; + let offsets = variable + .offsets + .as_deref() + .and_then(|encoding| encoding.compression.as_ref()) + .and_then(|compression| match compression { + Compression::Flat(flat) => Some(flat), + _ => None, + }) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse FSST offset encoding is malformed".into(), + ) + })?; + Self::validate_variable_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source("Sparse FSST value buffer is missing".into()) + })?, + num_values, + offsets.bits_per_value, + ) + } + Compression::ByteStreamSplit(split) => { + let bits = split + .values + .as_deref() + .and_then(|encoding| encoding.compression.as_ref()) + .and_then(|compression| match compression { + Compression::Flat(flat) => Some(flat.bits_per_value), + _ => None, + }) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse byte-stream-split encoding is malformed".into(), + ) + })?; + Self::validate_fixed_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse byte-stream-split buffer is missing".into(), + ) + })?, + num_values, + bits, + "byte-stream-split values", + ) + } + Compression::FixedSizeList(fsl) => { + let mut buffer_index = 0; + Self::validate_fsl_buffers(fsl, buffers, num_values, &mut buffer_index)?; + if buffer_index != buffers.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse fixed-size-list descriptor consumed {} of {} buffers", + buffer_index, + buffers.len() + ) + .into(), + )); + } + Ok(()) + } + Compression::PackedStruct(packed) => { + let bits = packed.bits_per_value.iter().try_fold(0_u64, |sum, bits| { + sum.checked_add(*bits).ok_or_else(|| { + Error::invalid_input_source( + "Sparse packed-struct bit width sum overflows".into(), + ) + }) + })?; + Self::validate_fixed_buffer( + buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse packed-struct value buffer is missing".into(), + ) + })?, + num_values, + bits, + "packed-struct values", + ) + } + Compression::Rle(rle) => { + if buffers.len() != 2 { + return Err(Error::invalid_input_source( + format!( + "Sparse RLE value chunk has {} buffers, expected 2", + buffers.len() + ) + .into(), + )); + } + SparseStructuralScheduler::validate_general_child_buffer( + SparseStructuralScheduler::require_encoding(&rle.values, "RLE values")?, + buffers + .first() + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse RLE value buffer is missing after count validation".into(), + ) + })? + .as_ref(), + "value chunk RLE values", + )?; + SparseStructuralScheduler::validate_general_child_buffer( + SparseStructuralScheduler::require_encoding( + &rle.run_lengths, + "RLE run lengths", + )?, + buffers + .get(1) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse RLE run-length buffer is missing after count validation" + .into(), + ) + })? + .as_ref(), + "value chunk RLE run lengths", + ) + } + Compression::General(general) => { + let buffer = buffers.first().ok_or_else(|| { + Error::invalid_input_source( + "Sparse general-compressed value chunk is missing its first buffer".into(), + ) + })?; + SparseStructuralScheduler::validate_general_buffer_header( + general, + buffer.as_ref(), + "value chunk", + ) + } + _ => Err(Error::invalid_input_source( + "Sparse value chunk uses an unsupported compression descriptor".into(), + )), + } + } + + fn loaded_chunk(&self, chunk_idx: usize) -> Result<&LoadedChunk> { + let index = self + .loaded_chunks + .binary_search_by_key(&chunk_idx, |chunk| chunk.chunk_idx) + .map_err(|_| { + Error::internal(format!( + "Sparse structural decode missing loaded value chunk {}", + chunk_idx + )) + })?; + self.loaded_chunks.get(index).ok_or_else(|| { + Error::internal(format!( + "Sparse structural loaded chunk index {} is missing", + index + )) + }) + } + + fn decode_value_chunk(&self, chunk: &LoadedChunk) -> Result { + let buf = &chunk.data; + if buf.len() < 2 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} is too small for its header: {} bytes", + chunk.chunk_idx, + buf.len() + ) + .into(), + )); + } + let num_levels = u16::from_le_bytes( + buf.as_ref() + .get(..2) + .and_then(|bytes| bytes.try_into().ok()) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} has a malformed level header", + chunk.chunk_idx + ) + .into(), + ) + })?, + ); + let mut offset: usize = 2; + if num_levels != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk unexpectedly contains {} rep/def levels", + num_levels + ) + .into(), + )); + } + + let size_width = if self.has_large_chunk { 4 } else { 2 }; + let num_buffers = usize::try_from(self.num_buffers).map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk has too many buffers: {}", + self.num_buffers + ) + .into(), + ) + })?; + let sizes_len = num_buffers.checked_mul(size_width).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural value chunk buffer-size header overflows".into(), + ) + })?; + let header_len = offset.checked_add(sizes_len).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural value chunk header length overflows".into(), + ) + })?; + if buf.len() < header_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} is too small for {} buffer sizes: {} bytes", + chunk.chunk_idx, + self.num_buffers, + buf.len() + ) + .into(), + )); + } + let buffer_sizes = (0..num_buffers) + .map(|_| Self::read_chunk_size(buf, &mut offset, size_width, chunk.chunk_idx)) + .collect::>>()?; + + offset = offset + .checked_add(pad_bytes::(offset)) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} padded header overflows", + chunk.chunk_idx + ) + .into(), + ) + })?; + if offset > buf.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} is missing padding after its header", + chunk.chunk_idx + ) + .into(), + )); + } + let buffers = buffer_sizes + .into_iter() + .map(|buf_size| { + let buf_size = buf_size as usize; + let end = offset.checked_add(buf_size).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} buffer size overflows", + chunk.chunk_idx + ) + .into(), + ) + })?; + if end > buf.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} buffer extends past chunk end", + chunk.chunk_idx + ) + .into(), + )); + } + let buffer = buf.slice_with_length(offset, buf_size); + offset = end; + offset = offset + .checked_add(pad_bytes::(offset)) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} padded buffer range overflows", + chunk.chunk_idx + ) + .into(), + ) + })?; + if offset > buf.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} padding extends past chunk end", + chunk.chunk_idx + ) + .into(), + )); + } + Ok(buffer) + }) + .collect::>>()?; + + if offset != buf.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} consumed {} of {} bytes", + chunk.chunk_idx, + offset, + buf.len() + ) + .into(), + )); + } + + self.validate_value_buffers(&buffers, chunk.items_in_chunk)?; + + let decoded = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.value_decompressor + .decompress(buffers, chunk.items_in_chunk) + })) + .map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} decompression panicked", + chunk.chunk_idx + ) + .into(), + ) + })? + .map_err(|error| { + Error::invalid_input_source( + format!( + "Sparse structural value chunk {} decompression failed: {error}", + chunk.chunk_idx + ) + .into(), + ) + })?; + if decoded.num_values() != chunk.items_in_chunk { + return Err(Error::invalid_input_source( + format!( + "Sparse structural value chunk {} decoded {} values, expected {}", + chunk.chunk_idx, + decoded.num_values(), + chunk.items_in_chunk + ) + .into(), + )); + } + Ok(decoded) + } + + fn append_value_range( + &self, + value_range: Range, + data_builder: &mut DataBlockBuilder, + chunk_cache: &mut Option<(usize, DataBlock)>, + ) -> Result<()> { + let mut value_start = value_range.start; + while value_start < value_range.end { + let chunk_idx = SparseStructuralScheduler::value_chunk_index( + &self.page_meta.chunk_value_offsets, + value_start, + )?; + let chunk_value_start = *self + .page_meta + .chunk_value_offsets + .get(chunk_idx) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse value chunk {} has no start offset", chunk_idx).into(), + ) + })?; + let chunk_value_end = *self + .page_meta + .chunk_value_offsets + .get(chunk_idx.checked_add(1).ok_or_else(|| { + Error::invalid_input_source("Sparse value chunk index overflows".into()) + })?) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse value chunk {} has no end offset", chunk_idx).into(), + ) + })?; + let take_end = value_range.end.min(chunk_value_end); + if value_start < chunk_value_start || take_end <= value_start { + return Err(Error::invalid_input_source( + format!( + "Sparse value range {}..{} does not make progress in chunk {} covering {}..{}", + value_range.start, + value_range.end, + chunk_idx, + chunk_value_start, + chunk_value_end + ) + .into(), + )); + } + + if !matches!(chunk_cache, Some((cached_idx, _)) if *cached_idx == chunk_idx) { + let chunk = self.loaded_chunk(chunk_idx)?; + *chunk_cache = Some((chunk_idx, self.decode_value_chunk(chunk)?)); + } + let values = &chunk_cache + .as_ref() + .ok_or_else(|| Error::internal("Sparse structural chunk cache is empty"))? + .1; + data_builder.append( + values, + value_start - chunk_value_start..take_end - chunk_value_start, + )?; + value_start = take_end; + } + Ok(()) + } + + fn decode_checked(self) -> Result { + let selection = slice_sparse_plan( + &self.page_meta.plan, + &self.row_ranges, + self.page_meta.row_domain, + )?; + let estimated_size_bytes = self + .loaded_chunks + .iter() + .map(|chunk| chunk.data.len()) + .try_fold(0_usize, |total, len| total.checked_add(len)) + .and_then(|total| total.checked_mul(2)) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural decode size estimate overflows".into(), + ) + })?; + let mut data_builder = DataBlockBuilder::with_capacity_estimate( + u64::try_from(estimated_size_bytes).map_err(|_| { + Error::invalid_input_source("Sparse structural decode size exceeds u64::MAX".into()) + })?, + ); + let mut chunk_cache: Option<(usize, DataBlock)> = None; + let mut appended_values = false; + for value_range in &selection.leaf_ranges { + self.append_value_range(value_range.clone(), &mut data_builder, &mut chunk_cache)?; + appended_values = true; + } + + let data = if appended_values { + data_builder.finish() + } else { + DataBlock::from_array(new_empty_array(&self.data_type)) + }; + let unraveler = RepDefUnraveler::new_sparse(selection.plan); + Ok(DecodedPage { + data, + repdef: unraveler, + }) + } +} + +impl DecodePageTask for DecodeSparseStructuralTask { + fn decode(self: Box) -> Result { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| (*self).decode_checked())) + .map_err(|_| { + Error::invalid_input_source( + "Sparse structural page decoding panicked on malformed input".into(), + ) + })? + } +} + +struct SparseStructuralSelection { + plan: SparseStructuralPlan, + leaf_ranges: Vec>, +} + +struct SparsePositionSelection { + positions: SparsePositionSet, + ordinal_ranges: Vec>, +} + +fn validate_slice_ranges(ranges: &[Range], domain_len: u64, label: &str) -> Result { + let mut total = 0_u64; + for range in ranges { + if range.start > range.end || range.end > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} slice {}..{} is outside domain {}", + range.start, range.end, domain_len + ) + .into(), + )); + } + total = total.checked_add(range.end - range.start).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} slice length overflows").into(), + ) + })?; + } + Ok(total) +} + +fn push_coalesced_range(ranges: &mut Vec>, range: Range) { + if range.is_empty() { + return; + } + if let Some(last) = ranges.last_mut() + && last.end == range.start + { + last.end = range.end; + return; + } + ranges.push(range); +} + +fn position_segments_to_set( + segments: Vec>, + output_domain: u64, + label: &str, +) -> Result { + let segments = coalesce_ranges(segments); + if segments.is_empty() { + return Ok(SparsePositionSet::empty()); + } + if segments.len() == 1 { + let segment = segments.first().ok_or_else(|| { + Error::internal("Sparse structural segment unexpectedly missing".to_string()) + })?; + if segment.start == 0 && segment.end == output_domain { + return Ok(SparsePositionSet::all(output_domain)); + } + return Ok(SparsePositionSet::range( + segment.start, + segment.end - segment.start, + )); + } + + let total_len = segments + .iter() + .map(|range| range.end - range.start) + .try_fold(0_u64, |sum, len| { + sum.checked_add(len).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} segment length overflows").into(), + ) + }) + })?; + let total_len = usize::try_from(total_len).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} segment length exceeds usize::MAX").into(), + ) + })?; + let mut positions = Vec::with_capacity(total_len); + for segment in segments { + positions.extend(segment); + } + SparsePositionSet::from_positions(positions, output_domain, label) +} + +fn select_position_set( + positions: &SparsePositionSet, + ranges: &[Range], + domain_len: u64, + label: &str, +) -> Result { + let output_domain = validate_slice_ranges(ranges, domain_len, label)?; + let mut segments = Vec::new(); + let mut ordinal_ranges = Vec::new(); + let mut output_base = 0_u64; + + match positions { + SparsePositionSet::Empty => {} + SparsePositionSet::All { len } => { + if *len != domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} all set length {} does not match domain {}", + len, domain_len + ) + .into(), + )); + } + for range in ranges { + let range_len = range.end - range.start; + let output_end = output_base.checked_add(range_len).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output domain overflows").into(), + ) + })?; + push_coalesced_range(&mut segments, output_base..output_end); + push_coalesced_range(&mut ordinal_ranges, range.clone()); + output_base = output_end; + } + } + SparsePositionSet::Range { start, len } => { + let source_end = start.checked_add(*len).ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} range overflows").into(), + ) + })?; + if source_end > domain_len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} range {}..{} is outside domain {}", + start, source_end, domain_len + ) + .into(), + )); + } + for range in ranges { + let intersect_start = range.start.max(*start); + let intersect_end = range.end.min(source_end); + if intersect_start < intersect_end { + let out_start = output_base + .checked_add(intersect_start - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output position overflows") + .into(), + ) + })?; + let out_end = output_base + .checked_add(intersect_end - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output position overflows") + .into(), + ) + })?; + push_coalesced_range(&mut segments, out_start..out_end); + push_coalesced_range( + &mut ordinal_ranges, + intersect_start - *start..intersect_end - *start, + ); + } + output_base = output_base + .checked_add(range.end - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output domain overflows").into(), + ) + })?; + } + } + SparsePositionSet::Explicit(source_positions) => { + let mut out_positions = Vec::new(); + for range in ranges { + let idx_start = + source_positions.partition_point(|position| *position < range.start); + let idx_end = source_positions.partition_point(|position| *position < range.end); + if idx_start < idx_end { + let selected_positions = + source_positions.get(idx_start..idx_end).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} explicit position slice is invalid" + ) + .into(), + ) + })?; + for position in selected_positions { + out_positions.push( + output_base + .checked_add(*position - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural {label} output position overflows" + ) + .into(), + ) + })?, + ); + } + push_coalesced_range(&mut ordinal_ranges, idx_start as u64..idx_end as u64); + } + output_base = output_base + .checked_add(range.end - range.start) + .ok_or_else(|| { + Error::invalid_input_source( + format!("Sparse structural {label} output domain overflows").into(), + ) + })?; + } + let positions = SparsePositionSet::from_positions(out_positions, output_domain, label)?; + return Ok(SparsePositionSelection { + positions, + ordinal_ranges, + }); + } + } + + Ok(SparsePositionSelection { + positions: position_segments_to_set(segments, output_domain, label)?, + ordinal_ranges, + }) +} + +fn select_validity_set( + validity: &SparseValiditySet, + ranges: &[Range], + domain_len: u64, + label: &str, +) -> Result { + Ok(SparseValiditySet { + meaning: validity.meaning, + positions: select_position_set(&validity.positions, ranges, domain_len, label)?.positions, + }) +} + +fn offsets_from_counts(counts: &[u64]) -> Result> { + let mut offsets = Vec::with_capacity(counts.len() + 1); + let mut offset = 0_u64; + offsets.push(offset); + for count in counts { + offset = offset.checked_add(*count).ok_or_else(|| { + Error::invalid_input_source("Sparse structural list count offsets overflow".into()) + })?; + offsets.push(offset); + } + Ok(offsets) +} + +fn coalesce_ranges(ranges: Vec>) -> Vec> { + let mut coalesced: Vec> = Vec::with_capacity(ranges.len()); + for range in ranges { + if range.is_empty() { + continue; + } + if let Some(last) = coalesced.last_mut() + && last.end == range.start + { + last.end = range.end; + continue; + } + coalesced.push(range); + } + coalesced +} + +fn slice_list_layer( + num_slots: u64, + num_child_slots: u64, + non_empty_positions: &SparsePositionSet, + counts: &SparseCountSet, + validity: &SparseValiditySet, + ranges: &[Range], +) -> Result<(SparseStructuralLayerPlan, Vec>)> { + if non_empty_positions.len() != counts.len() { + return Err(Error::invalid_input_source( + "Sparse structural list has mismatched non-empty positions and counts".into(), + )); + } + let count_sum = counts.sum()?; + if count_sum != num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match child slots {}", + count_sum, num_child_slots + ) + .into(), + )); + } + + let non_empty_selection = + select_position_set(non_empty_positions, ranges, num_slots, "list non-empty")?; + let out_counts = select_count_set( + counts, + &non_empty_selection.ordinal_ranges, + non_empty_selection.positions.len(), + )?; + let child_ranges = + child_ranges_from_counts(counts, num_child_slots, &non_empty_selection.ordinal_ranges)?; + let out_validity = select_validity_set(validity, ranges, num_slots, "list validity")?; + let out_num_slots = validate_slice_ranges(ranges, num_slots, "list")?; + let out_num_child_slots = out_counts.sum()?; + Ok(( + SparseStructuralLayerPlan::List { + num_slots: out_num_slots, + num_child_slots: out_num_child_slots, + non_empty_positions: non_empty_selection.positions, + counts: out_counts, + validity: out_validity, + }, + child_ranges, + )) +} + +fn select_count_set( + counts: &SparseCountSet, + ordinal_ranges: &[Range], + selected_len: u64, +) -> Result { + match counts { + SparseCountSet::Empty => { + if selected_len != 0 { + return Err(Error::invalid_input_source( + "Sparse structural selected non-empty positions but counts are empty".into(), + )); + } + Ok(SparseCountSet::Empty) + } + SparseCountSet::Constant { value, len } => { + validate_ordinal_ranges(ordinal_ranges, *len, "constant list counts")?; + Ok(SparseCountSet::constant(*value, selected_len)) + } + SparseCountSet::Explicit { + counts: source_counts, + .. + } => { + validate_ordinal_ranges( + ordinal_ranges, + source_counts.len() as u64, + "explicit list counts", + )?; + let selected_len = usize::try_from(selected_len).map_err(|_| { + Error::invalid_input_source( + "Sparse structural selected count length exceeds usize::MAX".into(), + ) + })?; + let mut out_counts = Vec::with_capacity(selected_len); + for range in ordinal_ranges { + let start = usize::try_from(range.start).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count ordinal exceeds usize::MAX".into(), + ) + })?; + let end = usize::try_from(range.end).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count ordinal exceeds usize::MAX".into(), + ) + })?; + out_counts.extend_from_slice(source_counts.get(start..end).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural explicit count slice is invalid".into(), + ) + })?); + } + SparseCountSet::from_counts(out_counts) + } + } +} + +fn validate_ordinal_ranges(ranges: &[Range], len: u64, label: &str) -> Result<()> { + for range in ranges { + if range.start > range.end || range.end > len { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} ordinal range {}..{} is outside {} values", + range.start, range.end, len + ) + .into(), + )); + } + } + Ok(()) +} + +fn child_ranges_from_counts( + counts: &SparseCountSet, + num_child_slots: u64, + ordinal_ranges: &[Range], +) -> Result>> { + if ordinal_ranges.is_empty() { + return Ok(Vec::new()); + } + let child_ranges = match counts { + SparseCountSet::Empty => { + return Err(Error::invalid_input_source( + "Sparse structural selected non-empty positions but counts are empty".into(), + )); + } + SparseCountSet::Constant { value, len } => { + let expected_child_slots = value.checked_mul(*len).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list constant count sum overflows child slots".into(), + ) + })?; + if expected_child_slots != num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list constant count sum {} does not match child slots {}", + expected_child_slots, num_child_slots + ) + .into(), + )); + } + validate_ordinal_ranges(ordinal_ranges, *len, "constant list counts")?; + ordinal_ranges + .iter() + .map(|range| { + let start = range.start.checked_mul(*value).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list child range start overflows".into(), + ) + })?; + let end = range.end.checked_mul(*value).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list child range end overflows".into(), + ) + })?; + Ok(start..end) + }) + .collect::>>()? + } + SparseCountSet::Explicit { + counts, + offsets: value_offsets, + } => { + let last_offset = *value_offsets.last().ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list count offsets are unexpectedly empty".into(), + ) + })?; + if last_offset != num_child_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list count sum {} does not match child slots {}", + last_offset, num_child_slots + ) + .into(), + )); + } + validate_ordinal_ranges(ordinal_ranges, counts.len() as u64, "explicit list counts")?; + ordinal_ranges + .iter() + .map(|range| { + let start = usize::try_from(range.start).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count ordinal exceeds usize::MAX".into(), + ) + })?; + let end = usize::try_from(range.end).map_err(|_| { + Error::invalid_input_source( + "Sparse structural count ordinal exceeds usize::MAX".into(), + ) + })?; + let start_offset = *value_offsets.get(start).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list start offset is missing".into(), + ) + })?; + let end_offset = *value_offsets.get(end).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list end offset is missing".into(), + ) + })?; + Ok(start_offset..end_offset) + }) + .collect::>>()? + } + }; + Ok(coalesce_ranges(child_ranges)) +} + +fn slice_sparse_plan( + plan: &SparseStructuralPlan, + row_ranges: &[Range], + row_domain: u64, +) -> Result { + plan.validate(row_domain)?; + let mut selected_ranges = row_ranges.to_vec(); + let mut selected_domain = row_domain; + let mut sliced_layers = Vec::with_capacity(plan.layers.len()); + + for layer in &plan.layers { + match layer { + SparseStructuralLayerPlan::Validity { + num_slots, + validity, + } => { + if selected_domain != *num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural validity slice domain {} does not match {} slots", + selected_domain, num_slots + ) + .into(), + )); + } + let out_num_slots = + validate_slice_ranges(&selected_ranges, *num_slots, "validity")?; + sliced_layers.push(SparseStructuralLayerPlan::Validity { + num_slots: out_num_slots, + validity: select_validity_set( + validity, + &selected_ranges, + *num_slots, + "validity", + )?, + }); + } + SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + .. + } => { + if selected_domain != *num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list slice domain {} does not match {} slots", + selected_domain, num_slots + ) + .into(), + )); + } + let (sliced_layer, child_ranges) = slice_list_layer( + *num_slots, + *num_child_slots, + non_empty_positions, + counts, + validity, + &selected_ranges, + )?; + sliced_layers.push(sliced_layer); + selected_ranges = child_ranges; + selected_domain = *num_child_slots; + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots, + dimension, + validity, + } => { + if selected_domain != *num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list slice domain {} does not match {} slots", + selected_domain, num_slots + ) + .into(), + )); + } + let num_child_slots = num_slots.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={}, dimension={}", + num_slots, dimension + ) + .into(), + ) + })?; + let out_num_slots = + validate_slice_ranges(&selected_ranges, *num_slots, "fixed-size-list")?; + let child_ranges = selected_ranges + .iter() + .map(|range| { + let start = range.start.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural fixed-size-list child range start overflows" + .into(), + ) + })?; + let end = range.end.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural fixed-size-list child range end overflows" + .into(), + ) + })?; + Ok(start..end) + }) + .collect::>>()?; + sliced_layers.push(SparseStructuralLayerPlan::FixedSizeList { + num_slots: out_num_slots, + dimension: *dimension, + validity: select_validity_set( + validity, + &selected_ranges, + *num_slots, + "fixed-size-list validity", + )?, + }); + selected_ranges = child_ranges; + selected_domain = num_child_slots; + } + } + } + + if selected_domain != plan.num_visible_items { + return Err(Error::invalid_input_source( + format!( + "Sparse structural selected terminal domain {} does not match {} visible items", + selected_domain, plan.num_visible_items + ) + .into(), + )); + } + let num_visible_items = + validate_slice_ranges(&selected_ranges, plan.num_visible_items, "visible value")?; + let num_items = SparseStructuralPlan::expected_num_items(&sliced_layers, num_visible_items)?; + Ok(SparseStructuralSelection { + plan: SparseStructuralPlan { + layers: sliced_layers, + num_items, + num_visible_items, + }, + leaf_ranges: coalesce_ranges(selected_ranges), + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use crate::{ + compression::DefaultDecompressionStrategy, + encodings::physical::block::{CompressionConfig, CompressionScheme}, + testing::SimulatedScheduler, + }; + + use super::*; + + async fn initialize_scheduler( + scheduler: &mut SparseStructuralScheduler, + io: &Arc, + ) -> Result> { + let initialization = scheduler.init_layout()?; + let buffers = io + .submit_request(initialization.required_ranges().cloned().collect(), 0) + .await?; + let buffers = initialization.with_buffers(buffers)?; + let state = scheduler.init_from_buffers(buffers, io).await?; + scheduler.try_load(&state)?; + Ok(state) + } + + fn position_set( + positions: pb21::sparse_position_set::Positions, + num_positions: u64, + ) -> Option { + Some(pb21::SparsePositionSet { + positions: Some(positions), + num_positions, + }) + } + + fn position_empty() -> Option { + position_set( + pb21::sparse_position_set::Positions::Empty(pb21::SparsePositionEmpty {}), + 0, + ) + } + + fn position_all(num_positions: u64) -> Option { + position_set( + pb21::sparse_position_set::Positions::All(pb21::SparsePositionAll {}), + num_positions, + ) + } + + fn position_explicit(num_positions: u64) -> Option { + position_set( + pb21::sparse_position_set::Positions::Explicit(ProtobufUtils21::flat(64, None)), + num_positions, + ) + } + + fn general_lz4(values: CompressiveEncoding) -> CompressiveEncoding { + ProtobufUtils21::wrapped(CompressionConfig::new(CompressionScheme::Lz4, None), values) + .unwrap() + } + + fn validity( + meaning: pb21::sparse_validity_set::Meaning, + positions: Option, + ) -> Option { + Some(pb21::SparseValiditySet { + meaning: meaning as i32, + positions, + }) + } + + fn null_positions( + positions: Option, + ) -> Option { + validity( + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions, + positions, + ) + } + + fn count_empty() -> Option { + Some(pb21::SparseCountSet { + counts: Some(pb21::sparse_count_set::Counts::Empty( + pb21::SparseCountEmpty {}, + )), + }) + } + + fn count_constant(value: u64) -> Option { + Some(pb21::SparseCountSet { + counts: Some(pb21::sparse_count_set::Counts::Constant( + pb21::SparseCountConstant { value }, + )), + }) + } + + fn sparse_layout() -> pb21::SparseLayout { + pb21::SparseLayout { + value_compression: Some(ProtobufUtils21::flat(32, None)), + num_buffers: 1, + num_items: 1, + num_visible_items: 1, + has_large_chunk: false, + structural_layers: Vec::new(), + } + } + + fn validity_layer( + num_slots: u64, + validity: Option, + ) -> pb21::SparseStructuralLayer { + pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::Validity( + pb21::SparseValidityLayer { + num_slots, + validity, + }, + )), + } + } + + fn list_layer( + num_slots: u64, + num_child_slots: u64, + non_empty_positions: Option, + counts: Option, + validity: Option, + ) -> pb21::SparseStructuralLayer { + pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::List( + pb21::SparseListLayer { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + }, + )), + } + } + + fn fixed_size_list_layer( + num_slots: u64, + dimension: u64, + validity: Option, + ) -> pb21::SparseStructuralLayer { + pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::FixedSizeList( + pb21::SparseFixedSizeListLayer { + num_slots, + dimension, + validity, + }, + )), + } + } + + fn assert_invalid_input_contains(err: Error, expected: &str) { + let message = err.to_string(); + assert!( + matches!(&err, Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + assert!( + message.contains(expected), + "expected error to contain {expected:?}, got {message}" + ); + } + + #[test] + fn rejects_missing_layer_variants_and_item_count_mismatches() { + let decompressors = DefaultDecompressionStrategy::default(); + + let mut layout = sparse_layout(); + layout + .structural_layers + .push(pb21::SparseStructuralLayer { layer: None }); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "missing its layer variant"); + + let mut layout = sparse_layout(); + layout.num_items = 2; + layout + .structural_layers + .push(validity_layer(1, null_positions(position_empty()))); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "layers imply 1"); + } + + #[test] + fn accepts_large_structural_descriptors() { + let explicit_values = 8_u64 * 1024 * 1024 + 1; + let metadata_size = explicit_values * 8; + let value_position = metadata_size; + let value_size = explicit_values * 16; + let structural_position = value_position + value_size; + let structural_size = explicit_values * std::mem::size_of::() as u64; + let mut layout = sparse_layout(); + layout.num_items = explicit_values; + layout.num_visible_items = explicit_values; + layout.structural_layers.push(validity_layer( + explicit_values, + null_positions(position_explicit(explicit_values)), + )); + + SparseStructuralScheduler::try_new( + &[ + (0, metadata_size), + (value_position, value_size), + (structural_position, structural_size), + ], + 0, + explicit_values, + DataType::Int32, + &layout, + &DefaultDecompressionStrategy::default(), + ) + .unwrap(); + } + + #[test] + fn accepts_deep_supported_value_encodings() { + let encoding = (0..300).fold(ProtobufUtils21::flat(32, None), |values, _| { + ProtobufUtils21::fsl(1, false, values) + }); + + assert_eq!( + SparseStructuralScheduler::validate_value_encoding(&encoding).unwrap(), + 1 + ); + } + + fn null_set(positions: SparsePositionSet) -> SparseValiditySet { + SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions, + } + } + + fn valid_set(positions: SparsePositionSet) -> SparseValiditySet { + SparseValiditySet { + meaning: SparseValidityMeaning::ValidPositions, + positions, + } + } + + #[test] + fn semantic_position_and_count_sets_project_without_materializing_ranges() { + let ranges = [1..4]; + assert_eq!( + select_position_set(&SparsePositionSet::Empty, &ranges, 5, "empty") + .unwrap() + .positions, + SparsePositionSet::Empty + ); + assert_eq!( + select_position_set(&SparsePositionSet::all(5), &ranges, 5, "all") + .unwrap() + .positions, + SparsePositionSet::all(3) + ); + assert_eq!( + select_position_set(&SparsePositionSet::range(1, 3), &ranges, 5, "range") + .unwrap() + .positions, + SparsePositionSet::all(3) + ); + assert_eq!( + select_position_set( + &SparsePositionSet::Explicit(vec![0, 2, 4]), + &[0..1, 4..5], + 5, + "explicit", + ) + .unwrap() + .positions, + SparsePositionSet::all(2) + ); + + assert_eq!( + select_count_set(&SparseCountSet::Empty, &[], 0).unwrap(), + SparseCountSet::Empty + ); + assert_eq!( + select_count_set(&SparseCountSet::constant(2, 3), &[0..1, 2..3], 2,).unwrap(), + SparseCountSet::constant(2, 2) + ); + assert_eq!( + select_count_set( + &SparseCountSet::from_counts(vec![1, 2, 3]).unwrap(), + &[0..1, 2..3], + 2, + ) + .unwrap(), + SparseCountSet::from_counts(vec![1, 3]).unwrap() + ); + } + + #[test] + fn validity_polarities_rebuild_the_same_arrow_domain() { + let mut null_builder = BooleanBufferBuilder::new(4); + null_set(SparsePositionSet::Explicit(vec![1, 3])) + .append_to(&mut null_builder, 4) + .unwrap(); + assert_eq!( + null_builder.finish().iter().collect::>(), + vec![true, false, true, false] + ); + + let mut valid_builder = BooleanBufferBuilder::new(4); + valid_set(SparsePositionSet::range(1, 2)) + .append_to(&mut valid_builder, 4) + .unwrap(); + assert_eq!( + valid_builder.finish().iter().collect::>(), + vec![false, true, true, false] + ); + } + + #[test] + fn schema_layer_mismatches_are_invalid_input() { + let mut missing = SparseStructuralUnraveler::new(SparseStructuralPlan { + layers: Vec::new(), + num_items: 1, + num_visible_items: 1, + }); + let mut validity = BooleanBufferBuilder::new(1); + let err = missing.unravel_validity(&mut validity).unwrap_err(); + assert_invalid_input_contains(err, "fewer layers than the Arrow schema"); + + let mut fixed_size_list = SparseStructuralUnraveler::new(SparseStructuralPlan { + layers: vec![SparseStructuralLayerPlan::FixedSizeList { + num_slots: 2, + dimension: 2, + validity: null_set(SparsePositionSet::empty()), + }], + num_items: 4, + num_visible_items: 4, + }); + let mut validity = BooleanBufferBuilder::new(2); + let err = fixed_size_list.unravel_validity(&mut validity).unwrap_err(); + assert_invalid_input_contains(err, "does not match the Arrow schema"); + + let extra = SparseStructuralUnraveler::new(SparseStructuralPlan { + layers: vec![SparseStructuralLayerPlan::Validity { + num_slots: 1, + validity: null_set(SparsePositionSet::empty()), + }], + num_items: 1, + num_visible_items: 1, + }); + let err = extra.ensure_exhausted().unwrap_err(); + assert_invalid_input_contains(err, "1 unconsumed layer"); + } + + fn nested_plan() -> SparseStructuralPlan { + SparseStructuralPlan { + layers: vec![ + SparseStructuralLayerPlan::Validity { + num_slots: 6, + validity: null_set(SparsePositionSet::Explicit(vec![1, 4])), + }, + SparseStructuralLayerPlan::List { + num_slots: 6, + num_child_slots: 5, + non_empty_positions: SparsePositionSet::Explicit(vec![0, 2, 5]), + counts: SparseCountSet::from_counts(vec![2, 1, 2]).unwrap(), + validity: null_set(SparsePositionSet::Explicit(vec![1, 4])), + }, + SparseStructuralLayerPlan::FixedSizeList { + num_slots: 5, + dimension: 2, + validity: valid_set(SparsePositionSet::range(1, 3)), + }, + ], + num_items: 13, + num_visible_items: 10, + } + } + + #[test] + fn discontiguous_projection_preserves_outer_to_inner_order() { + let selection = slice_sparse_plan(&nested_plan(), &[5..6, 0..1], 6).unwrap(); + assert_eq!(selection.leaf_ranges, vec![6..10, 0..4]); + assert_eq!(selection.plan.num_visible_items, 8); + selection.plan.validate(2).unwrap(); + + let SparseStructuralLayerPlan::List { + non_empty_positions, + counts, + .. + } = &selection.plan.layers[1] + else { + panic!("expected projected list layer"); + }; + assert_eq!(*non_empty_positions, SparsePositionSet::all(2)); + assert_eq!(*counts, SparseCountSet::constant(2, 2)); + } + + #[test] + fn no_value_projection_keeps_empty_and_null_list_structure() { + let selection = slice_sparse_plan(&nested_plan(), &[1..2, 3..4], 6).unwrap(); + assert!(selection.leaf_ranges.is_empty()); + assert_eq!(selection.plan.num_visible_items, 0); + selection.plan.validate(2).unwrap(); + + let SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + } = &selection.plan.layers[1] + else { + panic!("expected projected list layer"); + }; + assert_eq!((*num_slots, *num_child_slots), (2, 0)); + assert_eq!(*non_empty_positions, SparsePositionSet::Empty); + assert_eq!(*counts, SparseCountSet::Empty); + assert_eq!(*validity, null_set(SparsePositionSet::range(0, 1))); + } + + #[test] + fn rejects_missing_value_compression_and_buffer_count_mismatch() { + let decompressors = DefaultDecompressionStrategy::default(); + let mut layout = sparse_layout(); + layout.value_compression = None; + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "missing value compression"); + + let mut layout = sparse_layout(); + layout.num_buffers = 2; + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "declares 2 value buffers"); + } + + #[test] + fn rejects_inconsistent_chunk_count_before_io() { + let decompressors = DefaultDecompressionStrategy::default(); + + let layout = sparse_layout(); + let err = SparseStructuralScheduler::try_new( + &[(0, 16), (16, 8)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "declares 2 chunks for 1 visible items"); + } + + #[cfg(feature = "lz4")] + #[test] + fn accepts_large_general_decompression_headers() { + let encoding = general_lz4(ProtobufUtils21::flat(64, None)); + let Some(pb21::compressive_encoding::Compression::General(general)) = + encoding.compression.as_ref() + else { + panic!("expected General compression"); + }; + let declared_size = 65_u32 * 1024 * 1024; + + SparseStructuralScheduler::validate_general_buffer_header( + general, + &declared_size.to_le_bytes(), + "test", + ) + .unwrap(); + } + + #[cfg(feature = "lz4")] + #[tokio::test] + async fn rejects_malformed_general_structural_buffers_as_invalid_input() { + let mut layout = sparse_layout(); + layout.structural_layers.push(validity_layer( + 1, + null_positions(position_set( + pb21::sparse_position_set::Positions::Explicit(general_lz4(ProtobufUtils21::flat( + 64, None, + ))), + 1, + )), + )); + let decompressors = DefaultDecompressionStrategy::default(); + + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8), (16, 8)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + let mut data = Vec::new(); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&[0; 8]); + data.extend_from_slice(&8_u32.to_le_bytes()); + data.extend_from_slice(&[0xff; 4]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + let Err(err) = initialize_scheduler(&mut scheduler, &io).await else { + panic!("expected malformed General buffer to be rejected"); + }; + assert_invalid_input_contains(err, "decompression failed"); + } + + #[cfg(feature = "lz4")] + #[tokio::test] + async fn rejects_malformed_general_value_buffer() { + let mut layout = sparse_layout(); + layout.value_compression = Some(general_lz4(ProtobufUtils21::flat(32, None))); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + + let mut data = Vec::new(); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&0_u16.to_le_bytes()); + data.extend_from_slice(&0_u16.to_le_bytes()); + data.extend_from_slice(&[0; 4]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + initialize_scheduler(&mut scheduler, &io).await.unwrap(); + let mut page_tasks = scheduler.schedule_ranges(&[0..1], &io).unwrap(); + let mut decoder = page_tasks.pop().unwrap().decoder_fut.await.unwrap(); + let Err(err) = decoder.drain(1).unwrap().decode() else { + panic!("expected malformed General value buffer to be rejected"); + }; + assert_invalid_input_contains(err, "missing its length prefix"); + } + + #[test] + fn rejects_layer_domain_and_fixed_size_list_mismatches() { + let decompressors = DefaultDecompressionStrategy::default(); + let mut layout = sparse_layout(); + layout.num_items = 2; + layout.num_visible_items = 2; + layout + .structural_layers + .push(validity_layer(1, null_positions(position_empty()))); + layout + .structural_layers + .push(validity_layer(2, null_positions(position_empty()))); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "layer 1 has 2 slots, expected 1"); + + let mut layout = sparse_layout(); + layout.num_items = 4; + layout.num_visible_items = 4; + layout.structural_layers.push(fixed_size_list_layer( + 2, + 3, + null_positions(position_empty()), + )); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 2, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "terminal domain has 6 slots"); + } + + #[test] + fn rejects_invalid_validity_and_list_count_semantics() { + let decompressors = DefaultDecompressionStrategy::default(); + let mut layout = sparse_layout(); + layout.structural_layers.push(validity_layer( + 1, + validity( + pb21::sparse_validity_set::Meaning::SparseValidityUnspecified, + position_empty(), + ), + )); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "meaning is unspecified"); + + let mut layout = sparse_layout(); + layout.num_items = 3; + layout.num_visible_items = 3; + layout.structural_layers.push(list_layer( + 2, + 3, + position_all(2), + count_constant(2), + null_positions(position_empty()), + )); + let err = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 2, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap_err(); + assert_invalid_input_contains(err, "count sum 4 does not match child slots 3"); + } + + #[tokio::test] + async fn rejects_unordered_explicit_positions_after_decompression() { + let mut layout = sparse_layout(); + layout.num_items = 4; + layout.num_visible_items = 4; + layout + .structural_layers + .push(validity_layer(4, null_positions(position_explicit(2)))); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8), (16, 16)], + 0, + 4, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + + let mut data = Vec::new(); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&4_u32.to_le_bytes()); + data.extend_from_slice(&[0; 8]); + data.extend_from_slice(&3_u64.to_le_bytes()); + data.extend_from_slice(&0_u64.to_le_bytes()); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + + let Err(err) = initialize_scheduler(&mut scheduler, &io).await else { + panic!("expected unordered sparse positions to be rejected"); + }; + assert_invalid_input_contains(err, "positions must be strictly increasing"); + } + + #[tokio::test] + async fn rejects_chunk_metadata_value_and_byte_sum_mismatches() { + let layout = sparse_layout(); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + let mut data = Vec::new(); + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&[0; 8]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + let Err(err) = initialize_scheduler(&mut scheduler, &io).await else { + panic!("expected chunk byte sum mismatch"); + }; + assert_invalid_input_contains(err, "describes 16 value bytes"); + + let mut layout = sparse_layout(); + layout.num_items = 2; + layout.num_visible_items = 2; + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + 2, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + let mut data = Vec::new(); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&[0; 8]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + let Err(err) = initialize_scheduler(&mut scheduler, &io).await else { + panic!("expected chunk value sum mismatch"); + }; + assert_invalid_input_contains(err, "metadata has 1, layout has 2"); + } + + #[tokio::test] + async fn rejects_malformed_value_chunk_before_decompression() { + let layout = sparse_layout(); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, 8)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + let mut data = Vec::new(); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&[0; 8]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + initialize_scheduler(&mut scheduler, &io).await.unwrap(); + let mut page_tasks = scheduler.schedule_ranges(&[0..1], &io).unwrap(); + let page_task = page_tasks.pop().unwrap(); + let mut decoder = page_task.decoder_fut.await.unwrap(); + let decode_task = decoder.drain(1).unwrap(); + let Err(err) = decode_task.decode() else { + panic!("expected malformed value chunk to be rejected"); + }; + assert_invalid_input_contains(err, "flat values buffer has 0 bytes, expected 4"); + } + + #[tokio::test] + async fn accepts_value_chunks_larger_than_64_mib() { + let chunk_size = 65_u64 * 1024 * 1024; + let layout = sparse_layout(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 8), (8, chunk_size)], + 0, + 1, + DataType::Int32, + &layout, + &DefaultDecompressionStrategy::default(), + ) + .unwrap(); + let words_minus_one = u32::try_from(chunk_size / MINIBLOCK_ALIGNMENT as u64 - 1).unwrap(); + let mut metadata = Vec::new(); + metadata.extend_from_slice(&words_minus_one.to_le_bytes()); + metadata.extend_from_slice(&1_u32.to_le_bytes()); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(metadata))); + + initialize_scheduler(&mut scheduler, &io).await.unwrap(); + } + + #[tokio::test] + async fn rejects_value_chunks_above_the_miniblock_limit() { + let num_values = miniblock::MAX_CONFIGURABLE_MINIBLOCK_VALUES + 1; + let mut layout = sparse_layout(); + layout.num_items = num_values; + layout.num_visible_items = num_values; + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 16), (16, 16)], + 0, + num_values, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + let mut data = Vec::new(); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&(num_values as u32).to_le_bytes()); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&0_u32.to_le_bytes()); + data.extend_from_slice(&[0; 16]); + let io: Arc = Arc::new(SimulatedScheduler::new(Bytes::from(data))); + + let Err(err) = initialize_scheduler(&mut scheduler, &io).await else { + panic!("expected oversized value chunk to be rejected"); + }; + assert_invalid_input_contains(err, "exceeding the mini-block limit"); + } + + #[derive(Debug, Clone)] + struct RecordingIo { + data: Bytes, + calls: Arc>>>>, + } + + impl RecordingIo { + fn new(data: Bytes) -> Self { + Self { + data, + calls: Arc::new(Mutex::new(Vec::new())), + } + } + } + + impl EncodingsIo for RecordingIo { + fn submit_request( + &self, + ranges: Vec>, + _priority: u64, + ) -> BoxFuture<'static, Result>> { + self.calls.lock().unwrap().push(ranges.clone()); + let data = self.data.clone(); + async move { + ranges + .into_iter() + .map(|range| { + let start = usize_from_u64(range.start, "test range start")?; + let end = usize_from_u64(range.end, "test range end")?; + if start > end || end > data.len() { + return Err(Error::invalid_input_source( + "Test I/O range is outside fixture data".into(), + )); + } + Ok(data.slice(start..end)) + }) + .collect() + } + .boxed() + } + } + + #[tokio::test] + async fn selective_read_requests_only_intersecting_value_chunk() { + let mut layout = sparse_layout(); + layout.num_items = 4; + layout.num_visible_items = 4; + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 16), (16, 32)], + 0, + 4, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + + let mut data = Vec::new(); + for _ in 0..2 { + data.extend_from_slice(&1_u32.to_le_bytes()); + data.extend_from_slice(&2_u32.to_le_bytes()); + } + for values in [[10_i32, 20], [30, 40]] { + data.extend_from_slice(&0_u16.to_le_bytes()); + data.extend_from_slice(&8_u16.to_le_bytes()); + data.extend_from_slice(&[0; 4]); + for value in values { + data.extend_from_slice(&value.to_le_bytes()); + } + } + + let io = Arc::new(RecordingIo::new(Bytes::from(data))); + let trait_io: Arc = io.clone(); + initialize_scheduler(&mut scheduler, &trait_io) + .await + .unwrap(); + let mut page_tasks = scheduler.schedule_ranges(&[2..3], &trait_io).unwrap(); + let page_task = page_tasks.pop().unwrap(); + let mut decoder = page_task.decoder_fut.await.unwrap(); + let decoded = decoder.drain(1).unwrap().decode().unwrap(); + assert_eq!(decoded.data.num_values(), 1); + + let calls = io.calls.lock().unwrap(); + assert_eq!(calls.as_slice(), &[vec![0..16], vec![32..48]]); + } + + #[tokio::test] + async fn empty_leaf_selection_rebuilds_offsets_without_value_io() { + let mut layout = sparse_layout(); + layout.num_visible_items = 0; + layout.structural_layers.push(list_layer( + 1, + 0, + position_empty(), + count_empty(), + null_positions(position_empty()), + )); + let decompressors = DefaultDecompressionStrategy::default(); + let mut scheduler = SparseStructuralScheduler::try_new( + &[(0, 0), (0, 0)], + 0, + 1, + DataType::Int32, + &layout, + &decompressors, + ) + .unwrap(); + let io = Arc::new(RecordingIo::new(Bytes::new())); + let trait_io: Arc = io.clone(); + initialize_scheduler(&mut scheduler, &trait_io) + .await + .unwrap(); + let mut page_tasks = scheduler.schedule_ranges(&[0..1], &trait_io).unwrap(); + let page_task = page_tasks.pop().unwrap(); + let mut decoder = page_task.decoder_fut.await.unwrap(); + let decoded = decoder.drain(1).unwrap().decode().unwrap(); + assert_eq!(decoded.data.num_values(), 0); + + let mut repdef = CompositeRepDefUnraveler::new(vec![decoded.repdef]); + let (offsets, validity) = repdef.unravel_offsets::().unwrap(); + assert_eq!(offsets.as_ref(), &[0, 0]); + assert!(validity.is_none()); + + let calls = io.calls.lock().unwrap(); + assert_eq!(calls.len(), 2); + assert!(calls[1].is_empty(), "value payload must not be requested"); + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/vendor/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs new file mode 100644 index 000000000..5ec1d32b5 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -0,0 +1,2407 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Sparse structural planning and serialization. + +use std::iter; + +use arrow_buffer::BooleanBuffer; +use lance_core::{Error, Result, datatypes::Field, utils::bit::pad_bytes}; + +use crate::{ + buffer::LanceBuffer, + compression::{CompressionStrategy, compress_required_block}, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + decoder::PageEncoding, + encoder::EncodedPage, + format::pb21::{self, CompressiveEncoding}, + repdef::{NormalizedStructuralLayer, NormalizedStructuralPlan}, + statistics::ComputeStat, +}; + +use super::super::MiniblockChunkSize; + +use super::{ + SparseCountSet, SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan, + SparseValidityMeaning, SparseValiditySet, +}; +use crate::encodings::logical::primitive::{ + FILL_BYTE, MINIBLOCK_ALIGNMENT, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext}, +}; + +#[derive(Clone, Copy, Default)] +struct PositionSetStats { + count: u64, + first: u64, + last: u64, + is_contiguous: bool, +} + +impl PositionSetStats { + fn observe(&mut self, position: u64) { + if self.count == 0 { + self.first = position; + self.is_contiguous = true; + } else if self.last.checked_add(1) != Some(position) { + self.is_contiguous = false; + } + self.last = position; + self.count += 1; + } + + fn encoded_cost(&self, domain_len: u64) -> u64 { + if self.count == 0 || self.count == domain_len || self.is_contiguous { + 0 + } else { + self.count + } + } + + fn to_set( + self, + validity: &BooleanBuffer, + want_valid: bool, + domain_len: u64, + label: &str, + ) -> Result { + if self.count == 0 { + return Ok(SparsePositionSet::empty()); + } + if self.count == domain_len { + return Ok(SparsePositionSet::all(domain_len)); + } + if self.is_contiguous { + return Ok(SparsePositionSet::range(self.first, self.count)); + } + + let capacity = usize::try_from(self.count).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} positions exceed usize::MAX").into(), + ) + })?; + let mut positions = Vec::with_capacity(capacity); + for (index, is_valid) in validity.iter().enumerate() { + if is_valid == want_valid { + positions.push(u64::try_from(index).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} position exceeds u64::MAX").into(), + ) + })?); + } + } + SparsePositionSet::from_positions(positions, domain_len, label) + } +} + +fn usize_to_u64(value: usize, label: &str) -> Result { + u64::try_from(value).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural {label} {value} exceeds u64::MAX").into(), + ) + }) +} + +fn validity_set( + validity: Option<&BooleanBuffer>, + num_slots: usize, + label: &str, +) -> Result { + let Some(validity) = validity else { + return Ok(SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions: SparsePositionSet::empty(), + }); + }; + if validity.len() != num_slots { + return Err(Error::invalid_input_source( + format!( + "Sparse structural {label} validity length {} does not match {} slots", + validity.len(), + num_slots + ) + .into(), + )); + } + + let domain_len = usize_to_u64(num_slots, "validity domain")?; + let mut valid_stats = PositionSetStats::default(); + let mut null_stats = PositionSetStats::default(); + for (index, is_valid) in validity.iter().enumerate() { + let index = usize_to_u64(index, "validity position")?; + if is_valid { + valid_stats.observe(index); + } else { + null_stats.observe(index); + } + } + + if null_stats.count == 0 { + return Ok(SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions: SparsePositionSet::empty(), + }); + } + if valid_stats.count == 0 { + return Ok(SparseValiditySet { + meaning: SparseValidityMeaning::ValidPositions, + positions: SparsePositionSet::empty(), + }); + } + + let valid_cost = valid_stats.encoded_cost(domain_len); + let null_cost = null_stats.encoded_cost(domain_len); + if valid_cost < null_cost { + Ok(SparseValiditySet { + meaning: SparseValidityMeaning::ValidPositions, + positions: valid_stats.to_set(validity, true, domain_len, label)?, + }) + } else { + Ok(SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions: null_stats.to_set(validity, false, domain_len, label)?, + }) + } +} + +/// Builds the semantic sparse plan directly from the once-normalized Arrow layers. +pub(in crate::encodings::logical::primitive) fn plan( + normalized: &NormalizedStructuralPlan, + num_visible_items: u64, +) -> Result { + let mut layers = Vec::with_capacity(normalized.layers().len()); + let mut num_items = num_visible_items; + + for layer in normalized.layers() { + match layer { + NormalizedStructuralLayer::Validity { + validity, + num_slots, + } => { + layers.push(SparseStructuralLayerPlan::Validity { + num_slots: usize_to_u64(num_slots, "validity slot count")?, + validity: validity_set(validity, num_slots, "validity")?, + }); + } + NormalizedStructuralLayer::FixedSizeList { + validity, + dimension, + num_slots, + } => { + if dimension == 0 { + return Err(Error::invalid_input_source( + "Sparse structural fixed-size-list dimension is zero".into(), + )); + } + layers.push(SparseStructuralLayerPlan::FixedSizeList { + num_slots: usize_to_u64(num_slots, "fixed-size-list slot count")?, + dimension: usize_to_u64(dimension, "fixed-size-list dimension")?, + validity: validity_set(validity, num_slots, "fixed-size-list validity")?, + }); + } + NormalizedStructuralLayer::List { + offsets, + validity, + num_slots, + } => { + let expected_offsets = num_slots.checked_add(1).ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural list offset count overflows".into(), + ) + })?; + if offsets.len() != expected_offsets { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list has {} offsets for {} slots", + offsets.len(), + num_slots + ) + .into(), + )); + } + if offsets.first().copied() != Some(0) { + return Err(Error::invalid_input_source( + "Sparse structural list offsets must start at zero".into(), + )); + } + + let mut non_empty_positions = Vec::new(); + let mut counts = Vec::new(); + for slot in 0..num_slots { + let start = offsets[slot]; + let end = offsets[slot + 1]; + let count = end.checked_sub(start).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural list offsets decrease at slot {slot}: {start}..{end}" + ) + .into(), + ) + })?; + let is_valid = validity.is_none_or(|validity| validity.value(slot)); + if !is_valid && count != 0 { + return Err(Error::invalid_input_source( + format!( + "Sparse structural null list slot {slot} has {count} child slots" + ) + .into(), + )); + } + if is_valid && count > 0 { + non_empty_positions.push(usize_to_u64(slot, "list position")?); + counts.push(u64::try_from(count).map_err(|_| { + Error::invalid_input_source( + format!("Sparse structural list count {count} exceeds u64::MAX") + .into(), + ) + })?); + } + } + + let num_slots_u64 = usize_to_u64(num_slots, "list slot count")?; + let num_non_empty = usize_to_u64(non_empty_positions.len(), "list position count")?; + num_items = num_items + .checked_add(num_slots_u64 - num_non_empty) + .ok_or_else(|| { + Error::invalid_input_source( + "Sparse structural item count overflows u64".into(), + ) + })?; + let num_child_slots = offsets.last().copied().ok_or_else(|| { + Error::invalid_input_source("Sparse structural list has no offsets".into()) + })?; + let num_child_slots = u64::try_from(num_child_slots).map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse structural list child slot count {num_child_slots} is negative" + ) + .into(), + ) + })?; + layers.push(SparseStructuralLayerPlan::List { + num_slots: num_slots_u64, + num_child_slots, + non_empty_positions: SparsePositionSet::from_positions( + non_empty_positions, + num_slots_u64, + "list non-empty", + )?, + counts: SparseCountSet::from_counts(counts)?, + validity: validity_set(validity, num_slots, "list validity")?, + }); + } + } + } + + let row_domain = match layers.first() { + Some(SparseStructuralLayerPlan::Validity { num_slots, .. }) + | Some(SparseStructuralLayerPlan::List { num_slots, .. }) + | Some(SparseStructuralLayerPlan::FixedSizeList { num_slots, .. }) => *num_slots, + None => { + return Err(Error::invalid_input_source( + "Sparse structural encoding requires at least one Arrow structural layer".into(), + )); + } + }; + let plan = SparseStructuralPlan { + layers, + num_items, + num_visible_items, + }; + plan.validate(row_domain)?; + Ok(plan) +} + +/// ConstantLayout remains the canonical representation when no value payload exists. +pub(in crate::encodings::logical::primitive) fn uses_constant_layout( + plan: &SparseStructuralPlan, + field: &Field, +) -> bool { + if plan.num_visible_items == 0 { + return true; + } + if matches!(field.data_type(), arrow_schema::DataType::Struct(fields) if fields.is_empty()) { + return true; + } + + let Some(layer) = plan.layers.last() else { + return false; + }; + let (num_slots, validity) = match layer { + SparseStructuralLayerPlan::Validity { + num_slots, + validity, + } + | SparseStructuralLayerPlan::List { + num_slots, + validity, + .. + } + | SparseStructuralLayerPlan::FixedSizeList { + num_slots, + validity, + .. + } => (*num_slots, validity), + }; + match validity.meaning { + SparseValidityMeaning::NullPositions => validity.positions.len() == num_slots, + SparseValidityMeaning::ValidPositions => validity.positions.is_empty(), + } +} + +fn supports_fixed_size_list_values(data: &DataBlock) -> bool { + match data { + DataBlock::FixedWidth(_) => true, + DataBlock::FixedSizeList(list) => supports_fixed_size_list_values(list.child.as_ref()), + DataBlock::Nullable(nullable) => supports_fixed_size_list_values(nullable.data.as_ref()), + _ => false, + } +} + +/// Whether the sparse writer can encode this value block without changing its value path. +pub fn supports_value_block(data: &DataBlock) -> bool { + match data { + DataBlock::FixedWidth(_) | DataBlock::VariableWidth(_) => true, + DataBlock::Struct(data) => !data.has_variable_width_child(), + DataBlock::FixedSizeList(data) => supports_fixed_size_list_values(data.child.as_ref()), + DataBlock::Empty() + | DataBlock::Constant(_) + | DataBlock::AllNull(_) + | DataBlock::Nullable(_) + | DataBlock::Opaque(_) + | DataBlock::Dictionary(_) => false, + } +} + +struct SparseMiniBlockChunk { + buffer_sizes: Vec, + num_values: u32, +} + +struct SparseMiniBlockCompressed { + data: Vec, + chunks: Vec, +} + +struct SerializedValuePage { + num_buffers: u64, + data: LanceBuffer, + metadata: LanceBuffer, +} + +pub struct PreparedSparseValues { + num_values: u64, + value_compression: CompressiveEncoding, + values: SerializedValuePage, +} + +pub enum SparseValueInput { + Unprepared(DataBlock), + Prepared(PreparedSparseValues), +} + +struct EncodedStructuralPlan { + layers: Vec, + buffers: Vec, +} + +fn with_explicit_value_counts( + compressed: MiniBlockCompressed, +) -> Result { + let mut values_in_previous_chunks = 0_u64; + let mut chunks = Vec::with_capacity(compressed.chunks.len()); + for chunk in compressed.chunks { + let num_values = chunk.num_values(values_in_previous_chunks, compressed.num_values); + values_in_previous_chunks = values_in_previous_chunks + .checked_add(num_values) + .ok_or_else(|| Error::internal("Sparse value count overflows u64".to_string()))?; + chunks.push(SparseMiniBlockChunk { + buffer_sizes: chunk.buffer_sizes, + num_values: u32::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!( + "Sparse value chunk has {num_values} visible values, which exceeds the u32 metadata limit" + ) + .into(), + ) + })?, + }); + } + if values_in_previous_chunks != compressed.num_values { + return Err(Error::internal(format!( + "Sparse value chunks describe {values_in_previous_chunks} values, expected {}", + compressed.num_values + ))); + } + Ok(SparseMiniBlockCompressed { + data: compressed.data, + chunks, + }) +} + +fn serialize_value_chunks( + compressed: SparseMiniBlockCompressed, + miniblock_chunk_size: MiniblockChunkSize, +) -> Result { + let bytes_data = compressed.data.iter().map(LanceBuffer::len).sum::(); + let num_buffers = compressed.data.len(); + let mut data_buffer = Vec::with_capacity(bytes_data + 9 * num_buffers); + let mut metadata = Vec::with_capacity(compressed.chunks.len() * 8); + let mut buffer_offsets = vec![0_usize; num_buffers]; + + for chunk in compressed.chunks { + if chunk.buffer_sizes.len() != num_buffers { + return Err(Error::internal(format!( + "Sparse chunk has {} value buffer sizes, expected {num_buffers}", + chunk.buffer_sizes.len() + ))); + } + + let chunk_start = data_buffer.len(); + debug_assert_eq!(chunk_start % MINIBLOCK_ALIGNMENT, 0); + data_buffer.extend_from_slice(&0_u16.to_le_bytes()); + if miniblock_chunk_size == MiniblockChunkSize::U32 { + for buffer_size in &chunk.buffer_sizes { + data_buffer.extend_from_slice(&buffer_size.to_le_bytes()); + } + } else { + for buffer_size in &chunk.buffer_sizes { + let buffer_size = u16::try_from(*buffer_size).map_err(|_| { + Error::internal(format!( + "Sparse value buffer size ({buffer_size} bytes) exceeds 16-bit metadata" + )) + })?; + data_buffer.extend_from_slice(&buffer_size.to_le_bytes()); + } + } + let add_padding = |buffer: &mut Vec| { + let padding = pad_bytes::(buffer.len()); + buffer.extend(iter::repeat_n(FILL_BYTE, padding)); + }; + add_padding(&mut data_buffer); + + for (buffer_size, (buffer, buffer_offset)) in chunk + .buffer_sizes + .iter() + .zip(compressed.data.iter().zip(buffer_offsets.iter_mut())) + { + let start = *buffer_offset; + let end = start.checked_add(*buffer_size as usize).ok_or_else(|| { + Error::internal("Sparse value buffer range overflows".to_string()) + })?; + let bytes = buffer.as_ref().get(start..end).ok_or_else(|| { + Error::internal(format!( + "Sparse value chunk requests bytes {start}..{end} from a {}-byte buffer", + buffer.len() + )) + })?; + *buffer_offset = end; + data_buffer.extend_from_slice(bytes); + add_padding(&mut data_buffer); + } + + let chunk_bytes = data_buffer.len() - chunk_start; + if chunk_bytes == 0 || !chunk_bytes.is_multiple_of(MINIBLOCK_ALIGNMENT) { + return Err(Error::internal(format!( + "Sparse value chunk size {chunk_bytes} is not a positive multiple of {MINIBLOCK_ALIGNMENT}" + ))); + } + let words_minus_one = chunk_bytes / MINIBLOCK_ALIGNMENT - 1; + metadata.extend_from_slice( + &u32::try_from(words_minus_one) + .map_err(|_| { + Error::internal(format!( + "Sparse value chunk size {chunk_bytes} exceeds the metadata limit" + )) + })? + .to_le_bytes(), + ); + metadata.extend_from_slice(&chunk.num_values.to_le_bytes()); + } + + for (index, (consumed, buffer)) in buffer_offsets + .iter() + .zip(compressed.data.iter()) + .enumerate() + { + if *consumed != buffer.len() { + return Err(Error::internal(format!( + "Sparse value buffer {index} consumed {consumed} bytes, expected {}", + buffer.len() + ))); + } + } + + Ok(SerializedValuePage { + num_buffers: usize_to_u64(num_buffers, "value buffer count")?, + data: LanceBuffer::from(data_buffer), + metadata: LanceBuffer::from(metadata), + }) +} + +pub fn prepare_values( + field: &Field, + compression_strategy: &dyn CompressionStrategy, + data: DataBlock, + miniblock_chunk_size: MiniblockChunkSize, +) -> Result { + match &data { + DataBlock::AllNull(_) => { + return Err(Error::internal( + "All-null values must use ConstantLayout".to_string(), + )); + } + DataBlock::Dictionary(_) => { + return Err(Error::not_supported_source( + "Sparse layout does not support dictionary data blocks".into(), + )); + } + DataBlock::Struct(data) if data.has_variable_width_child() => { + return Err(Error::not_supported_source( + "Sparse layout does not support variable-width packed struct data blocks".into(), + )); + } + _ => {} + } + + let num_values = data.num_values(); + let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; + let support_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32; + let compression_context = MiniBlockCompressionContext::new(0, support_large_chunk, false); + let (compressed, value_compression) = compressor.compress(compression_context, data)?; + let values = serialize_value_chunks( + with_explicit_value_counts(compressed)?, + miniblock_chunk_size, + )?; + Ok(PreparedSparseValues { + num_values, + value_compression, + values, + }) +} + +fn encode_u64_values( + values: Vec, + compression_strategy: &dyn CompressionStrategy, +) -> Result<(LanceBuffer, CompressiveEncoding)> { + let num_values = usize_to_u64(values.len(), "u64 value count")?; + let mut block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(values), + bits_per_value: 64, + num_values, + block_info: BlockInfo::new(), + }); + block.compute_stat(); + let field = Field::new_arrow("", arrow_schema::DataType::UInt64, false)?; + compress_required_block(compression_strategy, &field, block) +} + +fn positions_to_deltas(positions: &[u64], label: &str) -> Result> { + let mut previous = 0_u64; + positions + .iter() + .copied() + .enumerate() + .map(|(index, position)| { + if index > 0 && position <= previous { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} positions must be strictly increasing") + .into(), + )); + } + let delta = if index == 0 { + position + } else { + position - previous + }; + previous = position; + Ok(delta) + }) + .collect() +} + +fn encode_position_set( + positions: &SparsePositionSet, + compression_strategy: &dyn CompressionStrategy, + label: &str, +) -> Result<(Option, pb21::SparsePositionSet)> { + let (buffer, positions_pb) = match positions { + SparsePositionSet::Empty => ( + None, + pb21::sparse_position_set::Positions::Empty(pb21::SparsePositionEmpty {}), + ), + SparsePositionSet::All { .. } => ( + None, + pb21::sparse_position_set::Positions::All(pb21::SparsePositionAll {}), + ), + SparsePositionSet::Range { start, len } => ( + None, + pb21::sparse_position_set::Positions::Range(pb21::SparsePositionRange { + start: *start, + length: *len, + }), + ), + SparsePositionSet::Explicit(positions) => { + if positions.is_empty() { + return Err(Error::internal(format!( + "Sparse structural {label} explicit set is empty" + ))); + } + let (buffer, encoding) = + encode_u64_values(positions_to_deltas(positions, label)?, compression_strategy)?; + ( + Some(buffer), + pb21::sparse_position_set::Positions::Explicit(encoding), + ) + } + }; + Ok(( + buffer, + pb21::SparsePositionSet { + positions: Some(positions_pb), + num_positions: positions.len(), + }, + )) +} + +fn encode_count_set( + counts: &SparseCountSet, + compression_strategy: &dyn CompressionStrategy, +) -> Result<(Option, pb21::SparseCountSet)> { + let (buffer, counts_pb) = match counts { + SparseCountSet::Empty => ( + None, + pb21::sparse_count_set::Counts::Empty(pb21::SparseCountEmpty {}), + ), + SparseCountSet::Constant { value, .. } => ( + None, + pb21::sparse_count_set::Counts::Constant(pb21::SparseCountConstant { value: *value }), + ), + SparseCountSet::Explicit { counts, .. } => { + if counts.is_empty() { + return Err(Error::internal( + "Sparse structural explicit count set is empty".to_string(), + )); + } + let (buffer, encoding) = encode_u64_values(counts.to_vec(), compression_strategy)?; + ( + Some(buffer), + pb21::sparse_count_set::Counts::Explicit(encoding), + ) + } + }; + Ok(( + buffer, + pb21::SparseCountSet { + counts: Some(counts_pb), + }, + )) +} + +fn encode_validity_set( + validity: &SparseValiditySet, + compression_strategy: &dyn CompressionStrategy, + label: &str, +) -> Result<(Option, pb21::SparseValiditySet)> { + let (buffer, positions) = + encode_position_set(&validity.positions, compression_strategy, label)?; + let meaning = match validity.meaning { + SparseValidityMeaning::NullPositions => { + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions + } + SparseValidityMeaning::ValidPositions => { + pb21::sparse_validity_set::Meaning::SparseValidityValidPositions + } + }; + Ok(( + buffer, + pb21::SparseValiditySet { + meaning: meaning as i32, + positions: Some(positions), + }, + )) +} + +fn encode_structural_plan( + plan: &SparseStructuralPlan, + compression_strategy: &dyn CompressionStrategy, +) -> Result { + let mut layers = Vec::with_capacity(plan.layers.len()); + let mut buffers = Vec::new(); + + for layer in &plan.layers { + match layer { + SparseStructuralLayerPlan::Validity { + num_slots, + validity, + } => { + let (validity_buffer, validity) = + encode_validity_set(validity, compression_strategy, "validity")?; + buffers.extend(validity_buffer); + layers.push(pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::Validity( + pb21::SparseValidityLayer { + num_slots: *num_slots, + validity: Some(validity), + }, + )), + }); + } + SparseStructuralLayerPlan::List { + num_slots, + num_child_slots, + non_empty_positions, + counts, + validity, + } => { + if non_empty_positions.len() != counts.len() { + return Err(Error::invalid_input_source( + format!( + "Sparse structural list has {} non-empty positions but {} counts", + non_empty_positions.len(), + counts.len() + ) + .into(), + )); + } + let (position_buffer, non_empty_positions) = encode_position_set( + non_empty_positions, + compression_strategy, + "list non-empty", + )?; + buffers.extend(position_buffer); + let (count_buffer, counts) = encode_count_set(counts, compression_strategy)?; + buffers.extend(count_buffer); + let (validity_buffer, validity) = + encode_validity_set(validity, compression_strategy, "list validity")?; + buffers.extend(validity_buffer); + layers.push(pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::List( + pb21::SparseListLayer { + num_slots: *num_slots, + num_child_slots: *num_child_slots, + non_empty_positions: Some(non_empty_positions), + counts: Some(counts), + validity: Some(validity), + }, + )), + }); + } + SparseStructuralLayerPlan::FixedSizeList { + num_slots, + dimension, + validity, + } => { + let (validity_buffer, validity) = encode_validity_set( + validity, + compression_strategy, + "fixed-size-list validity", + )?; + buffers.extend(validity_buffer); + num_slots.checked_mul(*dimension).ok_or_else(|| { + Error::invalid_input_source( + format!( + "Sparse structural fixed-size-list child slot count overflows: slots={num_slots}, dimension={dimension}" + ) + .into(), + ) + })?; + layers.push(pb21::SparseStructuralLayer { + layer: Some(pb21::sparse_structural_layer::Layer::FixedSizeList( + pb21::SparseFixedSizeListLayer { + num_slots: *num_slots, + dimension: *dimension, + validity: Some(validity), + }, + )), + }); + } + } + } + + Ok(EncodedStructuralPlan { layers, buffers }) +} + +#[allow(clippy::too_many_arguments)] +pub(in crate::encodings::logical::primitive) fn encode_page( + column_idx: u32, + field: &Field, + compression_strategy: &dyn CompressionStrategy, + values: SparseValueInput, + plan: SparseStructuralPlan, + row_number: u64, + num_rows: u64, + miniblock_chunk_size: MiniblockChunkSize, +) -> Result { + let PreparedSparseValues { + num_values, + value_compression, + values, + } = match values { + SparseValueInput::Unprepared(data) => { + prepare_values(field, compression_strategy, data, miniblock_chunk_size)? + } + SparseValueInput::Prepared(prepared) => prepared, + }; + if plan.num_visible_items != num_values { + return Err(Error::internal(format!( + "Sparse structural plan has {} visible items but data has {} values", + plan.num_visible_items, num_values + ))); + } + let structural = encode_structural_plan(&plan, compression_strategy)?; + let description = pb21::PageLayout { + layout: Some(pb21::page_layout::Layout::SparseLayout( + pb21::SparseLayout { + value_compression: Some(value_compression), + num_buffers: values.num_buffers, + num_items: plan.num_items, + num_visible_items: plan.num_visible_items, + has_large_chunk: miniblock_chunk_size == MiniblockChunkSize::U32, + structural_layers: structural.layers, + }, + )), + }; + + let mut page_data = Vec::with_capacity(2 + structural.buffers.len()); + page_data.push(values.metadata); + page_data.push(values.data); + page_data.extend(structural.buffers); + Ok(EncodedPage { + data: page_data, + description: PageEncoding::Structural(description), + num_rows, + row_number, + column_idx, + }) +} + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use arrow_array::{ + Array, ArrayRef, DictionaryArray, FixedSizeBinaryArray, FixedSizeListArray, Int8Array, + Int32Array, LargeListArray, ListArray, StringArray, StructArray, + builder::{Int32Builder, MapBuilder, StringBuilder}, + types::Int8Type, + }; + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field as ArrowField, Fields}; + + use crate::{ + constants::{ + PACKED_STRUCT_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, + STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_SPARSE, + }, + data::FixedSizeListBlock, + encoder::{ + ColumnIndexSequence, EncodingOptions, FieldEncoder, MIN_PAGE_BUFFER_ALIGNMENT, + OutOfLineBuffers, + }, + testing::{ + TestCases, TestEncoding, check_round_trip_encoding_of_data, test_encoding_strategy, + }, + }; + + use super::*; + + fn sparse_metadata() -> HashMap { + HashMap::from([( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_SPARSE.to_string(), + )]) + } + + fn structural_metadata(value: &str) -> HashMap { + HashMap::from([(STRUCTURAL_ENCODING_META_KEY.to_string(), value.to_string())]) + } + + fn null_buffer(validity: impl IntoIterator) -> NullBuffer { + NullBuffer::new(BooleanBuffer::from_iter(validity)) + } + + fn list_i32(offsets: Vec, validity: Option>) -> ArrayRef { + let num_values = offsets.last().copied().unwrap_or_default(); + let values = Arc::new(Int32Array::from_iter_values(0..num_values)) as ArrayRef; + Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + values, + validity.map(null_buffer), + ) + .unwrap(), + ) + } + + fn sparse_list_values( + num_rows: usize, + stride: usize, + values: ArrayRef, + item_field: Arc, + ) -> ArrayRef { + let mut offsets = Vec::with_capacity(num_rows + 1); + let mut num_values = 0_i32; + offsets.push(num_values); + for row in 0..num_rows { + if (row + 1).is_multiple_of(stride) || row + 1 == num_rows { + num_values += 1; + } + offsets.push(num_values); + } + assert_eq!(values.len(), num_values as usize); + Arc::new( + ListArray::try_new( + item_field, + OffsetBuffer::new(ScalarBuffer::from(offsets)), + values, + None, + ) + .unwrap(), + ) + } + + fn sparse_i32_list(num_rows: usize, stride: usize) -> ArrayRef { + let num_values = num_rows.div_ceil(stride); + sparse_list_values( + num_rows, + stride, + Arc::new(Int32Array::from_iter_values(0..num_values as i32)), + Arc::new(ArrowField::new("item", DataType::Int32, true)), + ) + } + + fn unsplittable_nested_list(values: ArrayRef, item_field: Arc) -> ArrayRef { + const NUM_INNER_LISTS: usize = 70_000; + + let mut inner_offsets = vec![0_i32; NUM_INNER_LISTS + 1]; + assert!(!values.is_empty()); + assert!(values.len() <= NUM_INNER_LISTS); + for value_index in 1..=values.len() { + inner_offsets[NUM_INNER_LISTS - values.len() + value_index] = value_index as i32; + } + let inner = Arc::new( + ListArray::try_new( + item_field, + OffsetBuffer::new(ScalarBuffer::from(inner_offsets)), + values, + None, + ) + .unwrap(), + ) as ArrayRef; + Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new("item", inner.data_type().clone(), true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, NUM_INNER_LISTS as i32])), + inner, + None, + ) + .unwrap(), + ) + } + + fn variable_packed_struct_values(num_values: usize) -> (ArrayRef, Arc) { + let fields = Fields::from(vec![ArrowField::new("value", DataType::Utf8, false)]); + let values = Arc::new(StructArray::new( + fields.clone(), + vec![Arc::new(StringArray::from_iter_values( + (0..num_values).map(|index| format!("value-{index}")), + ))], + None, + )) as ArrayRef; + let item_field = Arc::new( + ArrowField::new("item", DataType::Struct(fields), true).with_metadata(HashMap::from([ + (PACKED_STRUCT_META_KEY.to_string(), "true".to_string()), + ])), + ); + (values, item_field) + } + + fn dictionary_values(num_values: usize) -> (ArrayRef, Arc) { + let keys = Int8Array::from_iter_values((0..num_values).map(|index| (index % 2) as i8)); + let values = Arc::new(StringArray::from(vec!["value-0", "value-1"])); + let dictionary = Arc::new(DictionaryArray::::try_new(keys, values).unwrap()); + let item_field = Arc::new(ArrowField::new( + "item", + dictionary.data_type().clone(), + true, + )); + (dictionary, item_field) + } + + fn large_list_i32(offsets: Vec, validity: Option>) -> ArrayRef { + let num_values = offsets.last().copied().unwrap_or_default(); + let values = Arc::new(Int32Array::from_iter_values(0..num_values as i32)) as ArrayRef; + Arc::new( + LargeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + values, + validity.map(null_buffer), + ) + .unwrap(), + ) + } + + fn map_i32() -> ArrayRef { + let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + builder.keys().append_value("a"); + builder.values().append_value(1); + builder.append(true).unwrap(); + builder.append(false).unwrap(); + builder.append(true).unwrap(); + builder.keys().append_value("b"); + builder.values().append_null(); + builder.keys().append_value("c"); + builder.values().append_value(3); + builder.append(true).unwrap(); + Arc::new(builder.finish()) + } + + fn fixed_size_list_struct() -> ArrayRef { + let fields = Fields::from(vec![ArrowField::new("value", DataType::Int32, true)]); + let child = Arc::new(StructArray::new( + fields.clone(), + vec![Arc::new(Int32Array::from(vec![ + Some(0), + Some(1), + None, + Some(3), + Some(4), + Some(5), + None, + Some(7), + Some(8), + Some(9), + Some(10), + None, + ]))], + Some(null_buffer([ + true, true, false, true, true, true, true, false, true, true, true, true, + ])), + )) as ArrayRef; + Arc::new( + FixedSizeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Struct(fields), true)), + 2, + child, + Some(null_buffer([true, false, true, true, false, true])), + ) + .unwrap(), + ) + } + + fn list_fixed_size_list_struct() -> ArrayRef { + let struct_fields = Fields::from(vec![ArrowField::new("value", DataType::Int32, true)]); + let structs = Arc::new(StructArray::new( + struct_fields.clone(), + vec![Arc::new(Int32Array::from(vec![ + Some(0), + None, + Some(2), + Some(3), + None, + Some(5), + ]))], + Some(null_buffer([true, true, false, true, true, true])), + )) as ArrayRef; + let fixed_size_list = Arc::new( + FixedSizeListArray::try_new( + Arc::new(ArrowField::new( + "item", + DataType::Struct(struct_fields), + true, + )), + 2, + structs, + Some(null_buffer([true, false, true])), + ) + .unwrap(), + ) as ArrayRef; + Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new( + "item", + fixed_size_list.data_type().clone(), + true, + )), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1, 1, 3, 3])), + fixed_size_list, + Some(null_buffer([true, false, true, true])), + ) + .unwrap(), + ) + } + + fn nullable_struct() -> ArrayRef { + let fields = Fields::from(vec![ArrowField::new("value", DataType::Int32, true)]); + Arc::new(StructArray::new( + fields, + vec![Arc::new(Int32Array::from(vec![ + Some(10), + Some(20), + None, + Some(40), + Some(50), + ]))], + Some(null_buffer([true, false, true, true, false])), + )) + } + + fn deeply_nested() -> ArrayRef { + let leaf = Arc::new(Int32Array::from(vec![ + Some(0), + None, + Some(2), + Some(3), + None, + Some(5), + Some(6), + Some(7), + ])) as ArrayRef; + let inner = Arc::new( + LargeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i64, 2, 2, 3, 5, 5, 8])), + leaf, + Some(null_buffer([true, false, true, true, true, true])), + ) + .unwrap(), + ) as ArrayRef; + let struct_fields = Fields::from(vec![ArrowField::new( + "inner", + inner.data_type().clone(), + true, + )]); + let structs = Arc::new(StructArray::new( + struct_fields.clone(), + vec![inner], + Some(null_buffer([true, true, false, true, true, true])), + )) as ArrayRef; + Arc::new( + ListArray::try_new( + Arc::new(ArrowField::new( + "item", + DataType::Struct(struct_fields), + true, + )), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 2, 2, 4, 6])), + structs, + Some(null_buffer([true, false, true, true, true])), + ) + .unwrap(), + ) + } + + fn page_layout(page: &EncodedPage) -> &pb21::page_layout::Layout { + let PageEncoding::Structural(layout) = &page.description else { + panic!("expected structural page encoding"); + }; + layout.layout.as_ref().expect("page layout must be present") + } + + fn create_encoder( + array: &ArrayRef, + version: TestEncoding, + metadata: HashMap, + ) -> Result> { + let arrow_field = + ArrowField::new("values", array.data_type().clone(), true).with_metadata(metadata); + let field = Field::try_from(&arrow_field)?; + let strategy = test_encoding_strategy(version); + let options = EncodingOptions { + cache_bytes_per_column: 1, + ..Default::default() + }; + crate::testing::create_test_field_encoder( + strategy.as_ref(), + &field, + &mut ColumnIndexSequence::default(), + &options, + ) + } + + async fn encode_pages( + array: ArrayRef, + version: TestEncoding, + metadata: HashMap, + ) -> Result> { + encode_chunks(vec![array], version, metadata).await + } + + async fn encode_chunks( + arrays: Vec, + version: TestEncoding, + metadata: HashMap, + ) -> Result> { + let first = arrays + .first() + .ok_or_else(|| Error::internal("test input has no arrays".to_string()))?; + let mut encoder = create_encoder(first, version, metadata)?; + let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT); + let mut pages = Vec::new(); + let mut row_number = 0_u64; + for array in arrays { + let num_rows = array.len() as u64; + for task in encoder.maybe_encode( + array, + &mut external_buffers, + crate::repdef::RepDefBuilder::default(), + row_number, + num_rows, + )? { + pages.push(task.await?); + } + row_number += num_rows; + } + for task in encoder.flush(&mut external_buffers)? { + pages.push(task.await?); + } + for column in encoder.finish(&mut external_buffers).await? { + pages.extend(column.final_pages); + } + Ok(pages) + } + + fn sparse_layout(page: &EncodedPage) -> &pb21::SparseLayout { + let pb21::page_layout::Layout::SparseLayout(sparse) = page_layout(page) else { + panic!("expected SparseLayout, got {:?}", page_layout(page)); + }; + sparse + } + + fn list_layer(sparse: &pb21::SparseLayout) -> &pb21::SparseListLayer { + sparse + .structural_layers + .iter() + .find_map(|layer| match layer.layer.as_ref() { + Some(pb21::sparse_structural_layer::Layer::List(layer)) => Some(layer), + _ => None, + }) + .expect("expected sparse list layer") + } + + fn validity_layer(layer: &pb21::SparseStructuralLayer) -> Option<&pb21::SparseValidityLayer> { + match layer.layer.as_ref() { + Some(pb21::sparse_structural_layer::Layer::Validity(layer)) => Some(layer), + _ => None, + } + } + + fn layer_num_slots(layer: &pb21::SparseStructuralLayer) -> u64 { + match layer.layer.as_ref().expect("expected sparse layer variant") { + pb21::sparse_structural_layer::Layer::Validity(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::List(layer) => layer.num_slots, + pb21::sparse_structural_layer::Layer::FixedSizeList(layer) => layer.num_slots, + } + } + + fn fixed_size_list_dimension(layer: &pb21::SparseStructuralLayer) -> Option { + match layer.layer.as_ref() { + Some(pb21::sparse_structural_layer::Layer::FixedSizeList(layer)) => { + Some(layer.dimension) + } + _ => None, + } + } + + #[test] + fn test_sparse_value_block_eligibility() { + let fixed = DataBlock::from_array(Int32Array::from_iter_values(0..4)); + assert!(supports_value_block(&fixed)); + + let variable = DataBlock::from_array(StringArray::from(vec!["a", "b"])); + assert!(supports_value_block(&variable)); + + let nullable = DataBlock::from_array(Int32Array::from(vec![Some(1), None])); + assert!(!supports_value_block(&nullable)); + + let all_null = DataBlock::from_array(Int32Array::from(vec![None, None])); + assert!(!supports_value_block(&all_null)); + + let (dictionary, _) = dictionary_values(4); + assert!(!supports_value_block(&DataBlock::from_arrays( + &[dictionary], + 4 + ))); + + let fixed_fields = Fields::from(vec![ArrowField::new("value", DataType::Int32, false)]); + let fixed_struct = StructArray::new( + fixed_fields, + vec![Arc::new(Int32Array::from_iter_values(0..4))], + None, + ); + assert!(supports_value_block(&DataBlock::from_array(fixed_struct))); + + let variable_fields = Fields::from(vec![ArrowField::new("value", DataType::Utf8, false)]); + let variable_struct = StructArray::new( + variable_fields, + vec![Arc::new(StringArray::from(vec!["a", "b"]))], + None, + ); + assert!(!supports_value_block(&DataBlock::from_array( + variable_struct + ))); + + let fixed_size_list = FixedSizeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, false)), + 2, + Arc::new(Int32Array::from_iter_values(0..4)), + None, + ) + .unwrap(); + assert!(supports_value_block(&DataBlock::from_array( + fixed_size_list + ))); + + let unsupported_fixed_size_list = DataBlock::FixedSizeList(FixedSizeListBlock { + child: Box::new(DataBlock::from_array(StringArray::from(vec![ + "a", "b", "c", "d", + ]))), + dimension: 2, + }); + assert!(!supports_value_block(&unsupported_fixed_size_list)); + } + + fn planned_list( + offsets: Vec, + list_validity: Option>, + leaf_validity: Option>, + ) -> SparseStructuralPlan { + let num_values = u64::try_from(*offsets.last().unwrap()).unwrap(); + let mut builder = crate::repdef::RepDefBuilder::default(); + assert!(!builder.add_offsets( + OffsetBuffer::new(ScalarBuffer::from(offsets)), + list_validity.map(null_buffer), + )); + if let Some(leaf_validity) = leaf_validity { + builder.add_validity_bitmap(null_buffer(leaf_validity)); + } else { + builder.add_no_null(num_values as usize); + } + let normalized = crate::repdef::RepDefBuilder::normalize(vec![builder]); + plan(&normalized, num_values).unwrap() + } + + fn planned_list_layer(plan: &SparseStructuralPlan) -> &SparseStructuralLayerPlan { + plan.layers + .iter() + .find(|layer| matches!(layer, SparseStructuralLayerPlan::List { .. })) + .expect("expected planned list layer") + } + + #[test] + fn test_semantic_position_and_count_forms() { + let empty = planned_list(vec![0, 0, 0, 0], None, None); + assert_eq!(empty.num_items, 3); + assert_eq!(empty.num_visible_items, 0); + assert!(matches!( + planned_list_layer(&empty), + SparseStructuralLayerPlan::List { + non_empty_positions: SparsePositionSet::Empty, + counts: SparseCountSet::Empty, + .. + } + )); + + let all = planned_list(vec![0, 2, 4, 6], None, None); + assert_eq!(all.num_items, 6); + assert!(matches!( + planned_list_layer(&all), + SparseStructuralLayerPlan::List { + non_empty_positions: SparsePositionSet::All { len: 3 }, + counts: SparseCountSet::Constant { value: 2, len: 3 }, + .. + } + )); + + let range = planned_list(vec![0, 2, 4, 4, 4], None, None); + assert_eq!(range.num_items, 6); + assert!(matches!( + planned_list_layer(&range), + SparseStructuralLayerPlan::List { + non_empty_positions: SparsePositionSet::Range { start: 0, len: 2 }, + counts: SparseCountSet::Constant { value: 2, len: 2 }, + .. + } + )); + + let explicit = planned_list(vec![0, 1, 1, 4, 4, 6], None, None); + assert_eq!(explicit.num_items, 8); + assert!(matches!( + planned_list_layer(&explicit), + SparseStructuralLayerPlan::List { + non_empty_positions: SparsePositionSet::Explicit(positions), + counts: SparseCountSet::Explicit { counts, .. }, + .. + } if positions == &vec![0, 2, 4] && counts.as_ref() == [1, 3, 2] + )); + } + + #[test] + fn test_validity_polarity_uses_semantic_encoded_cost() { + let mostly_valid = BooleanBuffer::from_iter([true, false, true, true, false, true]); + let validity = validity_set(Some(&mostly_valid), mostly_valid.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::NullPositions); + assert!(matches!(validity.positions, SparsePositionSet::Explicit(ref p) if p == &[1, 4])); + + let mostly_null = BooleanBuffer::from_iter([false, true, false, false, true, false]); + let validity = validity_set(Some(&mostly_null), mostly_null.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::ValidPositions); + assert!(matches!(validity.positions, SparsePositionSet::Explicit(ref p) if p == &[1, 4])); + + let valid_island = BooleanBuffer::from_iter([false, false, true, true, false]); + let validity = validity_set(Some(&valid_island), valid_island.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::ValidPositions); + assert!(matches!( + validity.positions, + SparsePositionSet::Range { start: 2, len: 2 } + )); + + let all_valid = BooleanBuffer::from_iter([true, true, true]); + let validity = validity_set(Some(&all_valid), all_valid.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::NullPositions); + assert!(matches!(validity.positions, SparsePositionSet::Empty)); + + let all_null = BooleanBuffer::from_iter([false, false, false]); + let validity = validity_set(Some(&all_null), all_null.len(), "test").unwrap(); + assert_eq!(validity.meaning, SparseValidityMeaning::ValidPositions); + assert!(matches!(validity.positions, SparsePositionSet::Empty)); + } + + #[tokio::test] + async fn test_explicit_sparse_nullable_primitive_roundtrip() { + let array = Arc::new(Int32Array::from(vec![ + Some(10), + None, + Some(20), + Some(30), + None, + Some(40), + ])) as ArrayRef; + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert_eq!(pages.len(), 1); + let sparse = sparse_layout(&pages[0]); + assert_eq!(sparse.num_items, 6); + assert_eq!(sparse.num_visible_items, 6); + assert_eq!(sparse.structural_layers.len(), 1); + let validity = validity_layer(&sparse.structural_layers[0]) + .unwrap() + .validity + .as_ref() + .unwrap(); + assert_eq!( + validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions as i32 + ); + assert!(matches!( + validity.positions.as_ref().unwrap().positions, + Some(pb21::sparse_position_set::Positions::Explicit(_)) + )); + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_range(1..5) + .with_indices(vec![0, 2, 5]); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_nullable_struct_roundtrip() { + let array = nullable_struct(); + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert_eq!(pages.len(), 1); + let sparse = sparse_layout(&pages[0]); + assert_eq!(sparse.structural_layers.len(), 2); + assert!( + sparse + .structural_layers + .iter() + .all(|layer| validity_layer(layer).is_some()) + ); + let struct_validity = validity_layer(&sparse.structural_layers[0]) + .unwrap() + .validity + .as_ref() + .unwrap(); + assert_eq!( + struct_validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions as i32 + ); + assert!(matches!( + struct_validity.positions.as_ref().unwrap().positions, + Some(pb21::sparse_position_set::Positions::Explicit(_)) + )); + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_range(1..5) + .with_indices(vec![0, 2, 4]); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_struct_with_constant_and_sparse_children() { + let fields = Fields::from(vec![ + ArrowField::new("constant", DataType::Int32, true), + ArrowField::new("sparse", DataType::Int32, true), + ]); + let array = Arc::new(StructArray::new( + fields, + vec![ + Arc::new(Int32Array::from(vec![None::; 5])), + Arc::new(Int32Array::from(vec![ + Some(10), + Some(20), + None, + Some(40), + Some(50), + ])), + ], + Some(null_buffer([true, false, true, true, false])), + )) as ArrayRef; + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert!(pages.iter().any(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::ConstantLayout(_) + ))); + assert!(pages.iter().any(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::SparseLayout(_) + ))); + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_range(1..5) + .with_indices(vec![0, 2, 4]); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_emits_both_validity_polarities() { + let mostly_valid = Arc::new(Int32Array::from(vec![ + Some(0), + None, + Some(2), + Some(3), + None, + Some(5), + ])) as ArrayRef; + let mostly_null = Arc::new(Int32Array::from(vec![ + None, + Some(1), + None, + None, + Some(4), + None, + ])) as ArrayRef; + + let null_positions = encode_pages( + mostly_valid.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + let validity = validity_layer(&sparse_layout(&null_positions[0]).structural_layers[0]) + .unwrap() + .validity + .as_ref() + .unwrap(); + assert_eq!( + validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityNullPositions as i32 + ); + + let valid_positions = encode_pages( + mostly_null.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + let validity = validity_layer(&sparse_layout(&valid_positions[0]).structural_layers[0]) + .unwrap() + .validity + .as_ref() + .unwrap(); + assert_eq!( + validity.meaning, + pb21::sparse_validity_set::Meaning::SparseValidityValidPositions as i32 + ); + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_range(1..5) + .with_indices(vec![0, 2, 5]); + for array in [mostly_valid, mostly_null] { + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + } + + #[tokio::test] + async fn test_explicit_sparse_nested_page_boundaries_range_and_take() { + let nested = deeply_nested(); + let chunks = vec![nested.slice(0, 2), nested.slice(2, 3)]; + let pages = encode_chunks( + chunks.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert!(pages.len() >= 2, "expected multiple sparse pages"); + let sparse = pages + .iter() + .map(sparse_layout) + .collect::>(); + assert!(sparse.iter().any(|layout| { + layout + .structural_layers + .iter() + .filter(|layer| { + matches!( + layer.layer.as_ref(), + Some(pb21::sparse_structural_layer::Layer::List(_)) + ) + }) + .count() + >= 2 + })); + assert!(sparse.iter().any(|layout| { + layout + .structural_layers + .iter() + .any(|layer| validity_layer(layer).is_some()) + })); + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_batch_size(2) + .with_range(1..5) + .with_range(2..4) + .with_indices(vec![0, 2, 4]) + .with_indices(vec![1, 3]); + check_round_trip_encoding_of_data(chunks, &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_list_and_large_list_null_empty_roundtrip() { + let list = list_i32( + vec![0, 0, 0, 2, 3, 3], + Some(vec![false, true, true, true, true]), + ); + let large_list = large_list_i32( + vec![0, 0, 0, 2, 3, 3], + Some(vec![false, true, true, true, true]), + ); + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_range(0..5) + .with_range(1..4) + .with_indices(vec![0, 1, 4]) + .with_indices(vec![2, 3]); + + for array in [list, large_list] { + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert!(pages.iter().map(sparse_layout).all(|layout| { + layout.structural_layers.iter().any(|layer| { + matches!( + layer.layer.as_ref(), + Some(pb21::sparse_structural_layer::Layer::List(_)) + ) + }) + })); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + } + + #[tokio::test] + async fn test_explicit_sparse_map_and_fixed_size_list_roundtrip() { + let map = map_i32(); + let map_pages = encode_pages( + map.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert!(map_pages.iter().map(sparse_layout).any(|layout| { + layout.structural_layers.iter().any(|layer| { + matches!( + layer.layer.as_ref(), + Some(pb21::sparse_structural_layer::Layer::List(_)) + ) + }) + })); + let map_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_range(1..4) + .with_indices(vec![0, 2, 3]); + check_round_trip_encoding_of_data(vec![map], &map_cases, sparse_metadata()).await; + + let fsl = fixed_size_list_struct(); + let fsl_pages = encode_pages( + fsl.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + for page in &fsl_pages { + let layout = sparse_layout(page); + let outer_slots = layer_num_slots(layout.structural_layers.first().unwrap()); + let fixed_size_scale = layout + .structural_layers + .iter() + .filter_map(fixed_size_list_dimension) + .product::(); + assert_eq!(page.num_rows, outer_slots * fixed_size_scale); + } + assert!(fsl_pages.iter().map(sparse_layout).any(|layout| { + layout + .structural_layers + .iter() + .any(|layer| fixed_size_list_dimension(layer) == Some(2)) + })); + let fsl_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_range(1..6) + .with_indices(vec![0, 3, 5]); + check_round_trip_encoding_of_data(vec![fsl], &fsl_cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_list_fixed_size_list_struct_roundtrip() { + let array = list_fixed_size_list_struct(); + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert!(pages.iter().map(sparse_layout).any(|layout| { + let kinds = layout + .structural_layers + .iter() + .map(|layer| match layer.layer.as_ref().unwrap() { + pb21::sparse_structural_layer::Layer::Validity(_) => "validity", + pb21::sparse_structural_layer::Layer::List(_) => "list", + pb21::sparse_structural_layer::Layer::FixedSizeList(_) => "fixed-size-list", + }) + .collect::>(); + kinds.starts_with(&["list", "fixed-size-list"]) + })); + for page in &pages { + let layout = sparse_layout(page); + let outer_slots = layer_num_slots(layout.structural_layers.first().unwrap()); + assert_eq!(page.num_rows, outer_slots * 2); + } + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_range(1..4) + .with_indices(vec![0, 2, 3]) + .with_indices(vec![1, 3]); + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + + #[tokio::test] + async fn test_explicit_sparse_serializes_semantic_list_forms() { + let all_array = list_i32(vec![0, 2, 4, 6], None); + let all = encode_pages( + all_array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + let all_layer = list_layer(sparse_layout(&all[0])); + assert!(matches!( + all_layer.non_empty_positions.as_ref().unwrap().positions, + Some(pb21::sparse_position_set::Positions::All(_)) + )); + assert!(matches!( + all_layer.counts.as_ref().unwrap().counts, + Some(pb21::sparse_count_set::Counts::Constant( + pb21::SparseCountConstant { value: 2 } + )) + )); + assert_eq!(all[0].data.len(), 2); + + let range_array = list_i32(vec![0, 2, 4, 4, 4], None); + let range = encode_pages( + range_array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + let range_layer = list_layer(sparse_layout(&range[0])); + assert!(matches!( + range_layer.non_empty_positions.as_ref().unwrap().positions, + Some(pb21::sparse_position_set::Positions::Range( + pb21::SparsePositionRange { + start: 0, + length: 2 + } + )) + )); + assert!(matches!( + range_layer.counts.as_ref().unwrap().counts, + Some(pb21::sparse_count_set::Counts::Constant( + pb21::SparseCountConstant { value: 2 } + )) + )); + assert_eq!(range[0].data.len(), 2); + + let explicit_array = list_i32(vec![0, 1, 1, 4, 4, 6], None); + let explicit = encode_pages( + explicit_array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + let explicit_layer = list_layer(sparse_layout(&explicit[0])); + assert!(matches!( + explicit_layer + .non_empty_positions + .as_ref() + .unwrap() + .positions, + Some(pb21::sparse_position_set::Positions::Explicit(_)) + )); + assert!(matches!( + explicit_layer.counts.as_ref().unwrap().counts, + Some(pb21::sparse_count_set::Counts::Explicit(_)) + )); + assert_eq!(explicit[0].data.len(), 4); + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1]) + .with_range(1..3) + .with_indices(vec![0, 2]); + for array in [all_array, range_array, explicit_array] { + check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; + } + } + + #[tokio::test] + async fn test_constant_layout_boundary_is_explicit() { + let structural_only = encode_pages( + list_i32(vec![0, 0, 0, 0], None), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&structural_only[0]), + pb21::page_layout::Layout::ConstantLayout(_) + )); + + let empty_struct = Arc::new(StructArray::new_empty_fields(3, None)) as ArrayRef; + let empty_struct_pages = encode_pages( + empty_struct, + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&empty_struct_pages[0]), + pb21::page_layout::Layout::ConstantLayout(_) + )); + + let all_null = Arc::new(Int32Array::from(vec![None, None, None])) as ArrayRef; + let all_null_pages = + encode_pages(all_null, TestEncoding::StructuralSparse, sparse_metadata()) + .await + .unwrap(); + assert!(matches!( + page_layout(&all_null_pages[0]), + pb21::page_layout::Layout::ConstantLayout(_) + )); + + let constant = Arc::new(Int32Array::from(vec![7, 7, 7])) as ArrayRef; + let constant_pages = + encode_pages(constant, TestEncoding::StructuralSparse, sparse_metadata()) + .await + .unwrap(); + assert!(matches!( + page_layout(&constant_pages[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + } + + #[tokio::test] + async fn test_default_and_explicit_dense_layouts_are_unchanged() { + let array = Arc::new(Int32Array::from_iter_values(0..16)) as ArrayRef; + let v2_2_default = encode_pages(array.clone(), TestEncoding::StructuralU32, HashMap::new()) + .await + .unwrap(); + let v2_2_miniblock = encode_pages( + array.clone(), + TestEncoding::StructuralU32, + structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), + ) + .await + .unwrap(); + assert_eq!(v2_2_default.len(), v2_2_miniblock.len()); + for (default, explicit) in v2_2_default.iter().zip(v2_2_miniblock.iter()) { + assert!(matches!( + page_layout(default), + pb21::page_layout::Layout::MiniBlockLayout(_) + )); + assert_eq!(page_layout(default), page_layout(explicit)); + assert_eq!(default.data, explicit.data); + } + + let v2_3_default = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&v2_3_default[0]), + pb21::page_layout::Layout::MiniBlockLayout(_) + )); + + let v2_3_miniblock = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&v2_3_miniblock[0]), + pb21::page_layout::Layout::MiniBlockLayout(_) + )); + + let v2_3_fullzip = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + structural_metadata(STRUCTURAL_ENCODING_FULLZIP), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&v2_3_fullzip[0]), + pb21::page_layout::Layout::FullZipLayout(_) + )); + + let v2_3_sparse = encode_pages(array, TestEncoding::StructuralSparse, sparse_metadata()) + .await + .unwrap(); + assert!(matches!( + page_layout(&v2_3_sparse[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + } + + #[tokio::test] + async fn test_auto_sparse_for_split_required_page() { + const NUM_ROWS: usize = 70_000; + let array = sparse_i32_list(NUM_ROWS, 2_000); + + let within_budget = encode_pages( + sparse_i32_list(4_096, 1_024), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); + assert_eq!(within_budget.len(), 1); + assert!(matches!( + page_layout(&within_budget[0]), + pb21::page_layout::Layout::MiniBlockLayout(_) + )); + + let automatic = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); + assert_eq!(automatic.len(), 1); + assert!(matches!( + page_layout(&automatic[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + + let miniblock = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), + ) + .await + .unwrap(); + assert!(miniblock.len() > 1); + assert!(miniblock.iter().all(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::MiniBlockLayout(_) + ))); + + let fullzip = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + structural_metadata(STRUCTURAL_ENCODING_FULLZIP), + ) + .await + .unwrap(); + assert_eq!(fullzip.len(), miniblock.len()); + assert!(fullzip.iter().all(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::FullZipLayout(_) + ))); + + let sparse = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert_eq!(sparse.len(), 1); + assert!(matches!( + page_layout(&sparse[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + + let v2_2_default = encode_pages(array.clone(), TestEncoding::StructuralU32, HashMap::new()) + .await + .unwrap(); + let v2_2_miniblock = encode_pages( + array.clone(), + TestEncoding::StructuralU32, + structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), + ) + .await + .unwrap(); + assert_eq!(v2_2_default.len(), v2_2_miniblock.len()); + for (default, explicit) in v2_2_default.iter().zip(v2_2_miniblock.iter()) { + assert_eq!(page_layout(default), page_layout(explicit)); + assert_eq!(default.data, explicit.data); + } + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1024 * 1024]) + .with_batch_size(NUM_ROWS as u32) + .with_range(1_999..2_002) + .with_indices(vec![0, 1_999, 2_000, NUM_ROWS as u64 - 1]); + check_round_trip_encoding_of_data(vec![array], &cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_auto_sparse_for_single_row_over_budget() { + let array = unsplittable_nested_list( + Arc::new(Int32Array::from(vec![42, 43])), + Arc::new(ArrowField::new("item", DataType::Int32, true)), + ); + + let automatic = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); + assert_eq!(automatic.len(), 1); + assert!(matches!( + page_layout(&automatic[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + + let v2_2 = encode_pages(array.clone(), TestEncoding::StructuralU32, HashMap::new()) + .await + .unwrap(); + assert_eq!(v2_2.len(), 1); + assert!(matches!( + page_layout(&v2_2[0]), + pb21::page_layout::Layout::FullZipLayout(_) + )); + + let miniblock_error = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), + ) + .await + .unwrap_err(); + assert!( + miniblock_error + .to_string() + .contains("Mini-block cannot encode 70000 rep/def levels") + ); + + let fullzip = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + structural_metadata(STRUCTURAL_ENCODING_FULLZIP), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&fullzip[0]), + pb21::page_layout::Layout::FullZipLayout(_) + )); + + let explicit_sparse = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); + assert!(matches!( + page_layout(&explicit_sparse[0]), + pb21::page_layout::Layout::SparseLayout(_) + )); + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_range(0..1) + .with_indices(vec![0]); + check_round_trip_encoding_of_data(vec![array], &cases, HashMap::new()).await; + } + + #[tokio::test] + async fn test_auto_sparse_keeps_unsupported_values_dense() { + const NUM_ROWS: usize = 70_000; + let num_values = NUM_ROWS.div_ceil(2_000); + + let (dictionary, dictionary_field) = dictionary_values(num_values); + let dictionary_array = sparse_list_values(NUM_ROWS, 2_000, dictionary, dictionary_field); + let dictionary_pages = encode_pages( + dictionary_array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); + assert!(dictionary_pages.len() > 1); + assert!(dictionary_pages.iter().all(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::MiniBlockLayout(_) + ))); + + let (packed_values, packed_field) = variable_packed_struct_values(num_values); + let packed_array = sparse_list_values(NUM_ROWS, 2_000, packed_values, packed_field); + let packed_pages = encode_pages( + packed_array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); + assert!(packed_pages.len() > 1); + assert!(packed_pages.iter().all(|page| matches!( + page_layout(page), + pb21::page_layout::Layout::FullZipLayout(_) + ))); + + let (packed_values, packed_field) = variable_packed_struct_values(2); + let unsplittable_packed = unsplittable_nested_list(packed_values, packed_field); + let packed_fallback = encode_pages( + unsplittable_packed.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); + assert_eq!(packed_fallback.len(), 1); + assert!(matches!( + page_layout(&packed_fallback[0]), + pb21::page_layout::Layout::FullZipLayout(_) + )); + + let (dictionary, dictionary_field) = dictionary_values(2); + let unsplittable_dictionary = unsplittable_nested_list(dictionary, dictionary_field); + let v2_2_error = encode_pages( + unsplittable_dictionary.clone(), + TestEncoding::StructuralU32, + HashMap::new(), + ) + .await + .unwrap_err(); + let v2_3_error = encode_pages( + unsplittable_dictionary, + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap_err(); + assert_eq!(v2_3_error.to_string(), v2_2_error.to_string()); + assert!( + v2_3_error + .to_string() + .contains("Mini-block cannot encode 70000 rep/def levels") + ); + + let cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1024 * 1024]) + .with_batch_size(NUM_ROWS as u32) + .with_range(1_999..2_002) + .with_indices(vec![0, 2_000, NUM_ROWS as u64 - 1]); + check_round_trip_encoding_of_data(vec![dictionary_array], &cases, HashMap::new()).await; + let packed_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1024 * 1024]) + .with_batch_size(NUM_ROWS as u32); + check_round_trip_encoding_of_data(vec![packed_array], &packed_cases, HashMap::new()).await; + + let single_row_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralSparse) + .with_page_sizes(vec![1024 * 1024]) + .with_range(0..1) + .with_indices(vec![0]); + check_round_trip_encoding_of_data( + vec![unsplittable_packed], + &single_row_cases, + HashMap::new(), + ) + .await; + } + + #[tokio::test] + async fn test_auto_sparse_wide_values_keep_dense_fallback() { + let first = vec![0xAB_u8; 5_000]; + let second = vec![0xCD_u8; 5_000]; + let fixed_size_binary = Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [Some(first.as_slice()), Some(second.as_slice())].into_iter(), + 5_000, + ) + .unwrap(), + ) as ArrayRef; + + const FSL_DIMENSION: i32 = 2_048; + let fixed_size_list = Arc::new( + FixedSizeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Int32, true)), + FSL_DIMENSION, + Arc::new(Int32Array::from_iter_values(0..(FSL_DIMENSION * 2))), + None, + ) + .unwrap(), + ) as ArrayRef; + + for (label, values) in [ + ("fixed-size binary", fixed_size_binary), + ("fixed-size list", fixed_size_list), + ] { + let item_field = Arc::new(ArrowField::new("item", values.data_type().clone(), true)); + let array = unsplittable_nested_list(values, item_field); + + let v2_2 = encode_pages(array.clone(), TestEncoding::StructuralU32, HashMap::new()) + .await + .unwrap(); + let v2_3 = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); + for pages in [&v2_2, &v2_3] { + assert_eq!(pages.len(), 1, "unexpected {label} page count"); + assert!( + matches!( + page_layout(&pages[0]), + pb21::page_layout::Layout::FullZipLayout(_) + ), + "{label} should retain the dense full-zip fallback" + ); + } + + let explicit_error = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap_err(); + assert!( + explicit_error.to_string().contains("too wide"), + "explicit sparse should preserve the {label} value error: {explicit_error}" + ); + + let cases = TestCases::default() + .with_u32_structural_encodings() + .with_range(0..1) + .with_indices(vec![0]); + check_round_trip_encoding_of_data(vec![array], &cases, HashMap::new()).await; + } + } + + #[test] + fn test_explicit_sparse_rejects_lance_2_2() { + let array = Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef; + let Err(error) = create_encoder(&array, TestEncoding::StructuralU32, sparse_metadata()) + else { + panic!("expected Lance 2.2 to reject explicit sparse encoding"); + }; + assert!( + error + .to_string() + .contains("not enabled by the selected file format") + ); + + let structural_only = list_i32(vec![0, 0, 0], None); + let Err(error) = create_encoder( + &structural_only, + TestEncoding::StructuralU32, + sparse_metadata(), + ) else { + panic!("expected Lance 2.2 structural-only input to reject explicit sparse encoding"); + }; + assert!( + error + .to_string() + .contains("not enabled by the selected file format") + ); + } +} diff --git a/vendor/lance-encoding/src/encodings/logical/struct.rs b/vendor/lance-encoding/src/encodings/logical/struct.rs new file mode 100644 index 000000000..3f27ae6ea --- /dev/null +++ b/vendor/lance-encoding/src/encodings/logical/struct.rs @@ -0,0 +1,1180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ + collections::{BinaryHeap, VecDeque}, + ops::Range, + sync::Arc, +}; + +use super::{ + fixed_size_list::StructuralFixedSizeListDecoder, list::StructuralListDecoder, + map::StructuralMapDecoder, primitive::StructuralPrimitiveFieldDecoder, +}; +use crate::{ + decoder::{ + DecodedArray, FilterExpression, LoadedPageShard, NextDecodeTask, PageEncoding, + ScheduledScanLine, SchedulerContext, StructuralDecodeArrayTask, StructuralFieldDecoder, + StructuralFieldScheduler, StructuralSchedulingJob, + }, + encoder::{EncodeTask, EncodedColumn, EncodedPage, FieldEncoder, OutOfLineBuffers}, + format::pb, + repdef::{CompositeRepDefUnraveler, RepDefBuilder}, +}; +use arrow_array::{Array, ArrayRef, StructArray, cast::AsArray}; +use arrow_schema::{DataType, Fields}; +use futures::{ + FutureExt, StreamExt, TryStreamExt, + future::BoxFuture, + stream::{FuturesOrdered, FuturesUnordered}, +}; +use itertools::Itertools; +use lance_arrow::FieldExt; +use lance_arrow::{deepcopy::deep_copy_nulls, r#struct::StructArrayExt}; +use lance_core::{Error, Result, datatypes::validate_fixed_size_list_dimensions}; +use log::trace; + +#[derive(Debug)] +struct StructuralSchedulingJobWithStatus<'a> { + col_idx: u32, + col_name: &'a str, + job: Box, + rows_scheduled: u64, + rows_remaining: u64, + ready_scan_lines: VecDeque, +} + +impl PartialEq for StructuralSchedulingJobWithStatus<'_> { + fn eq(&self, other: &Self) -> bool { + self.col_idx == other.col_idx + } +} + +impl Eq for StructuralSchedulingJobWithStatus<'_> {} + +impl PartialOrd for StructuralSchedulingJobWithStatus<'_> { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for StructuralSchedulingJobWithStatus<'_> { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + // Note this is reversed to make it min-heap + other.rows_scheduled.cmp(&self.rows_scheduled) + } +} + +/// Scheduling job for struct data +/// +/// The order in which we schedule the children is important. We want to schedule the child +/// with the least amount of data first. +/// +/// This allows us to decode entire rows as quickly as possible +#[derive(Debug)] +struct RepDefStructSchedulingJob<'a> { + /// A min-heap whose key is the # of rows currently scheduled + children: BinaryHeap>, + rows_scheduled: u64, + num_rows: u64, +} + +impl<'a> RepDefStructSchedulingJob<'a> { + fn new( + scheduler: &'a StructuralStructScheduler, + children: Vec>, + num_rows: u64, + ) -> Self { + let children = children + .into_iter() + .enumerate() + .map(|(idx, job)| StructuralSchedulingJobWithStatus { + col_idx: idx as u32, + col_name: scheduler.child_fields[idx].name(), + job, + rows_scheduled: 0, + rows_remaining: num_rows, + ready_scan_lines: VecDeque::new(), + }) + .collect::>(); + Self { + children, + rows_scheduled: 0, + num_rows, + } + } +} + +impl StructuralSchedulingJob for RepDefStructSchedulingJob<'_> { + fn schedule_next( + &mut self, + mut context: &mut SchedulerContext, + ) -> Result> { + if self.children.is_empty() { + // Special path for empty structs + if self.rows_scheduled == self.num_rows { + return Ok(Vec::new()); + } + self.rows_scheduled = self.num_rows; + return Ok(vec![ScheduledScanLine { + decoders: Vec::new(), + rows_scheduled: self.num_rows, + }]); + } + + let mut decoders = Vec::new(); + let old_rows_scheduled = self.rows_scheduled; + // Schedule as many children as we need to until we have scheduled at least one + // complete row + while old_rows_scheduled == self.rows_scheduled { + if self.children.is_empty() { + // Early exit when schedulers are exhausted prematurely (TODO: does this still happen?) + return Ok(Vec::new()); + } + let mut next_child = self.children.pop().unwrap(); + if next_child.ready_scan_lines.is_empty() { + let scoped = context.push(next_child.col_name, next_child.col_idx); + let child_scans = next_child.job.schedule_next(scoped.context)?; + context = scoped.pop(); + if child_scans.is_empty() { + // Continue without pushing next_child back onto the heap (it is done) + continue; + } + next_child.ready_scan_lines.extend(child_scans); + } + let child_scan = next_child.ready_scan_lines.pop_front().unwrap(); + trace!( + "Scheduled {} rows for child {}", + child_scan.rows_scheduled, next_child.col_idx + ); + next_child.rows_scheduled += child_scan.rows_scheduled; + next_child.rows_remaining -= child_scan.rows_scheduled; + decoders.extend(child_scan.decoders); + self.children.push(next_child); + self.rows_scheduled = self.children.peek().unwrap().rows_scheduled; + } + let struct_rows_scheduled = self.rows_scheduled - old_rows_scheduled; + Ok(vec![ScheduledScanLine { + decoders, + rows_scheduled: struct_rows_scheduled, + }]) + } +} + +/// A scheduler for structs +/// +/// The implementation is actually a bit more tricky than one might initially think. We can't just +/// go through and schedule each column one after the other. This would mean our decode can't start +/// until nearly all the data has arrived (since we need data from each column to yield a batch) +/// +/// Instead, we schedule in row-major fashion +/// +/// Note: this scheduler is the starting point for all decoding. This is because we treat the top-level +/// record batch as a non-nullable struct. +#[derive(Debug)] +pub struct StructuralStructScheduler { + children: Vec>, + child_fields: Fields, +} + +impl StructuralStructScheduler { + pub fn new(children: Vec>, child_fields: Fields) -> Self { + Self { + children, + child_fields, + } + } +} + +impl StructuralFieldScheduler for StructuralStructScheduler { + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result> { + let num_rows = ranges.iter().map(|r| r.end - r.start).sum(); + + let child_schedulers = self + .children + .iter() + .map(|child| child.schedule_ranges(ranges, filter)) + .collect::>>()?; + + Ok(Box::new(RepDefStructSchedulingJob::new( + self, + child_schedulers, + num_rows, + ))) + } + + fn initialize<'a>( + &'a mut self, + requested_ranges: Option<&'a [Range]>, + filter: &'a FilterExpression, + context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>> { + // Struct children share the parent's row coordinates 1:1 (see + // `schedule_ranges`), so every child gets the same ranges. + let children_initialization = self + .children + .iter_mut() + .map(|child| child.initialize(requested_ranges, filter, context)) + .collect::>(); + async move { + children_initialization + .map(|res| res.map(|_| ())) + .try_collect::>() + .await?; + Ok(()) + } + .boxed() + } +} + +#[derive(Debug)] +pub struct StructuralStructDecoder { + children: Vec>, + data_type: DataType, + child_fields: Fields, + // The root decoder is slightly different because it cannot have nulls + is_root: bool, + nullable: bool, +} + +impl StructuralStructDecoder { + pub fn new( + fields: Fields, + should_validate: bool, + is_root: bool, + nullable: bool, + ) -> Result { + let children = fields + .iter() + .map(|field| Self::field_to_decoder(field, should_validate)) + .collect::>>()?; + let data_type = DataType::Struct(fields.clone()); + Ok(Self { + data_type, + children, + child_fields: fields, + is_root, + nullable: !is_root && nullable, + }) + } + + fn field_to_decoder( + field: &Arc, + should_validate: bool, + ) -> Result> { + match field.data_type() { + DataType::Struct(fields) => { + if field.is_packed_struct() || field.is_blob() { + let decoder = + StructuralPrimitiveFieldDecoder::new(&field.clone(), should_validate); + Ok(Box::new(decoder)) + } else { + Ok(Box::new(Self::new( + fields.clone(), + should_validate, + false, + field.is_nullable(), + )?)) + } + } + DataType::List(child_field) | DataType::LargeList(child_field) => { + let child_decoder = Self::field_to_decoder(child_field, should_validate)?; + Ok(Box::new(StructuralListDecoder::new( + child_decoder, + field.data_type().clone(), + ))) + } + DataType::FixedSizeList(child_field, _) + if matches!(child_field.data_type(), DataType::Struct(_)) => + { + // The scheduler factories run the same guard, but the decoder tree can be + // built independently (e.g. `create_decode_stream`) so a zero dimension from + // a malformed schema must be rejected here as well. Draining and unraveling + // validity both scale by the dimension and a zero would make that math + // degenerate. + validate_fixed_size_list_dimensions(field.name(), field.data_type())?; + // FixedSizeList containing Struct needs structural decoding + let child_decoder = Self::field_to_decoder(child_field, should_validate)?; + Ok(Box::new(StructuralFixedSizeListDecoder::new( + child_decoder, + field.data_type().clone(), + ))) + } + DataType::Map(entries_field, keys_sorted) => { + if *keys_sorted { + return Err(Error::not_supported_source( + "Map data type with keys_sorted=true is not supported yet" + .to_string() + .into(), + )); + } + let child_decoder = Self::field_to_decoder(entries_field, should_validate)?; + Ok(Box::new(StructuralMapDecoder::new( + child_decoder, + field.data_type().clone(), + ))) + } + DataType::RunEndEncoded(_, _) => todo!(), + DataType::ListView(_) | DataType::LargeListView(_) => todo!(), + DataType::Union(_, _) => todo!(), + _ => Ok(Box::new(StructuralPrimitiveFieldDecoder::new( + field, + should_validate, + ))), + } + } + + pub fn drain_batch_task(&mut self, num_rows: u64) -> Result { + let array_drain = self.drain(num_rows)?; + Ok(NextDecodeTask { + num_rows, + task: Box::new(array_drain), + }) + } +} + +impl StructuralFieldDecoder for StructuralStructDecoder { + fn accept_page(&mut self, mut child: LoadedPageShard) -> Result<()> { + // children with empty path should not be delivered to this method + let child_idx = child.path.pop_front().unwrap(); + // This decoder is intended for one of our children + self.children[child_idx as usize].accept_page(child)?; + Ok(()) + } + + fn drain(&mut self, num_rows: u64) -> Result> { + let child_tasks = self + .children + .iter_mut() + .map(|child| child.drain(num_rows)) + .collect::>>()?; + Ok(Box::new(RepDefStructDecodeTask { + children: child_tasks, + child_fields: self.child_fields.clone(), + is_root: self.is_root, + num_rows, + })) + } + + fn data_type(&self) -> &DataType { + &self.data_type + } + + fn plan_decoded_bytes(&self, rows_remaining: u64) -> lance_core::Result<[u64; 8]> { + use crate::decoder::CANDIDATE_BATCH_SIZES; + let mut out = [0u64; 8]; + for child in &self.children { + let child_bytes = child.plan_decoded_bytes(rows_remaining)?; + for (o, c) in out.iter_mut().zip(child_bytes) { + *o += c; + } + } + if self.nullable { + for (i, &candidate) in CANDIDATE_BATCH_SIZES.iter().enumerate() { + let rows = (candidate as u64).min(rows_remaining); + out[i] += rows.div_ceil(8); + } + } + Ok(out) + } +} + +#[derive(Debug)] +struct RepDefStructDecodeTask { + children: Vec>, + child_fields: Fields, + is_root: bool, + num_rows: u64, +} + +impl StructuralDecodeArrayTask for RepDefStructDecodeTask { + fn decode(self: Box) -> Result { + if self.children.is_empty() { + return Ok(DecodedArray { + array: Arc::new(StructArray::new_empty_fields(self.num_rows as usize, None)), + repdef: CompositeRepDefUnraveler::new(vec![]), + data_size: 0, + }); + } + + let arrays = self + .children + .into_iter() + .map(|task| task.decode()) + .collect::>>()?; + let mut children = Vec::with_capacity(arrays.len()); + let mut repdefs = Vec::with_capacity(arrays.len()); + let mut data_size = 0u64; + let mut arrays_iter = arrays.into_iter(); + let first_array = arrays_iter.next().ok_or_else(|| { + Error::internal("Struct decoder unexpectedly has no child arrays".to_string()) + })?; + let length = first_array.array.len(); + + // The repdef should be identical across all children at this point + repdefs.push(first_array.repdef); + data_size += first_array.data_size; + children.push(first_array.array); + + for array in arrays_iter { + if length != array.array.len() { + return Err(Error::invalid_input_source( + format!( + "Struct child array length {} does not match sibling length {}", + array.array.len(), + length + ) + .into(), + )); + } + data_size += array.data_size; + children.push(array.array); + repdefs.push(array.repdef); + } + + // Dense rep/def state can retain child-specific repetition information after a child + // decoder finishes, so comparing dense siblings is not meaningful. If any child carries + // sparse state, keep a sparse child as the canonical structural plan and compare it with + // every other sparse sibling so sparse metadata is never silently discarded. + let primary_repdef = repdefs + .iter() + .position(CompositeRepDefUnraveler::has_sparse) + .unwrap_or(0); + let mut repdef = repdefs.swap_remove(primary_repdef); + if repdef.has_sparse() { + for sibling in repdefs { + if sibling.has_sparse() { + repdef.add_compatibility_check(sibling); + } + } + } + + let validity = if self.is_root { + repdef.ensure_exhausted()?; + None + } else { + repdef.unravel_validity(length)? + }; + + let array = StructArray::try_new(self.child_fields, children, validity) + .map_err(|e| Error::invalid_input_source(e.to_string().into()))?; + Ok(DecodedArray { + array: Arc::new(array), + repdef, + data_size, + }) + } +} + +/// A structural encoder for struct fields +/// +/// The struct's validity is added to the rep/def builder +/// and the builder is cloned to all children. +pub struct StructStructuralEncoder { + keep_original_array: bool, + children: Vec>, +} + +impl StructStructuralEncoder { + pub fn new(keep_original_array: bool, children: Vec>) -> Self { + Self { + keep_original_array, + children, + } + } +} + +impl FieldEncoder for StructStructuralEncoder { + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + mut repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + let struct_array = array.as_struct(); + let mut struct_array = struct_array.normalize_slicing()?; + if let Some(validity) = struct_array.nulls() { + if self.keep_original_array { + repdef.add_validity_bitmap(validity.clone()) + } else { + repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap()) + } + struct_array = struct_array.pushdown_nulls()?; + } else { + repdef.add_no_null(struct_array.len()); + } + let child_tasks = self + .children + .iter_mut() + .zip(struct_array.columns().iter()) + .map(|(encoder, arr)| { + encoder.maybe_encode( + arr.clone(), + external_buffers, + repdef.clone(), + row_number, + num_rows, + ) + }) + .collect::>>()?; + Ok(child_tasks.into_iter().flatten().collect::>()) + } + + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + self.children + .iter_mut() + .map(|encoder| encoder.flush(external_buffers)) + .flatten_ok() + .collect::>>() + } + + fn num_columns(&self) -> u32 { + self.children + .iter() + .map(|child| child.num_columns()) + .sum::() + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + let mut child_columns = self + .children + .iter_mut() + .map(|child| child.finish(external_buffers)) + .collect::>(); + async move { + let mut encoded_columns = Vec::with_capacity(child_columns.len()); + while let Some(child_cols) = child_columns.next().await { + encoded_columns.extend(child_cols?); + } + Ok(encoded_columns) + } + .boxed() + } +} + +pub struct StructFieldEncoder { + children: Vec>, + column_index: u32, + num_rows_seen: u64, +} + +impl StructFieldEncoder { + pub fn new(children: Vec>, column_index: u32) -> Self { + Self { + children, + column_index, + num_rows_seen: 0, + } + } +} + +impl FieldEncoder for StructFieldEncoder { + fn maybe_encode( + &mut self, + array: ArrayRef, + external_buffers: &mut OutOfLineBuffers, + repdef: RepDefBuilder, + row_number: u64, + num_rows: u64, + ) -> Result> { + let struct_array = array.as_struct(); + self.num_rows_seen += array.len() as u64; + let child_tasks = self + .children + .iter_mut() + .zip(struct_array.columns().iter()) + .map(|(encoder, arr)| { + encoder.maybe_encode( + arr.clone(), + external_buffers, + repdef.clone(), + row_number, + num_rows, + ) + }) + .collect::>>()?; + Ok(child_tasks.into_iter().flatten().collect::>()) + } + + fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result> { + let child_tasks = self + .children + .iter_mut() + .map(|encoder| encoder.flush(external_buffers)) + .collect::>>()?; + Ok(child_tasks.into_iter().flatten().collect::>()) + } + + fn num_columns(&self) -> u32 { + self.children + .iter() + .map(|child| child.num_columns()) + .sum::() + + 1 + } + + fn finish( + &mut self, + external_buffers: &mut OutOfLineBuffers, + ) -> BoxFuture<'_, Result>> { + let mut child_columns = self + .children + .iter_mut() + .map(|child| child.finish(external_buffers)) + .collect::>(); + let num_rows_seen = self.num_rows_seen; + let column_index = self.column_index; + async move { + let mut columns = Vec::new(); + // Add a column for the struct header + let mut header = EncodedColumn::default(); + header.final_pages.push(EncodedPage { + data: Vec::new(), + description: PageEncoding::Legacy(pb::ArrayEncoding { + array_encoding: Some(pb::array_encoding::ArrayEncoding::Struct( + pb::SimpleStruct {}, + )), + }), + num_rows: num_rows_seen, + column_idx: column_index, + row_number: 0, // Not used by legacy encoding + }); + columns.push(header); + // Now run finish on the children + while let Some(child_cols) = child_columns.next().await { + columns.extend(child_cols?); + } + Ok(columns) + } + .boxed() + } +} + +#[cfg(test)] +mod tests { + + use std::{collections::HashMap, sync::Arc}; + + use arrow_array::{ + Array, ArrayRef, Float64Array, Int32Array, Int64Array, ListArray, StringArray, StructArray, + builder::{Int32Builder, ListBuilder}, + }; + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field, Fields}; + + use super::StructuralStructDecoder; + use crate::constants::{STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY}; + use crate::decoder::StructuralFieldDecoder; + use crate::testing::{ + TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data, + }; + + #[test] + fn test_zero_dimension_fsl_decoder_errors() { + // Simulates a stored schema declaring a zero-dimension FixedSizeList (writers reject + // it but old files may contain one). Building the decoder must fail cleanly instead + // of letting the zero dimension reach the rep/def decimation. + let item_fields = Fields::from(vec![Field::new("x", DataType::Int32, true)]); + let fields = Fields::from(vec![Field::new( + "vecs", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Struct(item_fields), true)), + 0, + ), + true, + )]); + + let err = StructuralStructDecoder::new( + fields, false, /*is_root=*/ true, /*nullable=*/ false, + ) + .unwrap_err(); + assert!(matches!(err, lance_core::Error::Schema { .. })); + assert!( + err.to_string() + .contains("dimension must be a positive integer"), + "unexpected error: {}", + err + ); + } + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_simple_struct( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let data_type = DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let field = Field::new("", data_type, false); + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[test_log::test(tokio::test)] + async fn test_nullable_struct() { + // Test data struct> + // - score: null + // location: + // x: 1 + // y: 6 + // - score: 12 + // location: + // x: 2 + // y: null + // - score: 13 + // location: + // x: 3 + // y: 8 + // - score: 14 + // location: null + // - null + // + let inner_fields = Fields::from(vec![ + Field::new("x", DataType::Int32, false), + Field::new("y", DataType::Int32, true), + ]); + let inner_struct = DataType::Struct(inner_fields.clone()); + let outer_fields = Fields::from(vec![ + Field::new("score", DataType::Int32, true), + Field::new("location", inner_struct, true), + ]); + + let x_vals = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]); + let y_vals = Int32Array::from(vec![Some(6), None, Some(8), Some(9), Some(10)]); + let scores = Int32Array::from(vec![None, Some(12), Some(13), Some(14), Some(15)]); + + let location_validity = NullBuffer::from(vec![true, true, true, false, true]); + let locations = StructArray::new( + inner_fields, + vec![Arc::new(x_vals), Arc::new(y_vals)], + Some(location_validity), + ); + + let rows_validity = NullBuffer::from(vec![true, true, true, true, false]); + let rows = StructArray::new( + outer_fields, + vec![Arc::new(scores), Arc::new(locations)], + Some(rows_validity), + ); + + let test_cases = TestCases::default().with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, HashMap::new()).await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_masked_nonempty_list() { + // [[1, 2], [NULL], [4], [], NULL, NULL-STRUCT] + // + let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4), Some(5), Some(6)]); + let offsets = OffsetBuffer::new(ScalarBuffer::::from(vec![0, 2, 3, 4, 4, 4, 5])); + let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]); + let list_array = ListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + offsets, + Arc::new(items), + Some(NullBuffer::new(list_validity)), + ); + let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]); + let struct_array = StructArray::new( + Fields::from(vec![Field::new( + "inner_list", + list_array.data_type().clone(), + true, + )]), + vec![Arc::new(list_array)], + Some(NullBuffer::new(struct_validity)), + ); + check_round_trip_encoding_of_data( + vec![Arc::new(struct_array)], + &TestCases::default().with_structural_encodings(), + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_struct_list() { + // [[1, 2], [NULL], [4], [], NULL, NULL-STRUCT] + // + let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]); + let offsets = OffsetBuffer::new(ScalarBuffer::::from(vec![0, 2, 3, 4, 4, 4, 4])); + let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]); + let list_array = ListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + offsets, + Arc::new(items), + Some(NullBuffer::new(list_validity)), + ); + let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]); + let struct_array = StructArray::new( + Fields::from(vec![Field::new( + "inner_list", + list_array.data_type().clone(), + true, + )]), + vec![Arc::new(list_array)], + Some(NullBuffer::new(struct_validity)), + ); + check_round_trip_encoding_of_data( + vec![Arc::new(struct_array)], + &TestCases::default().with_structural_encodings(), + HashMap::new(), + ) + .await; + } + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_struct_list( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let data_type = DataType::Struct(Fields::from(vec![ + Field::new( + "inner_list", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + ), + Field::new("outer_int", DataType::Int32, true), + ])); + let field = Field::new("row", data_type, false); + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_empty_struct( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + // It's technically legal for a struct to have 0 children, need to + // make sure we support that + let data_type = DataType::Struct(Fields::from(Vec::::default())); + let field = Field::new("row", data_type, false); + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_complicated_struct( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let data_type = DataType::Struct(Fields::from(vec![ + Field::new("int", DataType::Int32, true), + Field::new( + "inner", + DataType::Struct(Fields::from(vec![ + Field::new("inner_int", DataType::Int32, true), + Field::new( + "inner_list", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + ), + ])), + true, + ), + Field::new("outer_binary", DataType::Binary, true), + ])); + let field = Field::new("row", data_type, false); + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[test_log::test(tokio::test)] + async fn test_list_of_struct_with_null_struct_element() { + // Regression: a list containing structs where most struct elements are null + // causes a length mismatch during decoding with V2_2 encoding. + use arrow_array::StringArray; + + let tag_array = StringArray::from(vec![ + Some("valid"), + Some("null_struct"), + Some("valid"), + Some("valid"), + ]); + let struct_fields = Fields::from(vec![Field::new("tag", DataType::Utf8, true)]); + // 3 out of 4 struct elements are null + let struct_validity = NullBuffer::from(vec![false, true, false, false]); + let struct_array = StructArray::new( + struct_fields.clone(), + vec![Arc::new(tag_array)], + Some(struct_validity), + ); + + let offsets = OffsetBuffer::new(ScalarBuffer::::from(vec![0, 4])); + let list_field = Field::new("item", DataType::Struct(struct_fields), true); + let list_array = + ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None); + + check_round_trip_encoding_of_data( + vec![Arc::new(list_array)], + &TestCases::default().with_u32_structural_encodings(), + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_list_of_struct_with_all_null_child() { + use arrow_array::StringArray; + + let a_array = StringArray::from(vec![Some("w"), Some("x"), Some("y"), Some("z")]); + let b_array = StringArray::from(vec![None::<&str>, None, None, None]); + let struct_fields = Fields::from(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + ]); + let struct_array = StructArray::new( + struct_fields.clone(), + vec![Arc::new(a_array), Arc::new(b_array)], + None, + ); + + let offsets = OffsetBuffer::new(ScalarBuffer::::from(vec![0, 2, 4])); + let list_field = Field::new("item", DataType::Struct(struct_fields), true); + let list_array = + ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None); + + check_round_trip_encoding_of_data( + vec![Arc::new(list_array)], + &TestCases::default() + .with_range(0..1) + .with_range(1..2) + .with_indices(vec![0]) + .with_indices(vec![1]) + .with_structural_encodings(), + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_list_of_struct_with_constant_child_and_empty_lists() { + let item_fields = Fields::from(vec![ + Field::new("maneuver", DataType::Int64, true), + Field::new("remaining_dist", DataType::Float64, true), + ]); + let item_array = StructArray::new( + item_fields.clone(), + vec![ + Arc::new(Int64Array::from(vec![Some(3), Some(3), Some(3)])), + Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0), Some(3.0)])), + ], + None, + ); + + let list_field = Field::new("item", DataType::Struct(item_fields), true); + let list_array = ListArray::new( + Arc::new(list_field), + OffsetBuffer::new(ScalarBuffer::::from(vec![0, 1, 2, 3, 3, 3, 3])), + Arc::new(item_array), + None, + ); + + let data_fields = Fields::from(vec![Field::new( + "maneuvers", + list_array.data_type().clone(), + true, + )]); + let data_array = StructArray::new(data_fields.clone(), vec![Arc::new(list_array)], None); + let row_validity = NullBuffer::from(vec![true, true, true, true, true, false]); + let row_array = StructArray::new( + Fields::from(vec![Field::new( + "data", + DataType::Struct(data_fields), + true, + )]), + vec![Arc::new(data_array)], + Some(row_validity), + ); + + check_round_trip_encoding_of_data( + vec![Arc::new(row_array)], + &TestCases::default().with_structural_encodings(), + HashMap::new(), + ) + .await; + } + + #[test] + fn test_plan_decoded_bytes_struct_non_nullable_primitive() { + // Nullable struct containing a non-nullable i32 child. + // Only the struct-level bitmap contributes; the child has no bitmap. + // + // N == CANDIDATE_BATCH_SIZES[5] == 1024 so that all sizes are + // naturally aligned to Arrow's 64-byte allocation boundary. + use arrow_buffer::NullBuffer; + + const N: usize = 1024; + let x_vals = Int32Array::from_iter_values(0..N as i32); + let struct_nulls = NullBuffer::from((0..N).map(|i| i % 2 == 0).collect::>()); + let struct_field = Field::new("x", DataType::Int32, false); + let struct_array = StructArray::new( + Fields::from(vec![struct_field]), + vec![Arc::new(x_vals) as ArrayRef], + Some(struct_nulls), + ); + let expected = struct_array.get_buffer_memory_size() as u64; + + let child_field = Arc::new(Field::new("x", DataType::Int32, false)); + let fields = Fields::from(vec![child_field]); + let decoder = StructuralStructDecoder::new( + fields, false, /*is_root=*/ false, /*nullable=*/ true, + ) + .unwrap(); + + let bytes = decoder.plan_decoded_bytes(N as u64).unwrap(); + // bytes[5] corresponds to CANDIDATE_BATCH_SIZES[5] == 1024 == N + assert_eq!( + bytes[5], expected, + "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}", + bytes[5], + ); + } + + #[test] + fn test_plan_decoded_bytes_struct_nullable_primitive() { + // Nullable struct containing a nullable i32 child. + // Both the struct and the child contribute a validity bitmap. + use arrow_buffer::NullBuffer; + + const N: usize = 1024; + let y_vals = + Int32Array::from_iter((0..N as i32).map(|i| if i % 2 == 0 { Some(i) } else { None })); + let struct_nulls = NullBuffer::from((0..N).map(|i| i % 3 != 0).collect::>()); + let struct_field = Field::new("y", DataType::Int32, true); + let struct_array = StructArray::new( + Fields::from(vec![struct_field]), + vec![Arc::new(y_vals) as ArrayRef], + Some(struct_nulls), + ); + let expected = struct_array.get_buffer_memory_size() as u64; + + let child_field = Arc::new(Field::new("y", DataType::Int32, true)); + let fields = Fields::from(vec![child_field]); + let decoder = StructuralStructDecoder::new( + fields, false, /*is_root=*/ false, /*nullable=*/ true, + ) + .unwrap(); + + let bytes = decoder.plan_decoded_bytes(N as u64).unwrap(); + assert_eq!( + bytes[5], expected, + "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}", + bytes[5], + ); + } + + #[test_log::test(tokio::test)] + async fn test_ragged_scheduling() { + // This test covers scheduling when batches straddle page boundaries + + // Create a list with 10k nulls + let items_builder = Int32Builder::new(); + let mut list_builder = ListBuilder::new(items_builder); + for _ in 0..10000 { + list_builder.append_null(); + } + let list_array = Arc::new(list_builder.finish()); + let int_array = Arc::new(Int32Array::from_iter_values(0..10000)); + let fields = vec![ + Field::new("", list_array.data_type().clone(), true), + Field::new("", int_array.data_type().clone(), true), + ]; + let struct_array = Arc::new(StructArray::new( + Fields::from(fields), + vec![list_array, int_array], + None, + )) as ArrayRef; + let struct_arrays = (0..10000) + // Intentionally skip in some randomish amount to create more ragged scheduling + .step_by(437) + .map(|offset| struct_array.slice(offset, 437.min(10000 - offset))) + .collect::>(); + check_round_trip_encoding_of_data(struct_arrays, &TestCases::default(), HashMap::new()) + .await; + } + + /// A null list produces an "invisible" item: a rep/def entry that owns no value in the + /// leaf page. The full-zip decoder works out which items own a value from + /// `max_visible_def`, and adding a second nullable layer *above* the list (the struct + /// itself) must not move that threshold. When it does, the decoder hands a value to + /// every null list and each later value shifts by one. + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_full_zip_invisible_items_under_nullable_struct( + #[values(false, true)] variable_width_items: bool, + ) { + // struct> with both layers nullable: + // {x: [_, _]} / {x: null} / {x: [_]} / null / {x: [_, _]} + let (item_type, items): (DataType, ArrayRef) = if variable_width_items { + ( + DataType::Utf8, + Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])), + ) + } else { + ( + DataType::Int32, + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])), + ) + }; + let lists = ListArray::new( + Arc::new(Field::new("item", item_type, true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 2, 3, 3, 5])), + items, + Some(NullBuffer::from(vec![true, false, true, true, true])), + ); + let fields = Fields::from(vec![Field::new("x", lists.data_type().clone(), true)]); + let rows = StructArray::new( + fields, + vec![Arc::new(lists)], + Some(NullBuffer::from(vec![true, true, true, false, true])), + ); + + let metadata = HashMap::from([( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_FULLZIP.to_string(), + )]); + let test_cases = TestCases::default().with_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, metadata).await; + } +} diff --git a/vendor/lance-encoding/src/encodings/physical.rs b/vendor/lance-encoding/src/encodings/physical.rs new file mode 100644 index 000000000..0439c0216 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +pub mod binary; +#[cfg(feature = "bitpacking")] +pub mod bitpacking; +pub mod block; +pub mod byte_stream_split; +pub mod constant; +pub mod fsst; +pub mod general; +pub mod packed; +pub mod rle; +pub mod value; diff --git a/vendor/lance-encoding/src/encodings/physical/binary.rs b/vendor/lance-encoding/src/encodings/physical/binary.rs new file mode 100644 index 000000000..e07d586a4 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/binary.rs @@ -0,0 +1,1436 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Basic encodings for variable width data +//! +//! These are not compression but represent the "leaf" encodings for variable length data +//! where we simply match the data with the rules of the structural encoding. +//! +//! These encodings are transparent since we aren't actually doing any compression. No information +//! is needed in the encoding description. + +use arrow_array::OffsetSizeTrait; +use byteorder::{ByteOrder, LittleEndian}; +use core::panic; + +use crate::compression::{ + BlockCompressor, BlockDecompressor, MiniBlockDecompressor, VariablePerValueDecompressor, + require_block_payload, +}; + +use crate::buffer::LanceBuffer; +use crate::data::{BlockInfo, DataBlock, VariableWidthBlock}; +use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}; +use crate::encodings::logical::primitive::miniblock::{ + MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, + MiniBlockCompressor, +}; +use crate::format::pb21::CompressiveEncoding; +use crate::format::pb21::compressive_encoding::Compression; +use crate::format::{ProtobufUtils21, pb21}; + +use lance_core::utils::bit::pad_bytes_to; +use lance_core::{Error, Result}; + +#[derive(Debug)] +pub struct BinaryMiniBlockEncoder { + minichunk_size: i64, +} + +impl Default for BinaryMiniBlockEncoder { + fn default() -> Self { + Self { + minichunk_size: *AIM_MINICHUNK_SIZE, + } + } +} + +const DEFAULT_AIM_MINICHUNK_SIZE: i64 = 4 * 1024; + +pub static AIM_MINICHUNK_SIZE: std::sync::LazyLock = std::sync::LazyLock::new(|| { + std::env::var("LANCE_BINARY_MINIBLOCK_CHUNK_SIZE") + .unwrap_or_else(|_| DEFAULT_AIM_MINICHUNK_SIZE.to_string()) + .parse::() + .unwrap_or(DEFAULT_AIM_MINICHUNK_SIZE) +}); + +// Make it to support both u32 and u64 +fn chunk_offsets( + offsets: &[N], + data: &[u8], + alignment: usize, + minichunk_size: i64, +) -> (Vec, Vec) { + #[derive(Debug)] + struct ChunkInfo { + chunk_start_offset_in_orig_idx: usize, + chunk_last_offset_in_orig_idx: usize, + // the bytes in every chunk starts at `chunk.bytes_start_offset` + bytes_start_offset: usize, + // every chunk is padded to 8 bytes. + // we need to interpret every chunk as &[u32] so we need it to padded at least to 4 bytes, + // this field can actually be eliminated and I can use `num_bytes` in `MiniBlockChunk` to compute + // the `output_total_bytes`. + padded_chunk_size: usize, + } + + let byte_width: usize = N::get_byte_width(); + let mut chunks_info = vec![]; + let mut chunks = vec![]; + let mut last_offset_in_orig_idx = 0; + loop { + let this_last_offset_in_orig_idx = + search_next_offset_idx(offsets, last_offset_in_orig_idx, minichunk_size); + + let num_values_in_this_chunk = this_last_offset_in_orig_idx - last_offset_in_orig_idx; + let chunk_bytes = offsets[this_last_offset_in_orig_idx] - offsets[last_offset_in_orig_idx]; + let this_chunk_size = + (num_values_in_this_chunk + 1) * byte_width + chunk_bytes.to_usize().unwrap(); + + let padded_chunk_size = this_chunk_size.next_multiple_of(alignment); + debug_assert!(padded_chunk_size > 0); + + let this_chunk_bytes_start_offset = (num_values_in_this_chunk + 1) * byte_width; + chunks_info.push(ChunkInfo { + chunk_start_offset_in_orig_idx: last_offset_in_orig_idx, + chunk_last_offset_in_orig_idx: this_last_offset_in_orig_idx, + bytes_start_offset: this_chunk_bytes_start_offset, + padded_chunk_size, + }); + chunks.push(MiniBlockChunk { + log_num_values: if this_last_offset_in_orig_idx == offsets.len() - 1 { + 0 + } else { + num_values_in_this_chunk.trailing_zeros() as u8 + }, + buffer_sizes: vec![padded_chunk_size as u32], + }); + if this_last_offset_in_orig_idx == offsets.len() - 1 { + break; + } + last_offset_in_orig_idx = this_last_offset_in_orig_idx; + } + + let output_total_bytes = chunks_info + .iter() + .map(|chunk_info| chunk_info.padded_chunk_size) + .sum::(); + + let mut output: Vec = Vec::with_capacity(output_total_bytes); + + for chunk in chunks_info { + let this_chunk_offsets: Vec = offsets + [chunk.chunk_start_offset_in_orig_idx..=chunk.chunk_last_offset_in_orig_idx] + .iter() + .map(|offset| { + *offset - offsets[chunk.chunk_start_offset_in_orig_idx] + + N::from_usize(chunk.bytes_start_offset).unwrap() + }) + .collect(); + + let this_chunk_offsets = LanceBuffer::reinterpret_vec(this_chunk_offsets); + output.extend_from_slice(&this_chunk_offsets); + + let start_in_orig = offsets[chunk.chunk_start_offset_in_orig_idx] + .to_usize() + .unwrap(); + let end_in_orig = offsets[chunk.chunk_last_offset_in_orig_idx] + .to_usize() + .unwrap(); + output.extend_from_slice(&data[start_in_orig..end_in_orig]); + + // pad this chunk to make it align to desired bytes. + const PAD_BYTE: u8 = 72; + let pad_len = pad_bytes_to(output.len(), alignment); + + // Compare with usize literal to avoid type mismatch with N + if pad_len > 0_usize { + output.extend(std::iter::repeat_n(PAD_BYTE, pad_len)); + } + } + (vec![LanceBuffer::reinterpret_vec(output)], chunks) +} + +// search for the next offset index to cut the values into a chunk. +// this function incrementally peek the number of values in a chunk, +// each time multiplies the number of values by 2. +// It returns the offset_idx in `offsets` that belongs to this chunk. +fn search_next_offset_idx( + offsets: &[N], + last_offset_idx: usize, + minichunk_size: i64, +) -> usize { + // MiniBlockChunk uses `log_num_values == 0` as a sentinel for the final chunk. This means we + // must avoid creating 1-value chunks except for the final chunk, even if the configured + // `minichunk_size` is too small to fit more than one value. + let remaining_values = offsets.len().saturating_sub(last_offset_idx + 1); + if remaining_values <= 1 { + return offsets.len() - 1; + } + + let mut num_values = 2; + let mut new_num_values = num_values * 2; + loop { + if last_offset_idx + new_num_values >= offsets.len() { + let existing_bytes = offsets[offsets.len() - 1] - offsets[last_offset_idx]; + // existing bytes plus the new offset size + let new_size = existing_bytes + + N::from_usize((offsets.len() - last_offset_idx) * N::get_byte_width()).unwrap(); + if new_size.to_i64().unwrap() <= minichunk_size { + // case 1: can fit the rest of all data into a miniblock + return offsets.len() - 1; + } else { + // case 2: can only fit the last tried `num_values` into a miniblock + return last_offset_idx + num_values; + } + } + let existing_bytes = offsets[last_offset_idx + new_num_values] - offsets[last_offset_idx]; + let new_size = + existing_bytes + N::from_usize((new_num_values + 1) * N::get_byte_width()).unwrap(); + if new_size.to_i64().unwrap() <= minichunk_size { + if new_num_values * 2 > *MAX_MINIBLOCK_VALUES as usize { + // hit the max number of values limit + break; + } + num_values = new_num_values; + new_num_values *= 2; + } else { + break; + } + } + last_offset_idx + num_values +} + +impl BinaryMiniBlockEncoder { + pub fn new(minichunk_size: Option) -> Self { + Self { + minichunk_size: minichunk_size.unwrap_or(*AIM_MINICHUNK_SIZE), + } + } + + // put binary data into chunks, every chunk is less than or equal to `minichunk_size`. + // In each chunk, offsets are put first then followed by binary bytes data, each chunk is padded to 8 bytes. + // the offsets in the chunk points to the bytes offset in this chunk. + fn chunk_data(&self, data: VariableWidthBlock) -> (MiniBlockCompressed, CompressiveEncoding) { + // TODO: Support compression of offsets + // TODO: Support general compression of data + match data.bits_per_offset { + 32 => { + let offsets = data.offsets.borrow_to_typed_slice::(); + let (buffers, chunks) = + chunk_offsets(offsets.as_ref(), &data.data, 4, self.minichunk_size); + ( + MiniBlockCompressed { + data: buffers, + chunks, + num_values: data.num_values, + }, + ProtobufUtils21::variable(ProtobufUtils21::flat(32, None), None), + ) + } + 64 => { + let offsets = data.offsets.borrow_to_typed_slice::(); + let (buffers, chunks) = + chunk_offsets(offsets.as_ref(), &data.data, 8, self.minichunk_size); + ( + MiniBlockCompressed { + data: buffers, + chunks, + num_values: data.num_values, + }, + ProtobufUtils21::variable(ProtobufUtils21::flat(64, None), None), + ) + } + _ => panic!("Unsupported bits_per_offset={}", data.bits_per_offset), + } + } +} + +impl MiniBlockCompressor for BinaryMiniBlockEncoder { + fn compress( + &self, + _context: MiniBlockCompressionContext, + data: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + match data { + DataBlock::VariableWidth(variable_width) => Ok(self.chunk_data(variable_width)), + _ => Err(Error::invalid_input_source( + format!( + "Cannot compress a data block of type {} with BinaryMiniBlockEncoder", + data.name() + ) + .into(), + )), + } + } +} + +#[derive(Debug)] +pub struct BinaryMiniBlockDecompressor { + bits_per_offset: u8, +} + +impl BinaryMiniBlockDecompressor { + pub fn new(bits_per_offset: u8) -> Self { + assert!(bits_per_offset == 32 || bits_per_offset == 64); + Self { bits_per_offset } + } + + pub fn from_variable(variable: &pb21::Variable) -> Self { + if let Compression::Flat(flat) = variable + .offsets + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + { + Self { + bits_per_offset: flat.bits_per_value as u8, + } + } else { + panic!("Unsupported offsets compression: {:?}", variable.offsets); + } + } +} + +/// Cold path: pinpoint why the chunk-relative offsets of a binary mini-block +/// chunk failed validation. +fn chunk_offset_violation_error>(offsets: &[T], chunk_len: usize) -> Error { + let mut previous: u64 = offsets[0].into(); + for (position, &offset) in offsets.iter().enumerate().skip(1) { + let offset: u64 = offset.into(); + if offset < previous { + return Error::corrupt_file_named( + "binary mini-block", + format!( + "value offset at position {position} decreases: {offset} < {previous} \ + (chunk is {chunk_len} bytes)" + ), + ); + } + previous = offset; + } + Error::corrupt_file_named( + "binary mini-block", + format!("value offset {previous} is out of bounds for a chunk of {chunk_len} bytes"), + ) +} + +impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { + // decompress a MiniBlock of binary data, the num_values must be less than or equal + // to the number of values this MiniBlock has, BinaryMiniBlock doesn't store `the number of values` + // it has so assertion can not be done here and the caller of `decompress` must ensure + // `num_values` <= number of values in the chunk. + // + // The chunk-relative value offsets at the front of the chunk come straight + // from the file and are used to slice the chunk buffer, so corrupt values + // must surface as a typed error instead of a panic or an out-of-bounds + // read. The monotonicity check rides along the existing rebase loop (the + // `&=` accumulation keeps it branchless) so validation adds no extra pass. + fn decompress(&self, data: Vec, num_values: u64) -> Result { + assert_eq!(data.len(), 1); + let data = data.into_iter().next().unwrap(); + + let bytes_per_offset = self.bits_per_offset as usize / 8; + if !data.len().is_multiple_of(bytes_per_offset) { + return Err(Error::corrupt_file_named( + "binary mini-block", + format!( + "chunk size {} is not a multiple of the {}-byte offset width", + data.len(), + bytes_per_offset + ), + )); + } + let num_offsets = (num_values as usize).checked_add(1).ok_or_else(|| { + Error::corrupt_file_named( + "binary mini-block", + format!("cannot decode {num_values} values from a single chunk"), + ) + })?; + if data.len() / bytes_per_offset < num_offsets { + return Err(Error::corrupt_file_named( + "binary mini-block", + format!( + "chunk of {} bytes holds {} offsets but decoding {} values requires {}", + data.len(), + data.len() / bytes_per_offset, + num_values, + num_offsets + ), + )); + } + + // The value region must start past the offsets being decoded, otherwise + // the offset table itself aliases into the value bytes. A lower bound + // (not equality) because a prefix read of the chunk legitimately leaves + // unrequested offsets between the requested prefix and the values. + let min_value_region_start = num_offsets * bytes_per_offset; + let value_region_overlap_error = |first: u64| { + Error::corrupt_file_named( + "binary mini-block", + format!( + "value region starts at offset {first} which overlaps the {num_offsets} \ + requested offsets ({min_value_region_start} bytes)" + ), + ) + }; + + if self.bits_per_offset == 64 { + let offsets_buffer = data.borrow_to_typed_slice::(); + let offsets = &offsets_buffer.as_ref()[..num_offsets]; + + let first = offsets[0]; + if first < min_value_region_start as u64 { + return Err(value_region_overlap_error(first)); + } + let mut previous = first; + let mut is_monotonic = true; + let result_offsets = offsets + .iter() + .map(|&offset| { + is_monotonic &= previous <= offset; + previous = offset; + offset.wrapping_sub(first) + }) + .collect::>(); + let last = offsets[num_offsets - 1]; + if !is_monotonic || last as usize > data.len() { + return Err(chunk_offset_violation_error(offsets, data.len())); + } + + Ok(DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::from(data[first as usize..last as usize].to_vec()), + offsets: LanceBuffer::reinterpret_vec(result_offsets), + bits_per_offset: 64, + num_values, + block_info: BlockInfo::new(), + })) + } else { + let offsets_buffer = data.borrow_to_typed_slice::(); + let offsets = &offsets_buffer.as_ref()[..num_offsets]; + + let first = offsets[0]; + if (first as u64) < min_value_region_start as u64 { + return Err(value_region_overlap_error(first as u64)); + } + let mut previous = first; + let mut is_monotonic = true; + let result_offsets = offsets + .iter() + .map(|&offset| { + is_monotonic &= previous <= offset; + previous = offset; + offset.wrapping_sub(first) + }) + .collect::>(); + let last = offsets[num_offsets - 1]; + if !is_monotonic || last as usize > data.len() { + return Err(chunk_offset_violation_error(offsets, data.len())); + } + + Ok(DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::from(data[first as usize..last as usize].to_vec()), + offsets: LanceBuffer::reinterpret_vec(result_offsets), + bits_per_offset: 32, + num_values, + block_info: BlockInfo::new(), + })) + } + } +} + +/// Most basic encoding for variable-width data which does no compression at all +/// The DataBlock memory layout looks like below: +/// +/// | bits_per_offset | bytes_start_offset | offsets data | bytes data | +/// | ------------------------- | ------------------------- | ------------ | ---------- | +/// | /8 bytes | /8 bytes | offsets_len | data_len | +/// +/// It's used in VariableEncoder and BinaryBlockDecompressor +/// +#[derive(Debug, Default)] +pub struct VariableEncoder {} + +impl BlockCompressor for VariableEncoder { + fn compress(&self, mut data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + let bits_per_offset = match &data { + DataBlock::VariableWidth(data) => data.bits_per_offset, + _ => { + return Err(Error::invalid_input( + "BinaryBlockEncoder requires a variable-width block", + )); + } + }; + match data { + DataBlock::VariableWidth(ref mut variable_width_data) => { + match variable_width_data.bits_per_offset { + 32 => { + let offsets = variable_width_data.offsets.borrow_to_typed_slice::(); + let offsets = offsets.as_ref(); + // The first 4 bytes store the bits per offset, the next 4 bytes store the start + // offset of the bytes data, then offsets data, then bytes data. + let bytes_start_offset = 4 + 4 + std::mem::size_of_val(offsets) as u32; + + let output_total_bytes = + bytes_start_offset as usize + variable_width_data.data.len(); + let mut output: Vec = Vec::with_capacity(output_total_bytes); + + // Store bit_per_offset info + output.extend_from_slice(&(32_u32).to_le_bytes()); + + // store `bytes_start_offset` in the next 4 bytes of output buffer + output.extend_from_slice(&(bytes_start_offset).to_le_bytes()); + + // store offsets + output.extend_from_slice(&variable_width_data.offsets); + + // store bytes + output.extend_from_slice(&variable_width_data.data); + Ok(LanceBuffer::from(output)) + } + 64 => { + let offsets = variable_width_data.offsets.borrow_to_typed_slice::(); + let offsets = offsets.as_ref(); + // The first 8 bytes store the bits per offset, the next 8 bytes store the start + // offset of the bytes data, then offsets data, then bytes data. + let bytes_start_offset = 8 + 8 + std::mem::size_of_val(offsets) as u64; + + let output_total_bytes = + bytes_start_offset as usize + variable_width_data.data.len(); + let mut output: Vec = Vec::with_capacity(output_total_bytes); + + // Store bit_per_offset info + output.extend_from_slice(&(64_u64).to_le_bytes()); + + // store `bytes_start_offset` in the next 8 bytes of output buffer + output.extend_from_slice(&(bytes_start_offset).to_le_bytes()); + + // store offsets + output.extend_from_slice(&variable_width_data.offsets); + + // store bytes + output.extend_from_slice(&variable_width_data.data); + Ok(LanceBuffer::from(output)) + } + _ => Err(Error::invalid_input(format!( + "BinaryBlockEncoder does not support {}-bit offsets", + variable_width_data.bits_per_offset + ))), + } + } + _ => unreachable!("variable-width input was validated above"), + } + .map(|payload| { + ( + Some(payload), + ProtobufUtils21::variable( + ProtobufUtils21::flat(bits_per_offset as u64, None), + None, + ), + ) + }) + } +} + +impl PerValueCompressor for VariableEncoder { + fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> { + let DataBlock::VariableWidth(variable) = data else { + panic!("BinaryPerValueCompressor can only work with Variable Width DataBlock."); + }; + + let encoding = ProtobufUtils21::variable( + ProtobufUtils21::flat(variable.bits_per_offset as u64, None), + None, + ); + Ok((PerValueDataBlock::Variable(variable), encoding)) + } +} + +#[derive(Debug, Default)] +pub struct VariableDecoder {} + +impl VariablePerValueDecompressor for VariableDecoder { + fn decompress(&self, data: VariableWidthBlock) -> Result { + Ok(DataBlock::VariableWidth(data)) + } +} + +#[derive(Debug, Default)] +pub struct BinaryBlockDecompressor {} + +impl BlockDecompressor for BinaryBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Binary block")?; + // In older (not quite stable) versions we stored the bits per offset as a single byte and then the num_values + // as four bytes. However, this led to alignment problems and was wasteful since we already store the num_values + // in higher layers. + // + // In the standard scheme we use 4 bytes for the bits per offset and 4 bytes for the bytes_start_offset and we + // rely on the passed in num_values to be correct. + + // This isn't perfect but it's probably good enough and the best I think we can do. The bits per offset will + // never be more than 255 and it's little endian so the last 3 bytes will always be 0. These will be the least + // significant 3 bytes of the number of values in the old scheme. It's pretty unlikely these are all 0 (that would + // mean there are at least 16M values in a single page) so we'll use this to determine if the old scheme is used. + // + // The header fields and the offsets themselves come straight from the file. + // The structural checks below (all O(1)) reject blocks whose regions do not + // line up; the offset *values* are validated later, by the mandatory layout + // validation in `VariableWidthBlock::into_arrow`, so they are not rescanned + // here. + if data.len() < 4 { + return Err(Error::corrupt_file_named( + "variable-width block", + format!( + "block of {} bytes is too small to hold a header", + data.len() + ), + )); + } + let is_old_scheme = data[1] != 0 || data[2] != 0 || data[3] != 0; + + let ensure_header = |header_len: usize| { + if data.len() < header_len { + return Err(Error::corrupt_file_named( + "variable-width block", + format!( + "block of {} bytes is too small for a {} byte header", + data.len(), + header_len + ), + )); + } + Ok(()) + }; + let (bits_per_offset, bytes_start_offset, offset_start) = if is_old_scheme { + // Old scheme + let bits_per_offset = data[0]; + match bits_per_offset { + 32 => { + ensure_header(9)?; + debug_assert_eq!(LittleEndian::read_u32(&data[1..5]), num_values as u32); + let bytes_start_offset = LittleEndian::read_u32(&data[5..9]); + (bits_per_offset, bytes_start_offset as u64, 9_u64) + } + 64 => { + ensure_header(17)?; + debug_assert_eq!(LittleEndian::read_u64(&data[1..9]), num_values); + let bytes_start_offset = LittleEndian::read_u64(&data[9..17]); + (bits_per_offset, bytes_start_offset, 17) + } + _ => { + return Err(Error::invalid_input_source( + format!("Unsupported bits_per_offset={}", bits_per_offset).into(), + )); + } + } + } else { + // Standard scheme + let bits_per_offset = LittleEndian::read_u32(&data[0..4]) as u8; + match bits_per_offset { + 32 => { + ensure_header(8)?; + let bytes_start_offset = LittleEndian::read_u32(&data[4..8]); + (bits_per_offset, bytes_start_offset as u64, 8) + } + 64 => { + ensure_header(16)?; + let bytes_start_offset = LittleEndian::read_u64(&data[8..16]); + (bits_per_offset, bytes_start_offset, 16) + } + _ => { + return Err(Error::invalid_input_source( + format!("Unsupported bits_per_offset={}", bits_per_offset).into(), + )); + } + } + }; + + // The offsets region sits between the header and `bytes_start_offset` + // and must hold exactly `num_values + 1` offsets starting at zero. + let expected_offsets_bytes = num_values + .checked_add(1) + .and_then(|num_offsets| num_offsets.checked_mul(bits_per_offset as u64 / 8)) + .ok_or_else(|| { + Error::corrupt_file_named( + "variable-width block", + format!("offsets region size overflows for {num_values} values"), + ) + })?; + if bytes_start_offset < offset_start || bytes_start_offset > data.len() as u64 { + return Err(Error::corrupt_file_named( + "variable-width block", + format!( + "bytes start offset {} is outside the block (header: {} bytes, block: {} bytes)", + bytes_start_offset, + offset_start, + data.len() + ), + )); + } + if bytes_start_offset - offset_start != expected_offsets_bytes { + return Err(Error::corrupt_file_named( + "variable-width block", + format!( + "expected {} offset bytes for {} values but found {}", + expected_offsets_bytes, + num_values, + bytes_start_offset - offset_start + ), + )); + } + + // the next `bytes_start_offset - offset_start` stores the offsets. + let offsets = data.slice_with_length( + offset_start as usize, + (bytes_start_offset - offset_start) as usize, + ); + let first_offset = match bits_per_offset { + 32 => LittleEndian::read_u32(&offsets[0..4]) as u64, + _ => LittleEndian::read_u64(&offsets[0..8]), + }; + if first_offset != 0 { + return Err(Error::corrupt_file_named( + "variable-width block", + format!("first offset must be 0 but found {first_offset}"), + )); + } + + // the rest are the binary bytes. + let data = data.slice_with_length( + bytes_start_offset as usize, + data.len() - bytes_start_offset as usize, + ); + + Ok(DataBlock::VariableWidth(VariableWidthBlock { + data, + offsets, + bits_per_offset, + num_values, + block_info: BlockInfo::new(), + })) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::{ + ArrayRef, StringArray, + builder::{LargeStringBuilder, StringBuilder}, + }; + use arrow_schema::{DataType, Field}; + + use crate::{ + buffer::LanceBuffer, + constants::{ + COMPRESSION_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, + STRUCTURAL_ENCODING_MINIBLOCK, + }, + data::{BlockInfo, DataBlock, VariableWidthBlock}, + testing::check_specific_random, + }; + use rstest::rstest; + use std::{collections::HashMap, sync::Arc, vec}; + + use crate::testing::{ + FnArrayGeneratorProvider, TestCases, TestEncoding, check_basic_random_case, + check_round_trip_encoding_generated, check_round_trip_encoding_of_data, + }; + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_utf8_binary( + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let field = Field::new("", DataType::Utf8, false); + check_specific_random( + field, + TestCases::basic() + .with_encoding(encoding) + .with_page_sizes(vec![page_size]) + .with_slicing_modes([use_slicing]), + ) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_binary( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + #[values(DataType::Utf8, DataType::Binary)] data_type: DataType, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let field = Field::new("", data_type, false).with_metadata(field_metadata); + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_binary_fsst( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + #[values(DataType::Binary, DataType::Utf8)] data_type: DataType, + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + field_metadata.insert(COMPRESSION_META_KEY.to_string(), "fsst".into()); + let field = Field::new("", data_type, true).with_metadata(field_metadata); + // TODO (https://github.com/lance-format/lance/issues/4783) + let test_cases = TestCases::default() + .with_encoding(encoding) + .with_page_sizes(vec![page_size]) + .with_slicing_modes([use_slicing]); + check_specific_random(field, test_cases).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_fsst_large_binary( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + #[values(DataType::LargeBinary, DataType::LargeUtf8)] data_type: DataType, + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + field_metadata.insert(COMPRESSION_META_KEY.to_string(), "fsst".into()); + let field = Field::new("", data_type, true).with_metadata(field_metadata); + check_specific_random( + field, + TestCases::basic() + .with_encoding(encoding) + .with_page_sizes(vec![page_size]) + .with_slicing_modes([use_slicing]), + ) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_large_binary_types( + #[values(DataType::LargeBinary, DataType::LargeUtf8)] data_type: DataType, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let field = Field::new("", data_type, true); + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_small_strings( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + let field = Field::new("", DataType::Utf8, true).with_metadata(field_metadata); + check_round_trip_encoding_generated( + field, + Box::new(FnArrayGeneratorProvider::new(move || { + lance_datagen::array::utf8_prefix_plus_counter("user_", /*is_large=*/ false) + })), + TestCases::basic() + .with_encoding(encoding) + .with_page_sizes(vec![page_size]) + .with_slicing_modes([use_slicing]), + ) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_simple_binary( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + #[values(DataType::Utf8, DataType::Binary)] data_type: DataType, + ) { + let string_array = StringArray::from(vec![Some("abc"), None, Some("pqr"), None, Some("m")]); + let string_array = arrow_cast::cast(&string_array, &data_type).unwrap(); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default() + .with_range(0..2) + .with_range(0..3) + .with_range(1..3) + .with_indices(vec![0, 1, 3, 4]); + check_round_trip_encoding_of_data( + vec![Arc::new(string_array)], + &test_cases, + field_metadata, + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_sliced_utf8() { + let string_array = StringArray::from(vec![Some("abc"), Some("de"), None, Some("fgh")]); + let string_array = string_array.slice(1, 3); + + let test_cases = TestCases::default() + .with_range(0..1) + .with_range(0..2) + .with_range(1..2); + check_round_trip_encoding_of_data( + vec![Arc::new(string_array)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_value_bigger_than_max_page_size( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + // Create one value larger than the configured 1MiB page budget. + let big_string = String::from_iter((0..(2 * 1024 * 1024)).map(|_| '0')); + let string_array = StringArray::from(vec![ + Some(big_string), + Some("abc".to_string()), + None, + None, + Some("xyz".to_string()), + ]); + + // Drop the max page size to 1MiB + let test_cases = TestCases::default() + .with_max_page_size(1024 * 1024) + .with_encoding(encoding); + + check_round_trip_encoding_of_data( + vec![Arc::new(string_array)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_page_split_parts_do_not_evenly_divide_rows( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + // Regression: split 90 rows into four parts, where the part count does + // not evenly divide the row count. + let big_string = String::from_iter((0..45_000).map(|_| '0')); + let string_array = StringArray::from_iter_values((0..90).map(|_| big_string.clone())); + + check_round_trip_encoding_of_data( + vec![Arc::new(string_array)], + &TestCases::default() + .with_max_page_size(1024 * 1024) + .with_encoding(encoding), + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_empty_strings() { + // Scenario 1: Some strings are empty + + let values = [Some("abc"), Some(""), None]; + // Test empty list at beginning, middle, and end + for order in [[0, 1, 2], [1, 0, 2], [2, 0, 1]] { + let mut string_builder = StringBuilder::new(); + for idx in order { + string_builder.append_option(values[idx]); + } + let string_array = Arc::new(string_builder.finish()); + let test_cases = TestCases::default() + .with_indices(vec![1]) + .with_indices(vec![0]) + .with_indices(vec![2]) + .with_indices(vec![0, 1]); + check_round_trip_encoding_of_data( + vec![string_array.clone()], + &test_cases, + HashMap::new(), + ) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()) + .await; + } + + // Scenario 2: All strings are empty + + // When encoding an array of empty strings there are no bytes to encode + // which is strange and we want to ensure we handle it + let string_array = Arc::new(StringArray::from(vec![Some(""), None, Some("")])); + + let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]); + check_round_trip_encoding_of_data(vec![string_array.clone()], &test_cases, HashMap::new()) + .await; + let test_cases = test_cases.with_batch_size(1); + check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await; + } + + #[test_log::test(tokio::test)] + #[ignore] // This test is quite slow in debug mode + async fn test_jumbo_string() { + // This is an overflow test. We have a list of lists where each list + // has 1Mi items. We encode 5000 of these lists and so we have over 4Gi in the + // offsets range + let mut string_builder = LargeStringBuilder::new(); + // a 1 MiB string + let giant_string = String::from_iter((0..(1024 * 1024)).map(|_| '0')); + for _ in 0..5000 { + string_builder.append_option(Some(&giant_string)); + } + let giant_array = Arc::new(string_builder.finish()) as ArrayRef; + let arrs = vec![giant_array]; + + // // We can't validate because our validation relies on concatenating all input arrays + let test_cases = TestCases::default().without_validation(); + check_round_trip_encoding_of_data(arrs, &test_cases, HashMap::new()).await; + } + + #[rstest] + #[test_log::test(tokio::test)] + async fn test_binary_dictionary_encoding( + #[values(true, false)] with_nulls: bool, + #[values(100, 500, 35000)] dict_size: u32, + ) { + let test_cases = TestCases::default().with_structural_encodings(); + let strings = (0..dict_size) + .map(|i| i.to_string()) + .collect::>(); + + let repeated_strings: Vec<_> = strings + .iter() + .cycle() + .take(70000) + .enumerate() + .map(|(i, s)| { + if with_nulls && i % 7 == 0 { + None + } else { + Some(s.clone()) + } + }) + .collect(); + let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef; + check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await; + } + + #[test_log::test(tokio::test)] + async fn test_binary_encoding_verification() { + use lance_datagen::{ByteCount, RowCount}; + + let test_cases = TestCases::default() + .with_expected_encoding("variable") + .with_structural_encodings(); + + // Test both automatic selection and explicit configuration + // 1. Test automatic binary encoding selection (small strings that won't trigger FSST) + let arr_small = lance_datagen::gen_batch() + .anon_col(lance_datagen::array::rand_utf8(ByteCount::from(10), false)) + .into_batch_rows(RowCount::from(1000)) + .unwrap() + .column(0) + .clone(); + check_round_trip_encoding_of_data(vec![arr_small], &test_cases, HashMap::new()).await; + + // 2. Test explicit "none" compression to force binary encoding + let metadata_explicit = + HashMap::from([("lance-encoding:compression".to_string(), "none".to_string())]); + let arr_large = lance_datagen::gen_batch() + .anon_col(lance_datagen::array::rand_utf8(ByteCount::from(50), false)) + .into_batch_rows(RowCount::from(2000)) + .unwrap() + .column(0) + .clone(); + check_round_trip_encoding_of_data(vec![arr_large], &test_cases, metadata_explicit).await; + } + + #[test] + fn test_binary_miniblock_with_misaligned_buffer() { + use super::BinaryMiniBlockDecompressor; + use crate::buffer::LanceBuffer; + use crate::compression::MiniBlockDecompressor; + use crate::data::DataBlock; + + // Test case 1: u32 offsets + { + let decompressor = BinaryMiniBlockDecompressor { + bits_per_offset: 32, + }; + + // Create test data with u32 offsets + // BinaryMiniBlock format: all offsets followed by all string data + // Need to ensure total size is divisible by 4 for u32 + let mut test_data = Vec::new(); + + // Offsets section (3 offsets for 2 values + 1 end offset) + test_data.extend_from_slice(&12u32.to_le_bytes()); // offset to start of strings (after offsets) + test_data.extend_from_slice(&15u32.to_le_bytes()); // offset to second string + test_data.extend_from_slice(&20u32.to_le_bytes()); // offset to end + + // String data section + test_data.extend_from_slice(b"ABCXYZ"); // 6 bytes of string data + test_data.extend_from_slice(&[0, 0]); // 2 bytes padding to make total 20 bytes (divisible by 4) + + // Create a misaligned buffer by adding padding and slicing + let mut padded = Vec::with_capacity(test_data.len() + 1); + padded.push(0xFF); // Padding byte to misalign + padded.extend_from_slice(&test_data); + + let bytes = bytes::Bytes::from(padded); + let misaligned = bytes.slice(1..); // Skip first byte to create misalignment + + // Create LanceBuffer with bytes_per_value=1 to bypass alignment check + let buffer = LanceBuffer::from_bytes(misaligned, 1); + + // Verify the buffer is actually misaligned + let ptr = buffer.as_ref().as_ptr(); + assert_ne!( + ptr.align_offset(4), + 0, + "Test setup: buffer should be misaligned for u32" + ); + + // Decompress with misaligned buffer - should work with borrow_to_typed_slice + let result = decompressor.decompress(vec![buffer], 2); + assert!( + result.is_ok(), + "Decompression should succeed with misaligned buffer" + ); + + // Verify the data is correct + if let Ok(DataBlock::VariableWidth(block)) = result { + assert_eq!(block.num_values, 2); + // Data should be the strings (including padding from the original buffer) + assert_eq!(&block.data.as_ref()[..6], b"ABCXYZ"); + } else { + panic!("Expected VariableWidth block"); + } + } + + // Test case 2: u64 offsets + { + let decompressor = BinaryMiniBlockDecompressor { + bits_per_offset: 64, + }; + + // Create test data with u64 offsets + let mut test_data = Vec::new(); + + // Offsets section (3 offsets for 2 values + 1 end offset) + test_data.extend_from_slice(&24u64.to_le_bytes()); // offset to start of strings (after offsets) + test_data.extend_from_slice(&29u64.to_le_bytes()); // offset to second string + test_data.extend_from_slice(&40u64.to_le_bytes()); // offset to end (divisible by 8) + + // String data section + test_data.extend_from_slice(b"HelloWorld"); // 10 bytes of string data + test_data.extend_from_slice(&[0, 0, 0, 0, 0, 0]); // 6 bytes padding to make total 40 bytes (divisible by 8) + + // Create misaligned buffer + let mut padded = Vec::with_capacity(test_data.len() + 3); + padded.extend_from_slice(&[0xFF, 0xFF, 0xFF]); // 3 bytes padding for misalignment + padded.extend_from_slice(&test_data); + + let bytes = bytes::Bytes::from(padded); + let misaligned = bytes.slice(3..); // Skip 3 bytes + + let buffer = LanceBuffer::from_bytes(misaligned, 1); + + // Verify misalignment for u64 + let ptr = buffer.as_ref().as_ptr(); + assert_ne!( + ptr.align_offset(8), + 0, + "Test setup: buffer should be misaligned for u64" + ); + + // Decompress should succeed + let result = decompressor.decompress(vec![buffer], 2); + assert!( + result.is_ok(), + "Decompression should succeed with misaligned u64 buffer" + ); + + if let Ok(DataBlock::VariableWidth(block)) = result { + assert_eq!(block.num_values, 2); + // Data should be the strings (including padding from the original buffer) + assert_eq!(&block.data.as_ref()[..10], b"HelloWorld"); + } else { + panic!("Expected VariableWidth block"); + } + } + } + + #[test] + fn test_binary_miniblock_rejects_corrupt_offsets() { + use super::BinaryMiniBlockDecompressor; + use crate::compression::MiniBlockDecompressor; + use lance_core::Error; + + // Chunk layout mirrors the on-disk format for ["alpha", "beta", "gamma"]: + // LE u32 offsets [16, 21, 25, 30] followed by the value bytes, padded to + // a multiple of 8 bytes. + fn chunk_u32(offsets: &[u32], values: &[u8]) -> LanceBuffer { + let mut chunk = offsets + .iter() + .flat_map(|offset| offset.to_le_bytes()) + .collect::>(); + chunk.extend_from_slice(values); + chunk.resize(chunk.len().next_multiple_of(8), 0); + LanceBuffer::from(chunk) + } + + let decompressor = BinaryMiniBlockDecompressor::new(32); + + // The tail offset points past the end of the 32-byte chunk. + let err = decompressor + .decompress( + vec![chunk_u32(&[16, 21, 25, 100_000], b"alphabetagamma")], + 3, + ) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("out of bounds"), "{err}"); + + // Offsets go backwards, which would underflow the rebase subtraction. + let err = decompressor + .decompress(vec![chunk_u32(&[16, 25, 21, 30], b"alphabetagamma")], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("decreases"), "{err}"); + + // The first offset points inside the offset table, which would alias + // the serialized offsets into the value bytes. + let err = decompressor + .decompress(vec![chunk_u32(&[0, 21, 25, 30], b"alphabetagamma")], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("overlaps"), "{err}"); + + // The chunk stores fewer offsets than the requested value count needs. + let err = decompressor + .decompress(vec![chunk_u32(&[8, 8], &[])], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("requires 4"), "{err}"); + + // The chunk size is not a multiple of the offset width. + let err = decompressor + .decompress(vec![LanceBuffer::from(vec![0u8; 10])], 1) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("multiple"), "{err}"); + + // 64-bit offsets take the same validation path. + fn chunk_u64(offsets: &[u64], values: &[u8]) -> LanceBuffer { + let mut chunk = offsets + .iter() + .flat_map(|offset| offset.to_le_bytes()) + .collect::>(); + chunk.extend_from_slice(values); + chunk.resize(chunk.len().next_multiple_of(8), 0); + LanceBuffer::from(chunk) + } + let decompressor = BinaryMiniBlockDecompressor::new(64); + let err = decompressor + .decompress( + vec![chunk_u64(&[32, 37, 41, 100_000], b"alphabetagamma")], + 3, + ) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("out of bounds"), "{err}"); + let err = decompressor + .decompress(vec![chunk_u64(&[0, 37, 41, 46], b"alphabetagamma")], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("overlaps"), "{err}"); + + // A valid chunk still decodes: offsets rebase to [0, 5, 9, 14]. + let decompressor = BinaryMiniBlockDecompressor::new(32); + let block = decompressor + .decompress(vec![chunk_u32(&[16, 21, 25, 30], b"alphabetagamma")], 3) + .unwrap(); + let DataBlock::VariableWidth(block) = block else { + panic!("expected a variable-width block"); + }; + assert_eq!(block.data.as_ref(), b"alphabetagamma"); + assert_eq!( + block.offsets, + LanceBuffer::reinterpret_vec(vec![0_u32, 5, 9, 14]) + ); + } + + fn encoded_binary_block(bits_per_offset: u8) -> Vec { + use crate::compression::BlockCompressor; + + let offsets = match bits_per_offset { + 32 => LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 14]), + 64 => LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 14]), + _ => unreachable!(), + }; + let block = DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::copy_slice(b"alphabetagamma"), + offsets, + bits_per_offset, + num_values: 3, + block_info: BlockInfo::new(), + }); + BlockCompressor::compress(&super::VariableEncoder::default(), block) + .unwrap() + .0 + .as_ref() + .unwrap() + .to_vec() + } + + /// The block decompressor only checks the block structure (all O(1)); bad + /// offset values inside a structurally-sound block are rejected by the + /// mandatory layout validation when the block is converted to Arrow. + #[rstest] + #[case::i32_tail_out_of_bounds(32, 3, 100_000, "out of bounds")] + #[case::i64_tail_out_of_bounds(64, 3, 15, "out of bounds")] + #[case::i32_non_monotonic(32, 2, 4, "non-monotonic")] + #[case::i64_non_monotonic(64, 2, 4, "non-monotonic")] + fn test_binary_block_bad_offsets_rejected_at_arrow_conversion( + #[case] bits_per_offset: u8, + #[case] mutated_offset_index: usize, + #[case] mutated_offset_value: u64, + #[case] expected_message: &str, + ) { + use crate::compression::BlockDecompressor; + use lance_core::Error; + + let mut encoded = encoded_binary_block(bits_per_offset); + let bytes_per_offset = (bits_per_offset / 8) as usize; + // The standard scheme header is two offset-width fields. + let mutated_offset_start = bytes_per_offset * (2 + mutated_offset_index); + encoded[mutated_offset_start..mutated_offset_start + bytes_per_offset] + .copy_from_slice(&mutated_offset_value.to_le_bytes()[..bytes_per_offset]); + + let block = super::BinaryBlockDecompressor::default() + .decompress(Some(LanceBuffer::from(encoded)), 3) + .unwrap(); + let data_type = match bits_per_offset { + 32 => DataType::Binary, + _ => DataType::LargeBinary, + }; + let err = block.into_arrow(data_type, false).unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains(expected_message), "{err}"); + } + + #[test] + fn test_binary_block_rejects_corrupt_structure() { + use crate::compression::BlockDecompressor; + use lance_core::Error; + + let decompressor = super::BinaryBlockDecompressor::default(); + + // The first offset must be zero. + let mut encoded = encoded_binary_block(32); + encoded[8..12].copy_from_slice(&5_u32.to_le_bytes()); + let err = decompressor + .decompress(Some(LanceBuffer::from(encoded)), 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("first offset"), "{err}"); + + // The offsets region must hold exactly num_values + 1 offsets. + let encoded = encoded_binary_block(32); + let err = decompressor + .decompress(Some(LanceBuffer::from(encoded)), 4) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("offset bytes"), "{err}"); + + // A block too small to hold its header is rejected, not a panic. + let err = decompressor + .decompress(Some(LanceBuffer::from(vec![0_u8; 2])), 1) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("too small"), "{err}"); + } +} diff --git a/vendor/lance-encoding/src/encodings/physical/bitpacking.rs b/vendor/lance-encoding/src/encodings/physical/bitpacking.rs new file mode 100644 index 000000000..6c957e31e --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/bitpacking.rs @@ -0,0 +1,897 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Bitpacking encodings +//! +//! These encodings look for unused higher order bits and discard them. For example, if we +//! have a u32 array and all values are between 0 and 5000 then we only need 12 bits to store +//! each value. The encoding will discard the upper 20 bits and only store 12 bits. +//! +//! This is a simple encoding that works well for data that has a small range. +//! +//! In order to decode the values we need to know the bit width of the values. This can be stored +//! inline with the data (miniblock) or in the encoding description, out of line (full zip). +//! +//! The encoding is transparent because the output has a fixed width (just like the input) and +//! we can easily jump to the correct value. + +use arrow_array::types::UInt64Type; +use arrow_array::{Array, PrimitiveArray}; +use arrow_buffer::ArrowNativeType; +use lance_bitpacking::BitPackingUninit; + +use lance_core::{Error, Result}; + +use crate::buffer::LanceBuffer; +use crate::compression::{ + BlockCompressor, BlockDecompressor, MiniBlockDecompressor, require_block_payload, +}; +use crate::data::BlockInfo; +use crate::data::{DataBlock, FixedWidthDataBlock}; +use crate::encodings::logical::primitive::miniblock::{ + MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, +}; +use crate::format::pb21::CompressiveEncoding; +use crate::format::{ProtobufUtils21, pb21}; +use crate::statistics::{GetStat, Stat}; +use bytemuck::Pod; + +pub(crate) const LOG_ELEMS_PER_CHUNK: u8 = 10; +/// Number of values encoded in each inline bitpacking chunk. +pub const ELEMS_PER_CHUNK: u64 = 1 << LOG_ELEMS_PER_CHUNK; + +#[derive(Debug, Default)] +pub struct InlineBitpacking { + uncompressed_bit_width: u64, +} + +impl InlineBitpacking { + pub fn new(uncompressed_bit_width: u64) -> Self { + Self { + uncompressed_bit_width, + } + } + + pub fn from_description(description: &pb21::InlineBitpacking) -> Self { + Self { + uncompressed_bit_width: description.uncompressed_bits_per_value, + } + } + + /// The minimum number of bytes required to actually get compression + /// + /// We have to compress in blocks of 1024 values. For example, we can compress 500 2-byte (1000 bytes) + /// values into 1024 2-bit values (256 bytes) for a win but we don't want to compress 10 2-byte values + /// into 1024 2-bit values because that's not a win. + pub fn min_size_bytes(compressed_bit_width: u64) -> u64 { + (ELEMS_PER_CHUNK * compressed_bit_width).div_ceil(8) + } + + /// Bitpacks a FixedWidthDataBlock into compressed chunks of 1024 values + /// + /// Each chunk can have a different bit width + /// + /// Each chunk has the compressed bit width stored inline in the chunk itself. + fn bitpack_chunked( + data: FixedWidthDataBlock, + ) -> MiniBlockCompressed { + debug_assert!(data.num_values > 0); + let data_buffer = data.data.borrow_to_typed_slice::(); + let data_buffer = data_buffer.as_ref(); + + let bit_widths = data.expect_stat(Stat::BitWidth); + let bit_widths_array = bit_widths + .as_any() + .downcast_ref::>() + .unwrap(); + + let (packed_chunk_sizes, total_size) = bit_widths_array + .values() + .iter() + .map(|&bit_width| { + let chunk_size = ((1024 * bit_width) / data.bits_per_value) as usize; + (chunk_size, chunk_size + 1) + }) + .fold( + (Vec::with_capacity(bit_widths_array.len()), 0), + |(mut sizes, total), (size, inc)| { + sizes.push(size); + (sizes, total + inc) + }, + ); + + let mut output: Vec = Vec::with_capacity(total_size); + let mut chunks = Vec::with_capacity(bit_widths_array.len()); + + for (i, packed_chunk_size) in packed_chunk_sizes + .iter() + .enumerate() + .take(bit_widths_array.len() - 1) + { + let start_elem = i * ELEMS_PER_CHUNK as usize; + let bit_width = bit_widths_array.value(i) as usize; + output.push(T::from_usize(bit_width).unwrap()); + let output_len = output.len(); + unsafe { + BitPackingUninit::unchecked_pack_uninit( + bit_width, + &data_buffer[start_elem..][..ELEMS_PER_CHUNK as usize], + &mut output.spare_capacity_mut()[..*packed_chunk_size], + ); + // The bitpacking kernel initialized every reserved output word. + output.set_len(output_len + *packed_chunk_size); + } + chunks.push(MiniBlockChunk { + buffer_sizes: vec![((1 + *packed_chunk_size) * std::mem::size_of::()) as u32], + log_num_values: LOG_ELEMS_PER_CHUNK, + }); + } + + // Handle the last chunk + let last_chunk_elem_num = if data.num_values.is_multiple_of(ELEMS_PER_CHUNK) { + ELEMS_PER_CHUNK + } else { + data.num_values % ELEMS_PER_CHUNK + }; + let mut last_chunk: Vec = vec![T::from_usize(0).unwrap(); ELEMS_PER_CHUNK as usize]; + last_chunk[..last_chunk_elem_num as usize].clone_from_slice( + &data_buffer[data.num_values as usize - last_chunk_elem_num as usize..], + ); + let bit_width = bit_widths_array.value(bit_widths_array.len() - 1) as usize; + output.push(T::from_usize(bit_width).unwrap()); + let output_len = output.len(); + let packed_chunk_size = packed_chunk_sizes[bit_widths_array.len() - 1]; + unsafe { + BitPackingUninit::unchecked_pack_uninit( + bit_width, + &last_chunk, + &mut output.spare_capacity_mut()[..packed_chunk_size], + ); + // The bitpacking kernel initialized every reserved output word. + output.set_len(output_len + packed_chunk_size); + } + chunks.push(MiniBlockChunk { + buffer_sizes: vec![ + ((1 + packed_chunk_sizes[bit_widths_array.len() - 1]) * std::mem::size_of::()) + as u32, + ], + log_num_values: 0, + }); + + MiniBlockCompressed { + data: vec![LanceBuffer::reinterpret_vec(output)], + chunks, + num_values: data.num_values, + } + } + + fn chunk_data(&self, data: FixedWidthDataBlock) -> (MiniBlockCompressed, CompressiveEncoding) { + assert!(data.bits_per_value.is_multiple_of(8)); + assert_eq!(data.bits_per_value, self.uncompressed_bit_width); + let bits_per_value = data.bits_per_value; + let compressed = match bits_per_value { + 8 => Self::bitpack_chunked::(data), + 16 => Self::bitpack_chunked::(data), + 32 => Self::bitpack_chunked::(data), + 64 => Self::bitpack_chunked::(data), + _ => unreachable!(), + }; + ( + compressed, + ProtobufUtils21::inline_bitpacking( + bits_per_value, + // TODO: Could potentially compress the data here + None, + ), + ) + } + + fn unchunk( + data: LanceBuffer, + num_values: u64, + ) -> Result { + // This macro decompresses a chunk(1024 values) of bitpacked values. + let uncompressed_bit_width = std::mem::size_of::() * 8; + let word_size = std::mem::size_of::(); + + if data.len() < word_size { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking chunk is too small for {}-byte header: {} bytes", + word_size, + data.len() + ), + )); + } + if !data.len().is_multiple_of(word_size) { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking chunk size must be a multiple of {} bytes, got {} bytes", + word_size, + data.len() + ), + )); + } + if num_values > ELEMS_PER_CHUNK { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking chunk has {} values, expected at most {}", + num_values, ELEMS_PER_CHUNK + ), + )); + } + + let chunk_words = data.borrow_to_typed_view::(); + let bit_width_value = chunk_words[0].as_usize(); + if bit_width_value > uncompressed_bit_width { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking width {} exceeds {}-bit values", + bit_width_value, uncompressed_bit_width + ), + )); + } + let chunk = &chunk_words[1..]; + // bit_width_value has already been verified to be <= uncompressed_bit_width + // (8/16/32/64), so bit_width_value * ELEMS_PER_CHUNK (1024) can never + // overflow usize on supported targets. Keep checked_mul as defense in depth. + let expected_num_bits = bit_width_value + .checked_mul(ELEMS_PER_CHUNK as usize) + .ok_or_else(|| { + Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking width {} overflows chunk bit count", + bit_width_value + ), + ) + })?; + let expected_num_bytes = expected_num_bits / 8; + let actual_num_bytes = std::mem::size_of_val(chunk); + if actual_num_bytes != expected_num_bytes { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking payload has {} bytes, expected {} bytes for bit width {}", + actual_num_bytes, expected_num_bytes, bit_width_value + ), + )); + } + + let mut decompressed = Vec::with_capacity(ELEMS_PER_CHUNK as usize); + unsafe { + BitPackingUninit::unchecked_unpack_uninit( + bit_width_value, + chunk, + &mut decompressed.spare_capacity_mut()[..ELEMS_PER_CHUNK as usize], + ); + // The bitpacking kernel initialized all 1024 decoded values. + decompressed.set_len(ELEMS_PER_CHUNK as usize); + } + + decompressed.truncate(num_values as usize); + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(decompressed), + bits_per_value: uncompressed_bit_width as u64, + num_values, + block_info: BlockInfo::new(), + })) + } + + /// An empty fixed-width block, used for the `num_values == 0` short-circuit in + /// both decompressor entry points so empty blocks skip chunk validation entirely. + fn empty_block(&self) -> DataBlock { + DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::empty(), + bits_per_value: self.uncompressed_bit_width, + num_values: 0, + block_info: BlockInfo::new(), + }) + } +} + +impl MiniBlockCompressor for InlineBitpacking { + fn compress( + &self, + _context: MiniBlockCompressionContext, + chunk: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + match chunk { + DataBlock::FixedWidth(fixed_width) => Ok(self.chunk_data(fixed_width)), + _ => Err(Error::invalid_input_source( + format!( + "Cannot compress a data block of type {} with BitpackMiniBlockEncoder", + chunk.name() + ) + .into(), + )), + } + } +} + +impl BlockCompressor for InlineBitpacking { + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + let DataBlock::FixedWidth(fixed_width) = data else { + return Err(Error::invalid_input( + "Inline bitpacking requires fixed-width data", + )); + }; + if fixed_width.bits_per_value != self.uncompressed_bit_width { + return Err(Error::invalid_input(format!( + "Inline bitpacking expects {}-bit values, got {}", + self.uncompressed_bit_width, fixed_width.bits_per_value + ))); + } + let (chunked, _) = self.chunk_data(fixed_width); + let payload = + chunked.data.into_iter().next().ok_or_else(|| { + Error::internal("Inline bitpacking produced no payload".to_string()) + })?; + Ok(( + Some(payload), + ProtobufUtils21::inline_bitpacking(self.uncompressed_bit_width, None), + )) + } +} + +impl MiniBlockDecompressor for InlineBitpacking { + fn decompress(&self, data: Vec, num_values: u64) -> Result { + assert_eq!(data.len(), 1); + let data = data.into_iter().next().unwrap(); + if num_values == 0 { + // Empty mini-blocks have no inline bit-width header to decode. + return Ok(self.empty_block()); + } + match self.uncompressed_bit_width { + 8 => Self::unchunk::(data, num_values), + 16 => Self::unchunk::(data, num_values), + 32 => Self::unchunk::(data, num_values), + 64 => Self::unchunk::(data, num_values), + _ => unimplemented!("Bitpacking word size must be 8, 16, 32, or 64"), + } + } + + fn decoded_size_bytes(&self, num_values: u64) -> Option { + num_values + .checked_mul(self.uncompressed_bit_width) + .map(|bits| bits.div_ceil(8)) + } +} + +impl BlockDecompressor for InlineBitpacking { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Inline bitpacking")?; + if num_values == 0 { + // Empty blocks carry no inline bit-width header to decode; avoid + // spurious "too small for header" corrupt-file errors and mirror + // the MiniBlockDecompressor path. See #7794. + return Ok(self.empty_block()); + } + match self.uncompressed_bit_width { + 8 => Self::unchunk::(data, num_values), + 16 => Self::unchunk::(data, num_values), + 32 => Self::unchunk::(data, num_values), + 64 => Self::unchunk::(data, num_values), + _ => unimplemented!("Bitpacking word size must be 8, 16, 32, or 64"), + } + } +} + +/// Bitpacks a FixedWidthDataBlock with a given bit width. +/// +/// Each chunk of 1024 values is packed with a constant bit width. For the tail we compare the +/// cost of padding and packing against storing the raw values: if padding yields a smaller +/// representation we pack; otherwise we append the raw tail. +fn bitpack_out_of_line( + data: FixedWidthDataBlock, + compressed_bits_per_value: usize, +) -> LanceBuffer { + let data_buffer = data.data.borrow_to_typed_slice::(); + let data_buffer = data_buffer.as_ref(); + + let num_chunks = data_buffer.len().div_ceil(ELEMS_PER_CHUNK as usize); + let last_chunk_is_runt = data_buffer.len() % ELEMS_PER_CHUNK as usize != 0; + let words_per_chunk = (ELEMS_PER_CHUNK as usize * compressed_bits_per_value) + .div_ceil(data.bits_per_value as usize); + let mut output: Vec = Vec::with_capacity(num_chunks * words_per_chunk); + + let num_whole_chunks = if last_chunk_is_runt { + num_chunks - 1 + } else { + num_chunks + }; + + // Simple case for complete chunks + for i in 0..num_whole_chunks { + let input_start = i * ELEMS_PER_CHUNK as usize; + let input_end = input_start + ELEMS_PER_CHUNK as usize; + let output_start = output.len(); + unsafe { + BitPackingUninit::unchecked_pack_uninit( + compressed_bits_per_value, + &data_buffer[input_start..input_end], + &mut output.spare_capacity_mut()[..words_per_chunk], + ); + // The bitpacking kernel initialized this complete packed chunk. + output.set_len(output_start + words_per_chunk); + } + } + + if !last_chunk_is_runt { + return LanceBuffer::reinterpret_vec(output); + } + + let last_chunk_start = num_whole_chunks * ELEMS_PER_CHUNK as usize; + let remaining_items = data_buffer.len() - last_chunk_start; + + let uncompressed_bits = data.bits_per_value as usize; + let tail_bit_savings = uncompressed_bits.saturating_sub(compressed_bits_per_value); + let padding_cost = compressed_bits_per_value * (ELEMS_PER_CHUNK as usize - remaining_items); + let tail_pack_savings = tail_bit_savings.saturating_mul(remaining_items); + debug_assert!(remaining_items > 0, "remaining_items must be non-zero"); + debug_assert!(tail_bit_savings > 0, "tail_bit_savings must be non-zero"); + + if padding_cost < tail_pack_savings { + // Padding buys us more than it costs: pad to 1024 values and pack them as a normal chunk. + let mut last_chunk: Vec = vec![T::from_usize(0).unwrap(); ELEMS_PER_CHUNK as usize]; + last_chunk[..remaining_items].copy_from_slice(&data_buffer[last_chunk_start..]); + let start = output.len(); + unsafe { + BitPackingUninit::unchecked_pack_uninit( + compressed_bits_per_value, + &last_chunk, + &mut output.spare_capacity_mut()[..words_per_chunk], + ); + // The bitpacking kernel initialized the padded tail chunk. + output.set_len(start + words_per_chunk); + } + } else { + // Padding would waste space; append tail values as-is. + output.extend_from_slice(&data_buffer[last_chunk_start..]); + } + + LanceBuffer::reinterpret_vec(output) +} + +/// Unpacks a FixedWidthDataBlock that has been bitpacked with a constant bit width. +/// +/// The compressed bit width is provided while the uncompressed width comes from `T`. +/// Depending on the encoding decision the final chunk may be fully packed (with padding) +/// or stored as raw tail values. We infer the layout from the buffer length. +fn unpack_out_of_line( + data: FixedWidthDataBlock, + num_values: usize, + compressed_bits_per_value: usize, +) -> FixedWidthDataBlock { + let words_per_chunk = (ELEMS_PER_CHUNK as usize * compressed_bits_per_value) + .div_ceil(data.bits_per_value as usize); + let compressed_words = data.data.borrow_to_typed_slice::(); + + let num_whole_chunks = num_values / ELEMS_PER_CHUNK as usize; + let tail_values = num_values % ELEMS_PER_CHUNK as usize; + let expected_full_words = num_whole_chunks * words_per_chunk; + let expected_new_len = expected_full_words + tail_values; + let tail_is_raw = tail_values > 0 && compressed_words.len() == expected_new_len; + + let extra_tail_capacity = ELEMS_PER_CHUNK as usize; + let mut decompressed: Vec = + Vec::with_capacity(num_values.saturating_add(extra_tail_capacity)); + + for chunk_idx in 0..num_whole_chunks { + let input_start = chunk_idx * words_per_chunk; + let input_end = input_start + words_per_chunk; + let output_start = decompressed.len(); + unsafe { + BitPackingUninit::unchecked_unpack_uninit( + compressed_bits_per_value, + &compressed_words[input_start..input_end], + &mut decompressed.spare_capacity_mut()[..ELEMS_PER_CHUNK as usize], + ); + // The bitpacking kernel initialized this complete decoded chunk. + decompressed.set_len(output_start + ELEMS_PER_CHUNK as usize); + } + } + + if tail_values > 0 { + // The tail might be padded and bit packed or it might be appended raw. We infer the + // layout from the buffer length to decode appropriately. + if tail_is_raw { + let tail_start = expected_full_words; + decompressed.extend_from_slice(&compressed_words[tail_start..tail_start + tail_values]); + } else { + let tail_start = expected_full_words; + let output_start = decompressed.len(); + unsafe { + BitPackingUninit::unchecked_unpack_uninit( + compressed_bits_per_value, + &compressed_words[tail_start..tail_start + words_per_chunk], + &mut decompressed.spare_capacity_mut()[..ELEMS_PER_CHUNK as usize], + ); + // The kernel initialized a full chunk; only the requested tail stays visible. + decompressed.set_len(output_start + ELEMS_PER_CHUNK as usize); + } + decompressed.truncate(output_start + tail_values); + } + } + + debug_assert_eq!(decompressed.len(), num_values); + + FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(decompressed), + bits_per_value: data.bits_per_value, + num_values: num_values as u64, + block_info: BlockInfo::new(), + } +} + +/// A transparent compressor that bit packs data +/// +/// In order for the encoding to be transparent we must have a fixed bit width +/// across the entire array. Chunking within the buffer is not supported. This +/// means that we will be slightly less efficient than something like the mini-block +/// approach. +/// +/// This was an interesting experiment but it can't be used as a per-value compressor +/// at the moment. The resulting data IS transparent but it's not quite so simple. We +/// compress in blocks of 1024 and each block has a fixed size but also has some padding. +/// +/// We do use this as a block compressor currently. +/// +/// In other words, if we try the simple math to access the item at index `i` we will be +/// out of luck because `bits_per_value * i` is not the location. What we need is something +/// like: +/// +/// ```ignore +/// let chunk_idx = i / 1024; +/// let chunk_offset = i % 1024; +/// bits_per_chunk * chunk_idx + bits_per_value * chunk_offset +/// ``` +/// +/// However, this logic isn't expressible with the per-value traits we have today. We can +/// enhance these traits should we need to support it at some point in the future. +#[derive(Debug)] +pub struct OutOfLineBitpacking { + compressed_bit_width: u64, + uncompressed_bit_width: u64, +} + +impl OutOfLineBitpacking { + pub fn new(compressed_bit_width: u64, uncompressed_bit_width: u64) -> Self { + Self { + compressed_bit_width, + uncompressed_bit_width, + } + } +} + +impl BlockCompressor for OutOfLineBitpacking { + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + let DataBlock::FixedWidth(fixed_width) = data else { + return Err(Error::invalid_input( + "Out-of-line bitpacking requires fixed-width data", + )); + }; + if fixed_width.bits_per_value != self.uncompressed_bit_width { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking expects {}-bit values, got {}", + self.uncompressed_bit_width, fixed_width.bits_per_value + ))); + } + let compressed = match fixed_width.bits_per_value { + 8 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), + 16 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), + 32 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), + 64 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), + _ => { + return Err(Error::invalid_input(format!( + "Bitpacking word size must be 8, 16, 32, or 64, got {}", + fixed_width.bits_per_value + ))); + } + }; + Ok(( + Some(compressed), + ProtobufUtils21::out_of_line_bitpacking( + self.uncompressed_bit_width, + ProtobufUtils21::flat(self.compressed_bit_width, None), + ), + )) + } +} + +impl BlockDecompressor for OutOfLineBitpacking { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Out-of-line bitpacking")?; + let word_size = match self.uncompressed_bit_width { + 8 => std::mem::size_of::(), + 16 => std::mem::size_of::(), + 32 => std::mem::size_of::(), + 64 => std::mem::size_of::(), + _ => panic!("Bitpacking word size must be 8,16,32,64"), + }; + debug_assert_eq!(data.len() % word_size, 0); + let total_words = (data.len() / word_size) as u64; + let block = FixedWidthDataBlock { + data, + bits_per_value: self.uncompressed_bit_width, + num_values: total_words, + block_info: BlockInfo::new(), + }; + + let unpacked = match self.uncompressed_bit_width { + 8 => unpack_out_of_line::( + block, + num_values as usize, + self.compressed_bit_width as usize, + ), + 16 => unpack_out_of_line::( + block, + num_values as usize, + self.compressed_bit_width as usize, + ), + 32 => unpack_out_of_line::( + block, + num_values as usize, + self.compressed_bit_width as usize, + ), + 64 => unpack_out_of_line::( + block, + num_values as usize, + self.compressed_bit_width as usize, + ), + _ => unreachable!(), + }; + Ok(DataBlock::FixedWidth(unpacked)) + } +} + +#[cfg(test)] +mod test { + use std::{collections::HashMap, sync::Arc}; + + use arrow_array::{Array, Int8Array, Int64Array}; + use arrow_buffer::ArrowNativeType; + use arrow_schema::DataType; + use bytemuck::Pod; + use lance_bitpacking::{BitPacking, BitPackingUninit}; + use rstest::rstest; + + use super::{ELEMS_PER_CHUNK, InlineBitpacking, bitpack_out_of_line, unpack_out_of_line}; + use crate::{ + buffer::LanceBuffer, + compression::{BlockDecompressor, MiniBlockDecompressor}, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + testing::{TestCases, check_round_trip_encoding_of_data}, + }; + + #[rstest] + #[case::u8(8)] + #[case::u16(16)] + #[case::u32(32)] + #[case::u64(64)] + fn test_inline_bitpacking_decompress_empty_miniblock(#[case] bit_width: u64) { + let decompressor = InlineBitpacking::new(bit_width); + let decompressed = + MiniBlockDecompressor::decompress(&decompressor, vec![LanceBuffer::empty()], 0) + .unwrap(); + + let DataBlock::FixedWidth(block) = decompressed else { + panic!("Expected FixedWidth block"); + }; + assert_eq!(block.bits_per_value, bit_width); + assert_eq!(block.num_values, 0); + assert_eq!(block.data.len(), 0); + } + + // Regression test for #7794: the block-level decompressor must short-circuit + // on num_values == 0 the same way the mini-block decompressor does, instead + // of reporting a spurious "too small for header" corrupt-file error. + #[rstest] + #[case::u8(8)] + #[case::u16(16)] + #[case::u32(32)] + #[case::u64(64)] + fn test_inline_bitpacking_decompress_empty_block(#[case] bit_width: u64) { + let decompressor = InlineBitpacking::new(bit_width); + let decompressed = + BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::empty()), 0).unwrap(); + + let DataBlock::FixedWidth(block) = decompressed else { + panic!("Expected FixedWidth block"); + }; + assert_eq!(block.bits_per_value, bit_width); + assert_eq!(block.num_values, 0); + assert_eq!(block.data.len(), 0); + } + + fn roundtrip_unchunk(values: &[T], bit_width: usize) + where + T: ArrowNativeType + BitPackingUninit + Pod, + { + assert!(values.len() <= ELEMS_PER_CHUNK as usize); + let num_values = values.len() as u64; + + let mut padded = vec![T::from_usize(0).unwrap(); ELEMS_PER_CHUNK as usize]; + padded[..values.len()].copy_from_slice(values); + + let packed_words = ELEMS_PER_CHUNK as usize * bit_width / (std::mem::size_of::() * 8); + let mut chunk: Vec = Vec::with_capacity(1 + packed_words); + chunk.push(T::from_usize(bit_width).unwrap()); + let out_len = chunk.len(); + chunk.resize(out_len + packed_words, T::from_usize(0).unwrap()); + unsafe { + BitPacking::unchecked_pack(bit_width, &padded, &mut chunk[out_len..]); + } + + let data = LanceBuffer::reinterpret_vec(chunk); + let decoded = InlineBitpacking::unchunk::(data, num_values).unwrap(); + let DataBlock::FixedWidth(fixed) = decoded else { + panic!("expected FixedWidth DataBlock"); + }; + let decoded_values = fixed.data.borrow_to_typed_view::(); + assert_eq!(decoded_values.as_ref(), values); + } + + fn assert_corrupt_unchunk(data: LanceBuffer, num_values: u64, expected_message: &str) + where + T: ArrowNativeType + BitPackingUninit + Pod, + { + let err = InlineBitpacking::unchunk::(data, num_values).unwrap_err(); + assert!(matches!(err, lance_core::Error::CorruptFile { .. })); + let err = err.to_string(); + assert!( + err.contains(expected_message), + "expected error containing {expected_message:?}, got {err:?}" + ); + } + + #[test] + fn unchunk_u32_bw12_tail() { + let values: Vec = (0..500).map(|i| ((i * 7) % (1 << 12)) as u32).collect(); + roundtrip_unchunk(&values, 12); + } + + #[test] + fn unchunk_u64_bw23_full() { + let values: Vec = (0..1024).map(|i| ((i * 3) % (1 << 23)) as u64).collect(); + roundtrip_unchunk(&values, 23); + } + + #[rstest] + #[case::too_small_header(LanceBuffer::from(vec![1, 2, 3]), 1, "too small")] + #[case::misaligned_chunk_size(LanceBuffer::from(vec![0, 0, 0, 0, 0]), 1, "multiple")] + #[case::too_many_values( + LanceBuffer::reinterpret_vec(vec![0_u32]), + ELEMS_PER_CHUNK + 1, + "expected at most" + )] + #[case::payload_size_mismatch(LanceBuffer::reinterpret_vec(vec![12_u32]), 1, "payload")] + #[case::invalid_bit_width(LanceBuffer::reinterpret_vec(vec![33_u32]), 1, "exceeds")] + fn unchunk_rejects( + #[case] data: LanceBuffer, + #[case] num_values: u64, + #[case] expected_message: &str, + ) { + assert_corrupt_unchunk::(data, num_values, expected_message); + } + + #[test_log::test(tokio::test)] + async fn test_miniblock_bitpack() { + let test_cases = TestCases::default().with_structural_encodings(); + + let arrays = vec![ + Arc::new(Int8Array::from(vec![100; 1024])) as Arc, + Arc::new(Int8Array::from(vec![1; 1024])) as Arc, + Arc::new(Int8Array::from(vec![16; 1024])) as Arc, + Arc::new(Int8Array::from(vec![-1; 1024])) as Arc, + Arc::new(Int8Array::from(vec![5; 1])) as Arc, + ]; + check_round_trip_encoding_of_data(arrays, &test_cases, HashMap::new()).await; + + for data_type in [DataType::Int16, DataType::Int32, DataType::Int64] { + let int64_arrays = vec![ + Int64Array::from(vec![3; 1024]), + Int64Array::from(vec![8; 1024]), + Int64Array::from(vec![16; 1024]), + Int64Array::from(vec![100; 1024]), + Int64Array::from(vec![512; 1024]), + Int64Array::from(vec![1000; 1024]), + Int64Array::from(vec![2000; 1024]), + Int64Array::from(vec![-1; 10]), + ]; + + let mut arrays = vec![]; + for int64_array in int64_arrays { + arrays.push(arrow_cast::cast(&int64_array, &data_type).unwrap()); + } + + check_round_trip_encoding_of_data(arrays, &test_cases, HashMap::new()).await; + } + } + + #[test_log::test(tokio::test)] + async fn test_bitpack_encoding_verification() { + use arrow_array::Int32Array; + + // Test bitpacking encoding verification with varied small values that should trigger bitpacking + let test_cases = TestCases::default() + .with_expected_encoding("inline_bitpacking") + .with_structural_encodings(); + + // Generate data with varied small values to avoid RLE + // Mix different values but keep them small to trigger bitpacking + let mut values = Vec::new(); + for i in 0..2048 { + values.push(i % 16); // Values 0-15, varied enough to avoid RLE + } + + let arrays = vec![Arc::new(Int32Array::from(values)) as Arc]; + + // Explicitly disable BSS to ensure bitpacking is tested + let mut metadata = HashMap::new(); + metadata.insert("lance-encoding:bss".to_string(), "off".to_string()); + + check_round_trip_encoding_of_data(arrays, &test_cases, metadata.clone()).await; + } + + #[test_log::test(tokio::test)] + async fn test_miniblock_bitpack_zero_chunk_selection() { + use arrow_array::Int32Array; + + let test_cases = TestCases::default() + .with_expected_encoding("inline_bitpacking") + .with_structural_encodings(); + + // Build 2048 values: first 1024 all zeros (bit_width=0), + // next 1024 small varied values to avoid RLE and trigger bitpacking. + let mut vals = vec![0i32; 1024]; + for i in 0..1024 { + vals.push(i % 16); + } + + let arrays = vec![Arc::new(Int32Array::from(vals)) as Arc]; + + // Disable BSS and RLE to prefer bitpacking in selection + let mut metadata = HashMap::new(); + metadata.insert("lance-encoding:bss".to_string(), "off".to_string()); + metadata.insert("lance-encoding:rle-threshold".to_string(), "0".to_string()); + + check_round_trip_encoding_of_data(arrays, &test_cases, metadata).await; + } + + #[test] + fn test_out_of_line_bitpack_raw_tail_roundtrip() { + let bit_width = 8usize; + let word_bits = std::mem::size_of::() as u64 * 8; + let values: Vec = (0..1025).map(|i| (i % 200) as u32).collect(); + let input = FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(values.clone()), + bits_per_value: word_bits, + num_values: values.len() as u64, + block_info: BlockInfo::new(), + }; + + let compressed = bitpack_out_of_line::(input, bit_width); + let compressed_words = compressed.borrow_to_typed_slice::().to_vec(); + let words_per_chunk = (ELEMS_PER_CHUNK as usize * bit_width).div_ceil(word_bits as usize); + assert_eq!( + compressed_words.len(), + words_per_chunk + (values.len() - ELEMS_PER_CHUNK as usize), + ); + + let compressed_block = FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(compressed_words.clone()), + bits_per_value: word_bits, + num_values: compressed_words.len() as u64, + block_info: BlockInfo::new(), + }; + + let decoded = unpack_out_of_line::(compressed_block, values.len(), bit_width); + let decoded_values = decoded.data.borrow_to_typed_slice::(); + assert_eq!(decoded_values.as_ref(), values.as_slice()); + } +} diff --git a/vendor/lance-encoding/src/encodings/physical/block.rs b/vendor/lance-encoding/src/encodings/physical/block.rs new file mode 100644 index 000000000..9d97c9e72 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/block.rs @@ -0,0 +1,884 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Encodings based on traditional block compression schemes +//! +//! Traditional compressors take in a buffer and return a smaller buffer. All encoding +//! description is shoved into the compressed buffer and the entire buffer is needed to +//! decompress any of the data. +//! +//! These encodings are not transparent, which limits our ability to use them. In addition +//! they are often quite expensive in CPU terms. +//! +//! However, they are effective and useful for some cases. For example, when working with large +//! variable length values (e.g. source code files) they can be very effective. +//! +//! The module introduces the `[BufferCompressor]` trait which describes the interface for a +//! traditional block compressor. It is implemented for the most common compression schemes +//! (zstd, lz4, etc). +//! +//! There is not yet a mini-block variant of this compressor (but could easily be one) and the +//! full zip variant works by applying compression on a per-value basis (which allows it to be +//! transparent). + +use arrow_buffer::ArrowNativeType; +use lance_core::{Error, Result}; + +use std::str::FromStr; + +use crate::compression::{BlockCompressor, BlockDecompressor, require_block_payload}; +use crate::encodings::physical::binary::{BinaryBlockDecompressor, VariableEncoder}; +use crate::format::{ + ProtobufUtils21, + pb21::{self, CompressiveEncoding}, +}; +use crate::{ + buffer::LanceBuffer, + compression::VariablePerValueDecompressor, + data::{BlockInfo, DataBlock, VariableWidthBlock}, + encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}, +}; + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct CompressionConfig { + pub(crate) scheme: CompressionScheme, + pub(crate) level: Option, +} + +impl CompressionConfig { + /// Create a compression configuration for an encoding mechanism. + pub fn new(scheme: CompressionScheme, level: Option) -> Self { + Self { scheme, level } + } + + /// Return the selected compression scheme. + pub fn scheme(&self) -> CompressionScheme { + self.scheme + } + + /// Return the optional compression level. + pub fn level(&self) -> Option { + self.level + } +} + +impl Default for CompressionConfig { + fn default() -> Self { + Self { + scheme: CompressionScheme::Lz4, + level: Some(0), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum CompressionScheme { + None, + Fsst, + Zstd, + Lz4, +} + +impl TryFrom for pb21::CompressionScheme { + type Error = Error; + + fn try_from(scheme: CompressionScheme) -> Result { + match scheme { + CompressionScheme::Lz4 => Ok(Self::CompressionAlgorithmLz4), + CompressionScheme::Zstd => Ok(Self::CompressionAlgorithmZstd), + _ => Err(Error::invalid_input(format!( + "Unsupported compression scheme: {:?}", + scheme + ))), + } + } +} + +impl TryFrom for CompressionScheme { + type Error = Error; + + fn try_from(scheme: pb21::CompressionScheme) -> Result { + match scheme { + pb21::CompressionScheme::CompressionAlgorithmLz4 => Ok(Self::Lz4), + pb21::CompressionScheme::CompressionAlgorithmZstd => Ok(Self::Zstd), + _ => Err(Error::invalid_input(format!( + "Unsupported compression scheme: {:?}", + scheme + ))), + } + } +} + +impl std::fmt::Display for CompressionScheme { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + let scheme_str = match self { + Self::Fsst => "fsst", + Self::Zstd => "zstd", + Self::None => "none", + Self::Lz4 => "lz4", + }; + write!(f, "{}", scheme_str) + } +} + +impl FromStr for CompressionScheme { + type Err = Error; + + fn from_str(s: &str) -> Result { + match s { + "none" => Ok(Self::None), + "fsst" => Ok(Self::Fsst), + "zstd" => Ok(Self::Zstd), + "lz4" => Ok(Self::Lz4), + _ => Err(Error::invalid_input(format!( + "Unknown compression scheme: {}", + s + ))), + } + } +} + +pub trait BufferCompressor: std::fmt::Debug + Send + Sync { + fn compress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()>; + fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()>; + fn config(&self) -> CompressionConfig; +} + +#[cfg(feature = "zstd")] +mod zstd { + use std::io::{Cursor, Write}; + use std::sync::{Mutex, OnceLock}; + + use super::*; + + use ::zstd::bulk::{Compressor, decompress_to_buffer}; + use ::zstd::stream::copy_decode; + + /// A zstd buffer compressor that lazily creates and reuses compression contexts. + /// + /// The compression context is cached to enable reuse across chunks within a + /// page. It is lazily initialized to prevent it from getting initialized on + /// decode-only codepaths. + /// + /// Reuse is not implemented for decompression, only for compression: + /// * The single-threaded benefit of reuse was negligible when measured. + /// * Decompressors can get shared across threads, leading to mutex + /// contention if the same strategy is used as for compression here. This + /// should be mitigable with pooling but we can skip the complexity until a + /// need is demonstrated. The multithreaded decode benchmark effectively + /// demonstrates this scenario. + pub struct ZstdBufferCompressor { + compression_level: i32, + compressor: OnceLock>, String>>, + } + + impl std::fmt::Debug for ZstdBufferCompressor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ZstdBufferCompressor") + .field("compression_level", &self.compression_level) + .finish() + } + } + + impl ZstdBufferCompressor { + pub fn new(compression_level: i32) -> Self { + Self { + compression_level, + compressor: OnceLock::new(), + } + } + + fn get_compressor(&self) -> Result<&Mutex>> { + self.compressor + .get_or_init(|| { + Compressor::new(self.compression_level) + .map(Mutex::new) + .map_err(|e| e.to_string()) + }) + .as_ref() + .map_err(|e| Error::internal(format!("Failed to create zstd compressor: {}", e))) + } + + // https://datatracker.ietf.org/doc/html/rfc8878 + fn is_raw_stream_format(&self, input_buf: &[u8]) -> bool { + if input_buf.len() < 8 { + return true; // can't be length prefixed format if less than 8 bytes + } + // read the first 4 bytes as the magic number + let mut magic_buf = [0u8; 4]; + magic_buf.copy_from_slice(&input_buf[..4]); + let magic = u32::from_le_bytes(magic_buf); + + // see RFC 8878, section 3.1.1. Zstandard Frames, which defines the magic number + const ZSTD_MAGIC_NUMBER: u32 = 0xFD2FB528; + if magic == ZSTD_MAGIC_NUMBER { + // the compressed buffer starts like a Zstd frame. + // Per RFC 8878, the reserved bit (with Bit Number 3, the 4th bit) in the FHD (frame header descriptor) MUST be 0 + // see section 3.1.1.1.1. 'Frame_Header_Descriptor' and section 3.1.1.1.1.4. 'Reserved Bit' for details + const FHD_BYTE_INDEX: usize = 4; + let fhd_byte = input_buf[FHD_BYTE_INDEX]; + const FHD_RESERVED_BIT_MASK: u8 = 0b0001_0000; + let reserved_bit = fhd_byte & FHD_RESERVED_BIT_MASK; + + if reserved_bit != 0 { + // this bit is 1. This is NOT a valid zstd frame. + // therefore, it must be length prefixed format where the length coincidentally + // started with the magic number + false + } else { + // the reserved bit is 0. This is consistent with a valid Zstd frame. + // treat it as raw stream format + true + } + } else { + // doesn't start with the magic number, so it can't be the raw stream format + false + } + } + + fn decompress_length_prefixed_zstd( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + ) -> Result<()> { + const LENGTH_PREFIX_SIZE: usize = 8; + let mut len_buf = [0u8; LENGTH_PREFIX_SIZE]; + len_buf.copy_from_slice(&input_buf[..LENGTH_PREFIX_SIZE]); + + let uncompressed_len = u64::from_le_bytes(len_buf) as usize; + + let start = output_buf.len(); + output_buf.resize(start + uncompressed_len, 0); + + let compressed_data = &input_buf[LENGTH_PREFIX_SIZE..]; + decompress_to_buffer(compressed_data, &mut output_buf[start..])?; + Ok(()) + } + } + + impl BufferCompressor for ZstdBufferCompressor { + fn compress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + output_buf.write_all(&(input_buf.len() as u64).to_le_bytes())?; + + let max_compressed_size = ::zstd::zstd_safe::compress_bound(input_buf.len()); + let start_pos = output_buf.len(); + output_buf.resize(start_pos + max_compressed_size, 0); + + let compressed_size = self + .get_compressor()? + .lock() + .unwrap() + .compress_to_buffer(input_buf, &mut output_buf[start_pos..]) + .map_err(|e| Error::internal(format!("Zstd compression error: {}", e)))?; + + output_buf.truncate(start_pos + compressed_size); + Ok(()) + } + + fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + if input_buf.is_empty() { + return Ok(()); + } + + let is_raw_stream_format = self.is_raw_stream_format(input_buf); + if is_raw_stream_format { + copy_decode(Cursor::new(input_buf), output_buf)?; + } else { + self.decompress_length_prefixed_zstd(input_buf, output_buf)?; + } + + Ok(()) + } + + fn config(&self) -> CompressionConfig { + CompressionConfig { + scheme: CompressionScheme::Zstd, + level: Some(self.compression_level), + } + } + } +} + +#[cfg(feature = "lz4")] +mod lz4 { + use super::*; + + #[derive(Debug, Default)] + pub struct Lz4BufferCompressor {} + + impl BufferCompressor for Lz4BufferCompressor { + fn compress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + // Remember the starting position + let start_pos = output_buf.len(); + + // LZ4 needs space for the compressed data + let max_size = ::lz4::block::compress_bound(input_buf.len())?; + // Resize to ensure we have enough space (including 4 bytes for size header) + output_buf.resize(start_pos + max_size + 4, 0); + + let compressed_size = ::lz4::block::compress_to_buffer( + input_buf, + None, + true, + &mut output_buf[start_pos..], + ) + .map_err(|err| Error::internal(format!("LZ4 compression error: {}", err)))?; + + // Truncate to actual size + output_buf.truncate(start_pos + compressed_size); + Ok(()) + } + + fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + // When prepend_size is true, LZ4 stores the uncompressed size in the first 4 bytes + // We can read this to know exactly how much space we need + if input_buf.len() < 4 { + return Err(Error::internal("LZ4 compressed data too short".to_string())); + } + + // Read the uncompressed size from the first 4 bytes (little-endian) + let uncompressed_size = + u32::from_le_bytes([input_buf[0], input_buf[1], input_buf[2], input_buf[3]]) + as usize; + + // Remember the starting position + let start_pos = output_buf.len(); + + // Resize to ensure we have the exact space needed + output_buf.resize(start_pos + uncompressed_size, 0); + + // Now decompress directly into the buffer slice + let decompressed_size = + ::lz4::block::decompress_to_buffer(input_buf, None, &mut output_buf[start_pos..]) + .map_err(|err| Error::internal(format!("LZ4 decompression error: {}", err)))?; + + // Truncate to actual decompressed size (should be same as uncompressed_size) + output_buf.truncate(start_pos + decompressed_size); + + Ok(()) + } + + fn config(&self) -> CompressionConfig { + CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + } + } + } +} + +#[derive(Debug, Default)] +pub struct NoopBufferCompressor {} + +impl BufferCompressor for NoopBufferCompressor { + fn compress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + output_buf.extend_from_slice(input_buf); + Ok(()) + } + + fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + output_buf.extend_from_slice(input_buf); + Ok(()) + } + + fn config(&self) -> CompressionConfig { + CompressionConfig { + scheme: CompressionScheme::None, + level: None, + } + } +} + +pub struct GeneralBufferCompressor {} + +impl GeneralBufferCompressor { + pub fn get_compressor( + compression_config: CompressionConfig, + ) -> Result> { + match compression_config.scheme { + // FSST has its own compression path and isn't implemented as a generic buffer compressor + CompressionScheme::Fsst => Err(Error::invalid_input_source( + "fsst is not usable as a general buffer compressor".into(), + )), + CompressionScheme::Zstd => { + #[cfg(feature = "zstd")] + { + Ok(Box::new(zstd::ZstdBufferCompressor::new( + compression_config.level.unwrap_or(0), + ))) + } + #[cfg(not(feature = "zstd"))] + { + Err(Error::invalid_input_source( + "package was not built with zstd support".into(), + )) + } + } + CompressionScheme::Lz4 => { + #[cfg(feature = "lz4")] + { + Ok(Box::new(lz4::Lz4BufferCompressor::default())) + } + #[cfg(not(feature = "lz4"))] + { + Err(Error::invalid_input_source( + "package was not built with lz4 support".into(), + )) + } + } + CompressionScheme::None => Ok(Box::new(NoopBufferCompressor {})), + } + } +} + +/// A block decompressor that first applies general-purpose compression (LZ4/Zstd) +/// before delegating to an inner block decompressor. +#[derive(Debug)] +pub struct GeneralBlockDecompressor { + inner: Box, + compressor: Box, +} + +impl GeneralBlockDecompressor { + pub fn try_new( + inner: Box, + compression: CompressionConfig, + ) -> Result { + let compressor = GeneralBufferCompressor::get_compressor(compression)?; + Ok(Self { inner, compressor }) + } +} + +impl BlockDecompressor for GeneralBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "General block compression")?; + let mut decompressed = Vec::new(); + self.compressor.decompress(&data, &mut decompressed)?; + self.inner + .decompress(Some(LanceBuffer::from(decompressed)), num_values) + } +} + +// An encoder which uses generic compression, such as zstd/lz4 to encode buffers +#[derive(Debug)] +pub struct CompressedBufferEncoder { + pub(crate) compressor: Box, + // Runtime compressors normalize levels that they default or ignore. Block descriptors must + // retain the selected configuration so stable writers preserve those present/absent values. + block_compression: CompressionConfig, +} + +impl Default for CompressedBufferEncoder { + fn default() -> Self { + // Pick zstd if available, otherwise lz4, otherwise none + #[cfg(feature = "zstd")] + let (scheme, level) = (CompressionScheme::Zstd, Some(0)); + #[cfg(all(feature = "lz4", not(feature = "zstd")))] + let (scheme, level) = (CompressionScheme::Lz4, None); + #[cfg(not(any(feature = "zstd", feature = "lz4")))] + let (scheme, level) = (CompressionScheme::None, None); + + let block_compression = CompressionConfig { scheme, level }; + let compressor = GeneralBufferCompressor::get_compressor(block_compression).unwrap(); + Self { + compressor, + block_compression, + } + } +} + +impl CompressedBufferEncoder { + pub fn try_new(compression_config: CompressionConfig) -> Result { + let compressor = GeneralBufferCompressor::get_compressor(compression_config)?; + Ok(Self { + compressor, + block_compression: compression_config, + }) + } + + pub fn from_scheme(scheme: pb21::CompressionScheme) -> Result { + let scheme = CompressionScheme::try_from(scheme)?; + let block_compression = CompressionConfig { + scheme, + level: Some(0), + }; + Ok(Self { + compressor: GeneralBufferCompressor::get_compressor(block_compression)?, + block_compression, + }) + } +} + +impl CompressedBufferEncoder { + pub fn per_value_compress( + &self, + data: &[u8], + offsets: &[T], + compressed: &mut Vec, + ) -> Result { + let mut new_offsets: Vec = Vec::with_capacity(offsets.len()); + new_offsets.push(T::from_usize(0).unwrap()); + + for off in offsets.windows(2) { + let start = off[0].as_usize(); + let end = off[1].as_usize(); + self.compressor.compress(&data[start..end], compressed)?; + new_offsets.push(T::from_usize(compressed.len()).unwrap()); + } + + Ok(LanceBuffer::reinterpret_vec(new_offsets)) + } + + pub fn per_value_decompress( + &self, + data: &[u8], + offsets: &[T], + decompressed: &mut Vec, + ) -> Result { + let mut new_offsets: Vec = Vec::with_capacity(offsets.len()); + new_offsets.push(T::from_usize(0).unwrap()); + + for off in offsets.windows(2) { + let start = off[0].as_usize(); + let end = off[1].as_usize(); + self.compressor + .decompress(&data[start..end], decompressed)?; + new_offsets.push(T::from_usize(decompressed.len()).unwrap()); + } + + Ok(LanceBuffer::reinterpret_vec(new_offsets)) + } +} + +impl PerValueCompressor for CompressedBufferEncoder { + fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> { + let data_type = data.name(); + let data = data.as_variable_width().ok_or(Error::internal(format!( + "Attempt to use CompressedBufferEncoder on data of type {}", + data_type + )))?; + + let data_bytes = &data.data; + let mut compressed = Vec::with_capacity(data_bytes.len()); + + let new_offsets = match data.bits_per_offset { + 32 => self.per_value_compress::( + data_bytes, + &data.offsets.borrow_to_typed_slice::(), + &mut compressed, + )?, + 64 => self.per_value_compress::( + data_bytes, + &data.offsets.borrow_to_typed_slice::(), + &mut compressed, + )?, + _ => unreachable!(), + }; + + let compressed = PerValueDataBlock::Variable(VariableWidthBlock { + bits_per_offset: data.bits_per_offset, + data: LanceBuffer::from(compressed), + offsets: new_offsets, + num_values: data.num_values, + block_info: BlockInfo::new(), + }); + + // TODO: Support setting the level + // TODO: Support underlying compression of data (e.g. defer to binary encoding for offset bitpacking) + let encoding = ProtobufUtils21::wrapped( + self.compressor.config(), + ProtobufUtils21::variable( + ProtobufUtils21::flat(data.bits_per_offset as u64, None), + None, + ), + )?; + + Ok((compressed, encoding)) + } +} + +impl VariablePerValueDecompressor for CompressedBufferEncoder { + fn decompress(&self, data: VariableWidthBlock) -> Result { + let data_bytes = &data.data; + let mut decompressed = Vec::with_capacity(data_bytes.len() * 2); + + let new_offsets = match data.bits_per_offset { + 32 => self.per_value_decompress( + data_bytes, + &data.offsets.borrow_to_typed_slice::(), + &mut decompressed, + )?, + 64 => self.per_value_decompress( + data_bytes, + &data.offsets.borrow_to_typed_slice::(), + &mut decompressed, + )?, + _ => unreachable!(), + }; + Ok(DataBlock::VariableWidth(VariableWidthBlock { + bits_per_offset: data.bits_per_offset, + data: LanceBuffer::from(decompressed), + offsets: new_offsets, + num_values: data.num_values, + block_info: BlockInfo::new(), + })) + } +} + +impl BlockCompressor for CompressedBufferEncoder { + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + let (encoded, inner_encoding) = match data { + DataBlock::FixedWidth(fixed_width) => ( + fixed_width.data, + ProtobufUtils21::flat(fixed_width.bits_per_value, None), + ), + DataBlock::VariableWidth(variable_width) => { + // Wrap VariableEncoder to handle the encoding + let encoder = VariableEncoder::default(); + let (payload, encoding) = + BlockCompressor::compress(&encoder, DataBlock::VariableWidth(variable_width))?; + ( + payload.ok_or_else(|| { + Error::internal( + "VariableEncoder returned no payload for general compression" + .to_string(), + ) + })?, + encoding, + ) + } + _ => { + return Err(Error::invalid_input_source( + "Unsupported data block type".into(), + )); + } + }; + + let mut compressed = Vec::new(); + self.compressor.compress(&encoded, &mut compressed)?; + Ok(( + Some(LanceBuffer::from(compressed)), + ProtobufUtils21::wrapped(self.block_compression, inner_encoding)?, + )) + } +} + +impl BlockDecompressor for CompressedBufferEncoder { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Compressed variable block")?; + let mut decompressed = Vec::new(); + self.compressor.decompress(&data, &mut decompressed)?; + + // Delegate to BinaryBlockDecompressor which handles the inline metadata + let inner_decoder = BinaryBlockDecompressor::default(); + inner_decoder.decompress(Some(LanceBuffer::from(decompressed)), num_values) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::str::FromStr; + + use crate::encodings::physical::block::zstd::ZstdBufferCompressor; + + #[test] + fn test_compression_scheme_from_str() { + assert_eq!( + CompressionScheme::from_str("none").unwrap(), + CompressionScheme::None + ); + assert_eq!( + CompressionScheme::from_str("zstd").unwrap(), + CompressionScheme::Zstd + ); + } + + #[test] + fn test_compression_scheme_from_str_invalid() { + assert!(CompressionScheme::from_str("invalid").is_err()); + } + + #[cfg(feature = "zstd")] + mod zstd { + use std::io::Write; + + use super::*; + + #[test] + fn test_compress_zstd_with_length_prefixed() { + let compressor = ZstdBufferCompressor::new(0); + let input_data = b"Hello, world!"; + let mut compressed_data = Vec::new(); + + compressor + .compress(input_data, &mut compressed_data) + .unwrap(); + let mut decompressed_data = Vec::new(); + compressor + .decompress(&compressed_data, &mut decompressed_data) + .unwrap(); + assert_eq!(input_data, decompressed_data.as_slice()); + } + + #[test] + fn test_zstd_compress_decompress_multiple_times() { + let compressor = ZstdBufferCompressor::new(0); + let (input_data_1, input_data_2) = (b"Hello ", b"World"); + let mut compressed_data = Vec::new(); + + compressor + .compress(input_data_1, &mut compressed_data) + .unwrap(); + let compressed_length_1 = compressed_data.len(); + + compressor + .compress(input_data_2, &mut compressed_data) + .unwrap(); + + let mut decompressed_data = Vec::new(); + compressor + .decompress( + &compressed_data[..compressed_length_1], + &mut decompressed_data, + ) + .unwrap(); + + compressor + .decompress( + &compressed_data[compressed_length_1..], + &mut decompressed_data, + ) + .unwrap(); + + // the output should contain both input_data_1 and input_data_2 + assert_eq!( + decompressed_data.len(), + input_data_1.len() + input_data_2.len() + ); + assert_eq!( + &decompressed_data[..input_data_1.len()], + input_data_1, + "First part of decompressed data should match input_1" + ); + assert_eq!( + &decompressed_data[input_data_1.len()..], + input_data_2, + "Second part of decompressed data should match input_2" + ); + } + + #[test] + fn test_compress_zstd_raw_stream_format_and_decompress_with_length_prefixed() { + let compressor = ZstdBufferCompressor::new(0); + let input_data = b"Hello, world!"; + let mut compressed_data = Vec::new(); + + // compress using raw stream format + let mut encoder = ::zstd::Encoder::new(&mut compressed_data, 0).unwrap(); + encoder.write_all(input_data).unwrap(); + encoder.finish().expect("failed to encode data with zstd"); + + // decompress using length prefixed format + let mut decompressed_data = Vec::new(); + compressor + .decompress(&compressed_data, &mut decompressed_data) + .unwrap(); + assert_eq!(input_data, decompressed_data.as_slice()); + } + } + + #[cfg(feature = "lz4")] + mod lz4 { + use std::{collections::HashMap, sync::Arc}; + + use arrow_schema::{DataType, Field}; + use lance_datagen::array::{binary_prefix_plus_counter, utf8_prefix_plus_counter}; + + use super::*; + + use crate::constants::DICT_SIZE_RATIO_META_KEY; + use crate::{ + constants::{ + COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY, STRUCTURAL_ENCODING_FULLZIP, + STRUCTURAL_ENCODING_META_KEY, + }, + encodings::physical::block::lz4::Lz4BufferCompressor, + testing::{ + FnArrayGeneratorProvider, TestCases, TestEncoding, + check_round_trip_encoding_generated, + }, + }; + + #[test] + fn test_lz4_compress_decompress() { + let compressor = Lz4BufferCompressor::default(); + let input_data = b"Hello, world!"; + let mut compressed_data = Vec::new(); + + compressor + .compress(input_data, &mut compressed_data) + .unwrap(); + let mut decompressed_data = Vec::new(); + compressor + .decompress(&compressed_data, &mut decompressed_data) + .unwrap(); + assert_eq!(input_data, decompressed_data.as_slice()); + } + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_lz4_compress_round_trip( + #[values( + DataType::Utf8, + DataType::LargeUtf8, + DataType::Binary, + DataType::LargeBinary + )] + data_type: DataType, + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(false, true)] use_slicing: bool, + ) { + let field = Field::new("", data_type.clone(), false); + let mut field_meta = HashMap::new(); + field_meta.insert(COMPRESSION_META_KEY.to_string(), "lz4".to_string()); + // Some bad cardinality estimatation causes us to use dictionary encoding currently + // which causes the expected encoding check to fail. + field_meta.insert(DICT_DIVISOR_META_KEY.to_string(), "100000".to_string()); + field_meta.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.0001".to_string()); + // Also disable size-based dictionary encoding + field_meta.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_FULLZIP.to_string(), + ); + let field = field.with_metadata(field_meta); + let test_cases = TestCases::basic() + // Need to use large pages as small pages might be too small to compress + .with_page_sizes(vec![1024 * 1024]) + .with_expected_encoding("zstd") + .with_encoding(encoding) + .with_slicing_modes([use_slicing]); + + // Can't use the default random provider because random data isn't compressible + // and we will fallback to uncompressed encoding + let datagen = Box::new(FnArrayGeneratorProvider::new(move || match data_type { + DataType::Utf8 => utf8_prefix_plus_counter("compressme", false), + DataType::Binary => { + binary_prefix_plus_counter(Arc::from(b"compressme".to_owned()), false) + } + DataType::LargeUtf8 => utf8_prefix_plus_counter("compressme", true), + DataType::LargeBinary => { + binary_prefix_plus_counter(Arc::from(b"compressme".to_owned()), true) + } + _ => panic!("Unsupported data type: {:?}", data_type), + })); + + check_round_trip_encoding_generated(field, datagen, test_cases).await; + } + } +} diff --git a/vendor/lance-encoding/src/encodings/physical/byte_stream_split.rs b/vendor/lance-encoding/src/encodings/physical/byte_stream_split.rs new file mode 100644 index 000000000..d4dbd1045 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/byte_stream_split.rs @@ -0,0 +1,452 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! # Byte Stream Split (BSS) Miniblock Format +//! +//! Byte Stream Split is a data transformation technique that improves compression +//! by reorganizing multi-byte values to group bytes from the same position together. +//! This is particularly effective for data where some byte positions have low entropy. +//! +//! ## How It Works +//! +//! BSS splits multi-byte values by byte position, creating separate streams +//! for each byte position across all values. This transformation is most beneficial +//! when certain byte positions have low entropy (e.g., high-order bytes that are +//! mostly zeros, sign-extended bytes, or floating-point sign/exponent bytes that +//! cluster around common values). +//! +//! ### Example +//! +//! Input data (f32): `[1.0, 2.0, 3.0, 4.0]` +//! +//! In little-endian bytes: +//! - 1.0 = `[00, 00, 80, 3F]` +//! - 2.0 = `[00, 00, 00, 40]` +//! - 3.0 = `[00, 00, 40, 40]` +//! - 4.0 = `[00, 00, 80, 40]` +//! +//! After BSS transformation: +//! - Byte stream 0: `[00, 00, 00, 00]` (all first bytes) +//! - Byte stream 1: `[00, 00, 00, 00]` (all second bytes) +//! - Byte stream 2: `[80, 00, 40, 80]` (all third bytes) +//! - Byte stream 3: `[3F, 40, 40, 40]` (all fourth bytes) +//! +//! Output: `[00, 00, 00, 00, 00, 00, 00, 00, 80, 00, 40, 80, 3F, 40, 40, 40]` +//! +//! ## Compression Benefits +//! +//! BSS itself doesn't compress data - it reorders it. The compression benefit +//! comes when BSS is combined with general-purpose compression (e.g., LZ4): +//! +//! 1. **Timestamps**: Sequential timestamps have similar high-order bytes +//! 2. **Sensor data**: Readings often vary in a small range, sharing exponent bits +//! 3. **Financial data**: Prices may cluster around certain values +//! +//! ## Supported Types +//! +//! - 32-bit floating point (f32) +//! - 64-bit floating point (f64) +//! +//! ## Chunk Handling +//! +//! - Maximum chunk size depends on data type: +//! - f32: 1024 values (4KB per chunk) +//! - f64: 512 values (4KB per chunk) +//! - All chunks share a single global buffer +//! - Non-last chunks always contain power-of-2 values + +use std::fmt::Debug; + +use crate::buffer::LanceBuffer; +use crate::compression::MiniBlockDecompressor; +use crate::compression_config::BssMode; +use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock}; +use crate::encodings::logical::primitive::miniblock::{ + MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, +}; +use crate::format::ProtobufUtils21; +use crate::format::pb21::CompressiveEncoding; +use crate::statistics::{GetStat, Stat}; +use arrow_array::{cast::AsArray, types::UInt64Type}; +use lance_core::Result; + +/// Byte Stream Split encoder for floating point values +/// +/// This encoding splits floating point values by byte position and stores +/// each byte stream separately. This improves compression ratios for +/// floating point data with similar patterns. +#[derive(Debug, Clone)] +pub struct ByteStreamSplitEncoder { + bits_per_value: usize, +} + +impl ByteStreamSplitEncoder { + pub fn new(bits_per_value: usize) -> Self { + assert!( + bits_per_value == 32 || bits_per_value == 64, + "ByteStreamSplit only supports 32-bit (f32) or 64-bit (f64) values" + ); + Self { bits_per_value } + } + + fn bytes_per_value(&self) -> usize { + self.bits_per_value / 8 + } + + fn max_chunk_size(&self) -> usize { + // For ByteStreamSplit, total bytes = bytes_per_value * chunk_size + // MAX_MINIBLOCK_BYTES = 8186 + // For f32 (4 bytes): 8186 / 4 = 2046, so max chunk = 1024 (power of 2) + // For f64 (8 bytes): 8186 / 8 = 1023, so max chunk = 512 (power of 2) + match self.bits_per_value { + 32 => 1024, + 64 => 512, + _ => unreachable!("ByteStreamSplit only supports 32 or 64 bit values"), + } + } +} + +impl MiniBlockCompressor for ByteStreamSplitEncoder { + fn compress( + &self, + _context: MiniBlockCompressionContext, + page: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + match page { + DataBlock::FixedWidth(data) => { + let num_values = data.num_values; + let bytes_per_value = self.bytes_per_value(); + + if num_values == 0 { + return Ok(( + MiniBlockCompressed { + data: vec![], + chunks: vec![], + num_values: 0, + }, + ProtobufUtils21::byte_stream_split(ProtobufUtils21::flat( + self.bits_per_value as u64, + None, + )), + )); + } + + let total_size = num_values as usize * bytes_per_value; + let mut global_buffer = vec![0u8; total_size]; + + let mut chunks = Vec::new(); + let data_slice = data.data.as_ref(); + let mut processed_values = 0usize; + let max_chunk_size = self.max_chunk_size(); + + while processed_values < num_values as usize { + let chunk_size = (num_values as usize - processed_values).min(max_chunk_size); + let chunk_offset = processed_values * bytes_per_value; + + // Create chunk-local byte streams + for i in 0..chunk_size { + let src_offset = (processed_values + i) * bytes_per_value; + for j in 0..bytes_per_value { + // Store in chunk-local byte stream format + let dst_offset = chunk_offset + j * chunk_size + i; + global_buffer[dst_offset] = data_slice[src_offset + j]; + } + } + + let chunk_bytes = chunk_size * bytes_per_value; + let log_num_values = if processed_values + chunk_size == num_values as usize { + 0 // Last chunk + } else { + chunk_size.ilog2() as u8 + }; + + debug_assert!(chunk_bytes > 0); + chunks.push(MiniBlockChunk { + buffer_sizes: vec![chunk_bytes as u32], + log_num_values, + }); + + processed_values += chunk_size; + } + + let data_buffers = vec![LanceBuffer::from(global_buffer)]; + + // TODO: Should support underlying compression + let encoding = ProtobufUtils21::byte_stream_split(ProtobufUtils21::flat( + self.bits_per_value as u64, + None, + )); + + Ok(( + MiniBlockCompressed { + data: data_buffers, + chunks, + num_values, + }, + encoding, + )) + } + _ => Err(lance_core::Error::invalid_input_source( + "ByteStreamSplit encoding only supports FixedWidth data blocks".into(), + )), + } + } +} + +/// Byte Stream Split decompressor +#[derive(Debug)] +pub struct ByteStreamSplitDecompressor { + bits_per_value: usize, +} + +impl ByteStreamSplitDecompressor { + pub fn new(bits_per_value: usize) -> Self { + assert!( + bits_per_value == 32 || bits_per_value == 64, + "ByteStreamSplit only supports 32-bit (f32) or 64-bit (f64) values" + ); + Self { bits_per_value } + } + + fn bytes_per_value(&self) -> usize { + self.bits_per_value / 8 + } +} + +impl MiniBlockDecompressor for ByteStreamSplitDecompressor { + fn decompress(&self, data: Vec, num_values: u64) -> Result { + if num_values == 0 { + return Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::empty(), + bits_per_value: self.bits_per_value as u64, + num_values: 0, + block_info: BlockInfo::new(), + })); + } + + let bytes_per_value = self.bytes_per_value(); + let total_bytes = num_values as usize * bytes_per_value; + + if data.len() != 1 { + return Err(lance_core::Error::invalid_input_source( + format!( + "ByteStreamSplit decompression expects 1 buffer, but got {}", + data.len() + ) + .into(), + )); + } + + let input_buffer = &data[0]; + + if input_buffer.len() != total_bytes { + return Err(lance_core::Error::invalid_input_source( + format!( + "Expected {} bytes for decompression, but got {}", + total_bytes, + input_buffer.len() + ) + .into(), + )); + } + + let mut output = vec![0u8; total_bytes]; + + // Input buffer contains chunk-local byte streams + for i in 0..num_values as usize { + for j in 0..bytes_per_value { + let src_offset = j * num_values as usize + i; + output[i * bytes_per_value + j] = input_buffer[src_offset]; + } + } + + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(output), + bits_per_value: self.bits_per_value as u64, + num_values, + block_info: BlockInfo::new(), + })) + } + + fn decoded_size_bytes(&self, num_values: u64) -> Option { + num_values.checked_mul(self.bytes_per_value() as u64) + } +} + +/// Determine if BSS should be used based on mode and data characteristics +pub fn should_use_bss(data: &FixedWidthDataBlock, mode: BssMode) -> bool { + // Only support 32-bit and 64-bit values + // BSS is most effective for these common types (floats, timestamps, etc.) + // 16-bit values have limited benefit with only 2 streams + if data.bits_per_value != 32 && data.bits_per_value != 64 { + return false; + } + + let sensitivity = mode.to_sensitivity(); + + // Fast paths + if sensitivity <= 0.0 { + return false; + } + if sensitivity >= 1.0 { + return true; + } + + // Auto mode: check byte position entropy + evaluate_entropy_for_bss(data, sensitivity) +} + +/// Evaluate if BSS should be used based on byte position entropy +fn evaluate_entropy_for_bss(data: &FixedWidthDataBlock, sensitivity: f32) -> bool { + // Get the precomputed entropy statistics + let Some(entropy_stat) = data.get_stat(Stat::BytePositionEntropy) else { + return false; // No entropy data available + }; + + let entropies = entropy_stat.as_primitive::(); + if entropies.is_empty() { + return false; + } + + // Calculate average entropy across all byte positions + let sum: u64 = entropies.values().iter().sum(); + let avg_entropy = sum as f64 / entropies.len() as f64 / 1000.0; // Scale back from integer + + // Entropy threshold based on sensitivity + // sensitivity = 0.5 (default auto) -> threshold = 4.0 bits + // sensitivity = 0.0 (off) -> threshold = 0.0 (never use) + // sensitivity = 1.0 (on) -> threshold = 8.0 (always use) + let entropy_threshold = sensitivity as f64 * 8.0; + + // Use BSS if average entropy is below threshold + // Lower entropy means more repetitive byte patterns + avg_entropy < entropy_threshold +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_round_trip_f32() { + let encoder = ByteStreamSplitEncoder::new(32); + let decompressor = ByteStreamSplitDecompressor::new(32); + + // Test data + let values: Vec = vec![ + 1.0, + 2.5, + -3.7, + 4.2, + 0.0, + -0.0, + f32::INFINITY, + f32::NEG_INFINITY, + ]; + let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + + let data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(bytes), + bits_per_value: 32, + num_values: values.len() as u64, + block_info: BlockInfo::new(), + }); + + // Compress + let (compressed, _encoding) = encoder + .compress(MiniBlockCompressionContext::new(0, true, true), data_block) + .unwrap(); + + // Decompress + let decompressed = decompressor + .decompress(compressed.data, values.len() as u64) + .unwrap(); + let DataBlock::FixedWidth(decompressed_fixed) = &decompressed else { + panic!("Expected FixedWidth DataBlock") + }; + + // Verify + let result_bytes = decompressed_fixed.data.as_ref(); + let result_values: Vec = result_bytes + .chunks_exact(4) + .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap())) + .collect(); + + assert_eq!(values, result_values); + } + + #[test] + fn test_round_trip_f64() { + let encoder = ByteStreamSplitEncoder::new(64); + let decompressor = ByteStreamSplitDecompressor::new(64); + + // Test data + let values: Vec = vec![ + 1.0, + 2.5, + -3.7, + 4.2, + 0.0, + -0.0, + f64::INFINITY, + f64::NEG_INFINITY, + ]; + let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + + let data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(bytes), + bits_per_value: 64, + num_values: values.len() as u64, + block_info: BlockInfo::new(), + }); + + // Compress + let (compressed, _encoding) = encoder + .compress(MiniBlockCompressionContext::new(0, true, true), data_block) + .unwrap(); + + // Decompress + let decompressed = decompressor + .decompress(compressed.data, values.len() as u64) + .unwrap(); + let DataBlock::FixedWidth(decompressed_fixed) = &decompressed else { + panic!("Expected FixedWidth DataBlock") + }; + + // Verify + let result_bytes = decompressed_fixed.data.as_ref(); + let result_values: Vec = result_bytes + .chunks_exact(8) + .map(|chunk| f64::from_le_bytes(chunk.try_into().unwrap())) + .collect(); + + assert_eq!(values, result_values); + } + + #[test] + fn test_empty_data() { + let encoder = ByteStreamSplitEncoder::new(32); + let decompressor = ByteStreamSplitDecompressor::new(32); + + let data_block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::empty(), + bits_per_value: 32, + num_values: 0, + block_info: BlockInfo::new(), + }); + + // Compress empty data + let (compressed, _encoding) = encoder + .compress(MiniBlockCompressionContext::new(0, true, true), data_block) + .unwrap(); + + // Decompress empty data + let decompressed = decompressor.decompress(compressed.data, 0).unwrap(); + let DataBlock::FixedWidth(decompressed_fixed) = &decompressed else { + panic!("Expected FixedWidth DataBlock") + }; + + assert_eq!(decompressed_fixed.num_values, 0); + assert_eq!(decompressed_fixed.data.len(), 0); + } +} diff --git a/vendor/lance-encoding/src/encodings/physical/constant.rs b/vendor/lance-encoding/src/encodings/physical/constant.rs new file mode 100644 index 000000000..dd153f789 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/constant.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Routines for compressing and decompressing constant-encoded data + +use crate::{ + buffer::LanceBuffer, + compression::{BlockDecompressor, FixedPerValueDecompressor, require_no_block_payload}, + data::{AllNullDataBlock, ConstantDataBlock, DataBlock, FixedWidthDataBlock}, +}; + +use lance_core::Result; + +/// A decompressor for constant-encoded data +#[derive(Debug)] +pub struct ConstantDecompressor { + scalar: Option, +} + +impl ConstantDecompressor { + pub fn new(scalar: Option) -> Self { + Self { scalar } + } +} + +impl BlockDecompressor for ConstantDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + require_no_block_payload(data, "Constant")?; + if let Some(scalar) = self.scalar.clone() { + Ok(DataBlock::Constant(ConstantDataBlock { + data: scalar, + num_values, + })) + } else { + Ok(DataBlock::AllNull(AllNullDataBlock { num_values })) + } + } + + fn requires_payload(&self) -> bool { + false + } +} + +impl FixedPerValueDecompressor for ConstantDecompressor { + fn decompress(&self, _data: FixedWidthDataBlock, num_values: u64) -> Result { + if let Some(scalar) = self.scalar.clone() { + Ok(DataBlock::Constant(ConstantDataBlock { + data: scalar, + num_values, + })) + } else { + Ok(DataBlock::AllNull(AllNullDataBlock { num_values })) + } + } + + fn bits_per_value(&self) -> u64 { + self.scalar + .as_ref() + .map(|s| s.len() as u64 * 8) + .unwrap_or(0) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn block_constant_requires_no_payload() { + let decompressor = ConstantDecompressor::new(None); + + assert!(!decompressor.requires_payload()); + assert!(matches!( + BlockDecompressor::decompress(&decompressor, None, 3).unwrap(), + DataBlock::AllNull(AllNullDataBlock { num_values: 3 }) + )); + assert!( + BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::empty()), 3).is_err() + ); + } +} diff --git a/vendor/lance-encoding/src/encodings/physical/fsst.rs b/vendor/lance-encoding/src/encodings/physical/fsst.rs new file mode 100644 index 000000000..38acd24af --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/fsst.rs @@ -0,0 +1,469 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! FSST encoding +//! +//! FSST is a lightweight encoding for variable width data. This module includes +//! adapters for both miniblock and per-value encoding. +//! +//! FSST encoding creates a small symbol table that is needed for decoding. Currently +//! we create one symbol table per disk page and store it in the description. +//! +//! TODO: This seems to be potentially limiting. Perhaps we should create one symbol +//! table per mini-block chunk? In the per-value compression it may even make sense to +//! create multiple symbol tables for a single value! +//! +//! FSST encoding is transparent. + +use lance_core::{Error, Result}; + +use crate::{ + buffer::LanceBuffer, + compression::{MiniBlockDecompressor, VariablePerValueDecompressor}, + data::{BlockInfo, DataBlock, VariableWidthBlock}, + encodings::logical::primitive::{ + fullzip::{PerValueCompressor, PerValueDataBlock}, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor}, + }, + format::{ + ProtobufUtils21, + pb21::{self, CompressiveEncoding}, + }, +}; + +use super::binary::BinaryMiniBlockEncoder; + +pub(crate) fn map_fsst_error(err: std::io::Error) -> Error { + match err.kind() { + std::io::ErrorKind::InvalidData => Error::corrupt_file_named("fsst", err.to_string()), + _ => err.into(), + } +} + +struct FsstCompressed { + data: VariableWidthBlock, + symbol_table: Vec, +} + +impl FsstCompressed { + fn fsst_compress(data: DataBlock) -> Result { + match data { + DataBlock::VariableWidth(variable_width) => { + match variable_width.bits_per_offset { + 32 => { + let offsets = variable_width.offsets.borrow_to_typed_slice::(); + let offsets_slice = offsets.as_ref(); + let bytes_data = variable_width.data.into_buffer(); + + // prepare compression output buffer + let mut dest_offsets = vec![0_i32; offsets_slice.len() * 2]; + let mut dest_values = vec![0_u8; bytes_data.len() * 2]; + let mut symbol_table = vec![0_u8; fsst::fsst::FSST_SYMBOL_TABLE_SIZE]; + + // fsst compression + fsst::fsst::compress( + &mut symbol_table, + bytes_data.as_slice(), + offsets_slice, + &mut dest_values, + &mut dest_offsets, + )?; + + // construct `DataBlock` for BinaryMiniBlockEncoder, we may want some `DataBlock` construct methods later + let compressed = VariableWidthBlock { + data: LanceBuffer::reinterpret_vec(dest_values), + bits_per_offset: 32, + offsets: LanceBuffer::reinterpret_vec(dest_offsets), + num_values: variable_width.num_values, + block_info: BlockInfo::new(), + }; + + Ok(Self { + data: compressed, + symbol_table, + }) + } + 64 => { + let offsets = variable_width.offsets.borrow_to_typed_slice::(); + let offsets_slice = offsets.as_ref(); + let bytes_data = variable_width.data.into_buffer(); + + // prepare compression output buffer + let mut dest_offsets = vec![0_i64; offsets_slice.len() * 2]; + let mut dest_values = vec![0_u8; bytes_data.len() * 2]; + let mut symbol_table = vec![0_u8; fsst::fsst::FSST_SYMBOL_TABLE_SIZE]; + + // fsst compression + fsst::fsst::compress( + &mut symbol_table, + bytes_data.as_slice(), + offsets_slice, + &mut dest_values, + &mut dest_offsets, + )?; + + // construct `DataBlock` for BinaryMiniBlockEncoder, we may want some `DataBlock` construct methods later + let compressed = VariableWidthBlock { + data: LanceBuffer::reinterpret_vec(dest_values), + bits_per_offset: 64, + offsets: LanceBuffer::reinterpret_vec(dest_offsets), + num_values: variable_width.num_values, + block_info: BlockInfo::new(), + }; + + Ok(Self { + data: compressed, + symbol_table, + }) + } + _ => panic!( + "Unsupported offsets type {}", + variable_width.bits_per_offset + ), + } + } + _ => Err(Error::invalid_input_source( + format!( + "Cannot compress a data block of type {} with FsstEncoder", + data.name() + ) + .into(), + )), + } + } +} + +#[derive(Debug, Default)] +pub struct FsstMiniBlockEncoder { + minichunk_size: Option, +} + +impl FsstMiniBlockEncoder { + pub fn new(minichunk_size: Option) -> Self { + Self { minichunk_size } + } +} + +impl MiniBlockCompressor for FsstMiniBlockEncoder { + fn compress( + &self, + context: MiniBlockCompressionContext, + data: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + let compressed = FsstCompressed::fsst_compress(data)?; + + let data_block = DataBlock::VariableWidth(compressed.data); + + // compress the fsst compressed data using `BinaryMiniBlockEncoder` + let binary_compressor = Box::new(BinaryMiniBlockEncoder::new(self.minichunk_size)) + as Box; + + let (binary_miniblock_compressed, binary_array_encoding) = + binary_compressor.compress(context, data_block)?; + + Ok(( + binary_miniblock_compressed, + ProtobufUtils21::fsst(binary_array_encoding, compressed.symbol_table), + )) + } +} + +#[derive(Debug)] +pub struct FsstPerValueEncoder { + inner: Box, +} + +impl FsstPerValueEncoder { + pub fn new(inner: Box) -> Self { + Self { inner } + } +} + +impl PerValueCompressor for FsstPerValueEncoder { + fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> { + let compressed = FsstCompressed::fsst_compress(data)?; + + let data_block = DataBlock::VariableWidth(compressed.data); + + let (binary_compressed, binary_array_encoding) = self.inner.compress(data_block)?; + + Ok(( + binary_compressed, + ProtobufUtils21::fsst(binary_array_encoding, compressed.symbol_table), + )) + } +} + +#[derive(Debug)] +pub struct FsstPerValueDecompressor { + symbol_table: LanceBuffer, + inner_decompressor: Box, +} + +impl FsstPerValueDecompressor { + pub fn new( + symbol_table: LanceBuffer, + inner_decompressor: Box, + ) -> Self { + Self { + symbol_table, + inner_decompressor, + } + } +} + +impl VariablePerValueDecompressor for FsstPerValueDecompressor { + fn decompress(&self, data: VariableWidthBlock) -> Result { + // Step 1. Run inner decompressor + let compressed_variable_data = self + .inner_decompressor + .decompress(data)? + .as_variable_width() + .unwrap(); + + // Step 2. FSST decompress + let bytes = compressed_variable_data.data.borrow_to_typed_slice::(); + let bytes = bytes.as_ref(); + + match compressed_variable_data.bits_per_offset { + 32 => { + let offsets = compressed_variable_data + .offsets + .borrow_to_typed_slice::(); + let offsets = offsets.as_ref(); + let num_values = compressed_variable_data.num_values; + + // The data will expand at most 8 times + // The offsets will be the same size because we have the same # of strings + let mut decompress_bytes_buf = vec![0u8; bytes.len() * 8]; + let mut decompress_offset_buf = vec![0i32; offsets.len()]; + fsst::fsst::decompress( + &self.symbol_table, + bytes, + offsets, + &mut decompress_bytes_buf, + &mut decompress_offset_buf, + ) + .map_err(map_fsst_error)?; + + // Ensure the offsets array is trimmed to exactly num_values + 1 elements + decompress_offset_buf.truncate((num_values + 1) as usize); + + Ok(DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::from(decompress_bytes_buf), + offsets: LanceBuffer::reinterpret_vec(decompress_offset_buf), + bits_per_offset: 32, + num_values, + block_info: BlockInfo::new(), + })) + } + 64 => { + let offsets = compressed_variable_data + .offsets + .borrow_to_typed_slice::(); + let offsets = offsets.as_ref(); + let num_values = compressed_variable_data.num_values; + + // The data will expand at most 8 times + // The offsets will be the same size because we have the same # of strings + let mut decompress_bytes_buf = vec![0u8; bytes.len() * 8]; + let mut decompress_offset_buf = vec![0i64; offsets.len()]; + fsst::fsst::decompress( + &self.symbol_table, + bytes, + offsets, + &mut decompress_bytes_buf, + &mut decompress_offset_buf, + ) + .map_err(map_fsst_error)?; + + // Ensure the offsets array is trimmed to exactly num_values + 1 elements + decompress_offset_buf.truncate((num_values + 1) as usize); + + Ok(DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::from(decompress_bytes_buf), + offsets: LanceBuffer::reinterpret_vec(decompress_offset_buf), + bits_per_offset: 64, + num_values, + block_info: BlockInfo::new(), + })) + } + _ => panic!( + "Unsupported offset type {}", + compressed_variable_data.bits_per_offset, + ), + } + } +} + +#[derive(Debug)] +pub struct FsstMiniBlockDecompressor { + symbol_table: LanceBuffer, + inner_decompressor: Box, +} + +impl FsstMiniBlockDecompressor { + pub fn new( + description: &pb21::Fsst, + inner_decompressor: Box, + ) -> Self { + Self { + symbol_table: LanceBuffer::from_bytes(description.symbol_table.clone(), 1), + inner_decompressor, + } + } +} + +impl MiniBlockDecompressor for FsstMiniBlockDecompressor { + fn decompress(&self, data: Vec, num_values: u64) -> Result { + // Step 1. decompress data use `BinaryMiniBlockDecompressor` + // Extract the bits_per_offset from the binary encoding + let compressed_data_block = self.inner_decompressor.decompress(data, num_values)?; + let DataBlock::VariableWidth(compressed_data_block) = compressed_data_block else { + panic!("BinaryMiniBlockDecompressor should output VariableWidth DataBlock") + }; + + // Step 2. FSST decompress + let bytes = &compressed_data_block.data; + let (decompress_bytes_buf, decompress_offset_buf) = + if compressed_data_block.bits_per_offset == 64 { + let offsets = compressed_data_block.offsets.borrow_to_typed_slice::(); + let offsets = offsets.as_ref(); + + // The data will expand at most 8 times + // The offsets will be the same size because we have the same # of strings + let mut decompress_bytes_buf = vec![0u8; bytes.len() * 8]; + let mut decompress_offset_buf = vec![0i64; offsets.len()]; + fsst::fsst::decompress( + &self.symbol_table, + bytes.as_ref(), + offsets, + &mut decompress_bytes_buf, + &mut decompress_offset_buf, + ) + .map_err(map_fsst_error)?; + + // Ensure the offsets array is trimmed to exactly num_values + 1 elements + decompress_offset_buf.truncate((num_values + 1) as usize); + + ( + decompress_bytes_buf, + LanceBuffer::reinterpret_vec(decompress_offset_buf), + ) + } else { + let offsets = compressed_data_block.offsets.borrow_to_typed_slice::(); + let offsets = offsets.as_ref(); + + // The data will expand at most 8 times + // The offsets will be the same size because we have the same # of strings + let mut decompress_bytes_buf = vec![0u8; bytes.len() * 8]; + let mut decompress_offset_buf = vec![0i32; offsets.len()]; + fsst::fsst::decompress( + &self.symbol_table, + bytes.as_ref(), + offsets, + &mut decompress_bytes_buf, + &mut decompress_offset_buf, + ) + .map_err(map_fsst_error)?; + + // Ensure the offsets array is trimmed to exactly num_values + 1 elements + decompress_offset_buf.truncate((num_values + 1) as usize); + + ( + decompress_bytes_buf, + LanceBuffer::reinterpret_vec(decompress_offset_buf), + ) + }; + + Ok(DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::from(decompress_bytes_buf), + offsets: decompress_offset_buf, + bits_per_offset: compressed_data_block.bits_per_offset, + num_values, + block_info: BlockInfo::new(), + })) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use arrow_array::StringArray; + use fsst::fsst::{FSST_SYMBOL_TABLE_SIZE, compress, decompress}; + use lance_core::Error; + use lance_datagen::{ByteCount, RowCount}; + + use super::map_fsst_error; + use crate::testing::{TestCases, TestEncoding, check_round_trip_encoding_of_data}; + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_fsst( + #[values(false, true)] explicit: bool, + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + let test_cases = TestCases::default() + .with_expected_encoding("fsst") + .with_encoding(encoding); + + // Generate data suitable for FSST (large strings, total size > 32KB) + let arr = lance_datagen::gen_batch() + .anon_col(lance_datagen::array::rand_utf8(ByteCount::from(100), false)) + .into_batch_rows(RowCount::from(5000)) + .unwrap() + .column(0) + .clone(); + + let metadata = if explicit { + HashMap::from([("lance-encoding:compression".to_string(), "fsst".to_string())]) + } else { + // Automatic selection requires max_len >= 5 and total_size >= 32KB. + HashMap::new() + }; + check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + } + + #[test] + fn test_corrupt_fsst_symbol_table_is_corrupt_file() { + let input = "the rain in spain stays mainly in the plain ".repeat(2048); + let array = StringArray::from(vec![input.as_str()]); + let mut symbol_table = [0u8; FSST_SYMBOL_TABLE_SIZE]; + let mut compressed = vec![0u8; array.value_data().len().max(1)]; + let mut compressed_offsets = vec![0i32; array.value_offsets().len()]; + compress( + symbol_table.as_mut(), + array.value_data(), + array.value_offsets(), + &mut compressed, + &mut compressed_offsets, + ) + .unwrap(); + + let st_info = u64::from_ne_bytes(symbol_table[..8].try_into().unwrap()); + assert!(st_info & (1 << 24) != 0, "expected decoder_switch_on input"); + let n_symbols = (st_info & 255) as usize; + assert!(n_symbols > 0); + symbol_table[8 + n_symbols * 8] = 9; + + let mut out = vec![0u8; compressed.len() * 8]; + let mut out_offsets = vec![0i32; compressed_offsets.len()]; + let err = decompress( + &symbol_table, + &compressed, + &compressed_offsets, + &mut out, + &mut out_offsets, + ) + .map_err(map_fsst_error) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err}"); + assert!(err.to_string().contains("symbol length"), "{err}"); + } +} diff --git a/vendor/lance-encoding/src/encodings/physical/general.rs b/vendor/lance-encoding/src/encodings/physical/general.rs new file mode 100644 index 000000000..bea6a85cb --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/general.rs @@ -0,0 +1,683 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use log::trace; + +use crate::{ + Result, + buffer::LanceBuffer, + compression::MiniBlockDecompressor, + data::DataBlock, + encodings::{ + logical::primitive::miniblock::{ + MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, + }, + physical::block::{CompressionConfig, GeneralBufferCompressor}, + }, + format::{ProtobufUtils21, pb21::CompressiveEncoding}, +}; + +/// A miniblock compressor that wraps another miniblock compressor and applies +/// general-purpose compression (LZ4, Zstd) to the resulting buffers. +#[derive(Debug)] +pub struct GeneralMiniBlockCompressor { + inner: Box, + compression: CompressionConfig, +} + +impl GeneralMiniBlockCompressor { + pub fn new(inner: Box, compression: CompressionConfig) -> Self { + Self { inner, compression } + } +} + +/// Minimum buffer size to consider for compression +const MIN_BUFFER_SIZE_FOR_COMPRESSION: usize = 4 * 1024; + +use super::super::logical::primitive::miniblock::MiniBlockChunk; + +impl MiniBlockCompressor for GeneralMiniBlockCompressor { + fn compress( + &self, + context: MiniBlockCompressionContext, + page: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + // First, compress with the inner compressor + let (inner_compressed, inner_encoding) = self.inner.compress(context, page)?; + + // Return the original encoding without compression if there's no data or + // the first buffer is not large enough + if inner_compressed.data.is_empty() + || inner_compressed.data[0].len() < MIN_BUFFER_SIZE_FOR_COMPRESSION + { + return Ok((inner_compressed, inner_encoding)); + } + + // We'll compress each chunk's portion of the first buffer independently + let first_buffer = &inner_compressed.data[0]; + let mut compressed_first_buffer = Vec::new(); + let mut new_chunks = Vec::with_capacity(inner_compressed.chunks.iter().len()); + let mut offset = 0usize; + let mut total_original_size = 0usize; + + let compressor = GeneralBufferCompressor::get_compressor(self.compression)?; + + for chunk in &inner_compressed.chunks { + let chunk_first_buffer_size = chunk.buffer_sizes[0] as usize; + + let chunk_data = &first_buffer.as_ref()[offset..offset + chunk_first_buffer_size]; + total_original_size += chunk_first_buffer_size; + + let compressed_start = compressed_first_buffer.len(); + compressor.compress(chunk_data, &mut compressed_first_buffer)?; + let compressed_size = compressed_first_buffer.len() - compressed_start; + + // Create new chunk with updated first buffer size + let mut new_buffer_sizes = chunk.buffer_sizes.clone(); + new_buffer_sizes[0] = compressed_size as u32; + + new_chunks.push(MiniBlockChunk { + buffer_sizes: new_buffer_sizes, + log_num_values: chunk.log_num_values, + }); + + offset += chunk_first_buffer_size; + } + + // Check if compression was effective + let compressed_total_size = compressed_first_buffer.len(); + if compressed_total_size >= total_original_size { + // Compression didn't help, return original + return Ok((inner_compressed, inner_encoding)); + } + + trace!( + "First buffer compressed from {} to {} bytes (ratio: {:.2})", + total_original_size, + compressed_total_size, + compressed_total_size as f32 / total_original_size as f32 + ); + + // Build final buffers: compressed first buffer + remaining original buffers + let mut final_buffers = vec![LanceBuffer::from(compressed_first_buffer)]; + final_buffers.extend(inner_compressed.data.into_iter().skip(1)); + + let compressed_result = MiniBlockCompressed { + data: final_buffers, + chunks: new_chunks, + num_values: inner_compressed.num_values, + }; + + // Return compressed encoding + let encoding = ProtobufUtils21::wrapped(self.compression, inner_encoding)?; + Ok((compressed_result, encoding)) + } +} + +/// A miniblock decompressor that first decompresses buffers using general-purpose +/// compression (LZ4, Zstd) and then delegates to an inner miniblock decompressor. +#[derive(Debug)] +pub struct GeneralMiniBlockDecompressor { + inner: Box, + compression: CompressionConfig, +} + +impl GeneralMiniBlockDecompressor { + pub fn new(inner: Box, compression: CompressionConfig) -> Self { + Self { inner, compression } + } +} + +impl MiniBlockDecompressor for GeneralMiniBlockDecompressor { + fn decompress(&self, mut data: Vec, num_values: u64) -> Result { + let mut decompressed_buffer = Vec::new(); + + let decompressor = GeneralBufferCompressor::get_compressor(self.compression)?; + decompressor.decompress(&data[0], &mut decompressed_buffer)?; + data[0] = LanceBuffer::from(decompressed_buffer); + + self.inner.decompress(data, num_values) + } + + fn decoded_size_bytes(&self, num_values: u64) -> Option { + self.inner.decoded_size_bytes(num_values) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy}; + use crate::data::{BlockInfo, FixedWidthDataBlock}; + use crate::encodings::physical::block::CompressionScheme; + use crate::encodings::physical::rle::RleEncoder; + use crate::encodings::physical::value::ValueEncoder; + use crate::format::pb21; + use crate::format::pb21::compressive_encoding::Compression; + use arrow_array::{Float64Array, Int32Array}; + + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) + } + + #[derive(Debug)] + struct TestCase { + name: &'static str, + inner_encoder: Box, + compression: CompressionConfig, + data: DataBlock, + expected_compressed: bool, // Whether we expect compression to be applied + min_compression_ratio: f32, // Minimum compression ratio if compressed + } + + fn create_test_cases() -> Vec { + vec![ + // Small data with RLE - should not compress due to size threshold + TestCase { + name: "small_rle_data", + inner_encoder: Box::new(RleEncoder::new()), + compression: CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + }, + data: create_repeated_i32_block(vec![1, 1, 1, 1, 2, 2, 2, 2]), + expected_compressed: false, + min_compression_ratio: 1.0, + }, + // Large repeated data with RLE + LZ4 + TestCase { + name: "large_rle_lz4", + inner_encoder: Box::new(RleEncoder::new()), + compression: CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + }, + data: create_pattern_i32_block(2048, |i| (i / 8) as i32), + expected_compressed: false, // RLE already compresses well, additional LZ4 may not help + min_compression_ratio: 1.0, + }, + // Large repeated data with RLE + Zstd + TestCase { + name: "large_rle_zstd", + inner_encoder: Box::new(RleEncoder::new()), + compression: CompressionConfig { + scheme: CompressionScheme::Zstd, + level: Some(3), + }, + data: create_pattern_i32_block(8192, |i| (i / 16) as i32), + expected_compressed: true, // Zstd might provide additional compression + min_compression_ratio: 0.9, // But not as much since RLE already compressed + }, + // Sequential data with ValueEncoder + LZ4 + TestCase { + name: "sequential_value_lz4", + inner_encoder: Box::new(ValueEncoder {}), + compression: CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + }, + data: create_pattern_i32_block(1024, |i| i as i32), + expected_compressed: false, // Sequential data doesn't compress well + min_compression_ratio: 1.0, + }, + // Float data with ValueEncoder + Zstd + TestCase { + name: "float_value_zstd", + inner_encoder: Box::new(ValueEncoder {}), + compression: CompressionConfig { + scheme: CompressionScheme::Zstd, + level: Some(3), + }, + data: create_pattern_f64_block(1024, |i| i as f64 * 0.1), + expected_compressed: true, + min_compression_ratio: 0.9, + }, + ] + } + + fn create_repeated_i32_block(values: Vec) -> DataBlock { + let array = Int32Array::from(values); + DataBlock::from_array(array) + } + + fn create_pattern_i32_block(size: usize, pattern: F) -> DataBlock + where + F: Fn(usize) -> i32, + { + let values: Vec = (0..size).map(pattern).collect(); + let array = Int32Array::from(values); + DataBlock::from_array(array) + } + + fn create_pattern_f64_block(size: usize, pattern: F) -> DataBlock + where + F: Fn(usize) -> f64, + { + let values: Vec = (0..size).map(pattern).collect(); + let array = Float64Array::from(values); + DataBlock::from_array(array) + } + + fn run_round_trip_test(test_case: TestCase) { + let compressor = + GeneralMiniBlockCompressor::new(test_case.inner_encoder, test_case.compression); + + // Compress the data + let (compressed, encoding) = compressor + .compress(miniblock_context(), test_case.data) + .unwrap(); + + // Check if compression was applied as expected + match &encoding.compression { + Some(Compression::General(cm)) => { + assert!( + test_case.expected_compressed, + "{}: Expected compression to be applied", + test_case.name + ); + assert_eq!( + CompressionScheme::try_from(cm.compression.as_ref().unwrap().scheme()).unwrap(), + test_case.compression.scheme + ); + } + _ => { + // Could be RLE or other encoding if compression didn't help + if test_case.expected_compressed { + // Check if it's RLE encoding (which means compression didn't help) + match &encoding.compression { + Some(Compression::Rle(_)) => { + // RLE encoding returned - compression didn't help + } + Some(Compression::Flat(_)) => { + // Flat encoding returned - compression didn't help + } + _ => { + panic!( + "{}: Expected GeneralMiniBlock but got {:?}", + test_case.name, encoding.compression + ); + } + } + } + } + } + + // Verify chunks are created correctly + assert!( + !compressed.chunks.is_empty(), + "{}: No chunks created", + test_case.name + ); + + // Test decompression by simulating the actual miniblock decoding process + let decompressed_data = decompress_miniblock_chunks(&compressed, &encoding); + + // Verify round trip by checking data size + // We expect the decompressed data to match the original number of values + // The bytes per value depends on the test case + let bytes_per_value = if test_case.name.contains("float") { + 8 // f64 + } else { + 4 // i32 + }; + let expected_bytes = compressed.num_values as usize * bytes_per_value; + assert_eq!( + expected_bytes, + decompressed_data.len(), + "{}: Data size mismatch", + test_case.name + ); + + // Check compression ratio if applicable + if test_case.expected_compressed { + let compression_ratio = compressed.data[0].len() as f32 / expected_bytes as f32; + assert!( + compression_ratio <= test_case.min_compression_ratio, + "{}: Compression ratio {} > expected {}", + test_case.name, + compression_ratio, + test_case.min_compression_ratio + ); + } + } + + fn decompress_miniblock_chunks( + compressed: &MiniBlockCompressed, + encoding: &CompressiveEncoding, + ) -> Vec { + let mut decompressed_data = Vec::new(); + let mut offsets = vec![0usize; compressed.data.len()]; // Track offset for each buffer + let decompression_strategy = DefaultDecompressionStrategy::default(); + + for chunk in &compressed.chunks { + let chunk_values = if chunk.log_num_values > 0 { + 1u64 << chunk.log_num_values + } else { + // Last chunk - calculate remaining values + let decompressed_values = + decompressed_data.len() as u64 / get_bytes_per_value(compressed) as u64; + compressed.num_values.saturating_sub(decompressed_values) + }; + + // Extract buffers for this chunk + let mut chunk_buffers = Vec::new(); + for (i, &size) in chunk.buffer_sizes.iter().enumerate() { + if i < compressed.data.len() { + let buffer_data = + compressed.data[i].slice_with_length(offsets[i], size as usize); + chunk_buffers.push(buffer_data); + offsets[i] += size as usize; + } + } + + // Create a decompressor for this chunk + let decompressor = decompression_strategy + .create_miniblock_decompressor(encoding, &decompression_strategy) + .unwrap(); + + // Decompress the chunk + let chunk_decompressed = decompressor + .decompress(chunk_buffers, chunk_values) + .unwrap(); + + match chunk_decompressed { + DataBlock::FixedWidth(ref block) => { + decompressed_data.extend_from_slice(block.data.as_ref()); + } + _ => panic!("Expected FixedWidth block"), + } + } + + decompressed_data + } + + fn get_bytes_per_value(compressed: &MiniBlockCompressed) -> usize { + // This is a simplification - in reality we'd need to know the data type + // For our tests, we mostly use i32 (4 bytes) or f64 (8 bytes) + // We can try to guess based on the data size + if compressed.num_values == 0 { + return 4; // Default to i32 + } + + // For float tests, the number is usually 1024 and we use f64 + if compressed.num_values == 1024 { + return 8; // Likely f64 + } + + 4 // Default to i32 + } + + #[test] + fn test_compressed_mini_block_table_driven() { + for test_case in create_test_cases() { + run_round_trip_test(test_case); + } + } + + #[test] + fn test_compressed_mini_block_threshold() { + // Test that small buffers don't get compressed + let small_test = TestCase { + name: "small_buffer_no_compression", + inner_encoder: Box::new(RleEncoder::new()), + compression: CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + }, + data: create_repeated_i32_block(vec![1, 1, 2, 2]), + expected_compressed: false, + min_compression_ratio: 1.0, + }; + run_round_trip_test(small_test); + } + + #[test] + fn test_compressed_mini_block_with_doubles() { + // Test with large sequential doubles that should compress well with Zstd + // The test focuses on verifying that GeneralMiniBlock works correctly + // when wrapping a simple ValueEncoder + let test_case = TestCase { + name: "float_values_with_zstd", + inner_encoder: Box::new(ValueEncoder {}), + compression: CompressionConfig { + scheme: CompressionScheme::Zstd, + level: Some(3), + }, + // Create enough data to ensure compression is applied + data: create_pattern_f64_block(1024, |i| (i / 10) as f64), + expected_compressed: true, + min_compression_ratio: 0.5, // Zstd should achieve good compression on repetitive data + }; + + run_round_trip_test(test_case); + } + + #[test] + fn test_compressed_mini_block_large_buffers() { + // Use value encoding which doesn't compress data, ensuring large buffers + // Create 1024 i32 values (4KB of data) + let values: Vec = (0..1024).collect(); + let data = LanceBuffer::from_bytes( + bytemuck::cast_slice(&values).to_vec().into(), + std::mem::align_of::() as u64, + ); + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 32, + data, + num_values: 1024, + block_info: BlockInfo::new(), + }); + + // Create compressor with ValueEncoder (no compression) and Zstd wrapper + let inner = Box::new(ValueEncoder {}); + let compression = CompressionConfig { + scheme: CompressionScheme::Zstd, + level: Some(3), + }; + let compressor = GeneralMiniBlockCompressor::new(inner, compression); + + // Compress the data + let (compressed, encoding) = compressor.compress(miniblock_context(), block).unwrap(); + + // Should get GeneralMiniBlock encoding since buffer is 4KB + match &encoding.compression { + Some(Compression::General(cm)) => { + assert!(cm.values.is_some()); + assert_eq!( + cm.compression.as_ref().unwrap().scheme(), + pb21::CompressionScheme::CompressionAlgorithmZstd + ); + assert_eq!(cm.compression.as_ref().unwrap().level, Some(3)); + + // Verify inner encoding is Flat (from ValueEncoder) + match &cm.values.as_ref().unwrap().compression { + Some(Compression::Flat(flat)) => { + assert_eq!(flat.bits_per_value, 32); + } + _ => panic!("Expected Flat inner encoding"), + } + } + _ => panic!("Expected GeneralMiniBlock encoding"), + } + + assert_eq!(compressed.num_values, 1024); + // ValueEncoder produces 1 buffer, so compressed result also has 1 buffer + assert_eq!(compressed.data.len(), 1); + } + + // Special test cases that don't fit the table-driven pattern + + #[test] + fn test_compressed_mini_block_rle_multiple_buffers() { + // RLE produces 2 buffers (values and lengths), test that both are handled correctly + let data = create_repeated_i32_block(vec![1; 100]); + let compressor = GeneralMiniBlockCompressor::new( + Box::new(RleEncoder::new()), + CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + }, + ); + + let (compressed, _) = compressor.compress(miniblock_context(), data).unwrap(); + // RLE produces 2 buffers, but only the first one is compressed + assert_eq!(compressed.data.len(), 2); + } + + #[test] + fn test_rle_with_general_miniblock_wrapper() { + // Test that RLE encoding with bits_per_value >= 32 is automatically wrapped + // in GeneralMiniBlock with LZ4 compression + + // This test directly tests the RLE encoder behavior + // When bits_per_value >= 32, RLE should be wrapped in GeneralMiniBlock with LZ4 + + // Test case 1: 32-bit RLE data + let test_32 = TestCase { + name: "rle_32bit_with_general_wrapper", + inner_encoder: Box::new(RleEncoder::new()), + compression: CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + }, + data: create_repeated_i32_block(vec![1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), + expected_compressed: false, // RLE already compresses well, LZ4 might not help much + min_compression_ratio: 1.0, + }; + + // For 32-bit RLE, the compression strategy should automatically wrap it + // Let's directly test the compressor + let compressor = GeneralMiniBlockCompressor::new( + Box::new(RleEncoder::new()), + CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + }, + ); + + let (_compressed, encoding) = compressor + .compress(miniblock_context(), test_32.data) + .unwrap(); + + // Verify the encoding structure + match &encoding.compression { + Some(Compression::General(cm)) => { + // Check inner encoding is RLE + match &cm.values.as_ref().unwrap().compression { + Some(Compression::Rle(rle)) => { + let Compression::Flat(values) = + rle.values.as_ref().unwrap().compression.as_ref().unwrap() + else { + panic!("Expected flat for RLE values") + }; + let Compression::Flat(run_lengths) = rle + .run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + else { + panic!("Expected flat for RLE run lengths") + }; + assert_eq!(values.bits_per_value, 32); + assert_eq!(run_lengths.bits_per_value, 8); + } + _ => panic!("Expected RLE as inner encoding"), + } + // Check compression is LZ4 + assert_eq!( + cm.compression.as_ref().unwrap().scheme(), + pb21::CompressionScheme::CompressionAlgorithmLz4 + ); + } + Some(Compression::Rle(_)) => { + // Also acceptable if compression didn't help + } + _ => panic!("Expected GeneralMiniBlock or Rle encoding"), + } + + // Test case 2: 64-bit RLE data + let values_64: Vec = vec![100i64; 50] + .into_iter() + .chain(vec![200i64; 50]) + .chain(vec![300i64; 50]) + .collect(); + let array_64 = arrow_array::Int64Array::from(values_64); + let block_64 = DataBlock::from_array(array_64); + + let compressor_64 = GeneralMiniBlockCompressor::new( + Box::new(RleEncoder::new()), + CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + }, + ); + + let (_compressed_64, encoding_64) = compressor_64 + .compress(miniblock_context(), block_64) + .unwrap(); + + // Verify the encoding structure for 64-bit + match &encoding_64.compression { + Some(Compression::General(cm)) => { + // Check inner encoding is RLE + match &cm.values.as_ref().unwrap().compression { + Some(Compression::Rle(rle)) => { + let Compression::Flat(values) = + rle.values.as_ref().unwrap().compression.as_ref().unwrap() + else { + panic!("Expected flat for RLE values") + }; + let Compression::Flat(run_lengths) = rle + .run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + else { + panic!("Expected flat for RLE run lengths") + }; + assert_eq!(values.bits_per_value, 64); + assert_eq!(run_lengths.bits_per_value, 8); + } + _ => panic!("Expected RLE as inner encoding for 64-bit"), + } + // Check compression is LZ4 + assert_eq!( + cm.compression.as_ref().unwrap().scheme(), + pb21::CompressionScheme::CompressionAlgorithmLz4 + ); + } + Some(Compression::Rle(_)) => { + // Also acceptable if compression didn't help + } + _ => panic!("Expected GeneralMiniBlock or Rle encoding for 64-bit"), + } + } + + #[test] + fn test_compressed_mini_block_empty_data() { + let empty_array = Int32Array::from(vec![] as Vec); + let empty_block = DataBlock::from_array(empty_array); + + let compressor = GeneralMiniBlockCompressor::new( + Box::new(ValueEncoder {}), + CompressionConfig { + scheme: CompressionScheme::Lz4, + level: None, + }, + ); + + let result = compressor.compress(miniblock_context(), empty_block); + match result { + Ok((compressed, _)) => { + assert_eq!(compressed.num_values, 0); + } + Err(_) => { + // Empty data might not be supported by ValueEncoder + } + } + } +} diff --git a/vendor/lance-encoding/src/encodings/physical/packed.rs b/vendor/lance-encoding/src/encodings/physical/packed.rs new file mode 100644 index 000000000..3ae3dc4d6 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/packed.rs @@ -0,0 +1,1571 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Packed encoding +//! +//! These encodings take struct data and compress it in a way that all fields are collected +//! together. +//! +//! This encoding can be transparent or opaque. In order to be transparent we must use transparent +//! compression on all children. Then we can zip together the compressed children. + +use std::{convert::TryInto, sync::Arc}; + +use arrow_array::types::UInt64Type; +use lance_core::{Error, Result, datatypes::Field}; + +use crate::{ + buffer::LanceBuffer, + compression::{ + CompressionStrategy, FixedPerValueDecompressor, MiniBlockDecompressor, + VariablePerValueDecompressor, + }, + data::{ + BlockInfo, DataBlock, DataBlockBuilder, FixedWidthDataBlock, StructDataBlock, + VariableWidthBlock, + }, + encodings::logical::primitive::{ + fullzip::{PerValueCompressor, PerValueDataBlock}, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor}, + }, + format::{ + ProtobufUtils21, + pb21::{CompressiveEncoding, PackedStruct, compressive_encoding::Compression}, + }, + statistics::{GetStat, Stat}, +}; + +use super::value::{ValueDecompressor, ValueEncoder}; + +// Transforms a `StructDataBlock` into a row major `FixedWidthDataBlock`. +// Only fields with fixed-width fields are supported for now, and the +// assumption that all fields has `bits_per_value % 8 == 0` is made. +fn struct_data_block_to_fixed_width_data_block( + struct_data_block: StructDataBlock, + bits_per_values: &[u64], +) -> DataBlock { + let data_size = struct_data_block.expect_single_stat::(Stat::DataSize); + let mut output = Vec::with_capacity(data_size as usize); + let num_values = struct_data_block.children[0].num_values(); + + for i in 0..num_values as usize { + for (j, child) in struct_data_block.children.iter().enumerate() { + let bytes_per_value = (bits_per_values[j] / 8) as usize; + let this_data = child + .as_fixed_width_ref() + .unwrap() + .data + .slice_with_length(bytes_per_value * i, bytes_per_value); + output.extend_from_slice(&this_data); + } + } + + DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: bits_per_values.iter().copied().sum(), + data: LanceBuffer::from(output), + num_values, + block_info: BlockInfo::default(), + }) +} + +#[derive(Debug, Default)] +pub struct PackedStructFixedWidthMiniBlockEncoder {} + +impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder { + fn compress( + &self, + context: MiniBlockCompressionContext, + data: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + match data { + DataBlock::Struct(struct_data_block) => { + let bits_per_values = struct_data_block.children.iter().map(|data_block| data_block.as_fixed_width_ref().unwrap().bits_per_value).collect::>(); + + // transform struct datablock to fixed-width data block. + let data_block = struct_data_block_to_fixed_width_data_block(struct_data_block, &bits_per_values); + + // store and transformed fixed-width data block. + let value_miniblock_compressor = Box::new(ValueEncoder::default()) as Box; + let (value_miniblock_compressed, value_array_encoding) = + value_miniblock_compressor.compress(context, data_block)?; + + Ok(( + value_miniblock_compressed, + ProtobufUtils21::packed_struct(value_array_encoding, bits_per_values), + )) + } + _ => Err(Error::invalid_input_source(format!( + "Cannot compress a data block of type {} with PackedStructFixedWidthBlockEncoder", + data.name() + ) + .into())), + } + } +} + +#[derive(Debug)] +pub struct PackedStructFixedWidthMiniBlockDecompressor { + bits_per_values: Vec, + array_encoding: Box, +} + +impl PackedStructFixedWidthMiniBlockDecompressor { + pub fn new(description: &PackedStruct) -> Self { + let array_encoding: Box = match description + .values + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + { + Compression::Flat(flat) => Box::new(ValueDecompressor::from_flat(flat)), + _ => panic!( + "Currently only `ArrayEncoding::Flat` is supported in packed struct encoding in Lance 2.1." + ), + }; + Self { + bits_per_values: description.bits_per_value.clone(), + array_encoding, + } + } +} + +impl MiniBlockDecompressor for PackedStructFixedWidthMiniBlockDecompressor { + fn decompress(&self, data: Vec, num_values: u64) -> Result { + assert_eq!(data.len(), 1); + let encoded_data_block = self.array_encoding.decompress(data, num_values)?; + let DataBlock::FixedWidth(encoded_data_block) = encoded_data_block else { + panic!("ValueDecompressor should output FixedWidth DataBlock") + }; + + let bytes_per_values = self + .bits_per_values + .iter() + .map(|bits_per_value| *bits_per_value as usize / 8) + .collect::>(); + + assert!(encoded_data_block.bits_per_value % 8 == 0); + let encoded_bytes_per_row = (encoded_data_block.bits_per_value / 8) as usize; + + // use a prefix_sum vector as a helper to reconstruct to `StructDataBlock`. + let mut prefix_sum = vec![0; self.bits_per_values.len()]; + for i in 0..(self.bits_per_values.len() - 1) { + prefix_sum[i + 1] = prefix_sum[i] + bytes_per_values[i]; + } + + let mut children_data_block = vec![]; + for i in 0..self.bits_per_values.len() { + let child_buf_size = bytes_per_values[i] * num_values as usize; + let mut child_buf: Vec = Vec::with_capacity(child_buf_size); + + for j in 0..num_values as usize { + // the start of the data at this row is `j * encoded_bytes_per_row`, and the offset for this field is `prefix_sum[i]`, this field has length `bytes_per_values[i]`. + let this_value = encoded_data_block.data.slice_with_length( + prefix_sum[i] + (j * encoded_bytes_per_row), + bytes_per_values[i], + ); + + child_buf.extend_from_slice(&this_value); + } + + let child = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(child_buf), + bits_per_value: self.bits_per_values[i], + num_values, + block_info: BlockInfo::default(), + }); + children_data_block.push(child); + } + Ok(DataBlock::Struct(StructDataBlock { + children: children_data_block, + block_info: BlockInfo::default(), + validity: None, + })) + } +} + +#[derive(Debug)] +struct FixedPackedFieldData { + block: FixedWidthDataBlock, +} + +impl FixedPackedFieldData { + fn append_row_bytes(&self, row_idx: usize, output: &mut Vec) -> Result<()> { + let bits_per_value = self.block.bits_per_value; + if !bits_per_value.is_multiple_of(8) { + return Err(Error::invalid_input( + "Packed struct encoding requires byte-aligned fixed-width children", + )); + } + let bytes_per_value = (bits_per_value / 8) as usize; + let start = row_idx + .checked_mul(bytes_per_value) + .ok_or_else(|| Error::invalid_input("Packed struct row size overflow"))?; + let end = start.checked_add(bytes_per_value).ok_or_else(|| { + Error::invalid_input(format!( + "Packed struct fixed child range overflow: row_idx={row_idx}, \ + bytes_per_value={bytes_per_value}" + )) + })?; + let data = self.block.data.as_ref(); + if end > data.len() { + return Err(Error::invalid_input( + "Packed struct fixed child out of bounds", + )); + } + output.extend_from_slice(&data[start..end]); + Ok(()) + } +} + +#[derive(Debug)] +enum VariablePackedFieldData { + Fixed(FixedPackedFieldData), + Variable { + block: VariableWidthBlock, + bits_per_length: u64, + }, +} + +impl VariablePackedFieldData { + fn append_row_bytes(&self, row_idx: usize, output: &mut Vec) -> Result<()> { + match self { + Self::Fixed(fixed_data) => fixed_data.append_row_bytes(row_idx, output), + Self::Variable { + block, + bits_per_length, + } => { + if !bits_per_length.is_multiple_of(8) { + return Err(Error::invalid_input( + "Packed struct variable children must have byte-aligned length prefixes", + )); + } + let prefix_bytes = (*bits_per_length / 8) as usize; + if !(prefix_bytes == 4 || prefix_bytes == 8) { + return Err(Error::invalid_input( + "Packed struct variable children must use 32 or 64-bit length prefixes", + )); + } + match block.bits_per_offset { + 32 => { + let offsets = block.offsets.borrow_to_typed_slice::(); + let start = offsets[row_idx] as usize; + let end = offsets[row_idx + 1] as usize; + if end > block.data.len() { + return Err(Error::invalid_input( + "Packed struct variable child offsets out of bounds", + )); + } + let len = (end - start) as u32; + if prefix_bytes != std::mem::size_of::() { + return Err(Error::invalid_input( + "Packed struct variable child length prefix mismatch", + )); + } + output.extend_from_slice(&len.to_le_bytes()); + output.extend_from_slice(&block.data[start..end]); + Ok(()) + } + 64 => { + let offsets = block.offsets.borrow_to_typed_slice::(); + let start = offsets[row_idx] as usize; + let end = offsets[row_idx + 1] as usize; + if end > block.data.len() { + return Err(Error::invalid_input( + "Packed struct variable child offsets out of bounds", + )); + } + let len = (end - start) as u64; + if prefix_bytes != std::mem::size_of::() { + return Err(Error::invalid_input( + "Packed struct variable child length prefix mismatch", + )); + } + output.extend_from_slice(&len.to_le_bytes()); + output.extend_from_slice(&block.data[start..end]); + Ok(()) + } + _ => Err(Error::invalid_input( + "Packed struct variable child must use 32 or 64-bit offsets", + )), + } + } + } + } +} + +fn check_struct_validity(data: DataBlock, field_length: usize) -> Result { + let DataBlock::Struct(struct_block) = data else { + return Err(Error::invalid_input( + "Packed struct encoder requires Struct data block", + )); + }; + + if struct_block.children.is_empty() { + return Err(Error::invalid_input( + "Packed struct encoder requires at least one child field", + )); + } + if struct_block.children.len() != field_length { + return Err(Error::invalid_input( + "Struct field metadata does not match number of children", + )); + } + + let num_values = struct_block.children[0].num_values(); + for child in struct_block.children.iter() { + if child.num_values() != num_values { + return Err(Error::invalid_input( + "Packed struct children must have matching value counts", + )); + } + } + Ok(struct_block) +} + +#[derive(Debug)] +pub struct PackedStructVariablePerValueEncoder { + strategy: Arc, + fields: Vec, +} + +impl PackedStructVariablePerValueEncoder { + pub fn new(strategy: Arc, fields: Vec) -> Self { + Self { strategy, fields } + } +} + +impl PerValueCompressor for PackedStructVariablePerValueEncoder { + fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> { + let struct_block = check_struct_validity(data, self.fields.len())?; + let num_values = struct_block.children[0].num_values(); + let mut field_data = Vec::with_capacity(self.fields.len()); + let mut field_metadata = Vec::with_capacity(self.fields.len()); + + for (field, child_block) in self.fields.iter().zip(struct_block.children) { + let compressor = self.strategy.create_per_value(field, &child_block)?; + let (compressed, encoding) = compressor.compress(child_block)?; + match compressed { + PerValueDataBlock::Fixed(block) => { + field_metadata.push(ProtobufUtils21::packed_struct_field_fixed( + encoding, + block.bits_per_value, + )); + let block = FixedPackedFieldData { block }; + field_data.push(VariablePackedFieldData::Fixed(block)); + } + PerValueDataBlock::Variable(block) => { + let bits_per_length = block.bits_per_offset as u64; + field_metadata.push(ProtobufUtils21::packed_struct_field_variable( + encoding, + bits_per_length, + )); + field_data.push(VariablePackedFieldData::Variable { + block, + bits_per_length, + }); + } + } + } + + let mut row_data: Vec = Vec::new(); + let mut row_offsets: Vec = Vec::with_capacity(num_values as usize + 1); + row_offsets.push(0); + let mut total_bytes: usize = 0; + let mut max_row_len: usize = 0; + for row in 0..num_values as usize { + let start = row_data.len(); + for field in &field_data { + field.append_row_bytes(row, &mut row_data)?; + } + let end = row_data.len(); + let row_len = end - start; + max_row_len = max_row_len.max(row_len); + total_bytes = total_bytes + .checked_add(row_len) + .ok_or_else(|| Error::invalid_input("Packed struct row data size overflow"))?; + row_offsets.push(end as u64); + } + debug_assert_eq!(total_bytes, row_data.len()); + + let use_u32_offsets = total_bytes <= u32::MAX as usize && max_row_len <= u32::MAX as usize; + let bits_per_offset = if use_u32_offsets { 32 } else { 64 }; + let offsets_buffer = if use_u32_offsets { + let offsets_u32 = row_offsets + .iter() + .map(|&offset| offset as u32) + .collect::>(); + LanceBuffer::reinterpret_vec(offsets_u32) + } else { + LanceBuffer::reinterpret_vec(row_offsets) + }; + + let data_block = VariableWidthBlock { + data: LanceBuffer::from(row_data), + bits_per_offset, + offsets: offsets_buffer, + num_values, + block_info: BlockInfo::new(), + }; + + Ok(( + PerValueDataBlock::Variable(data_block), + ProtobufUtils21::packed_struct_variable(field_metadata), + )) + } +} + +#[derive(Debug)] +pub(crate) struct PackedStructFixedPerValueEncoder { + field_len: usize, +} + +impl PackedStructFixedPerValueEncoder { + pub(crate) fn new(fields: Vec) -> Self { + Self { + field_len: fields.len(), + } + } +} + +impl PerValueCompressor for PackedStructFixedPerValueEncoder { + fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> { + let struct_block = check_struct_validity(data, self.field_len)?; + + if struct_block.has_variable_width_child() { + return Err(Error::invalid_input( + "Packed struct fixed encoding requires all children to be fixed-width", + )); + } + + let num_values = struct_block.children[0].num_values(); + // Fixed length - supporting only flat encoding + let compressor = Box::new(ValueEncoder::default()) as Box; + let mut field_data = Vec::with_capacity(self.field_len); + let mut field_bits_per_value = Vec::with_capacity(self.field_len); + let mut bits_per_row: u64 = 0; + + for child_block in struct_block.children.into_iter() { + let (compressed, ..) = compressor.compress(child_block)?; + match compressed { + PerValueDataBlock::Fixed(block) => { + bits_per_row = bits_per_row + .checked_add(block.bits_per_value) + .ok_or_else(|| Error::invalid_input("Packed struct row width overflow"))?; + field_bits_per_value.push(block.bits_per_value); + field_data.push(FixedPackedFieldData { block }); + } + _ => { + return Err(Error::invalid_input( + "Packed struct fixed encoding requires all children to be fixed-width", + )); + } + } + } + + // Children are validated byte-aligned in `append_row_bytes`, so the row width + // is an exact number of bytes. + let bytes_per_row = (bits_per_row / 8) as usize; + let mut row_data: Vec = + Vec::with_capacity(bytes_per_row.saturating_mul(num_values as usize)); + for row in 0..num_values as usize { + for field in &field_data { + field.append_row_bytes(row, &mut row_data)?; + } + debug_assert_eq!(row_data.len(), bytes_per_row * (row + 1)); + } + + let data_block = FixedWidthDataBlock { + data: LanceBuffer::from(row_data), + bits_per_value: bits_per_row, + num_values, + block_info: BlockInfo::new(), + }; + + Ok(( + PerValueDataBlock::Fixed(data_block), + ProtobufUtils21::packed_struct( + ProtobufUtils21::flat(bits_per_row, None), + field_bits_per_value, + ), + )) + } +} + +#[derive(Debug)] +pub(crate) enum VariablePackedStructFieldKind { + Fixed { + bits_per_value: u64, + decompressor: Arc, + }, + Variable { + bits_per_length: u64, + decompressor: Arc, + }, +} + +#[derive(Debug)] +pub(crate) struct VariablePackedStructFieldDecoder { + pub(crate) kind: VariablePackedStructFieldKind, +} + +#[derive(Debug)] +pub struct PackedStructVariablePerValueDecompressor { + fields: Vec, +} + +impl PackedStructVariablePerValueDecompressor { + pub(crate) fn new(fields: Vec) -> Self { + Self { fields } + } +} + +#[derive(Debug)] +struct FixedFieldAccumulator { + builder: DataBlockBuilder, + bits_per_value: u64, + empty_value: DataBlock, +} + +impl FixedFieldAccumulator { + fn append_empty(&mut self) -> Result<()> { + self.builder.append(&self.empty_value, 0..1) + } + + fn new(bits_per_value: u64, num_values: u64) -> Result { + if !bits_per_value.is_multiple_of(8) { + return Err(Error::invalid_input( + "Packed struct fixed child must be byte-aligned", + )); + } + + let bytes_per_value = bits_per_value.checked_div(8).ok_or_else(|| { + Error::invalid_input("Invalid bits per value for packed struct field") + })?; + + let estimate = bytes_per_value + .checked_mul(num_values) + .ok_or_else(|| Error::invalid_input("Packed struct fixed child allocation overflow"))?; + + let empty_value = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(vec![0_u8; bytes_per_value as usize]), + bits_per_value, + num_values: 1, + block_info: BlockInfo::new(), + }); + + Ok(Self { + builder: DataBlockBuilder::with_capacity_estimate(estimate), + bits_per_value, + empty_value, + }) + } + + fn finish(self) -> Result { + let DataBlock::FixedWidth(block) = self.builder.finish() else { + return Err(Error::invalid_input( + "Expected fixed-width datablock from builder", + )); + }; + Ok(block) + } +} + +enum FieldAccumulator { + Fixed(FixedFieldAccumulator), + Variable32 { + builder: DataBlockBuilder, + empty_value: DataBlock, + }, + Variable64 { + builder: DataBlockBuilder, + empty_value: DataBlock, + }, +} + +impl FieldAccumulator { + // In full-zip variable packed decoding, rep/def may produce a visible row + // with an empty payload (e.g. null/invalid item). We still need to append + // one placeholder per child so child row counts remain aligned. + fn append_empty(&mut self) -> Result<()> { + match self { + Self::Fixed(fixed_field_accumulator) => fixed_field_accumulator.append_empty(), + Self::Variable32 { + builder, + empty_value, + } => builder.append(empty_value, 0..1), + Self::Variable64 { + builder, + empty_value, + } => builder.append(empty_value, 0..1), + } + } +} + +impl VariablePerValueDecompressor for PackedStructVariablePerValueDecompressor { + fn decompress(&self, data: VariableWidthBlock) -> Result { + let num_values = data.num_values; + let offsets_u64 = match data.bits_per_offset { + 32 => data + .offsets + .borrow_to_typed_slice::() + .iter() + .map(|v| *v as u64) + .collect::>(), + 64 => data + .offsets + .borrow_to_typed_slice::() + .as_ref() + .to_vec(), + _ => { + return Err(Error::invalid_input( + "Packed struct row offsets must be 32 or 64 bits", + )); + } + }; + + if offsets_u64.len() != num_values as usize + 1 { + return Err(Error::invalid_input( + "Packed struct row offsets length mismatch", + )); + } + + let mut accumulators = Vec::with_capacity(self.fields.len()); + for field in &self.fields { + match &field.kind { + VariablePackedStructFieldKind::Fixed { bits_per_value, .. } => { + let accumulator = FixedFieldAccumulator::new(*bits_per_value, num_values)?; + accumulators.push(FieldAccumulator::Fixed(accumulator)); + } + VariablePackedStructFieldKind::Variable { + bits_per_length, .. + } => match bits_per_length { + 32 => accumulators.push(FieldAccumulator::Variable32 { + builder: DataBlockBuilder::with_capacity_estimate(data.data.len() as u64), + empty_value: DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::empty(), + bits_per_offset: 32, + offsets: LanceBuffer::reinterpret_vec(vec![0_u32, 0_u32]), + num_values: 1, + block_info: BlockInfo::new(), + }), + }), + 64 => accumulators.push(FieldAccumulator::Variable64 { + builder: DataBlockBuilder::with_capacity_estimate(data.data.len() as u64), + empty_value: DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::empty(), + bits_per_offset: 64, + offsets: LanceBuffer::reinterpret_vec(vec![0_u64, 0_u64]), + num_values: 1, + block_info: BlockInfo::new(), + }), + }), + _ => { + return Err(Error::invalid_input( + "Packed struct variable child must use 32 or 64-bit length prefixes", + )); + } + }, + } + } + + for row_idx in 0..num_values as usize { + let row_start = offsets_u64[row_idx] as usize; + let row_end = offsets_u64[row_idx + 1] as usize; + if row_end > data.data.len() || row_start > row_end { + return Err(Error::invalid_input( + "Packed struct row bounds exceed buffer", + )); + } + if row_start == row_end { + for accumulator in accumulators.iter_mut() { + accumulator.append_empty()?; + } + continue; + } + let mut cursor = row_start; + for (field, accumulator) in self.fields.iter().zip(accumulators.iter_mut()) { + match (&field.kind, accumulator) { + ( + VariablePackedStructFieldKind::Fixed { bits_per_value, .. }, + FieldAccumulator::Fixed(fixed_accumulator), + ) => { + debug_assert_eq!(*bits_per_value, fixed_accumulator.bits_per_value); + let bytes_per_value = (*bits_per_value / 8) as usize; + let end = cursor + bytes_per_value; + if end > row_end { + return Err(Error::invalid_input( + "Packed struct fixed child exceeds row bounds", + )); + } + let value_block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(data.data[cursor..end].to_vec()), + bits_per_value: *bits_per_value, + num_values: 1, + block_info: BlockInfo::new(), + }); + fixed_accumulator.builder.append(&value_block, 0..1)?; + cursor = end; + } + ( + VariablePackedStructFieldKind::Variable { + bits_per_length, .. + }, + FieldAccumulator::Variable32 { builder, .. }, + ) => { + if *bits_per_length != 32 { + return Err(Error::invalid_input( + "Packed struct length prefix size mismatch", + )); + } + let end = cursor + std::mem::size_of::(); + if end > row_end { + return Err(Error::invalid_input( + "Packed struct variable child length prefix out of bounds", + )); + } + let len = u32::from_le_bytes( + data.data[cursor..end] + .try_into() + .expect("slice has exact length"), + ) as usize; + cursor = end; + let value_end = cursor + len; + if value_end > row_end { + return Err(Error::invalid_input( + "Packed struct variable child exceeds row bounds", + )); + } + let value_block = DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::from(data.data[cursor..value_end].to_vec()), + bits_per_offset: 32, + offsets: LanceBuffer::reinterpret_vec(vec![0_u32, len as u32]), + num_values: 1, + block_info: BlockInfo::new(), + }); + builder.append(&value_block, 0..1)?; + cursor = value_end; + } + ( + VariablePackedStructFieldKind::Variable { + bits_per_length, .. + }, + FieldAccumulator::Variable64 { builder, .. }, + ) => { + if *bits_per_length != 64 { + return Err(Error::invalid_input( + "Packed struct length prefix size mismatch", + )); + } + let end = cursor + std::mem::size_of::(); + if end > row_end { + return Err(Error::invalid_input( + "Packed struct variable child length prefix out of bounds", + )); + } + let len = u64::from_le_bytes( + data.data[cursor..end] + .try_into() + .expect("slice has exact length"), + ) as usize; + cursor = end; + let value_end = cursor + len; + if value_end > row_end { + return Err(Error::invalid_input( + "Packed struct variable child exceeds row bounds", + )); + } + let value_block = DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::from(data.data[cursor..value_end].to_vec()), + bits_per_offset: 64, + offsets: LanceBuffer::reinterpret_vec(vec![0_u64, len as u64]), + num_values: 1, + block_info: BlockInfo::new(), + }); + builder.append(&value_block, 0..1)?; + cursor = value_end; + } + _ => { + return Err(Error::invalid_input( + "Packed struct accumulator kind mismatch", + )); + } + } + } + if cursor != row_end { + return Err(Error::invalid_input( + "Packed struct row parsing did not consume full row", + )); + } + } + + let mut children = Vec::with_capacity(self.fields.len()); + for (field, accumulator) in self.fields.iter().zip(accumulators) { + match (field, accumulator) { + ( + VariablePackedStructFieldDecoder { + kind: VariablePackedStructFieldKind::Fixed { decompressor, .. }, + }, + FieldAccumulator::Fixed(fixed_accumulator), + ) => { + let finished_accumulator = fixed_accumulator.finish()?; + let decoded = decompressor.decompress(finished_accumulator, num_values)?; + children.push(decoded); + } + ( + VariablePackedStructFieldDecoder { + kind: + VariablePackedStructFieldKind::Variable { + bits_per_length, + decompressor, + }, + }, + FieldAccumulator::Variable32 { builder, .. }, + ) => { + let DataBlock::VariableWidth(mut block) = builder.finish() else { + panic!("Expected variable-width datablock from builder"); + }; + debug_assert_eq!(block.bits_per_offset, 32); + block.bits_per_offset = (*bits_per_length) as u8; + let decoded = decompressor.decompress(block)?; + children.push(decoded); + } + ( + VariablePackedStructFieldDecoder { + kind: + VariablePackedStructFieldKind::Variable { + bits_per_length, + decompressor, + }, + }, + FieldAccumulator::Variable64 { builder, .. }, + ) => { + let DataBlock::VariableWidth(mut block) = builder.finish() else { + panic!("Expected variable-width datablock from builder"); + }; + debug_assert_eq!(block.bits_per_offset, 64); + block.bits_per_offset = (*bits_per_length) as u8; + let decoded = decompressor.decompress(block)?; + children.push(decoded); + } + _ => { + return Err(Error::invalid_input( + "Packed struct accumulator mismatch during finalize", + )); + } + } + } + + Ok(DataBlock::Struct(StructDataBlock { + children, + block_info: BlockInfo::new(), + validity: None, + })) + } +} + +#[derive(Debug)] +struct PackedStructFixedFieldDecoder { + bits_per_value: u64, + decompressor: Box, +} + +#[derive(Debug)] +pub(crate) struct PackedStructFixedPerValueDecompressor { + decoders: Vec, +} + +impl PackedStructFixedPerValueDecompressor { + pub(crate) fn new(description: &PackedStruct) -> Result { + let compression = description + .values + .as_ref() + .ok_or_else(|| Error::invalid_input("PackedStruct missing values encoding"))? + .compression + .as_ref() + .ok_or_else(|| { + Error::invalid_input("PackedStruct values missing compression encoding") + })?; + + // The encoder always flat-encodes each child, so that is the only layout we can decode. + if !matches!(compression, Compression::Flat(..)) { + return Err(Error::invalid_input( + "PackedStruct fixed encoding currently requires flat compression", + )); + } + + let decoders = description + .bits_per_value + .iter() + .map(|&bits_per_value| { + let flat = crate::format::pb21::Flat { + bits_per_value, + data: None, + }; + PackedStructFixedFieldDecoder { + bits_per_value, + decompressor: Box::new(ValueDecompressor::from_flat(&flat)), + } + }) + .collect(); + Ok(Self { decoders }) + } +} + +impl FixedPerValueDecompressor for PackedStructFixedPerValueDecompressor { + fn decompress(&self, data: FixedWidthDataBlock, num_values: u64) -> Result { + if !data.bits_per_value.is_multiple_of(8) { + return Err(Error::invalid_input( + "Packed struct fixed encoding requires byte-aligned children", + )); + } + let bytes_per_row = (data.bits_per_value / 8) as usize; + + // Byte offset of each child within a packed row (a running prefix sum of the + // child widths). The final offset must equal the packed row width. + let mut child_bytes = Vec::with_capacity(self.decoders.len()); + for decoder in &self.decoders { + if !decoder.bits_per_value.is_multiple_of(8) { + return Err(Error::invalid_input( + "Packed struct fixed child must be byte-aligned", + )); + } + child_bytes.push((decoder.bits_per_value / 8) as usize); + } + if child_bytes.iter().sum::() != bytes_per_row { + return Err(Error::invalid_input( + "Packed struct child widths do not sum to the packed row width", + )); + } + if bytes_per_row.saturating_mul(num_values as usize) > data.data.len() { + return Err(Error::invalid_input( + "Packed struct row bounds exceed buffer", + )); + } + + // Un-zip the row-major buffer one child at a time by gathering that child's + // slice out of every row, then hand the column to the child decompressor. + let bytes = data.data.as_ref(); + let mut children = Vec::with_capacity(self.decoders.len()); + let mut field_offset = 0; + for (decoder, &field_bytes) in self.decoders.iter().zip(child_bytes.iter()) { + let mut child_buf = Vec::with_capacity(field_bytes * num_values as usize); + for row_idx in 0..num_values as usize { + let start = row_idx * bytes_per_row + field_offset; + child_buf.extend_from_slice(&bytes[start..start + field_bytes]); + } + let child_block = FixedWidthDataBlock { + data: LanceBuffer::from(child_buf), + bits_per_value: decoder.bits_per_value, + num_values, + block_info: BlockInfo::new(), + }; + children.push(decoder.decompressor.decompress(child_block, num_values)?); + field_offset += field_bytes; + } + + Ok(DataBlock::Struct(StructDataBlock { + children, + block_info: BlockInfo::new(), + validity: None, + })) + } + + fn bits_per_value(&self) -> u64 { + self.decoders + .iter() + .map(|decoder| decoder.bits_per_value) + .sum() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + compression::CompressionStrategy, + compression::DefaultDecompressionStrategy, + compression_config::CompressionParams, + constants::{ + PACKED_STRUCT_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, + }, + statistics::ComputeStat, + testing::{ + TestCases, TestEncoding, check_round_trip_encoding_of_data, test_compression_strategy, + }, + }; + use arrow_array::{ + Array, ArrayRef, BinaryArray, Int32Array, Int64Array, LargeStringArray, StringArray, + StructArray, UInt32Array, + }; + use arrow_schema::{DataType, Field as ArrowField, Fields}; + use std::collections::HashMap; + use std::sync::Arc; + + fn fixed_block_from_array(array: Int64Array) -> FixedWidthDataBlock { + let num_values = array.len() as u64; + let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values); + match block { + DataBlock::FixedWidth(block) => block, + _ => panic!("Expected fixed-width data block"), + } + } + + fn fixed_i32_block_from_array(array: Int32Array) -> FixedWidthDataBlock { + let num_values = array.len() as u64; + let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values); + match block { + DataBlock::FixedWidth(block) => block, + _ => panic!("Expected fixed-width data block"), + } + } + + fn variable_block_from_string_array(array: StringArray) -> VariableWidthBlock { + let num_values = array.len() as u64; + let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values); + match block { + DataBlock::VariableWidth(block) => block, + _ => panic!("Expected variable-width block"), + } + } + + fn variable_block_from_large_string_array(array: LargeStringArray) -> VariableWidthBlock { + let num_values = array.len() as u64; + let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values); + match block { + DataBlock::VariableWidth(block) => block, + _ => panic!("Expected variable-width block"), + } + } + + fn variable_block_from_binary_array(array: BinaryArray) -> VariableWidthBlock { + let num_values = array.len() as u64; + let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values); + match block { + DataBlock::VariableWidth(block) => block, + _ => panic!("Expected variable-width block"), + } + } + + #[test] + fn variable_packed_struct_round_trip() -> Result<()> { + let arrow_fields: Fields = vec![ + ArrowField::new("id", DataType::UInt32, false), + ArrowField::new("name", DataType::Utf8, true), + ] + .into(); + let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false); + let struct_field = Field::try_from(&arrow_struct)?; + + let ids = vec![1_u32, 2, 42]; + let id_bytes = ids + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let mut id_block = FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(ids), + bits_per_value: 32, + num_values: 3, + block_info: BlockInfo::new(), + }; + id_block.compute_stat(); + let id_block = DataBlock::FixedWidth(id_block); + + let name_offsets = vec![0_i32, 1, 4, 4]; + let name_bytes = b"abcz".to_vec(); + let mut name_block = VariableWidthBlock { + data: LanceBuffer::from(name_bytes.clone()), + bits_per_offset: 32, + offsets: LanceBuffer::reinterpret_vec(name_offsets.clone()), + num_values: 3, + block_info: BlockInfo::new(), + }; + name_block.compute_stat(); + let name_block = DataBlock::VariableWidth(name_block); + + let struct_block = StructDataBlock { + children: vec![id_block, name_block], + block_info: BlockInfo::new(), + validity: None, + }; + + let data_block = DataBlock::Struct(struct_block); + + let compression_strategy = + test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default()); + let compressor = crate::compression::CompressionStrategy::create_per_value( + compression_strategy.as_ref(), + &struct_field, + &data_block, + )?; + let (compressed, encoding) = compressor.compress(data_block)?; + + let PerValueDataBlock::Variable(zipped) = compressed else { + panic!("expected variable-width packed struct output"); + }; + + let decompression_strategy = DefaultDecompressionStrategy::default(); + let decompressor = + crate::compression::DecompressionStrategy::create_variable_per_value_decompressor( + &decompression_strategy, + &encoding, + )?; + let decoded = decompressor.decompress(zipped)?; + + let DataBlock::Struct(decoded_struct) = decoded else { + panic!("expected struct datablock after decode"); + }; + + let decoded_id = decoded_struct.children[0].as_fixed_width_ref().unwrap(); + assert_eq!(decoded_id.bits_per_value, 32); + assert_eq!(decoded_id.data.as_ref(), id_bytes.as_slice()); + + let decoded_name = decoded_struct.children[1].as_variable_width_ref().unwrap(); + assert_eq!(decoded_name.bits_per_offset, 32); + let decoded_offsets = decoded_name.offsets.borrow_to_typed_slice::(); + assert_eq!(decoded_offsets.as_ref(), name_offsets.as_slice()); + assert_eq!(decoded_name.data.as_ref(), name_bytes.as_slice()); + + Ok(()) + } + + #[test] + fn variable_packed_struct_large_utf8_round_trip() -> Result<()> { + let arrow_fields: Fields = vec![ + ArrowField::new("value", DataType::Int64, false), + ArrowField::new("text", DataType::LargeUtf8, false), + ] + .into(); + let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false); + let struct_field = Field::try_from(&arrow_struct)?; + + let id_block = fixed_block_from_array(Int64Array::from(vec![10, 20, 30, 40])); + let payload_array = LargeStringArray::from(vec![ + "alpha", + "a considerably longer payload for testing", + "mid", + "z", + ]); + let payload_block = variable_block_from_large_string_array(payload_array); + + let struct_block = StructDataBlock { + children: vec![ + DataBlock::FixedWidth(id_block.clone()), + DataBlock::VariableWidth(payload_block.clone()), + ], + block_info: BlockInfo::new(), + validity: None, + }; + + let data_block = DataBlock::Struct(struct_block); + + let compression_strategy = + test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default()); + let compressor = crate::compression::CompressionStrategy::create_per_value( + compression_strategy.as_ref(), + &struct_field, + &data_block, + )?; + let (compressed, encoding) = compressor.compress(data_block)?; + + let PerValueDataBlock::Variable(zipped) = compressed else { + panic!("expected variable-width packed struct output"); + }; + + let decompression_strategy = DefaultDecompressionStrategy::default(); + let decompressor = + crate::compression::DecompressionStrategy::create_variable_per_value_decompressor( + &decompression_strategy, + &encoding, + )?; + let decoded = decompressor.decompress(zipped)?; + + let DataBlock::Struct(decoded_struct) = decoded else { + panic!("expected struct datablock after decode"); + }; + + let decoded_id = decoded_struct.children[0].as_fixed_width_ref().unwrap(); + assert_eq!(decoded_id.bits_per_value, 64); + assert_eq!(decoded_id.data.as_ref(), id_block.data.as_ref()); + + let decoded_payload = decoded_struct.children[1].as_variable_width_ref().unwrap(); + assert_eq!(decoded_payload.bits_per_offset, 64); + assert_eq!( + decoded_payload + .offsets + .borrow_to_typed_slice::() + .as_ref(), + payload_block + .offsets + .borrow_to_typed_slice::() + .as_ref() + ); + assert_eq!(decoded_payload.data.as_ref(), payload_block.data.as_ref()); + + Ok(()) + } + + #[tokio::test] + async fn variable_packed_struct_utf8_round_trip() { + // schema: Struct + let fields = Fields::from(vec![ + Arc::new(ArrowField::new("id", DataType::UInt32, false)), + Arc::new(ArrowField::new("uri", DataType::Utf8, false)), + Arc::new(ArrowField::new("long_text", DataType::LargeUtf8, false)), + ]); + + // mark struct as packed + let mut meta = HashMap::new(); + meta.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string()); + + let array = Arc::new(StructArray::from(vec![ + ( + fields[0].clone(), + Arc::new(UInt32Array::from(vec![1, 2, 3])) as ArrayRef, + ), + ( + fields[1].clone(), + Arc::new(StringArray::from(vec![ + Some("a"), + Some("b"), + Some("/tmp/x"), + ])) as ArrayRef, + ), + ( + fields[2].clone(), + Arc::new(LargeStringArray::from(vec![ + Some("alpha"), + Some("a considerably longer payload for testing"), + Some("mid"), + ])) as ArrayRef, + ), + ])); + + let test_cases = TestCases::default() + .with_u32_structural_encodings() + .with_expected_encoding("variable_packed_struct"); + + check_round_trip_encoding_of_data(vec![array], &test_cases, meta).await; + } + + #[test] + fn variable_packed_struct_multi_variable_round_trip() -> Result<()> { + let arrow_fields: Fields = vec![ + ArrowField::new("category", DataType::Utf8, false), + ArrowField::new("payload", DataType::Binary, false), + ArrowField::new("count", DataType::Int32, false), + ] + .into(); + let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false); + let struct_field = Field::try_from(&arrow_struct)?; + + let category_array = StringArray::from(vec!["red", "blue", "green", "red"]); + let category_block = variable_block_from_string_array(category_array); + let payload_values: Vec> = + vec![vec![0x01, 0x02], vec![], vec![0x05, 0x06, 0x07], vec![0xff]]; + let payload_array = + BinaryArray::from_iter_values(payload_values.iter().map(|v| v.as_slice())); + let payload_block = variable_block_from_binary_array(payload_array); + let count_block = fixed_i32_block_from_array(Int32Array::from(vec![1, 2, 3, 4])); + + let struct_block = StructDataBlock { + children: vec![ + DataBlock::VariableWidth(category_block.clone()), + DataBlock::VariableWidth(payload_block.clone()), + DataBlock::FixedWidth(count_block.clone()), + ], + block_info: BlockInfo::new(), + validity: None, + }; + + let data_block = DataBlock::Struct(struct_block); + + let compression_strategy = + test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default()); + let compressor = crate::compression::CompressionStrategy::create_per_value( + compression_strategy.as_ref(), + &struct_field, + &data_block, + )?; + let (compressed, encoding) = compressor.compress(data_block)?; + + let PerValueDataBlock::Variable(zipped) = compressed else { + panic!("expected variable-width packed struct output"); + }; + + let decompression_strategy = DefaultDecompressionStrategy::default(); + let decompressor = + crate::compression::DecompressionStrategy::create_variable_per_value_decompressor( + &decompression_strategy, + &encoding, + )?; + let decoded = decompressor.decompress(zipped)?; + + let DataBlock::Struct(decoded_struct) = decoded else { + panic!("expected struct datablock after decode"); + }; + + let decoded_category = decoded_struct.children[0].as_variable_width_ref().unwrap(); + assert_eq!(decoded_category.bits_per_offset, 32); + assert_eq!( + decoded_category + .offsets + .borrow_to_typed_slice::() + .as_ref(), + category_block + .offsets + .borrow_to_typed_slice::() + .as_ref() + ); + assert_eq!(decoded_category.data.as_ref(), category_block.data.as_ref()); + + let decoded_payload = decoded_struct.children[1].as_variable_width_ref().unwrap(); + assert_eq!(decoded_payload.bits_per_offset, 32); + assert_eq!( + decoded_payload + .offsets + .borrow_to_typed_slice::() + .as_ref(), + payload_block + .offsets + .borrow_to_typed_slice::() + .as_ref() + ); + assert_eq!(decoded_payload.data.as_ref(), payload_block.data.as_ref()); + + let decoded_count = decoded_struct.children[2].as_fixed_width_ref().unwrap(); + assert_eq!(decoded_count.bits_per_value, 32); + assert_eq!(decoded_count.data.as_ref(), count_block.data.as_ref()); + + Ok(()) + } + + #[test] + fn variable_packed_struct_requires_v22() { + let arrow_fields: Fields = vec![ + ArrowField::new("value", DataType::Int64, false), + ArrowField::new("text", DataType::Utf8, false), + ] + .into(); + let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false); + let struct_field = Field::try_from(&arrow_struct).unwrap(); + + let value_block = fixed_block_from_array(Int64Array::from(vec![1, 2, 3])); + let text_block = + variable_block_from_string_array(StringArray::from(vec!["a", "bb", "ccc"])); + + let struct_block = StructDataBlock { + children: vec![ + DataBlock::FixedWidth(value_block), + DataBlock::VariableWidth(text_block), + ], + block_info: BlockInfo::new(), + validity: None, + }; + + let compression_strategy = + test_compression_strategy(TestEncoding::StructuralU16, CompressionParams::default()); + let result = + compression_strategy.create_per_value(&struct_field, &DataBlock::Struct(struct_block)); + + assert!(matches!(result, Err(Error::NotSupported { .. }))); + } + + #[test] + fn variable_packed_struct_decompress_empty_row() -> Result<()> { + let strategy = DefaultDecompressionStrategy::default(); + let fixed_decompressor = Arc::from( + crate::compression::DecompressionStrategy::create_fixed_per_value_decompressor( + &strategy, + &ProtobufUtils21::flat(32, None), + )?, + ); + let variable_decompressor = Arc::from( + crate::compression::DecompressionStrategy::create_variable_per_value_decompressor( + &strategy, + &ProtobufUtils21::variable(ProtobufUtils21::flat(32, None), None), + )?, + ); + + let decompressor = PackedStructVariablePerValueDecompressor::new(vec![ + VariablePackedStructFieldDecoder { + kind: VariablePackedStructFieldKind::Fixed { + bits_per_value: 32, + decompressor: fixed_decompressor, + }, + }, + VariablePackedStructFieldDecoder { + kind: VariablePackedStructFieldKind::Variable { + bits_per_length: 32, + decompressor: variable_decompressor, + }, + }, + ]); + + let mut row_data = Vec::new(); + row_data.extend_from_slice(&1_u32.to_le_bytes()); + row_data.extend_from_slice(&1_u32.to_le_bytes()); + row_data.extend_from_slice(b"a"); + row_data.extend_from_slice(&2_u32.to_le_bytes()); + row_data.extend_from_slice(&0_u32.to_le_bytes()); + + let input = VariableWidthBlock { + data: LanceBuffer::from(row_data), + bits_per_offset: 32, + offsets: LanceBuffer::reinterpret_vec(vec![0_u32, 9_u32, 9_u32, 17_u32]), + num_values: 3, + block_info: BlockInfo::new(), + }; + + let decoded = decompressor.decompress(input)?; + let DataBlock::Struct(decoded_struct) = decoded else { + panic!("expected struct output"); + }; + + let fixed = decoded_struct.children[0].as_fixed_width_ref().unwrap(); + assert_eq!(fixed.bits_per_value, 32); + assert_eq!( + fixed.data.borrow_to_typed_slice::().as_ref(), + &[1, 0, 2] + ); + + let variable = decoded_struct.children[1].as_variable_width_ref().unwrap(); + assert_eq!(variable.bits_per_offset, 32); + assert_eq!( + variable.offsets.borrow_to_typed_slice::().as_ref(), + &[0_u32, 1_u32, 1_u32, 1_u32] + ); + assert_eq!(variable.data.as_ref(), b"a"); + + Ok(()) + } + + #[test] + fn fixed_packed_struct_round_trip() -> Result<()> { + let arrow_fields: Fields = vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("value", DataType::Int64, false), + ] + .into(); + let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false); + let struct_field = Field::try_from(&arrow_struct)?; + + let id_block = fixed_i32_block_from_array(Int32Array::from(vec![1, 2, 3, 4])); + let value_block = fixed_block_from_array(Int64Array::from(vec![10, 20, 30, 40])); + + let struct_block = StructDataBlock { + children: vec![ + DataBlock::FixedWidth(id_block.clone()), + DataBlock::FixedWidth(value_block.clone()), + ], + block_info: BlockInfo::new(), + validity: None, + }; + + let data_block = DataBlock::Struct(struct_block); + + let compression_strategy = + test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default()); + let compressor = CompressionStrategy::create_per_value( + compression_strategy.as_ref(), + &struct_field, + &data_block, + )?; + let (compressed, encoding) = compressor.compress(data_block)?; + + let PerValueDataBlock::Fixed(zipped) = compressed else { + panic!("expected fixed-width packed struct output"); + }; + + let decompression_strategy = DefaultDecompressionStrategy::default(); + let decompressor = + crate::compression::DecompressionStrategy::create_fixed_per_value_decompressor( + &decompression_strategy, + &encoding, + )?; + let decoded = decompressor.decompress(zipped, 4)?; + + let DataBlock::Struct(decoded_struct) = decoded else { + panic!("expected struct datablock after decode"); + }; + + let decoded_id = decoded_struct.children[0].as_fixed_width_ref().unwrap(); + assert_eq!(decoded_id.bits_per_value, 32); + assert_eq!(decoded_id.data.as_ref(), id_block.data.as_ref()); + + let decoded_value = decoded_struct.children[1].as_fixed_width_ref().unwrap(); + assert_eq!(decoded_value.bits_per_value, 64); + assert_eq!(decoded_value.data.as_ref(), value_block.data.as_ref()); + + Ok(()) + } + + // End-to-end round trip through the file writer. Requesting full-zip on an + // all-fixed-width struct routes to `PackedStructFixedPerValueEncoder` (mini-block + // is only chosen for narrow structs), exercising the writer/reader wiring rather + // than just the block-level compress/decompress above. + #[tokio::test] + async fn fixed_packed_struct_full_zip_round_trip() { + let fields = Fields::from(vec![ + Arc::new(ArrowField::new("id", DataType::Int32, false)), + Arc::new(ArrowField::new("value", DataType::Int64, false)), + ]); + + let mut meta = HashMap::new(); + meta.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string()); + meta.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_FULLZIP.to_string(), + ); + + let array = Arc::new(StructArray::from(vec![ + ( + fields[0].clone(), + Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef, + ), + ( + fields[1].clone(), + Arc::new(Int64Array::from(vec![10, 20, 30, 40])) as ArrayRef, + ), + ])); + + let test_cases = TestCases::default() + .with_u32_structural_encodings() + .with_expected_encoding("packed_struct"); + + check_round_trip_encoding_of_data(vec![array], &test_cases, meta).await; + } + + #[test] + fn fixed_packed_struct_rejects_variable_child() -> Result<()> { + let arrow_fields: Fields = vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("name", DataType::Utf8, false), + ] + .into(); + let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false); + let struct_field = Field::try_from(&arrow_struct)?; + + let struct_block = DataBlock::Struct(StructDataBlock { + children: vec![ + DataBlock::FixedWidth(fixed_i32_block_from_array(Int32Array::from(vec![1, 2]))), + DataBlock::VariableWidth(variable_block_from_string_array(StringArray::from( + vec!["a", "bb"], + ))), + ], + block_info: BlockInfo::new(), + validity: None, + }); + + let encoder = PackedStructFixedPerValueEncoder::new(struct_field.children); + let err = encoder.compress(struct_block).unwrap_err(); + assert!(matches!(&err, Error::InvalidInput { .. })); + assert!( + err.to_string().contains("fixed-width"), + "unexpected error: {err}" + ); + + Ok(()) + } +} diff --git a/vendor/lance-encoding/src/encodings/physical/rle.rs b/vendor/lance-encoding/src/encodings/physical/rle.rs new file mode 100644 index 000000000..8d4bb2003 --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/rle.rs @@ -0,0 +1,3408 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! # RLE (Run-Length Encoding) +//! +//! RLE compression for Lance, optimized for data with repeated values. +//! +//! ## Encoding Format +//! +//! RLE uses a dual-buffer format to store compressed data: +//! +//! - **Values Buffer**: Stores unique values in their original data type +//! - **Lengths Buffer**: Stores the repeat count for each value as u8, u16, or u32 +//! +//! ### Example +//! +//! Input data: `[1, 1, 1, 2, 2, 3, 3, 3, 3]` +//! +//! Encoded as: +//! - Values buffer: `[1, 2, 3]` (3 × 4 bytes for i32) +//! - Lengths buffer: `[3, 2, 4]` (3 × 1 byte for u8 in compatibility mode) +//! +//! ### Long Run Handling +//! +//! In compatibility mode, when a run exceeds 255 values, it is split into multiple +//! runs of 255 followed by a final run with the remainder. RLE v2 can use u16 or +//! u32 run lengths to reduce this splitting. +//! +//! ## Supported Types +//! +//! RLE supports all fixed-width primitive types: +//! - 8-bit: u8, i8 +//! - 16-bit: u16, i16 +//! - 32-bit: u32, i32, f32 +//! - 64-bit: u64, i64, f64 +//! +//! ## Compression Strategy +//! +//! RLE is automatically selected when: +//! - The run count (number of value transitions) < 50% of total values +//! - This indicates sufficient repetition for RLE to be effective +//! +//! ## MiniBlock Chunk Handling +//! +//! When used in the miniblock path, all chunks share two global buffers (values and lengths). +//! Each chunk's `buffer_sizes` identifies its slice within those global buffers. Non-last chunks +//! contain a power-of-2 number of values. +//! +//! NOTE: The current encoder uses a 2048-value cap per chunk as a workaround for +//! . +//! +//! ## Block Format +//! +//! When used in the block compression path, the encoded output is a single buffer: +//! `[8-byte header: values buffer size][values buffer][run_lengths buffer]`. + +use arrow_buffer::{ArrowNativeType, ScalarBuffer}; +use log::trace; + +use crate::buffer::LanceBuffer; +use crate::compression::{ + BlockCompressor, BlockDecompressor, MiniBlockDecompressor, require_block_payload, +}; +use crate::data::DataBlock; +use crate::data::{BlockInfo, FixedWidthDataBlock}; +use crate::encodings::logical::primitive::miniblock::{ + MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, + MiniBlockCompressionContext, MiniBlockCompressor, +}; +use crate::encodings::physical::block::{CompressionConfig, GeneralBufferCompressor}; +use crate::format::ProtobufUtils21; +use crate::format::pb21::CompressiveEncoding; + +use lance_core::{Error, Result}; + +/// Width used to encode RLE run lengths. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RunLengthWidth { + /// Compatibility mode. Runs longer than 255 values are split. + U8, + /// RLE v2 mode for runs up to 65,535 values per entry. + U16, + /// RLE v2 mode for runs up to 4,294,967,295 values per entry. + U32, +} + +impl RunLengthWidth { + pub(crate) fn from_bits(bits_per_value: u64) -> Option { + match bits_per_value { + 8 => Some(Self::U8), + 16 => Some(Self::U16), + 32 => Some(Self::U32), + _ => None, + } + } + + pub(crate) fn bits_per_value(self) -> u64 { + match self { + Self::U8 => 8, + Self::U16 => 16, + Self::U32 => 32, + } + } + + fn bytes_per_value(self) -> usize { + match self { + Self::U8 => 1, + Self::U16 => 2, + Self::U32 => 4, + } + } + + fn max_run_length(self) -> u64 { + match self { + Self::U8 => u8::MAX as u64, + Self::U16 => u16::MAX as u64, + Self::U32 => u32::MAX as u64, + } + } + + fn write_length(self, length: u64, dst: &mut Vec) { + match self { + Self::U8 => dst.push(length as u8), + Self::U16 => dst.extend_from_slice(&(length as u16).to_le_bytes()), + Self::U32 => dst.extend_from_slice(&(length as u32).to_le_bytes()), + } + } + + fn read_length(self, bytes: &[u8]) -> u64 { + match self { + Self::U8 => bytes[0] as u64, + Self::U16 => u16::from_le_bytes([bytes[0], bytes[1]]) as u64, + Self::U32 => u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as u64, + } + } +} + +const RUN_LENGTH_WIDTHS: [RunLengthWidth; 3] = + [RunLengthWidth::U8, RunLengthWidth::U16, RunLengthWidth::U32]; + +/// Select the lowest-cost run length width from precomputed entry counts. +pub(crate) fn select_run_length_width_from_entries( + entries: &[u64], + bits_per_value: u64, +) -> Result<(RunLengthWidth, u128)> { + if entries.len() != RUN_LENGTH_WIDTHS.len() { + return Err(Error::invalid_input_source( + format!( + "RLE run length entry statistics must have {} values, got {}", + RUN_LENGTH_WIDTHS.len(), + entries.len() + ) + .into(), + )); + } + + if !matches!(bits_per_value, 8 | 16 | 32 | 64) { + return Err(Error::invalid_input_source( + format!("RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}") + .into(), + )); + } + + let mut best_width = RUN_LENGTH_WIDTHS[0]; + let mut best_cost = rle_encoded_size_from_entries(entries[0], bits_per_value, best_width); + for (&width, &entry_count) in RUN_LENGTH_WIDTHS.iter().zip(entries.iter()).skip(1) { + let cost = rle_encoded_size_from_entries(entry_count, bits_per_value, width); + if cost < best_cost { + best_width = width; + best_cost = cost; + } + } + + Ok((best_width, best_cost)) +} + +pub(crate) fn rle_encoded_size_from_entries( + entry_count: u64, + bits_per_value: u64, + run_length_width: RunLengthWidth, +) -> u128 { + let bytes_per_value = (bits_per_value / 8) as u128; + let bytes_per_length = run_length_width.bytes_per_value() as u128; + (entry_count as u128) * (bytes_per_value + bytes_per_length) +} + +pub(crate) fn run_length_width_index(run_length_width: RunLengthWidth) -> usize { + match run_length_width { + RunLengthWidth::U8 => 0, + RunLengthWidth::U16 => 1, + RunLengthWidth::U32 => 2, + } +} + +pub(crate) fn select_run_length_width( + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + max_segment_values: Option, +) -> Result<(RunLengthWidth, u128)> { + let entries = collect_run_length_entries(data, num_values, bits_per_value, max_segment_values)?; + select_run_length_width_from_entries(&entries, bits_per_value) +} + +pub(crate) fn rle_encoded_size( + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + max_segment_values: Option, + run_length_width: RunLengthWidth, +) -> Result { + let entries = collect_run_length_entries(data, num_values, bits_per_value, max_segment_values)?; + let width_idx = run_length_width_index(run_length_width); + Ok(rle_encoded_size_from_entries( + entries[width_idx], + bits_per_value, + run_length_width, + )) +} + +fn collect_run_length_entries( + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + max_segment_values: Option, +) -> Result<[u64; 3]> { + let num_values = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + + macro_rules! collect_entries { + ($ty:ty) => {{ + let type_size = std::mem::size_of::<$ty>(); + let expected_bytes = num_values.checked_mul(type_size).ok_or_else(|| { + Error::invalid_input_source( + format!( + "RLE input byte length overflow: {num_values} values of {type_size} bytes" + ) + .into(), + ) + })?; + if data.len() != expected_bytes { + return Err(Error::invalid_input_source( + format!( + "RLE input data size mismatch: {} bytes for {} values of {} bytes", + data.len(), + num_values, + type_size + ) + .into(), + )); + } + let values = data.borrow_to_typed_slice::<$ty>(); + let values = values.get(..num_values).ok_or_else(|| { + Error::invalid_input_source( + format!( + "RLE data has {} values but {} were expected", + values.len(), + num_values + ) + .into(), + ) + })?; + Ok(collect_run_length_entries_from_slice( + values, + max_segment_values, + )) + }}; + } + + match bits_per_value { + 8 => collect_entries!(u8), + 16 => collect_entries!(u16), + 32 => collect_entries!(u32), + 64 => collect_entries!(u64), + _ => Err(Error::invalid_input_source( + format!("RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}") + .into(), + )), + } +} + +fn collect_run_length_entries_from_slice( + values: &[T], + max_segment_values: Option, +) -> [u64; 3] { + if values.is_empty() { + return [0; 3]; + } + + let mut entries = [0u64; 3]; + let mut prev = values[0]; + let mut current_length = 1u64; + + for &value in &values[1..] { + if value != prev { + accumulate_run_length_entries(current_length, max_segment_values, &mut entries); + prev = value; + current_length = 1; + } else { + current_length += 1; + } + } + accumulate_run_length_entries(current_length, max_segment_values, &mut entries); + + entries +} + +pub(crate) fn accumulate_run_length_entries( + run_length: u64, + max_segment_values: Option, + entries: &mut [u64; 3], +) { + let max_segment_values = max_segment_values.unwrap_or(run_length).max(1); + let mut remaining = run_length; + while remaining > 0 { + let segment = remaining.min(max_segment_values); + for (idx, width) in RUN_LENGTH_WIDTHS.iter().enumerate() { + let entry_count = segment.div_ceil(width.max_run_length()); + entries[idx] = entries[idx].saturating_add(entry_count); + } + remaining -= segment; + } +} + +/// RLE encoder for miniblock format +#[derive(Debug)] +pub struct RleEncoder { + run_length_width: RunLengthWidth, + values_compression: Option, + run_lengths_compression: Option, + use_child_bitpacking: bool, +} + +#[derive(Clone)] +struct RleChildCandidate { + encoding: CompressiveEncoding, + data: LanceBuffer, + chunk_sizes: Vec, + size: usize, + requires_num_values: bool, +} + +impl Default for RleEncoder { + fn default() -> Self { + Self::new() + } +} + +impl RleEncoder { + pub fn new() -> Self { + Self { + run_length_width: RunLengthWidth::U8, + values_compression: None, + run_lengths_compression: None, + use_child_bitpacking: false, + } + } + + pub(crate) fn with_run_length_width(run_length_width: RunLengthWidth) -> Self { + Self { + run_length_width, + values_compression: None, + run_lengths_compression: None, + use_child_bitpacking: false, + } + } + + pub(crate) fn with_child_encoding( + run_length_width: RunLengthWidth, + values_compression: Option, + run_lengths_compression: Option, + use_child_bitpacking: bool, + ) -> Self { + Self { + run_length_width, + values_compression, + run_lengths_compression, + use_child_bitpacking, + } + } + + fn encode_data( + &self, + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + ) -> Result<(Vec, Vec)> { + if num_values == 0 { + return Ok((Vec::new(), Vec::new())); + } + + let num_values = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + let bytes_per_value = (bits_per_value / 8) as usize; + let bytes_per_length = self.run_length_width.bytes_per_value(); + + // Pre-allocate global buffers with estimated capacity + // Assume average compression ratio of ~10:1 (10 values per run) + let estimated_runs = num_values / 10; + let mut all_values = Vec::with_capacity(estimated_runs * bytes_per_value); + let mut all_lengths = Vec::with_capacity(estimated_runs * bytes_per_length); + let mut chunks = Vec::new(); + + let mut offset = 0usize; + let mut values_remaining = num_values; + + while values_remaining > 0 { + let values_start = all_values.len(); + let lengths_start = all_lengths.len(); + + let (_num_runs, values_processed, is_last_chunk) = match bits_per_value { + 8 => self.encode_chunk_rolling::( + data, + offset, + values_remaining, + &mut all_values, + &mut all_lengths, + ), + 16 => self.encode_chunk_rolling::( + data, + offset, + values_remaining, + &mut all_values, + &mut all_lengths, + ), + 32 => self.encode_chunk_rolling::( + data, + offset, + values_remaining, + &mut all_values, + &mut all_lengths, + ), + 64 => self.encode_chunk_rolling::( + data, + offset, + values_remaining, + &mut all_values, + &mut all_lengths, + ), + _ => { + return Err(Error::invalid_input_source( + format!( + "RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}" + ) + .into(), + )); + } + }; + + if values_processed == 0 { + // A non-final chunk needs at least two values because log_num_values == 0 + // identifies the final chunk. Report an error instead of returning partial data. + return Err(Error::internal(format!( + "RLE encoder made no progress: values_remaining={values_remaining}, \ + offset={offset}, data_len={}, bits_per_value={bits_per_value}, \ + max_miniblock_values={}", + data.len(), + *MAX_MINIBLOCK_VALUES + ))); + } + + let log_num_values = if is_last_chunk { + 0 + } else { + assert!( + values_processed.is_power_of_two(), + "Non-last chunk must have power-of-2 values" + ); + values_processed.ilog2() as u8 + }; + + let values_size = all_values.len() - values_start; + let lengths_size = all_lengths.len() - lengths_start; + + let chunk = MiniBlockChunk { + buffer_sizes: vec![values_size as u32, lengths_size as u32], + log_num_values, + }; + + chunks.push(chunk); + + offset += values_processed; + values_remaining -= values_processed; + } + + // Return exactly two buffers: values and lengths + Ok(( + vec![ + LanceBuffer::from(all_values), + LanceBuffer::from(all_lengths), + ], + chunks, + )) + } + + fn encode_block_data( + &self, + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + ) -> Result> { + match bits_per_value { + 8 => self.encode_block_data_generic::(data, num_values), + 16 => self.encode_block_data_generic::(data, num_values), + 32 => self.encode_block_data_generic::(data, num_values), + 64 => self.encode_block_data_generic::(data, num_values), + _ => Err(Error::invalid_input_source( + format!( + "RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}" + ) + .into(), + )), + } + } + + fn encode_block_data_generic( + &self, + data: &LanceBuffer, + num_values: u64, + ) -> Result> + where + T: bytemuck::Pod + PartialEq + Copy + ArrowNativeType, + { + let num_values = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + let type_size = std::mem::size_of::(); + let expected_bytes = num_values.checked_mul(type_size).ok_or_else(|| { + Error::invalid_input_source( + format!("RLE input byte length overflow: {num_values} values of {type_size} bytes") + .into(), + ) + })?; + if data.len() != expected_bytes { + return Err(Error::invalid_input_source( + format!( + "RLE input data size mismatch: {} bytes for {} values of {} bytes", + data.len(), + num_values, + type_size + ) + .into(), + )); + } + if num_values == 0 { + return Ok(vec![LanceBuffer::empty(), LanceBuffer::empty()]); + } + + let values_ref = data.borrow_to_typed_slice::(); + let values = values_ref.as_ref(); + let estimated_runs = num_values / 10; + let mut all_values = Vec::with_capacity(estimated_runs * type_size); + let mut all_lengths = + Vec::with_capacity(estimated_runs * self.run_length_width.bytes_per_value()); + self.encode_values(values, &mut all_values, &mut all_lengths); + Ok(vec![ + LanceBuffer::from(all_values), + LanceBuffer::from(all_lengths), + ]) + } + + /// Encodes the largest valid mini-block prefix from `offset`. + fn encode_chunk_rolling( + &self, + data: &LanceBuffer, + offset: usize, + values_remaining: usize, + all_values: &mut Vec, + all_lengths: &mut Vec, + ) -> (usize, usize, bool) + where + T: bytemuck::Pod + PartialEq + Copy + std::fmt::Debug + ArrowNativeType, + { + let type_size = std::mem::size_of::(); + let chunk_start = offset * type_size; + let max_by_count = *MAX_MINIBLOCK_VALUES as usize; + let max_values = values_remaining.min(max_by_count); + let chunk_end = chunk_start + max_values * type_size; + + if chunk_start >= data.len() { + return (0, 0, false); + } + + let chunk_len = chunk_end.min(data.len()) - chunk_start; + let chunk_buffer = data.slice_with_length(chunk_start, chunk_len); + let typed_data_ref = chunk_buffer.borrow_to_typed_slice::(); + let typed_data: &[T] = typed_data_ref.as_ref(); + let max_values = max_values.min(typed_data.len()); + + if typed_data.is_empty() { + return (0, 0, false); + } + + let values_start = all_values.len(); + let all_remaining_values_fit = values_remaining <= max_by_count; + let encoded_size = self.encoded_size(&typed_data[..max_values]); + let (values_to_encode, is_last_chunk) = if all_remaining_values_fit + && encoded_size <= MAX_MINIBLOCK_BYTES as usize + { + (max_values, true) + } else if let Some(values_to_encode) = self.largest_power_of_two_prefix::(typed_data) { + (values_to_encode, false) + } else { + return (0, 0, false); + }; + + self.encode_values(&typed_data[..values_to_encode], all_values, all_lengths); + + let num_runs = (all_values.len() - values_start) / type_size; + (num_runs, values_to_encode, is_last_chunk) + } + + fn largest_power_of_two_prefix(&self, values: &[T]) -> Option + where + T: bytemuck::Pod + PartialEq + Copy, + { + let max_prefix = values.len().min(*MAX_MINIBLOCK_VALUES as usize); + let mut prefix = 1usize << max_prefix.ilog2(); + while prefix > 1 { + if self.encoded_size(&values[..prefix]) <= MAX_MINIBLOCK_BYTES as usize { + return Some(prefix); + } + prefix >>= 1; + } + None + } + + fn encoded_size(&self, values: &[T]) -> usize + where + T: bytemuck::Pod + PartialEq + Copy, + { + if values.is_empty() { + return 0; + } + + let mut current_value = values[0]; + let mut current_length = 1u64; + let mut encoded_size = 0usize; + + for &value in values.iter().skip(1) { + if value == current_value { + current_length += 1; + } else { + encoded_size += self.run_size::(current_length); + current_value = value; + current_length = 1; + } + } + encoded_size += self.run_size::(current_length); + encoded_size + } + + fn run_size(&self, length: u64) -> usize + where + T: bytemuck::Pod, + { + let type_size = std::mem::size_of::(); + let run_chunks = length.div_ceil(self.run_length_width.max_run_length()) as usize; + run_chunks * (type_size + self.run_length_width.bytes_per_value()) + } + + fn encode_values(&self, values: &[T], all_values: &mut Vec, all_lengths: &mut Vec) + where + T: bytemuck::Pod + PartialEq + Copy, + { + if values.is_empty() { + return; + } + + let mut current_value = values[0]; + let mut current_length = 1u64; + + for &value in values.iter().skip(1) { + if value == current_value { + current_length += 1; + } else { + self.add_run(¤t_value, current_length, all_values, all_lengths); + current_value = value; + current_length = 1; + } + } + self.add_run(¤t_value, current_length, all_values, all_lengths); + } + + fn add_run( + &self, + value: &T, + length: u64, + all_values: &mut Vec, + all_lengths: &mut Vec, + ) -> usize + where + T: bytemuck::Pod, + { + let value_bytes = bytemuck::bytes_of(value); + let type_size = std::mem::size_of::(); + let max_run_length = self.run_length_width.max_run_length(); + let num_full_chunks = (length / max_run_length) as usize; + let remainder = length % max_run_length; + + let total_chunks = num_full_chunks + if remainder > 0 { 1 } else { 0 }; + all_values.reserve(total_chunks * type_size); + all_lengths.reserve(total_chunks * self.run_length_width.bytes_per_value()); + + for _ in 0..num_full_chunks { + all_values.extend_from_slice(value_bytes); + self.run_length_width + .write_length(max_run_length, all_lengths); + } + + if remainder > 0 { + all_values.extend_from_slice(value_bytes); + self.run_length_width.write_length(remainder, all_lengths); + } + + total_chunks * (type_size + self.run_length_width.bytes_per_value()) + } + + fn flat_child_candidate( + buffers: &[LanceBuffer], + chunks: &[MiniBlockChunk], + buffer_index: usize, + bits_per_value: u64, + ) -> RleChildCandidate { + RleChildCandidate { + encoding: ProtobufUtils21::flat(bits_per_value, None), + data: buffers[buffer_index].clone(), + chunk_sizes: chunks + .iter() + .map(|chunk| chunk.buffer_sizes[buffer_index]) + .collect(), + size: buffers[buffer_index].len(), + requires_num_values: false, + } + } + + fn general_child_candidate( + buffers: &[LanceBuffer], + chunks: &[MiniBlockChunk], + buffer_index: usize, + bits_per_value: u64, + compression: CompressionConfig, + ) -> Result> { + if buffers.is_empty() || buffers[buffer_index].is_empty() { + return Ok(None); + }; + + let compressor = GeneralBufferCompressor::get_compressor(compression)?; + let original = &buffers[buffer_index]; + let mut compressed = Vec::new(); + let mut offset = 0usize; + let mut total_original_size = 0usize; + let mut compressed_sizes = Vec::with_capacity(chunks.len()); + + for chunk in chunks.iter() { + let chunk_size = chunk.buffer_sizes[buffer_index] as usize; + let end = offset.checked_add(chunk_size).ok_or_else(|| { + Error::invalid_input_source("RLE child buffer offset overflow".into()) + })?; + if end > original.len() { + return Err(Error::invalid_input_source( + format!( + "RLE child buffer {} chunk size exceeds buffer length: end {}, len {}", + buffer_index, + end, + original.len() + ) + .into(), + )); + } + + let start = compressed.len(); + compressor.compress(&original.as_ref()[offset..end], &mut compressed)?; + let compressed_size = compressed.len() - start; + let compressed_size = u32::try_from(compressed_size).map_err(|_| { + Error::invalid_input_source( + format!( + "RLE child buffer {} compressed chunk is too large: {} bytes", + buffer_index, compressed_size + ) + .into(), + ) + })?; + compressed_sizes.push(compressed_size); + total_original_size += chunk_size; + offset = end; + } + + if compressed.len() >= total_original_size { + return Ok(None); + } + + let encoding = + ProtobufUtils21::wrapped(compression, ProtobufUtils21::flat(bits_per_value, None))?; + Ok(Some( + RleChildCandidate { + encoding, + data: LanceBuffer::from(compressed), + chunk_sizes: compressed_sizes, + size: 0, + requires_num_values: false, + } + .with_size_from_data(), + )) + } + + #[cfg(feature = "bitpacking")] + fn bitpacked_child_candidate( + buffers: &[LanceBuffer], + chunks: &[MiniBlockChunk], + buffer_index: usize, + bits_per_value: u64, + ) -> Result> { + let original = &buffers[buffer_index]; + if original.is_empty() { + return Ok(None); + } + let packed_bits = Self::required_bits(original, bits_per_value)?; + if packed_bits >= bits_per_value { + return Ok(None); + } + + let compressor = crate::encodings::physical::bitpacking::OutOfLineBitpacking::new( + packed_bits, + bits_per_value, + ); + let mut packed = Vec::new(); + let mut offset = 0usize; + let mut packed_sizes = Vec::with_capacity(chunks.len()); + let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| { + Error::invalid_input_source( + format!("RLE child bit width is too large: {bits_per_value}").into(), + ) + })?; + + for chunk in chunks { + let chunk_size = chunk.buffer_sizes[buffer_index] as usize; + let end = offset.checked_add(chunk_size).ok_or_else(|| { + Error::invalid_input_source("RLE child buffer offset overflow".into()) + })?; + if end > original.len() { + return Err(Error::invalid_input_source( + format!( + "RLE child buffer {} chunk size exceeds buffer length: end {}, len {}", + buffer_index, + end, + original.len() + ) + .into(), + )); + } + if bytes_per_value == 0 || !chunk_size.is_multiple_of(bytes_per_value) { + return Err(Error::invalid_input_source( + format!( + "RLE child buffer {} chunk has invalid size {} for {} bits per value", + buffer_index, chunk_size, bits_per_value + ) + .into(), + )); + } + + let child_values = (chunk_size / bytes_per_value) as u64; + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value, + data: original.slice_with_length(offset, chunk_size), + num_values: child_values, + block_info: BlockInfo::default(), + }); + let (chunk_packed, _) = BlockCompressor::compress(&compressor, block)?; + let chunk_packed = chunk_packed.ok_or_else(|| { + Error::internal("RLE bitpacking child returned no payload".to_string()) + })?; + let packed_size = u32::try_from(chunk_packed.len()).map_err(|_| { + Error::invalid_input_source( + format!( + "RLE child buffer {} bitpacked chunk is too large: {} bytes", + buffer_index, + chunk_packed.len() + ) + .into(), + ) + })?; + packed_sizes.push(packed_size); + packed.extend_from_slice(chunk_packed.as_ref()); + offset = end; + } + + if packed.len() >= original.len() { + return Ok(None); + } + + Ok(Some( + RleChildCandidate { + encoding: ProtobufUtils21::out_of_line_bitpacking( + bits_per_value, + ProtobufUtils21::flat(packed_bits, None), + ), + data: LanceBuffer::from(packed), + chunk_sizes: packed_sizes, + size: 0, + requires_num_values: true, + } + .with_size_from_data(), + )) + } + + #[cfg(feature = "bitpacking")] + fn required_bits(buffer: &LanceBuffer, bits_per_value: u64) -> Result { + let max_value = match bits_per_value { + 8 => buffer.as_ref().iter().map(|value| *value as u64).max(), + 16 => buffer + .as_ref() + .chunks_exact(2) + .map(|value| u16::from_le_bytes(value.try_into().unwrap()) as u64) + .max(), + 32 => buffer + .as_ref() + .chunks_exact(4) + .map(|value| u32::from_le_bytes(value.try_into().unwrap()) as u64) + .max(), + 64 => buffer + .as_ref() + .chunks_exact(8) + .map(|value| u64::from_le_bytes(value.try_into().unwrap())) + .max(), + _ => { + return Err(Error::invalid_input_source( + format!( + "RLE child bitpacking only supports 8, 16, 32, or 64-bit values, got {bits_per_value}" + ) + .into(), + )); + } + } + .unwrap_or(0); + Ok((u64::BITS - max_value.leading_zeros()).max(1) as u64) + } + + fn child_candidates( + buffers: &[LanceBuffer], + chunks: &[MiniBlockChunk], + buffer_index: usize, + bits_per_value: u64, + compression: Option, + use_child_bitpacking: bool, + ) -> Result> { + #[cfg(not(feature = "bitpacking"))] + let _ = use_child_bitpacking; + let mut candidates = vec![Self::flat_child_candidate( + buffers, + chunks, + buffer_index, + bits_per_value, + )]; + if let Some(compression) = compression + && let Some(candidate) = Self::general_child_candidate( + buffers, + chunks, + buffer_index, + bits_per_value, + compression, + )? + { + candidates.push(candidate); + } + #[cfg(feature = "bitpacking")] + { + if use_child_bitpacking + && let Some(candidate) = + Self::bitpacked_child_candidate(buffers, chunks, buffer_index, bits_per_value)? + { + candidates.push(candidate); + } + } + Ok(candidates) + } + + fn select_child_candidates( + values: Vec, + run_lengths: Vec, + ) -> (RleChildCandidate, RleChildCandidate) { + let mut best: Option<(usize, usize, usize)> = None; + for (value_idx, value) in values.iter().enumerate() { + for (length_idx, length) in run_lengths.iter().enumerate() { + if value.requires_num_values && length.requires_num_values { + continue; + } + let size = value.size + length.size; + if best.is_none_or(|(_, _, best_size)| size < best_size) { + best = Some((value_idx, length_idx, size)); + } + } + } + let (value_idx, length_idx, _) = + best.expect("flat RLE child candidates should always be selectable"); + (values[value_idx].clone(), run_lengths[length_idx].clone()) + } + + pub(crate) fn selected_payload_size(&self, data: &FixedWidthDataBlock) -> Result { + let (all_buffers, chunks) = + self.encode_data(&data.data, data.num_values, data.bits_per_value)?; + if all_buffers.is_empty() { + return Ok(0); + } + + let values_candidates = Self::child_candidates( + &all_buffers, + &chunks, + 0, + data.bits_per_value, + self.values_compression, + self.use_child_bitpacking, + )?; + let run_lengths_candidates = Self::child_candidates( + &all_buffers, + &chunks, + 1, + self.run_length_width.bits_per_value(), + self.run_lengths_compression, + self.use_child_bitpacking, + )?; + let (values, run_lengths) = + Self::select_child_candidates(values_candidates, run_lengths_candidates); + Ok((values.size as u128).saturating_add(run_lengths.size as u128)) + } +} + +impl RleChildCandidate { + fn with_size_from_data(mut self) -> Self { + self.size = self.data.len(); + self + } +} + +impl MiniBlockCompressor for RleEncoder { + fn compress( + &self, + _context: MiniBlockCompressionContext, + data: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + match data { + DataBlock::FixedWidth(fixed_width) => { + let num_values = fixed_width.num_values; + let bits_per_value = fixed_width.bits_per_value; + + let (all_buffers, chunks) = + self.encode_data(&fixed_width.data, num_values, bits_per_value)?; + if all_buffers.is_empty() { + let compressed = MiniBlockCompressed { + data: all_buffers, + chunks, + num_values, + }; + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::flat(bits_per_value, None), + ProtobufUtils21::flat(self.run_length_width.bits_per_value(), None), + ); + return Ok((compressed, encoding)); + } + + let values_candidates = Self::child_candidates( + &all_buffers, + &chunks, + 0, + bits_per_value, + self.values_compression, + self.use_child_bitpacking, + )?; + let run_lengths_candidates = Self::child_candidates( + &all_buffers, + &chunks, + 1, + self.run_length_width.bits_per_value(), + self.run_lengths_compression, + self.use_child_bitpacking, + )?; + let (values, run_lengths) = + Self::select_child_candidates(values_candidates, run_lengths_candidates); + let chunks = chunks + .into_iter() + .enumerate() + .map(|(idx, chunk)| MiniBlockChunk { + buffer_sizes: vec![values.chunk_sizes[idx], run_lengths.chunk_sizes[idx]], + log_num_values: chunk.log_num_values, + }) + .collect(); + + let compressed = MiniBlockCompressed { + data: vec![values.data, run_lengths.data], + chunks, + num_values, + }; + + let encoding = ProtobufUtils21::rle(values.encoding, run_lengths.encoding); + + Ok((compressed, encoding)) + } + _ => Err(Error::invalid_input_source( + "RLE encoding only supports FixedWidth data blocks".into(), + )), + } + } +} + +impl BlockCompressor for RleEncoder { + // Block format: [8-byte header: values buffer size][values buffer][run_lengths buffer] + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + match data { + DataBlock::FixedWidth(fixed_width) => { + let num_values = fixed_width.num_values; + let bits_per_value = fixed_width.bits_per_value; + + let all_buffers = + self.encode_block_data(&fixed_width.data, num_values, bits_per_value)?; + + let values_size = all_buffers[0].len() as u64; + + let mut combined = Vec::new(); + combined.extend_from_slice(&values_size.to_le_bytes()); + combined.extend_from_slice(&all_buffers[0]); + combined.extend_from_slice(&all_buffers[1]); + Ok(( + Some(LanceBuffer::from(combined)), + ProtobufUtils21::rle( + ProtobufUtils21::flat(bits_per_value, None), + ProtobufUtils21::flat(self.run_length_width.bits_per_value(), None), + ), + )) + } + _ => Err(Error::invalid_input_source( + "RLE encoding only supports FixedWidth data blocks".into(), + )), + } + } +} + +/// RLE decompressor for miniblock format +#[derive(Debug)] +pub struct RleDecompressor { + bits_per_value: u64, + run_length_width: RunLengthWidth, + values: RleChildDecompressor, + run_lengths: RleChildDecompressor, +} + +#[derive(Debug)] +pub(crate) struct RleChildDecompressor { + bits_per_value: u64, + inner: RleChildDecompressorInner, +} + +#[derive(Debug)] +enum RleChildDecompressorInner { + Flat, + Block { + decompressor: Box, + requires_num_values: bool, + }, +} + +impl RleChildDecompressor { + pub(crate) fn flat(bits_per_value: u64) -> Self { + Self { + bits_per_value, + inner: RleChildDecompressorInner::Flat, + } + } + + pub(crate) fn block( + bits_per_value: u64, + decompressor: Box, + requires_num_values: bool, + ) -> Self { + Self { + bits_per_value, + inner: RleChildDecompressorInner::Block { + decompressor, + requires_num_values, + }, + } + } + + pub(crate) fn bits_per_value(&self) -> u64 { + self.bits_per_value + } + + pub(crate) fn requires_num_values(&self) -> bool { + match &self.inner { + RleChildDecompressorInner::Flat => false, + RleChildDecompressorInner::Block { + requires_num_values, + .. + } => *requires_num_values, + } + } + + pub(crate) fn is_identity(&self) -> bool { + matches!(self.inner, RleChildDecompressorInner::Flat) + } + + fn decode( + &self, + data: LanceBuffer, + num_values: Option, + label: &str, + ) -> Result { + match &self.inner { + RleChildDecompressorInner::Flat => Ok(data), + RleChildDecompressorInner::Block { + decompressor, + requires_num_values, + } => { + let num_values = if *requires_num_values { + num_values.ok_or_else(|| { + Error::invalid_input_source( + format!("RLE {label} child compression requires the run count").into(), + ) + })? + } else { + num_values.unwrap_or(0) + }; + let decoded = decompressor.decompress(Some(data), num_values)?; + self.extract_fixed_width(decoded, num_values, label) + } + } + } + + fn extract_fixed_width( + &self, + data: DataBlock, + expected_num_values: u64, + label: &str, + ) -> Result { + match data { + DataBlock::FixedWidth(block) => { + if block.bits_per_value != self.bits_per_value { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child decoded {}-bit values, expected {}", + block.bits_per_value, self.bits_per_value + ) + .into(), + )); + } + if expected_num_values != 0 && block.num_values != expected_num_values { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child decoded {} values, expected {}", + block.num_values, expected_num_values + ) + .into(), + )); + } + Ok(block.data) + } + _ => Err(Error::invalid_input_source( + format!("RLE {label} child decoded to a non fixed-width block").into(), + )), + } + } +} + +impl RleDecompressor { + pub fn new(bits_per_value: u64) -> Self { + Self { + bits_per_value, + run_length_width: RunLengthWidth::U8, + values: RleChildDecompressor::flat(bits_per_value), + run_lengths: RleChildDecompressor::flat(RunLengthWidth::U8.bits_per_value()), + } + } + + pub(crate) fn with_run_length_width( + bits_per_value: u64, + run_length_width: RunLengthWidth, + ) -> Self { + Self { + bits_per_value, + run_length_width, + values: RleChildDecompressor::flat(bits_per_value), + run_lengths: RleChildDecompressor::flat(run_length_width.bits_per_value()), + } + } + + pub(crate) fn with_child_decompressors( + bits_per_value: u64, + run_length_width: RunLengthWidth, + values: RleChildDecompressor, + run_lengths: RleChildDecompressor, + ) -> Self { + Self { + bits_per_value, + run_length_width, + values, + run_lengths, + } + } + + fn decode_data( + &self, + data: Vec, + num_values: u64, + clamp_overflow: bool, + ) -> Result { + if num_values == 0 { + return Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.bits_per_value, + data: LanceBuffer::from(vec![]), + num_values: 0, + block_info: BlockInfo::default(), + })); + } + + if data.len() != 2 { + return Err(Error::invalid_input_source( + format!( + "RLE decompressor expects exactly 2 buffers, got {}", + data.len() + ) + .into(), + )); + } + + let mut data_iter = data.into_iter(); + let values_buffer = data_iter.next().unwrap(); + let lengths_buffer = data_iter.next().unwrap(); + let (values_buffer, lengths_buffer) = + self.decode_child_buffers(values_buffer, lengths_buffer)?; + + self.decode_child_data(&values_buffer, &lengths_buffer, num_values, clamp_overflow) + } + + fn decode_child_data( + &self, + values_buffer: &LanceBuffer, + lengths_buffer: &LanceBuffer, + num_values: u64, + clamp_overflow: bool, + ) -> Result { + let decoded_data = match self.bits_per_value { + 8 => self.decode_generic::( + values_buffer, + lengths_buffer, + num_values, + clamp_overflow, + )?, + 16 => self.decode_generic::( + values_buffer, + lengths_buffer, + num_values, + clamp_overflow, + )?, + 32 => self.decode_generic::( + values_buffer, + lengths_buffer, + num_values, + clamp_overflow, + )?, + 64 => self.decode_generic::( + values_buffer, + lengths_buffer, + num_values, + clamp_overflow, + )?, + _ => { + return Err(Error::invalid_input_source( + format!( + "RLE decoding bits_per_value must be 8, 16, 32, or 64, got {}", + self.bits_per_value + ) + .into(), + )); + } + }; + + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.bits_per_value, + data: decoded_data, + num_values, + block_info: BlockInfo::default(), + })) + } + + fn sum_run_lengths( + &self, + values_buffer: &LanceBuffer, + lengths_buffer: &LanceBuffer, + ) -> Result { + let (value_size, value_type) = match self.bits_per_value { + 8 => (1, "u8"), + 16 => (2, "u16"), + 32 => (4, "u32"), + 64 => (8, "u64"), + _ => { + return Err(Error::invalid_input_source( + format!( + "RLE decoding bits_per_value must be 8, 16, 32, or 64, got {}", + self.bits_per_value + ) + .into(), + )); + } + }; + let length_size = + self.validate_buffer_sizes(values_buffer, lengths_buffer, value_size, value_type)?; + + lengths_buffer + .chunks_exact(length_size) + .try_fold(0_u64, |num_values, length_bytes| { + let length = self.run_length_width.read_length(length_bytes); + if length == 0 { + return Err(Error::invalid_input_source( + "RLE decoding encountered a zero run length".into(), + )); + } + num_values.checked_add(length).ok_or_else(|| { + Error::invalid_input_source("RLE run length sum overflowed u64".into()) + }) + }) + } + + fn validate_buffer_sizes( + &self, + values_buffer: &LanceBuffer, + lengths_buffer: &LanceBuffer, + value_size: usize, + value_type: &str, + ) -> Result { + let length_size = self.run_length_width.bytes_per_value(); + if !values_buffer.len().is_multiple_of(value_size) + || !lengths_buffer.len().is_multiple_of(length_size) + { + return Err(Error::invalid_input_source(format!( + "Invalid buffer sizes for RLE {value_type} decoding: values {} bytes (not divisible by {}), lengths {} bytes (not divisible by {})", + values_buffer.len(), + value_size, + lengths_buffer.len(), + length_size + ) + .into())); + } + + let num_runs = values_buffer.len() / value_size; + let num_length_entries = lengths_buffer.len() / length_size; + if num_runs != num_length_entries { + return Err(Error::invalid_input_source( + format!( + "Inconsistent RLE buffers: {} runs but {} length entries", + num_runs, num_length_entries + ) + .into(), + )); + } + Ok(length_size) + } + + fn decode_child_buffers( + &self, + values_buffer: LanceBuffer, + lengths_buffer: LanceBuffer, + ) -> Result<(LanceBuffer, LanceBuffer)> { + let values_requires_num_runs = self.values.requires_num_values(); + let lengths_requires_num_runs = self.run_lengths.requires_num_values(); + if values_requires_num_runs && lengths_requires_num_runs { + return Err(Error::invalid_input_source( + "RLE values and run lengths child compression both require the run count".into(), + )); + } + + if values_requires_num_runs { + let lengths_buffer = self + .run_lengths + .decode(lengths_buffer, None, "run lengths")?; + let num_runs = Self::num_child_values( + &lengths_buffer, + self.run_lengths.bits_per_value(), + "run lengths", + )?; + let values_buffer = self + .values + .decode(values_buffer, Some(num_runs), "values")?; + Ok((values_buffer, lengths_buffer)) + } else if lengths_requires_num_runs { + let values_buffer = self.values.decode(values_buffer, None, "values")?; + let num_runs = + Self::num_child_values(&values_buffer, self.values.bits_per_value(), "values")?; + let lengths_buffer = + self.run_lengths + .decode(lengths_buffer, Some(num_runs), "run lengths")?; + Ok((values_buffer, lengths_buffer)) + } else { + let values_buffer = self.values.decode(values_buffer, None, "values")?; + let lengths_buffer = self + .run_lengths + .decode(lengths_buffer, None, "run lengths")?; + Ok((values_buffer, lengths_buffer)) + } + } + + fn num_child_values(buffer: &LanceBuffer, bits_per_value: u64, label: &str) -> Result { + let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| { + Error::invalid_input_source( + format!("RLE {label} child bit width is too large: {bits_per_value}").into(), + ) + })?; + if bytes_per_value == 0 || !buffer.len().is_multiple_of(bytes_per_value) { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child decoded to {} bytes, not divisible by {}", + buffer.len(), + bytes_per_value + ) + .into(), + )); + } + Ok((buffer.len() / bytes_per_value) as u64) + } + + fn decode_generic( + &self, + values_buffer: &LanceBuffer, + lengths_buffer: &LanceBuffer, + num_values: u64, + clamp_overflow: bool, + ) -> Result + where + T: bytemuck::Pod + Copy + std::fmt::Debug + ArrowNativeType, + { + let type_size = std::mem::size_of::(); + + if values_buffer.is_empty() || lengths_buffer.is_empty() { + if num_values == 0 { + return Ok(LanceBuffer::empty()); + } else { + return Err(Error::invalid_input_source( + format!("Empty buffers but expected {} values", num_values).into(), + )); + } + } + + let length_size = self.validate_buffer_sizes( + values_buffer, + lengths_buffer, + type_size, + std::any::type_name::(), + )?; + + let values_ref = values_buffer.borrow_to_typed_slice::(); + let values: &[T] = values_ref.as_ref(); + let lengths = lengths_buffer.as_ref(); + + let expected_value_count = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + // Legacy miniblock encoders rolled back to a power-of-2 checkpoint after a run + // had already crossed it, so a chunk's run lengths can sum past its declared + // value count (the excess values are re-encoded at the start of the next chunk). + // The pre-run-length-width decoder truncated the excess, so miniblock decoding + // clamps rather than rejects to keep those files readable. Block payloads never + // legitimately overflow, so they decode strictly. + let mut decoded: Vec = Vec::new(); + decoded + .try_reserve_exact(expected_value_count) + .map_err(|_| { + Error::invalid_input_source( + format!("RLE decoding cannot allocate {expected_value_count} values").into(), + ) + })?; + for (value, length_bytes) in values.iter().zip(lengths.chunks_exact(length_size)) { + let length = self.run_length_width.read_length(length_bytes); + if length == 0 { + return Err(Error::invalid_input_source( + "RLE decoding encountered a zero run length".into(), + )); + } + let length = usize::try_from(length).map_err(|_| { + Error::invalid_input_source( + format!("RLE run length does not fit in usize: {length}").into(), + ) + })?; + let remaining = expected_value_count - decoded.len(); + if length > remaining { + if !clamp_overflow { + return Err(Error::invalid_input_source( + format!( + "RLE decoding overflowed expected value count: produced at least {}, expected {}", + decoded.len() + length, + expected_value_count + ) + .into(), + )); + } + decoded.resize(expected_value_count, *value); + break; + } + decoded.resize(decoded.len() + length, *value); + } + + if decoded.len() != expected_value_count { + return Err(Error::invalid_input_source( + format!( + "RLE decoding produced {} values, expected {}", + decoded.len(), + expected_value_count + ) + .into(), + )); + } + + trace!( + "RLE decoded {} {} values", + num_values, + std::any::type_name::() + ); + Ok(LanceBuffer::reinterpret_vec(decoded)) + } +} + +impl MiniBlockDecompressor for RleDecompressor { + fn decompress(&self, data: Vec, num_values: u64) -> Result { + self.decode_data(data, num_values, true) + } + + fn decoded_size_bytes(&self, num_values: u64) -> Option { + num_values + .checked_mul(self.bits_per_value) + .map(|bits| bits.div_ceil(8)) + } +} + +impl BlockDecompressor for RleDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "RLE")?; + let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?; + self.decode_data(vec![values_buffer, lengths_buffer], num_values, false) + } + + fn infer_num_values(&self, data: &LanceBuffer) -> Result> { + // Pylance 6.0.1 used this exact RLE signature for structural levels. Newer RLE + // variants are not part of that compatibility case and may contain much wider, + // untrusted run lengths. + if self.bits_per_value != 16 + || self.run_length_width != RunLengthWidth::U8 + || !self.values.is_identity() + || !self.run_lengths.is_identity() + { + return Ok(None); + } + let (values_buffer, lengths_buffer) = parse_rle_block_frame(data)?; + self.sum_run_lengths(&values_buffer, &lengths_buffer) + .map(Some) + } +} + +/// Split an RLE block-format buffer into its `(values, lengths)` sub-buffers. +/// Frame: `[values_size: u64-le][values bytes][run-length bytes]`. +fn parse_rle_block_frame(data: &LanceBuffer) -> Result<(LanceBuffer, LanceBuffer)> { + // fetch the values_size + if data.len() < 8 { + return Err(Error::invalid_input_source( + format!("Insufficient data size: {}", data.len()).into(), + )); + } + + let values_size_bytes: [u8; 8] = data[..8].try_into().expect("slice length already checked"); + let values_size: usize = u64::from_le_bytes(values_size_bytes) + .try_into() + .map_err(|_| { + Error::invalid_input_source( + format!( + "Invalid values buffer size: {}", + u64::from_le_bytes(values_size_bytes) + ) + .into(), + ) + })?; + + // parse values + let values_start: usize = 8; + let lengths_start = values_start + .checked_add(values_size) + .ok_or_else(|| Error::invalid_input_source("Invalid RLE values buffer size".into()))?; + + if data.len() < lengths_start { + return Err(Error::invalid_input_source( + format!("Insufficient data size: {}", data.len()).into(), + )); + } + + let values_buffer = data.slice_with_length(values_start, values_size); + let lengths_buffer = data.slice_with_length(lengths_start, data.len() - lengths_start); + Ok((values_buffer, lengths_buffer)) +} + +#[derive(Clone, Debug)] +enum RleRunLengths { + U8(ScalarBuffer), + U16(ScalarBuffer), + U32(ScalarBuffer), +} + +impl RleRunLengths { + fn try_new(buffer: LanceBuffer, width: RunLengthWidth) -> Result { + let width_bytes = width.bytes_per_value(); + if !buffer.len().is_multiple_of(width_bytes) { + return Err(Error::invalid_input_source( + format!( + "Invalid RLE run lengths buffer: {} bytes (not divisible by {})", + buffer.len(), + width_bytes + ) + .into(), + )); + } + Ok(match width { + RunLengthWidth::U8 => Self::U8(buffer.borrow_to_typed_slice()), + RunLengthWidth::U16 => Self::U16(buffer.borrow_to_typed_slice()), + RunLengthWidth::U32 => Self::U32(buffer.borrow_to_typed_slice()), + }) + } + + fn len(&self) -> usize { + match self { + Self::U8(lengths) => lengths.len(), + Self::U16(lengths) => lengths.len(), + Self::U32(lengths) => lengths.len(), + } + } + + fn get(&self, index: usize) -> usize { + match self { + Self::U8(lengths) => lengths[index] as usize, + Self::U16(lengths) => lengths[index] as usize, + Self::U32(lengths) => lengths[index] as usize, + } + } + + fn owned_size(&self) -> usize { + match self { + Self::U8(lengths) => std::mem::size_of_val(lengths.as_ref()), + Self::U16(lengths) => std::mem::size_of_val(lengths.as_ref()), + Self::U32(lengths) => std::mem::size_of_val(lengths.as_ref()), + } + } + + fn into_owned(self) -> Self { + match self { + Self::U8(lengths) => Self::U8(ScalarBuffer::from(lengths.as_ref().to_vec())), + Self::U16(lengths) => Self::U16(ScalarBuffer::from(lengths.as_ref().to_vec())), + Self::U32(lengths) => Self::U32(ScalarBuffer::from(lengths.as_ref().to_vec())), + } + } + + fn deep_size(&self) -> usize { + match self { + Self::U8(lengths) => lengths.inner().capacity(), + Self::U16(lengths) => lengths.inner().capacity(), + Self::U32(lengths) => lengths.inner().capacity(), + } + } +} + +/// Validated physical RLE runs for `u16` values. +/// +/// The values and original-width lengths remain unexpanded. The constructor +/// verifies that every length is non-zero and that the runs cover exactly +/// `num_values` logical values. +#[derive(Clone, Debug)] +pub(crate) struct RleRuns { + values: ScalarBuffer, + lengths: RleRunLengths, + num_values: usize, + coalesced_runs: usize, +} + +impl RleRuns { + fn try_new( + values_buffer: LanceBuffer, + lengths_buffer: LanceBuffer, + run_length_width: RunLengthWidth, + num_values: u64, + ) -> Result { + let num_values = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + let type_size = std::mem::size_of::(); + if !values_buffer.len().is_multiple_of(type_size) { + return Err(Error::invalid_input_source( + format!( + "Invalid RLE u16 values buffer: {} bytes (not divisible by {})", + values_buffer.len(), + type_size + ) + .into(), + )); + } + + let values = values_buffer.borrow_to_typed_slice::(); + let lengths = RleRunLengths::try_new(lengths_buffer, run_length_width)?; + if values.len() != lengths.len() { + return Err(Error::invalid_input_source( + format!( + "Inconsistent RLE buffers: {} runs but {} length entries", + values.len(), + lengths.len() + ) + .into(), + )); + } + if values.is_empty() && num_values != 0 { + return Err(Error::invalid_input_source( + format!("Empty RLE buffers but expected {num_values} values").into(), + )); + } + + let mut decoded_values = 0usize; + let mut coalesced_runs = 0usize; + let mut previous_value = None; + for run in 0..values.len() { + let length = lengths.get(run); + if length == 0 { + return Err(Error::invalid_input_source( + "RLE decoding encountered a zero run length".into(), + )); + } + decoded_values = decoded_values.checked_add(length).ok_or_else(|| { + Error::invalid_input_source("RLE run length sum overflowed usize".into()) + })?; + if decoded_values > num_values { + return Err(Error::invalid_input_source( + format!( + "RLE decoding overflowed expected value count: produced at least {}, expected {}", + decoded_values, num_values + ) + .into(), + )); + } + if previous_value != Some(values[run]) { + coalesced_runs += 1; + previous_value = Some(values[run]); + } + } + if decoded_values != num_values { + return Err(Error::invalid_input_source( + format!( + "RLE decoding produced {} values, expected {}", + decoded_values, num_values + ) + .into(), + )); + } + + Ok(Self { + values, + lengths, + num_values, + coalesced_runs, + }) + } + + pub(crate) fn num_values(&self) -> usize { + self.num_values + } + + pub(crate) fn num_runs(&self) -> usize { + self.values.len() + } + + pub(crate) fn coalesced_runs(&self) -> usize { + self.coalesced_runs + } + + pub(crate) fn owned_size(&self) -> usize { + std::mem::size_of_val(self.values.as_ref()) + self.lengths.owned_size() + } + + pub(crate) fn into_owned(self) -> Self { + Self { + values: ScalarBuffer::from(self.values.as_ref().to_vec()), + lengths: self.lengths.into_owned(), + num_values: self.num_values, + coalesced_runs: self.coalesced_runs, + } + } + + pub(crate) fn deep_size(&self) -> usize { + self.values.inner().capacity() + self.lengths.deep_size() + } + + pub(crate) fn value(&self, run: usize) -> u16 { + self.values[run] + } + + pub(crate) fn length(&self, run: usize) -> usize { + self.lengths.get(run) + } + + pub(crate) fn iter(&self) -> impl ExactSizeIterator + '_ { + (0..self.num_runs()).map(|run| (self.value(run), self.length(run))) + } +} + +impl RleDecompressor { + /// Decode and validate a block frame while preserving its physical runs. + pub(crate) fn decode_u16_runs(&self, data: LanceBuffer, num_values: u64) -> Result { + if self.bits_per_value != 16 { + return Err(Error::invalid_input_source( + format!( + "RLE level values must be 16 bits, got {}", + self.bits_per_value + ) + .into(), + )); + } + let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?; + let (values_buffer, lengths_buffer) = + self.decode_child_buffers(values_buffer, lengths_buffer)?; + RleRuns::try_new( + values_buffer, + lengths_buffer, + self.run_length_width, + num_values, + ) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::compression::{ + DecompressionStrategy, DefaultDecompressionStrategy, create_rle_decompressor, + }; + use crate::data::DataBlock; + use crate::encodings::logical::primitive::miniblock::MAX_MINIBLOCK_VALUES; + use crate::encodings::physical::block::{CompressionConfig, CompressionScheme}; + use crate::{ + buffer::LanceBuffer, + compression::{BlockCompressor, BlockDecompressor}, + }; + use arrow_array::Int32Array; + use rstest::rstest; + + fn compress_miniblock( + compressor: &dyn MiniBlockCompressor, + data: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + compressor.compress(MiniBlockCompressionContext::new(0, true, true), data) + } + + fn expand_u16_runs(runs: &RleRuns) -> Vec { + let mut expanded = Vec::with_capacity(runs.num_values()); + for (value, length) in runs.iter() { + expanded.extend(std::iter::repeat_n(value, length)); + } + expanded + } + + #[test] + fn decode_u16_runs_matches_eager() { + // Near-constant u16 levels (the all-null shape): a few long runs. + let mut levels: Vec = vec![1u16; 1000]; + levels.extend(std::iter::repeat_n(0u16, 500)); + levels.extend(std::iter::repeat_n(2u16, 300)); + let num_values = levels.len() as u64; + + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_slice(Arc::from(levels.clone())), + bits_per_value: 16, + num_values, + block_info: BlockInfo::new(), + }); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .0 + .unwrap(); + + let eager = BlockDecompressor::decompress( + &RleDecompressor::new(16), + Some(frame.clone()), + num_values, + ) + .unwrap(); + let DataBlock::FixedWidth(eager) = eager else { + panic!("expected fixed-width block"); + }; + assert_eq!( + eager.data.borrow_to_typed_slice::().as_ref(), + levels.as_slice() + ); + + // Lazy run form preserves boundaries and expands identically. + let runs = RleDecompressor::new(16) + .decode_u16_runs(frame, num_values) + .unwrap(); + assert_eq!(runs.num_values(), num_values as usize); + // The encoder splits each value into <=255-length runs (4 + 2 + 2 = 8 + // on-disk runs here); the scan identifies 3 coalesced logical runs. + assert_eq!(runs.num_runs(), 8); + assert_eq!(runs.coalesced_runs(), 3); + assert_eq!(expand_u16_runs(&runs), levels); + } + + #[test] + fn decode_u16_runs_empty() { + let mut empty_frame = Vec::new(); + empty_frame.extend_from_slice(&0u64.to_le_bytes()); + let runs = RleDecompressor::new(16) + .decode_u16_runs(LanceBuffer::from(empty_frame), 0) + .unwrap(); + assert_eq!(runs.num_values(), 0); + assert_eq!(runs.num_runs(), 0); + assert_eq!(runs.coalesced_runs(), 0); + + let error = RleDecompressor::new(16) + .decode_u16_runs(LanceBuffer::empty(), 0) + .unwrap_err(); + assert!(error.to_string().contains("Insufficient data size: 0")); + } + + #[test] + fn legacy_block_rle_infers_value_count_without_materializing() { + let num_values = u64::from(u16::MAX) + 8; + let full_runs = num_values / u64::from(u8::MAX); + let remainder = num_values % u64::from(u8::MAX); + let num_runs = full_runs + u64::from(remainder != 0); + let mut frame = Vec::new(); + frame.extend_from_slice(&(num_runs * 2).to_le_bytes()); + frame.extend(std::iter::repeat_n(7_u16, num_runs as usize).flat_map(u16::to_le_bytes)); + frame.extend(std::iter::repeat_n(u8::MAX, full_runs as usize)); + if remainder != 0 { + frame.push(remainder as u8); + } + + let inferred_num_values = RleDecompressor::new(16) + .infer_num_values(&LanceBuffer::from(frame)) + .unwrap(); + assert_eq!(inferred_num_values, Some(num_values)); + } + + #[test] + fn newer_block_rle_does_not_infer_untrusted_run_sum() { + let mut frame = Vec::new(); + frame.extend_from_slice(&2_u64.to_le_bytes()); + frame.extend_from_slice(&7_u16.to_le_bytes()); + frame.extend_from_slice(&u32::MAX.to_le_bytes()); + let frame = LanceBuffer::from(frame); + let decompressor = RleDecompressor::with_run_length_width(16, RunLengthWidth::U32); + + assert_eq!(decompressor.infer_num_values(&frame).unwrap(), None); + let error = BlockDecompressor::decompress(&decompressor, Some(frame), u64::from(u16::MAX)) + .unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("RLE decoding overflowed expected value count") + ); + } + + #[rstest] + #[case::zero(0, 1, "zero run length")] + #[case::underflow(1, 2, "produced 1 values, expected 2")] + #[case::overflow(2, 1, "overflowed expected value count")] + #[case::nonempty_for_empty(1, 0, "overflowed expected value count")] + fn decode_u16_runs_rejects_invalid_coverage( + #[case] run_length: u8, + #[case] num_values: u64, + #[case] expected_message: &str, + ) { + let mut frame = Vec::new(); + frame.extend_from_slice(&2u64.to_le_bytes()); + frame.extend_from_slice(&7u16.to_le_bytes()); + frame.push(run_length); + + let error = RleDecompressor::new(16) + .decode_u16_runs(LanceBuffer::from(frame), num_values) + .unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!(error.to_string().contains(expected_message)); + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn decode_u16_runs_supports_compressed_values_child() { + let levels: Vec = (0..1024) + .flat_map(|run| std::iter::repeat_n((run % 8) as u16, 4)) + .collect(); + let num_values = levels.len() as u64; + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_slice(Arc::from(levels.clone())), + bits_per_value: 16, + num_values, + block_info: BlockInfo::new(), + }); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .0 + .unwrap(); + let (values, lengths) = parse_rle_block_frame(&frame).unwrap(); + + let compression = test_general_compression(); + let compressor = GeneralBufferCompressor::get_compressor(compression).unwrap(); + let mut compressed_values = Vec::new(); + compressor + .compress(values.as_ref(), &mut compressed_values) + .unwrap(); + let mut compressed_frame = Vec::new(); + compressed_frame.extend_from_slice(&(compressed_values.len() as u64).to_le_bytes()); + compressed_frame.extend_from_slice(&compressed_values); + compressed_frame.extend_from_slice(lengths.as_ref()); + + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::wrapped(compression, ProtobufUtils21::flat(16, None)).unwrap(), + ProtobufUtils21::flat(8, None), + ); + let decompressor = create_rle_decompressor( + expect_rle(&encoding), + &DefaultDecompressionStrategy::default(), + ) + .unwrap(); + let runs = decompressor + .decode_u16_runs(LanceBuffer::from(compressed_frame), num_values) + .unwrap(); + assert_eq!(expand_u16_runs(&runs), levels); + } + + #[test] + fn decode_u16_runs_counts_coalesced_runs() { + // A logically constant page is emitted as ceil(N / 255) equal-valued + // runs (the encoder caps run lengths at 255); the validated view records + // that they can collapse to a single logical run. + let num_values = 5000u64; + let constant: Vec = vec![7u16; num_values as usize]; + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_slice(Arc::from(constant)), + bits_per_value: 16, + num_values, + block_info: BlockInfo::new(), + }); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .0 + .unwrap(); + let runs = RleDecompressor::new(16) + .decode_u16_runs(frame, num_values) + .unwrap(); + assert_eq!( + runs.num_runs(), + num_values.div_ceil(u8::MAX as u64) as usize + ); + assert_eq!(runs.coalesced_runs(), 1); + assert_eq!(expand_u16_runs(&runs), vec![7u16; num_values as usize]); + + // Distinct adjacent values must not be merged: alternating single-value + // runs stay separate and still expand to the original. + let alternating: Vec = (0..200u16).map(|i| i % 2).collect(); + let n = alternating.len() as u64; + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_slice(Arc::from(alternating.clone())), + bits_per_value: 16, + num_values: n, + block_info: BlockInfo::new(), + }); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .0 + .unwrap(); + let runs = RleDecompressor::new(16).decode_u16_runs(frame, n).unwrap(); + assert_eq!( + runs.coalesced_runs() as u64, + n, + "no two adjacent values are equal" + ); + assert_eq!(expand_u16_runs(&runs), alternating); + } + + // ========== Core Functionality Tests ========== + + #[test] + fn test_basic_miniblock_rle_encoding() { + let encoder = RleEncoder::new(); + + // Test basic RLE pattern: [1, 1, 1, 2, 2, 3, 3, 3, 3] + let array = Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 3, 3]); + let data_block = DataBlock::from_array(array); + + let (compressed, _) = compress_miniblock(&encoder, data_block).unwrap(); + + assert_eq!(compressed.num_values, 9); + assert_eq!(compressed.chunks.len(), 1); + + // Verify compression happened (3 runs instead of 9 values) + let values_buffer = &compressed.data[0]; + let lengths_buffer = &compressed.data[1]; + assert_eq!(values_buffer.len(), 12); // 3 i32 values + assert_eq!(lengths_buffer.len(), 3); // 3 u8 lengths + } + + #[test] + fn test_long_run_splitting() { + let encoder = RleEncoder::new(); + + // Create a run longer than 255 to test splitting + let mut data = vec![42i32; 1000]; // Will be split into 255+255+255+235 + data.extend(&[100i32; 300]); // Will be split into 255+45 + + let array = Int32Array::from(data); + let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); + + // Should have 6 runs total (4 for first value, 2 for second) + let lengths_buffer = &compressed.data[1]; + assert_eq!(lengths_buffer.len(), 6); + } + + #[test] + fn test_rle_v2_u16_miniblock_encoding() { + let encoder = RleEncoder::with_run_length_width(RunLengthWidth::U16); + + let data = vec![42i32; 1000]; + let array = Int32Array::from(data); + let (compressed, encoding) = + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); + + assert_eq!(compressed.data[0].len(), 4); + assert_eq!(compressed.data[1].len(), 2); + assert_eq!(compressed.data[1].as_ref(), &1000u16.to_le_bytes()); + + let rle = match encoding.compression.as_ref().unwrap() { + crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle, + other => panic!("expected RLE encoding, got {other:?}"), + }; + let run_lengths = rle.run_lengths.as_ref().unwrap(); + let flat = match run_lengths.compression.as_ref().unwrap() { + crate::format::pb21::compressive_encoding::Compression::Flat(flat) => flat, + other => panic!("expected flat run lengths, got {other:?}"), + }; + assert_eq!(flat.bits_per_value, 16); + + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let decompressed = MiniBlockDecompressor::decompress( + &decompressor, + compressed.data, + compressed.num_values, + ) + .unwrap(); + match decompressed { + DataBlock::FixedWidth(block) => { + let values = block.data.borrow_to_typed_slice::(); + assert_eq!(values.as_ref(), vec![42i32; 1000]); + } + _ => panic!("Expected FixedWidth block"), + } + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_rle_miniblock_compressed_values_child() { + let compression = test_general_compression(); + let encoder = + RleEncoder::with_child_encoding(RunLengthWidth::U8, Some(compression), None, false); + let array = Int32Array::from(repeating_runs(1024, 4)); + let (compressed, encoding) = + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); + + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::General(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::Flat(_) + )); + + let decompressor = DefaultDecompressionStrategy::default() + .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default()) + .unwrap(); + let decoded = + MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4) + .unwrap(); + assert_decoded_i32_eq(decoded, &repeating_runs(1024, 4)); + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_rle_miniblock_compressed_run_lengths_child() { + let compression = test_general_compression(); + let encoder = + RleEncoder::with_child_encoding(RunLengthWidth::U8, None, Some(compression), false); + let expected = repeating_runs(1024, 4); + let (compressed, encoding) = compress_miniblock( + &encoder, + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::Flat(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::General(_) + )); + + let decompressor = DefaultDecompressionStrategy::default() + .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default()) + .unwrap(); + let decoded = + MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4) + .unwrap(); + assert_decoded_i32_eq(decoded, &expected); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_bitpacked_run_lengths_child() { + use crate::encodings::physical::bitpacking::OutOfLineBitpacking; + + let expected = repeating_runs(1024, 4); + let (compressed, _) = compress_miniblock( + &RleEncoder::new(), + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + let run_lengths = compressed.data[1].clone(); + let num_runs = run_lengths.len() as u64; + let run_lengths_block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 8, + data: run_lengths, + num_values: num_runs, + block_info: BlockInfo::default(), + }); + let bitpacked_run_lengths = + BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block) + .unwrap() + .0 + .unwrap(); + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)), + ); + + let decompressor = DefaultDecompressionStrategy::default() + .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default()) + .unwrap(); + let decoded = MiniBlockDecompressor::decompress( + decompressor.as_ref(), + vec![compressed.data[0].clone(), bitpacked_run_lengths], + expected.len() as u64, + ) + .unwrap(); + assert_decoded_i32_eq(decoded, &expected); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_rejects_two_count_dependent_child_encodings() { + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::out_of_line_bitpacking(32, ProtobufUtils21::flat(3, None)), + ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)), + ); + + let err = DefaultDecompressionStrategy::default() + .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default()) + .unwrap_err(); + assert!( + err.to_string() + .contains("cannot both require the run count") + ); + } + + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_general_compression() -> CompressionConfig { + if cfg!(feature = "zstd") { + CompressionConfig::new(CompressionScheme::Zstd, Some(3)) + } else { + CompressionConfig::new(CompressionScheme::Lz4, None) + } + } + + fn repeating_runs(num_runs: usize, run_length: usize) -> Vec { + let mut values = Vec::with_capacity(num_runs * run_length); + for run in 0..num_runs { + values.extend(std::iter::repeat_n((run % 8) as i32, run_length)); + } + values + } + + fn expect_rle(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle { + match encoding.compression.as_ref().unwrap() { + crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle, + other => panic!("expected RLE encoding, got {other:?}"), + } + } + + fn assert_decoded_i32_eq(decoded: DataBlock, expected: &[i32]) { + match decoded { + DataBlock::FixedWidth(block) => { + let values = block.data.borrow_to_typed_slice::(); + assert_eq!(values.as_ref(), expected); + } + _ => panic!("Expected FixedWidth block"), + } + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_rle_miniblock_compressed_children_multiple_chunks() { + let compression = test_general_compression(); + let encoder = RleEncoder::with_child_encoding( + RunLengthWidth::U8, + Some(compression), + Some(compression), + false, + ); + let expected = repeating_runs(8192, 4); + let (compressed, encoding) = compress_miniblock( + &encoder, + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + + assert!(compressed.chunks.len() > 1); + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::General(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::General(_) + )); + + let decoded = decompress_i32_chunks(&compressed, &encoding); + assert_eq!(decoded, expected); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_bitpacks_values_child_when_smaller() { + let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); + let expected = monotonic_runs(2048, 4); + let (compressed, encoding) = compress_miniblock( + &encoder, + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::Flat(_) + )); + + let decoded = decompress_i32_chunks(&compressed, &encoding); + assert_eq!(decoded, expected); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_bitpacks_run_lengths_when_values_do_not_shrink() { + let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); + let expected = high_entropy_runs(2048, 4); + let (compressed, encoding) = compress_miniblock( + &encoder, + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::Flat(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_) + )); + + let decoded = decompress_i32_chunks(&compressed, &encoding); + assert_eq!(decoded, expected); + } + + fn decompress_i32_chunks( + compressed: &MiniBlockCompressed, + encoding: &CompressiveEncoding, + ) -> Vec { + let strategy = DefaultDecompressionStrategy::default(); + let decompressor = strategy + .create_miniblock_decompressor(encoding, &strategy) + .unwrap(); + let mut offsets = vec![0usize; compressed.data.len()]; + let mut values_processed = 0u64; + let mut decoded_values = Vec::new(); + + for chunk in &compressed.chunks { + let chunk_values = chunk.num_values(values_processed, compressed.num_values); + let mut chunk_buffers = Vec::with_capacity(chunk.buffer_sizes.len()); + for (idx, size) in chunk.buffer_sizes.iter().enumerate() { + let size = *size as usize; + chunk_buffers.push(compressed.data[idx].slice_with_length(offsets[idx], size)); + offsets[idx] += size; + } + + let decoded = decompressor + .decompress(chunk_buffers, chunk_values) + .unwrap(); + match decoded { + DataBlock::FixedWidth(block) => { + let values = block.data.borrow_to_typed_slice::(); + decoded_values.extend_from_slice(values.as_ref()); + } + _ => panic!("Expected FixedWidth block"), + } + values_processed += chunk_values; + } + + assert_eq!(values_processed, compressed.num_values); + decoded_values + } + + #[cfg(feature = "bitpacking")] + fn monotonic_runs(num_runs: usize, run_length: usize) -> Vec { + let mut values = Vec::with_capacity(num_runs * run_length); + for run in 0..num_runs { + values.extend(std::iter::repeat_n(run as i32, run_length)); + } + values + } + + #[cfg(feature = "bitpacking")] + fn high_entropy_runs(num_runs: usize, run_length: usize) -> Vec { + let mut values = Vec::with_capacity(num_runs * run_length); + let mut state = 7u64; + for _ in 0..num_runs { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + values.extend(std::iter::repeat_n((state >> 32) as i32, run_length)); + } + values + } + + #[test] + fn test_select_run_length_width_prefers_u16_for_long_runs() { + let mut entries = [0u64; 3]; + accumulate_run_length_entries(300, Some(*MAX_MINIBLOCK_VALUES), &mut entries); + let (width, _) = select_run_length_width_from_entries(&entries, 32).unwrap(); + assert_eq!(width, RunLengthWidth::U16); + } + + // ========== Round-trip Tests for Different Types ========== + + #[test] + fn test_round_trip_all_types() { + // Test u8 + test_round_trip_helper(vec![42u8, 42, 42, 100, 100, 255, 255, 255, 255], 8); + + // Test u16 + test_round_trip_helper(vec![1000u16, 1000, 2000, 2000, 2000, 3000], 16); + + // Test i32 + test_round_trip_helper(vec![100i32, 100, 100, -200, -200, 300, 300, 300, 300], 32); + + // Test u64 + test_round_trip_helper(vec![1_000_000_000u64; 5], 64); + } + + fn test_round_trip_helper(data: Vec, bits_per_value: u64) + where + T: bytemuck::Pod + PartialEq + std::fmt::Debug, + { + let encoder = RleEncoder::new(); + let bytes: Vec = data + .iter() + .flat_map(|v| bytemuck::bytes_of(v)) + .copied() + .collect(); + + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value, + data: LanceBuffer::from(bytes), + num_values: data.len() as u64, + block_info: BlockInfo::default(), + }); + + let (compressed, _) = compress_miniblock(&encoder, block).unwrap(); + let decompressor = RleDecompressor::new(bits_per_value); + let decompressed = MiniBlockDecompressor::decompress( + &decompressor, + compressed.data, + compressed.num_values, + ) + .unwrap(); + + match decompressed { + DataBlock::FixedWidth(ref block) => { + // Verify the decompressed data length matches expected + assert_eq!(block.data.len(), data.len() * std::mem::size_of::()); + } + _ => panic!("Expected FixedWidth block"), + } + } + + // ========== Chunk Boundary Tests ========== + + #[test] + fn test_power_of_two_chunking() { + let encoder = RleEncoder::new(); + + // Create data that will require multiple chunks + let test_sizes = vec![1000, 2500, 5000, 10000]; + + for size in test_sizes { + let data: Vec = (0..size) + .map(|i| i / 50) // Create runs of 50 + .collect(); + + let array = Int32Array::from(data); + let (compressed, _) = + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); + + // Verify all non-last chunks have power-of-2 values + for (i, chunk) in compressed.chunks.iter().enumerate() { + if i < compressed.chunks.len() - 1 { + assert!(chunk.log_num_values > 0); + let chunk_values = 1u64 << chunk.log_num_values; + assert!(chunk_values.is_power_of_two()); + assert!(chunk_values <= *MAX_MINIBLOCK_VALUES); + } else { + assert_eq!(chunk.log_num_values, 0); + } + } + } + } + + #[rstest] + #[case::u8_lengths(RunLengthWidth::U8)] + #[case::u16_lengths(RunLengthWidth::U16)] + #[case::u32_lengths(RunLengthWidth::U32)] + fn test_miniblock_chunk_counts_match_encoded_runs(#[case] run_length_width: RunLengthWidth) { + // This pattern crosses the 2,048-value boundary in the middle of a two-value run. + let levels = (0..4098) + .map(|index| if index % 3 == 0 { 1u16 } else { 0u16 }) + .collect::>(); + let num_values = levels.len() as u64; + let encoder = RleEncoder::with_run_length_width(run_length_width); + let (buffers, chunks) = encoder + .encode_data( + &LanceBuffer::reinterpret_vec(levels), + num_values, + u16::BITS as u64, + ) + .unwrap(); + + assert_eq!(buffers.len(), 2); + let bytes_per_length = run_length_width.bytes_per_value(); + let mut values_offset = 0usize; + let mut lengths_offset = 0usize; + let mut values_processed = 0u64; + + for chunk in &chunks { + let values_size = chunk.buffer_sizes[0] as usize; + let lengths_size = chunk.buffer_sizes[1] as usize; + let lengths_end = lengths_offset + lengths_size; + let chunk_lengths = &buffers[1].as_ref()[lengths_offset..lengths_end]; + let length_chunks = chunk_lengths.chunks_exact(bytes_per_length); + assert!(length_chunks.remainder().is_empty()); + let num_runs = length_chunks.len(); + let encoded_values = length_chunks + .map(|bytes| run_length_width.read_length(bytes)) + .sum::(); + let declared_values = chunk.num_values(values_processed, num_values); + + assert_eq!(values_size, num_runs * size_of::()); + assert_eq!(encoded_values, declared_values); + + values_offset += values_size; + lengths_offset = lengths_end; + values_processed += declared_values; + } + + assert_eq!(values_processed, num_values); + assert_eq!(values_offset, buffers[0].len()); + assert_eq!(lengths_offset, buffers[1].len()); + } + + // ========== Error Handling Tests ========== + + #[test] + fn test_encoder_rejects_zero_progress() { + let error = RleEncoder::new() + .encode_data(&LanceBuffer::empty(), 1, u16::BITS as u64) + .unwrap_err(); + + assert!( + matches!(&error, Error::Internal { .. }), + "expected internal error, got: {error:?}" + ); + assert!(error.to_string().contains("made no progress")); + assert!(error.to_string().contains("values_remaining=1")); + } + + #[test] + fn test_invalid_buffer_count() { + let decompressor = RleDecompressor::new(32); + let result = MiniBlockDecompressor::decompress( + &decompressor, + vec![LanceBuffer::from(vec![1, 2, 3, 4])], + 10, + ); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("expects exactly 2 buffers") + ); + } + + #[test] + fn test_buffer_consistency() { + let decompressor = RleDecompressor::new(32); + let values = LanceBuffer::from(vec![1, 0, 0, 0]); // 1 i32 value + let lengths = LanceBuffer::from(vec![5, 10]); // 2 lengths - mismatch! + let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 15); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("Inconsistent RLE buffers") + ); + } + + #[test] + fn test_u16_length_buffer_must_be_aligned() { + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let values = LanceBuffer::from(vec![1, 0, 0, 0]); + let lengths = LanceBuffer::from(vec![5]); + let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 5); + assert!(matches!(&result, Err(Error::InvalidInput { .. }))); + assert!( + result + .unwrap_err() + .to_string() + .contains("not divisible by 2") + ); + } + + #[test] + fn test_rle_rejects_underflow_and_zero_lengths_and_clamps_overflow() { + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let value = LanceBuffer::from(1i32.to_le_bytes().to_vec()); + + let underflow = MiniBlockDecompressor::decompress( + &decompressor, + vec![ + value.clone(), + LanceBuffer::from(4u16.to_le_bytes().to_vec()), + ], + 5, + ) + .unwrap_err(); + assert!(underflow.to_string().contains("produced 4 values")); + + let overflow = MiniBlockDecompressor::decompress( + &decompressor, + vec![ + value.clone(), + LanceBuffer::from(6u16.to_le_bytes().to_vec()), + ], + 5, + ) + .unwrap(); + match overflow { + DataBlock::FixedWidth(block) => { + assert_eq!(block.num_values, 5); + let decoded = block.data.borrow_to_typed_slice::(); + assert_eq!(decoded.as_ref(), &[1i32; 5]); + } + _ => panic!("Expected FixedWidth block"), + } + + let zero = MiniBlockDecompressor::decompress( + &decompressor, + vec![value, LanceBuffer::from(0u16.to_le_bytes().to_vec())], + 5, + ) + .unwrap_err(); + assert!(zero.to_string().contains("zero run length")); + } + + #[test] + fn test_block_rle_rejects_overflow() { + // Block payloads have no chunk boundaries, so run lengths summing past + // num_values can only be corruption and must stay a hard error. + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let values = 1i32.to_le_bytes(); + let lengths = 6u16.to_le_bytes(); + let mut payload = Vec::new(); + payload.extend_from_slice(&(values.len() as u64).to_le_bytes()); + payload.extend_from_slice(&values); + payload.extend_from_slice(&lengths); + + let error = + BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::from(payload)), 5) + .unwrap_err(); + assert!(matches!(&error, Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("overflowed expected value count") + ); + } + + #[test] + fn test_rle_truncates_legacy_chunk_boundary_overflow() { + // Legacy encoders emitted chunks declaring 2048 values whose final run crossed + // the checkpoint boundary (e.g. run lengths summing to 2080); the excess values + // are duplicated at the start of the next chunk and must be ignored here. + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let mut values = Vec::new(); + values.extend_from_slice(&7i32.to_le_bytes()); + values.extend_from_slice(&8i32.to_le_bytes()); + let mut lengths = Vec::new(); + lengths.extend_from_slice(&2000u16.to_le_bytes()); + lengths.extend_from_slice(&80u16.to_le_bytes()); + + let decoded = MiniBlockDecompressor::decompress( + &decompressor, + vec![LanceBuffer::from(values), LanceBuffer::from(lengths)], + 2048, + ) + .unwrap(); + match decoded { + DataBlock::FixedWidth(block) => { + assert_eq!(block.num_values, 2048); + let decoded = block.data.borrow_to_typed_slice::(); + let decoded = decoded.as_ref(); + assert_eq!(decoded.len(), 2048); + assert!(decoded[..2000].iter().all(|&v| v == 7)); + assert!(decoded[2000..].iter().all(|&v| v == 8)); + } + _ => panic!("Expected FixedWidth block"), + } + } + + #[test] + fn test_empty_data_handling() { + let encoder = RleEncoder::new(); + + // Test empty block + let empty_block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::from(vec![]), + num_values: 0, + block_info: BlockInfo::default(), + }); + + let (compressed, _) = compress_miniblock(&encoder, empty_block).unwrap(); + assert_eq!(compressed.num_values, 0); + assert!(compressed.data.is_empty()); + + // Test decompression of empty data + let decompressor = RleDecompressor::new(32); + let decompressed = MiniBlockDecompressor::decompress(&decompressor, vec![], 0).unwrap(); + + match decompressed { + DataBlock::FixedWidth(ref block) => { + assert_eq!(block.num_values, 0); + assert_eq!(block.data.len(), 0); + } + _ => panic!("Expected FixedWidth block"), + } + } + + // ========== Integration Test ========== + + #[test] + fn test_multi_chunk_round_trip() { + let encoder = RleEncoder::new(); + + // Create data that spans multiple chunks with mixed patterns + let mut data = Vec::new(); + + // High compression section + data.extend(vec![999i32; 2000]); + // Low compression section + data.extend(0..1000); + // Another high compression section + data.extend(vec![777i32; 2000]); + + let array = Int32Array::from(data.clone()); + let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); + + // Manually decompress all chunks + let mut reconstructed = Vec::new(); + let mut values_offset = 0usize; + let mut lengths_offset = 0usize; + let mut values_processed = 0u64; + + // We now have exactly 2 global buffers + assert_eq!(compressed.data.len(), 2); + let global_values = &compressed.data[0]; + let global_lengths = &compressed.data[1]; + + for chunk in &compressed.chunks { + let chunk_values = if chunk.log_num_values > 0 { + 1u64 << chunk.log_num_values + } else { + compressed.num_values - values_processed + }; + + // Extract chunk buffers from global buffers using buffer_sizes + let values_size = chunk.buffer_sizes[0] as usize; + let lengths_size = chunk.buffer_sizes[1] as usize; + + let chunk_values_buffer = global_values.slice_with_length(values_offset, values_size); + let chunk_lengths_buffer = + global_lengths.slice_with_length(lengths_offset, lengths_size); + + let decompressor = RleDecompressor::new(32); + let chunk_data = MiniBlockDecompressor::decompress( + &decompressor, + vec![chunk_values_buffer, chunk_lengths_buffer], + chunk_values, + ) + .unwrap(); + + values_offset += values_size; + lengths_offset += lengths_size; + values_processed += chunk_values; + + match chunk_data { + DataBlock::FixedWidth(ref block) => { + let values: &[i32] = bytemuck::cast_slice(block.data.as_ref()); + reconstructed.extend_from_slice(values); + } + _ => panic!("Expected FixedWidth block"), + } + } + + assert_eq!(reconstructed, data); + } + + #[test] + fn test_1024_boundary_conditions() { + // Comprehensive test for various boundary conditions at 1024 values + // This consolidates multiple bug tests that were previously separate + let encoder = RleEncoder::new(); + let decompressor = RleDecompressor::new(32); + + let test_cases = [ + ("runs_of_2", { + let mut data = Vec::new(); + for i in 0..512 { + data.push(i); + data.push(i); + } + data + }), + ("single_run_1024", vec![42i32; 1024]), + ("alternating_values", { + let mut data = Vec::new(); + for i in 0..1024 { + data.push(i % 2); + } + data + }), + ("run_boundary_255s", { + let mut data = Vec::new(); + data.extend(vec![1i32; 255]); + data.extend(vec![2i32; 255]); + data.extend(vec![3i32; 255]); + data.extend(vec![4i32; 255]); + data.extend(vec![5i32; 4]); + data + }), + ("unique_values_1024", (0..1024).collect::>()), + ("unique_plus_duplicate", { + // 1023 unique values followed by one duplicate (regression test) + let mut data = Vec::new(); + for i in 0..1023 { + data.push(i); + } + data.push(1022i32); // Last value same as second-to-last + data + }), + ("bug_4092_pattern", { + // Test exact scenario that produces 4092 bytes instead of 4096 + let mut data = Vec::new(); + for i in 0..1022 { + data.push(i); + } + data.push(999999i32); + data.push(999999i32); + data + }), + ]; + + for (test_name, data) in test_cases.iter() { + assert_eq!(data.len(), 1024, "Test case {} has wrong length", test_name); + + // Compress the data + let array = Int32Array::from(data.clone()); + let (compressed, _) = + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); + + // Decompress and verify + match MiniBlockDecompressor::decompress( + &decompressor, + compressed.data, + compressed.num_values, + ) { + Ok(decompressed) => match decompressed { + DataBlock::FixedWidth(ref block) => { + let values: &[i32] = bytemuck::cast_slice(block.data.as_ref()); + assert_eq!( + values.len(), + 1024, + "Test case {} got {} values, expected 1024", + test_name, + values.len() + ); + assert_eq!( + block.data.len(), + 4096, + "Test case {} got {} bytes, expected 4096", + test_name, + block.data.len() + ); + assert_eq!(values, &data[..], "Test case {} data mismatch", test_name); + } + _ => panic!("Test case {} expected FixedWidth block", test_name), + }, + Err(e) => { + if e.to_string().contains("4092") { + panic!("Test case {} found bug 4092: {}", test_name, e); + } + panic!("Test case {} failed with error: {}", test_name, e); + } + } + } + } + + #[test] + fn test_low_repetition_50pct_bug() { + // Test case that reproduces the 4092 bytes bug with low repetition (50%) + // This simulates the 1M benchmark case + let encoder = RleEncoder::new(); + + // Create 1M values with low repetition (50% chance of change) + let num_values = 1_048_576; // 1M values + let mut data = Vec::with_capacity(num_values); + let mut value = 0i32; + let mut rng = 12345u64; // Simple deterministic RNG + + for _ in 0..num_values { + data.push(value); + // Simple LCG for deterministic randomness + rng = rng.wrapping_mul(1664525).wrapping_add(1013904223); + // 50% chance to increment value + if (rng >> 16) & 1 == 1 { + value += 1; + } + } + + let bytes: Vec = data.iter().flat_map(|v| v.to_le_bytes()).collect(); + + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::from(bytes), + num_values: num_values as u64, + block_info: BlockInfo::default(), + }); + + let (compressed, _) = compress_miniblock(&encoder, block).unwrap(); + + // Debug first few chunks + for (i, chunk) in compressed.chunks.iter().take(5).enumerate() { + let _chunk_values = if chunk.log_num_values > 0 { + 1 << chunk.log_num_values + } else { + // Last chunk - calculate remaining + let prev_total: usize = compressed.chunks[..i] + .iter() + .map(|c| 1usize << c.log_num_values) + .sum(); + num_values - prev_total + }; + } + + // Try to decompress + let decompressor = RleDecompressor::new(32); + match MiniBlockDecompressor::decompress( + &decompressor, + compressed.data, + compressed.num_values, + ) { + Ok(decompressed) => match decompressed { + DataBlock::FixedWidth(ref block) => { + assert_eq!( + block.data.len(), + num_values * 4, + "Expected {} bytes but got {}", + num_values * 4, + block.data.len() + ); + } + _ => panic!("Expected FixedWidth block"), + }, + Err(e) => { + if e.to_string().contains("4092") { + panic!("Bug reproduced! {}", e); + } else { + panic!("Unexpected error: {}", e); + } + } + } + } + + // ========== Encoding Verification Tests ========== + + #[test_log::test(tokio::test)] + async fn test_rle_encoding_verification() { + use crate::testing::{TestCases, check_round_trip_encoding_of_data}; + use arrow_array::{Array, Int32Array}; + use lance_datagen::{ArrayGenerator, RowCount}; + use std::collections::HashMap; + use std::sync::Arc; + + let test_cases = TestCases::default() + .with_expected_encoding("rle") + .with_structural_encodings(); + + // Test both explicit metadata and automatic selection + // 1. Test with explicit RLE threshold metadata (also disable BSS) + let mut metadata_explicit = HashMap::new(); + metadata_explicit.insert( + "lance-encoding:rle-threshold".to_string(), + "0.8".to_string(), + ); + metadata_explicit.insert("lance-encoding:bss".to_string(), "off".to_string()); + + let mut generator = RleDataGenerator::new(vec![ + i32::MIN, + i32::MIN, + i32::MIN, + i32::MIN, + i32::MIN + 1, + i32::MIN + 1, + i32::MIN + 1, + i32::MIN + 1, + i32::MIN + 2, + i32::MIN + 2, + i32::MIN + 2, + i32::MIN + 2, + ]); + let data_explicit = generator.generate_default(RowCount::from(10000)).unwrap(); + check_round_trip_encoding_of_data(vec![data_explicit], &test_cases, metadata_explicit) + .await; + + // 2. Test automatic RLE selection based on data characteristics + // 80% repetition should trigger RLE (> default 50% threshold). + // + // Use values with the high bit set so bitpacking can't shrink the values. + // Explicitly disable BSS to ensure RLE is tested + let mut metadata = HashMap::new(); + metadata.insert("lance-encoding:bss".to_string(), "off".to_string()); + + let mut values = vec![i32::MIN; 8000]; // 80% repetition + values.extend( + [ + i32::MIN + 1, + i32::MIN + 2, + i32::MIN + 3, + i32::MIN + 4, + i32::MIN + 5, + ] + .repeat(400), + ); // 20% variety + let arr = Arc::new(Int32Array::from(values)) as Arc; + check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + + #[cfg(any(feature = "lz4", feature = "zstd"))] + { + let mut metadata = HashMap::new(); + metadata.insert( + "lance-encoding:rle-threshold".to_string(), + "0.8".to_string(), + ); + metadata.insert("lance-encoding:bss".to_string(), "off".to_string()); + metadata.insert( + "lance-encoding:compression".to_string(), + if cfg!(feature = "zstd") { + "zstd".to_string() + } else { + "lz4".to_string() + }, + ); + let mut values = Vec::with_capacity(2048 * 4); + for run in 0..2048 { + values.extend(std::iter::repeat_n(i32::MIN + (run % 8), 4)); + } + let arr = Arc::new(Int32Array::from(values)) as Arc; + check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + } + } + + /// Generator that produces repetitive patterns suitable for RLE + #[derive(Debug)] + struct RleDataGenerator { + pattern: Vec, + idx: usize, + } + + impl RleDataGenerator { + fn new(pattern: Vec) -> Self { + Self { pattern, idx: 0 } + } + } + + impl lance_datagen::ArrayGenerator for RleDataGenerator { + fn generate( + &mut self, + _length: lance_datagen::RowCount, + _rng: &mut rand_xoshiro::Xoshiro256PlusPlus, + ) -> std::result::Result, arrow_schema::ArrowError> + { + use arrow_array::Int32Array; + use std::sync::Arc; + + // Generate enough repetitive data to trigger RLE + let mut values = Vec::new(); + for _ in 0..10000 { + values.push(self.pattern[self.idx]); + self.idx = (self.idx + 1) % self.pattern.len(); + } + Ok(Arc::new(Int32Array::from(values))) + } + + fn data_type(&self) -> &arrow_schema::DataType { + &arrow_schema::DataType::Int32 + } + + fn element_size_bytes(&self) -> Option { + Some(lance_datagen::ByteCount::from(4)) + } + } + + // ========== Block Related tests ========== + #[test] + fn test_block_decompressor_rejects_overflowing_values_size() { + let decompressor = RleDecompressor::new(32); + + let mut data = Vec::new(); + data.extend_from_slice(&u64::MAX.to_le_bytes()); + let result = BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::from(data)), 1); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("Invalid RLE values buffer size") + ); + } + + #[test] + fn test_block_decompressor_too_small() { + let decompressor = RleDecompressor::new(32); + let result = BlockDecompressor::decompress( + &decompressor, + Some(LanceBuffer::from(vec![1, 2, 3])), + 10, + ); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("Insufficient data size: 3") + ); + } + + #[test] + fn test_block_compressor_header_format() { + let encoder = RleEncoder::new(); + + let data = vec![1i32, 1, 1]; + let array = Int32Array::from(data); + let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)) + .unwrap() + .0 + .unwrap(); + + // Verify header format: first 8 bytes should be values_size as u64 + assert!(compressed.len() >= 8); + let values_size_bytes: [u8; 8] = compressed.as_ref()[..8].try_into().unwrap(); + let values_size = u64::from_le_bytes(values_size_bytes); + + // Values buffer should contain 1 i32 value (4 bytes) + assert_eq!(values_size, 4); + + // Total size should be: 8 (header) + 4 (values) + 1 (lengths) + assert_eq!(compressed.len(), 13); + } + + #[test] + fn test_block_compressor_round_trip() { + let encoder = RleEncoder::new(); + let decompressor = RleDecompressor::new(32); + + // Test basic pattern + let data = vec![1i32, 1, 1, 2, 2, 3, 3, 3, 3]; + let array = Int32Array::from(data.clone()); + let data_block = DataBlock::from_array(array); + + let compressed = BlockCompressor::compress(&encoder, data_block).unwrap().0; + let decompressed = + BlockDecompressor::decompress(&decompressor, compressed, data.len() as u64).unwrap(); + + match decompressed { + DataBlock::FixedWidth(block) => { + let values: &[i32] = bytemuck::cast_slice(block.data.as_ref()); + assert_eq!(values, &data[..]); + } + _ => panic!("Expected FixedWidth block"), + } + } + + #[test] + fn test_block_compressor_large_data() { + let encoder = RleEncoder::new(); + let decompressor = RleDecompressor::new(32); + + // Create data that will span multiple chunks + // Each chunks can handle ~2048 values, so use 10K values + let mut data = Vec::new(); + data.extend(vec![999i32; 3000]); // First ~2 chunks + data.extend(vec![777i32; 3000]); // Next ~2 chunks + data.extend(vec![555i32; 4000]); // Final ~2 chunks + + let total_values = data.len(); + assert_eq!(total_values, 10000); + + let array = Int32Array::from(data.clone()); + let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)) + .unwrap() + .0; + let decompressed = + BlockDecompressor::decompress(&decompressor, compressed, total_values as u64).unwrap(); + + match decompressed { + DataBlock::FixedWidth(block) => { + let values: &[i32] = bytemuck::cast_slice(block.data.as_ref()); + assert_eq!(values.len(), total_values); + assert_eq!(values, &data[..]); + } + _ => panic!("Expected FixedWidth block"), + } + } +} diff --git a/vendor/lance-encoding/src/encodings/physical/value.rs b/vendor/lance-encoding/src/encodings/physical/value.rs new file mode 100644 index 000000000..0dafa478f --- /dev/null +++ b/vendor/lance-encoding/src/encodings/physical/value.rs @@ -0,0 +1,1396 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow_buffer::{BooleanBufferBuilder, bit_util}; + +use crate::buffer::LanceBuffer; +use crate::compression::{ + BlockCompressor, BlockDecompressor, FixedPerValueDecompressor, MiniBlockDecompressor, + require_block_payload, +}; +use crate::data::{ + AllNullDataBlock, BlockInfo, DataBlock, FixedSizeListBlock, FixedWidthDataBlock, + NullableDataBlock, +}; +use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}; +use crate::encodings::logical::primitive::miniblock::{ + MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, + MiniBlockCompressionContext, MiniBlockCompressor, +}; +use crate::format::ProtobufUtils21; +use crate::format::pb21::compressive_encoding::Compression; +use crate::format::pb21::{self, CompressiveEncoding}; + +use lance_core::{Error, Result}; + +/// A compression strategy that writes fixed-width data as-is (no compression) +#[derive(Debug, Default)] +pub struct ValueEncoder {} + +impl ValueEncoder { + /// Use the largest chunk we can smaller than 4KiB + fn find_log_vals_per_chunk(bytes_per_word: u64, values_per_word: u64) -> Result<(u64, u64)> { + let mut size_bytes = 2 * bytes_per_word; + let (mut log_num_vals, mut num_vals) = match values_per_word { + 1 => (1, 2), + 8 => (3, 8), + _ => unreachable!(), + }; + + if size_bytes >= MAX_MINIBLOCK_BYTES { + let num_values = 2 * values_per_word; + return Err(Error::invalid_input(format!( + "Value is too wide for miniblock encoding: {} values require {} bytes but a \ + miniblock chunk is limited to {} bytes.", + num_values, size_bytes, MAX_MINIBLOCK_BYTES + ))); + } + + while 2 * size_bytes < MAX_MINIBLOCK_BYTES && 2 * num_vals <= *MAX_MINIBLOCK_VALUES { + log_num_vals += 1; + size_bytes *= 2; + num_vals *= 2; + } + + Ok((log_num_vals, num_vals)) + } + + fn chunk_data(data: FixedWidthDataBlock) -> Result { + // Usually there are X bytes per value. However, when working with boolean + // or FSL we might have some number of bits per value that isn't + // divisible by 8. In this case, to avoid chunking in the middle of a byte + // we calculate how many 8-value words we can fit in a chunk. + let (bytes_per_word, values_per_word) = if data.bits_per_value.is_multiple_of(8) { + (data.bits_per_value / 8, 1) + } else { + (data.bits_per_value, 8) + }; + + // Aim for 4KiB chunks + let (log_vals_per_chunk, vals_per_chunk) = + Self::find_log_vals_per_chunk(bytes_per_word, values_per_word)?; + let num_chunks = bit_util::ceil(data.num_values as usize, vals_per_chunk as usize); + debug_assert_eq!(vals_per_chunk % values_per_word, 0); + let bytes_per_chunk = bytes_per_word * (vals_per_chunk / values_per_word); + let bytes_per_chunk = u32::try_from(bytes_per_chunk).unwrap(); + debug_assert!(bytes_per_chunk > 0); + + let data_buffer = data.data; + + let mut row_offset = 0; + let mut chunks = Vec::with_capacity(num_chunks); + + let mut bytes_counter = 0; + loop { + if row_offset + vals_per_chunk <= data.num_values { + // We can make a full chunk + chunks.push(MiniBlockChunk { + log_num_values: log_vals_per_chunk as u8, + buffer_sizes: vec![bytes_per_chunk], + }); + row_offset += vals_per_chunk; + bytes_counter += bytes_per_chunk as u64; + } else if row_offset < data.num_values { + // Final chunk, special values + let num_bytes = data_buffer.len() as u64 - bytes_counter; + let num_bytes = u32::try_from(num_bytes).unwrap(); + chunks.push(MiniBlockChunk { + log_num_values: 0, + buffer_sizes: vec![num_bytes], + }); + break; + } else { + // If we get here then all chunks were full chunks and we have no remainder chunk + break; + } + } + + debug_assert_eq!(chunks.len(), num_chunks); + + Ok(MiniBlockCompressed { + chunks, + data: vec![data_buffer], + num_values: data.num_values, + }) + } +} + +#[derive(Debug)] +struct MiniblockFslLayer { + validity: Option, + dimension: u64, +} + +/// This impl deals with encoding FSL>>> data as a mini-block compressor. +/// The tricky part of FSL data is that we want to include inner validity buffers (we don't want these +/// to be part of the rep-def because that usually ends up being more expensive). +/// +/// The resulting mini-block will, instead of having a single buffer, have X + 1 buffers where X is +/// the number of FSL layers that contain validity. +/// +/// In the simple case where there is no validity inside the FSL layers, all we are doing here is flattening +/// the FSL layers into a single buffer. +/// +/// Also: We don't allow a row to be broken across chunks. This typically isn't too big of a deal since we +/// are usually dealing with relatively small vectors if we are using mini-block. +/// +/// Note: when we do have validity we have to make copies of the validity buffers because they are bit buffers +/// and we need to bit slice them which requires copies or offsets. Paying the price at write time to make +/// the copies is better than paying the price at read time to do the bit slicing. +impl ValueEncoder { + fn make_fsl_encoding(layers: &[MiniblockFslLayer], bits_per_value: u64) -> CompressiveEncoding { + let mut encoding = ProtobufUtils21::flat(bits_per_value, None); + for layer in layers.iter().rev() { + let has_validity = layer.validity.is_some(); + let dimension = layer.dimension; + encoding = ProtobufUtils21::fsl(dimension, has_validity, encoding); + } + encoding + } + + fn extract_fsl_chunk( + data: &FixedWidthDataBlock, + layers: &[MiniblockFslLayer], + row_offset: usize, + num_rows: usize, + validity_buffers: &mut [Vec], + ) -> Vec { + let mut row_offset = row_offset; + let mut num_values = num_rows; + let mut buffer_counter = 0; + let mut buffer_sizes = Vec::with_capacity(validity_buffers.len() + 1); + for layer in layers { + row_offset *= layer.dimension as usize; + num_values *= layer.dimension as usize; + if let Some(validity) = &layer.validity { + let validity_slice = validity + .clone() + .bit_slice_le_with_length(row_offset, num_values); + validity_buffers[buffer_counter].extend_from_slice(&validity_slice); + buffer_sizes.push(validity_slice.len() as u32); + buffer_counter += 1; + } + } + + let bits_in_chunk = data.bits_per_value * num_values as u64; + let bytes_in_chunk = bits_in_chunk.div_ceil(8); + let bytes_in_chunk = u32::try_from(bytes_in_chunk).unwrap(); + debug_assert!(bytes_in_chunk > 0); + buffer_sizes.push(bytes_in_chunk); + + buffer_sizes + } + + fn chunk_fsl( + data: FixedWidthDataBlock, + layers: Vec, + num_rows: u64, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + // Count size to calculate rows per chunk + let mut ceil_bytes_validity = 0; + let mut cum_dim = 1; + let mut num_validity_buffers = 0; + for layer in &layers { + cum_dim *= layer.dimension; + if layer.validity.is_some() { + ceil_bytes_validity += cum_dim.div_ceil(8); + num_validity_buffers += 1; + } + } + // It's an estimate because validity buffers may have some padding bits + let cum_bits_per_value = data.bits_per_value * cum_dim; + let (cum_bytes_per_word, vals_per_word) = if cum_bits_per_value.is_multiple_of(8) { + (cum_bits_per_value / 8, 1) + } else { + (cum_bits_per_value, 8) + }; + let est_bytes_per_word = (ceil_bytes_validity * vals_per_word) + cum_bytes_per_word; + let (log_rows_per_chunk, rows_per_chunk) = + Self::find_log_vals_per_chunk(est_bytes_per_word, vals_per_word)?; + + let num_chunks = num_rows.div_ceil(rows_per_chunk) as usize; + + // Allocate buffers for validity, these will be slightly bigger than the input validity buffers + let mut chunks = Vec::with_capacity(num_chunks); + let mut validity_buffers: Vec> = Vec::with_capacity(num_validity_buffers); + cum_dim = 1; + for layer in &layers { + cum_dim *= layer.dimension; + if let Some(validity) = &layer.validity { + let layer_bytes_validity = cum_dim.div_ceil(8); + let validity_with_padding = + layer_bytes_validity as usize * num_chunks * rows_per_chunk as usize; + debug_assert!(validity_with_padding >= validity.len()); + validity_buffers.push(Vec::with_capacity( + layer_bytes_validity as usize * num_chunks, + )); + } + } + + // Now go through and extract validity buffers + let mut row_offset = 0; + while row_offset + rows_per_chunk <= num_rows { + let buffer_sizes = Self::extract_fsl_chunk( + &data, + &layers, + row_offset as usize, + rows_per_chunk as usize, + &mut validity_buffers, + ); + row_offset += rows_per_chunk; + chunks.push(MiniBlockChunk { + log_num_values: log_rows_per_chunk as u8, + buffer_sizes, + }) + } + let rows_in_chunk = num_rows - row_offset; + if rows_in_chunk > 0 { + let buffer_sizes = Self::extract_fsl_chunk( + &data, + &layers, + row_offset as usize, + rows_in_chunk as usize, + &mut validity_buffers, + ); + chunks.push(MiniBlockChunk { + log_num_values: 0, + buffer_sizes, + }); + } + + let encoding = Self::make_fsl_encoding(&layers, data.bits_per_value); + // Finally, add the data buffer + let buffers = validity_buffers + .into_iter() + .map(LanceBuffer::from) + .chain(std::iter::once(data.data)) + .collect::>(); + + Ok(( + MiniBlockCompressed { + chunks, + data: buffers, + num_values: num_rows, + }, + encoding, + )) + } + + fn miniblock_fsl(data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + let num_rows = data.num_values(); + let fsl = data.as_fixed_size_list().unwrap(); + let mut layers = Vec::new(); + let mut child = *fsl.child; + let mut cur_layer = MiniblockFslLayer { + validity: None, + dimension: fsl.dimension, + }; + loop { + if let DataBlock::Nullable(nullable) = child { + cur_layer.validity = Some(nullable.nulls); + child = *nullable.data; + } + match child { + DataBlock::FixedSizeList(inner) => { + layers.push(cur_layer); + cur_layer = MiniblockFslLayer { + validity: None, + dimension: inner.dimension, + }; + child = *inner.child; + } + DataBlock::FixedWidth(inner) => { + layers.push(cur_layer); + return Self::chunk_fsl(inner, layers, num_rows); + } + _ => unreachable!("Unexpected data block type in value encoder's miniblock_fsl"), + } + } + } +} + +struct PerValueFslValidityIter { + buffer: LanceBuffer, + bits_per_row: usize, + offset: usize, +} + +/// In this section we deal with per-value encoding of FSL>>> data. +/// +/// It's easier than mini-block. All we need to do is flatten the FSL layers into a single buffer. +/// This includes any validity buffers we encounter on the way. +impl ValueEncoder { + fn fsl_to_encoding(fsl: &FixedSizeListBlock) -> CompressiveEncoding { + let mut inner = fsl.child.as_ref(); + let mut has_validity = false; + if let DataBlock::Nullable(nullable) = inner { + has_validity = true; + inner = nullable.data.as_ref(); + } + let inner_encoding = match inner { + // All inner values are null. Reserve validity bits (has_validity=true) so that + // the decoder knows one validity byte per cum_dim items is stored per row. + // constant(None) signals that every stored item decodes to null. + DataBlock::AllNull(_) => { + has_validity = true; + ProtobufUtils21::constant(None) + } + DataBlock::FixedWidth(fixed_width) => { + ProtobufUtils21::flat(fixed_width.bits_per_value, None) + } + DataBlock::FixedSizeList(inner) => Self::fsl_to_encoding(inner), + _ => unreachable!( + "Unexpected data block type in value encoder's fsl_to_encoding: {}", + inner.name() + ), + }; + ProtobufUtils21::fsl(fsl.dimension, has_validity, inner_encoding) + } + + fn simple_per_value_fsl(fsl: FixedSizeListBlock) -> (PerValueDataBlock, CompressiveEncoding) { + // The simple case is zero-copy, we just return the flattened inner buffer + let encoding = Self::fsl_to_encoding(&fsl); + let num_values = fsl.num_values(); + let mut child = *fsl.child; + let mut cum_dim = 1; + loop { + cum_dim *= fsl.dimension; + match child { + DataBlock::Nullable(nullable) => { + child = *nullable.data; + } + DataBlock::FixedSizeList(inner) => { + child = *inner.child; + } + DataBlock::FixedWidth(inner) => { + let data = FixedWidthDataBlock { + bits_per_value: inner.bits_per_value * cum_dim, + num_values, + data: inner.data, + block_info: BlockInfo::new(), + }; + return (PerValueDataBlock::Fixed(data), encoding); + } + _ => unreachable!( + "Unexpected data block type in value encoder's simple_per_value_fsl" + ), + } + } + } + + fn nullable_per_value_fsl(fsl: FixedSizeListBlock) -> (PerValueDataBlock, CompressiveEncoding) { + // If there are nullable inner values then we need to zip the validity with the values + let encoding = Self::fsl_to_encoding(&fsl); + let num_values = fsl.num_values(); + let mut bytes_per_row = 0; + let mut cum_dim = 1; + let mut current = fsl; + let mut validity_iters: Vec = Vec::new(); + let data_bytes_per_row: usize; + let data_buffer: LanceBuffer; + loop { + cum_dim *= current.dimension; + let mut child = *current.child; + if let DataBlock::Nullable(nullable) = child { + // Each item will need this many bytes of validity prepended to it + bytes_per_row += cum_dim.div_ceil(8) as usize; + validity_iters.push(PerValueFslValidityIter { + buffer: nullable.nulls, + bits_per_row: cum_dim as usize, + offset: 0, + }); + child = *nullable.data; + }; + match child { + DataBlock::FixedSizeList(inner) => { + current = inner; + } + DataBlock::FixedWidth(fixed_width) => { + data_bytes_per_row = + (fixed_width.bits_per_value.div_ceil(8) * cum_dim) as usize; + bytes_per_row += data_bytes_per_row; + data_buffer = fixed_width.data; + break; + } + DataBlock::AllNull(_) => { + // All inner values are null. Add all-zero validity bits so that + // bytes_per_row > 0 and the FullZip layout doesn't write bits_per_value=0, + // which would crash the reader when there are also no ctrl-word bytes. + bytes_per_row += cum_dim.div_ceil(8) as usize; + validity_iters.push(PerValueFslValidityIter { + buffer: LanceBuffer::from(vec![ + 0u8; + cum_dim.div_ceil(8) as usize + * num_values as usize + ]), + bits_per_row: cum_dim as usize, + offset: 0, + }); + data_bytes_per_row = 0; + data_buffer = LanceBuffer::empty(); + break; + } + _ => unreachable!( + "Unexpected data block type in value encoder's nullable_per_value_fsl: {:?}", + child + ), + } + } + + let bytes_needed = bytes_per_row * num_values as usize; + let mut zipped = Vec::with_capacity(bytes_needed); + let data_slice = &data_buffer; + // Hopefully values are pretty large so we don't iterate this loop _too_ many times + for i in 0..num_values as usize { + for validity in validity_iters.iter_mut() { + let validity_slice = validity + .buffer + .bit_slice_le_with_length(validity.offset, validity.bits_per_row); + zipped.extend_from_slice(&validity_slice); + validity.offset += validity.bits_per_row; + } + let start = i * data_bytes_per_row; + let end = start + data_bytes_per_row; + zipped.extend_from_slice(&data_slice[start..end]); + } + + let zipped = LanceBuffer::from(zipped); + let data = PerValueDataBlock::Fixed(FixedWidthDataBlock { + bits_per_value: bytes_per_row as u64 * 8, + num_values, + data: zipped, + block_info: BlockInfo::new(), + }); + (data, encoding) + } + + fn per_value_fsl(fsl: FixedSizeListBlock) -> (PerValueDataBlock, CompressiveEncoding) { + if !fsl.child.is_nullable() { + Self::simple_per_value_fsl(fsl) + } else { + Self::nullable_per_value_fsl(fsl) + } + } +} + +impl BlockCompressor for ValueEncoder { + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + let DataBlock::FixedWidth(fixed_width) = data else { + return Err(Error::invalid_input(format!( + "ValueEncoder cannot compress a {} block", + data.name() + ))); + }; + Ok(( + Some(fixed_width.data), + ProtobufUtils21::flat(fixed_width.bits_per_value, None), + )) + } +} + +impl MiniBlockCompressor for ValueEncoder { + fn compress( + &self, + _context: MiniBlockCompressionContext, + chunk: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + match chunk { + DataBlock::FixedWidth(fixed_width) => { + let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); + Ok((Self::chunk_data(fixed_width)?, encoding)) + } + DataBlock::FixedSizeList(_) => Self::miniblock_fsl(chunk), + _ => Err(Error::invalid_input_source( + format!( + "Cannot compress a data block of type {} with ValueEncoder", + chunk.name() + ) + .into(), + )), + } + } +} + +#[derive(Debug)] +struct ValueFslDesc { + dimension: u64, + has_validity: bool, +} + +/// A decompressor for fixed-width data that has +/// been written, as-is, to disk in single contiguous array +#[derive(Debug)] +pub struct ValueDecompressor { + /// How many bits are in each inner-most item (e.g. FSL would be 32) + bits_per_item: u64, + /// How many bits are in each value (e.g. FSL would be 3200) + /// + /// This number is a little trickier to compute because we also have to include bytes + /// of any inner validity + bits_per_value: u64, + /// How many items are in each value (e.g. FSL would be 100) + items_per_value: u64, + layers: Vec, +} + +impl ValueDecompressor { + pub fn from_flat(description: &pb21::Flat) -> Self { + Self { + bits_per_item: description.bits_per_value, + bits_per_value: description.bits_per_value, + items_per_value: 1, + layers: Vec::default(), + } + } + + pub fn from_fsl(mut description: &pb21::FixedSizeList) -> Result { + let mut layers = Vec::new(); + let mut cum_dim = 1; + let mut bytes_per_value = 0; + loop { + layers.push(ValueFslDesc { + has_validity: description.has_validity, + dimension: description.items_per_value, + }); + cum_dim *= description.items_per_value; + if description.has_validity { + bytes_per_value += cum_dim.div_ceil(8); + } + let encoding = description + .values + .as_ref() + .ok_or_else(|| Error::invalid_input("FSL encoding missing inner values field"))? + .compression + .as_ref() + .ok_or_else(|| { + Error::invalid_input("FSL encoding missing inner compression field") + })?; + match encoding { + Compression::FixedSizeList(inner) => { + description = inner; + } + Compression::Flat(flat) => { + let mut bits_per_value = bytes_per_value * 8; + bits_per_value += flat.bits_per_value * cum_dim; + return Ok(Self { + bits_per_item: flat.bits_per_value, + bits_per_value, + items_per_value: cum_dim, + layers, + }); + } + // All inner values are null: only validity bytes are stored per row. + // bits_per_item=0 signals unzip_decompress to emit AllNull for the values. + Compression::Constant(_) => { + return Ok(Self { + bits_per_item: 0, + bits_per_value: bytes_per_value * 8, + items_per_value: cum_dim, + layers, + }); + } + _ => { + return Err(Error::invalid_input(format!( + "Unexpected inner encoding type in FSL descriptor: {:?}", + encoding + ))); + } + } + } + } + + fn buffer_to_block(&self, data: LanceBuffer, num_values: u64) -> DataBlock { + DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.bits_per_item, + num_values, + data, + block_info: BlockInfo::new(), + }) + } +} + +impl BlockDecompressor for ValueDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Flat block")?; + let block = self.buffer_to_block(data, num_values); + assert_eq!(block.num_values(), num_values); + Ok(block) + } +} + +impl MiniBlockDecompressor for ValueDecompressor { + fn decompress(&self, data: Vec, num_values: u64) -> Result { + let num_items = num_values * self.items_per_value; + let mut buffer_iter = data.into_iter().rev(); + + // Always at least 1 buffer + let data_buf = buffer_iter.next().unwrap(); + let items = self.buffer_to_block(data_buf, num_items); + let mut lists = items; + + for layer in self.layers.iter().rev() { + if layer.has_validity { + let validity_buf = buffer_iter.next().unwrap(); + lists = DataBlock::Nullable(NullableDataBlock { + data: Box::new(lists), + nulls: validity_buf, + block_info: BlockInfo::default(), + }); + } + lists = DataBlock::FixedSizeList(FixedSizeListBlock { + child: Box::new(lists), + dimension: layer.dimension, + }) + } + + assert_eq!(lists.num_values(), num_values); + Ok(lists) + } + + fn decoded_size_bytes(&self, num_values: u64) -> Option { + if self.has_validity() { + return None; + } + num_values + .checked_mul(self.bits_per_value) + .map(|bits| bits.div_ceil(8)) + } +} + +struct FslDecompressorValidityBuilder { + buffer: BooleanBufferBuilder, + bits_per_row: usize, + bytes_per_row: usize, +} + +// Helper methods for per-value decompression +impl ValueDecompressor { + fn has_validity(&self) -> bool { + self.layers.iter().any(|layer| layer.has_validity) + } + + // If there is no validity then decompression is zero-copy, we just need to restore any FSL layers + fn simple_decompress(&self, data: FixedWidthDataBlock, num_rows: u64) -> DataBlock { + let mut cum_dim = 1; + for layer in &self.layers { + cum_dim *= layer.dimension; + } + debug_assert_eq!(self.bits_per_item, data.bits_per_value / cum_dim); + let mut block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.bits_per_item, + num_values: num_rows * cum_dim, + data: data.data, + block_info: BlockInfo::new(), + }); + for layer in self.layers.iter().rev() { + block = DataBlock::FixedSizeList(FixedSizeListBlock { + child: Box::new(block), + dimension: layer.dimension, + }); + } + debug_assert_eq!(num_rows, block.num_values()); + block + } + + // If there is validity then it has been zipped in with the values and we must unzip it + fn unzip_decompress(&self, data: FixedWidthDataBlock, num_rows: usize) -> DataBlock { + // No support for full-zip on per-value encodings + assert_eq!(self.bits_per_item % 8, 0); + let bytes_per_item = self.bits_per_item / 8; + let mut buffer_builders = Vec::with_capacity(self.layers.len()); + let mut cum_dim = 1; + let mut total_size_bytes = 0; + // First, go through the layers, setup our builders, allocate space + for layer in &self.layers { + cum_dim *= layer.dimension as usize; + if layer.has_validity { + let validity_size_bits = cum_dim; + let validity_size_bytes = validity_size_bits.div_ceil(8); + total_size_bytes += num_rows * validity_size_bytes; + buffer_builders.push(FslDecompressorValidityBuilder { + buffer: BooleanBufferBuilder::new(validity_size_bits * num_rows), + bits_per_row: cum_dim, + bytes_per_row: validity_size_bytes, + }) + } + } + let num_items = num_rows * cum_dim; + let data_size = num_items * bytes_per_item as usize; + total_size_bytes += data_size; + let mut data_buffer = Vec::with_capacity(data_size); + + assert_eq!(data.data.len(), total_size_bytes); + + let bytes_per_value = bytes_per_item as usize; + let data_bytes_per_row = bytes_per_value * cum_dim; + + // Next, unzip + let mut data_offset = 0; + while data_offset < total_size_bytes { + for builder in buffer_builders.iter_mut() { + let start = data_offset * 8; + let end = start + builder.bits_per_row; + builder.buffer.append_packed_range(start..end, &data.data); + data_offset += builder.bytes_per_row; + } + let end = data_offset + data_bytes_per_row; + data_buffer.extend_from_slice(&data.data[data_offset..end]); + data_offset += data_bytes_per_row; + } + + // Finally, restore the structure + // + // bits_per_item=0 means the terminal encoding is constant-null (all inner values + // are null and only validity bits were stored). Use AllNull so that into_arrow + // produces a properly typed all-null array without a zero-width data buffer. + let mut block = if self.bits_per_item == 0 { + debug_assert!(data_buffer.is_empty()); + DataBlock::AllNull(AllNullDataBlock { + num_values: num_items as u64, + }) + } else { + DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.bits_per_item, + num_values: num_items as u64, + data: LanceBuffer::from(data_buffer), + block_info: BlockInfo::new(), + }) + }; + + let mut validity_bufs = buffer_builders + .into_iter() + .rev() + .map(|mut b| LanceBuffer::from(b.buffer.finish().into_inner())); + for layer in self.layers.iter().rev() { + if layer.has_validity { + let nullable = NullableDataBlock { + data: Box::new(block), + nulls: validity_bufs.next().unwrap(), + block_info: BlockInfo::new(), + }; + block = DataBlock::Nullable(nullable); + } + block = DataBlock::FixedSizeList(FixedSizeListBlock { + child: Box::new(block), + dimension: layer.dimension, + }); + } + + assert_eq!(num_rows, block.num_values() as usize); + + block + } +} + +impl FixedPerValueDecompressor for ValueDecompressor { + fn decompress(&self, data: FixedWidthDataBlock, num_rows: u64) -> Result { + if self.has_validity() { + Ok(self.unzip_decompress(data, num_rows as usize)) + } else { + Ok(self.simple_decompress(data, num_rows)) + } + } + + fn bits_per_value(&self) -> u64 { + self.bits_per_value + } + + fn decoded_size_bytes(&self, num_values: u64) -> Option { + if self.has_validity() { + return None; + } + num_values + .checked_mul(self.bits_per_value) + .map(|bits| bits.div_ceil(8)) + } +} + +impl PerValueCompressor for ValueEncoder { + fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> { + let (data, encoding) = match data { + DataBlock::FixedWidth(fixed_width) => { + let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); + (PerValueDataBlock::Fixed(fixed_width), encoding) + } + DataBlock::FixedSizeList(fixed_size_list) => Self::per_value_fsl(fixed_size_list), + _ => unimplemented!( + "Cannot compress block of type {} with ValueEncoder", + data.name() + ), + }; + Ok((data, encoding)) + } +} + +// public tests module because we share the PRIMITIVE_TYPES constant with fixed_size_list +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use arrow_array::{ + Array, ArrayRef, Decimal128Array, FixedSizeListArray, Int32Array, ListArray, UInt8Array, + make_array, new_null_array, types::UInt32Type, + }; + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field, TimeUnit}; + use lance_datagen::{ArrayGeneratorExt, Dimension, RowCount, Seed, array, gen_batch}; + + use crate::{ + compression::{FixedPerValueDecompressor, MiniBlockDecompressor}, + data::DataBlock, + encodings::{ + logical::primitive::{ + fullzip::{PerValueCompressor, PerValueDataBlock}, + miniblock::{MiniBlockCompressionContext, MiniBlockCompressor}, + }, + physical::value::ValueDecompressor, + }, + format::pb21::compressive_encoding::Compression, + testing::{ + TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data, + }, + }; + + use super::ValueEncoder; + + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) + } + + const PRIMITIVE_TYPES: &[DataType] = &[ + DataType::Null, + DataType::FixedSizeBinary(2), + DataType::Date32, + DataType::Date64, + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + DataType::UInt8, + DataType::UInt16, + DataType::UInt32, + DataType::UInt64, + DataType::Float16, + DataType::Float32, + DataType::Float64, + DataType::Decimal128(10, 10), + DataType::Decimal256(10, 10), + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Time32(TimeUnit::Second), + DataType::Time64(TimeUnit::Nanosecond), + DataType::Duration(TimeUnit::Second), + // The Interval type is supported by the reader but the writer works with Lance schema + // at the moment and Lance schema can't parse interval + // DataType::Interval(IntervalUnit::DayTime), + ]; + + #[test_log::test(tokio::test)] + async fn test_simple_value() { + let items = Arc::new(Int32Array::from(vec![ + Some(0), + None, + Some(2), + Some(3), + Some(4), + Some(5), + ])); + + let test_cases = TestCases::default() + .with_range(0..3) + .with_range(0..2) + .with_range(1..3) + .with_indices(vec![0, 1, 2]) + .with_indices(vec![1]) + .with_indices(vec![2]) + .with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![items], &test_cases, HashMap::default()).await; + } + + #[test_log::test(tokio::test)] + async fn test_simple_range() { + let items = Arc::new(Int32Array::from_iter( + (0..5000).map(|i| if i % 2 == 0 { Some(i) } else { None }), + )); + + let test_cases = TestCases::default().with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![items], &test_cases, HashMap::default()).await; + } + + #[test_log::test(tokio::test)] + async fn test_value_primitive() { + const NUM_ROWS: u32 = 1025; + + let test_cases = TestCases::default() + .with_batch_size(NUM_ROWS) + .with_page_sizes(vec![4096]) + .with_expected_encoding("flat"); + let value_metadata = + HashMap::from([("lance-encoding:compression".to_string(), "none".to_string())]); + + for (seed, data_type) in PRIMITIVE_TYPES.iter().enumerate() { + log::info!("Testing encoding for {:?}", data_type); + let data = gen_batch() + .with_seed(Seed::from(seed as u64)) + .anon_col(array::rand_type(data_type)) + .into_batch_rows(RowCount::from(NUM_ROWS as u64)) + .unwrap() + .column(0) + .clone(); + + check_round_trip_encoding_of_data(vec![data], &test_cases, value_metadata.clone()) + .await; + } + } + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_large_primitive( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + #[values(4096, 1024 * 1024)] page_size: u64, + #[values(false, true)] use_slicing: bool, + ) { + let data_type = + DataType::FixedSizeList(Arc::new(Field::new("", DataType::Int32, false)), 128); + let field = Field::new("", data_type, false); + check_basic_random_case(field, encoding, page_size, use_slicing).await; + } + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_decimal128_dictionary_encoding( + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + let test_cases = TestCases::default() + .with_encoding(encoding) + .with_expected_encoding("dictionary"); + let decimals: Vec = (0..100).collect(); + let repeated_strings: Vec<_> = decimals + .iter() + .cycle() + .take(decimals.len() * 1000) + .map(|&v| Some(v as i128)) + .collect(); + let decimal_array = Arc::new(Decimal128Array::from(repeated_strings)) as ArrayRef; + check_round_trip_encoding_of_data(vec![decimal_array], &test_cases, HashMap::new()).await; + } + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_miniblock_stress( + #[values(false, true)] mixed_validity: bool, + #[values(10, 100, 1500, 15000)] batch_size: u32, + #[values(1000, 2000, 3000, 60000)] page_size: u64, + #[values( + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + // Tests for strange page sizes and batch sizes and validity scenarios for miniblock + + // 10K integers, 100 per array, either all valid or mixed validity. + let data = (0..100) + .map(|_| { + if mixed_validity { + Arc::new(Int32Array::from_iter( + (0..100).map(|i| if i % 2 == 0 { Some(i) } else { None }), + )) as Arc + } else { + Arc::new(Int32Array::from_iter_values(0..100)) as Arc + } + }) + .collect::>(); + + let test_cases = TestCases::default() + .with_page_sizes(vec![page_size]) + .with_batch_size(batch_size) + .with_encoding(encoding); + + check_round_trip_encoding_of_data(data, &test_cases, HashMap::new()).await; + } + + fn create_simple_fsl() -> FixedSizeListArray { + // [[0, 1], NULL], [NULL, NULL], [[8, 9], [NULL, 11]] + let items = Arc::new(Int32Array::from(vec![ + Some(0), + Some(1), + Some(2), + Some(3), + None, + None, + None, + None, + Some(8), + Some(9), + None, + Some(11), + ])); + let items_field = Arc::new(Field::new("item", DataType::Int32, true)); + let inner_list_nulls = BooleanBuffer::from(vec![true, false, false, false, true, true]); + let inner_list = Arc::new(FixedSizeListArray::new( + items_field.clone(), + 2, + items, + Some(NullBuffer::new(inner_list_nulls)), + )); + let inner_list_field = Arc::new(Field::new( + "item", + DataType::FixedSizeList(items_field, 2), + true, + )); + FixedSizeListArray::new(inner_list_field, 2, inner_list, None) + } + + #[test] + fn test_fsl_value_compression_miniblock() { + let sample_list = create_simple_fsl(); + + let starting_data = DataBlock::from_array(sample_list.clone()); + + let encoder = ValueEncoder::default(); + let (data, compression) = + MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data).unwrap(); + + assert_eq!(data.num_values, 3); + assert_eq!(data.data.len(), 3); + assert_eq!(data.chunks.len(), 1); + assert_eq!(data.chunks[0].buffer_sizes, vec![1, 2, 48]); + assert_eq!(data.chunks[0].log_num_values, 0); + + let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { + panic!() + }; + + let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap(); + + let decompressed = + MiniBlockDecompressor::decompress(&decompressor, data.data, data.num_values).unwrap(); + + let decompressed = make_array( + decompressed + .into_arrow(sample_list.data_type().clone(), true) + .unwrap(), + ); + + assert_eq!(decompressed.as_ref(), &sample_list); + } + + fn wide_fixed_size_binary() -> ArrayRef { + let wide_value = vec![0xABu8; 5000]; + Arc::new( + arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size( + std::iter::repeat_n(Some(wide_value.as_slice()), 4), + 5000, + ) + .unwrap(), + ) + } + + fn wide_fixed_size_list_bool() -> ArrayRef { + // A wide FSL is sub-byte, so it chunks eight values per word and the + // smallest unit is 16 values rather than 2. + let dimension = 4095; + let values = arrow_array::BooleanArray::from(vec![false; dimension * 2]); + let field = Arc::new(Field::new("item", DataType::Boolean, true)); + Arc::new(FixedSizeListArray::new( + field, + dimension as i32, + Arc::new(values), + None, + )) + } + + #[rstest::rstest] + #[case::fixed_size_binary(wide_fixed_size_binary(), 2)] + #[case::fixed_size_list_bool(wide_fixed_size_list_bool(), 16)] + fn test_wide_value_miniblock_returns_error( + #[case] array: ArrayRef, + #[case] expected_min_values: u64, + ) { + let starting_data = DataBlock::from_array(array); + + let encoder = ValueEncoder::default(); + let result = MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data); + + let err = result.expect_err("wide values should not be encodable as miniblock"); + assert!( + matches!(err, lance_core::Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + let msg = err.to_string(); + assert!( + msg.contains("too wide for miniblock encoding"), + "unexpected error message: {msg}" + ); + assert!( + msg.contains(&format!("{expected_min_values} values require")), + "unexpected error message: {msg}" + ); + } + + #[test] + fn test_fsl_value_compression_per_value() { + let sample_list = create_simple_fsl(); + + let starting_data = DataBlock::from_array(sample_list.clone()); + + let encoder = ValueEncoder::default(); + let (data, compression) = PerValueCompressor::compress(&encoder, starting_data).unwrap(); + + let PerValueDataBlock::Fixed(data) = data else { + panic!() + }; + + assert_eq!(data.bits_per_value, 144); + assert_eq!(data.num_values, 3); + assert_eq!(data.data.len(), 18 * 3); + + let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { + panic!() + }; + + let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap(); + + let num_values = data.num_values; + assert_eq!( + FixedPerValueDecompressor::decoded_size_bytes(&decompressor, num_values), + None, + "nullable FSL output uses multiple buffers and requires the fallback estimate" + ); + let decompressed = + FixedPerValueDecompressor::decompress(&decompressor, data, num_values).unwrap(); + + let decompressed = make_array( + decompressed + .into_arrow(sample_list.data_type().clone(), true) + .unwrap(), + ); + + assert_eq!(decompressed.as_ref(), &sample_list); + } + + #[test_log::test(tokio::test)] + async fn test_fsl_all_null() { + let items = new_null_array(&DataType::Int32, 12); + let items_field = Arc::new(Field::new("item", DataType::Int32, true)); + let list_nulls = BooleanBuffer::from(vec![true, false, false, false, true, true]); + let list_array = + FixedSizeListArray::new(items_field, 2, items, Some(NullBuffer::new(list_nulls))); + + let test_cases = TestCases::default().with_structural_encodings(); + + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) + .await; + } + + // Regression: nullable_per_value_fsl wrote bits_per_value=0 when the child block was + // AllNull but the FSL had no outer null buffer (no ctrl-word bytes), making + // total_bytes_per_value=0 and crashing the reader with "per-row byte width must be > 0". + #[test_log::test(tokio::test)] + async fn test_fsl_nullable_child_all_null_no_outer_nulls() { + // FSL, 2 outer rows, no outer null buffer, all 8 child + // Float32 values are null. DataBlock::from_arrays returns AllNull for the child, + // which routes into nullable_per_value_fsl. Without this fix the encoder writes + // bits_per_value=0 and the decoder errors on read. + let items = new_null_array(&DataType::Float32, 8); + let items_field = Arc::new(Field::new("item", DataType::Float32, true)); + let fsl = FixedSizeListArray::new(items_field, 4, items, None); + + let test_cases = TestCases::default().with_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(fsl)], &test_cases, HashMap::new()).await; + } + + #[test_log::test(tokio::test)] + async fn regress_list_fsl() { + // This regresses a case where rows are large lists that span multiple + // mini-block chunks which gives us some all-premable mini-block chunks. + let offsets = ScalarBuffer::::from(vec![0, 393, 755, 1156, 1536]); + let data = UInt8Array::from(vec![0; 1536 * 16]); + let fsl_field = Arc::new(Field::new("item", DataType::UInt8, true)); + let fsl = FixedSizeListArray::new(fsl_field, 16, Arc::new(data), None); + let list_field = Arc::new(Field::new("item", fsl.data_type().clone(), false)); + let list_arr = ListArray::new(list_field, OffsetBuffer::new(offsets), Arc::new(fsl), None); + + let test_cases = TestCases::default() + .with_structural_encodings() + .with_batch_size(1); + + check_round_trip_encoding_of_data(vec![Arc::new(list_arr)], &test_cases, HashMap::new()) + .await; + } + + fn create_random_fsl() -> Arc { + // Several levels of def and multiple pages + let inner = array::rand_type(&DataType::Int32).with_random_nulls(0.1); + let list_one = array::cycle_vec(inner, Dimension::from(4)).with_random_nulls(0.1); + let list_two = array::cycle_vec(list_one, Dimension::from(4)).with_random_nulls(0.1); + let list_three = array::cycle_vec(list_two, Dimension::from(2)); + + // Should be 256Ki rows ~ 1MiB of data + let batch = gen_batch() + .anon_col(list_three) + .into_batch_rows(RowCount::from(8 * 1024)) + .unwrap(); + batch.column(0).clone() + } + + #[test] + fn fsl_value_miniblock_stress() { + let sample_array = create_random_fsl(); + + let starting_data = DataBlock::from_arrays( + std::slice::from_ref(&sample_array), + sample_array.len() as u64, + ); + + let encoder = ValueEncoder::default(); + let (data, compression) = + MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data).unwrap(); + + let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { + panic!() + }; + + let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap(); + + let decompressed = + MiniBlockDecompressor::decompress(&decompressor, data.data, data.num_values).unwrap(); + + let decompressed = make_array( + decompressed + .into_arrow(sample_array.data_type().clone(), true) + .unwrap(), + ); + + assert_eq!(decompressed.as_ref(), sample_array.as_ref()); + } + + #[test] + fn fsl_value_per_value_stress() { + let sample_array = create_random_fsl(); + + let starting_data = DataBlock::from_arrays( + std::slice::from_ref(&sample_array), + sample_array.len() as u64, + ); + + let encoder = ValueEncoder::default(); + let (data, compression) = PerValueCompressor::compress(&encoder, starting_data).unwrap(); + + let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { + panic!() + }; + + let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap(); + + let PerValueDataBlock::Fixed(data) = data else { + panic!() + }; + + let num_values = data.num_values; + let decompressed = + FixedPerValueDecompressor::decompress(&decompressor, data, num_values).unwrap(); + + let decompressed = make_array( + decompressed + .into_arrow(sample_array.data_type().clone(), true) + .unwrap(), + ); + + assert_eq!(decompressed.as_ref(), sample_array.as_ref()); + } + + #[rstest::rstest] + #[test_log::test(tokio::test)] + async fn test_fsl_nullable_items( + #[values( + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse + )] + encoding: TestEncoding, + ) { + let mut generator = + gen_batch() + .with_seed(Seed::from(0)) + .anon_col(array::rand_vec_nullable::( + Dimension::from(128), + 0.5, + )); + generator.with_random_nulls(0.2); + let source = generator + .into_batch_rows(RowCount::from(1026)) + .unwrap() + .column(0) + .clone(); + let test_cases = TestCases::default() + .with_page_sizes(vec![4096]) + .with_encoding(encoding) + .with_batch_size(257) + .with_range(510..515) + .with_indices(vec![0, 511, 512, 1024]); + + check_round_trip_encoding_of_data( + vec![source.slice(1, 512), source.slice(513, 513)], + &test_cases, + HashMap::new(), + ) + .await; + } + + #[test_log::test(tokio::test)] + async fn test_value_encoding_verification() { + use std::collections::HashMap; + + let test_cases = TestCases::default() + .with_expected_encoding("flat") + .with_structural_encodings(); + + // Test both explicit configuration and automatic fallback scenarios + // 1. Test explicit "none" compression to force flat encoding + // Also explicitly disable BSS to ensure value encoding is tested + let mut metadata_explicit = HashMap::new(); + metadata_explicit.insert("lance-encoding:compression".to_string(), "none".to_string()); + metadata_explicit.insert("lance-encoding:bss".to_string(), "off".to_string()); + + let arr_explicit = + Arc::new(Int32Array::from((0..1000).collect::>())) as Arc; + check_round_trip_encoding_of_data(vec![arr_explicit], &test_cases, metadata_explicit).await; + + // 2. Test automatic fallback to flat encoding when bitpacking conditions aren't met + // Use unique values to avoid RLE encoding + // Explicitly disable BSS to ensure value encoding is tested + let mut metadata = HashMap::new(); + metadata.insert("lance-encoding:bss".to_string(), "off".to_string()); + + let arr_fallback = Arc::new(Int32Array::from( + (0..100).map(|i| i * 73 + 19).collect::>(), + )) as Arc; + check_round_trip_encoding_of_data(vec![arr_fallback], &test_cases, metadata).await; + } + + #[test_log::test(tokio::test)] + async fn test_mixed_page_validity() { + let no_nulls = Arc::new(Int32Array::from_iter_values([1, 2])); + let has_nulls = Arc::new(Int32Array::from_iter([Some(3), None, Some(5)])); + + let test_cases = TestCases::default().with_page_sizes(vec![1]); + check_round_trip_encoding_of_data(vec![no_nulls, has_nulls], &test_cases, HashMap::new()) + .await; + } +} diff --git a/vendor/lance-encoding/src/format.rs b/vendor/lance-encoding/src/format.rs new file mode 100644 index 000000000..f37e69b02 --- /dev/null +++ b/vendor/lance-encoding/src/format.rs @@ -0,0 +1,773 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +/// Protobuf definitions for encodings +/// +/// These are the messages used for describing encoding in the 2.0 format +pub mod pb { + #![allow(clippy::all)] + #![allow(non_upper_case_globals)] + #![allow(non_camel_case_types)] + #![allow(non_snake_case)] + #![allow(unused)] + #![allow(improper_ctypes)] + #![allow(clippy::upper_case_acronyms)] + #![allow(clippy::use_self)] + include!(concat!(env!("OUT_DIR"), "/lance.encodings.rs")); +} + +/// Protobuf definitions for encodings21 +/// +/// These are the messages used for describing encoding in the 2.1 format +/// and any newer formats. +pub mod pb21 { + #![allow(clippy::all)] + #![allow(non_upper_case_globals)] + #![allow(non_camel_case_types)] + #![allow(non_snake_case)] + #![allow(unused)] + #![allow(improper_ctypes)] + #![allow(clippy::upper_case_acronyms)] + #![allow(clippy::use_self)] + include!(concat!(env!("OUT_DIR"), "/lance.encodings21.rs")); +} + +use pb::{ + ArrayEncoding, Binary, Bitpacked, BitpackedForNonNeg, Block, Dictionary, FixedSizeBinary, + FixedSizeList, Flat, Fsst, InlineBitpacking, Nullable, OutOfLineBitpacking, PackedStruct, + PackedStructFixedWidthMiniBlock, Rle, Variable, + array_encoding::ArrayEncoding as ArrayEncodingEnum, + buffer::BufferType, + nullable::{AllNull, NoNull, Nullability, SomeNull}, +}; + +use crate::{encodings::physical::block::CompressionConfig, repdef::DefinitionInterpretation}; + +use self::pb::Constant; +use lance_core::Result; + +// Utility functions for creating complex protobuf objects +pub struct ProtobufUtils {} + +impl ProtobufUtils { + pub fn constant(value: Vec) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Constant(Constant { + value: value.into(), + })), + } + } + + pub fn basic_all_null_encoding() -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Nullable(Box::new(Nullable { + nullability: Some(Nullability::AllNulls(AllNull {})), + }))), + } + } + + pub fn basic_some_null_encoding( + validity: ArrayEncoding, + values: ArrayEncoding, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Nullable(Box::new(Nullable { + nullability: Some(Nullability::SomeNulls(Box::new(SomeNull { + validity: Some(Box::new(validity)), + values: Some(Box::new(values)), + }))), + }))), + } + } + + pub fn basic_no_null_encoding(values: ArrayEncoding) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Nullable(Box::new(Nullable { + nullability: Some(Nullability::NoNulls(Box::new(NoNull { + values: Some(Box::new(values)), + }))), + }))), + } + } + + pub fn block(scheme: &str) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Block(Block { + scheme: scheme.to_string(), + })), + } + } + + pub fn flat_encoding( + bits_per_value: u64, + buffer_index: u32, + compression: Option, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Flat(Flat { + bits_per_value, + buffer: Some(pb::Buffer { + buffer_index, + buffer_type: BufferType::Page as i32, + }), + compression: compression.map(|compression_config| pb::Compression { + scheme: compression_config.scheme.to_string(), + level: compression_config.level, + }), + })), + } + } + + pub fn fsl_encoding(dimension: u64, items: ArrayEncoding, has_validity: bool) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::FixedSizeList(Box::new(FixedSizeList { + dimension: dimension.try_into().unwrap(), + items: Some(Box::new(items)), + has_validity, + }))), + } + } + + pub fn bitpacked_encoding( + compressed_bits_per_value: u64, + uncompressed_bits_per_value: u64, + buffer_index: u32, + signed: bool, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Bitpacked(Bitpacked { + compressed_bits_per_value, + buffer: Some(pb::Buffer { + buffer_index, + buffer_type: BufferType::Page as i32, + }), + uncompressed_bits_per_value, + signed, + })), + } + } + + pub fn bitpacked_for_non_neg_encoding( + compressed_bits_per_value: u64, + uncompressed_bits_per_value: u64, + buffer_index: u32, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::BitpackedForNonNeg(BitpackedForNonNeg { + compressed_bits_per_value, + buffer: Some(pb::Buffer { + buffer_index, + buffer_type: BufferType::Page as i32, + }), + uncompressed_bits_per_value, + })), + } + } + pub fn inline_bitpacking(uncompressed_bits_per_value: u64) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::InlineBitpacking(InlineBitpacking { + uncompressed_bits_per_value, + })), + } + } + pub fn out_of_line_bitpacking( + uncompressed_bits_per_value: u64, + compressed_bits_per_value: u64, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::OutOfLineBitpacking( + OutOfLineBitpacking { + uncompressed_bits_per_value, + compressed_bits_per_value, + }, + )), + } + } + + pub fn variable(bits_per_offset: u8) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Variable(Variable { + bits_per_offset: bits_per_offset as u32, + })), + } + } + + // Construct a `FsstMiniBlock` ArrayEncoding, the inner `binary_mini_block` encoding is actually + // not used and `FsstMiniBlockDecompressor` constructs a `binary_mini_block` in a `hard-coded` fashion. + // This can be an optimization later. + pub fn fsst(data: ArrayEncoding, symbol_table: Vec) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Fsst(Box::new(Fsst { + binary: Some(Box::new(data)), + symbol_table: symbol_table.into(), + }))), + } + } + + pub fn rle(bits_per_value: u64) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Rle(Rle { bits_per_value })), + } + } + + pub fn byte_stream_split(bits_per_value: u64) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::ByteStreamSplit(pb::ByteStreamSplit { + bits_per_value, + })), + } + } + + pub fn general_mini_block( + inner: ArrayEncoding, + compression: CompressionConfig, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::GeneralMiniBlock(Box::new( + pb::GeneralMiniBlock { + inner: Some(Box::new(inner)), + compression: Some(pb::Compression { + scheme: compression.scheme.to_string(), + level: compression.level, + }), + }, + ))), + } + } + + pub fn packed_struct( + child_encodings: Vec, + packed_buffer_index: u32, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::PackedStruct(PackedStruct { + inner: child_encodings, + buffer: Some(pb::Buffer { + buffer_index: packed_buffer_index, + buffer_type: BufferType::Page as i32, + }), + })), + } + } + + pub fn packed_struct_fixed_width_mini_block( + data: ArrayEncoding, + bits_per_values: Vec, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::PackedStructFixedWidthMiniBlock( + Box::new(PackedStructFixedWidthMiniBlock { + flat: Some(Box::new(data)), + bits_per_values, + }), + )), + } + } + + pub fn binary( + indices_encoding: ArrayEncoding, + bytes_encoding: ArrayEncoding, + null_adjustment: u64, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Binary(Box::new(Binary { + bytes: Some(Box::new(bytes_encoding)), + indices: Some(Box::new(indices_encoding)), + null_adjustment, + }))), + } + } + + pub fn dict_encoding( + indices: ArrayEncoding, + items: ArrayEncoding, + num_items: u32, + ) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::Dictionary(Box::new(Dictionary { + indices: Some(Box::new(indices)), + items: Some(Box::new(items)), + num_dictionary_items: num_items, + }))), + } + } + + pub fn fixed_size_binary(data: ArrayEncoding, byte_width: u32) -> ArrayEncoding { + ArrayEncoding { + array_encoding: Some(ArrayEncodingEnum::FixedSizeBinary(Box::new( + FixedSizeBinary { + bytes: Some(Box::new(data)), + byte_width, + }, + ))), + } + } +} + +macro_rules! impl_common_protobuf_utils { + ($module:ident, $struct_name:ident) => { + pub struct $struct_name {} + + impl $struct_name { + pub fn flat( + bits_per_value: u64, + values_compression: Option, + ) -> crate::format::$module::CompressiveEncoding { + crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::Flat( + crate::format::$module::Flat { + bits_per_value, + data: values_compression, + }, + ), + ), + } + } + + pub fn constant( + value: Option, + ) -> crate::format::$module::CompressiveEncoding { + crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::Constant( + crate::format::$module::Constant { value }, + ), + ), + } + } + + pub fn fsl( + items_per_value: u64, + has_validity: bool, + values: crate::format::$module::CompressiveEncoding, + ) -> crate::format::$module::CompressiveEncoding { + crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::FixedSizeList( + Box::new(crate::format::$module::FixedSizeList { + items_per_value, + has_validity, + values: Some(Box::new(values)), + }), + ), + ), + } + } + + pub fn variable( + offsets_desc: crate::format::$module::CompressiveEncoding, + values_compression: Option, + ) -> crate::format::$module::CompressiveEncoding { + crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::Variable( + Box::new(crate::format::$module::Variable { + offsets: Some(Box::new(offsets_desc)), + values: values_compression, + }), + ), + ), + } + } + + pub fn inline_bitpacking( + uncompressed_bits_per_value: u64, + values_compression: Option, + ) -> crate::format::$module::CompressiveEncoding { + crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::InlineBitpacking( + crate::format::$module::InlineBitpacking { + uncompressed_bits_per_value, + values: values_compression, + }, + ), + ), + } + } + + pub fn out_of_line_bitpacking( + uncompressed_bits_per_value: u64, + values: crate::format::$module::CompressiveEncoding, + ) -> crate::format::$module::CompressiveEncoding { + crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::OutOfLineBitpacking( + Box::new(crate::format::$module::OutOfLineBitpacking { + uncompressed_bits_per_value, + values: Some(Box::new(values)), + }), + ), + ), + } + } + + pub fn buffer_compression( + compression: CompressionConfig, + ) -> Result { + Ok(crate::format::$module::BufferCompression { + scheme: crate::format::$module::CompressionScheme::try_from( + compression.scheme, + )? as i32, + level: compression.level, + }) + } + + pub fn wrapped( + compression: CompressionConfig, + values: crate::format::$module::CompressiveEncoding, + ) -> Result { + Ok(crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::General( + Box::new(crate::format::$module::General { + compression: Some(Self::buffer_compression(compression)?), + values: Some(Box::new(values)), + }), + ), + ), + }) + } + + pub fn rle( + values: crate::format::$module::CompressiveEncoding, + run_lengths: crate::format::$module::CompressiveEncoding, + ) -> crate::format::$module::CompressiveEncoding { + crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::Rle(Box::new( + crate::format::$module::Rle { + values: Some(Box::new(values)), + run_lengths: Some(Box::new(run_lengths)), + }, + )), + ), + } + } + + pub fn byte_stream_split( + values: crate::format::$module::CompressiveEncoding, + ) -> crate::format::$module::CompressiveEncoding { + crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::ByteStreamSplit( + Box::new(crate::format::$module::ByteStreamSplit { + values: Some(Box::new(values)), + }), + ), + ), + } + } + + pub fn fsst( + data: crate::format::$module::CompressiveEncoding, + symbol_table: Vec, + ) -> crate::format::$module::CompressiveEncoding { + crate::format::$module::CompressiveEncoding { + compression: Some( + crate::format::$module::compressive_encoding::Compression::Fsst( + Box::new(crate::format::$module::Fsst { + symbol_table: symbol_table.into(), + values: Some(Box::new(data)), + }), + ), + ), + } + } + + fn def_inter_to_repdef_layer(def: DefinitionInterpretation) -> i32 { + match def { + DefinitionInterpretation::AllValidItem => { + crate::format::$module::RepDefLayer::RepdefAllValidItem as i32 + } + DefinitionInterpretation::AllValidList => { + crate::format::$module::RepDefLayer::RepdefAllValidList as i32 + } + DefinitionInterpretation::NullableItem => { + crate::format::$module::RepDefLayer::RepdefNullableItem as i32 + } + DefinitionInterpretation::NullableList => { + crate::format::$module::RepDefLayer::RepdefNullableList as i32 + } + DefinitionInterpretation::EmptyableList => { + crate::format::$module::RepDefLayer::RepdefEmptyableList as i32 + } + DefinitionInterpretation::NullableAndEmptyableList => { + crate::format::$module::RepDefLayer::RepdefNullAndEmptyList as i32 + } + } + } + + pub fn repdef_layer_to_def_interp( + layer: i32, + ) -> DefinitionInterpretation { + let layer = crate::format::$module::RepDefLayer::try_from(layer).unwrap(); + match layer { + crate::format::$module::RepDefLayer::RepdefAllValidItem => { + DefinitionInterpretation::AllValidItem + } + crate::format::$module::RepDefLayer::RepdefAllValidList => { + DefinitionInterpretation::AllValidList + } + crate::format::$module::RepDefLayer::RepdefNullableItem => { + DefinitionInterpretation::NullableItem + } + crate::format::$module::RepDefLayer::RepdefNullableList => { + DefinitionInterpretation::NullableList + } + crate::format::$module::RepDefLayer::RepdefEmptyableList => { + DefinitionInterpretation::EmptyableList + } + crate::format::$module::RepDefLayer::RepdefNullAndEmptyList => { + DefinitionInterpretation::NullableAndEmptyableList + } + crate::format::$module::RepDefLayer::RepdefUnspecified => { + panic!("Unspecified repdef layer") + } + } + } + + #[allow(clippy::too_many_arguments)] + pub fn miniblock_layout( + rep_encoding: Option, + def_encoding: Option, + value_encoding: crate::format::$module::CompressiveEncoding, + repetition_index_depth: u32, + num_buffers: u64, + dictionary_encoding: Option<( + crate::format::$module::CompressiveEncoding, + u64, + )>, + def_meaning: &[DefinitionInterpretation], + num_items: u64, + has_large_chunk: bool, + ) -> crate::format::$module::PageLayout { + assert!(!def_meaning.is_empty()); + let (dictionary, num_dictionary_items) = dictionary_encoding + .map(|(d, i)| (Some(d), i)) + .unwrap_or((None, 0)); + crate::format::$module::PageLayout { + layout: Some( + crate::format::$module::page_layout::Layout::MiniBlockLayout( + crate::format::$module::MiniBlockLayout { + def_compression: def_encoding, + rep_compression: rep_encoding, + value_compression: Some(value_encoding), + repetition_index_depth, + num_buffers, + dictionary, + num_dictionary_items, + layers: def_meaning + .iter() + .map(|&def| Self::def_inter_to_repdef_layer(def)) + .collect(), + num_items, + has_large_chunk, + }, + ), + ), + } + } + + fn full_zip_layout( + bits_rep: u8, + bits_def: u8, + details: crate::format::$module::full_zip_layout::Details, + value_encoding: crate::format::$module::CompressiveEncoding, + def_meaning: &[DefinitionInterpretation], + num_items: u32, + num_visible_items: u32, + ) -> crate::format::$module::PageLayout { + crate::format::$module::PageLayout { + layout: Some( + crate::format::$module::page_layout::Layout::FullZipLayout( + crate::format::$module::FullZipLayout { + bits_rep: bits_rep as u32, + bits_def: bits_def as u32, + details: Some(details), + value_compression: Some(value_encoding), + num_items, + num_visible_items, + layers: def_meaning + .iter() + .map(|&def| Self::def_inter_to_repdef_layer(def)) + .collect(), + }, + ), + ), + } + } + + pub fn fixed_full_zip_layout( + bits_rep: u8, + bits_def: u8, + bits_per_value: u32, + value_encoding: crate::format::$module::CompressiveEncoding, + def_meaning: &[DefinitionInterpretation], + num_items: u32, + num_visible_items: u32, + ) -> crate::format::$module::PageLayout { + Self::full_zip_layout( + bits_rep, + bits_def, + crate::format::$module::full_zip_layout::Details::BitsPerValue( + bits_per_value, + ), + value_encoding, + def_meaning, + num_items, + num_visible_items, + ) + } + + pub fn variable_full_zip_layout( + bits_rep: u8, + bits_def: u8, + bits_per_offset: u32, + value_encoding: crate::format::$module::CompressiveEncoding, + def_meaning: &[DefinitionInterpretation], + num_items: u32, + num_visible_items: u32, + ) -> crate::format::$module::PageLayout { + Self::full_zip_layout( + bits_rep, + bits_def, + crate::format::$module::full_zip_layout::Details::BitsPerOffset( + bits_per_offset, + ), + value_encoding, + def_meaning, + num_items, + num_visible_items, + ) + } + + pub fn blob_layout( + inner_layout: crate::format::$module::PageLayout, + def_meaning: &[DefinitionInterpretation], + ) -> crate::format::$module::PageLayout { + crate::format::$module::PageLayout { + layout: Some( + crate::format::$module::page_layout::Layout::BlobLayout(Box::new( + crate::format::$module::BlobLayout { + inner_layout: Some(Box::new(inner_layout)), + layers: def_meaning + .iter() + .map(|&def| Self::def_inter_to_repdef_layer(def)) + .collect(), + }, + )), + ), + } + } + + + } + }; +} + +impl_common_protobuf_utils!(pb21, ProtobufUtils21); + +impl ProtobufUtils21 { + pub fn constant_layout( + def_meaning: &[DefinitionInterpretation], + inline_value: Option>, + ) -> crate::format::pb21::PageLayout { + crate::format::pb21::PageLayout { + layout: Some(crate::format::pb21::page_layout::Layout::ConstantLayout( + crate::format::pb21::ConstantLayout { + inline_value: inline_value.map(bytes::Bytes::from), + rep_compression: None, + def_compression: None, + num_rep_values: 0, + num_def_values: 0, + layers: def_meaning + .iter() + .map(|&def| Self::def_inter_to_repdef_layer(def)) + .collect(), + }, + )), + } + } + + pub fn compressed_all_null_constant_layout( + def_meaning: &[DefinitionInterpretation], + rep_compression: Option, + def_compression: Option, + num_rep_values: u64, + num_def_values: u64, + ) -> crate::format::pb21::PageLayout { + crate::format::pb21::PageLayout { + layout: Some(crate::format::pb21::page_layout::Layout::ConstantLayout( + crate::format::pb21::ConstantLayout { + inline_value: None, + rep_compression, + def_compression, + num_rep_values, + num_def_values, + layers: def_meaning + .iter() + .map(|&def| Self::def_inter_to_repdef_layer(def)) + .collect(), + }, + )), + } + } + + pub fn packed_struct( + values: crate::format::pb21::CompressiveEncoding, + bits_per_values: Vec, + ) -> crate::format::pb21::CompressiveEncoding { + crate::format::pb21::CompressiveEncoding { + compression: Some( + crate::format::pb21::compressive_encoding::Compression::PackedStruct(Box::new( + crate::format::pb21::PackedStruct { + bits_per_value: bits_per_values, + values: Some(Box::new(values)), + }, + )), + ), + } + } + + pub fn packed_struct_variable( + fields: Vec, + ) -> crate::format::pb21::CompressiveEncoding { + crate::format::pb21::CompressiveEncoding { + compression: Some( + crate::format::pb21::compressive_encoding::Compression::VariablePackedStruct( + crate::format::pb21::VariablePackedStruct { fields }, + ), + ), + } + } + + pub fn packed_struct_field_fixed( + value_encoding: crate::format::pb21::CompressiveEncoding, + bits_per_value: u64, + ) -> crate::format::pb21::variable_packed_struct::FieldEncoding { + crate::format::pb21::variable_packed_struct::FieldEncoding { + value: Some(value_encoding), + layout: Some( + crate::format::pb21::variable_packed_struct::field_encoding::Layout::BitsPerValue( + bits_per_value, + ), + ), + } + } + + pub fn packed_struct_field_variable( + value_encoding: crate::format::pb21::CompressiveEncoding, + bits_per_length: u64, + ) -> crate::format::pb21::variable_packed_struct::FieldEncoding { + crate::format::pb21::variable_packed_struct::FieldEncoding { + value: Some(value_encoding), + layout: Some( + crate::format::pb21::variable_packed_struct::field_encoding::Layout::BitsPerLength( + bits_per_length, + ), + ), + } + } +} diff --git a/vendor/lance-encoding/src/lib.rs b/vendor/lance-encoding/src/lib.rs new file mode 100644 index 000000000..338028b79 --- /dev/null +++ b/vendor/lance-encoding/src/lib.rs @@ -0,0 +1,134 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ops::Range, sync::Arc}; + +use bytes::Bytes; +use futures::{FutureExt, TryFutureExt, future::BoxFuture}; + +use lance_core::Result; + +mod array_encoding; +pub mod buffer; +pub mod compression; +pub mod compression_config; +pub mod constants; +pub mod data; +pub mod decoder; +pub mod encoder; +pub mod encodings; +pub mod format; +pub mod repdef; +pub mod statistics; +#[cfg(test)] +pub mod testing; +pub mod utils; + +// We can definitely add support for big-endian machines someday. However, it's not a priority and +// would involve extensive testing (probably through emulation) to ensure that the encodings are +// correct. +#[cfg(not(target_endian = "little"))] +compile_error!("Lance encodings only support little-endian systems."); + +/// A trait for an I/O service +/// +/// This represents the I/O API that the encoders and decoders need in order to operate. +/// We specify this as a trait so that lance-encodings does not need to depend on lance-io +/// +/// In general, it is assumed that this trait will be implemented by some kind of "file reader" +/// or "file scheduler". The encodings here are all limited to accessing a single file. +pub trait EncodingsIo: std::fmt::Debug + Send + Sync { + /// Submit an I/O request + /// + /// The response must contain a `Bytes` object for each range requested even if the underlying + /// I/O was coalesced into fewer actual requests. + /// + /// # Arguments + /// + /// * `ranges` - the byte ranges to request + /// * `priority` - the priority of the request + /// + /// Priority should be set to the lowest row number that this request is delivering data for. + /// This is important in cases where indirect I/O causes high priority requests to be submitted + /// after low priority requests. We want to fulfill the indirect I/O more quickly so that we + /// can decode as quickly as possible. + /// + /// The implementation should be able to handle empty ranges, and should return an empty + /// byte buffer for each empty range. + fn submit_request( + &self, + range: Vec>, + priority: u64, + ) -> BoxFuture<'static, Result>>; + + /// Submit an I/O request with a single range + /// + /// This is just a utitliy function that wraps [`EncodingsIo::submit_request`] for the common + /// case of a single range request. + fn submit_single( + &self, + range: std::ops::Range, + priority: u64, + ) -> BoxFuture<'static, lance_core::Result> { + self.submit_request(vec![range], priority) + .map_ok(|mut v| v.pop().unwrap()) + .boxed() + } + + /// Returns a version of this I/O service that bypasses backpressure for all requests. + /// + /// This is intended for indirect I/O (e.g. fetching items after decoding offsets) where + /// blocking on backpressure could cause deadlocks or excessive latency. + /// + /// Returns `None` if this implementation does not support bypass (e.g. in-memory or test + /// schedulers), in which case the caller should fall back to using self. + fn with_bypass_backpressure(&self) -> Option> { + None + } + + /// Returns a version of this I/O service that additionally records the I/O it + /// performs into `stats`, on top of any global accounting. This is the seam + /// used to measure exact per-scope (e.g. per-query) I/O without re-opening + /// files: wrap a reader's I/O service, perform the reads, then inspect the + /// recorder. + /// + /// Returns `None` if this implementation does not support per-scope I/O + /// statistics (e.g. in-memory or test schedulers), in which case the caller + /// should fall back to using self (and no statistics are recorded). + fn with_io_stats( + &self, + _stats: Arc, + ) -> Option> { + None + } +} + +/// An implementation of EncodingsIo that serves data from an in-memory buffer +#[derive(Debug)] +pub struct BufferScheduler { + data: Bytes, +} + +impl BufferScheduler { + pub fn new(data: Bytes) -> Self { + Self { data } + } + + fn satisfy_request(&self, req: Range) -> Bytes { + self.data.slice(req.start as usize..req.end as usize) + } +} + +impl EncodingsIo for BufferScheduler { + fn submit_request( + &self, + ranges: Vec>, + _priority: u64, + ) -> BoxFuture<'static, Result>> { + std::future::ready(Ok(ranges + .into_iter() + .map(|range| self.satisfy_request(range)) + .collect::>())) + .boxed() + } +} diff --git a/vendor/lance-encoding/src/repdef.rs b/vendor/lance-encoding/src/repdef.rs new file mode 100644 index 000000000..1c102d1e7 --- /dev/null +++ b/vendor/lance-encoding/src/repdef.rs @@ -0,0 +1,4083 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Utilities for rep-def levels +//! +//! Repetition and definition levels are a way to encode multipile validity / offsets arrays +//! into a single buffer. They are a form of "zipping" buffers together that takes advantage +//! of the fact that, if the outermost array is invalid, then the validity of the inner items +//! is irrelevant. +//! +//! Note: the concept of repetition & definition levels comes from the Dremel paper and has +//! been implemented in Apache Parquet. However, the implementation here is not necessarily +//! compatible with Parquet. For example, we use 0 to represent the "inner-most" item and +//! Parquet uses 0 to represent the "outer-most" item. +//! +//! # Repetition Levels +//! +//! With repetition levels we convert a sparse array of offsets into a dense array of levels. +//! These levels are marked non-zero whenever a new list begins. In other words, given the +//! list array with 3 rows [{<0,1>, <>, <2>}, {<3>}, {}], [], [{<4>}] we would have three +//! offsets arrays: +//! +//! Outer-most ([]): [0, 3, 3, 4] +//! Middle ({}): [0, 3, 4, 4, 5] +//! Inner (<>): [0, 2, 2, 3, 4, 5] +//! Values : [0, 1, 2, 3, 4] +//! +//! We can convert these into repetition levels as follows: +//! +//! | Values | Repetition | +//! | ------ | ---------- | +//! | 0 | 3 | // Start of outer-most list +//! | 1 | 0 | // Continues inner-most list (no new lists) +//! | - | 1 | // Start of new inner-most list (empty list) +//! | 2 | 1 | // Start of new inner-most list +//! | 3 | 2 | // Start of new middle list +//! | - | 2 | // Start of new inner-most list (empty list) +//! | - | 3 | // Start of new outer-most list (empty list) +//! | 4 | 0 | // Start of new outer-most list +//! +//! Note: We actually have MORE repetition levels than values. This is because the repetition +//! levels need to be able to represent empty lists. +//! +//! # Definition Levels +//! +//! Definition levels are simpler. We can think of them as zipping together various validity bitmaps +//! (from different levels of nesting) into a single buffer. For example, we could zip the arrays +//! [1, 1, 0, 0] and [1, 0, 1, 0] into [11, 10, 01, 00]. However, 00 and 01 are redundant. If the +//! outer level is null then the validity of the inner levels is irrelevant. To save space we instead +//! encode a "level" which is the "depth" of the null. Let's look at a more complete example: +//! +//! Array: [{"middle": {"inner": 1]}}, NULL, {"middle": NULL}, {"middle": {"inner": NULL}}] +//! +//! In Arrow we would have the following validity arrays: +//! Outer validity : 1, 0, 1, 1 +//! Middle validity: 1, ?, 0, 1 +//! Inner validity : 1, ?, ?, 0 +//! Values : 1, ?, ?, ? +//! +//! The ? values are undefined in the Arrow format. We can convert these into definition levels as follows: +//! +//! | Values | Definition | +//! | ------ | ---------- | +//! | 1 | 0 | // Valid at all levels +//! | - | 3 | // Null at outer level +//! | - | 2 | // Null at middle level +//! | - | 1 | // Null at inner level +//! +//! # Compression +//! +//! Note that we only need 2 bits of definition levels to represent 3 levels of nesting. Definition +//! levels are always more compact than the input validity arrays. However, compressed levels are not +//! necessarily more compact than the compressed validity arrays. +//! +//! Repetition levels are more complex. If there are very large lists then a sparse array of offsets +//! (which has one element per list) might be more compact than a dense array of repetition levels +//! (which has one element per list value, possibly even more if there are empty lists). +//! +//! However, both repetition levels and definition levels are typically very compressible with RLE. +//! +//! However, in Lance we don't always take advantage of that compression because we want to be able +//! to zip rep-def levels together with our values. This gives us fewer IOPS when accessing row values. +//! +//! # Utilities in this Module +//! +//! - `RepDefBuilder` - Extracts validity and offset information from Arrow arrays. We use this as we +//! shred the incoming data into primitive leaf arrays. We don't immediately convert into rep-def because +//! we need to share and cheaply clone the builder when we have structs (each struct child shares some parent +//! validity / offset information) The `serialize` method is called once all data has been received to create +//! the final rep-def levels. +//! +//! - `SerializerContext` - This is an internal utility that helps with serializing rep-def levels. +//! +//! - `CompositeRepDefUnraveler` - This structure is used to reverse the process. It starts with a set of +//! rep-def levels and then uses the `unravel_validity` and `unravel_offsets` methods to produce validity +//! buffers and offset buffers. It is "composite" because we may be combining sets of rep-def buffers from +//! multiple locations (e.g. multiple blocks in a mini-block encoded file). +//! +//! - `RepDefSlicer` - This is a utility that helps with slicing rep-def buffers. These buffers are "kind of" +//! transparent (maps 1:1 with the values in the array) but not exactly because of the special (empty/null) lists. +//! The slicer helps with this issue and is used when slicing a rep-def buffer into mini-blocks. +//! +//! - `build_control_word_iterator` - This takes in rep-def levels and returns an iterator that returns byte-padded +//! "control words" which are used when creating full-zip encoded data. +//! +//! - `ControlWordParser` - This parser can parse the control words returned by `build_control_word_iterator` and is +//! used when decoding full-zip encoded data. + +use std::{ + iter::{Copied, Zip}, + ops::Range, + sync::Arc, +}; + +use arrow_array::OffsetSizeTrait; +use arrow_buffer::{ + ArrowNativeType, BooleanBuffer, BooleanBufferBuilder, NullBuffer, OffsetBuffer, ScalarBuffer, +}; +use lance_core::{Error, Result, utils::bit::log_2_ceil}; + +use crate::{ + buffer::LanceBuffer, + encodings::logical::primitive::sparse::{SparseStructuralPlan, SparseStructuralUnraveler}, +}; + +pub type LevelBuffer = Vec; + +/// A top-level-row range whose dense rep/def stream fits one mini-block page. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct MiniBlockRepDefSplit { + /// Top-level row offset, relative to the original unsplit page. + pub(crate) row_start: u64, + /// Number of top-level rows in this split. + pub(crate) num_rows: u64, + /// Rep/def level range, relative to the original unsplit page. + pub(crate) level_range: Range, + /// Visible value offset, relative to the original unsplit page. + pub(crate) value_start: u64, + /// Number of visible values in this split. + pub(crate) num_values: u64, +} + +/// Dense mini-block rep/def budget result for one accumulated page. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum MiniBlockRepDefBudget { + /// The dense rep/def stream fits one mini-block structural page. + WithinBudget, + /// The dense rep/def stream fits after splitting on top-level row boundaries. + RequiresPageSplit(Vec), + /// A single top-level row has this many rep/def levels and exceeds the budget. + SingleRowOverBudget(u64), +} + +// As we build def levels we add this to special values to indicate that they +// are special so that we can skip over them when processing lower levels. +// +// We assume 16 bits is good enough for rep-def levels. This _would_ give +// us 65536 levels of struct nesting and list nesting. However, we cut that +// in half for SPECIAL_THRESHOLD because we use the top bit to indicate if an +// item is a special value (null list / empty list) during construction. +// +// We subtract this off at the end of construction to get the actual definition +// levels. +const SPECIAL_THRESHOLD: u16 = u16::MAX / 2; + +/// Represents information that we extract from a list array as we are +/// encoding +#[derive(Clone, Debug)] +struct OffsetDesc { + offsets: Arc<[i64]>, + validity: Option, + has_empty_lists: bool, + num_values: usize, + num_specials: usize, +} + +/// Represents validity information that we extract from non-list arrays (that +/// have nulls) as we are encoding +#[derive(Clone, Debug)] +struct ValidityDesc { + validity: Option, + num_values: usize, +} + +/// Represents validity information that we extract from FSL arrays. This is +/// just validity (no offsets) but we also record the dimension of the FSL array +/// as that will impact the next layer +#[derive(Clone, Debug)] +struct FslDesc { + validity: Option, + dimension: usize, + num_values: usize, +} + +// As we build up rep/def from arrow arrays we record a +// series of RawRepDef objects. Each one corresponds to layer +// in the array structure +#[derive(Clone, Debug)] +enum RawRepDef { + Offsets(OffsetDesc), + Validity(ValidityDesc), + Fsl(FslDesc), +} + +/// A normalized Arrow structural layer shared by dense and sparse serializers. +#[derive(Clone, Copy, Debug)] +pub(crate) enum NormalizedStructuralLayer<'a> { + List { + offsets: &'a [i64], + validity: Option<&'a BooleanBuffer>, + num_slots: usize, + }, + Validity { + validity: Option<&'a BooleanBuffer>, + num_slots: usize, + }, + FixedSizeList { + validity: Option<&'a BooleanBuffer>, + dimension: usize, + num_slots: usize, + }, +} + +/// Structural layers concatenated across input batches exactly once. +/// +/// Dense rep/def serialization and sparse metadata planning both consume this +/// representation so the Arrow nesting is not independently reconstructed. +#[derive(Debug)] +pub(crate) struct NormalizedStructuralPlan { + layers: Vec, + dense_all_valid: bool, +} + +impl NormalizedStructuralPlan { + pub(crate) fn layers(&self) -> impl ExactSizeIterator> { + self.layers.iter().map(|layer| match layer { + RawRepDef::Offsets(OffsetDesc { + offsets, + validity, + num_values, + .. + }) => NormalizedStructuralLayer::List { + offsets, + validity: validity.as_ref(), + num_slots: *num_values, + }, + RawRepDef::Validity(ValidityDesc { + validity, + num_values, + }) => NormalizedStructuralLayer::Validity { + validity: validity.as_ref(), + num_slots: *num_values, + }, + RawRepDef::Fsl(FslDesc { + validity, + dimension, + num_values, + }) => NormalizedStructuralLayer::FixedSizeList { + validity: validity.as_ref(), + dimension: *dimension, + num_slots: *num_values, + }, + }) + } + + fn to_serializer(&self) -> (SerializerContext, Option) { + if self.dense_all_valid { + let def_meaning = self + .layers + .iter() + .map(|_| DefinitionInterpretation::AllValidItem) + .collect::>(); + return ( + SerializerContext { + def_meaning, + rep_levels: LevelBuffer::default(), + spare_rep: LevelBuffer::default(), + def_levels: LevelBuffer::default(), + spare_def: LevelBuffer::default(), + current_rep: 0, + current_def: 0, + current_len: 0, + current_num_specials: 0, + has_fsl: false, + }, + None, + ); + } + + let total_len = self.layers.last().map_or(0, RawRepDef::num_values) + + self + .layers + .iter() + .map(RawRepDef::num_specials) + .sum::(); + let max_rep = self.layers.iter().map(RawRepDef::max_rep).sum::(); + let max_def = self.layers.iter().map(RawRepDef::max_def).sum::(); + let bits_per_rep = if max_rep > 0 { + u64::from(u16::BITS - max_rep.leading_zeros()) + } else { + 0 + }; + let bits_per_def = if max_def > 0 { + u64::from(u16::BITS - max_def.leading_zeros()) + } else { + 0 + }; + let bits_per_level = + (bits_per_rep + bits_per_def > 0).then_some(bits_per_rep + bits_per_def); + + let num_layers = self.layers.len(); + let mut context = SerializerContext::new(total_len, num_layers, max_rep, max_def); + for layer in &self.layers { + match layer { + RawRepDef::Validity(def) => context.record_validity(def), + RawRepDef::Offsets(rep) => context.record_offsets(rep), + RawRepDef::Fsl(fsl) => context.record_fsl(fsl), + } + } + (context, bits_per_level) + } + + pub(crate) fn serialize(&self) -> SerializedRepDefs { + self.to_serializer().0.build() + } + + pub(crate) fn serialize_with_miniblock_repdef_budget( + &self, + max_levels_for_bits: impl FnOnce(u64) -> u64, + num_rows: u64, + num_values: u64, + ) -> Result<(SerializedRepDefs, MiniBlockRepDefBudget)> { + let (context, bits_per_level) = self.to_serializer(); + context.build_with_miniblock_repdef_budget( + bits_per_level.map(max_levels_for_bits), + num_rows, + num_values, + ) + } +} + +impl RawRepDef { + // Are there any nulls in this layer + fn has_nulls(&self) -> bool { + match self { + Self::Offsets(OffsetDesc { validity, .. }) => validity.is_some(), + Self::Validity(ValidityDesc { validity, .. }) => validity.is_some(), + Self::Fsl(FslDesc { validity, .. }) => validity.is_some(), + } + } + + // How many values are in this layer + fn num_values(&self) -> usize { + match self { + Self::Offsets(OffsetDesc { num_values, .. }) => *num_values, + Self::Validity(ValidityDesc { num_values, .. }) => *num_values, + Self::Fsl(FslDesc { num_values, .. }) => *num_values, + } + } + + /// How many empty/null lists are in this layer + fn num_specials(&self) -> usize { + match self { + Self::Offsets(OffsetDesc { num_specials, .. }) => *num_specials, + _ => 0, + } + } + + /// How many definition levels do we need for this layer + fn max_def(&self) -> u16 { + match self { + Self::Offsets(OffsetDesc { + has_empty_lists, + validity, + .. + }) => { + let mut max_def = 0; + if *has_empty_lists { + max_def += 1; + } + if validity.is_some() { + max_def += 1; + } + max_def + } + Self::Validity(ValidityDesc { validity: None, .. }) => 0, + Self::Validity(ValidityDesc { .. }) => 1, + Self::Fsl(FslDesc { validity: None, .. }) => 0, + Self::Fsl(FslDesc { .. }) => 1, + } + } + + /// How many repetition levels do we need for this layer + fn max_rep(&self) -> u16 { + match self { + Self::Offsets(_) => 1, + _ => 0, + } + } +} + +/// Represents repetition and definition levels that have been +/// serialized into a pair of (optional) level buffers +#[derive(Debug)] +pub struct SerializedRepDefs { + /// The repetition levels, one per item + /// + /// If None, there are no lists + pub repetition_levels: Option>, + /// The definition levels, one per item + /// + /// If None, there are no nulls + pub definition_levels: Option>, + /// The meaning of each definition level + pub def_meaning: Vec, + /// The maximum level that is "visible" from the lowest level + /// + /// This is the last level before we encounter a list level of some kind. Once we've + /// hit a list level then nulls in any level beyond do not map to actual items. + /// + /// This is None if there are no lists + pub max_visible_level: Option, + has_fsl: bool, +} + +/// The highest definition level whose items still map to a value in the leaf array +/// +/// `def_meaning` is ordered leaf-first, so every layer ahead of the first list layer sits +/// underneath that list. Once a list boundary is crossed, a null (or empty) list at any +/// outer layer has no value in the leaf at all and its item is "invisible". +/// +/// Both the writer and every decoder decide whether an item owns a value by comparing +/// `def <= max_visible_level`, so the two sides must derive it the same way. When there +/// are no lists this sums every layer, i.e. `max_def`, which makes all items visible. +pub fn max_visible_level(def_meaning: &[DefinitionInterpretation]) -> u16 { + def_meaning + .iter() + .take_while(|level| !level.is_list()) + .map(|level| level.num_def_levels()) + .sum() +} + +impl SerializedRepDefs { + /// The same rule as [`max_visible_level`], but `None` when there are no lists and so + /// no item can be invisible in the first place. + fn max_visible_level(def_meaning: &[DefinitionInterpretation]) -> Option { + def_meaning + .iter() + .any(|level| level.is_list()) + .then(|| max_visible_level(def_meaning)) + } + + pub fn new( + repetition_levels: Option, + definition_levels: Option, + def_meaning: Vec, + ) -> Self { + Self::new_with_fixed_size_list_levels( + repetition_levels, + definition_levels, + def_meaning, + false, + ) + } + + pub(crate) fn new_with_fixed_size_list_levels( + repetition_levels: Option, + definition_levels: Option, + def_meaning: Vec, + has_fsl: bool, + ) -> Self { + let max_visible_level = Self::max_visible_level(&def_meaning); + Self { + repetition_levels: repetition_levels.map(Arc::from), + definition_levels: definition_levels.map(Arc::from), + def_meaning, + max_visible_level, + has_fsl, + } + } + + /// Creates an empty SerializedRepDefs (no repetition, all valid) + pub fn empty(def_meaning: Vec) -> Self { + Self { + repetition_levels: None, + definition_levels: None, + def_meaning, + max_visible_level: None, + has_fsl: false, + } + } + + pub fn rep_slicer(&self) -> Option> { + self.repetition_levels + .as_ref() + .map(|rep| RepDefSlicer::new(self, rep.clone())) + } + + pub fn def_slicer(&self) -> Option> { + self.definition_levels + .as_ref() + .map(|def| RepDefSlicer::new(self, def.clone())) + } + + pub(crate) fn has_fixed_size_list_levels(&self) -> bool { + self.has_fsl + } +} + +/// Slices a level buffer into pieces +/// +/// This is needed to handle the fact that a level buffer may have more +/// levels than values due to special (empty/null) lists. +/// +/// As a result, a call to `slice_next(10)` may return 10 levels or it may +/// return more than 10 levels if any special values are encountered. +#[derive(Debug)] +pub struct RepDefSlicer<'a> { + repdef: &'a SerializedRepDefs, + to_slice: LanceBuffer, + current: usize, +} + +// TODO: All of this logic will need some changing when we compress rep/def levels. +impl<'a> RepDefSlicer<'a> { + fn new(repdef: &'a SerializedRepDefs, levels: Arc<[u16]>) -> Self { + Self { + repdef, + to_slice: LanceBuffer::reinterpret_slice(levels), + current: 0, + } + } + + pub fn num_levels(&self) -> usize { + self.to_slice.len() / 2 + } + + pub fn num_levels_remaining(&self) -> usize { + self.num_levels() - self.current + } + + pub fn all_levels(&self) -> &LanceBuffer { + &self.to_slice + } + + /// Returns the rest of the levels not yet sliced + /// + /// This must be called instead of `slice_next` on the final iteration. + /// This is because anytime we slice there may be empty/null lists on the + /// boundary that are "free" and the current behavior in `slice_next` is to + /// leave them for the next call. + /// + /// `slice_rest` will slice all remaining levels and return them. + pub fn slice_rest(&mut self) -> LanceBuffer { + let start = self.current; + let remaining = self.num_levels_remaining(); + self.current = self.num_levels(); + self.to_slice.slice_with_length(start * 2, remaining * 2) + } + + /// Returns enough levels to satisfy the next `num_values` values + pub fn slice_next(&mut self, num_values: usize) -> LanceBuffer { + let start = self.current; + let Some(max_visible_level) = self.repdef.max_visible_level else { + // No lists, should be 1:1 mapping from levels to values + self.current = start + num_values; + return self.to_slice.slice_with_length(start * 2, num_values * 2); + }; + if let Some(def) = self.repdef.definition_levels.as_ref() { + // There are lists and there are def levels. That means there may be + // more rep/def levels than values. We need to scan the def levels to figure + // out which items are "invisible" and skip over them + let mut def_itr = def[start..].iter(); + let mut num_taken = 0; + let mut num_passed = 0; + while num_taken < num_values { + let def_level = *def_itr.next().unwrap(); + if def_level <= max_visible_level { + num_taken += 1; + } + num_passed += 1; + } + self.current = start + num_passed; + self.to_slice.slice_with_length(start * 2, num_passed * 2) + } else { + // No def levels, should be 1:1 mapping from levels to values + self.current = start + num_values; + self.to_slice.slice_with_length(start * 2, num_values * 2) + } + } +} + +/// This tells us how an array handles definition. Given a stack of +/// these and a nested array and a set of definition levels we can calculate +/// how we should interpret the definition levels. +/// +/// For example, if the interpretation is [AllValidItem, NullableItem] then +/// a 0 means "valid item" and a 1 means "null struct". If the interpretation +/// is [NullableItem, NullableItem] then a 0 means "valid item" and a 1 means +/// "null item" and a 2 means "null struct". +/// +/// Lists are tricky because we might use up to two definition levels for a +/// single layer of list nesting because we need one value to indicate "empty list" +/// and another value to indicate "null list". +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum DefinitionInterpretation { + AllValidItem, + AllValidList, + NullableItem, + NullableList, + EmptyableList, + NullableAndEmptyableList, +} + +impl DefinitionInterpretation { + /// How many definition levels do we need for this layer + pub fn num_def_levels(&self) -> u16 { + match self { + Self::AllValidItem => 0, + Self::AllValidList => 0, + Self::NullableItem => 1, + Self::NullableList => 1, + Self::EmptyableList => 1, + Self::NullableAndEmptyableList => 2, + } + } + + /// Does this layer have nulls? + pub fn is_all_valid(&self) -> bool { + matches!( + self, + Self::AllValidItem | Self::AllValidList | Self::EmptyableList + ) + } + + /// Does this layer represent a list? + pub fn is_list(&self) -> bool { + matches!( + self, + Self::AllValidList + | Self::NullableList + | Self::EmptyableList + | Self::NullableAndEmptyableList + ) + } +} + +/// The RepDefBuilder is used to collect offsets & validity buffers +/// from arrow structures. Once we have those we use the SerializerContext +/// to build the actual repetition and definition levels. +/// +/// We know ahead of time how many rep/def levels we will need (number of items +/// in inner-most array + the number of empty/null lists in any parent arrays). +/// +/// As a result we try and avoid any re-allocations by pre-allocating the buffers +/// up front. We allocate two copies of each buffer which allows us to avoid unsafe +/// code caused by reading and writing to the same buffer (also, it's unavoidable +/// because there are times we need to write 'faster' than we read) +#[derive(Debug)] +struct SerializerContext { + // This is built from outer-to-inner and then reversed at the end + def_meaning: Vec, + rep_levels: LevelBuffer, + spare_rep: LevelBuffer, + def_levels: LevelBuffer, + spare_def: LevelBuffer, + current_rep: u16, + current_def: u16, + current_len: usize, + current_num_specials: usize, + has_fsl: bool, +} + +impl SerializerContext { + fn new(len: usize, num_layers: usize, max_rep: u16, max_def: u16) -> Self { + let def_meaning = Vec::with_capacity(num_layers); + Self { + rep_levels: if max_rep > 0 { + vec![0; len] + } else { + LevelBuffer::default() + }, + spare_rep: if max_rep > 0 { + vec![0; len] + } else { + LevelBuffer::default() + }, + def_levels: if max_def > 0 { + vec![0; len] + } else { + LevelBuffer::default() + }, + spare_def: if max_def > 0 { + vec![0; len] + } else { + LevelBuffer::default() + }, + def_meaning, + current_rep: max_rep, + current_def: max_def, + current_len: 0, + current_num_specials: 0, + has_fsl: false, + } + } + + fn checkout_def(&mut self, meaning: DefinitionInterpretation) -> u16 { + let def = self.current_def; + self.current_def -= meaning.num_def_levels(); + self.def_meaning.push(meaning); + def + } + + fn record_offsets(&mut self, offset_desc: &OffsetDesc) { + let rep_level = self.current_rep; + let (null_list_level, empty_list_level) = + match (offset_desc.validity.is_some(), offset_desc.has_empty_lists) { + (true, true) => { + let level = + self.checkout_def(DefinitionInterpretation::NullableAndEmptyableList); + (level - 1, level) + } + (true, false) => (self.checkout_def(DefinitionInterpretation::NullableList), 0), + (false, true) => ( + 0, + self.checkout_def(DefinitionInterpretation::EmptyableList), + ), + (false, false) => { + self.checkout_def(DefinitionInterpretation::AllValidList); + (0, 0) + } + }; + self.current_rep -= 1; + + if let Some(validity) = &offset_desc.validity { + self.do_record_validity(validity, null_list_level); + } + + // We write into the spare buffers and read from the active buffers + // and then swap at the end. This way we don't write over what we + // are reading. + + let mut new_len = 0; + let expected_len = offset_desc.num_values + self.current_num_specials; + if expected_len == 0 { + // Offsets [0] mean no list values, so no levels. + self.current_len = 0; + return; + } + assert!(self.rep_levels.len() >= expected_len - 1); + if self.def_levels.is_empty() { + let mut write_itr = self.spare_rep.iter_mut(); + let mut read_iter = self.rep_levels.iter().copied(); + for w in offset_desc.offsets.windows(2) { + let len = w[1] - w[0]; + // len can't be 0 because then we'd have def levels + assert!(len > 0); + let rep = read_iter.next().unwrap(); + let list_level = if rep == 0 { rep_level } else { rep }; + *write_itr.next().unwrap() = list_level; + + for _ in 1..len { + *write_itr.next().unwrap() = 0; + } + new_len += len as usize; + } + std::mem::swap(&mut self.rep_levels, &mut self.spare_rep); + } else { + assert!(self.def_levels.len() >= expected_len - 1); + let mut def_write_itr = self.spare_def.iter_mut(); + let mut rep_write_itr = self.spare_rep.iter_mut(); + let mut rep_read_itr = self.rep_levels.iter().copied(); + let mut def_read_itr = self.def_levels.iter().copied(); + let specials_to_pass = self.current_num_specials; + let mut specials_passed = 0; + + for w in offset_desc.offsets.windows(2) { + let mut def = def_read_itr.next().unwrap(); + // Copy over any higher-level special values in place + while def > SPECIAL_THRESHOLD { + *def_write_itr.next().unwrap() = def; + *rep_write_itr.next().unwrap() = rep_read_itr.next().unwrap(); + def = def_read_itr.next().unwrap(); + new_len += 1; + specials_passed += 1; + } + + let len = w[1] - w[0]; + let rep = rep_read_itr.next().unwrap(); + + // If the rep_level is 0 then we are the first list level + // otherwise we are starting a higher level list so keep + // existing rep level + let list_level = if rep == 0 { rep_level } else { rep }; + + if def == 0 && len > 0 { + // New valid list, write a rep level and then add new 0/0 items + *def_write_itr.next().unwrap() = 0; + *rep_write_itr.next().unwrap() = list_level; + + for _ in 1..len { + *def_write_itr.next().unwrap() = 0; + *rep_write_itr.next().unwrap() = 0; + } + + new_len += len as usize; + } else if def == 0 { + // Empty list, insert new special + *def_write_itr.next().unwrap() = empty_list_level + SPECIAL_THRESHOLD; + *rep_write_itr.next().unwrap() = list_level; + new_len += 1; + } else { + // Either the list is null or one of its struct parents + // is null. Promote it to a special value. + *def_write_itr.next().unwrap() = def + SPECIAL_THRESHOLD; + *rep_write_itr.next().unwrap() = list_level; + new_len += 1; + } + } + + // If we have any special values at the end, we need to copy them over + while specials_passed < specials_to_pass { + *def_write_itr.next().unwrap() = def_read_itr.next().unwrap(); + *rep_write_itr.next().unwrap() = rep_read_itr.next().unwrap(); + new_len += 1; + specials_passed += 1; + } + std::mem::swap(&mut self.def_levels, &mut self.spare_def); + std::mem::swap(&mut self.rep_levels, &mut self.spare_rep); + } + + self.current_len = new_len; + self.current_num_specials += offset_desc.num_specials; + } + + fn do_record_validity(&mut self, validity: &BooleanBuffer, null_level: u16) { + assert!(self.def_levels.len() >= validity.len() + self.current_num_specials); + debug_assert!( + self.current_len == 0 || self.current_len == validity.len() + self.current_num_specials + ); + self.current_len = validity.len() + self.current_num_specials; + + let mut def_read_itr = self.def_levels.iter().copied(); + let mut def_write_itr = self.spare_def.iter_mut(); + + let specials_to_pass = self.current_num_specials; + let mut specials_passed = 0; + + for incoming_validity in validity.iter() { + let mut def = def_read_itr.next().unwrap(); + while def > SPECIAL_THRESHOLD { + *def_write_itr.next().unwrap() = def; + def = def_read_itr.next().unwrap(); + specials_passed += 1; + } + if def == 0 && !incoming_validity { + *def_write_itr.next().unwrap() = null_level; + } else { + *def_write_itr.next().unwrap() = def; + } + } + + while specials_passed < specials_to_pass { + *def_write_itr.next().unwrap() = def_read_itr.next().unwrap(); + specials_passed += 1; + } + + std::mem::swap(&mut self.def_levels, &mut self.spare_def); + } + + fn multiply_levels(&mut self, multiplier: usize) { + let old_len = self.current_len; + // All non-special values will be broadcasted by the multiplier. Special values are copied as-is. + self.current_len = + (self.current_len - self.current_num_specials) * multiplier + self.current_num_specials; + + if self.rep_levels.is_empty() && self.def_levels.is_empty() { + // All valid with no rep/def levels, nothing to do + return; + } else if self.rep_levels.is_empty() { + assert!(self.def_levels.len() >= self.current_len); + // No rep levels, just multiply the def levels + let mut def_read_itr = self.def_levels.iter().copied(); + let mut def_write_itr = self.spare_def.iter_mut(); + for _ in 0..old_len { + let mut def = def_read_itr.next().unwrap(); + while def > SPECIAL_THRESHOLD { + *def_write_itr.next().unwrap() = def; + def = def_read_itr.next().unwrap(); + } + for _ in 0..multiplier { + *def_write_itr.next().unwrap() = def; + } + } + } else if self.def_levels.is_empty() { + assert!(self.rep_levels.len() >= self.current_len); + // No def levels, just multiply the rep levels + let mut rep_read_itr = self.rep_levels.iter().copied(); + let mut rep_write_itr = self.spare_rep.iter_mut(); + for _ in 0..old_len { + let rep = rep_read_itr.next().unwrap(); + for _ in 0..multiplier { + *rep_write_itr.next().unwrap() = rep; + } + } + } else { + assert!(self.rep_levels.len() >= self.current_len); + assert!(self.def_levels.len() >= self.current_len); + let mut rep_read_itr = self.rep_levels.iter().copied(); + let mut def_read_itr = self.def_levels.iter().copied(); + let mut rep_write_itr = self.spare_rep.iter_mut(); + let mut def_write_itr = self.spare_def.iter_mut(); + for _ in 0..old_len { + let mut def = def_read_itr.next().unwrap(); + while def > SPECIAL_THRESHOLD { + *def_write_itr.next().unwrap() = def; + *rep_write_itr.next().unwrap() = rep_read_itr.next().unwrap(); + def = def_read_itr.next().unwrap(); + } + let rep = rep_read_itr.next().unwrap(); + for _ in 0..multiplier { + *def_write_itr.next().unwrap() = def; + *rep_write_itr.next().unwrap() = rep; + } + } + } + std::mem::swap(&mut self.def_levels, &mut self.spare_def); + std::mem::swap(&mut self.rep_levels, &mut self.spare_rep); + } + + fn record_validity_buf(&mut self, validity: &Option) { + if let Some(validity) = validity { + let def_level = self.checkout_def(DefinitionInterpretation::NullableItem); + self.do_record_validity(validity, def_level); + } else { + self.checkout_def(DefinitionInterpretation::AllValidItem); + } + } + + fn record_validity(&mut self, validity_desc: &ValidityDesc) { + self.record_validity_buf(&validity_desc.validity) + } + + fn record_fsl(&mut self, fsl_desc: &FslDesc) { + self.has_fsl = true; + self.record_validity_buf(&fsl_desc.validity); + self.multiply_levels(fsl_desc.dimension); + } + + fn normalize_specials(&mut self) { + for def in self.def_levels.iter_mut() { + if *def > SPECIAL_THRESHOLD { + *def -= SPECIAL_THRESHOLD; + } + } + } + + fn normalize_specials_and_plan_splits( + &mut self, + def_meaning: &[DefinitionInterpretation], + max_levels_per_page: Option, + num_rows: u64, + num_values: u64, + ) -> Result { + // Extremely sparse lists can have many rep/def levels for very few + // visible leaf values. If this ratio becomes too skewed then a + // mini-block rep/def chunk can exceed its packed metadata budget even + // though the value buffers are small. We detect that case while + // normalizing special def levels and split on top-level row boundaries + // so each emitted dense mini-block page stays within the budget. + if self.def_levels.is_empty() { + return Ok(MiniBlockRepDefBudget::WithinBudget); + } + + if self.rep_levels.is_empty() { + self.normalize_specials(); + return Ok(MiniBlockRepDefBudget::WithinBudget); + } + + if self.rep_levels.len() != self.def_levels.len() { + return Err(Error::internal(format!( + "Cannot plan structural page splits with mismatched rep/def lengths: rep={}, def={}", + self.rep_levels.len(), + self.def_levels.len() + ))); + } + + let Some(max_levels_per_page) = max_levels_per_page else { + self.normalize_specials(); + return Ok(MiniBlockRepDefBudget::WithinBudget); + }; + + if num_values == 0 { + self.normalize_specials(); + return Ok(MiniBlockRepDefBudget::WithinBudget); + } + + let max_schema_rep = def_meaning.iter().filter(|level| level.is_list()).count() as u16; + let max_visible_level = SerializedRepDefs::max_visible_level(def_meaning); + let should_plan = !self.has_fsl && max_schema_rep > 0 && max_visible_level.is_some(); + + if !should_plan { + self.normalize_specials(); + return Ok(MiniBlockRepDefBudget::WithinBudget); + } + + let max_visible_level = max_visible_level.unwrap(); + let mut splits = Vec::new(); + let mut counted_rows = 0u64; + let mut counted_values = 0u64; + let mut saw_structural_overhead = false; + let mut single_row_over_budget_levels = None; + + let mut current_row_level_start = None; + let mut current_row_num_values = 0u64; + + let mut current_page_row_start = 0u64; + let mut current_page_num_rows = 0u64; + let mut current_page_level_start = 0usize; + let mut current_page_level_end = 0usize; + let mut current_page_value_start = 0u64; + let mut current_page_num_values = 0u64; + let mut current_page_num_levels = 0u64; + let mut current_page_has_structural_overhead = false; + + let mut finish_row = + |row_level_start: usize, row_level_end: usize, row_num_values: u64| -> Result<()> { + let row_num_levels = (row_level_end - row_level_start) as u64; + let row_has_structural_overhead = row_num_levels > row_num_values; + saw_structural_overhead |= row_has_structural_overhead; + + if row_has_structural_overhead && row_num_levels > max_levels_per_page { + single_row_over_budget_levels = Some(row_num_levels); + } + + if current_page_num_rows > 0 + && (current_page_has_structural_overhead || row_has_structural_overhead) + && current_page_num_levels + row_num_levels > max_levels_per_page + { + splits.push(MiniBlockRepDefSplit { + row_start: current_page_row_start, + num_rows: current_page_num_rows, + level_range: current_page_level_start..current_page_level_end, + value_start: current_page_value_start, + num_values: current_page_num_values, + }); + current_page_row_start = counted_rows; + current_page_num_rows = 0; + current_page_level_start = row_level_start; + current_page_value_start = counted_values; + current_page_num_values = 0; + current_page_num_levels = 0; + current_page_has_structural_overhead = false; + } + + if current_page_num_rows == 0 { + current_page_level_start = row_level_start; + } + current_page_num_rows += 1; + current_page_level_end = row_level_end; + current_page_num_values += row_num_values; + current_page_num_levels += row_num_levels; + current_page_has_structural_overhead |= row_has_structural_overhead; + counted_rows += 1; + counted_values += row_num_values; + Ok(()) + }; + + for (idx, (rep_level, def_level)) in self + .rep_levels + .iter() + .copied() + .zip(self.def_levels.iter_mut()) + .enumerate() + { + if *def_level > SPECIAL_THRESHOLD { + *def_level -= SPECIAL_THRESHOLD; + } + + if rep_level == max_schema_rep { + if let Some(level_start) = current_row_level_start { + finish_row(level_start, idx, current_row_num_values)?; + current_row_num_values = 0; + } else if idx != 0 { + return Err(Error::internal(format!( + "Cannot plan structural page splits: first top-level row starts at level {}, expected 0", + idx + ))); + } + current_row_level_start = Some(idx); + } + + if current_row_level_start.is_none() { + return Err(Error::internal( + "Cannot plan structural page splits: found levels before the first top-level row start", + )); + } + if *def_level <= max_visible_level { + current_row_num_values += 1; + } + } + + let Some(level_start) = current_row_level_start else { + return Err(Error::internal( + "Cannot plan structural page splits: found no top-level row starts", + )); + }; + finish_row(level_start, self.rep_levels.len(), current_row_num_values)?; + + if counted_rows != num_rows { + return Err(Error::internal(format!( + "Cannot plan structural page splits: expected {} top-level row starts, found {}", + num_rows, counted_rows + ))); + } + if counted_values != num_values { + return Err(Error::internal(format!( + "Cannot plan structural page splits: counted {} visible values, expected {}", + counted_values, num_values + ))); + } + if !saw_structural_overhead { + return Ok(MiniBlockRepDefBudget::WithinBudget); + } + if let Some(row_num_levels) = single_row_over_budget_levels { + return Ok(MiniBlockRepDefBudget::SingleRowOverBudget(row_num_levels)); + } + + if current_page_num_rows > 0 { + splits.push(MiniBlockRepDefSplit { + row_start: current_page_row_start, + num_rows: current_page_num_rows, + level_range: current_page_level_start..current_page_level_end, + value_start: current_page_value_start, + num_values: current_page_num_values, + }); + } + + if splits.len() > 1 { + Ok(MiniBlockRepDefBudget::RequiresPageSplit(splits)) + } else { + Ok(MiniBlockRepDefBudget::WithinBudget) + } + } + + fn build(mut self) -> SerializedRepDefs { + if self.current_len == 0 { + return SerializedRepDefs::new_with_fixed_size_list_levels( + None, + None, + self.def_meaning, + self.has_fsl, + ); + } + + self.normalize_specials(); + + let definition_levels = if self.def_levels.is_empty() { + None + } else { + Some(self.def_levels) + }; + let repetition_levels = if self.rep_levels.is_empty() { + None + } else { + Some(self.rep_levels) + }; + + // Need to reverse the def meaning since we build rep / def levels in reverse + let def_meaning = self.def_meaning.into_iter().rev().collect::>(); + + SerializedRepDefs::new_with_fixed_size_list_levels( + repetition_levels, + definition_levels, + def_meaning, + self.has_fsl, + ) + } + + fn build_with_miniblock_repdef_budget( + mut self, + max_levels_per_page: Option, + num_rows: u64, + num_values: u64, + ) -> Result<(SerializedRepDefs, MiniBlockRepDefBudget)> { + if self.current_len == 0 { + return Ok(( + SerializedRepDefs::new_with_fixed_size_list_levels( + None, + None, + self.def_meaning, + self.has_fsl, + ), + MiniBlockRepDefBudget::WithinBudget, + )); + } + + // Need to reverse the def meaning since we build rep / def levels in reverse + let def_meaning = std::mem::take(&mut self.def_meaning) + .into_iter() + .rev() + .collect::>(); + let budget = self.normalize_specials_and_plan_splits( + &def_meaning, + max_levels_per_page, + num_rows, + num_values, + )?; + + let definition_levels = if self.def_levels.is_empty() { + None + } else { + Some(self.def_levels) + }; + let repetition_levels = if self.rep_levels.is_empty() { + None + } else { + Some(self.rep_levels) + }; + + Ok(( + SerializedRepDefs::new_with_fixed_size_list_levels( + repetition_levels, + definition_levels, + def_meaning, + self.has_fsl, + ), + budget, + )) + } +} + +/// A structure used to collect validity buffers and offsets from arrow +/// arrays and eventually create repetition and definition levels +/// +/// As we are encoding the structural encoders are given this struct and +/// will record the arrow information into it. Once we hit a leaf node we +/// serialize the data into rep/def levels and write these into the page. +#[derive(Clone, Default, Debug)] +pub struct RepDefBuilder { + // The rep/def info we have collected so far + repdefs: Vec, + // The current length, can get larger as we traverse lists (e.g. an + // array might have 5 lists which results in 50 items) + // + // Starts uninitialized until we see the first rep/def item + len: Option, +} + +impl RepDefBuilder { + fn check_validity_len(&mut self, incoming_len: usize) { + if let Some(len) = self.len { + assert_eq!(incoming_len, len); + } else { + // First validity buffer we've seen + self.len = Some(incoming_len); + } + } + + fn num_layers(&self) -> usize { + self.repdefs.len() + } + + /// The builder is "empty" if there is no repetition and no nulls. In this case we don't need + /// to store anything to disk (except the description) + pub fn is_empty(&self) -> bool { + self.repdefs + .iter() + .all(|r| matches!(r, RawRepDef::Validity(ValidityDesc { validity: None, .. }))) + } + + /// Returns true if there is only a single layer of definition + pub fn is_simple_validity(&self) -> bool { + self.repdefs.len() == 1 && matches!(self.repdefs[0], RawRepDef::Validity(_)) + } + + /// Registers a nullable validity bitmap + pub fn add_validity_bitmap(&mut self, validity: NullBuffer) { + self.check_validity_len(validity.len()); + if validity.null_count() == 0 { + self.add_no_null(validity.len()); + return; + } + self.repdefs.push(RawRepDef::Validity(ValidityDesc { + num_values: validity.len(), + validity: Some(validity.into_inner()), + })); + } + + /// Registers an all-valid validity layer + pub fn add_no_null(&mut self, len: usize) { + self.check_validity_len(len); + self.repdefs.push(RawRepDef::Validity(ValidityDesc { + validity: None, + num_values: len, + })); + } + + pub fn add_fsl(&mut self, validity: Option, dimension: usize, num_values: usize) { + if let Some(len) = self.len { + assert_eq!(num_values, len); + } + self.len = Some(num_values * dimension); + debug_assert!(validity.is_none() || validity.as_ref().unwrap().len() == num_values); + self.repdefs.push(RawRepDef::Fsl(FslDesc { + num_values, + validity: validity.map(|v| v.into_inner()), + dimension, + })) + } + + fn check_offset_len(&mut self, offsets: &[i64]) { + if let Some(len) = self.len { + assert!(offsets.len() == len + 1); + } + self.len = Some(offsets[offsets.len() - 1] as usize); + } + + fn do_add_offsets( + &mut self, + lengths: impl Iterator, + validity: Option, + capacity: usize, + ) -> bool { + let mut num_specials = 0; + let mut has_empty_lists = false; + let mut has_garbage_values = false; + let mut last_off: i64 = 0; + + let mut normalized_offsets = Vec::with_capacity(capacity); + normalized_offsets.push(0); + + if let Some(ref validity) = validity { + for (len, is_valid) in lengths.zip(validity.iter()) { + match (is_valid, len == 0) { + (false, is_empty) => { + num_specials += 1; + has_garbage_values |= !is_empty; + } + (true, true) => { + num_specials += 1; + has_empty_lists = true; + } + _ => { + last_off += len; + } + } + normalized_offsets.push(last_off); + } + } else { + for len in lengths { + if len == 0 { + num_specials += 1; + has_empty_lists = true; + } + last_off += len; + normalized_offsets.push(last_off); + } + } + + self.check_offset_len(&normalized_offsets); + self.repdefs.push(RawRepDef::Offsets(OffsetDesc { + num_values: normalized_offsets.len() - 1, + offsets: normalized_offsets.into(), + validity: validity.map(|v| v.into_inner()), + has_empty_lists, + num_specials: num_specials as usize, + })); + + has_garbage_values + } + + /// Adds a layer of offsets + /// + /// Offsets are casted to a common type (i64) and also normalized. Null lists are + /// always represented by a zero-length (identical) pair of offsets and so the caller + /// should filter out any garbage items before encoding them. To assist with this the + /// method will return true if any non-empty null lists were found. + pub fn add_offsets( + &mut self, + offsets: OffsetBuffer, + validity: Option, + ) -> bool { + let inner = offsets.into_inner(); + let buffer_len = inner.len(); + + if O::IS_LARGE { + let i64_buff = ScalarBuffer::::new(inner.into_inner(), 0, buffer_len); + let lengths = i64_buff.windows(2).map(|off| off[1] - off[0]); + self.do_add_offsets(lengths, validity, buffer_len) + } else { + let i32_buff = ScalarBuffer::::new(inner.into_inner(), 0, buffer_len); + let lengths = i32_buff.windows(2).map(|off| (off[1] - off[0]) as i64); + self.do_add_offsets(lengths, validity, buffer_len) + } + } + + // When we are encoding data it arrives in batches. For each batch we create a RepDefBuilder and collect the + // various validity buffers and offset buffers from that batch. Once we have enough batches to write a page we + // need to take this collection of RepDefBuilders and concatenate them and then serialize them into rep/def levels. + // + // TODO: In the future, we may concatenate and serialize at the same time? + // + // This method takes care of the concatenation part. First we collect all of layer 0 from each builder, then we + // call this method. Then we collect all of layer 1 from each builder and call this method. And so on. + // + // That means this method should get a collection of `RawRepDef` where each item is the same kind (all validity or + // all offsets) though the nullability / lengths may be different in each layer. + fn concat_layers<'a>( + layers: impl Iterator, + num_layers: usize, + ) -> RawRepDef { + enum LayerKind { + Validity, + Fsl, + Offsets, + } + + // We make two passes through the layers. The first determines if we need to pay the cost of allocating + // buffers. The second pass actually adds the values. + let mut collected = Vec::with_capacity(num_layers); + let mut has_nulls = false; + let mut layer_kind = LayerKind::Validity; + let mut total_num_specials = 0; + let mut all_dimension = 0; + let mut all_has_empty_lists = false; + let mut all_num_values = 0; + for layer in layers { + has_nulls |= layer.has_nulls(); + match layer { + RawRepDef::Validity(_) => { + layer_kind = LayerKind::Validity; + } + RawRepDef::Offsets(OffsetDesc { + num_specials, + has_empty_lists, + .. + }) => { + all_has_empty_lists |= *has_empty_lists; + layer_kind = LayerKind::Offsets; + total_num_specials += num_specials; + } + RawRepDef::Fsl(FslDesc { dimension, .. }) => { + layer_kind = LayerKind::Fsl; + all_dimension = *dimension; + } + } + collected.push(layer); + all_num_values += layer.num_values(); + } + + // Shortcut if there are no nulls + if !has_nulls { + match layer_kind { + LayerKind::Validity => { + return RawRepDef::Validity(ValidityDesc { + validity: None, + num_values: all_num_values, + }); + } + LayerKind::Fsl => { + return RawRepDef::Fsl(FslDesc { + validity: None, + num_values: all_num_values, + dimension: all_dimension, + }); + } + LayerKind::Offsets => {} + } + } + + // Only allocate if needed + let mut validity_builder = if has_nulls { + BooleanBufferBuilder::new(all_num_values) + } else { + BooleanBufferBuilder::new(0) + }; + let mut all_offsets = if matches!(layer_kind, LayerKind::Offsets) { + let mut all_offsets = Vec::with_capacity(all_num_values); + all_offsets.push(0); + all_offsets + } else { + Vec::new() + }; + + for layer in collected { + match layer { + RawRepDef::Validity(ValidityDesc { + validity: Some(validity), + .. + }) => { + validity_builder.append_buffer(validity); + } + RawRepDef::Validity(ValidityDesc { + validity: None, + num_values, + }) => { + validity_builder.append_n(*num_values, true); + } + RawRepDef::Fsl(FslDesc { + validity, + num_values, + .. + }) => { + if let Some(validity) = validity { + validity_builder.append_buffer(validity); + } else { + validity_builder.append_n(*num_values, true); + } + } + RawRepDef::Offsets(OffsetDesc { + offsets, + validity: Some(validity), + has_empty_lists, + .. + }) => { + all_has_empty_lists |= has_empty_lists; + validity_builder.append_buffer(validity); + let last = *all_offsets.last().unwrap(); + all_offsets.extend(offsets.iter().skip(1).map(|off| *off + last)); + } + RawRepDef::Offsets(OffsetDesc { + offsets, + validity: None, + has_empty_lists, + num_values, + .. + }) => { + all_has_empty_lists |= has_empty_lists; + if has_nulls { + validity_builder.append_n(*num_values, true); + } + let last = *all_offsets.last().unwrap(); + all_offsets.extend(offsets.iter().skip(1).map(|off| *off + last)); + } + } + } + let validity = if has_nulls { + Some(validity_builder.finish()) + } else { + None + }; + match layer_kind { + LayerKind::Fsl => RawRepDef::Fsl(FslDesc { + validity, + num_values: all_num_values, + dimension: all_dimension, + }), + LayerKind::Validity => RawRepDef::Validity(ValidityDesc { + validity, + num_values: all_num_values, + }), + LayerKind::Offsets => RawRepDef::Offsets(OffsetDesc { + offsets: all_offsets.into(), + validity, + has_empty_lists: all_has_empty_lists, + num_values: all_num_values, + num_specials: total_num_specials, + }), + } + } + + /// Converts the validity / offsets buffers that have been gathered so far + /// into repetition and definition levels + pub fn serialize(builders: Vec) -> SerializedRepDefs { + Self::normalize(builders).serialize() + } + + pub(crate) fn normalize(builders: Vec) -> NormalizedStructuralPlan { + assert!(!builders.is_empty()); + let num_layers = builders[0].num_layers(); + debug_assert!( + builders + .iter() + .all(|builder| builder.num_layers() == num_layers) + ); + let layers = (0..num_layers) + .map(|layer_index| { + Self::concat_layers( + builders.iter().map(|b| &b.repdefs[layer_index]), + builders.len(), + ) + }) + .collect::>(); + NormalizedStructuralPlan { + layers, + dense_all_valid: builders.iter().all(Self::is_empty), + } + } +} + +/// Starts with serialized repetition and definition levels and unravels +/// them into validity buffers and offsets buffers +/// +/// This is used during decoding to create the necessary arrow structures +#[derive(Debug)] +pub struct RepDefUnraveler { + sparse: Option, + rep_levels: Option, + def_levels: Option, + // Maps from definition level to the rep level at which that definition level is visible + levels_to_rep: Vec, + def_meaning: Arc<[DefinitionInterpretation]>, + // Current definition level to compare to. + current_def_cmp: u16, + // Current rep level, determines which specials we can see + current_rep_cmp: u16, + // Current layer index, 0 means inner-most layer and it counts up from there. Used to index + // into special_defs + current_layer: usize, + // Number of items in the inner-most layer (needed if the definition levels are not present) + num_items: u64, +} + +impl RepDefUnraveler { + /// Creates a new unraveler from serialized repetition and definition information + pub fn new( + rep_levels: Option, + def_levels: Option, + def_meaning: Arc<[DefinitionInterpretation]>, + num_items: u64, + ) -> Self { + let mut levels_to_rep = Vec::with_capacity(def_meaning.len()); + let mut rep_counter = 0; + // Level=0 is always visible and means valid item + levels_to_rep.push(0); + for meaning in def_meaning.as_ref() { + match meaning { + DefinitionInterpretation::AllValidItem | DefinitionInterpretation::AllValidList => { + // There is no corresponding level, so nothing to put in levels_to_rep + } + DefinitionInterpretation::NullableItem => { + // Some null structs are not visible at inner rep levels in cases like LIST>> + levels_to_rep.push(rep_counter); + } + DefinitionInterpretation::NullableList => { + rep_counter += 1; + levels_to_rep.push(rep_counter); + } + DefinitionInterpretation::EmptyableList => { + rep_counter += 1; + levels_to_rep.push(rep_counter); + } + DefinitionInterpretation::NullableAndEmptyableList => { + rep_counter += 1; + levels_to_rep.push(rep_counter); + levels_to_rep.push(rep_counter); + } + } + } + Self { + sparse: None, + rep_levels, + def_levels, + current_def_cmp: 0, + current_rep_cmp: 0, + levels_to_rep, + current_layer: 0, + def_meaning, + num_items, + } + } + + pub(crate) fn new_sparse(plan: SparseStructuralPlan) -> Self { + Self { + sparse: Some(SparseStructuralUnraveler::new(plan)), + rep_levels: None, + def_levels: None, + levels_to_rep: Vec::new(), + def_meaning: Arc::new([]), + current_def_cmp: 0, + current_rep_cmp: 0, + current_layer: 0, + num_items: 0, + } + } + + fn ensure_exhausted(&self) -> Result<()> { + if let Some(sparse) = &self.sparse { + sparse.ensure_exhausted()?; + } + Ok(()) + } + + fn is_sparse(&self) -> bool { + self.sparse.is_some() + } + + pub fn is_all_valid(&self) -> bool { + if let Some(sparse) = &self.sparse { + return sparse.is_all_valid(); + } + self.def_levels.is_none() || self.def_meaning[self.current_layer].is_all_valid() + } + + /// If the current level is a repetition layer then this returns the number of lists + /// at this level. + /// + /// This is not valid to call when the current level is a struct/primitive layer because + /// in some cases there may be no rep or def information to know this. + pub fn max_lists(&self) -> Result { + if let Some(sparse) = &self.sparse { + return sparse.max_lists(); + } + debug_assert!( + self.def_meaning[self.current_layer] != DefinitionInterpretation::NullableItem + ); + Ok(self + .rep_levels + .as_ref() + // Worst case every rep item is max_rep and a new list + .map(|levels| levels.len()) + .unwrap_or(0)) + } + + /// Unravels a layer of offsets from the unraveler into the given offset width + /// + /// When decoding a list the caller should first unravel the offsets and then + /// unravel the validity (this is the opposite order used during encoding) + pub fn unravel_offsets( + &mut self, + offsets: &mut Vec, + validity: Option<&mut BooleanBufferBuilder>, + ) -> Result<()> { + if let Some(sparse) = self.sparse.as_mut() { + return sparse.unravel_offsets(offsets, validity); + } + let rep_levels = self + .rep_levels + .as_mut() + .expect("Expected repetition level but data didn't contain repetition"); + let valid_level = self.current_def_cmp; + let (null_level, empty_level) = match self.def_meaning[self.current_layer] { + DefinitionInterpretation::NullableList => { + self.current_def_cmp += 1; + (valid_level + 1, 0) + } + DefinitionInterpretation::EmptyableList => { + self.current_def_cmp += 1; + (0, valid_level + 1) + } + DefinitionInterpretation::NullableAndEmptyableList => { + self.current_def_cmp += 2; + (valid_level + 1, valid_level + 2) + } + DefinitionInterpretation::AllValidList => (0, 0), + _ => unreachable!(), + }; + self.current_layer += 1; + + // This is the highest def level that is still visible. Once we hit a list then + // we stop looking because any null / empty list (or list masked by a higher level + // null) will not be visible + let mut max_level = null_level.max(empty_level).max(valid_level); + // Anything higher than this (but less than max_level) is a null struct masking our + // list. We will materialize this is a null list. + let upper_null = max_level; + for level in self.def_meaning[self.current_layer..].iter() { + match level { + DefinitionInterpretation::NullableItem => { + max_level += 1; + } + DefinitionInterpretation::AllValidItem => {} + _ => { + break; + } + } + } + + let mut curlen: usize = offsets.last().map(|o| o.as_usize()).unwrap_or(0); + + // If offsets is empty this is a no-op. If offsets is not empty that means we already + // added a set of offsets. For example, we might have added [0, 3, 5] (2 lists). Now + // say we want to add [0, 1, 4] (2 lists). We should get [0, 3, 5, 6, 9] (4 lists). If + // we don't pop here we get [0, 3, 5, 5, 6, 9] which is wrong. + // + // Or, to think about it another way, if every unraveler adds the starting 0 and the trailing + // length then we have N + unravelers.len() values instead of N + 1. + offsets.pop(); + + let to_offset = |val: usize| { + T::from_usize(val) + .ok_or_else(|| Error::invalid_input("A single batch had more than i32::MAX values and so a large container type is required")) + }; + self.current_rep_cmp += 1; + if let Some(def_levels) = &mut self.def_levels { + assert!(rep_levels.len() == def_levels.len()); + // It's possible validity is None even if we have def levels. For example, we might have + // empty lists (which require def levels) but no nulls. + let mut push_validity: Box = if let Some(validity) = validity { + Box::new(|is_valid| validity.append(is_valid)) + } else { + Box::new(|_| {}) + }; + // This is a strange access pattern. We are iterating over the rep/def levels and + // at the same time writing the rep/def levels. This means we need both a mutable + // and immutable reference to the rep/def levels. + let mut read_idx = 0; + let mut write_idx = 0; + while read_idx < rep_levels.len() { + // SAFETY: We assert that rep_levels and def_levels have the same + // len and read_idx and write_idx can never go past the end. + unsafe { + let rep_val = *rep_levels.get_unchecked(read_idx); + if rep_val != 0 { + let def_val = *def_levels.get_unchecked(read_idx); + // Copy over + *rep_levels.get_unchecked_mut(write_idx) = rep_val - 1; + *def_levels.get_unchecked_mut(write_idx) = def_val; + write_idx += 1; + + if def_val == 0 { + // This is a valid list + offsets.push(to_offset(curlen)?); + curlen += 1; + push_validity(true); + } else if def_val > max_level { + // This is not visible at this rep level, do not add to offsets, but keep in repdef + } else if def_val == null_level || def_val > upper_null { + // This is a null list (or a list masked by a null struct) + offsets.push(to_offset(curlen)?); + push_validity(false); + } else if def_val == empty_level { + // This is an empty list + offsets.push(to_offset(curlen)?); + push_validity(true); + } else { + // New valid list starting with null item + offsets.push(to_offset(curlen)?); + curlen += 1; + push_validity(true); + } + } else { + curlen += 1; + } + read_idx += 1; + } + } + offsets.push(to_offset(curlen)?); + rep_levels.truncate(write_idx); + def_levels.truncate(write_idx); + Ok(()) + } else { + // SAFETY: See above loop + let mut read_idx = 0; + let mut write_idx = 0; + let old_offsets_len = offsets.len(); + while read_idx < rep_levels.len() { + // SAFETY: read_idx / write_idx cannot go past rep_levels.len() + unsafe { + let rep_val = *rep_levels.get_unchecked(read_idx); + if rep_val != 0 { + // Finish the current list + offsets.push(to_offset(curlen)?); + *rep_levels.get_unchecked_mut(write_idx) = rep_val - 1; + write_idx += 1; + } + curlen += 1; + read_idx += 1; + } + } + let num_new_lists = offsets.len() - old_offsets_len; + offsets.push(to_offset(curlen)?); + // Truncate to the number of lists THIS unraveler produced (write_idx), + // not `offsets.len() - 1` — the latter includes offsets contributed by + // earlier unravelers in a multi-page read, which would leave too many + // rep levels for the next (outer) layer and over-count its lists. + rep_levels.truncate(write_idx); + if let Some(validity) = validity { + // Even though we don't have validity it is possible another unraveler did and so we need + // to push all valids + validity.append_n(num_new_lists, true); + } + Ok(()) + } + } + + pub fn skip_validity(&mut self) -> Result<()> { + if let Some(sparse) = self.sparse.as_mut() { + return sparse.skip_validity(); + } + debug_assert!(self.is_all_valid()); + self.current_layer += 1; + Ok(()) + } + + /// Unravels a layer of validity from the definition levels + pub fn unravel_validity(&mut self, validity: &mut BooleanBufferBuilder) -> Result<()> { + if let Some(sparse) = self.sparse.as_mut() { + return sparse.unravel_validity(validity); + } + let meaning = self.def_meaning[self.current_layer]; + if meaning == DefinitionInterpretation::AllValidItem || self.def_levels.is_none() { + self.current_layer += 1; + validity.append_n(self.num_items as usize, true); + return Ok(()); + } + + self.current_layer += 1; + let def_levels = &self.def_levels.as_ref().unwrap(); + + let current_def_cmp = self.current_def_cmp; + self.current_def_cmp += 1; + + for is_valid in def_levels.iter().filter_map(|&level| { + if self.levels_to_rep[level as usize] <= self.current_rep_cmp { + Some(level <= current_def_cmp) + } else { + None + } + }) { + validity.append(is_valid); + } + Ok(()) + } + + /// Removes all but the first definition level of each fixed-size-list slot + /// + /// The definition levels arrive with one entry per item. A fixed-size-list + /// layer has a single definition level per slot (all `dimension` items in a + /// slot share it) so we keep every `dimension`-th level and drop the rest. + /// + /// `dimension` must be non-zero. A zero dimension can only come from a + /// malformed schema (writers reject it, see + /// [`lance_core::datatypes::validate_fixed_size_list_dimensions`]) and is + /// rejected here rather than allowed to run off the end of the buffer. + pub fn decimate(&mut self, dimension: usize) -> Result<()> { + if dimension == 0 { + return Err(Error::invalid_input( + "Cannot decimate repetition/definition levels with a fixed-size-list dimension of 0; dimension must be a positive integer", + )); + } + if let Some(sparse) = self.sparse.as_mut() { + return sparse.decimate(dimension); + } + if self.rep_levels.is_some() { + // If we need to support this then I think we need to walk through the rep def levels to find + // the spots at which we keep. E.g. if we have: + // rep: 1 0 0 1 0 1 0 0 0 1 0 0 + // def: 1 1 1 0 1 0 1 1 0 1 1 0 + // dimension: 2 + // + // The output should be: + // rep: 1 0 0 1 0 0 0 + // def: 1 1 1 0 1 1 0 + // + // Maybe there's some special logic for empty/null lists? I'll save the headache for future me. + todo!("Not yet supported FSL<...List<...>>"); + } + let Some(def_levels) = self.def_levels.as_mut() else { + return Ok(()); + }; + let mut read_idx = 0; + let mut write_idx = 0; + while read_idx < def_levels.len() { + // SAFETY: `read_idx` is checked against the length by the loop condition and + // `dimension >= 1` (checked above) means `write_idx <= read_idx`, so both + // indices are in bounds. + unsafe { + *def_levels.get_unchecked_mut(write_idx) = *def_levels.get_unchecked(read_idx); + } + write_idx += 1; + read_idx += dimension; + } + def_levels.truncate(write_idx); + Ok(()) + } +} + +/// As we decode we may extract rep/def information from multiple pages (or multiple +/// chunks within a page). +/// +/// For each chunk we create an unraveler. Each unraveler can have a completely different +/// interpretation (e.g. one page might contain null items but no null structs and the next +/// page might have null structs but no null items). +/// +/// Concatenating these unravelers would be tricky and expensive so instead we have a +/// composite unraveler which unravels across multiple unravelers. +/// +/// Note: this class should be used even if there is only one page / unraveler. This is +/// because the `RepDefUnraveler`'s API is more complex (it's meant to be called by this +/// class) +#[derive(Debug)] +pub struct CompositeRepDefUnraveler { + unravelers: Vec, + comparisons: Vec, +} + +impl CompositeRepDefUnraveler { + pub fn new(unravelers: Vec) -> Self { + Self { + unravelers, + comparisons: Vec::new(), + } + } + + pub(crate) fn add_compatibility_check(&mut self, other: Self) { + self.comparisons.push(other); + } + + pub(crate) fn has_sparse(&self) -> bool { + self.unravelers.iter().any(RepDefUnraveler::is_sparse) + || self.comparisons.iter().any(Self::has_sparse) + } + + pub(crate) fn ensure_exhausted(&self) -> Result<()> { + for unraveler in &self.unravelers { + unraveler.ensure_exhausted()?; + } + for comparison in &self.comparisons { + comparison.ensure_exhausted()?; + } + Ok(()) + } + + fn null_buffers_equal( + left: &Option, + right: &Option, + expected_len: usize, + ) -> bool { + match (left, right) { + (None, None) => true, + (Some(left), Some(right)) => { + left.len() == expected_len + && right.len() == expected_len + && left.iter().eq(right.iter()) + } + (None, Some(right)) => right.len() == expected_len && right.null_count() == 0, + (Some(left), None) => left.len() == expected_len && left.null_count() == 0, + } + } + + fn decimate(&mut self, dimension: usize) -> Result<()> { + for unraveler in &mut self.unravelers { + unraveler.decimate(dimension)?; + } + for comparison in &mut self.comparisons { + comparison.decimate(dimension)?; + } + Ok(()) + } + + /// Unravels a layer of validity + /// + /// Returns None if there are no null items in this layer + pub fn unravel_validity(&mut self, num_values: usize) -> Result> { + let is_all_valid = self + .unravelers + .iter() + .all(|unraveler| unraveler.is_all_valid()); + + let validity = if is_all_valid { + for unraveler in self.unravelers.iter_mut() { + unraveler.skip_validity()?; + } + None + } else { + let mut validity = BooleanBufferBuilder::new(num_values); + for unraveler in self.unravelers.iter_mut() { + unraveler.unravel_validity(&mut validity)?; + } + Some(NullBuffer::new(validity.finish())) + }; + for comparison in &mut self.comparisons { + let other = comparison.unravel_validity(num_values)?; + if !Self::null_buffers_equal(&validity, &other, num_values) { + return Err(Error::invalid_input_source( + format!( + "Structural sibling fields have incompatible validity metadata for {num_values} values" + ) + .into(), + )); + } + } + Ok(validity) + } + + pub fn unravel_fsl_validity( + &mut self, + num_values: usize, + dimension: usize, + ) -> Result> { + self.decimate(dimension)?; + self.unravel_validity(num_values) + } + + /// Unravels a layer of offsets (and the validity for that layer) + pub fn unravel_offsets( + &mut self, + ) -> Result<(OffsetBuffer, Option)> { + let mut is_all_valid = true; + let mut max_num_lists: usize = 0; + for unraveler in self.unravelers.iter() { + is_all_valid &= unraveler.is_all_valid(); + max_num_lists = max_num_lists + .checked_add(unraveler.max_lists()?) + .ok_or_else(|| { + Error::invalid_input_source( + "Combined repetition/definition list count exceeds usize::MAX".into(), + ) + })?; + } + + let mut validity = if is_all_valid { + None + } else { + // Note: This is probably an over-estimate and potentially even an under-estimate. We only know + // right now how many items we have and not how many rows. (TODO: Shouldn't we know the # of rows?) + Some(BooleanBufferBuilder::new(max_num_lists)) + }; + + let mut offsets = Vec::with_capacity(max_num_lists + 1); + + for unraveler in self.unravelers.iter_mut() { + unraveler.unravel_offsets(&mut offsets, validity.as_mut())?; + } + + let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets)); + let validity = validity.map(|mut v| NullBuffer::new(v.finish())); + for comparison in &mut self.comparisons { + let (other_offsets, other_validity) = comparison.unravel_offsets::()?; + if offsets.as_ref() != other_offsets.as_ref() + || !Self::null_buffers_equal( + &validity, + &other_validity, + offsets.len().saturating_sub(1), + ) + { + return Err(Error::invalid_input_source( + format!( + "Structural sibling fields have incompatible list metadata for {} slots", + offsets.len().saturating_sub(1) + ) + .into(), + )); + } + } + + Ok((offsets, validity)) + } +} + +/// A [`ControlWordIterator`] when there are both repetition and definition levels +/// +/// The iterator will put the repetition level in the upper bits and the definition +/// level in the lower bits. The number of bits used for each level is determined +/// by the width of the repetition and definition levels. +#[derive(Debug)] +pub struct BinaryControlWordIterator, W> { + repdef: I, + def_width: usize, + max_rep: u16, + max_visible_def: u16, + rep_mask: u16, + def_mask: u16, + bits_rep: u8, + bits_def: u8, + phantom: std::marker::PhantomData, +} + +impl> BinaryControlWordIterator { + fn append_next(&mut self, buf: &mut Vec) -> Option { + let next = self.repdef.next()?; + let control_word: u8 = + (((next.0 & self.rep_mask) as u8) << self.def_width) + ((next.1 & self.def_mask) as u8); + buf.push(control_word); + let is_new_row = next.0 == self.max_rep; + let is_visible = next.1 <= self.max_visible_def; + let is_valid_item = next.1 == 0; + Some(ControlWordDesc { + is_new_row, + is_visible, + is_valid_item, + }) + } +} + +impl> BinaryControlWordIterator { + fn append_next(&mut self, buf: &mut Vec) -> Option { + let next = self.repdef.next()?; + let control_word: u16 = + ((next.0 & self.rep_mask) << self.def_width) + (next.1 & self.def_mask); + let control_word = control_word.to_le_bytes(); + buf.push(control_word[0]); + buf.push(control_word[1]); + let is_new_row = next.0 == self.max_rep; + let is_visible = next.1 <= self.max_visible_def; + let is_valid_item = next.1 == 0; + Some(ControlWordDesc { + is_new_row, + is_visible, + is_valid_item, + }) + } +} + +impl> BinaryControlWordIterator { + fn append_next(&mut self, buf: &mut Vec) -> Option { + let next = self.repdef.next()?; + let control_word: u32 = (((next.0 & self.rep_mask) as u32) << self.def_width) + + ((next.1 & self.def_mask) as u32); + let control_word = control_word.to_le_bytes(); + buf.push(control_word[0]); + buf.push(control_word[1]); + buf.push(control_word[2]); + buf.push(control_word[3]); + let is_new_row = next.0 == self.max_rep; + let is_visible = next.1 <= self.max_visible_def; + let is_valid_item = next.1 == 0; + Some(ControlWordDesc { + is_new_row, + is_visible, + is_valid_item, + }) + } +} + +/// A [`ControlWordIterator`] when there are only definition levels or only repetition levels +#[derive(Debug)] +pub struct UnaryControlWordIterator, W> { + repdef: I, + level_mask: u16, + bits_rep: u8, + bits_def: u8, + max_rep: u16, + phantom: std::marker::PhantomData, +} + +impl> UnaryControlWordIterator { + fn append_next(&mut self, buf: &mut Vec) -> Option { + let next = self.repdef.next()?; + buf.push((next & self.level_mask) as u8); + let is_new_row = self.max_rep == 0 || next == self.max_rep; + let is_valid_item = next == 0 || self.bits_def == 0; + Some(ControlWordDesc { + is_new_row, + // Either there is no rep, in which case there are no invisible items + // or there is no def, in which case there are no invisible items + is_visible: true, + is_valid_item, + }) + } +} + +impl> UnaryControlWordIterator { + fn append_next(&mut self, buf: &mut Vec) -> Option { + let next = self.repdef.next().unwrap() & self.level_mask; + let control_word = next.to_le_bytes(); + buf.push(control_word[0]); + buf.push(control_word[1]); + let is_new_row = self.max_rep == 0 || next == self.max_rep; + let is_valid_item = next == 0 || self.bits_def == 0; + Some(ControlWordDesc { + is_new_row, + is_visible: true, + is_valid_item, + }) + } +} + +impl> UnaryControlWordIterator { + fn append_next(&mut self, buf: &mut Vec) -> Option { + let next = self.repdef.next()?; + let next = (next & self.level_mask) as u32; + let control_word = next.to_le_bytes(); + buf.push(control_word[0]); + buf.push(control_word[1]); + buf.push(control_word[2]); + buf.push(control_word[3]); + let is_new_row = self.max_rep == 0 || next as u16 == self.max_rep; + let is_valid_item = next == 0 || self.bits_def == 0; + Some(ControlWordDesc { + is_new_row, + is_visible: true, + is_valid_item, + }) + } +} + +/// A [`ControlWordIterator`] when there are no repetition or definition levels +#[derive(Debug)] +pub struct NilaryControlWordIterator { + len: usize, + idx: usize, +} + +impl NilaryControlWordIterator { + fn append_next(&mut self) -> Option { + if self.idx == self.len { + None + } else { + self.idx += 1; + Some(ControlWordDesc { + is_new_row: true, + is_visible: true, + is_valid_item: true, + }) + } + } +} + +/// Helper function to get a bit mask of the given width +fn get_mask(width: u16) -> u16 { + (1 << width) - 1 +} + +// We're really going out of our way to avoid boxing here but this will be called on a per-value basis +// so it is in the critical path. +type SpecificBinaryControlWordIterator<'a, T> = BinaryControlWordIterator< + Zip>, Copied>>, + T, +>; + +/// An iterator that generates control words from repetition and definition levels +/// +/// "Control word" is just a fancy term for a single u8/u16/u32 that contains both +/// the repetition and definition in it. +/// +/// In the large majority of case we only need a single byte to represent both the +/// repetition and definition levels. However, if there is deep nesting then we may +/// need two bytes. In the worst case we need 4 bytes though this suggests hundreds of +/// levels of nesting which seems unlikely to encounter in practice. +#[derive(Debug)] +pub enum ControlWordIterator<'a> { + Binary8(SpecificBinaryControlWordIterator<'a, u8>), + Binary16(SpecificBinaryControlWordIterator<'a, u16>), + Binary32(SpecificBinaryControlWordIterator<'a, u32>), + Unary8(UnaryControlWordIterator>, u8>), + Unary16(UnaryControlWordIterator>, u16>), + Unary32(UnaryControlWordIterator>, u32>), + Nilary(NilaryControlWordIterator), +} + +/// Describes the properties of a control word +#[derive(Debug)] +pub struct ControlWordDesc { + pub is_new_row: bool, + pub is_visible: bool, + pub is_valid_item: bool, +} + +impl ControlWordIterator<'_> { + /// Appends the next control word to the buffer + /// + /// Returns true if this is the start of a new item (i.e. the repetition level is maxed out) + pub fn append_next(&mut self, buf: &mut Vec) -> Option { + match self { + Self::Binary8(iter) => iter.append_next(buf), + Self::Binary16(iter) => iter.append_next(buf), + Self::Binary32(iter) => iter.append_next(buf), + Self::Unary8(iter) => iter.append_next(buf), + Self::Unary16(iter) => iter.append_next(buf), + Self::Unary32(iter) => iter.append_next(buf), + Self::Nilary(iter) => iter.append_next(), + } + } + + /// Return true if the control word iterator has repetition levels + pub fn has_repetition(&self) -> bool { + match self { + Self::Binary8(_) | Self::Binary16(_) | Self::Binary32(_) => true, + Self::Unary8(iter) => iter.bits_rep > 0, + Self::Unary16(iter) => iter.bits_rep > 0, + Self::Unary32(iter) => iter.bits_rep > 0, + Self::Nilary(_) => false, + } + } + + /// Returns the number of bytes per control word + pub fn bytes_per_word(&self) -> usize { + match self { + Self::Binary8(_) => 1, + Self::Binary16(_) => 2, + Self::Binary32(_) => 4, + Self::Unary8(_) => 1, + Self::Unary16(_) => 2, + Self::Unary32(_) => 4, + Self::Nilary(_) => 0, + } + } + + /// Returns the number of bits used for the repetition level + pub fn bits_rep(&self) -> u8 { + match self { + Self::Binary8(iter) => iter.bits_rep, + Self::Binary16(iter) => iter.bits_rep, + Self::Binary32(iter) => iter.bits_rep, + Self::Unary8(iter) => iter.bits_rep, + Self::Unary16(iter) => iter.bits_rep, + Self::Unary32(iter) => iter.bits_rep, + Self::Nilary(_) => 0, + } + } + + /// Returns the number of bits used for the definition level + pub fn bits_def(&self) -> u8 { + match self { + Self::Binary8(iter) => iter.bits_def, + Self::Binary16(iter) => iter.bits_def, + Self::Binary32(iter) => iter.bits_def, + Self::Unary8(iter) => iter.bits_def, + Self::Unary16(iter) => iter.bits_def, + Self::Unary32(iter) => iter.bits_def, + Self::Nilary(_) => 0, + } + } +} + +/// Builds a [`ControlWordIterator`] from repetition and definition levels +/// by first calculating the width needed and then creating the iterator +/// with the appropriate width +pub fn build_control_word_iterator<'a>( + rep: Option<&'a [u16]>, + max_rep: u16, + def: Option<&'a [u16]>, + max_def: u16, + max_visible_def: u16, + len: usize, +) -> ControlWordIterator<'a> { + let rep_width = if max_rep == 0 { + 0 + } else { + log_2_ceil(max_rep as u32) as u16 + }; + let rep_mask = if max_rep == 0 { 0 } else { get_mask(rep_width) }; + let def_width = if max_def == 0 { + 0 + } else { + log_2_ceil(max_def as u32) as u16 + }; + let def_mask = if max_def == 0 { 0 } else { get_mask(def_width) }; + let total_width = rep_width + def_width; + // A levels array can be non-empty while the corresponding max level is 0: + // e.g. a page holding only valid rows after a page split has all zero + // definition levels. With a zero total width the decoder reads a NIL + // layout (zero bytes per control word), so the writer must also emit the + // NIL writer instead of a Unary/Binary writer that would write one + // all-zero byte per word. + if total_width == 0 { + return ControlWordIterator::Nilary(NilaryControlWordIterator { len, idx: 0 }); + } + match (rep, def) { + (Some(rep), Some(def)) => { + let iter = rep.iter().copied().zip(def.iter().copied()); + let def_width = def_width as usize; + if total_width <= 8 { + ControlWordIterator::Binary8(BinaryControlWordIterator { + repdef: iter, + rep_mask, + def_mask, + def_width, + max_rep, + max_visible_def, + bits_rep: rep_width as u8, + bits_def: def_width as u8, + phantom: std::marker::PhantomData, + }) + } else if total_width <= 16 { + ControlWordIterator::Binary16(BinaryControlWordIterator { + repdef: iter, + rep_mask, + def_mask, + def_width, + max_rep, + max_visible_def, + bits_rep: rep_width as u8, + bits_def: def_width as u8, + phantom: std::marker::PhantomData, + }) + } else { + ControlWordIterator::Binary32(BinaryControlWordIterator { + repdef: iter, + rep_mask, + def_mask, + def_width, + max_rep, + max_visible_def, + bits_rep: rep_width as u8, + bits_def: def_width as u8, + phantom: std::marker::PhantomData, + }) + } + } + (Some(lev), None) => { + let iter = lev.iter().copied(); + if total_width <= 8 { + ControlWordIterator::Unary8(UnaryControlWordIterator { + repdef: iter, + level_mask: rep_mask, + bits_rep: total_width as u8, + bits_def: 0, + max_rep, + phantom: std::marker::PhantomData, + }) + } else if total_width <= 16 { + ControlWordIterator::Unary16(UnaryControlWordIterator { + repdef: iter, + level_mask: rep_mask, + bits_rep: total_width as u8, + bits_def: 0, + max_rep, + phantom: std::marker::PhantomData, + }) + } else { + ControlWordIterator::Unary32(UnaryControlWordIterator { + repdef: iter, + level_mask: rep_mask, + bits_rep: total_width as u8, + bits_def: 0, + max_rep, + phantom: std::marker::PhantomData, + }) + } + } + (None, Some(lev)) => { + let iter = lev.iter().copied(); + if total_width <= 8 { + ControlWordIterator::Unary8(UnaryControlWordIterator { + repdef: iter, + level_mask: def_mask, + bits_rep: 0, + bits_def: total_width as u8, + max_rep: 0, + phantom: std::marker::PhantomData, + }) + } else if total_width <= 16 { + ControlWordIterator::Unary16(UnaryControlWordIterator { + repdef: iter, + level_mask: def_mask, + bits_rep: 0, + bits_def: total_width as u8, + max_rep: 0, + phantom: std::marker::PhantomData, + }) + } else { + ControlWordIterator::Unary32(UnaryControlWordIterator { + repdef: iter, + level_mask: def_mask, + bits_rep: 0, + bits_def: total_width as u8, + max_rep: 0, + phantom: std::marker::PhantomData, + }) + } + } + (None, None) => ControlWordIterator::Nilary(NilaryControlWordIterator { len, idx: 0 }), + } +} + +/// A parser to unwrap control words into repetition and definition levels +/// +/// This is the inverse of the [`ControlWordIterator`]. +#[derive(Copy, Clone, Debug)] +pub enum ControlWordParser { + // First item is the bits to shift, second is the mask to apply (the mask can be + // calculated from the bits to shift but we don't want to calculate it each time) + BOTH8(u8, u32), + BOTH16(u8, u32), + BOTH32(u8, u32), + REP8, + REP16, + REP32, + DEF8, + DEF16, + DEF32, + NIL, +} + +impl ControlWordParser { + fn parse_both( + src: &[u8], + dst_rep: &mut Vec, + dst_def: &mut Vec, + bits_to_shift: u8, + mask_to_apply: u32, + ) { + match WORD_SIZE { + 1 => { + let word = src[0]; + let rep = word >> bits_to_shift; + let def = word & (mask_to_apply as u8); + dst_rep.push(rep as u16); + dst_def.push(def as u16); + } + 2 => { + let word = u16::from_le_bytes([src[0], src[1]]); + let rep = word >> bits_to_shift; + let def = word & mask_to_apply as u16; + dst_rep.push(rep); + dst_def.push(def); + } + 4 => { + let word = u32::from_le_bytes([src[0], src[1], src[2], src[3]]); + let rep = word >> bits_to_shift; + let def = word & mask_to_apply; + dst_rep.push(rep as u16); + dst_def.push(def as u16); + } + _ => unreachable!(), + } + } + + fn parse_desc_both( + src: &[u8], + bits_to_shift: u8, + mask_to_apply: u32, + max_rep: u16, + max_visible_def: u16, + ) -> ControlWordDesc { + match WORD_SIZE { + 1 => { + let word = src[0]; + let rep = word >> bits_to_shift; + let def = word & (mask_to_apply as u8); + let is_visible = def as u16 <= max_visible_def; + let is_new_row = rep as u16 == max_rep; + let is_valid_item = def == 0; + ControlWordDesc { + is_visible, + is_new_row, + is_valid_item, + } + } + 2 => { + let word = u16::from_le_bytes([src[0], src[1]]); + let rep = word >> bits_to_shift; + let def = word & mask_to_apply as u16; + let is_visible = def <= max_visible_def; + let is_new_row = rep == max_rep; + let is_valid_item = def == 0; + ControlWordDesc { + is_visible, + is_new_row, + is_valid_item, + } + } + 4 => { + let word = u32::from_le_bytes([src[0], src[1], src[2], src[3]]); + let rep = word >> bits_to_shift; + let def = word & mask_to_apply; + let is_visible = def as u16 <= max_visible_def; + let is_new_row = rep as u16 == max_rep; + let is_valid_item = def == 0; + ControlWordDesc { + is_visible, + is_new_row, + is_valid_item, + } + } + _ => unreachable!(), + } + } + + fn parse_one(src: &[u8], dst: &mut Vec) { + match WORD_SIZE { + 1 => { + let word = src[0]; + dst.push(word as u16); + } + 2 => { + let word = u16::from_le_bytes([src[0], src[1]]); + dst.push(word); + } + 4 => { + let word = u32::from_le_bytes([src[0], src[1], src[2], src[3]]); + dst.push(word as u16); + } + _ => unreachable!(), + } + } + + fn parse_rep_desc_one(src: &[u8], max_rep: u16) -> ControlWordDesc { + match WORD_SIZE { + 1 => ControlWordDesc { + is_new_row: src[0] as u16 == max_rep, + is_visible: true, + is_valid_item: true, + }, + 2 => ControlWordDesc { + is_new_row: u16::from_le_bytes([src[0], src[1]]) == max_rep, + is_visible: true, + is_valid_item: true, + }, + 4 => ControlWordDesc { + is_new_row: u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as u16 == max_rep, + is_visible: true, + is_valid_item: true, + }, + _ => unreachable!(), + } + } + + fn parse_def_desc_one(src: &[u8]) -> ControlWordDesc { + match WORD_SIZE { + 1 => ControlWordDesc { + is_new_row: true, + is_visible: true, + is_valid_item: src[0] == 0, + }, + 2 => ControlWordDesc { + is_new_row: true, + is_visible: true, + is_valid_item: u16::from_le_bytes([src[0], src[1]]) == 0, + }, + 4 => ControlWordDesc { + is_new_row: true, + is_visible: true, + is_valid_item: u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as u16 == 0, + }, + _ => unreachable!(), + } + } + + /// Returns the number of bytes per control word + pub fn bytes_per_word(&self) -> usize { + match self { + Self::BOTH8(..) => 1, + Self::BOTH16(..) => 2, + Self::BOTH32(..) => 4, + Self::REP8 => 1, + Self::REP16 => 2, + Self::REP32 => 4, + Self::DEF8 => 1, + Self::DEF16 => 2, + Self::DEF32 => 4, + Self::NIL => 0, + } + } + + /// Appends the next control word to the rep & def buffers + /// + /// `src` should be pointing at the first byte (little endian) of the control word + /// + /// `dst_rep` and `dst_def` are the buffers to append the rep and def levels to. + /// They will not be appended to if not needed. + pub fn parse(&self, src: &[u8], dst_rep: &mut Vec, dst_def: &mut Vec) { + match self { + Self::BOTH8(bits_to_shift, mask_to_apply) => { + Self::parse_both::<1>(src, dst_rep, dst_def, *bits_to_shift, *mask_to_apply) + } + Self::BOTH16(bits_to_shift, mask_to_apply) => { + Self::parse_both::<2>(src, dst_rep, dst_def, *bits_to_shift, *mask_to_apply) + } + Self::BOTH32(bits_to_shift, mask_to_apply) => { + Self::parse_both::<4>(src, dst_rep, dst_def, *bits_to_shift, *mask_to_apply) + } + Self::REP8 => Self::parse_one::<1>(src, dst_rep), + Self::REP16 => Self::parse_one::<2>(src, dst_rep), + Self::REP32 => Self::parse_one::<4>(src, dst_rep), + Self::DEF8 => Self::parse_one::<1>(src, dst_def), + Self::DEF16 => Self::parse_one::<2>(src, dst_def), + Self::DEF32 => Self::parse_one::<4>(src, dst_def), + Self::NIL => {} + } + } + + /// Return true if the control words contain repetition information + pub fn has_rep(&self) -> bool { + match self { + Self::BOTH8(..) + | Self::BOTH16(..) + | Self::BOTH32(..) + | Self::REP8 + | Self::REP16 + | Self::REP32 => true, + Self::DEF8 | Self::DEF16 | Self::DEF32 | Self::NIL => false, + } + } + + /// Temporarily parses the control word to inspect its properties but does not append to any buffers + pub fn parse_desc(&self, src: &[u8], max_rep: u16, max_visible_def: u16) -> ControlWordDesc { + match self { + Self::BOTH8(bits_to_shift, mask_to_apply) => Self::parse_desc_both::<1>( + src, + *bits_to_shift, + *mask_to_apply, + max_rep, + max_visible_def, + ), + Self::BOTH16(bits_to_shift, mask_to_apply) => Self::parse_desc_both::<2>( + src, + *bits_to_shift, + *mask_to_apply, + max_rep, + max_visible_def, + ), + Self::BOTH32(bits_to_shift, mask_to_apply) => Self::parse_desc_both::<4>( + src, + *bits_to_shift, + *mask_to_apply, + max_rep, + max_visible_def, + ), + Self::REP8 => Self::parse_rep_desc_one::<1>(src, max_rep), + Self::REP16 => Self::parse_rep_desc_one::<2>(src, max_rep), + Self::REP32 => Self::parse_rep_desc_one::<4>(src, max_rep), + Self::DEF8 => Self::parse_def_desc_one::<1>(src), + Self::DEF16 => Self::parse_def_desc_one::<2>(src), + Self::DEF32 => Self::parse_def_desc_one::<4>(src), + Self::NIL => ControlWordDesc { + is_new_row: true, + is_valid_item: true, + is_visible: true, + }, + } + } + + /// Creates a new parser from the number of bits used for the repetition and definition levels + pub fn new(bits_rep: u8, bits_def: u8) -> Self { + let total_bits = bits_rep + bits_def; + + enum WordSize { + One, + Two, + Four, + } + + let word_size = if total_bits <= 8 { + WordSize::One + } else if total_bits <= 16 { + WordSize::Two + } else { + WordSize::Four + }; + + match (bits_rep > 0, bits_def > 0, word_size) { + (false, false, _) => Self::NIL, + (false, true, WordSize::One) => Self::DEF8, + (false, true, WordSize::Two) => Self::DEF16, + (false, true, WordSize::Four) => Self::DEF32, + (true, false, WordSize::One) => Self::REP8, + (true, false, WordSize::Two) => Self::REP16, + (true, false, WordSize::Four) => Self::REP32, + (true, true, WordSize::One) => Self::BOTH8(bits_def, get_mask(bits_def as u16) as u32), + (true, true, WordSize::Two) => Self::BOTH16(bits_def, get_mask(bits_def as u16) as u32), + (true, true, WordSize::Four) => { + Self::BOTH32(bits_def, get_mask(bits_def as u16) as u32) + } + } + } +} + +#[cfg(test)] +mod tests { + use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; + + use crate::encodings::logical::primitive::sparse::{ + SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan, SparseValidityMeaning, + SparseValiditySet, + }; + use crate::repdef::{ + CompositeRepDefUnraveler, DefinitionInterpretation, RepDefUnraveler, SerializedRepDefs, + max_visible_level, + }; + + use super::RepDefBuilder; + + /// `def_meaning` is leaf-first, so only the layers ahead of the first list sit + /// underneath it. Counting nullable layers *past* the list makes a decoder treat + /// invisible items as though they owned a leaf value. + #[rstest::rstest] + // No lists: nothing can be invisible, so every level is visible + #[case::no_lists( + &[DefinitionInterpretation::NullableItem, DefinitionInterpretation::NullableItem], + 2 + )] + // A nullable leaf under a list keeps its own level visible + #[case::nullable_leaf_under_list( + &[DefinitionInterpretation::NullableItem, DefinitionInterpretation::NullableList], + 1 + )] + // Regression: a nullable struct above the list must not raise the threshold + #[case::nullable_struct_above_list( + &[ + DefinitionInterpretation::AllValidItem, + DefinitionInterpretation::NullableList, + DefinitionInterpretation::NullableItem, + ], + 0 + )] + fn test_max_visible_level( + #[case] def_meaning: &[DefinitionInterpretation], + #[case] expected: u16, + ) { + assert_eq!(max_visible_level(def_meaning), expected); + } + + fn validity(values: &[bool]) -> NullBuffer { + NullBuffer::from_iter(values.iter().copied()) + } + + fn offsets_32(values: &[i32]) -> OffsetBuffer { + OffsetBuffer::::new(ScalarBuffer::from_iter(values.iter().copied())) + } + + fn offsets_64(values: &[i64]) -> OffsetBuffer { + OffsetBuffer::::new(ScalarBuffer::from_iter(values.iter().copied())) + } + + #[test] + fn sparse_sibling_validity_mismatch_is_invalid_input() { + let sparse = |positions| { + RepDefUnraveler::new_sparse(SparseStructuralPlan { + layers: vec![SparseStructuralLayerPlan::Validity { + num_slots: 2, + validity: SparseValiditySet { + meaning: SparseValidityMeaning::NullPositions, + positions, + }, + }], + num_items: 2, + num_visible_items: 2, + }) + }; + let mut repdef = CompositeRepDefUnraveler::new(vec![sparse(SparsePositionSet::Empty)]); + repdef.add_compatibility_check(CompositeRepDefUnraveler::new(vec![sparse( + SparsePositionSet::Explicit(vec![0]), + )])); + + let err = repdef.unravel_validity(2).unwrap_err(); + assert!(matches!(err, lance_core::Error::InvalidInput { .. })); + assert!(err.to_string().contains("incompatible validity metadata")); + } + + #[test] + fn test_repdef_empty_offsets() { + // Empty offsets should serialize without panicking. + let mut builder = RepDefBuilder::default(); + builder.add_offsets(offsets_32(&[0]), None); + let repdefs = RepDefBuilder::serialize(vec![builder]); + assert!(repdefs.repetition_levels.is_none()); + assert!(repdefs.definition_levels.is_none()); + } + + #[test] + fn test_repdef_basic() { + // Basic case, rep & def + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_64(&[0, 2, 2, 5]), + Some(validity(&[true, false, true])), + ); + builder.add_offsets( + offsets_64(&[0, 1, 3, 5, 5, 9]), + Some(validity(&[true, true, true, false, true])), + ); + builder.add_validity_bitmap(validity(&[ + true, true, true, false, false, false, true, true, false, + ])); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!(vec![0, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1], *def); + assert_eq!(vec![2, 1, 0, 2, 2, 0, 1, 1, 0, 0, 0], *rep); + + // [[I], [I, I]], NULL, [[NULL, NULL], NULL, [NULL, I, I, NULL]] + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + Some(rep.as_ref().to_vec()), + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 9, + )]); + + // Note: validity doesn't exactly round-trip because repdef normalizes some of the + // redundant validity values + assert_eq!( + unraveler.unravel_validity(9).unwrap(), + Some(validity(&[ + true, true, true, false, false, false, true, true, false + ])) + ); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 1, 3, 5, 5, 9]).inner()); + assert_eq!(val, Some(validity(&[true, true, true, false, true]))); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 2, 2, 5]).inner()); + assert_eq!(val, Some(validity(&[true, false, true]))); + } + + #[test] + fn test_repdef_simple_null_empty_list() { + let check = |repdefs: SerializedRepDefs, last_def: DefinitionInterpretation| { + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([1, 0, 1, 1, 0, 0], *rep); + assert_eq!([0, 0, 2, 0, 1, 0], *def); + assert_eq!( + vec![DefinitionInterpretation::NullableItem, last_def,], + repdefs.def_meaning + ); + }; + + // Null list and empty list should be serialized mostly the same + + // Null case + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_32(&[0, 2, 2, 5]), + Some(validity(&[true, false, true])), + ); + builder.add_validity_bitmap(validity(&[true, true, true, false, true])); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + check(repdefs, DefinitionInterpretation::NullableList); + + // Empty case + let mut builder = RepDefBuilder::default(); + builder.add_offsets(offsets_32(&[0, 2, 2, 5]), None); + builder.add_validity_bitmap(validity(&[true, true, true, false, true])); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + check(repdefs, DefinitionInterpretation::EmptyableList); + } + + #[test] + fn test_repdef_empty_list_at_end() { + // Regresses a failure we encountered when the last item was an empty list + let mut builder = RepDefBuilder::default(); + builder.add_offsets(offsets_32(&[0, 2, 5, 5]), None); + builder.add_validity_bitmap(validity(&[true, true, true, false, true])); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([1, 0, 1, 0, 0, 1], *rep); + assert_eq!([0, 0, 0, 1, 0, 2], *def); + assert_eq!( + vec![ + DefinitionInterpretation::NullableItem, + DefinitionInterpretation::EmptyableList, + ], + repdefs.def_meaning + ); + } + + #[test] + fn test_repdef_abnormal_nulls() { + // List nulls are allowed to have non-empty offsets and garbage values + // and the add_offsets call should normalize this + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_32(&[0, 2, 5, 8]), + Some(validity(&[true, false, true])), + ); + // Note: we pass 5 here and not 8. If add_offsets tells us there is garbage nulls they + // should be removed before continuing + builder.add_no_null(5); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([1, 0, 1, 1, 0, 0], *rep); + assert_eq!([0, 0, 1, 0, 0, 0], *def); + + assert_eq!( + vec![ + DefinitionInterpretation::AllValidItem, + DefinitionInterpretation::NullableList, + ], + repdefs.def_meaning + ); + } + + #[test] + fn test_repdef_fsl() { + let mut builder = RepDefBuilder::default(); + builder.add_fsl(Some(validity(&[true, false])), 2, 2); + builder.add_fsl(None, 2, 4); + builder.add_validity_bitmap(validity(&[ + true, false, true, false, true, false, true, false, + ])); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + assert_eq!( + vec![ + DefinitionInterpretation::NullableItem, + DefinitionInterpretation::AllValidItem, + DefinitionInterpretation::NullableItem + ], + repdefs.def_meaning + ); + + assert!(repdefs.repetition_levels.is_none()); + + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([0, 1, 0, 1, 2, 2, 2, 2], *def); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + None, + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 8, + )]); + + assert_eq!( + unraveler.unravel_validity(8).unwrap(), + Some(validity(&[ + true, false, true, false, false, false, false, false + ])) + ); + assert_eq!(unraveler.unravel_fsl_validity(4, 2).unwrap(), None); + assert_eq!( + unraveler.unravel_fsl_validity(2, 2).unwrap(), + Some(validity(&[true, false])) + ); + } + + #[test] + fn test_repdef_fsl_zero_dimension_is_invalid_input() { + // A zero dimension can only reach us from a malformed schema. Decimating with it + // used to loop forever, writing past the end of the definition levels buffer. + let mut builder = RepDefBuilder::default(); + builder.add_fsl(Some(validity(&[true, false])), 2, 2); + builder.add_validity_bitmap(validity(&[true, false, true, false])); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + let def = repdefs.definition_levels.unwrap(); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + None, + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 4, + )]); + // Consume the item layer so the fixed-size-list layer is next + unraveler.unravel_validity(4).unwrap(); + + let err = unraveler.unravel_fsl_validity(2, 0).unwrap_err(); + assert!(matches!(err, lance_core::Error::InvalidInput { .. })); + assert!( + err.to_string() + .contains("dimension must be a positive integer"), + "unexpected error: {}", + err + ); + } + + #[test] + fn test_repdef_fsl_allvalid_item() { + let mut builder = RepDefBuilder::default(); + builder.add_fsl(Some(validity(&[true, false])), 2, 2); + builder.add_fsl(None, 2, 4); + builder.add_no_null(8); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + assert_eq!( + vec![ + DefinitionInterpretation::AllValidItem, + DefinitionInterpretation::AllValidItem, + DefinitionInterpretation::NullableItem + ], + repdefs.def_meaning + ); + + assert!(repdefs.repetition_levels.is_none()); + + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([0, 0, 0, 0, 1, 1, 1, 1], *def); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + None, + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 8, + )]); + + assert_eq!(unraveler.unravel_validity(8).unwrap(), None); + assert_eq!(unraveler.unravel_fsl_validity(4, 2).unwrap(), None); + assert_eq!( + unraveler.unravel_fsl_validity(2, 2).unwrap(), + Some(validity(&[true, false])) + ); + } + + #[test] + fn test_repdef_sliced_offsets() { + // Sliced lists may have offsets that don't start with zero. The + // add_offsets call needs to normalize these to operate correctly. + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_32(&[5, 7, 7, 10]), + Some(validity(&[true, false, true])), + ); + builder.add_no_null(5); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([1, 0, 1, 1, 0, 0], *rep); + assert_eq!([0, 0, 1, 0, 0, 0], *def); + + assert_eq!( + vec![ + DefinitionInterpretation::AllValidItem, + DefinitionInterpretation::NullableList, + ], + repdefs.def_meaning + ); + } + + #[test] + fn test_repdef_complex_null_empty() { + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_32(&[0, 4, 4, 4, 6]), + Some(validity(&[true, false, true, true])), + ); + builder.add_offsets( + offsets_32(&[0, 1, 1, 2, 2, 2, 3]), + Some(validity(&[true, false, true, false, true, true])), + ); + builder.add_no_null(3); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([2, 1, 1, 1, 2, 2, 2, 1], *rep); + assert_eq!([0, 1, 0, 1, 3, 4, 2, 0], *def); + } + + #[test] + fn test_repdef_empty_list_no_null() { + // Tests when we have some empty lists but no null lists. This case + // caused some bugs because we have definition but no nulls + let mut builder = RepDefBuilder::default(); + builder.add_offsets(offsets_32(&[0, 4, 4, 4, 6]), None); + builder.add_no_null(6); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep); + assert_eq!([0, 0, 0, 0, 1, 1, 0, 0], *def); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + Some(rep.as_ref().to_vec()), + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 8, + )]); + + assert_eq!(unraveler.unravel_validity(6).unwrap(), None); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner()); + assert_eq!(val, None); + } + + #[test] + fn test_repdef_all_valid() { + let mut builder = RepDefBuilder::default(); + builder.add_offsets(offsets_64(&[0, 2, 3, 5]), None); + builder.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None); + builder.add_no_null(9); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + let rep = repdefs.repetition_levels.unwrap(); + assert!(repdefs.definition_levels.is_none()); + + assert_eq!([2, 1, 0, 2, 0, 2, 0, 1, 0], *rep); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + Some(rep.as_ref().to_vec()), + None, + repdefs.def_meaning.into(), + 9, + )]); + + assert_eq!(unraveler.unravel_validity(9).unwrap(), None); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 1, 3, 5, 7, 9]).inner()); + assert_eq!(val, None); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 2, 3, 5]).inner()); + assert_eq!(val, None); + } + + #[test] + fn test_repdef_nested_list_multibatch_matches_single() { + // Single builder: List>, 3 rows. + // outer [0,2,3,5] -> rows have 2,1,2 inner lists + // inner [0,1,3,5,7,9] -> 5 inner lists, lengths 1,2,2,2,2 (9 leaf) + let mut single = RepDefBuilder::default(); + single.add_offsets(offsets_64(&[0, 2, 3, 5]), None); + single.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None); + single.add_no_null(9); + let single_rep = RepDefBuilder::serialize(vec![single]) + .repetition_levels + .unwrap(); + + // Same logical data split into two batches: + // batch0 = rows 0,1 : outer [0,2,3], inner [0,1,3,5] (3 inner, 5 leaf) + // batch1 = row 2 : outer [0,2], inner [0,2,4] (2 inner, 4 leaf) + let mut b0 = RepDefBuilder::default(); + b0.add_offsets(offsets_64(&[0, 2, 3]), None); + b0.add_offsets(offsets_64(&[0, 1, 3, 5]), None); + b0.add_no_null(5); + let mut b1 = RepDefBuilder::default(); + b1.add_offsets(offsets_64(&[0, 2]), None); + b1.add_offsets(offsets_64(&[0, 2, 4]), None); + b1.add_no_null(4); + let multi_rep = RepDefBuilder::serialize(vec![b0, b1]) + .repetition_levels + .unwrap(); + + assert_eq!( + *single_rep, *multi_rep, + "multi-batch nested-list rep levels must equal single-batch" + ); + } + + #[test] + fn test_only_empty_lists() { + let mut builder = RepDefBuilder::default(); + builder.add_offsets(offsets_32(&[0, 4, 4, 4, 6]), None); + builder.add_no_null(6); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep); + assert_eq!([0, 0, 0, 0, 1, 1, 0, 0], *def); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + Some(rep.as_ref().to_vec()), + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 8, + )]); + + assert_eq!(unraveler.unravel_validity(6).unwrap(), None); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner()); + assert_eq!(val, None); + } + + #[test] + fn test_only_null_lists() { + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_32(&[0, 4, 4, 4, 6]), + Some(validity(&[true, false, false, true])), + ); + builder.add_no_null(6); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep); + assert_eq!([0, 0, 0, 0, 1, 1, 0, 0], *def); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + Some(rep.as_ref().to_vec()), + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 8, + )]); + + assert_eq!(unraveler.unravel_validity(6).unwrap(), None); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner()); + assert_eq!(val, Some(validity(&[true, false, false, true]))); + } + + #[test] + fn test_null_and_empty_lists() { + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_32(&[0, 4, 4, 4, 6]), + Some(validity(&[true, false, true, true])), + ); + builder.add_no_null(6); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep); + assert_eq!([0, 0, 0, 0, 1, 2, 0, 0], *def); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + Some(rep.as_ref().to_vec()), + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 8, + )]); + + assert_eq!(unraveler.unravel_validity(6).unwrap(), None); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner()); + assert_eq!(val, Some(validity(&[true, false, true, true]))); + } + + #[test] + fn test_repdef_nullable_struct_in_null_and_empty_lists() { + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_32(&[0, 0, 3, 3]), + Some(validity(&[false, true, true])), + ); + builder.add_validity_bitmap(validity(&[false, true, true])); + builder.add_validity_bitmap(validity(&[true, false, true])); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + repdefs.repetition_levels.map(|levels| levels.to_vec()), + repdefs.definition_levels.map(|levels| levels.to_vec()), + repdefs.def_meaning.into(), + 3, + )]); + + assert_eq!( + unraveler.unravel_validity(3).unwrap(), + Some(validity(&[false, false, true])) + ); + assert_eq!( + unraveler.unravel_validity(3).unwrap(), + Some(validity(&[false, true, true])) + ); + let (offsets, nulls) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(offsets.inner(), offsets_32(&[0, 0, 3, 3]).inner()); + assert_eq!(nulls, Some(validity(&[false, true, true]))); + } + + #[test] + fn test_repdef_null_struct_valid_list() { + // This regresses a bug + + let rep = vec![1, 0, 0, 0]; + let def = vec![2, 0, 2, 2]; + // AllValidList> + let def_meaning = vec![ + DefinitionInterpretation::NullableItem, + DefinitionInterpretation::NullableItem, + DefinitionInterpretation::AllValidList, + ]; + let num_items = 4; + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + Some(rep), + Some(def), + def_meaning.into(), + num_items, + )]); + + assert_eq!( + unraveler.unravel_validity(4).unwrap(), + Some(validity(&[false, true, false, false])) + ); + assert_eq!( + unraveler.unravel_validity(4).unwrap(), + Some(validity(&[false, true, false, false])) + ); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 4]).inner()); + assert_eq!(val, None); + } + + #[test] + fn test_repdef_no_rep() { + let mut builder = RepDefBuilder::default(); + builder.add_no_null(5); + builder.add_validity_bitmap(validity(&[false, false, true, true, true])); + builder.add_validity_bitmap(validity(&[false, true, true, true, false])); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + assert!(repdefs.repetition_levels.is_none()); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([2, 2, 0, 0, 1], *def); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + None, + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 5, + )]); + + assert_eq!( + unraveler.unravel_validity(5).unwrap(), + Some(validity(&[false, false, true, true, false])) + ); + assert_eq!( + unraveler.unravel_validity(5).unwrap(), + Some(validity(&[false, false, true, true, true])) + ); + assert_eq!(unraveler.unravel_validity(5).unwrap(), None); + } + + #[test] + fn test_composite_unravel() { + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_64(&[0, 2, 2, 5]), + Some(validity(&[true, false, true])), + ); + builder.add_no_null(5); + let repdef1 = RepDefBuilder::serialize(vec![builder]); + + let mut builder = RepDefBuilder::default(); + builder.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None); + builder.add_no_null(9); + let repdef2 = RepDefBuilder::serialize(vec![builder]); + + let rep1 = repdef1.repetition_levels.clone().unwrap(); + let def1 = repdef1.definition_levels.clone().unwrap(); + let rep2 = repdef2.repetition_levels.clone().unwrap(); + assert!(repdef2.definition_levels.is_none()); + + assert_eq!([1, 0, 1, 1, 0, 0], *rep1); + assert_eq!([0, 0, 1, 0, 0, 0], *def1); + assert_eq!([1, 1, 0, 1, 0, 1, 0, 1, 0], *rep2); + + let unravel1 = RepDefUnraveler::new( + repdef1.repetition_levels.map(|l| l.to_vec()), + repdef1.definition_levels.map(|l| l.to_vec()), + repdef1.def_meaning.into(), + 5, + ); + let unravel2 = RepDefUnraveler::new( + repdef2.repetition_levels.map(|l| l.to_vec()), + repdef2.definition_levels.map(|l| l.to_vec()), + repdef2.def_meaning.into(), + 9, + ); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![unravel1, unravel2]); + + assert!(unraveler.unravel_validity(9).unwrap().is_none()); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!( + off.inner(), + offsets_32(&[0, 2, 2, 5, 6, 8, 10, 12, 14]).inner() + ); + assert_eq!( + val, + Some(validity(&[true, false, true, true, true, true, true, true])) + ); + } + + #[test] + fn test_repdef_multiple_builders() { + // Basic case, rep & def + let mut builder1 = RepDefBuilder::default(); + builder1.add_offsets(offsets_64(&[0, 2]), None); + builder1.add_offsets(offsets_64(&[0, 1, 3]), None); + builder1.add_validity_bitmap(validity(&[true, true, true])); + + let mut builder2 = RepDefBuilder::default(); + builder2.add_offsets(offsets_64(&[0, 0, 3]), Some(validity(&[false, true]))); + builder2.add_offsets( + offsets_64(&[0, 2, 2, 6]), + Some(validity(&[true, false, true])), + ); + builder2.add_validity_bitmap(validity(&[false, false, false, true, true, false])); + + let repdefs = RepDefBuilder::serialize(vec![builder1, builder2]); + + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([2, 1, 0, 2, 2, 0, 1, 1, 0, 0, 0], *rep); + assert_eq!([0, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1], *def); + } + + #[test] + fn test_all_valid_validity_bitmap_serializes_as_no_null() { + let mut from_bitmap = RepDefBuilder::default(); + from_bitmap.add_validity_bitmap(validity(&[true, true, true, true])); + + let mut from_no_null = RepDefBuilder::default(); + from_no_null.add_no_null(4); + + let from_bitmap = RepDefBuilder::serialize(vec![from_bitmap]); + let from_no_null = RepDefBuilder::serialize(vec![from_no_null]); + + assert!(from_bitmap.repetition_levels.is_none()); + assert!(from_bitmap.definition_levels.is_none()); + assert_eq!(from_bitmap.def_meaning, from_no_null.def_meaning); + assert_eq!( + from_bitmap.max_visible_level, + from_no_null.max_visible_level + ); + } + + #[test] + fn test_slicer() { + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_64(&[0, 2, 2, 30, 30]), + Some(validity(&[true, false, true, true])), + ); + builder.add_no_null(30); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + + let mut rep_slicer = repdefs.rep_slicer().unwrap(); + + // First 5 items include a null list so we get 6 levels (12 bytes) + assert_eq!(rep_slicer.slice_next(5).len(), 12); + // Next 20 are all plain + assert_eq!(rep_slicer.slice_next(20).len(), 40); + // Last 5 include an empty list so we get 6 levels (12 bytes) + assert_eq!(rep_slicer.slice_rest().len(), 12); + + let mut def_slicer = repdefs.rep_slicer().unwrap(); + + // First 5 items include a null list so we get 6 levels (12 bytes) + assert_eq!(def_slicer.slice_next(5).len(), 12); + // Next 20 are all plain + assert_eq!(def_slicer.slice_next(20).len(), 40); + // Last 5 include an empty list so we get 6 levels (12 bytes) + assert_eq!(def_slicer.slice_rest().len(), 12); + } + + #[test] + fn test_control_words() { + // Convert to control words, verify expected, convert back, verify same as original + fn check( + rep: &[u16], + def: &[u16], + expected_values: Vec, + expected_bytes_per_word: usize, + expected_bits_rep: u8, + expected_bits_def: u8, + ) { + let num_vals = rep.len().max(def.len()); + let max_rep = rep.iter().max().copied().unwrap_or(0); + let max_def = def.iter().max().copied().unwrap_or(0); + + let in_rep = if rep.is_empty() { None } else { Some(rep) }; + let in_def = if def.is_empty() { None } else { Some(def) }; + + let mut iter = super::build_control_word_iterator( + in_rep, + max_rep, + in_def, + max_def, + max_def + 1, + expected_values.len(), + ); + assert_eq!(iter.bytes_per_word(), expected_bytes_per_word); + assert_eq!(iter.bits_rep(), expected_bits_rep); + assert_eq!(iter.bits_def(), expected_bits_def); + let mut cw_vec = Vec::with_capacity(num_vals * iter.bytes_per_word()); + + for _ in 0..num_vals { + iter.append_next(&mut cw_vec); + } + assert!(iter.append_next(&mut cw_vec).is_none()); + + assert_eq!(expected_values, cw_vec); + + let parser = super::ControlWordParser::new(expected_bits_rep, expected_bits_def); + + let mut rep_out = Vec::with_capacity(num_vals); + let mut def_out = Vec::with_capacity(num_vals); + + if expected_bytes_per_word > 0 { + for slice in cw_vec.chunks_exact(expected_bytes_per_word) { + parser.parse(slice, &mut rep_out, &mut def_out); + } + } + + assert_eq!(rep, rep_out.as_slice()); + assert_eq!(def, def_out.as_slice()); + } + + // Each will need 4 bits and so we should get 1-byte control words + let rep = &[0_u16, 7, 3, 2, 9, 8, 12, 5]; + let def = &[5_u16, 3, 1, 2, 12, 15, 0, 2]; + let expected = vec![ + 0b00000101, // 0, 5 + 0b01110011, // 7, 3 + 0b00110001, // 3, 1 + 0b00100010, // 2, 2 + 0b10011100, // 9, 12 + 0b10001111, // 8, 15 + 0b11000000, // 12, 0 + 0b01010010, // 5, 2 + ]; + check(rep, def, expected, 1, 4, 4); + + // Now we need 5 bits for def so we get 2-byte control words + let rep = &[0_u16, 7, 3, 2, 9, 8, 12, 5]; + let def = &[5_u16, 3, 1, 2, 12, 22, 0, 2]; + let expected = vec![ + 0b00000101, 0b00000000, // 0, 5 + 0b11100011, 0b00000000, // 7, 3 + 0b01100001, 0b00000000, // 3, 1 + 0b01000010, 0b00000000, // 2, 2 + 0b00101100, 0b00000001, // 9, 12 + 0b00010110, 0b00000001, // 8, 22 + 0b10000000, 0b00000001, // 12, 0 + 0b10100010, 0b00000000, // 5, 2 + ]; + check(rep, def, expected, 2, 4, 5); + + // Just rep, 4 bits so 1 byte each + let levels = &[0_u16, 7, 3, 2, 9, 8, 12, 5]; + let expected = vec![ + 0b00000000, // 0 + 0b00000111, // 7 + 0b00000011, // 3 + 0b00000010, // 2 + 0b00001001, // 9 + 0b00001000, // 8 + 0b00001100, // 12 + 0b00000101, // 5 + ]; + check(levels, &[], expected.clone(), 1, 4, 0); + + // Just def + check(&[], levels, expected, 1, 0, 4); + + // No rep, no def, no bytes + check(&[], &[], Vec::default(), 0, 0, 0); + } + + #[test] + fn test_control_words_def_only_all_zero_levels() { + // A page split can produce a page whose rows are all valid, so its + // definition levels are all zero and the max value in the (non-empty) + // buffer is 0. The decoder reads a zero-bit layout as NIL (zero bytes + // per control word), so the writer must emit the NIL iterator too: + // zero bytes per word, no writes, one new-row/visible item per level. + // Writing one all-zero byte per word (the pre-fix Unary behavior) + // would disagree with the decoder and corrupt the stream. + let def = [0_u16; 5]; + + let mut iter = super::build_control_word_iterator( + None, + 0, + Some(&def), + /*max_def=*/ 0, + /*max_visible_def=*/ u16::MAX, + def.len(), + ); + assert_eq!(iter.bytes_per_word(), 0); + assert_eq!(iter.bits_rep(), 0); + assert_eq!(iter.bits_def(), 0); + + let mut cw_vec = Vec::new(); + for _ in 0..def.len() { + let word_desc = iter.append_next(&mut cw_vec).unwrap(); + assert!(word_desc.is_new_row); + assert!(word_desc.is_visible); + } + assert!(iter.append_next(&mut cw_vec).is_none()); + // The NIL layout writes nothing, matching what ControlWordParser::new(0, 0) + // reads back (Self::NIL parses nothing). + assert!(cw_vec.is_empty()); + + // The parser side of the same layout must also be NIL: zero bits on + // both levels parse zero bytes and yield nothing. + let parser = super::ControlWordParser::new(0, 0); + let mut rep_out = Vec::new(); + let mut def_out = Vec::new(); + parser.parse(&[], &mut rep_out, &mut def_out); + assert!(rep_out.is_empty()); + assert!(def_out.is_empty()); + } + + #[test] + fn test_control_words_rep_only_all_zero_levels() { + // Same NIL-layout requirement as the definition-only case, for the + // repetition-only branches. Repetition levels are normally never all + // zero (the first entry of every list row is the schema max rep), so + // this is a defensive check that the writer matches the decoder if a + // zero-width rep buffer is ever produced. + let rep = [0_u16; 5]; + let mut iter = super::build_control_word_iterator( + Some(&rep), + /*max_rep=*/ 0, + None, + 0, + /*max_visible_def=*/ u16::MAX, + rep.len(), + ); + assert_eq!(iter.bytes_per_word(), 0); + let mut cw_vec = Vec::new(); + for _ in 0..rep.len() { + assert!(iter.append_next(&mut cw_vec).unwrap().is_new_row); + } + assert!(iter.append_next(&mut cw_vec).is_none()); + assert!(cw_vec.is_empty()); + } + + #[test] + fn test_control_words_both_levels_all_zero() { + // Both levels non-empty but all zero: the Binary writer would emit one + // zero byte per word while the decoder reads a NIL layout (zero bytes + // per word). The unified zero-width early return avoids that + // write/read asymmetry. + let rep = [0_u16; 5]; + let def = [0_u16; 5]; + let mut iter = super::build_control_word_iterator( + Some(&rep), + /*max_rep=*/ 0, + Some(&def), + /*max_def=*/ 0, + /*max_visible_def=*/ u16::MAX, + rep.len(), + ); + assert_eq!(iter.bytes_per_word(), 0); + assert_eq!(iter.bits_rep(), 0); + assert_eq!(iter.bits_def(), 0); + let mut cw_vec = Vec::new(); + for _ in 0..rep.len() { + assert!(iter.append_next(&mut cw_vec).unwrap().is_new_row); + } + assert!(iter.append_next(&mut cw_vec).is_none()); + assert!(cw_vec.is_empty()); + } + + #[test] + fn test_control_words_rep_index() { + fn check( + rep: &[u16], + def: &[u16], + expected_new_rows: Vec, + expected_is_visible: Vec, + ) { + let num_vals = rep.len().max(def.len()); + let max_rep = rep.iter().max().copied().unwrap_or(0); + let max_def = def.iter().max().copied().unwrap_or(0); + + let in_rep = if rep.is_empty() { None } else { Some(rep) }; + let in_def = if def.is_empty() { None } else { Some(def) }; + + let mut iter = super::build_control_word_iterator( + in_rep, + max_rep, + in_def, + max_def, + /*max_visible_def=*/ 2, + expected_new_rows.len(), + ); + + let mut cw_vec = Vec::with_capacity(num_vals * iter.bytes_per_word()); + let mut expected_new_rows = expected_new_rows.iter().copied(); + let mut expected_is_visible = expected_is_visible.iter().copied(); + for _ in 0..expected_new_rows.len() { + let word_desc = iter.append_next(&mut cw_vec).unwrap(); + assert_eq!(word_desc.is_new_row, expected_new_rows.next().unwrap()); + assert_eq!(word_desc.is_visible, expected_is_visible.next().unwrap()); + } + assert!(iter.append_next(&mut cw_vec).is_none()); + } + + // 2 means new list + let rep = &[2_u16, 1, 0, 2, 2, 0, 1, 1, 0, 2, 0]; + // These values don't matter for this test + let def = &[0_u16, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1]; + + // Rep & def + check( + rep, + def, + vec![ + true, false, false, true, true, false, false, false, false, true, false, + ], + vec![ + true, true, true, false, true, true, true, true, true, true, true, + ], + ); + // Rep only + check( + rep, + &[], + vec![ + true, false, false, true, true, false, false, false, false, true, false, + ], + vec![true; 11], + ); + // No repetition + check( + &[], + def, + vec![ + true, true, true, true, true, true, true, true, true, true, true, + ], + vec![true; 11], + ); + // No repetition, no definition + check( + &[], + &[], + vec![ + true, true, true, true, true, true, true, true, true, true, true, + ], + vec![true; 11], + ); + } + + #[test] + fn regress_empty_list_case() { + // This regresses a case where we had 3 null lists inside a struct + let mut builder = RepDefBuilder::default(); + builder.add_validity_bitmap(validity(&[true, false, true])); + builder.add_offsets( + offsets_32(&[0, 0, 0, 0]), + Some(validity(&[false, false, false])), + ); + builder.add_no_null(0); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([1, 1, 1], *rep); + assert_eq!([1, 2, 1], *def); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + Some(rep.as_ref().to_vec()), + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 0, + )]); + + assert_eq!(unraveler.unravel_validity(0).unwrap(), None); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 0, 0, 0]).inner()); + assert_eq!(val, Some(validity(&[false, false, false]))); + let val = unraveler.unravel_validity(3).unwrap().unwrap(); + assert_eq!(val.inner(), validity(&[true, false, true]).inner()); + } + + #[test] + fn regress_list_ends_null_case() { + let mut builder = RepDefBuilder::default(); + builder.add_offsets( + offsets_64(&[0, 1, 2, 2]), + Some(validity(&[true, true, false])), + ); + builder.add_offsets(offsets_64(&[0, 1, 1]), Some(validity(&[true, false]))); + builder.add_no_null(1); + + let repdefs = RepDefBuilder::serialize(vec![builder]); + let rep = repdefs.repetition_levels.unwrap(); + let def = repdefs.definition_levels.unwrap(); + + assert_eq!([2, 2, 2], *rep); + assert_eq!([0, 1, 2], *def); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new( + Some(rep.as_ref().to_vec()), + Some(def.as_ref().to_vec()), + repdefs.def_meaning.into(), + 1, + )]); + + assert_eq!(unraveler.unravel_validity(1).unwrap(), None); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 1, 1]).inner()); + assert_eq!(val, Some(validity(&[true, false]))); + let (off, val) = unraveler.unravel_offsets::().unwrap(); + assert_eq!(off.inner(), offsets_32(&[0, 1, 2, 2]).inner()); + assert_eq!(val, Some(validity(&[true, true, false]))); + } + + #[test] + fn test_mixed_unraveler() { + // This tests cases where the validity is different between two different pages + // because one page has nulls and the other doesn't. + + // Simple case with one layer of validity and no repetition + let mut unraveler = CompositeRepDefUnraveler::new(vec![ + RepDefUnraveler::new( + None, + Some(vec![0, 1, 0, 1]), + vec![DefinitionInterpretation::NullableItem].into(), + 4, + ), + RepDefUnraveler::new( + None, + None, + vec![DefinitionInterpretation::AllValidItem].into(), + 4, + ), + ]); + + assert_eq!( + unraveler.unravel_validity(8).unwrap(), + Some(validity(&[ + true, false, true, false, true, true, true, true + ])) + ); + + // More complex case with two layers of validity and repetition + let def1 = Some(vec![0, 1, 2]); + let rep1 = Some(vec![1, 0, 1]); + + let def2 = Some(vec![1, 0, 0]); + let rep2 = Some(vec![1, 1, 0]); + + let mut unraveler = CompositeRepDefUnraveler::new(vec![ + RepDefUnraveler::new( + rep1, + def1, + vec![ + DefinitionInterpretation::NullableItem, + DefinitionInterpretation::EmptyableList, + ] + .into(), + 2, + ), + RepDefUnraveler::new( + rep2, + def2, + vec![ + DefinitionInterpretation::AllValidItem, + DefinitionInterpretation::NullableList, + ] + .into(), + 2, + ), + ]); + + assert_eq!( + unraveler.unravel_validity(4).unwrap(), + Some(validity(&[true, false, true, true])) + ); + assert_eq!( + unraveler.unravel_offsets::().unwrap(), + ( + offsets_32(&[0, 2, 2, 2, 4]), + Some(validity(&[true, true, false, true])) + ) + ); + } + + #[test] + fn test_mixed_unraveler_nullable_without_def_levels() { + // A page can keep nullable layer metadata even when all definition levels are 0 + // and no definition buffer needs to be materialized. This should decode as all-valid. + let mut unraveler = CompositeRepDefUnraveler::new(vec![ + RepDefUnraveler::new( + None, + Some(vec![0, 1, 0, 1]), + vec![DefinitionInterpretation::NullableItem].into(), + 4, + ), + RepDefUnraveler::new( + None, + None, + vec![DefinitionInterpretation::NullableItem].into(), + 4, + ), + ]); + + assert_eq!( + unraveler.unravel_validity(8).unwrap(), + Some(validity(&[ + true, false, true, false, true, true, true, true + ])) + ); + } +} diff --git a/vendor/lance-encoding/src/statistics.rs b/vendor/lance-encoding/src/statistics.rs new file mode 100644 index 000000000..e312bc513 --- /dev/null +++ b/vendor/lance-encoding/src/statistics.rs @@ -0,0 +1,1269 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ + fmt::{self}, + hash::{Hash, RandomState}, + sync::Arc, +}; + +use arrow_array::{Array, ArrowPrimitiveType, UInt64Array, cast::AsArray, types::UInt64Type}; +use hyperloglogplus::{HyperLogLog, HyperLogLogPlus}; +use num_traits::PrimInt; + +use crate::data::{ + AllNullDataBlock, DataBlock, DictionaryDataBlock, FixedSizeListBlock, FixedWidthDataBlock, + NullableDataBlock, OpaqueBlock, StructDataBlock, VariableWidthBlock, +}; + +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +pub enum Stat { + BitWidth, + DataSize, + Cardinality, + FixedSize, + NullCount, + MaxLength, + RunCount, + BytePositionEntropy, +} + +impl fmt::Debug for Stat { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::BitWidth => write!(f, "BitWidth"), + Self::DataSize => write!(f, "DataSize"), + Self::Cardinality => write!(f, "Cardinality"), + Self::FixedSize => write!(f, "FixedSize"), + Self::NullCount => write!(f, "NullCount"), + Self::MaxLength => write!(f, "MaxLength"), + Self::RunCount => write!(f, "RunCount"), + Self::BytePositionEntropy => write!(f, "BytePositionEntropy"), + } + } +} + +impl fmt::Display for Stat { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub trait ComputeStat { + fn compute_stat(&mut self); +} + +impl ComputeStat for DataBlock { + fn compute_stat(&mut self) { + match self { + Self::Empty() => {} + Self::Constant(_) => {} + Self::AllNull(_) => {} + Self::Nullable(data_block) => data_block.data.compute_stat(), + Self::FixedWidth(data_block) => data_block.compute_stat(), + Self::FixedSizeList(data_block) => data_block.compute_stat(), + Self::VariableWidth(data_block) => data_block.compute_stat(), + Self::Opaque(data_block) => data_block.compute_stat(), + Self::Struct(data_block) => data_block.compute_stat(), + Self::Dictionary(_) => {} + } + } +} + +impl ComputeStat for VariableWidthBlock { + fn compute_stat(&mut self) { + if !self.block_info.0.read().unwrap().is_empty() { + panic!("compute_stat should only be called once during DataBlock construction"); + } + let data_size = self.data_size(); + let data_size_array = Arc::new(UInt64Array::from(vec![data_size])); + + let max_length_array = self.max_length(); + + let mut info = self.block_info.0.write().unwrap(); + info.insert(Stat::DataSize, data_size_array); + info.insert(Stat::MaxLength, max_length_array); + } +} + +impl ComputeStat for FixedWidthDataBlock { + fn compute_stat(&mut self) { + // compute this datablock's data_size + let data_size = self.data_size(); + let data_size_array = Arc::new(UInt64Array::from(vec![data_size])); + + // compute this datablock's max_bit_width + let max_bit_widths = self.max_bit_widths(); + + // the MaxLength of FixedWidthDataBlock is it's self.bits_per_value / 8 + let max_len = self.bits_per_value / 8; + let max_len_array = Arc::new(UInt64Array::from(vec![max_len])); + + // compute run count + let run_count_array = self.run_count(); + + // compute byte position entropy + let byte_position_entropy = self.byte_position_entropy(); + + let mut info = self.block_info.0.write().unwrap(); + info.insert(Stat::DataSize, data_size_array); + info.insert(Stat::BitWidth, max_bit_widths); + info.insert(Stat::MaxLength, max_len_array); + info.insert(Stat::RunCount, run_count_array); + info.insert(Stat::BytePositionEntropy, byte_position_entropy); + } +} + +impl ComputeStat for FixedSizeListBlock { + fn compute_stat(&mut self) { + // We leave the child stats unchanged. This may seem odd (e.g. should bit width be the + // bit width of the child * dimension?) but it's because we use these stats to determine + // compression and we are currently just compressing the child data. + // + // There is a potential opportunity here to do better. For example, if we have a FSL of + // 4 32-bit integers then we should probably treat them as a single 128-bit integer or maybe + // even 4 columns of 32-bit integers. This might yield better compression. + self.child.compute_stat(); + } +} + +impl ComputeStat for OpaqueBlock { + fn compute_stat(&mut self) { + // compute this datablock's data_size + let data_size = self.data_size(); + let data_size_array = Arc::new(UInt64Array::from(vec![data_size])); + let mut info = self.block_info.0.write().unwrap(); + info.insert(Stat::DataSize, data_size_array); + } +} + +pub trait GetStat: fmt::Debug { + fn get_stat(&self, stat: Stat) -> Option>; + + fn expect_stat(&self, stat: Stat) -> Arc { + self.get_stat(stat) + .unwrap_or_else(|| panic!("{:?} DataBlock does not have `{}` statistics.", self, stat)) + } + + fn expect_single_stat(&self, stat: Stat) -> T::Native { + let stat_value = self.expect_stat(stat); + let stat_value = stat_value.as_primitive::(); + if stat_value.len() != 1 { + panic!( + "{:?} DataBlock does not have exactly one value for `{} statistics.", + self, stat + ); + } + stat_value.value(0) + } +} + +impl GetStat for DataBlock { + fn get_stat(&self, stat: Stat) -> Option> { + match self { + Self::Empty() => None, + Self::Constant(_) => None, + Self::AllNull(data_block) => data_block.get_stat(stat), + Self::Nullable(data_block) => data_block.get_stat(stat), + Self::FixedWidth(data_block) => data_block.get_stat(stat), + Self::FixedSizeList(data_block) => data_block.get_stat(stat), + Self::VariableWidth(data_block) => data_block.get_stat(stat), + Self::Opaque(data_block) => data_block.get_stat(stat), + Self::Struct(data_block) => data_block.get_stat(stat), + Self::Dictionary(data_block) => data_block.get_stat(stat), + } + } +} + +// NullableDataBlock will be deprecated in Lance 2.1. +impl GetStat for NullableDataBlock { + // This function simply returns the statistics of the inner `DataBlock` of `NullableDataBlock`, + // this is not accurate but `NullableDataBlock` is going to be deprecated in Lance 2.1 anyway. + fn get_stat(&self, stat: Stat) -> Option> { + self.data.get_stat(stat) + } +} + +impl GetStat for VariableWidthBlock { + fn get_stat(&self, stat: Stat) -> Option> { + { + let block_info = self.block_info.0.read().unwrap(); + if block_info.is_empty() { + panic!("get_stat should be called after statistics are computed."); + } + if let Some(stat_value) = block_info.get(&stat) { + return Some(stat_value.clone()); + } + } + + if stat != Stat::Cardinality { + return None; + } + + let computed = self.compute_cardinality(); + let mut block_info = self.block_info.0.write().unwrap(); + if block_info.is_empty() { + panic!("get_stat should be called after statistics are computed."); + } + Some( + block_info + .entry(stat) + .or_insert_with(|| computed.clone()) + .clone(), + ) + } +} + +impl GetStat for FixedSizeListBlock { + fn get_stat(&self, stat: Stat) -> Option> { + let child_stat = self.child.get_stat(stat); + match stat { + Stat::MaxLength => child_stat.map(|max_length| { + // this is conservative when working with variable length data as we shouldn't assume + // that we have a list of all max-length elements but it's cheap and easy to calculate + let max_length = max_length.as_primitive::().value(0); + Arc::new(UInt64Array::from(vec![max_length * self.dimension])) as Arc + }), + _ => child_stat, + } + } +} + +impl VariableWidthBlock { + // Caveat: the computation here assumes VariableWidthBlock.offsets maps directly to VariableWidthBlock.data + // without any adjustment(for example, no null_adjustment for offsets) + fn compute_cardinality(&self) -> Arc { + const PRECISION: u8 = 4; + // The default hasher (currently sip hash 1-3) does not seem to give good results + // with HLL. + // + // In particular, when using randomly generated 12-byte strings, the HLL count was + // suggested a cardinality of 500 (out of 1000 unique items and hashes) at least 10% + // of the time. + // + // Using xxhash3 consistently gives better results. + let mut hll: HyperLogLogPlus<&[u8], xxhash_rust::xxh3::Xxh3Builder> = + HyperLogLogPlus::new(PRECISION, xxhash_rust::xxh3::Xxh3Builder::default()).unwrap(); + + match self.bits_per_offset { + 32 => { + let offsets_ref = self.offsets.borrow_to_typed_slice::(); + let offsets: &[u32] = offsets_ref.as_ref(); + + offsets + .iter() + .zip(offsets.iter().skip(1)) + .for_each(|(&start, &end)| { + hll.insert(&self.data[start as usize..end as usize]); + }); + let cardinality = hll.count() as u64; + Arc::new(UInt64Array::from(vec![cardinality])) + } + 64 => { + let offsets_ref = self.offsets.borrow_to_typed_slice::(); + let offsets: &[u64] = offsets_ref.as_ref(); + + offsets + .iter() + .zip(offsets.iter().skip(1)) + .for_each(|(&start, &end)| { + hll.insert(&self.data[start as usize..end as usize]); + }); + + let cardinality = hll.count() as u64; + Arc::new(UInt64Array::from(vec![cardinality])) + } + _ => { + unreachable!("the bits_per_offset of VariableWidthBlock can only be 32 or 64") + } + } + } + + fn max_length(&mut self) -> Arc { + match self.bits_per_offset { + 32 => { + let offsets = self.offsets.borrow_to_typed_slice::(); + let offsets = offsets.as_ref(); + let max_len = offsets + .windows(2) + .map(|pair| pair[1] - pair[0]) + .max() + .unwrap_or(0); + Arc::new(UInt64Array::from(vec![max_len as u64])) + } + 64 => { + let offsets = self.offsets.borrow_to_typed_slice::(); + let offsets = offsets.as_ref(); + let max_len = offsets + .windows(2) + .map(|pair| pair[1] - pair[0]) + .max() + .unwrap_or(0); + Arc::new(UInt64Array::from(vec![max_len])) + } + _ => { + unreachable!("the type of offsets in VariableWidth can only be u32 or u64"); + } + } + } +} + +impl GetStat for AllNullDataBlock { + fn get_stat(&self, stat: Stat) -> Option> { + match stat { + Stat::NullCount => { + let null_count = self.num_values; + Some(Arc::new(UInt64Array::from(vec![null_count]))) + } + Stat::DataSize => Some(Arc::new(UInt64Array::from(vec![0]))), + _ => None, + } + } +} + +impl GetStat for FixedWidthDataBlock { + fn get_stat(&self, stat: Stat) -> Option> { + { + let block_info = self.block_info.0.read().unwrap(); + + if block_info.is_empty() { + panic!("get_stat should be called after statistics are computed."); + } + + if let Some(stat_value) = block_info.get(&stat) { + return Some(stat_value.clone()); + } + } + + if stat == Stat::Cardinality && (self.bits_per_value == 64 || self.bits_per_value == 128) { + let computed = self.cardinality(); + let mut block_info = self.block_info.0.write().unwrap(); + Some( + block_info + .entry(stat) + .or_insert_with(|| computed.clone()) + .clone(), + ) + } else { + None + } + } +} + +impl FixedWidthDataBlock { + fn max_bit_widths(&mut self) -> Arc { + if self.num_values == 0 { + return Arc::new(UInt64Array::from(vec![0u64])); + } + + const CHUNK_SIZE: usize = 1024; + + fn calculate_max_bit_width(slice: &[T], bits_per_value: u64) -> Vec { + slice + .chunks(CHUNK_SIZE) + .map(|chunk| { + let max_value = chunk.iter().fold(T::zero(), |acc, &x| acc | x); + bits_per_value - max_value.leading_zeros() as u64 + }) + .collect() + } + + match self.bits_per_value { + 8 => { + let u8_slice = self.data.borrow_to_typed_slice::(); + let u8_slice = u8_slice.as_ref(); + Arc::new(UInt64Array::from(calculate_max_bit_width( + u8_slice, + self.bits_per_value, + ))) + } + 16 => { + let u16_slice = self.data.borrow_to_typed_slice::(); + let u16_slice = u16_slice.as_ref(); + Arc::new(UInt64Array::from(calculate_max_bit_width( + u16_slice, + self.bits_per_value, + ))) + } + 32 => { + let u32_slice = self.data.borrow_to_typed_slice::(); + let u32_slice = u32_slice.as_ref(); + Arc::new(UInt64Array::from(calculate_max_bit_width( + u32_slice, + self.bits_per_value, + ))) + } + 64 => { + let u64_slice = self.data.borrow_to_typed_slice::(); + let u64_slice = u64_slice.as_ref(); + Arc::new(UInt64Array::from(calculate_max_bit_width( + u64_slice, + self.bits_per_value, + ))) + } + _ => Arc::new(UInt64Array::from(vec![self.bits_per_value])), + } + } + + fn cardinality(&self) -> Arc { + match self.bits_per_value { + 64 => { + let u64_slice_ref = self.data.borrow_to_typed_slice::(); + let u64_slice = u64_slice_ref.as_ref(); + + const PRECISION: u8 = 4; + let mut hll: HyperLogLogPlus = + HyperLogLogPlus::new(PRECISION, xxhash_rust::xxh3::Xxh3Builder::default()) + .unwrap(); + for val in u64_slice { + hll.insert(val); + } + let cardinality = hll.count() as u64; + Arc::new(UInt64Array::from(vec![cardinality])) + } + 128 => { + let u128_slice_ref = self.data.borrow_to_typed_slice::(); + let u128_slice = u128_slice_ref.as_ref(); + + const PRECISION: u8 = 4; + let mut hll: HyperLogLogPlus = + HyperLogLogPlus::new(PRECISION, RandomState::new()).unwrap(); + for val in u128_slice { + hll.insert(val); + } + let cardinality = hll.count() as u64; + Arc::new(UInt64Array::from(vec![cardinality])) + } + _ => unreachable!(), + } + } + + /// Counts the number of runs (consecutive sequences of equal values) in the data. + /// + /// A "run" is defined as a sequence of one or more consecutive equal values. + /// For example: + /// - `[1, 1, 2, 2, 2, 3]` has 3 runs: [1,1], [2,2,2], and [3] + /// - `[1, 2, 3, 4]` has 4 runs (each value is its own run) + /// - `[5, 5, 5, 5]` has 1 run + /// + /// This count is used to determine if RLE compression would be effective. + /// Fewer runs relative to the total number of values indicates better RLE compression potential. + fn run_count(&mut self) -> Arc { + if self.num_values == 0 { + return Arc::new(UInt64Array::from(vec![0u64])); + } + + // Inner function to count runs in typed data + fn count_runs(slice: &[T]) -> u64 { + if slice.is_empty() { + return 0; + } + + // Start with 1 run (the first value) + let mut runs = 1u64; + let mut prev = slice[0]; + + // Count value transitions (each transition indicates a new run) + for &val in &slice[1..] { + if val != prev { + runs += 1; + prev = val; + } + } + + runs + } + + let run_count = match self.bits_per_value { + 8 => { + let u8_slice = self.data.borrow_to_typed_slice::(); + count_runs(u8_slice.as_ref()) + } + 16 => { + let u16_slice = self.data.borrow_to_typed_slice::(); + count_runs(u16_slice.as_ref()) + } + 32 => { + let u32_slice = self.data.borrow_to_typed_slice::(); + count_runs(u32_slice.as_ref()) + } + 64 => { + let u64_slice = self.data.borrow_to_typed_slice::(); + count_runs(u64_slice.as_ref()) + } + 128 => { + let u128_slice = self.data.borrow_to_typed_slice::(); + count_runs(u128_slice.as_ref()) + } + _ => self.num_values, // For other bit widths, assume no runs + }; + + Arc::new(UInt64Array::from(vec![run_count])) + } + + /// Calculates entropy for each byte position. + /// Returns an array with entropy values for each byte position (scaled by 1000 for integer storage). + /// Lower entropy in specific byte positions indicates better suitability for BSS. + fn byte_position_entropy(&mut self) -> Arc { + const SAMPLE_SIZE: usize = 64; // Sample more values for better entropy estimation + + // Get sample size (min of data length and SAMPLE_SIZE) + let sample_count = (self.num_values as usize).min(SAMPLE_SIZE); + + if sample_count == 0 { + // Return empty array for empty data + return Arc::new(UInt64Array::from(vec![] as Vec)); + } + + let bytes_per_value = (self.bits_per_value / 8) as usize; + let mut entropies = Vec::with_capacity(bytes_per_value); + + // Calculate entropy for each byte position + for pos in 0..bytes_per_value { + let mut byte_counts = [0u32; 256]; + + // Count occurrences of each byte value at this position + for i in 0..sample_count { + let byte_offset = i * bytes_per_value + pos; + if byte_offset < self.data.len() { + byte_counts[self.data[byte_offset] as usize] += 1; + } + } + + // Calculate Shannon entropy for this position + let mut entropy = 0.0f64; + let total = sample_count as f64; + + for &count in &byte_counts { + if count > 0 { + let p = count as f64 / total; + entropy -= p * p.log2(); + } + } + + // Scale by 1000 and store as integer for efficient storage + entropies.push((entropy * 1000.0) as u64); + } + + Arc::new(UInt64Array::from(entropies)) + } +} + +impl GetStat for OpaqueBlock { + fn get_stat(&self, stat: Stat) -> Option> { + let block_info = self.block_info.0.read().unwrap(); + + if block_info.is_empty() { + panic!("get_stat should be called after statistics are computed."); + } + block_info.get(&stat).cloned() + } +} + +impl GetStat for DictionaryDataBlock { + fn get_stat(&self, _stat: Stat) -> Option> { + None + } +} + +impl GetStat for StructDataBlock { + fn get_stat(&self, stat: Stat) -> Option> { + let block_info = self.block_info.0.read().unwrap(); + if block_info.is_empty() { + panic!("get_stat should be called after statistics are computed.") + } + block_info.get(&stat).cloned() + } +} + +impl ComputeStat for StructDataBlock { + fn compute_stat(&mut self) { + let data_size = self.data_size(); + let data_size_array = Arc::new(UInt64Array::from(vec![data_size])); + + let max_len = self + .children + .iter() + .map(|child| child.expect_single_stat::(Stat::MaxLength)) + .sum::(); + let max_len_array = Arc::new(UInt64Array::from(vec![max_len])); + + let mut info = self.block_info.0.write().unwrap(); + info.insert(Stat::DataSize, data_size_array); + info.insert(Stat::MaxLength, max_len_array); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{ + ArrayRef, Int8Array, Int16Array, Int32Array, Int64Array, LargeStringArray, StringArray, + UInt8Array, UInt16Array, UInt32Array, UInt64Array, + }; + use arrow_schema::{DataType, Field}; + use lance_arrow::DataTypeExt; + use lance_datagen::{ArrayGeneratorExt, DEFAULT_SEED, RowCount, array}; + use rand::SeedableRng; + + use crate::statistics::{GetStat, Stat}; + + use super::DataBlock; + + use arrow_array::{ + Array, + cast::AsArray, + types::{Int32Type, UInt64Type}, + }; + use arrow_select::concat::concat; + #[test] + fn test_data_size_stat() { + let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0); + let mut genn = array::rand::().with_nulls(&[false, false, false]); + let arr1 = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let arr2 = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let arr3 = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_arrays(&[arr1.clone(), arr2.clone(), arr3.clone()], 9); + + let concatenated_array = concat(&[ + &*Arc::new(arr1.clone()) as &dyn Array, + &*Arc::new(arr2.clone()) as &dyn Array, + &*Arc::new(arr3.clone()) as &dyn Array, + ]) + .unwrap(); + + let data_size = block.expect_single_stat::(Stat::DataSize); + + let total_buffer_size: usize = concatenated_array + .to_data() + .buffers() + .iter() + .map(|buffer| buffer.len()) + .sum(); + assert!(data_size == total_buffer_size as u64); + + // test DataType::Binary + let mut genn = lance_datagen::array::rand_type(&DataType::Binary); + let arr = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + let data_size = block.expect_single_stat::(Stat::DataSize); + + let total_buffer_size: usize = arr + .to_data() + .buffers() + .iter() + .map(|buffer| buffer.len()) + .sum(); + assert!(data_size == total_buffer_size as u64); + + // test DataType::Struct + let fields = vec![ + Arc::new(Field::new("int_field", DataType::Int32, false)), + Arc::new(Field::new("float_field", DataType::Float32, false)), + ] + .into(); + + let mut genn = lance_datagen::array::rand_type(&DataType::Struct(fields)); + let arr = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + let (_, arr_parts, _) = arr.as_struct().clone().into_parts(); + let total_buffer_size: usize = arr_parts + .iter() + .map(|arr| { + arr.to_data() + .buffers() + .iter() + .map(|buffer| buffer.len()) + .sum::() + }) + .sum(); + let data_size = block.expect_single_stat::(Stat::DataSize); + assert!(data_size == total_buffer_size as u64); + + // test DataType::Dictionary + let mut genn = array::rand_type(&DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8), + )); + let arr = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + assert!(block.get_stat(Stat::DataSize).is_none()); + + let mut genn = array::rand::().with_nulls(&[false, true, false]); + let arr = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + let data_size = block.expect_single_stat::(Stat::DataSize); + let total_buffer_size: usize = arr + .to_data() + .buffers() + .iter() + .map(|buffer| buffer.len()) + .sum(); + + assert!(data_size == total_buffer_size as u64); + } + + #[test] + fn test_bit_width_stat_for_integers() { + let int8_array = Int8Array::from(vec![1, 2, 3]); + let array_ref: ArrayRef = Arc::new(int8_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![2])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref(),); + + let int8_array = Int8Array::from(vec![0x1, 0x2, 0x3, 0x7F]); + let array_ref: ArrayRef = Arc::new(int8_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![7])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref(),); + + let int8_array = Int8Array::from(vec![0x1, 0x2, 0x3, 0xF, 0x1F]); + let array_ref: ArrayRef = Arc::new(int8_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![5])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref(),); + + let int8_array = Int8Array::from(vec![-1, 2, 3]); + let array_ref: ArrayRef = Arc::new(int8_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![8])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int16_array = Int16Array::from(vec![1, 2, 3]); + let array_ref: ArrayRef = Arc::new(int16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![2])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int16_array = Int16Array::from(vec![0x1, 0x2, 0x3, 0x7F]); + let array_ref: ArrayRef = Arc::new(int16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![7])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int16_array = Int16Array::from(vec![0x1, 0x2, 0x3, 0xFF]); + let array_ref: ArrayRef = Arc::new(int16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![8])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int16_array = Int16Array::from(vec![0x1, 0x2, 0x3, 0x1FF]); + let array_ref: ArrayRef = Arc::new(int16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![9])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int16_array = Int16Array::from(vec![0x1, 0x2, 0x3, 0xF, 0x1F]); + let array_ref: ArrayRef = Arc::new(int16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![5])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int16_array = Int16Array::from(vec![-1, 2, 3]); + let array_ref: ArrayRef = Arc::new(int16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![16])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int32_array = Int32Array::from(vec![1, 2, 3]); + let array_ref: ArrayRef = Arc::new(int32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![2])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int32_array = Int32Array::from(vec![0x1, 0x2, 0x3, 0xFF]); + let array_ref: ArrayRef = Arc::new(int32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![8])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int32_array = Int32Array::from(vec![0x1, 0x2, 0x3, 0xFF, 0x1FF]); + let array_ref: ArrayRef = Arc::new(int32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![9])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int32_array = Int32Array::from(vec![-1, 2, 3]); + let array_ref: ArrayRef = Arc::new(int32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![32])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int32_array = Int32Array::from(vec![-1, 2, 3, -88]); + let array_ref: ArrayRef = Arc::new(int32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![32])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int64_array = Int64Array::from(vec![1, 2, 3]); + let array_ref: ArrayRef = Arc::new(int64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![2])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int64_array = Int64Array::from(vec![0x1, 0x2, 0x3, 0xFF]); + let array_ref: ArrayRef = Arc::new(int64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![8])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int64_array = Int64Array::from(vec![0x1, 0x2, 0x3, 0xFF, 0x1FF]); + let array_ref: ArrayRef = Arc::new(int64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![9])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int64_array = Int64Array::from(vec![-1, 2, 3]); + let array_ref: ArrayRef = Arc::new(int64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![64])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let int64_array = Int64Array::from(vec![-1, 2, 3, -88]); + let array_ref: ArrayRef = Arc::new(int64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![64])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint8_array = UInt8Array::from(vec![1, 2, 3]); + let array_ref: ArrayRef = Arc::new(uint8_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![2])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint8_array = UInt8Array::from(vec![0x1, 0x2, 0x3, 0x7F]); + let array_ref: ArrayRef = Arc::new(uint8_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![7])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint8_array = UInt8Array::from(vec![0x1, 0x2, 0x3, 0xF, 0x1F]); + let array_ref: ArrayRef = Arc::new(uint8_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![5])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint8_array = UInt8Array::from(vec![1, 2, 3, 0xF]); + let array_ref: ArrayRef = Arc::new(uint8_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![4])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint16_array = UInt16Array::from(vec![1, 2, 3]); + let array_ref: ArrayRef = Arc::new(uint16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![2])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint16_array = UInt16Array::from(vec![0x1, 0x2, 0x3, 0x7F]); + let array_ref: ArrayRef = Arc::new(uint16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![7])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint16_array = UInt16Array::from(vec![0x1, 0x2, 0x3, 0xFF]); + let array_ref: ArrayRef = Arc::new(uint16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![8])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint16_array = UInt16Array::from(vec![0x1, 0x2, 0x3, 0x1FF]); + let array_ref: ArrayRef = Arc::new(uint16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![9])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint16_array = UInt16Array::from(vec![0x1, 0x2, 0x3, 0xF, 0x1F]); + let array_ref: ArrayRef = Arc::new(uint16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![5])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint16_array = UInt16Array::from(vec![1, 2, 3, 0xFFFF]); + let array_ref: ArrayRef = Arc::new(uint16_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![16])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint32_array = UInt32Array::from(vec![1, 2, 3]); + let array_ref: ArrayRef = Arc::new(uint32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![2])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint32_array = UInt32Array::from(vec![0x1, 0x2, 0x3, 0xFF]); + let array_ref: ArrayRef = Arc::new(uint32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![8])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref(),); + + let uint32_array = UInt32Array::from(vec![0x1, 0x2, 0x3, 0xFF, 0x1FF]); + let array_ref: ArrayRef = Arc::new(uint32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![9])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint32_array = UInt32Array::from(vec![1, 2, 3, 0xF]); + let array_ref: ArrayRef = Arc::new(uint32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![4])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint32_array = UInt32Array::from(vec![1, 2, 3, 0x77]); + let array_ref: ArrayRef = Arc::new(uint32_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![7])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint64_array = UInt64Array::from(vec![1, 2, 3]); + let array_ref: ArrayRef = Arc::new(uint64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![2])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint64_array = UInt64Array::from(vec![0x1, 0x2, 0x3, 0xFF]); + let array_ref: ArrayRef = Arc::new(uint64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![8])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint64_array = UInt64Array::from(vec![0x1, 0x2, 0x3, 0xFF, 0x1FF]); + let array_ref: ArrayRef = Arc::new(uint64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![9])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint64_array = UInt64Array::from(vec![0, 2, 3, 0xFFFF]); + let array_ref: ArrayRef = Arc::new(uint64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![16])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + + let uint64_array = UInt64Array::from(vec![1, 2, 3, 0xFFFF_FFFF_FFFF_FFFF]); + let array_ref: ArrayRef = Arc::new(uint64_array); + let block = DataBlock::from_array(array_ref); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![64])) as ArrayRef; + let actual_bit_width = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_width.as_ref(), expected_bit_width.as_ref()); + } + + #[test] + fn test_bit_width_stat_more_than_1024() { + for data_type in [ + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + ] { + let array1 = Int64Array::from(vec![3; 1024]); + let array2 = Int64Array::from(vec![8; 1024]); + let array3 = Int64Array::from(vec![-1; 10]); + let array1 = arrow_cast::cast(&array1, &data_type).unwrap(); + let array2 = arrow_cast::cast(&array2, &data_type).unwrap(); + let array3 = arrow_cast::cast(&array3, &data_type).unwrap(); + + let arrays: Vec<&dyn arrow_array::Array> = + vec![array1.as_ref(), array2.as_ref(), array3.as_ref()]; + let concatenated = concat(&arrays).unwrap(); + let block = DataBlock::from_array(concatenated.clone()); + + let expected_bit_width = Arc::new(UInt64Array::from(vec![ + 2, + 4, + (data_type.byte_width() * 8) as u64, + ])) as ArrayRef; + let actual_bit_widths = block.expect_stat(Stat::BitWidth); + assert_eq!(actual_bit_widths.as_ref(), expected_bit_width.as_ref(),); + } + } + + #[test] + fn test_bit_width_when_none() { + let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0); + let mut genn = lance_datagen::array::rand_type(&DataType::Binary); + let arr = genn.generate(RowCount::from(3), &mut rng).unwrap(); + let block = DataBlock::from_array(arr.clone()); + assert!(block.get_stat(Stat::BitWidth).is_none(),); + } + + #[test] + fn test_cardinality_variable_width_datablock() { + let string_array = StringArray::from(vec![Some("hello"), Some("world")]); + let block = DataBlock::from_array(string_array); + let expected_cardinality = 2; + let actual_cardinality = block.expect_single_stat::(Stat::Cardinality); + assert_eq!(actual_cardinality, expected_cardinality,); + + let string_array = StringArray::from(vec![ + Some("to be named by variables"), + Some("to be passed as arguments to procedures"), + Some("to be returned as values of procedures"), + ]); + let block = DataBlock::from_array(string_array); + let expected_cardinality = 3; + let actual_cardinality = block.expect_single_stat::(Stat::Cardinality); + + assert_eq!(actual_cardinality, expected_cardinality,); + + let string_array = StringArray::from(vec![ + Some("Samuel Eilenberg"), + Some("Saunders Mac Lane"), + Some("Samuel Eilenberg"), + ]); + let block = DataBlock::from_array(string_array); + let expected_cardinality = 2; + let actual_cardinality = block.expect_single_stat::(Stat::Cardinality); + assert_eq!(actual_cardinality, expected_cardinality,); + + let string_array = LargeStringArray::from(vec![Some("hello"), Some("world")]); + let block = DataBlock::from_array(string_array); + let expected_cardinality = 2; + let actual_cardinality = block.expect_single_stat::(Stat::Cardinality); + assert_eq!(actual_cardinality, expected_cardinality,); + + let string_array = LargeStringArray::from(vec![ + Some("to be named by variables"), + Some("to be passed as arguments to procedures"), + Some("to be returned as values of procedures"), + ]); + let block = DataBlock::from_array(string_array); + let expected_cardinality = 3; + let actual_cardinality = block.expect_single_stat::(Stat::Cardinality); + assert_eq!(actual_cardinality, expected_cardinality,); + + let string_array = LargeStringArray::from(vec![ + Some("Samuel Eilenberg"), + Some("Saunders Mac Lane"), + Some("Samuel Eilenberg"), + ]); + let block = DataBlock::from_array(string_array); + let expected_cardinality = 2; + let actual_cardinality = block.expect_single_stat::(Stat::Cardinality); + assert_eq!(actual_cardinality, expected_cardinality,); + } + + #[test] + fn test_max_length_variable_width_datablock() { + let string_array = StringArray::from(vec![Some("hello"), Some("world")]); + let block = DataBlock::from_array(string_array.clone()); + let expected_max_length = string_array.value_length(0) as u64; + let actual_max_length = block.expect_single_stat::(Stat::MaxLength); + assert_eq!(actual_max_length, expected_max_length); + + let string_array = StringArray::from(vec![ + Some("to be named by variables"), + Some("to be passed as arguments to procedures"), // string that has max length + Some("to be returned as values of procedures"), + ]); + let block = DataBlock::from_array(string_array.clone()); + let expected_max_length = string_array.value_length(1) as u64; + let actual_max_length = block.expect_single_stat::(Stat::MaxLength); + assert_eq!(actual_max_length, expected_max_length); + + let string_array = StringArray::from(vec![ + Some("Samuel Eilenberg"), + Some("Saunders Mac Lane"), // string that has max length + Some("Samuel Eilenberg"), + ]); + let block = DataBlock::from_array(string_array.clone()); + let expected_max_length = string_array.value_length(1) as u64; + let actual_max_length = block.expect_single_stat::(Stat::MaxLength); + assert_eq!(actual_max_length, expected_max_length); + + let string_array = LargeStringArray::from(vec![Some("hello"), Some("world")]); + let block = DataBlock::from_array(string_array.clone()); + let expected_max_length = string_array.value_length(1) as u64; + let actual_max_length = block.expect_single_stat::(Stat::MaxLength); + assert_eq!(actual_max_length, expected_max_length); + + let string_array = LargeStringArray::from(vec![ + Some("to be named by variables"), + Some("to be passed as arguments to procedures"), // string that has max length + Some("to be returned as values of procedures"), + ]); + let block = DataBlock::from_array(string_array.clone()); + let expected_max_length = string_array.value(1).len() as u64; + let actual_max_length = block.expect_single_stat::(Stat::MaxLength); + + assert_eq!(actual_max_length, expected_max_length); + } + + #[test] + fn test_run_count_stat() { + // Test with highly repetitive data + let int32_array = Int32Array::from(vec![1, 1, 1, 2, 2, 2, 3, 3, 3]); + let block = DataBlock::from_array(int32_array); + let expected_run_count = 3; + let actual_run_count = block.expect_single_stat::(Stat::RunCount); + assert_eq!(actual_run_count, expected_run_count); + + // Test with no repetition + let int32_array = Int32Array::from(vec![1, 2, 3, 4, 5]); + let block = DataBlock::from_array(int32_array); + let expected_run_count = 5; + let actual_run_count = block.expect_single_stat::(Stat::RunCount); + assert_eq!(actual_run_count, expected_run_count); + + // Test with mixed pattern + let int32_array = Int32Array::from(vec![1, 1, 2, 3, 3, 3, 4, 5, 5]); + let block = DataBlock::from_array(int32_array); + let expected_run_count = 5; + let actual_run_count = block.expect_single_stat::(Stat::RunCount); + assert_eq!(actual_run_count, expected_run_count); + + // Test with single value + let int32_array = Int32Array::from(vec![42, 42, 42, 42, 42]); + let block = DataBlock::from_array(int32_array); + let expected_run_count = 1; + let actual_run_count = block.expect_single_stat::(Stat::RunCount); + assert_eq!(actual_run_count, expected_run_count); + + // Test with different data types + let uint8_array = UInt8Array::from(vec![1, 1, 2, 2, 3, 3]); + let block = DataBlock::from_array(uint8_array); + let expected_run_count = 3; + let actual_run_count = block.expect_single_stat::(Stat::RunCount); + assert_eq!(actual_run_count, expected_run_count); + + let int64_array = Int64Array::from(vec![100, 100, 200, 300, 300]); + let block = DataBlock::from_array(int64_array); + let expected_run_count = 3; + let actual_run_count = block.expect_single_stat::(Stat::RunCount); + assert_eq!(actual_run_count, expected_run_count); + } + + #[test] + fn test_fixed_width_cardinality_is_lazy() { + let int64_array = Int64Array::from(vec![1, 2, 3, 1, 2, 3, 1]); + let block = DataBlock::from_array(int64_array); + + let DataBlock::FixedWidth(fixed) = &block else { + panic!("Expected FixedWidth datablock"); + }; + + let info = fixed.block_info.0.read().unwrap(); + assert!(info.contains_key(&Stat::DataSize)); + assert!(info.contains_key(&Stat::BitWidth)); + assert!(!info.contains_key(&Stat::Cardinality)); + } + + #[test] + fn test_fixed_width_cardinality_computed_on_demand() { + let int64_array = Int64Array::from(vec![1, 2, 3, 1, 2, 3, 1]); + let block = DataBlock::from_array(int64_array); + + let cardinality = block.expect_single_stat::(Stat::Cardinality); + assert_eq!(cardinality, 3); + + let DataBlock::FixedWidth(fixed) = &block else { + panic!("Expected FixedWidth datablock"); + }; + + let info = fixed.block_info.0.read().unwrap(); + assert!(info.contains_key(&Stat::Cardinality)); + } + + #[test] + fn test_variable_width_cardinality_is_lazy() { + let string_array = StringArray::from(vec!["a", "b", "a"]); + let block = DataBlock::from_array(string_array); + + let DataBlock::VariableWidth(var) = &block else { + panic!("Expected VariableWidth datablock"); + }; + + { + let info = var.block_info.0.read().unwrap(); + assert!(info.contains_key(&Stat::DataSize)); + assert!(info.contains_key(&Stat::MaxLength)); + assert!(!info.contains_key(&Stat::Cardinality)); + } + + let cardinality = block.expect_single_stat::(Stat::Cardinality); + assert_eq!(cardinality, 2); + + let info = var.block_info.0.read().unwrap(); + assert!(info.contains_key(&Stat::Cardinality)); + } +} diff --git a/vendor/lance-encoding/src/testing.rs b/vendor/lance-encoding/src/testing.rs new file mode 100644 index 000000000..22f7efe8e --- /dev/null +++ b/vendor/lance-encoding/src/testing.rs @@ -0,0 +1,1588 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{cmp::Ordering, collections::HashMap, ops::Range, sync::Arc}; + +use crate::{ + decoder::DecoderConfig, + encodings::physical::block::CompressionScheme, + format::pb21::{ + BufferCompression, CompressiveEncoding, PageLayout, compressive_encoding::Compression, + }, +}; + +use arrow_array::{Array, StructArray, UInt64Array, make_array}; +use arrow_data::transform::{Capacities, MutableArrayData}; +use arrow_ord::ord::make_comparator; +use arrow_schema::{DataType, Field, Field as ArrowField, FieldRef, Schema, SortOptions}; +use arrow_select::concat::concat; +use bytes::{Bytes, BytesMut}; +use futures::{FutureExt, StreamExt, future::BoxFuture}; +use log::{debug, info, trace}; +use tokio::sync::mpsc::{self, UnboundedSender}; + +use lance_core::{Error, Result, datatypes::Field as LanceField, utils::bit::pad_bytes}; +use lance_datagen::{ArrayGenerator, RowCount, Seed, array, gen_batch}; + +use crate::compression::try_packed_struct_per_value; +use crate::{ + EncodingsIo, + buffer::LanceBuffer, + compression::{ + BlockCompressor, CompressionStrategy, field_metadata_params, finalize_miniblock_compressor, + reject_packed_struct_per_value, try_bitpacking_block, try_bitpacking_miniblock, + try_byte_stream_split_miniblock, try_child_rle_miniblock, + try_fixed_packed_struct_miniblock, try_fixed_u8_rle_block, try_fixed_u8_rle_miniblock, + try_general_block, try_raw_block, try_raw_fixed_size_list_miniblock, + try_raw_fixed_width_miniblock, try_raw_per_value, try_uncompressed_fixed_width_miniblock, + try_variable_rle_block, try_variable_width_miniblock, try_variable_width_per_value, + }, + compression_config::{CompressionFieldParams, CompressionParams}, + data::DataBlock, + decoder::{ + ColumnInfo, DecodeBatchScheduler, DecoderMessage, DecoderPlugins, FilterExpression, + PageInfo, create_decode_stream, + }, + encoder::{ + ColumnIndexSequence, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, + FieldEncodingContext, FieldEncodingStrategy, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers, + structural::{ + PrimitiveFieldEncoding, PrimitivePageEncoding, try_create_binary_blob, try_create_list, + try_create_map, try_create_struct, try_create_structural_blob, + try_create_structural_fixed_size_list, + }, + }, + encodings::logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, + repdef::RepDefBuilder, +}; + +const MAX_PAGE_BYTES: u64 = 32 * 1024 * 1024; +const TEST_ALIGNMENT: usize = MIN_PAGE_BUFFER_ALIGNMENT as usize; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TestEncoding { + Array, + StructuralU16, + StructuralU32, + StructuralSparse, +} + +impl TestEncoding { + fn all() -> impl Iterator { + [ + Self::Array, + Self::StructuralU16, + Self::StructuralU32, + Self::StructuralSparse, + ] + .into_iter() + } + + fn is_structural(self) -> bool { + self != Self::Array + } +} + +impl std::fmt::Display for TestEncoding { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Array => write!(f, "array"), + Self::StructuralU16 => write!(f, "structural-u16"), + Self::StructuralU32 => write!(f, "structural-u32"), + Self::StructuralSparse => write!(f, "structural-sparse"), + } + } +} + +#[derive(Debug, Clone)] +struct TestCompressionStrategy { + encoding: TestEncoding, + params: CompressionParams, +} + +impl TestCompressionStrategy { + fn field_params(&self, field: &LanceField) -> CompressionFieldParams { + let mut params = self + .params + .get_field_params(&field.name, &field.data_type()); + let mut metadata = field_metadata_params(field); + if self.encoding == TestEncoding::StructuralU16 + && metadata + .minichunk_size + .is_some_and(|size| size >= 32 * 1024) + { + metadata.minichunk_size = None; + } + params.merge(&metadata); + params + } +} + +impl CompressionStrategy for TestCompressionStrategy { + fn create_miniblock_compressor( + &self, + field: &LanceField, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + let compressor = + if let Some(compressor) = try_uncompressed_fixed_width_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_byte_stream_split_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = match self.encoding { + TestEncoding::StructuralSparse => try_child_rle_miniblock(data, ¶ms), + TestEncoding::Array | TestEncoding::StructuralU16 | TestEncoding::StructuralU32 => { + try_fixed_u8_rle_miniblock(data, ¶ms) + } + } { + compressor + } else if let Some(compressor) = try_bitpacking_miniblock(data) { + compressor + } else if let Some(compressor) = try_raw_fixed_width_miniblock(data) { + compressor + } else if let Some(compressor) = try_variable_width_miniblock(field, data, ¶ms)? { + compressor + } else if let Some(compressor) = try_fixed_packed_struct_miniblock(data)? { + compressor + } else if let Some(compressor) = try_raw_fixed_size_list_miniblock(data) { + compressor + } else { + return Err(lance_core::Error::not_supported_source( + format!( + "Mini-block compression not yet supported for block type {}", + data.name() + ) + .into(), + )); + }; + finalize_miniblock_compressor(data, compressor, ¶ms) + } + + fn create_per_value( + &self, + field: &LanceField, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + if let Some(compressor) = try_raw_per_value(data) { + return Ok(compressor); + } + let packed = match self.encoding { + TestEncoding::Array | TestEncoding::StructuralU16 => { + reject_packed_struct_per_value(field, data)? + } + TestEncoding::StructuralU32 | TestEncoding::StructuralSparse => { + try_packed_struct_per_value(Arc::new(self.clone()), field, data)? + } + }; + if let Some(compressor) = packed { + return Ok(compressor); + } + if let Some(compressor) = try_variable_width_per_value(field, data, ¶ms)? { + return Ok(compressor); + } + Err(lance_core::Error::not_supported_source( + format!( + "Per-value compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } + + fn create_block_compressor( + &self, + field: &LanceField, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + let rle = match self.encoding { + TestEncoding::Array | TestEncoding::StructuralU16 => None, + TestEncoding::StructuralU32 => try_fixed_u8_rle_block(data, ¶ms)?, + TestEncoding::StructuralSparse => try_variable_rle_block(data, ¶ms)?, + }; + if let Some(compressor) = rle { + return Ok(compressor); + } + if let Some(compressor) = try_bitpacking_block(data) { + return Ok(compressor); + } + if matches!( + self.encoding, + TestEncoding::StructuralU32 | TestEncoding::StructuralSparse + ) && let Some(compressor) = try_general_block(data, ¶ms)? + { + return Ok(compressor); + } + if let Some(compressor) = try_raw_block(data) { + return Ok(compressor); + } + Err(lance_core::Error::not_supported_source( + format!( + "Block compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } +} + +pub fn test_compression_strategy( + encoding: TestEncoding, + params: CompressionParams, +) -> Arc { + Arc::new(TestCompressionStrategy { encoding, params }) +} + +#[derive(Debug)] +struct TestFieldEncodingStrategy { + encoding: TestEncoding, + primitive: PrimitiveFieldEncoding, +} + +impl FieldEncodingStrategy for TestFieldEncodingStrategy { + fn create_field_encoder( + &self, + field: &LanceField, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result> { + if let Some(encoder) = + try_create_binary_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if self.encoding != TestEncoding::StructuralU16 + && let Some(encoder) = + try_create_structural_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if field.is_blob() { + return Err(Error::invalid_input_source( + format!( + "Blob encoding is not available for field '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )); + } + if self.encoding != TestEncoding::StructuralU16 { + if let Some(encoder) = try_create_map(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = + try_create_structural_fixed_size_list(field, column_index, context)? + { + return Ok(encoder); + } + } + if let Some(encoder) = self.primitive.try_create(field, column_index, context)? { + return Ok(encoder); + } + if self.encoding == TestEncoding::StructuralU16 { + if matches!( + field.data_type(), + DataType::FixedSizeList(item, _) + if matches!(item.data_type(), DataType::Struct(_)) + ) { + return Err(Error::not_supported_source( + "FixedSizeList is not enabled by the selected file format".into(), + )); + } + if matches!(field.data_type(), DataType::Map(_, _)) { + return Err(Error::not_supported_source( + "Map data type is not enabled by the selected file format".into(), + )); + } + } + if let Some(encoder) = try_create_list(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_struct(field, column_index, context)? { + return Ok(encoder); + } + Err(Error::not_supported_source( + format!( + "{} has no field encoding for '{}' with data type {}", + self.encoding, + field.name, + field.data_type() + ) + .into(), + )) + } +} + +pub fn test_encoding_strategy(encoding: TestEncoding) -> Box { + if encoding == TestEncoding::Array { + return Box::new(crate::array_encoding::ArrayFieldEncodingStrategy::new()); + } + + let compression = test_compression_strategy(encoding, CompressionParams::default()); + let page_encodings = match encoding { + TestEncoding::StructuralU16 => vec![ + PrimitivePageEncoding::reject_sparse(), + PrimitivePageEncoding::dense_u16(compression), + ], + TestEncoding::StructuralU32 => vec![ + PrimitivePageEncoding::reject_sparse(), + PrimitivePageEncoding::constant(), + PrimitivePageEncoding::dense_u32(compression), + ], + TestEncoding::StructuralSparse => vec![ + PrimitivePageEncoding::sparse(compression.clone()), + PrimitivePageEncoding::constant(), + PrimitivePageEncoding::dense_u32(compression), + ], + TestEncoding::Array => unreachable!(), + }; + Box::new(TestFieldEncodingStrategy { + encoding, + primitive: PrimitiveFieldEncoding::new(page_encodings), + }) +} + +pub fn create_test_field_encoder( + strategy: &dyn FieldEncodingStrategy, + field: &lance_core::datatypes::Field, + column_index: &mut ColumnIndexSequence, + options: &EncodingOptions, +) -> Result> { + let context = FieldEncodingContext { + strategy, + options, + root_field_metadata: &field.metadata, + }; + strategy.create_field_encoder(field, column_index, &context) +} + +#[derive(Debug)] +pub(crate) struct SimulatedScheduler { + data: Bytes, +} + +impl SimulatedScheduler { + pub fn new(data: Bytes) -> Self { + Self { data } + } +} + +impl EncodingsIo for SimulatedScheduler { + fn submit_request( + &self, + ranges: Vec>, + priority: u64, + ) -> BoxFuture<'static, Result>> { + let data = ranges + .into_iter() + .map(|range| self.data.slice(range.start as usize..range.end as usize)) + .collect(); + + log::trace!("Scheduled request with priority {}", priority); + std::future::ready(data) + .map(move |data| { + log::trace!("Decoded request with priority {}", priority); + Ok(data) + }) + .boxed() + } +} + +fn column_indices_from_schema_helper( + fields: &[FieldRef], + column_indices: &mut Vec, + column_counter: &mut u32, + is_structural_encoding: bool, +) { + // In the old style, every field except FSL gets its own column. In the new style only primitive + // leaf fields get their own column. + for field in fields { + if is_structural_encoding && field.metadata().contains_key("lance-encoding:packed") { + column_indices.push(*column_counter); + *column_counter += 1; + continue; + } + + match field.data_type() { + DataType::Struct(fields) => { + if !is_structural_encoding { + column_indices.push(*column_counter); + *column_counter += 1; + } + column_indices_from_schema_helper( + fields.as_ref(), + column_indices, + column_counter, + is_structural_encoding, + ); + } + DataType::List(inner) => { + if !is_structural_encoding { + column_indices.push(*column_counter); + *column_counter += 1; + } + column_indices_from_schema_helper( + std::slice::from_ref(inner), + column_indices, + column_counter, + is_structural_encoding, + ); + } + DataType::LargeList(inner) => { + if !is_structural_encoding { + column_indices.push(*column_counter); + *column_counter += 1; + } + column_indices_from_schema_helper( + std::slice::from_ref(inner), + column_indices, + column_counter, + is_structural_encoding, + ); + } + DataType::Map(entries, _) => { + column_indices_from_schema_helper( + std::slice::from_ref(entries), + column_indices, + column_counter, + is_structural_encoding, + ); + } + DataType::FixedSizeList(inner, _) => { + // FSL(primitive) does not get its own column in either approach + column_indices_from_schema_helper( + std::slice::from_ref(inner), + column_indices, + column_counter, + is_structural_encoding, + ); + } + _ => { + column_indices.push(*column_counter); + *column_counter += 1; + + column_indices_from_schema_helper( + &[], + column_indices, + column_counter, + is_structural_encoding, + ); + } + } + } +} + +fn column_indices_from_schema(schema: &Schema, is_structural_encoding: bool) -> Vec { + let mut column_indices = Vec::new(); + let mut column_counter = 0; + column_indices_from_schema_helper( + schema.fields(), + &mut column_indices, + &mut column_counter, + is_structural_encoding, + ); + column_indices +} + +#[allow(clippy::too_many_arguments)] +async fn test_decode( + num_rows: u64, + batch_size: u32, + schema: &Schema, + column_infos: &[Arc], + expected: Option>, + io: Arc, + is_structural_encoding: bool, + // The rows this read will schedule, used to exercise range-scoped page + // initialization. `None` initializes every page (the eager path). + requested_ranges: Option]>>, + schedule_fn: impl FnOnce( + DecodeBatchScheduler, + UnboundedSender>, + ) -> BoxFuture<'static, ()>, +) { + let lance_schema = lance_core::datatypes::Schema::try_from(schema).unwrap(); + let cache = Arc::new(lance_core::cache::LanceCache::with_capacity( + 128 * 1024 * 1024, + )); + let column_indices = column_indices_from_schema(schema, is_structural_encoding); + let decode_scheduler = DecodeBatchScheduler::try_new_with_ranges( + &lance_schema, + &column_indices, + column_infos, + &Vec::new(), + num_rows, + Arc::::default(), + io, + cache, + requested_ranges.as_deref(), + &FilterExpression::no_filter(), + &DecoderConfig::default(), + ) + .await + .unwrap(); + + let (tx, rx) = mpsc::unbounded_channel(); + + let scheduler_fut = schedule_fn(decode_scheduler, tx); + + scheduler_fut.await; + + let mut decode_stream = create_decode_stream( + &lance_schema, + num_rows, + batch_size, + is_structural_encoding, + /*should_validate=*/ true, + /*spawn_structural_batch_decode_tasks=*/ is_structural_encoding, + rx, + /*batch_size_bytes=*/ None, + ) + .unwrap(); + + let mut offset = 0; + while let Some(batch) = decode_stream.next().await { + let batch = batch.task.await.unwrap(); + if let Some(expected) = expected.as_ref() { + let actual = batch.column(0); + let expected_size = (batch_size as usize).min(expected.len() - offset); + let expected = expected.slice(offset, expected_size); + assert_eq!(expected.data_type(), actual.data_type()); + if expected.len() != actual.len() { + panic!( + "Mismatch in length (at offset={}) expected {} but got {}", + offset, + expected.len(), + actual.len() + ); + } + if &expected != actual { + if let Ok(comparator) = make_comparator(&expected, &actual, SortOptions::default()) + { + // We can't just assert_eq! because the error message is not very helpful. This gives us a bit + // more information about where the mismatch is. + for i in 0..expected.len() { + if !matches!(comparator(i, i), Ordering::Equal) { + panic!( + "Mismatch at index {} (offset={}) expected {:?} but got {:?} first mismatch is expected {:?} but got {:?}", + i, + offset, + expected, + actual, + expected.slice(i, 1), + actual.slice(i, 1) + ); + } + } + } else { + // Some arrays (like the null type) don't have a comparator so we just re-run the normal comparison + // and let it assert + assert_eq!(&expected, actual); + } + } + } + offset += batch.num_rows(); + } + if let Some(expected) = expected.as_ref() { + assert_eq!(offset, expected.len()); + } +} + +pub trait ArrayGeneratorProvider { + fn provide(&self) -> Box; + fn copy(&self) -> Box; +} +struct RandomArrayGeneratorProvider { + field: Field, +} + +impl ArrayGeneratorProvider for RandomArrayGeneratorProvider { + fn provide(&self) -> Box { + array::rand_type(self.field.data_type()) + } + + fn copy(&self) -> Box { + Box::new(Self { + field: self.field.clone(), + }) + } +} + +/// Given a field this will test the round trip encoding and decoding of random data +pub async fn check_basic_random(field: Field) { + check_specific_random(field, TestCases::basic()).await; +} + +/// Runs one independently schedulable slice of [`check_basic_random`]. +/// +/// The complete matrix is the Cartesian product of all encodings, page sizes, +/// and slicing modes. Keeping these axes outside the helper lets expensive +/// data types preserve the full matrix without concentrating it in one test. +pub async fn check_basic_random_case( + field: Field, + encoding: TestEncoding, + page_size: u64, + use_slicing: bool, +) { + check_specific_random( + field, + TestCases::basic() + .with_encoding(encoding) + .with_page_sizes(vec![page_size]) + .with_slicing_modes([use_slicing]), + ) + .await; +} + +pub async fn check_specific_random(field: Field, test_cases: TestCases) { + let array_generator_provider = RandomArrayGeneratorProvider { + field: field.clone(), + }; + check_round_trip_encoding_generated(field, Box::new(array_generator_provider), test_cases) + .await; +} + +pub struct FnArrayGeneratorProvider Box + Clone + 'static> { + provider_fn: F, +} + +impl Box + Clone + 'static> FnArrayGeneratorProvider { + pub fn new(provider_fn: F) -> Self { + Self { provider_fn } + } +} + +impl Box + Clone + 'static> ArrayGeneratorProvider + for FnArrayGeneratorProvider +{ + fn provide(&self) -> Box { + (self.provider_fn)() + } + + fn copy(&self) -> Box { + Box::new(Self { + provider_fn: self.provider_fn.clone(), + }) + } +} + +pub async fn check_basic_generated( + field: Field, + array_generator_provider: Box, +) { + check_round_trip_encoding_generated(field, array_generator_provider, TestCases::basic()).await; +} + +pub async fn check_round_trip_encoding_generated( + field: Field, + array_generator_provider: Box, + test_cases: TestCases, +) { + let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap(); + for page_size in test_cases.page_sizes.iter().copied() { + debug!("Testing random data with a page size of {}", page_size); + let encoder_factory = |encoding: TestEncoding| { + let encoding_strategy = test_encoding_strategy(encoding); + let mut column_index_seq = ColumnIndexSequence::default(); + let encoding_options = EncodingOptions { + max_page_bytes: MAX_PAGE_BYTES, + cache_bytes_per_column: page_size, + keep_original_array: true, + buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT, + }; + create_test_field_encoder( + encoding_strategy.as_ref(), + &lance_field, + &mut column_index_seq, + &encoding_options, + ) + .unwrap() + }; + + check_round_trip_random( + encoder_factory, + field.clone(), + array_generator_provider.copy(), + &test_cases, + ) + .await + } +} + +fn supports_nulls(data_type: &DataType, encoding: TestEncoding) -> bool { + if let DataType::Struct(fields) = data_type { + if encoding == TestEncoding::Array { + // 2.0 doesn't support nullability for structs + false + } else if fields.is_empty() { + // Even in 2.1 we don't support nulls for struct if there are no children because + // we have no spot to stick the repdef info (there is no column) + false + } else { + true + } + } else { + true + } +} + +type EncodingVerificationFn = dyn Fn(&[EncodedColumn], &TestEncoding); + +// The default will just test the full read +#[derive(Clone)] +pub struct TestCases { + ranges: Vec>, + indices: Vec>, + batch_size: u32, + skip_validation: bool, + max_page_size: Option, + page_sizes: Vec, + slicing_modes: Vec, + ingest_batch_counts: Vec, + encodings: Vec, + verify_encoding: Option>, + expected_encoding: Option>, +} + +impl Default for TestCases { + fn default() -> Self { + Self { + batch_size: 100, + ranges: Vec::new(), + indices: Vec::new(), + skip_validation: false, + max_page_size: None, + page_sizes: vec![4096, 1024 * 1024], + slicing_modes: vec![false, true], + ingest_batch_counts: vec![1, 5, 10], + encodings: TestEncoding::all().collect(), + verify_encoding: None, + expected_encoding: None, + } + } +} + +impl TestCases { + pub fn basic() -> Self { + Self::default() + .with_range(0..500) + .with_range(100..1100) + .with_range(8000..8500) + .with_indices(vec![100]) + .with_indices(vec![0]) + .with_indices(vec![9999]) + .with_indices(vec![100, 1100, 5000]) + .with_indices(vec![1000, 2000, 3000]) + .with_indices(vec![2000, 2001, 2002, 2003, 2004]) + // Big take that spans multiple pages and generates multiple output batches + .with_indices((100..500).map(|i| i * 3).collect::>()) + } + + pub fn with_range(mut self, range: Range) -> Self { + self.ranges.push(range); + self + } + + pub fn with_indices(mut self, indices: Vec) -> Self { + self.indices.push(indices); + self + } + + pub fn with_batch_size(mut self, batch_size: u32) -> Self { + self.batch_size = batch_size; + self + } + + pub fn without_validation(mut self) -> Self { + self.skip_validation = true; + self + } + + pub fn with_encoding(mut self, encoding: TestEncoding) -> Self { + self.encodings = vec![encoding]; + self + } + + pub fn with_encodings(mut self, encodings: impl IntoIterator) -> Self { + self.encodings = encodings.into_iter().collect(); + self + } + + pub fn with_structural_encodings(self) -> Self { + self.with_encodings([ + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse, + ]) + } + + pub fn with_u32_structural_encodings(self) -> Self { + self.with_encodings([TestEncoding::StructuralU32, TestEncoding::StructuralSparse]) + } + + pub fn with_dense_encodings(self) -> Self { + self.with_encodings([ + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + ]) + } + + pub fn with_array_and_u16_encodings(self) -> Self { + self.with_encodings([TestEncoding::Array, TestEncoding::StructuralU16]) + } + + pub fn with_page_sizes(mut self, page_sizes: Vec) -> Self { + self.page_sizes = page_sizes; + self + } + + pub fn with_slicing_modes(mut self, slicing_modes: impl IntoIterator) -> Self { + self.slicing_modes = slicing_modes.into_iter().collect(); + self + } + + pub fn with_ingest_batch_counts( + mut self, + ingest_batch_counts: impl IntoIterator, + ) -> Self { + self.ingest_batch_counts = ingest_batch_counts.into_iter().collect(); + self + } + + pub fn with_max_page_size(mut self, max_page_size: u64) -> Self { + self.max_page_size = Some(max_page_size); + self + } + + fn get_max_page_size(&self) -> u64 { + self.max_page_size.unwrap_or(MAX_PAGE_BYTES) + } + + fn encodings(&self) -> impl Iterator + '_ { + self.encodings.iter().copied() + } + + pub fn with_verify_encoding(mut self, verify_encoding: Arc) -> Self { + self.verify_encoding = Some(verify_encoding); + self + } + + fn verify_encoding(&self, columns: &[EncodedColumn], encoding: &TestEncoding) { + if let Some(verify_encoding) = self.verify_encoding.as_ref() { + verify_encoding(columns, encoding); + } + } + + pub fn with_expected_encoding(mut self, encoding: impl Into) -> Self { + self.expected_encoding = Some(vec![encoding.into()]); + self + } + + pub fn with_expected_encoding_chain, T: Into>( + mut self, + encodings: I, + ) -> Self { + self.expected_encoding = Some(encodings.into_iter().map(Into::into).collect()); + self + } +} + +/// Maps encoding enum variant to its string tag +/// Uses exhaustive match to ensure compile-time checking when variants are added/removed +fn tag(e: &Compression) -> &'static str { + use Compression::*; + + match e { + Flat(_) => "flat", + Variable(_) => "variable", + Constant(_) => "constant", + OutOfLineBitpacking(_) => "out_of_line_bitpacking", + InlineBitpacking(_) => "inline_bitpacking", + Fsst(_) => "fsst", + Dictionary(_) => "dictionary", + Rle(_) => "rle", + ByteStreamSplit(_) => "byte_stream_split", + General(_) => "general", + FixedSizeList(_) => "fixed_size_list", + PackedStruct(_) => "packed_struct", + VariablePackedStruct(_) => "variable_packed_struct", + } +} + +/// Returns any buffer outputs of this encoding +fn buffer(c: &Compression) -> Option> { + use Compression::*; + + match c { + Flat(f) => f.data.as_ref().map(|b| vec![b]), + Variable(v) => v.values.as_ref().map(|b| vec![b]), + InlineBitpacking(i) => i.values.as_ref().map(|b| vec![b]), + General(g) => g.compression.as_ref().map(|c| vec![c]), + _ => None, + } +} + +/// Returns the child encoding if this variant contains nested ArrayEncoding +fn child(c: &Compression) -> Option> { + use Compression::*; + + match c { + Variable(v) => v.offsets.as_ref().map(|b| vec![b.as_ref()]), + OutOfLineBitpacking(o) => o.values.as_ref().map(|b| vec![b.as_ref()]), + Fsst(f) => f.values.as_ref().map(|b| vec![b.as_ref()]), + ByteStreamSplit(b) => b.values.as_ref().map(|b| vec![b.as_ref()]), + General(g) => g.values.as_ref().map(|b| vec![b.as_ref()]), + Dictionary(d) => { + let mut children = Vec::new(); + if let Some(values) = d.items.as_ref() { + children.push(values.as_ref()); + } + if let Some(indices) = d.indices.as_ref() { + children.push(indices.as_ref()); + } + Some(children) + } + Rle(r) => { + let mut children = Vec::new(); + if let Some(values) = r.values.as_ref() { + children.push(values.as_ref()); + } + if let Some(run_lengths) = r.run_lengths.as_ref() { + children.push(run_lengths.as_ref()); + } + Some(children) + } + FixedSizeList(f) => f.values.as_ref().map(|b| vec![b.as_ref()]), + VariablePackedStruct(v) => { + let nested: Vec<&CompressiveEncoding> = v + .fields + .iter() + .filter_map(|field| field.value.as_ref()) + .collect(); + if nested.is_empty() { + None + } else { + Some(nested) + } + } + _ => None, + } +} + +/// Extract encoding types from array encoding (helper for nested encodings) +/// Returns the encoding chain including compression schemes for Block variants +pub fn extract_array_encoding_chain(enc: &CompressiveEncoding) -> Vec { + let mut chain = Vec::with_capacity(8); + let mut stack = vec![enc]; + + while let Some(cur) = stack.pop() { + if let Some(inner) = &cur.compression { + // 1. Add current layer's tag + chain.push(tag(inner).to_string()); + + // 2. Extract any buffer output + if let Some(buffer) = buffer(inner) { + chain.extend(buffer.into_iter().map(|b| { + let scheme = CompressionScheme::try_from(b.scheme()).unwrap(); + scheme.to_string() + })); + } + + // 3. Process child encoding if exists + if let Some(children) = child(inner) { + stack.extend(children); + } + } + } + chain +} + +fn collect_page_encoding(layout: &PageLayout, actual_chain: &mut Vec) -> Result<()> { + // Extract encodings from the page layout + use crate::format::pb21::page_layout::Layout; + if let Some(ref layout_type) = layout.layout { + match layout_type { + Layout::MiniBlockLayout(mini_block) => { + if mini_block.dictionary.is_some() { + actual_chain.push("dictionary".to_string()); + } + // Check value compression + if let Some(ref value_comp) = mini_block.value_compression { + let chain = extract_array_encoding_chain(value_comp); + actual_chain.extend(chain); + } + } + Layout::FullZipLayout(full_zip) => { + // Check value compression in full zip layout + if let Some(ref value_comp) = full_zip.value_compression { + let chain = extract_array_encoding_chain(value_comp); + actual_chain.extend(chain); + } + } + Layout::ConstantLayout(_) => { + // Constant layout does not describe a value encoding chain. + } + Layout::BlobLayout(blob) => { + if let Some(inner_layout) = &blob.inner_layout { + collect_page_encoding(inner_layout.as_ref(), actual_chain)? + } + } + Layout::SparseLayout(sparse) => { + if let Some(value_compression) = &sparse.value_compression { + actual_chain.extend(extract_array_encoding_chain(value_compression)); + } + } + } + } + + Ok(()) +} + +/// Verify that a single page contains the expected encoding +fn verify_page_encoding( + page: &EncodedPage, + expected_chain: &[String], + col_idx: usize, +) -> Result<()> { + use crate::decoder::PageEncoding; + use lance_core::Error; + + let mut actual_chain = Vec::new(); + + match &page.description { + PageEncoding::Structural(layout) => { + collect_page_encoding(layout, &mut actual_chain)?; + + // All-null structural pages may legitimately contain no encodings to verify. + // This can happen even when compression is configured because there is no value data + // (and rep/def compression is not currently described in the page layout). + if actual_chain.is_empty() + && page.data.is_empty() + && let Some(crate::format::pb21::page_layout::Layout::ConstantLayout(cl)) = + layout.layout.as_ref() + && cl.inline_value.is_none() + { + return Ok(()); + } + } + PageEncoding::Legacy(_) => { + // We don't need to care about the v2.0 array encoding. + } + } + + // Check that all expected encodings appear in the actual chain + for expected in expected_chain { + if !actual_chain.iter().any(|actual| actual.contains(expected)) { + return Err(Error::invalid_input_source( + format!( + "Column {} expected encoding chain {:?} but got {:?}", + col_idx, expected_chain, actual_chain + ) + .into(), + )); + } + } + Ok(()) +} + +/// Given specific data and test cases we check round trip encoding and decoding +/// +/// Note that the input `data` is a `Vec` to simulate multiple calls to `maybe_encode`. +/// In other words, these are multiple chunks of one long array and not multiple columns +/// in a record batch. To feed a "record batch" you should first convert the record batch +/// to a struct array. +pub async fn check_round_trip_encoding_of_data( + data: Vec>, + test_cases: &TestCases, + metadata: HashMap, +) { + check_round_trip_encoding_of_data_with_expected(data, None, test_cases, metadata).await +} + +pub async fn check_round_trip_encoding_of_data_with_expected( + data: Vec>, + expected_override: Option>, + test_cases: &TestCases, + metadata: HashMap, +) { + let example_data = data.first().expect("Data must have at least one array"); + let mut field = Field::new("", example_data.data_type().clone(), true); + field = field.with_metadata(metadata); + let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap(); + for encoding in test_cases.encodings() { + for page_size in test_cases.page_sizes.iter() { + let encoding_strategy = test_encoding_strategy(encoding); + let mut column_index_seq = ColumnIndexSequence::default(); + let encoding_options = EncodingOptions { + cache_bytes_per_column: *page_size, + max_page_bytes: test_cases.get_max_page_size(), + keep_original_array: true, + buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT, + }; + let encoder = create_test_field_encoder( + encoding_strategy.as_ref(), + &lance_field, + &mut column_index_seq, + &encoding_options, + ) + .unwrap(); + info!( + "Testing round trip encoding of data with test encoding {} and page size {}", + encoding, page_size + ); + check_round_trip_encoding_inner( + encoder, + &field, + data.clone(), + expected_override.clone(), + test_cases, + encoding, + ) + .await + } + } +} + +struct SimulatedWriter { + page_infos: Vec>, + encoded_data: BytesMut, +} + +impl SimulatedWriter { + fn new(num_columns: u32) -> Self { + let mut page_infos = Vec::with_capacity(num_columns as usize); + page_infos.resize_with(num_columns as usize, Default::default); + Self { + page_infos, + encoded_data: BytesMut::new(), + } + } + + fn write_buffer(&mut self, buffer: LanceBuffer) -> (u64, u64) { + let offset = self.encoded_data.len() as u64; + self.encoded_data.extend_from_slice(&buffer); + let size = self.encoded_data.len() as u64 - offset; + let pad_bytes = pad_bytes::(self.encoded_data.len()); + self.encoded_data.extend(std::iter::repeat_n(0, pad_bytes)); + (offset, size) + } + + fn write_lance_buffer(&mut self, buffer: LanceBuffer) { + self.encoded_data.extend_from_slice(&buffer); + let pad_bytes = pad_bytes::(self.encoded_data.len()); + self.encoded_data.extend(std::iter::repeat_n(0, pad_bytes)); + } + + fn write_page(&mut self, encoded_page: EncodedPage) { + trace!("Encoded page {:?}", encoded_page); + let page_buffers = encoded_page.data; + let page_encoding = encoded_page.description; + let buffer_offsets_and_sizes = page_buffers + .into_iter() + .map(|b| { + let (offset, size) = self.write_buffer(b); + trace!("Encoded buffer offset={} size={}", offset, size); + (offset, size) + }) + .collect::>(); + + let page_info = PageInfo { + num_rows: encoded_page.num_rows, + encoding: page_encoding, + buffer_offsets_and_sizes: Arc::from(buffer_offsets_and_sizes), + priority: encoded_page.row_number, + }; + + let col_idx = encoded_page.column_idx as usize; + self.page_infos[col_idx].push(page_info); + } + + fn new_external_buffers(&self) -> OutOfLineBuffers { + OutOfLineBuffers::new(self.encoded_data.len() as u64, MIN_PAGE_BUFFER_ALIGNMENT) + } +} + +/// This is the inner-most check function that actually runs the round trip and tests it +async fn check_round_trip_encoding_inner( + mut encoder: Box, + field: &Field, + data: Vec>, + expected_override: Option>, + test_cases: &TestCases, + encoding: TestEncoding, +) { + let mut writer = SimulatedWriter::new(encoder.num_columns()); + + let log_page = |encoded_page: &EncodedPage| { + debug!( + "Encoded page on column {} with {} rows and start row {} and buffer sizes [{}]", + encoded_page.column_idx, + encoded_page.num_rows, + encoded_page.row_number, + encoded_page + .data + .iter() + .map(|buf| buf.len().to_string()) + .collect::>() + .join(", ") + ); + }; + + let mut row_number = 0; + for arr in &data { + let mut external_buffers = writer.new_external_buffers(); + let repdef = RepDefBuilder::default(); + let num_rows = arr.len() as u64; + let encode_tasks = encoder + .maybe_encode( + arr.clone(), + &mut external_buffers, + repdef, + row_number, + num_rows, + ) + .unwrap(); + for buffer in external_buffers.take_buffers() { + writer.write_lance_buffer(buffer); + } + for encode_task in encode_tasks { + let encoded_page = encode_task.await.unwrap(); + log_page(&encoded_page); + + // For V2.1, verify encoding in the page if expected + if encoding.is_structural() + && let Some(ref expected) = test_cases.expected_encoding + { + verify_page_encoding(&encoded_page, expected, encoded_page.column_idx as usize) + .unwrap(); + } + + writer.write_page(encoded_page); + } + row_number += arr.len() as u64; + } + + let mut external_buffers = writer.new_external_buffers(); + let encode_tasks = encoder.flush(&mut external_buffers).unwrap(); + for buffer in external_buffers.take_buffers() { + writer.write_lance_buffer(buffer); + } + for task in encode_tasks { + let encoded_page = task.await.unwrap(); + log_page(&encoded_page); + + // For V2.1, verify encoding in the page if expected + if encoding.is_structural() + && let Some(ref expected) = test_cases.expected_encoding + { + verify_page_encoding(&encoded_page, expected, encoded_page.column_idx as usize) + .unwrap(); + } + + writer.write_page(encoded_page); + } + + let mut external_buffers = writer.new_external_buffers(); + let encoded_columns = encoder.finish(&mut external_buffers).await.unwrap(); + test_cases.verify_encoding(&encoded_columns, &encoding); + for buffer in external_buffers.take_buffers() { + writer.write_lance_buffer(buffer); + } + let mut column_infos = Vec::new(); + for (col_idx, encoded_column) in encoded_columns.into_iter().enumerate() { + // Keep track of pages for encoding verification + for page in encoded_column.final_pages { + writer.write_page(page); + } + + let col_buffer_off_and_size = encoded_column + .column_buffers + .into_iter() + .map(|b| writer.write_buffer(b)) + .collect::>(); + + let column_info = ColumnInfo::new( + col_idx as u32, + Arc::from(std::mem::take(&mut writer.page_infos[col_idx])), + col_buffer_off_and_size, + encoded_column.encoding, + ); + + column_infos.push(Arc::new(column_info)); + } + + let encoded_data = writer.encoded_data.freeze(); + + let scheduler = Arc::new(SimulatedScheduler::new(encoded_data)) as Arc; + + let num_rows = data.iter().map(|arr| arr.len() as u64).sum::(); + let concat_data = if test_cases.skip_validation { + None + } else if let Some(DataType::Struct(_)) = data.first().map(|datum| datum.data_type()) { + // TODO(tsaucer) When arrow upgrades to 56, remove this if statement + // This is due to a check for concat_struct in arrow-rs. See https://github.com/lance-format/lance/pull/4598 + let capacities = Capacities::Array(num_rows as usize); + let array_data: Vec<_> = data.iter().map(|a| a.to_data()).collect::>(); + let array_data = array_data.iter().collect(); + let mut mutable = MutableArrayData::with_capacities(array_data, false, capacities); + + for (i, a) in data.iter().enumerate() { + mutable.extend(i, 0, a.len()) + } + + Some(make_array(mutable.freeze())) + } else { + Some(concat(&data.iter().map(|arr| arr.as_ref()).collect::>()).unwrap()) + }; + + let expected_data = expected_override.clone().or_else(|| concat_data.clone()); + + let is_structural_encoding = encoding.is_structural(); + + let decode_field = if is_structural_encoding { + let mut lance_field = lance_core::datatypes::Field::try_from(field).unwrap(); + if lance_field.is_blob() && matches!(lance_field.data_type(), DataType::Struct(_)) { + lance_field.unloaded_mut(); + let mut arrow_field = ArrowField::from(&lance_field); + let mut metadata = arrow_field.metadata().clone(); + metadata.insert("lance-encoding:packed".to_string(), "true".to_string()); + arrow_field = arrow_field.with_metadata(metadata); + arrow_field + } else { + field.clone() + } + } else { + field.clone() + }; + + let schema = Schema::new(vec![decode_field]); + + debug!("Testing full decode"); + let scheduler_copy = scheduler.clone(); + test_decode( + num_rows, + test_cases.batch_size, + &schema, + &column_infos, + expected_data.clone(), + scheduler_copy.clone(), + is_structural_encoding, + // Full scan exercises the eager (initialize-everything) path. + None, + |mut decode_scheduler, tx| { + async move { + decode_scheduler.schedule_range( + 0..num_rows, + &FilterExpression::no_filter(), + tx, + scheduler_copy, + ) + } + .boxed() + }, + ) + .await; + + // Test range scheduling + for range in &test_cases.ranges { + debug!("Testing decode of range {:?}", range); + let num_rows = range.end - range.start; + let expected = expected_data + .as_ref() + .map(|arr| arr.slice(range.start as usize, num_rows as usize)); + let scheduler = scheduler.clone(); + let range = range.clone(); + // Range reads exercise the range-scoped (lazy) path: only the pages this + // contiguous range overlaps should be initialized. + let requested_ranges = Some(Arc::<[Range]>::from(vec![range.clone()])); + test_decode( + num_rows, + test_cases.batch_size, + &schema, + &column_infos, + expected, + scheduler.clone(), + is_structural_encoding, + requested_ranges, + |mut decode_scheduler, tx| { + async move { + decode_scheduler.schedule_range( + range, + &FilterExpression::no_filter(), + tx, + scheduler, + ) + } + .boxed() + }, + ) + .await; + } + + // Test take scheduling + for indices in &test_cases.indices { + if indices.len() == 1 { + debug!("Testing decode of index {}", indices[0]); + } else { + debug!( + "Testing decode of {} indices spread across range [{}..{}]", + indices.len(), + indices[0], + indices[indices.len() - 1] + ); + } + let num_rows = indices.len() as u64; + let indices_arr = UInt64Array::from(indices.clone()); + + // There is a bug in arrow_select::take::take that causes it to return empty arrays + // if the data type is an empty struct. This is a workaround for that. + let is_empty_struct = if let DataType::Struct(fields) = field.data_type() { + fields.is_empty() + } else { + false + }; + + let expected = if is_empty_struct { + Some(Arc::new(StructArray::new_empty_fields(indices_arr.len(), None)) as Arc) + } else { + concat_data.as_ref().map(|concat_data| { + arrow_select::take::take(&concat_data, &indices_arr, None).unwrap() + }) + }; + + let scheduler = scheduler.clone(); + let indices = indices.clone(); + // Take reads exercise the lazy path with scattered rows. One range per + // index is a superset of the pages `schedule_take` later touches. + let requested_ranges = Some(Arc::<[Range]>::from( + indices.iter().map(|&i| i..i + 1).collect::>(), + )); + test_decode( + num_rows, + test_cases.batch_size, + &schema, + &column_infos, + expected, + scheduler.clone(), + is_structural_encoding, + requested_ranges, + |mut decode_scheduler, tx| { + async move { + decode_scheduler.schedule_take( + &indices, + &FilterExpression::no_filter(), + tx, + scheduler, + ) + } + .boxed() + }, + ) + .await; + } +} + +const NUM_RANDOM_ROWS: u32 = 10000; + +/// Generates random data (parameterized by null rate, slicing, and # ingest batches) +/// and tests with that against default test cases. +/// +/// To test specific test cases use the +async fn check_round_trip_random( + encoder_factory: impl Fn(TestEncoding) -> Box, + field: Field, + array_generator_provider: Box, + test_cases: &TestCases, +) { + for null_rate in [None, Some(0.5), Some(1.0)] { + for use_slicing in test_cases.slicing_modes.iter().copied() { + for encoding in test_cases.encodings() { + if null_rate != Some(1.0) && matches!(field.data_type(), DataType::Null) { + continue; + } + + let field = if null_rate.is_some() { + if !supports_nulls(field.data_type(), encoding) { + continue; + } + field.clone().with_nullable(true) + } else { + field.clone().with_nullable(false) + }; + + for num_ingest_batches in test_cases.ingest_batch_counts.iter().copied() { + let rows_per_batch = NUM_RANDOM_ROWS / num_ingest_batches; + let mut data = Vec::new(); + + // Test both ingesting one big array sliced into smaller arrays and smaller + // arrays independently generated. These behave slightly differently. For + // example, a list array sliced into smaller arrays will have arrays whose + // starting offset is not 0. + if use_slicing { + let mut generator = + gen_batch().anon_col(array_generator_provider.provide()); + if let Some(null_rate) = null_rate { + // The null generator is the only generator that already inserts nulls + // and attempting to do so again makes arrow-rs grumpy + if !matches!(field.data_type(), DataType::Null) { + generator.with_random_nulls(null_rate); + } + } + let all_data = generator + .into_batch_rows(RowCount::from(10000)) + .unwrap() + .column(0) + .clone(); + let mut offset = 0; + for _ in 0..num_ingest_batches { + data.push(all_data.slice(offset, rows_per_batch as usize)); + offset += rows_per_batch as usize; + } + } else { + for i in 0..num_ingest_batches { + let mut generator = gen_batch() + .with_seed(Seed::from(i as u64)) + .anon_col(array_generator_provider.provide()); + if let Some(null_rate) = null_rate { + // The null generator is the only generator that already inserts nulls + // and attempting to do so again makes arrow-rs grumpy + if !matches!(field.data_type(), DataType::Null) { + generator.with_random_nulls(null_rate); + } + } + let arr = generator + .into_batch_rows(RowCount::from(rows_per_batch as u64)) + .unwrap() + .column(0) + .clone(); + data.push(arr); + } + } + + info!( + "Testing encoding {} with {} rows divided across {} batches for {} rows per batch with null_rate={:?} and use_slicing={}", + encoding, + NUM_RANDOM_ROWS, + num_ingest_batches, + rows_per_batch, + null_rate, + use_slicing + ); + check_round_trip_encoding_inner( + encoder_factory(encoding), + &field, + data, + None, + test_cases, + encoding, + ) + .await + } + } + } + } +} diff --git a/vendor/lance-encoding/src/utils.rs b/vendor/lance-encoding/src/utils.rs new file mode 100644 index 000000000..4163277d2 --- /dev/null +++ b/vendor/lance-encoding/src/utils.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Miscellaneous utility functions that don't have a home elsewhere. + +pub mod accumulation; +pub mod bytepack; diff --git a/vendor/lance-encoding/src/utils/accumulation.rs b/vendor/lance-encoding/src/utils/accumulation.rs new file mode 100644 index 000000000..f6255d12e --- /dev/null +++ b/vendor/lance-encoding/src/utils/accumulation.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! An accumulation queue accumulates arrays until we have enough data to flush. + +use arrow_array::ArrayRef; +use lance_arrow::deepcopy::deep_copy_array; +use log::{debug, trace}; + +#[derive(Debug)] +pub struct AccumulationQueue { + cache_bytes: u64, + keep_original_array: bool, + buffered_arrays: Vec, + current_bytes: u64, + // Row number of the first item in buffered_arrays, reset on flush + row_number: u64, + // Number of top level rows represented in buffered_arrays, reset on flush + num_rows: u64, + // This is only for logging / debugging purposes + column_index: u32, +} + +impl AccumulationQueue { + pub fn new(cache_bytes: u64, column_index: u32, keep_original_array: bool) -> Self { + Self { + cache_bytes, + buffered_arrays: Vec::new(), + current_bytes: 0, + column_index, + keep_original_array, + row_number: u64::MAX, + num_rows: 0, + } + } + + /// Adds an array to the queue, if there is enough data then the queue is flushed + /// and returned + pub fn insert( + &mut self, + array: ArrayRef, + row_number: u64, + num_rows: u64, + ) -> Option<(Vec, u64, u64)> { + if self.row_number == u64::MAX { + self.row_number = row_number; + } + self.num_rows += num_rows; + self.current_bytes += array.get_array_memory_size() as u64; + if self.current_bytes > self.cache_bytes { + debug!( + "Flushing column {} page of size {} bytes (unencoded)", + self.column_index, self.current_bytes + ); + // Push into buffered_arrays without copy since we are about to flush anyways + self.buffered_arrays.push(array); + self.current_bytes = 0; + let row_number = self.row_number; + self.row_number = u64::MAX; + let num_rows = self.num_rows; + self.num_rows = 0; + Some(( + std::mem::take(&mut self.buffered_arrays), + row_number, + num_rows, + )) + } else { + trace!( + "Accumulating data for column {}. Now at {} bytes", + self.column_index, self.current_bytes + ); + if self.keep_original_array { + self.buffered_arrays.push(array); + } else { + self.buffered_arrays.push(deep_copy_array(array.as_ref())) + } + None + } + } + + pub fn flush(&mut self) -> Option<(Vec, u64, u64)> { + if self.buffered_arrays.is_empty() { + trace!( + "No final flush since no data at column {}", + self.column_index + ); + None + } else { + trace!( + "Final flush of column {} which has {} bytes", + self.column_index, self.current_bytes + ); + self.current_bytes = 0; + let row_number = self.row_number; + self.row_number = u64::MAX; + let num_rows = self.num_rows; + self.num_rows = 0; + Some(( + std::mem::take(&mut self.buffered_arrays), + row_number, + num_rows, + )) + } + } +} diff --git a/vendor/lance-encoding/src/utils/bytepack.rs b/vendor/lance-encoding/src/utils/bytepack.rs new file mode 100644 index 000000000..7c92b8f86 --- /dev/null +++ b/vendor/lance-encoding/src/utils/bytepack.rs @@ -0,0 +1,306 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Utilities for byte (not bit) packing for situations where saving a few +//! bits is less important than simplicity and speed. + +use lance_core::{Error, Result}; + +pub struct U8BytePacker { + data: Vec, +} + +impl U8BytePacker { + fn with_capacity(capacity: usize) -> Self { + Self { + data: Vec::with_capacity(capacity), + } + } + + fn append(&mut self, value: u8) { + self.data.push(value); + } +} + +pub struct U16BytePacker { + data: Vec, +} + +impl U16BytePacker { + fn with_capacity(capacity: usize) -> Self { + Self { + data: Vec::with_capacity(capacity * 2), + } + } + + fn append(&mut self, value: u16) { + self.data.extend_from_slice(&value.to_le_bytes()); + } +} + +pub struct U32BytePacker { + data: Vec, +} + +impl U32BytePacker { + fn with_capacity(capacity: usize) -> Self { + Self { + data: Vec::with_capacity(capacity * 4), + } + } + + fn append(&mut self, value: u32) { + self.data.extend_from_slice(&value.to_le_bytes()); + } +} + +pub struct U64BytePacker { + data: Vec, +} + +impl U64BytePacker { + fn with_capacity(capacity: usize) -> Self { + Self { + data: Vec::with_capacity(capacity * 8), + } + } + + fn append(&mut self, value: u64) { + self.data.extend_from_slice(&value.to_le_bytes()); + } +} + +/// A bytepacked integer encoder that automatically chooses the smallest +/// possible integer type to store the given values. +/// +/// This is byte packing (not bit packing). Not even that, we only fit things into +/// sizes of 1,2,4,8 bytes. It's simple, fast, and easy but doesn't provide the +/// maximum possible compression. +/// +/// Still, it's useful for things like offsets which are often small and fit into a +/// u16 or u32 but sometimes might need the full u64 range. +/// +/// In the future we can investigate replacing this with something more sophisticated. +pub enum BytepackedIntegerEncoder { + U8(U8BytePacker), + U16(U16BytePacker), + U32(U32BytePacker), + U64(U64BytePacker), + Zero, +} + +impl BytepackedIntegerEncoder { + /// Create a new encoder with the given capacity and maximum value. + pub fn with_capacity(capacity: usize, max_value: u64) -> Self { + if max_value == 0 { + Self::Zero + } else if max_value <= u8::MAX as u64 { + Self::U8(U8BytePacker::with_capacity(capacity)) + } else if max_value <= u16::MAX as u64 { + Self::U16(U16BytePacker::with_capacity(capacity)) + } else if max_value <= u32::MAX as u64 { + Self::U32(U32BytePacker::with_capacity(capacity)) + } else { + Self::U64(U64BytePacker::with_capacity(capacity)) + } + } + + /// Append a value to the encoder. + /// + /// # Errors + /// + /// Returns an error if `value` does not fit in the width selected at + /// construction time. + pub fn append(&mut self, value: u64) -> Result<()> { + match self { + Self::U8(_) if value > u8::MAX as u64 => { + return Err(Error::invalid_input(format!( + "value {value} does not fit in bytepacked u8" + ))); + } + Self::U16(_) if value > u16::MAX as u64 => { + return Err(Error::invalid_input(format!( + "value {value} does not fit in bytepacked u16" + ))); + } + Self::U32(_) if value > u32::MAX as u64 => { + return Err(Error::invalid_input(format!( + "value {value} does not fit in bytepacked u32" + ))); + } + _ => {} + } + self.append_trusted(value); + Ok(()) + } + + /// Append a value whose range is guaranteed by the caller's construction. + pub(crate) fn append_trusted(&mut self, value: u64) { + match self { + Self::U8(packer) => { + debug_assert!(u8::try_from(value).is_ok()); + packer.append(value as u8); + } + Self::U16(packer) => { + debug_assert!(u16::try_from(value).is_ok()); + packer.append(value as u16); + } + Self::U32(packer) => { + debug_assert!(u32::try_from(value).is_ok()); + packer.append(value as u32); + } + Self::U64(packer) => packer.append(value), + Self::Zero => {} + } + } + + /// Convert the encoder into a vector of bytes. + pub fn into_data(self) -> Vec { + match self { + Self::U8(packer) => packer.data, + Self::U16(packer) => packer.data, + Self::U32(packer) => packer.data, + Self::U64(packer) => packer.data, + Self::Zero => Vec::new(), + } + } +} + +/// An iterator that unpacks bytes into integers (currently only u64) +pub enum ByteUnpacker> { + U8(I), + U16(I), + U32(I), + U64(I), +} + +impl> ByteUnpacker { + #[allow(clippy::new_ret_no_self)] + pub fn new>(data: I, size: usize) -> impl Iterator { + match size { + 1 => Self::U8(data.into_iter()), + 2 => Self::U16(data.into_iter()), + 4 => Self::U32(data.into_iter()), + 8 => Self::U64(data.into_iter()), + _ => panic!("Invalid size"), + } + } +} + +impl> Iterator for ByteUnpacker { + type Item = u64; + + fn next(&mut self) -> Option { + match self { + Self::U8(iter) => iter.next().map(|v| v as u64), + Self::U16(iter) => { + let first_byte = iter.next()?; + Some(u16::from_le_bytes([first_byte, iter.next().unwrap()]) as u64) + } + Self::U32(iter) => { + let first_byte = iter.next()?; + Some(u32::from_le_bytes([ + first_byte, + iter.next().unwrap(), + iter.next().unwrap(), + iter.next().unwrap(), + ]) as u64) + } + Self::U64(iter) => { + let first_byte = iter.next()?; + Some(u64::from_le_bytes([ + first_byte, + iter.next().unwrap(), + iter.next().unwrap(), + iter.next().unwrap(), + iter.next().unwrap(), + iter.next().unwrap(), + iter.next().unwrap(), + iter.next().unwrap(), + ])) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_bytepacked_integer_encoder() { + // Fits in u8 + let mut encoder = BytepackedIntegerEncoder::with_capacity(10, 100); + encoder.append(50).unwrap(); + encoder.append(20).unwrap(); + encoder.append(30).unwrap(); + let data = encoder.into_data(); + assert_eq!(data, vec![50, 20, 30]); + + assert_eq!( + ByteUnpacker::new(data, 1).collect::>(), + vec![50, 20, 30] + ); + + // Requires u16 + let mut encoder = BytepackedIntegerEncoder::with_capacity(10, 1000); + encoder.append(500).unwrap(); + encoder.append(200).unwrap(); + encoder.append(300).unwrap(); + let data = encoder.into_data(); + assert_eq!(data, vec![244, 1, 200, 0, 44, 1]); + + assert_eq!( + ByteUnpacker::new(data, 2).collect::>(), + vec![500, 200, 300] + ); + + // Requires u32 + let mut encoder = BytepackedIntegerEncoder::with_capacity(10, 1000000); + encoder.append(500000).unwrap(); + encoder.append(200000).unwrap(); + encoder.append(300000).unwrap(); + let data = encoder.into_data(); + assert_eq!(data, vec![32, 161, 7, 0, 64, 13, 3, 0, 224, 147, 4, 0]); + + assert_eq!( + ByteUnpacker::new(data, 4).collect::>(), + vec![500000, 200000, 300000] + ); + + // Requires u64 + let mut encoder = BytepackedIntegerEncoder::with_capacity(10, 0x10000000000); + encoder.append(0x5000000000).unwrap(); + encoder.append(0x2000000000).unwrap(); + encoder.append(0x3000000000).unwrap(); + let data = encoder.into_data(); + assert_eq!( + data, + vec![ + 0, 0, 0, 0, 80, 0, 0, 0, 0, 0, 0, 0, 32, 0, 0, 0, 0, 0, 0, 0, 48, 0, 0, 0 + ] + ); + + assert_eq!( + ByteUnpacker::new(data, 8).collect::>(), + vec![0x5000000000, 0x2000000000, 0x3000000000] + ); + } + + #[test] + fn test_bytepacked_integer_encoder_rejects_overflow() { + for (max_value, invalid_value, expected_width) in [ + (u8::MAX as u64, u8::MAX as u64 + 1, "u8"), + (u16::MAX as u64, u16::MAX as u64 + 1, "u16"), + (u32::MAX as u64, u32::MAX as u64 + 1, "u32"), + ] { + let mut encoder = BytepackedIntegerEncoder::with_capacity(1, max_value); + let error = encoder.append(invalid_value).unwrap_err(); + assert!(error.to_string().contains(expected_width), "{error}"); + } + + let mut disabled = BytepackedIntegerEncoder::with_capacity(1, 0); + disabled.append(u64::MAX).unwrap(); + assert!(disabled.into_data().is_empty()); + } +}