rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-12 23:02:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(python): avoid full-width bitpacking panic

Browse Source
This commit is contained in:
Gatefixer
2026-09-24 16:51:56 +00:00
parent 604b1b583b
commit 26565662a0
76 changed files with 70977 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1 -1
View File
@@ -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",
Cargo.toml
+4
View File
@@ -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"
python/Cargo.toml
+3
View File
@@ -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
python/python/tests/test_table.py
+16
View File
@@ -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(
[
vendor/lance-encoding/Cargo.toml
Vendored
+55
View File
@@ -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"]
vendor/lance-encoding/LICENSE
Vendored
+255
View File
@@ -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
vendor/lance-encoding/PATCH.md
Vendored
+11
View File
@@ -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.
vendor/lance-encoding/README.md
Vendored
+6
View File
@@ -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**.
vendor/lance-encoding/build.rs
Vendored
+23
View File
@@ -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(())
}
vendor/lance-encoding/protos/encodings_v2_0.proto
Vendored
+358
View File
@@ -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;
}
}
vendor/lance-encoding/protos/encodings_v2_1.proto
Vendored
+675
View File
@@ -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<LIST<INT32>>.
*
* 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;
}
}
vendor/lance-encoding/src/array_encoding.rs
Vendored
+13
View File
@@ -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;
vendor/lance-encoding/src/array_encoding/logical.rs
Vendored
+8
View File
@@ -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;
vendor/lance-encoding/src/array_encoding/logical/binary.rs
Vendored
+262
View File
@@ -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<dyn SchedulingJob + 'a>,
}
impl SchedulingJob for BinarySchedulingJob<'_> {
fn schedule_next(
&mut self,
context: &mut SchedulerContext,
priority: &dyn PriorityRange,
) -> Result<ScheduledScanLine> {
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::<Vec<_>>();
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<u8>`
#[derive(Debug)]
pub struct BinaryFieldScheduler {
varbin_scheduler: Arc<dyn FieldScheduler>,
data_type: DataType,
}
impl BinaryFieldScheduler {
// Create a new ListPageScheduler
pub fn new(varbin_scheduler: Arc<dyn FieldScheduler>, data_type: DataType) -> Self {
Self {
varbin_scheduler,
data_type,
}
}
}
impl FieldScheduler for BinaryFieldScheduler {
fn schedule_ranges<'a>(
&'a self,
ranges: &[std::ops::Range<u64>],
filter: &FilterExpression,
) -> Result<Box<dyn SchedulingJob + 'a>> {
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<dyn LogicalPageDecoder>,
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<NextDecodeTask> {
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<DrainLimit> {
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<dyn DecodeArrayTask>,
data_type: DataType,
}
impl BinaryArrayDecoder {
fn from_list_array<T: ByteArrayType>(array: &GenericListArray<T::Offset>) -> Result<ArrayRef> {
let values = array
.values()
.as_primitive::<UInt8Type>()
.values()
.inner()
.clone();
let offsets = array.offsets().clone();
let array = GenericByteArray::<T>::try_new(offsets, values, array.nulls().cloned())?;
Ok(Arc::new(array))
}
}
impl DecodeArrayTask for BinaryArrayDecoder {
fn decode(self: Box<Self>) -> 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::<BinaryType>(arr.as_list::<i32>())?,
DataType::LargeBinary => {
Self::from_list_array::<LargeBinaryType>(arr.as_list::<i64>())?
}
DataType::Utf8 => Self::from_list_array::<Utf8Type>(arr.as_list::<i32>())?,
DataType::LargeUtf8 => Self::from_list_array::<LargeUtf8Type>(arr.as_list::<i64>())?,
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<Self>) -> 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")
);
}
}
vendor/lance-encoding/src/array_encoding/logical/blob.rs
Vendored
+544
View File
@@ -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<dyn FieldScheduler>,
}
impl BlobFieldScheduler {
pub fn new(descriptions_scheduler: Arc<dyn FieldScheduler>) -> Self {
Self {
descriptions_scheduler,
}
}
}
#[derive(Debug)]
struct BlobFieldSchedulingJob<'a> {
descriptions_job: Box<dyn SchedulingJob + 'a>,
}
impl SchedulingJob for BlobFieldSchedulingJob<'_> {
fn schedule_next(
&mut self,
context: &mut SchedulerContext,
priority: &dyn PriorityRange,
) -> Result<ScheduledScanLine> {
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::<UInt64Type>::from_iter_values(vec![]),
sizes: PrimitiveArray::<UInt64Type>::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<u64>],
filter: &FilterExpression,
) -> Result<Box<dyn SchedulingJob + 'a>> {
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<dyn EncodingsIo>,
unloaded_descriptions: Option<BoxFuture<'static, Result<ArrayRef>>>,
positions: PrimitiveArray<UInt64Type>,
sizes: PrimitiveArray<UInt64Type>,
num_rows: u64,
loaded: VecDeque<Bytes>,
validity: VecDeque<BooleanBuffer>,
rows_loaded: u64,
rows_drained: u64,
base_priority: u64,
}
impl BlobFieldDecoder {
fn drain_validity(&mut self, num_values: usize) -> Result<Option<NullBuffer>> {
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::<Vec<_>>();
let validity = positions
.iter()
.zip(sizes.iter())
.map(|(p, s)| *p != 1 || *s != 0)
.collect::<BooleanBuffer>();
// 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<NextDecodeTask> {
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::<Vec<_>>();
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<Bytes>,
validity: Option<NullBuffer>,
}
impl BlobArrayDecodeTask {
fn new(bytes: Vec<Bytes>, validity: Option<NullBuffer>) -> Self {
Self { bytes, validity }
}
}
impl DecodeArrayTask for BlobArrayDecodeTask {
fn decode(self: Box<Self>) -> Result<(ArrayRef, u64)> {
let num_bytes = self.bytes.iter().map(|b| b.len()).sum::<usize>();
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::<Vec<_>>();
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<dyn FieldEncoder>,
}
impl BlobFieldEncoder {
pub fn new(description_encoder: Box<dyn FieldEncoder>) -> Self {
Self {
description_encoder,
}
}
fn write_bins(array: ArrayRef, external_buffers: &mut OutOfLineBuffers) -> Result<ArrayRef> {
let binarray = array.as_binary_opt::<i64>().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<Vec<EncodeTask>> {
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<Vec<EncodeTask>> {
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<Vec<crate::encoder::EncodedColumn>>> {
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<HashMap<String, String>> = LazyLock::new(|| {
[(BLOB_META_KEY.to_string(), "true".to_string())]
.iter()
.cloned()
.collect::<HashMap<_, _>>()
});
#[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<Range<u64>>,
_priority: u64,
) -> BoxFuture<'static, Result<Vec<Bytes>>> {
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::<UInt64Type>::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());
}
}
vendor/lance-encoding/src/array_encoding/logical/list.rs
Vendored
+1505
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/array_encoding/logical/primitive.rs
Vendored
+589
View File
@@ -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<dyn PageScheduler>,
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<PrimitivePage>,
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::<Vec<_>>();
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<Range<u64>>,
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<Range<u64>>) -> 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<ScheduledScanLine> {
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<u64>],
// TODO: Could potentially use filter to simplify decode, something of a micro-optimization probably
_filter: &FilterExpression,
) -> Result<Box<dyn SchedulingJob + 'a>> {
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<BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>>>,
physical_decoder: Option<Arc<dyn PrimitivePageDecoder>>,
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<dyn PrimitivePageDecoder>,
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<dyn PrimitivePageDecoder>,
data_type: DataType,
}
impl DecodeArrayTask for PrimitiveFieldDecodeTask {
fn decode(self: Box<Self>) -> 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<NextDecodeTask> {
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<DrainLimit> {
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<dyn ArrayEncodingStrategy>,
column_index: u32,
field: Field,
max_page_bytes: u64,
}
impl PrimitiveFieldEncoder {
pub fn try_new(
options: &EncodingOptions,
array_encoding_strategy: Arc<dyn ArrayEncodingStrategy>,
column_index: u32,
field: Field,
) -> Result<Self> {
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<ArrayRef>) -> Result<EncodeTask> {
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<ArrayRef>) -> Result<Vec<EncodeTask>> {
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<Vec<EncodeTask>> {
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<Vec<EncodeTask>> {
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<Vec<crate::encoder::EncodedColumn>>> {
std::future::ready(Ok(vec![EncodedColumn::default()])).boxed()
}
}
vendor/lance-encoding/src/array_encoding/logical/struct.rs
Vendored
+661
View File
@@ -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<dyn SchedulingJob + 'a>,
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<std::cmp::Ordering> {
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<Self>) -> 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::<Field>::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<NextDecodeTask> {
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<ScheduledScanLine> {
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<SchedulingJobWithStatus<'a>>,
rows_scheduled: u64,
num_rows: u64,
initialized: bool,
}
impl<'a> SimpleStructSchedulerJob<'a> {
fn new(
scheduler: &'a SimpleStructScheduler,
children: Vec<Box<dyn SchedulingJob + 'a>>,
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::<BinaryHeap<_>>();
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<ScheduledScanLine> {
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<Arc<dyn FieldScheduler>>,
child_fields: Fields,
num_rows: u64,
}
impl SimpleStructScheduler {
pub fn new(
children: Vec<Arc<dyn FieldScheduler>>,
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<u64>],
filter: &FilterExpression,
) -> Result<Box<dyn SchedulingJob + 'a>> {
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::<Result<Vec<_>>>()?;
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::<FuturesUnordered<_>>();
async move {
futures
.map(|res| res.map(|_| ()))
.try_collect::<Vec<_>>()
.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<Box<dyn LogicalPageDecoder>>,
// 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<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<CompositeDecodeTask> {
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<DrainLimit> {
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<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
#[derive(Debug)]
pub struct SimpleStructDecoder {
children: Vec<ChildState>,
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::<BinaryHeap<_>>();
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<NextDecodeTask> {
let child_tasks = self
.children
.iter_mut()
.map(|child| child.drain(num_rows))
.collect::<Result<Vec<_>>>()?;
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<DrainLimit> {
// 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<Box<dyn DecodeArrayTask>>,
num_rows: u64,
has_more: bool,
}
impl CompositeDecodeTask {
fn decode(self) -> Result<ArrayRef> {
let arrays = self
.tasks
.into_iter()
.map(|task| task.decode().map(|(arr, _)| arr))
.collect::<Result<Vec<_>>>()?;
let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
// 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<CompositeDecodeTask>,
child_fields: Fields,
}
impl DecodeArrayTask for SimpleStructDecodeTask {
fn decode(self: Box<Self>) -> Result<(ArrayRef, u64)> {
let child_arrays = self
.children
.into_iter()
.map(|child| child.decode())
.collect::<Result<Vec<_>>>()?;
// 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,
))
}
}
vendor/lance-encoding/src/array_encoding/physical.rs
Vendored
+329
View File
@@ -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<dyn PageScheduler> {
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<dyn PageScheduler> {
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<dyn PageScheduler> {
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<dyn PageScheduler> {
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::<Vec<_>>();
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<dyn PageScheduler> {
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) } }"
);
}
}
vendor/lance-encoding/src/array_encoding/physical/basic.rs
Vendored
+246
View File
@@ -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<dyn PrimitivePageDecoder>,
values: Box<dyn PrimitivePageDecoder>,
}
enum DataNullStatus {
// Neither validity nor values
All,
// Values only
None(Box<dyn PrimitivePageDecoder>),
// Validity and values
Some(DataDecoders),
}
#[derive(Debug)]
struct DataSchedulers {
validity: Box<dyn PageScheduler>,
values: Box<dyn PageScheduler>,
}
#[derive(Debug)]
enum SchedulerNullStatus {
// Values only
None(Box<dyn PageScheduler>),
// 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<dyn PageScheduler>,
values_decoder: Box<dyn PageScheduler>,
) -> 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<dyn PageScheduler>) -> 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<u64>],
scheduler: &Arc<dyn EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
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<dyn PrimitivePageDecoder>)
}
.boxed()
}
}
struct BasicPageDecoder {
mode: DataNullStatus,
}
impl PrimitivePageDecoder for BasicPageDecoder {
fn variable_width_bytes(&self, rows_to_skip: u64, num_rows: u64) -> Result<Option<u64>> {
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<DataBlock> {
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<dyn ArrayEncoder>,
}
impl BasicEncoder {
pub fn new(values_encoder: Box<dyn ArrayEncoder>) -> Self {
Self { values_encoder }
}
}
impl ArrayEncoder for BasicEncoder {
fn encode(
&self,
data: DataBlock,
data_type: &DataType,
buffer_index: &mut u32,
) -> Result<EncodedArray> {
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,
})
}
}
}
}
vendor/lance-encoding/src/array_encoding/physical/binary.rs
Vendored
+873
View File
@@ -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<u64>,
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<UInt64Type>, 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<u64>, BooleanBuffer) {
(self.indices, self.validity.finish())
}
}
#[derive(Debug)]
pub struct BinaryPageScheduler {
indices_scheduler: Arc<dyn PageScheduler>,
bytes_scheduler: Arc<dyn PageScheduler>,
offsets_type: DataType,
null_adjustment: u64,
}
impl BinaryPageScheduler {
pub fn new(
indices_scheduler: Arc<dyn PageScheduler>,
bytes_scheduler: Arc<dyn PageScheduler>,
offsets_type: DataType,
null_adjustment: u64,
) -> Self {
Self {
indices_scheduler,
bytes_scheduler,
offsets_type,
null_adjustment,
}
}
fn decode_indices(decoder: Arc<dyn PrimitivePageDecoder>, num_rows: u64) -> Result<ArrayRef> {
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<Box<dyn PrimitivePageDecoder>>>,
}
impl PageScheduler for BinaryPageScheduler {
fn schedule_ranges(
&self,
ranges: &[std::ops::Range<u64>],
scheduler: &Arc<dyn EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
// 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::<Vec<std::ops::Range<u64>>>();
// 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::<u64>();
let indices_num_rows = indices_ranges.iter().map(|r| r.end - r.start).sum::<u64>();
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::<UInt64Type>();
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, &copy_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<dyn PrimitivePageDecoder>)
}
})
.boxed()
}
}
struct BinaryPageDecoder {
decoded_indices: UInt64Array,
offsets_type: DataType,
validity: BooleanBuffer,
bytes_decoder: Box<dyn PrimitivePageDecoder>,
}
impl PrimitivePageDecoder for BinaryPageDecoder {
fn variable_width_bytes(&self, rows_to_skip: u64, num_rows: u64) -> Result<Option<u64>> {
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<DataBlock> {
// 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<dyn ArrayEncoder>,
compression_config: Option<CompressionConfig>,
buffer_compressor: Option<Box<dyn BufferCompressor>>,
}
impl BinaryEncoder {
pub fn try_new(
indices_encoder: Box<dyn ArrayEncoder>,
compression_config: Option<CompressionConfig>,
) -> Result<Self> {
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<LanceBuffer>,
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::<i32>();
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::<i64>();
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<EncodedArray> {
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<DataBlock> {
Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
bits_per_value: 8,
data: LanceBuffer::empty(),
num_values: 0,
block_info: BlockInfo::new(),
}))
}
}
#[derive(Debug)]
struct BytesDecoder {
bytes: Vec<u8>,
}
impl PrimitivePageDecoder for BytesDecoder {
fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result<DataBlock> {
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::<i64>().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<DataBlock> {
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::<u64>()
}
#[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<DataBlock> {
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));
}
}
vendor/lance-encoding/src/array_encoding/physical/bitmap.rs
Vendored
+223
View File
@@ -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<u64>],
scheduler: &Arc<dyn EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
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::<Vec<_>>();
let byte_ranges = chunk_reqs
.iter()
.map(|(range, _, _)| range.clone())
.collect::<Vec<_>>();
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::<Vec<_>>();
Ok(Box::new(BitmapDecoder { chunks }) as Box<dyn PrimitivePageDecoder>)
}
.boxed()
}
}
struct BitmapData {
data: Bytes,
bit_offset: u64,
length: u64,
}
struct BitmapDecoder {
chunks: Vec<BitmapData>,
}
impl PrimitivePageDecoder for BitmapDecoder {
fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result<DataBlock> {
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]);
}
}
vendor/lance-encoding/src/array_encoding/physical/bitpack.rs
Vendored
+728
View File
@@ -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<u64>],
scheduler: &Arc<dyn crate::EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
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<dyn PrimitivePageDecoder>)
}
.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<DataBlock> {
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<u8> = 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<std::ops::Range<u64>>],
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<u64>],
scheduler: &Arc<dyn crate::EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
let mut min = u64::MAX;
let mut max = 0;
let mut buffer_bit_start_offsets: Vec<u8> = vec![];
let mut buffer_bit_end_offsets: Vec<Option<u8>> = 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::<Vec<_>>();
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<dyn PrimitivePageDecoder>)
}
.boxed()
}
}
#[derive(Debug)]
struct BitpackedPageDecoder {
// bit offsets of the first value within each buffer
buffer_bit_start_offsets: Vec<u8>,
// 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<Option<u8>>,
// 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<Bytes>,
}
impl PrimitivePageDecoder for BitpackedPageDecoder {
fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result<DataBlock> {
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<u8>,
) -> 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<u8>,
) -> 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<T: ArrowNativeType>(block: DataBlock) -> Vec<T> {
let DataBlock::FixedWidth(FixedWidthDataBlock { data, .. }) = block else {
panic!("expected fixed-width data");
};
data.borrow_to_typed_slice::<T>().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<dyn crate::EncodingsIo> = 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::<u16>(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<dyn crate::EncodingsIo> = 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::<i32>(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::<Vec<_>>();
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<dyn crate::EncodingsIo> = 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::<Vec<_>>();
assert_eq!(fixed_width_values::<u32>(decoded), expected);
}
}
vendor/lance-encoding/src/array_encoding/physical/block.rs
Vendored
+75
View File
@@ -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<EncodedArray> {
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);
}
}
vendor/lance-encoding/src/array_encoding/physical/dictionary.rs
Vendored
+651
View File
@@ -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<dyn PageScheduler>,
items_scheduler: Arc<dyn PageScheduler>,
// 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<dyn PageScheduler>,
items_scheduler: Arc<dyn PageScheduler>,
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<u64>],
scheduler: &Arc<dyn EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
// 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<dyn PrimitivePageDecoder> =
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<dyn PrimitivePageDecoder> = indices_page_decoder.await?;
Ok(Box::new(DictionaryPageDecoder {
decoded_dict,
indices_decoder,
}) as Box<dyn PrimitivePageDecoder>)
})
.map(|join_handle| join_handle.unwrap())
.boxed()
} else {
let num_dictionary_items = self.num_dictionary_items;
tokio::spawn(async move {
let items_decoder: Arc<dyn PrimitivePageDecoder> =
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<dyn PrimitivePageDecoder>)
})
.map(|join_handle| join_handle.unwrap())
.boxed()
}
}
}
struct DirectDictionaryPageDecoder {
decoded_dict: DataBlock,
indices_decoder: Box<dyn PrimitivePageDecoder>,
}
impl PrimitivePageDecoder for DirectDictionaryPageDecoder {
fn variable_width_bytes(&self, _rows_to_skip: u64, _num_rows: u64) -> Result<Option<u64>> {
// 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<DataBlock> {
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<dyn Array>,
indices_decoder: Box<dyn PrimitivePageDecoder>,
}
impl PrimitivePageDecoder for DictionaryPageDecoder {
fn variable_width_bytes(&self, rows_to_skip: u64, num_rows: u64) -> Result<Option<u64>> {
// 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::<StringArray>() 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<DataBlock> {
// 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::<UInt8Type>();
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::<UInt8Type>::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::<StringArray>().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<dyn ArrayEncoder>,
items_encoder: Box<dyn ArrayEncoder>,
}
impl AlreadyDictionaryEncoder {
pub fn new(
indices_encoder: Box<dyn ArrayEncoder>,
items_encoder: Box<dyn ArrayEncoder>,
) -> Self {
Self {
indices_encoder,
items_encoder,
}
}
}
impl ArrayEncoder for AlreadyDictionaryEncoder {
fn encode(
&self,
data: DataBlock,
data_type: &DataType,
buffer_index: &mut u32,
) -> Result<EncodedArray> {
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<dyn ArrayEncoder>,
items_encoder: Box<dyn ArrayEncoder>,
}
impl DictionaryEncoder {
pub fn new(
indices_encoder: Box<dyn ArrayEncoder>,
items_encoder: Box<dyn ArrayEncoder>,
) -> 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<EncodedArray> {
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;
}
}
vendor/lance-encoding/src/array_encoding/physical/fixed_size_binary.rs
Vendored
+366
View File
@@ -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<dyn PageScheduler>,
byte_width: u32,
bytes_per_offset: u32,
}
impl FixedSizeBinaryPageScheduler {
pub fn new(
bytes_scheduler: Box<dyn PageScheduler>,
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<u64>],
scheduler: &Arc<dyn EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
let expanded_ranges = ranges
.iter()
.map(|range| {
(range.start * self.byte_width as u64)..(range.end * self.byte_width as u64)
})
.collect::<Vec<_>>();
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<dyn PrimitivePageDecoder>)
}
.boxed()
}
}
pub struct FixedSizeBinaryDecoder {
bytes_decoder: Box<dyn PrimitivePageDecoder>,
byte_width: u64,
bytes_per_offset: u32,
}
impl PrimitivePageDecoder for FixedSizeBinaryDecoder {
fn variable_width_bytes(&self, _rows_to_skip: u64, num_rows: u64) -> Result<Option<u64>> {
Ok(Some(num_rows * self.byte_width))
}
fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result<DataBlock> {
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::<Vec<_>>();
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::<Vec<_>>();
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<DataBlock> {
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<DataBlock> {
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 });
}
}
vendor/lance-encoding/src/array_encoding/physical/fixed_size_list.rs
Vendored
+275
View File
@@ -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<dyn PageScheduler>,
dimension: u32,
}
impl FixedListScheduler {
pub fn new(items_scheduler: Box<dyn PageScheduler>, dimension: u32) -> Self {
Self {
items_scheduler,
dimension,
}
}
}
impl PageScheduler for FixedListScheduler {
fn schedule_ranges(
&self,
ranges: &[std::ops::Range<u64>],
scheduler: &Arc<dyn EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
let expanded_ranges = ranges
.iter()
.map(|range| (range.start * self.dimension as u64)..(range.end * self.dimension as u64))
.collect::<Vec<_>>();
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<dyn PrimitivePageDecoder>)
}
.boxed()
}
}
pub struct FixedListDecoder {
items_decoder: Box<dyn PrimitivePageDecoder>,
dimension: u64,
}
impl PrimitivePageDecoder for FixedListDecoder {
fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result<DataBlock> {
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<dyn ArrayEncoder>,
dimension: u32,
}
impl FslEncoder {
pub fn new(items_encoder: Box<dyn ArrayEncoder>, dimension: u32) -> Self {
Self {
items_encoder,
dimension,
}
}
}
impl ArrayEncoder for FslEncoder {
fn encode(
&self,
data: DataBlock,
data_type: &DataType,
buffer_index: &mut u32,
) -> Result<EncodedArray> {
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::<Int32Type>().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;
}
}
vendor/lance-encoding/src/array_encoding/physical/fsst.rs
Vendored
+208
View File
@@ -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<dyn PageScheduler>,
symbol_table: LanceBuffer,
}
impl FsstPageScheduler {
pub fn new(inner_scheduler: Box<dyn PageScheduler>, symbol_table: LanceBuffer) -> Self {
Self {
inner_scheduler,
symbol_table,
}
}
}
impl PageScheduler for FsstPageScheduler {
fn schedule_ranges(
&self,
ranges: &[Range<u64>],
scheduler: &Arc<dyn EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
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<dyn PrimitivePageDecoder>)
}
.boxed()
}
}
struct FsstPageDecoder {
inner_decoder: Box<dyn PrimitivePageDecoder>,
symbol_table: LanceBuffer,
}
impl PrimitivePageDecoder for FsstPageDecoder {
fn variable_width_bytes(&self, rows_to_skip: u64, num_rows: u64) -> Result<Option<u64>> {
// 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<DataBlock> {
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::<i32>::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::<i32>::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<dyn ArrayEncoder>,
}
impl FsstArrayEncoder {
pub fn new(inner_encoder: Box<dyn ArrayEncoder>) -> Self {
Self { inner_encoder }
}
}
impl ArrayEncoder for FsstArrayEncoder {
fn encode(
&self,
data: DataBlock,
data_type: &DataType,
buffer_index: &mut u32,
) -> lance_core::Result<EncodedArray> {
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::<i32>();
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,
})
}
}
vendor/lance-encoding/src/array_encoding/physical/packed_struct.rs
Vendored
+350
View File
@@ -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<Box<dyn PageScheduler>>,
fields: Fields,
buffer_offset: u64,
}
impl PackedStructPageScheduler {
pub fn new(
_inner_schedulers: Vec<Box<dyn PageScheduler>>,
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<u64>],
scheduler: &Arc<dyn EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
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::<Vec<_>>();
// 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<dyn PrimitivePageDecoder>)
})
.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<DataBlock> {
// 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<BytesMut>
// 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<Box<dyn ArrayEncoder>>,
}
impl PackedStructEncoder {
pub fn new(inner_encoders: Vec<Box<dyn ArrayEncoder>>) -> Self {
Self { inner_encoders }
}
}
impl ArrayEncoder for PackedStructEncoder {
fn encode(
&self,
data: DataBlock,
data_type: &DataType,
buffer_index: &mut u32,
) -> Result<EncodedArray> {
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::<Result<Vec<_>>>()?;
let total_bits_per_value = fixed_fields.iter().map(|f| f.bits_per_value).sum::<u64>();
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::<Vec<_>>();
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;
}
}
vendor/lance-encoding/src/array_encoding/physical/value.rs
Vendored
+243
View File
@@ -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<u64>],
scheduler: &Arc<dyn EncodingsIo>,
top_level_row: u64,
) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
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::<Vec<_>>()
} 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::<Vec<_>>()
} 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<dyn PrimitivePageDecoder>)
}
.boxed()
}
}
struct ValuePageDecoder {
bytes_per_value: u64,
data: Vec<Bytes>,
uncompressed_data: Arc<Mutex<Option<Vec<Bytes>>>>,
uncompressed_range_offsets: Vec<std::ops::Range<usize>>,
compression_config: CompressionConfig,
}
impl ValuePageDecoder {
fn decompress(&self) -> Result<Vec<Bytes>> {
// for compressed page, it is guaranteed that only one range is passed
let bytes_u8: Vec<u8> = self.data[0].to_vec();
let buffer_compressor = GeneralBufferCompressor::get_compressor(self.compression_config)?;
let mut uncompressed_bytes: Vec<u8> = Vec::new();
buffer_compressor.decompress(&bytes_u8, &mut uncompressed_bytes)?;
let mut bytes_in_ranges: Vec<Bytes> =
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<Arc<Mutex<Option<Vec<Bytes>>>>> {
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<Item = &'a Bytes>,
mut bytes_to_skip: u64,
mut bytes_to_take: u64,
) -> LanceBuffer {
let mut dest: Option<Vec<u8>> = 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<DataBlock> {
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<EncodedArray> {
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 })
}
}
vendor/lance-encoding/src/array_encoding/strategy.rs
Vendored
+1101
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/buffer.rs
Vendored
+503
View File
@@ -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::<u8>::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<T>` as a LanceBuffer
///
/// This is a zero-copy operation. We can safely reinterpret `Vec<T>` into `&[u8]` which is what happens here.
/// However, we cannot safely reinterpret a `Vec<T>` into a `Vec<u8>` 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<T: ArrowNativeType>(vec: Vec<T>) -> 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<T>`
///
/// The same alignment constraints apply
pub fn reinterpret_slice<T: ArrowNativeType + RefUnwindSafe>(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<T>`
///
/// 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<T: ArrowNativeType>(&self) -> ScalarBuffer<T> {
let align = std::mem::align_of::<T>();
let is_aligned = self.as_ptr().align_offset(align) == 0;
if !self.len().is_multiple_of(std::mem::size_of::<T>()) {
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::<T>(),
self.len()
);
}
if is_aligned {
ScalarBuffer::<T>::from(self.clone().into_buffer())
} else {
let num_values = self.len() / std::mem::size_of::<T>();
let vec = Vec::<T>::with_capacity(num_values);
let mut bytes = MutableBuffer::from(vec);
bytes.extend_from_slice(self);
ScalarBuffer::<T>::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<T: ArrowNativeType + bytemuck::Pod>(&self) -> Cow<'_, [T]> {
let align = std::mem::align_of::<T>();
if !self.len().is_multiple_of(std::mem::size_of::<T>()) {
panic!(
"attempt to view data type of size {} but we have {} bytes which isn't evenly divisible",
std::mem::size_of::<T>(),
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>) -> 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<Self> {
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::<Result<Vec<_>>>()?;
let total_bytes_per_value = bytes_per_value.iter().sum::<u64>();
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::<Vec<_>>();
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<const N: usize>(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<Vec<u8>> for LanceBuffer {
fn from(buffer: Vec<u8>) -> Self {
Self(Buffer::from_vec(buffer))
}
}
impl From<Buffer> 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<Self::Item> {
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::<u32>().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::<u16>();
}
#[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::<u16>();
let view_ptr = borrow.as_ref().as_ptr();
let borrow2 = buf.borrow_to_typed_slice::<u16>();
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::<u16>();
let view_ptr = borrow.as_ref().as_ptr();
let borrow2 = buf.borrow_to_typed_slice::<u16>();
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]);
}
}
vendor/lance-encoding/src/compression.rs
Vendored
+2806
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/compression_config.rs
Vendored
+299
View File
@@ -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<Self> {
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<String, CompressionFieldParams>,
/// Type-level parameters: data type name -> parameters
pub types: HashMap<String, CompressionFieldParams>,
}
/// 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<f64>,
/// General compression scheme: "lz4", "zstd", "none"
pub compression: Option<String>,
/// Compression level (only for schemes that support it, e.g., zstd)
pub compression_level: Option<i32>,
/// Byte stream split mode for floating point data
pub bss: Option<BssMode>,
/// Minichunk size threshold for encoding
pub minichunk_size: Option<i64>,
}
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()));
}
}
vendor/lance-encoding/src/constants.rs
Vendored
+65
View File
@@ -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";
vendor/lance-encoding/src/data.rs
Vendored
+2898
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/decoder.rs
Vendored
+4203
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/encoder.rs
Vendored
+603
View File
@@ -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<LanceBuffer>, 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<EncodedArray>;
}
/// Selects an `ArrayEncoder` for one page.
pub trait ArrayEncodingStrategy: Send + Sync + std::fmt::Debug {
fn create_array_encoder(
&self,
arrays: &[ArrayRef],
field: &Field,
) -> Result<Box<dyn ArrayEncoder>>;
}
/// An encoded page of data
///
/// Maps to a top-level array
///
/// For example, `FixedSizeList<Int32>` will have two EncodedArray instances and one EncodedPage
#[derive(Debug)]
pub struct EncodedPage {
// The encoded page buffers
pub data: Vec<LanceBuffer>,
// 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<LanceBuffer>,
pub encoding: pb::ColumnEncoding,
pub final_pages: Vec<EncodedPage>,
}
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<LanceBuffer>,
}
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<LanceBuffer> {
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<EncodedPage>>;
/// 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<ArrayRef> {
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<Vec<EncodeTask>>;
/// 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<Vec<EncodeTask>>;
/// 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<Vec<EncodedColumn>>>;
/// 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<Box<dyn FieldEncoder>>;
}
/// 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<String, String>,
}
/// 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<Box<dyn FieldEncoder>>,
pub field_id_to_column_index: Vec<(u32, u32)>,
}
impl BatchEncoder {
pub fn try_new(
schema: &Schema,
strategy: &dyn FieldEncodingStrategy,
options: &EncodingOptions,
) -> Result<Self> {
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::<Result<Vec<_>>>()?;
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::<u32>()
}
}
/// 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<Arc<ColumnInfo>>,
pub schema: Arc<Schema>,
pub top_level_columns: Vec<u32>,
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<Schema>,
encoding_strategy: &dyn FieldEncodingStrategy,
options: &EncodingOptions,
) -> Result<EncodedBatch> {
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::<Result<Vec<_>>>()?;
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::<u32, Vec<PageInfo>>::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<Struct> should be rejected for file version 2.1"
);
let err = result.err().unwrap();
assert!(
err.to_string()
.contains("FixedSizeList<Struct> is not enabled by the selected file format")
);
}
}
vendor/lance-encoding/src/encoder/structural.rs
Vendored
+285
View File
@@ -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<Item = PrimitivePageEncoding>) -> 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<Box<dyn FieldEncoder>> {
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<Option<Box<dyn FieldEncoder>>> {
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<Option<Box<dyn FieldEncoder>>> {
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<Option<Box<dyn FieldEncoder>>> {
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<Option<Box<dyn FieldEncoder>>> {
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<Option<Box<dyn FieldEncoder>>> {
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<Option<Box<dyn FieldEncoder>>> {
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<key, value>".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<Option<Box<dyn FieldEncoder>>> {
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::<Result<Vec<_>>>()?;
Ok(Some(Box::new(StructStructuralEncoder::new(
context.options.keep_original_array,
children_encoders,
))))
}
vendor/lance-encoding/src/encodings.rs
Vendored
+8
View File
@@ -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;
vendor/lance-encoding/src/encodings/fuzz_tests.rs
Vendored
+559
View File
@@ -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<Value = EncodingTestConfig> {
(
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::<bool>(), // 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<Value = FixedWidthType> {
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<Value = VariableWidthType> {
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<Arc<dyn Array>> {
let mut batch_gen = gen_batch().with_seed(Seed::from(seed));
// Generate base values
let generator: Box<dyn ArrayGenerator> = 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<dyn ArrayGenerator> = 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::<u64>());
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<dyn Array>;
let single_string = Arc::new(StringArray::from(vec!["test"])) as Arc<dyn Array>;
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<Option<i32>>)) as Arc<dyn Array>;
let all_nulls_string = Arc::new(StringArray::from(
vec![None, None, None] as Vec<Option<&str>>
)) as Arc<dyn Array>;
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::<u64>()
) {
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<dyn Array>;
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::<u64>());
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::<Option<&[u8]>>::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;
}
}
vendor/lance-encoding/src/encodings/logical.rs
Vendored
+9
View File
@@ -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;
vendor/lance-encoding/src/encodings/logical/blob.rs
Vendored
+904
View File
@@ -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<dyn FieldEncoder>,
// Set when we first see data
def_meaning: Option<Arc<[DefinitionInterpretation]>>,
}
impl BlobStructuralEncoder {
pub fn new(
field: &Field,
make_descriptor_encoder: impl FnOnce(Field) -> Result<Box<dyn FieldEncoder>>,
) -> Result<Self> {
// Create descriptor field: struct<position: u64, size: u64>
// 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<EncodeTask>,
def_meaning: Arc<[DefinitionInterpretation]>,
) -> Vec<EncodeTask> {
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::<Vec<_>>()
}
}
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<Vec<EncodeTask>> {
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::<i64>().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<Vec<EncodeTask>> {
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<Vec<EncodedColumn>>> {
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<dyn FieldEncoder>,
}
impl BlobV2StructuralEncoder {
pub fn new(
field: &Field,
make_descriptor_encoder: impl FnOnce(Field) -> Result<Box<dyn FieldEncoder>>,
) -> Result<Self> {
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<Vec<EncodeTask>> {
let struct_arr = array
.as_any()
.downcast_ref::<StructArray>()
.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::<UInt8Type>();
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::<i64>();
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::<i32>();
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::<UInt32Type>();
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::<UInt64Type>();
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::<UInt64Type>();
let row_count = struct_arr.len();
let mut kind_builder = PrimitiveBuilder::<UInt8Type>::with_capacity(row_count);
let mut position_builder = PrimitiveBuilder::<UInt64Type>::with_capacity(row_count);
let mut size_builder = PrimitiveBuilder::<UInt64Type>::with_capacity(row_count);
let mut blob_id_builder = PrimitiveBuilder::<UInt32Type>::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<ArrayRef> = 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<Vec<EncodeTask>> {
self.descriptor_encoder.flush(external_buffers)
}
fn finish(
&mut self,
external_buffers: &mut OutOfLineBuffers,
) -> BoxFuture<'_, Result<Vec<EncodedColumn>>> {
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<Option<&[u8]>> =
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;
}
}
vendor/lance-encoding/src/encodings/logical/fixed_size_list.rs
Vendored
+804
View File
@@ -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<Int32>`) 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<dyn FieldEncoder>,
}
impl FixedSizeListStructuralEncoder {
pub fn new(keep_original_array: bool, child: Box<dyn FieldEncoder>) -> 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<Vec<EncodeTask>> {
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<Vec<EncodeTask>> {
self.child.flush(external_buffers)
}
fn num_columns(&self) -> u32 {
self.child.num_columns()
}
fn finish(
&mut self,
external_buffers: &mut OutOfLineBuffers,
) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
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<dyn StructuralFieldScheduler>,
dimension: u64,
}
impl StructuralFixedSizeListScheduler {
pub fn new(child: Box<dyn StructuralFieldScheduler>, dimension: i32) -> Self {
Self {
child,
dimension: dimension as u64,
}
}
fn child_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Range<u64>>> {
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<u64>],
filter: &FilterExpression,
) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
// 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<u64>]>,
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<dyn StructuralSchedulingJob + 'a>,
dimension: u64,
}
impl<'a> StructuralFixedSizeListSchedulingJob<'a> {
fn new(child: Box<dyn StructuralSchedulingJob + 'a>, dimension: u64) -> Self {
Self { child, dimension }
}
}
impl StructuralSchedulingJob for StructuralFixedSizeListSchedulingJob<'_> {
fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
// 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<dyn StructuralFieldDecoder>,
data_type: DataType,
}
impl StructuralFixedSizeListDecoder {
pub fn new(child: Box<dyn StructuralFieldDecoder>, 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<Box<dyn StructuralDecodeArrayTask>> {
// 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<dyn StructuralDecodeArrayTask>,
data_type: DataType,
num_rows: u64,
}
impl StructuralFixedSizeListDecodeTask {
fn new(
child_task: Box<dyn StructuralDecodeArrayTask>,
data_type: DataType,
num_rows: u64,
) -> Self {
Self {
child_task,
data_type,
num_rows,
}
}
}
impl StructuralDecodeArrayTask for StructuralFixedSizeListDecodeTask {
fn decode(self: Box<Self>) -> Result<DecodedArray> {
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::<i32>(), is_garbage),
DataType::LargeList(_) => filter_list_garbage(array.as_list::<i64>(), 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<ArrayRef> = 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<bool> {
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<bool> {
fsl_validity.iter().map(|valid| !valid).collect()
}
fn filter_list_garbage<O: OffsetSizeTrait>(
list_arr: &GenericListArray<O>,
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<O> = Vec::with_capacity(list_arr.len() + 1);
let mut values_to_keep: Vec<usize> = 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<i32> = Vec::with_capacity(map_arr.len() + 1);
let mut values_to_keep: Vec<usize> = 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<List<Int32>> 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::<i32>();
assert_eq!(
(0..6).map(|i| child_list.is_valid(i)).collect::<Vec<_>>(),
vec![true, true, false, false, true, true]
);
}
#[test]
fn test_filter_nested_fsl_no_list_child() {
// FSL<Int32> - 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);
}
}
vendor/lance-encoding/src/encodings/logical/list.rs
Vendored
+1662
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/encodings/logical/map.rs
Vendored
+801
View File
@@ -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<Struct<key, value>>
/// 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<dyn FieldEncoder>,
}
impl MapStructuralEncoder {
pub fn new(keep_original_array: bool, child: Box<dyn FieldEncoder>) -> 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<Vec<EncodeTask>> {
let map_array = array
.as_any()
.downcast_ref::<MapArray>()
.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<Vec<EncodeTask>> {
self.child.flush(external_buffers)
}
fn num_columns(&self) -> u32 {
self.child.num_columns()
}
fn finish(
&mut self,
external_buffers: &mut OutOfLineBuffers,
) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
self.child.finish(external_buffers)
}
}
#[derive(Debug)]
pub struct StructuralMapScheduler {
child: Box<dyn StructuralFieldScheduler>,
}
impl StructuralMapScheduler {
pub fn new(child: Box<dyn StructuralFieldScheduler>) -> Self {
Self { child }
}
}
impl StructuralFieldScheduler for StructuralMapScheduler {
fn schedule_ranges<'a>(
&'a self,
ranges: &[Range<u64>],
filter: &FilterExpression,
) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
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<u64>]>,
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<dyn StructuralSchedulingJob + 'a>,
}
impl<'a> StructuralMapSchedulingJob<'a> {
fn new(child: Box<dyn StructuralSchedulingJob + 'a>) -> Self {
Self { child }
}
}
impl StructuralSchedulingJob for StructuralMapSchedulingJob<'_> {
fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
self.child.schedule_next(context)
}
}
#[derive(Debug)]
pub struct StructuralMapDecoder {
child: Box<dyn StructuralFieldDecoder>,
data_type: DataType,
}
impl StructuralMapDecoder {
pub fn new(child: Box<dyn StructuralFieldDecoder>, 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<Box<dyn StructuralDecodeArrayTask>> {
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<dyn StructuralDecodeArrayTask>,
data_type: DataType,
}
impl StructuralMapDecodeTask {
fn new(child_task: Box<dyn StructuralDecodeArrayTask>, data_type: DataType) -> Self {
Self {
child_task,
data_type,
}
}
}
impl StructuralDecodeArrayTask for StructuralMapDecodeTask {
fn decode(self: Box<Self>) -> Result<DecodedArray> {
let DecodedArray {
array,
mut repdef,
data_size,
} = self.child_task.decode()?;
// Decode the offsets from RepDef
let (offsets, validity) = repdef.unravel_offsets::<i32>()?;
// 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::<arrow_array::StructArray>()
.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<Self>) -> lance_core::Result<DecodedArray> {
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<String, Int32>
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<String, Int32> 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<id: Int32, properties: Map<String, String>>
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<Map> 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<dyn Array> = 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<Map<String, Int32>>
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<String, Map<String, Int32>>
// 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::<i32>::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<Int32, String> (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
);
}
}
vendor/lance-encoding/src/encodings/logical/primitive.rs
Vendored
+12571
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/encodings/logical/primitive/blob.rs
Vendored
+552
View File
@@ -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<dyn StructuralPageScheduler>,
def_meaning: Arc<[DefinitionInterpretation]>,
}
impl BlobDescriptionPageScheduler {
pub fn new(
inner_scheduler: Box<dyn StructuralPageScheduler>,
def_meaning: Arc<[DefinitionInterpretation]>,
) -> Self {
Self {
inner_scheduler,
def_meaning,
}
}
fn wrap_decoder_fut(
decoder_fut: BoxFuture<'static, Result<Box<dyn StructuralPageDecoder>>>,
def_meaning: Arc<[DefinitionInterpretation]>,
) -> BoxFuture<'static, Result<Box<dyn StructuralPageDecoder>>> {
async move {
let decoder = decoder_fut.await?;
Ok(
Box::new(BlobDescriptionPageDecoder::new(decoder, def_meaning))
as Box<dyn StructuralPageDecoder>,
)
}
.boxed()
}
}
impl StructuralPageScheduler for BlobDescriptionPageScheduler {
fn needs_initialization(&self) -> bool {
self.inner_scheduler.needs_initialization()
}
fn init_layout(&self) -> Result<PageInitialization> {
self.inner_scheduler.init_layout()
}
fn init_from_buffers<'a>(
&'a mut self,
buffers: PageInitializationBuffers,
io: &Arc<dyn EncodingsIo>,
) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
self.inner_scheduler.init_from_buffers(buffers, io)
}
fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> Result<()> {
self.inner_scheduler.try_load(data)
}
fn cache_view(&self) -> PageCacheView {
PageCacheView::BlobDescriptor
}
fn schedule_ranges(
&self,
ranges: &[Range<u64>],
io: &Arc<dyn EncodingsIo>,
) -> Result<Vec<PageLoadTask>> {
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<dyn StructuralPageDecoder>,
def_meaning: Arc<[DefinitionInterpretation]>,
}
impl BlobDescriptionPageDecoder {
fn new(
inner: Box<dyn StructuralPageDecoder>,
def_meaning: Arc<[DefinitionInterpretation]>,
) -> Self {
Self { inner, def_meaning }
}
}
impl StructuralPageDecoder for BlobDescriptionPageDecoder {
fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
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<dyn DecodePageTask>,
def_meaning: Arc<[DefinitionInterpretation]>,
}
impl BlobDescriptionDecodePageTask {
fn new(inner: Box<dyn DecodePageTask>, def_meaning: Arc<[DefinitionInterpretation]>) -> Self {
Self { inner, def_meaning }
}
}
impl DecodePageTask for BlobDescriptionDecodePageTask {
fn decode(self: Box<Self>) -> Result<DecodedPage> {
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::<u64>();
let sizes = sizes.data.borrow_to_typed_slice::<u64>();
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<UInt64Array>,
sizes: Arc<UInt64Array>,
}
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<Self>) -> Arc<dyn std::any::Any + Send + Sync + 'static> {
self
}
}
#[derive(Debug)]
pub(super) struct BlobPageScheduler {
inner_scheduler: Box<dyn StructuralPageScheduler>,
row_number: u64,
num_rows: u64,
def_meaning: Arc<[DefinitionInterpretation]>,
positions: Option<Arc<UInt64Array>>,
sizes: Option<Arc<UInt64Array>>,
}
impl BlobPageScheduler {
pub fn new(
inner_scheduler: Box<dyn StructuralPageScheduler>,
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<Range<u64>>,
mut loaded_blobs: Vec<LoadedBlob>,
first_row_number: u64,
io: &dyn EncodingsIo,
def_meaning: Arc<[DefinitionInterpretation]>,
) -> Result<PageLoadTask> {
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<dyn StructuralPageDecoder>)
}
.boxed();
Ok(PageLoadTask {
decoder_fut,
num_rows,
})
}
}
impl StructuralPageScheduler for BlobPageScheduler {
fn init_layout(&self) -> Result<PageInitialization> {
self.inner_scheduler.init_layout()
}
fn init_from_buffers<'a>(
&'a mut self,
buffers: PageInitializationBuffers,
io: &Arc<dyn EncodingsIo>,
) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
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::<UInt64Array>()
.unwrap()
.clone(),
);
let sizes = Arc::new(
descs
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap()
.clone(),
);
let state = Arc::new(BlobCacheableState { positions, sizes });
Ok(state as Arc<dyn CachedPageData>)
}
.boxed()
}
fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> Result<()> {
let blob_state = data
.clone()
.as_arc_any()
.downcast::<BlobCacheableState>()
.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<u64>],
io: &Arc<dyn EncodingsIo>,
) -> Result<Vec<PageLoadTask>> {
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<Bytes>,
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<LoadedBlob>,
def_meaning: Arc<[DefinitionInterpretation]>,
num_rows: u64,
}
impl BlobPageDecoder {
fn new(blobs: Vec<LoadedBlob>, 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<Box<dyn DecodePageTask>> {
let blobs = self.blobs.drain(0..num_rows as usize).collect::<Vec<_>>();
Ok(Box::new(BlobDecodePageTask::new(
blobs,
self.def_meaning.clone(),
)))
}
fn num_rows(&self) -> u64 {
self.num_rows
}
}
#[derive(Debug)]
struct BlobDecodePageTask {
blobs: Vec<LoadedBlob>,
def_meaning: Arc<[DefinitionInterpretation]>,
}
impl BlobDecodePageTask {
fn new(blobs: Vec<LoadedBlob>, def_meaning: Arc<[DefinitionInterpretation]>) -> Self {
Self { blobs, def_meaning }
}
}
impl DecodePageTask for BlobDecodePageTask {
fn decode(self: Box<Self>) -> Result<DecodedPage> {
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::<usize>();
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),
})
}
}
vendor/lance-encoding/src/encodings/logical/primitive/chunk_index.rs
Vendored
+536
View File
@@ -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<u32>),
U64(Vec<u64>),
}
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<Item = u64>, 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<u64>) -> 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<u8>,
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<u64> {
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::<u64>();
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<RefBlock> {
let buffer = LanceBuffer::from(rep_bytes.to_vec());
let u64_slice = buffer.borrow_to_typed_slice::<u64>();
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<u8> {
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<u64>> = 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}]");
}
}
}
}
vendor/lance-encoding/src/encodings/logical/primitive/constant.rs
Vendored
+522
View File
@@ -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<EncodedPage> {
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::<Vec<_>>()
} 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<ScalarBuffer<u16>>,
def: Option<ScalarBuffer<u16>>,
}
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<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
self
}
}
#[derive(Debug, Clone)]
enum ScalarSource {
Inline(Vec<u8>),
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<Range<u64>>,
rep: Option<Range<u64>>,
def: Option<Range<u64>>,
}
#[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<Arc<CachedConstantState>>,
}
impl ConstantPageScheduler {
pub fn try_new(
buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
inline_value: Option<Bytes>,
data_type: DataType,
def_meaning: Arc<[DefinitionInterpretation]>,
) -> Result<Self> {
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<PageInitialization> {
// 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<dyn EncodingsIo>,
) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
// 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::<u16>()
});
let def = def_buffer.map(|def| {
let def = LanceBuffer::from_bytes(def, 2);
def.borrow_to_typed_slice::<u16>()
});
let cached = Arc::new(CachedConstantState { scalar, rep, def });
Ok(cached as Arc<dyn CachedPageData>)
}
.boxed()
}
fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> Result<()> {
let cached: Arc<CachedConstantState> =
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<u64>],
_io: &Arc<dyn EncodingsIo>,
) -> Result<Vec<PageLoadTask>> {
let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
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<dyn crate::encodings::logical::primitive::StructuralPageDecoder>;
Ok(vec![PageLoadTask {
decoder_fut: std::future::ready(Ok(decoder)).boxed(),
num_rows,
}])
}
}
#[derive(Debug)]
struct ConstantPageDecoder {
ranges: VecDeque<Range<u64>>,
scalar: ArrayRef,
rep: Option<ScalarBuffer<u16>>,
def: Option<ScalarBuffer<u16>>,
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<Range<u64>> {
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<usize>, 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<Box<dyn crate::decoder::DecodePageTask>> {
let drained_ranges = self.drain_ranges(num_rows);
let mut level_slices: Vec<Range<usize>> = 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<ScalarBuffer<u16>>,
def: Option<ScalarBuffer<u16>>,
level_slices: Vec<Range<usize>>,
visible_items_total: u64,
def_meaning: Arc<[DefinitionInterpretation]>,
max_visible_def: u16,
}
impl DecodeConstantTask {
fn slice_levels(
levels: &Option<ScalarBuffer<u16>>,
slices: &[Range<usize>],
) -> Option<Vec<u16>> {
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<ArrayRef> {
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<Self>) -> Result<crate::decoder::DecodedPage> {
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<dyn CachedPageData> = 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"));
}
}
vendor/lance-encoding/src/encodings/logical/primitive/dict.rs
Vendored
+632
View File
@@ -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<K: ArrowDictionaryKeyType>(
array: Arc<dyn Array>,
) -> Result<Arc<dyn Array>> {
// TODO: Fast path when there is only one null index? (common case)
let dict_array = array.as_dictionary_opt::<K>().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::<K>::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<dyn Array>)
}
/// 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<dyn Array>) -> Result<Arc<dyn Array>> {
match array.data_type() {
DataType::Dictionary(key_type, _) => match key_type.as_ref() {
DataType::UInt8 => normalize_dict_nulls_impl::<UInt8Type>(array),
DataType::UInt16 => normalize_dict_nulls_impl::<UInt16Type>(array),
DataType::UInt32 => normalize_dict_nulls_impl::<UInt32Type>(array),
DataType::UInt64 => normalize_dict_nulls_impl::<UInt64Type>(array),
DataType::Int8 => normalize_dict_nulls_impl::<Int8Type>(array),
DataType::Int16 => normalize_dict_nulls_impl::<Int16Type>(array),
DataType::Int32 => normalize_dict_nulls_impl::<Int32Type>(array),
DataType::Int64 => normalize_dict_nulls_impl::<Int64Type>(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<K: ArrowDictionaryKeyType>(
array: Arc<dyn Array>,
) -> Result<Arc<dyn Array>> {
let dict_array = array.as_dictionary_opt::<K>().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::<K>::from_iter(dict_array.keys().iter());
let values = dict_array.values().clone();
Ok(Arc::new(DictionaryArray::<K>::try_new(keys, values)?) as Arc<dyn Array>)
}
/// 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<dyn Array>) -> Result<Arc<dyn Array>> {
match array.data_type() {
DataType::Dictionary(key_type, _) => match key_type.as_ref() {
DataType::UInt8 => clear_out_of_range_null_keys_impl::<UInt8Type>(array),
DataType::UInt16 => clear_out_of_range_null_keys_impl::<UInt16Type>(array),
DataType::UInt32 => clear_out_of_range_null_keys_impl::<UInt32Type>(array),
DataType::UInt64 => clear_out_of_range_null_keys_impl::<UInt64Type>(array),
DataType::Int8 => clear_out_of_range_null_keys_impl::<Int8Type>(array),
DataType::Int16 => clear_out_of_range_null_keys_impl::<Int16Type>(array),
DataType::Int32 => clear_out_of_range_null_keys_impl::<Int32Type>(array),
DataType::Int64 => clear_out_of_range_null_keys_impl::<Int64Type>(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<T>(
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<T>,
{
use std::collections::hash_map::Entry;
let mut map = HashMap::new();
let offsets = variable_width_data_block
.offsets
.borrow_to_typed_slice::<T>();
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::<UInt64Type>().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<u8> = 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::<u64>();
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::<u128>();
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::<u32>(
variable_width_data_block,
32,
max_dict_entries,
max_encoded_size,
),
64 => dict_encode_variable_width::<u64>(
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<dyn Array>);
// 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<dyn Array>);
// 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<dyn Array>);
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<dyn Array>);
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<dyn Array>);
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());
}
}
vendor/lance-encoding/src/encodings/logical/primitive/fullzip.rs
Vendored
+54
View File
@@ -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)>;
}
vendor/lance-encoding/src/encodings/logical/primitive/layout.rs
Vendored
+72
View File
@@ -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<T>(
requested_encoding: Option<&str>,
dense_budget: &MiniBlockRepDefBudget,
candidate: impl FnOnce() -> Result<Option<T>>,
) -> Result<Option<T>> {
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());
}
}
}
vendor/lance-encoding/src/encodings/logical/primitive/miniblock.rs
Vendored
+201
View File
@@ -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<u64> =
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<LanceBuffer>,
/// Describes the size of each chunk
pub chunks: Vec<MiniBlockChunk>,
/// 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<u32>,
// 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") };
}
}
vendor/lance-encoding/src/encodings/logical/primitive/sparse.rs
Vendored
+5628
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs
Vendored
+2407
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/encodings/logical/struct.rs
Vendored
+1180
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/encodings/physical.rs
Vendored
+14
View File
@@ -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;
vendor/lance-encoding/src/encodings/physical/binary.rs
Vendored
+1436
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/encodings/physical/bitpacking.rs
Vendored
+897
View File
@@ -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<T: ArrowNativeType + BitPackingUninit>(
data: FixedWidthDataBlock,
) -> MiniBlockCompressed {
debug_assert!(data.num_values > 0);
let data_buffer = data.data.borrow_to_typed_slice::<T>();
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::<PrimitiveArray<UInt64Type>>()
.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<T> = 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::<T>()) 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<T> = 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::<T>())
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::<u8>(data),
16 => Self::bitpack_chunked::<u16>(data),
32 => Self::bitpack_chunked::<u32>(data),
64 => Self::bitpack_chunked::<u64>(data),
_ => unreachable!(),
};
(
compressed,
ProtobufUtils21::inline_bitpacking(
bits_per_value,
// TODO: Could potentially compress the data here
None,
),
)
}
fn unchunk<T: ArrowNativeType + BitPackingUninit + Pod>(
data: LanceBuffer,
num_values: u64,
) -> Result<DataBlock> {
// This macro decompresses a chunk(1024 values) of bitpacked values.
let uncompressed_bit_width = std::mem::size_of::<T>() * 8;
let word_size = std::mem::size_of::<T>();
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::<T>();
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<LanceBuffer>, 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<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
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::<u8>(data, num_values),
16 => Self::unchunk::<u16>(data, num_values),
32 => Self::unchunk::<u32>(data, num_values),
64 => Self::unchunk::<u64>(data, num_values),
_ => unimplemented!("Bitpacking word size must be 8, 16, 32, or 64"),
}
}
fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
num_values
.checked_mul(self.uncompressed_bit_width)
.map(|bits| bits.div_ceil(8))
}
}
impl BlockDecompressor for InlineBitpacking {
fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
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::<u8>(data, num_values),
16 => Self::unchunk::<u16>(data, num_values),
32 => Self::unchunk::<u32>(data, num_values),
64 => Self::unchunk::<u64>(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<T: ArrowNativeType + BitPackingUninit>(
data: FixedWidthDataBlock,
compressed_bits_per_value: usize,
) -> LanceBuffer {
let data_buffer = data.data.borrow_to_typed_slice::<T>();
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<T> = 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<T> = 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<T: ArrowNativeType + BitPackingUninit>(
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::<T>();
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<T> =
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<LanceBuffer>, 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::<u8>(fixed_width, self.compressed_bit_width as usize),
16 => bitpack_out_of_line::<u16>(fixed_width, self.compressed_bit_width as usize),
32 => bitpack_out_of_line::<u32>(fixed_width, self.compressed_bit_width as usize),
64 => bitpack_out_of_line::<u64>(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<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
let data = require_block_payload(data, "Out-of-line bitpacking")?;
let word_size = match self.uncompressed_bit_width {
8 => std::mem::size_of::<u8>(),
16 => std::mem::size_of::<u16>(),
32 => std::mem::size_of::<u32>(),
64 => std::mem::size_of::<u64>(),
_ => 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::<u8>(
block,
num_values as usize,
self.compressed_bit_width as usize,
),
16 => unpack_out_of_line::<u16>(
block,
num_values as usize,
self.compressed_bit_width as usize,
),
32 => unpack_out_of_line::<u32>(
block,
num_values as usize,
self.compressed_bit_width as usize,
),
64 => unpack_out_of_line::<u64>(
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<T>(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::<T>() * 8);
let mut chunk: Vec<T> = 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::<T>(data, num_values).unwrap();
let DataBlock::FixedWidth(fixed) = decoded else {
panic!("expected FixedWidth DataBlock");
};
let decoded_values = fixed.data.borrow_to_typed_view::<T>();
assert_eq!(decoded_values.as_ref(), values);
}
fn assert_corrupt_unchunk<T>(data: LanceBuffer, num_values: u64, expected_message: &str)
where
T: ArrowNativeType + BitPackingUninit + Pod,
{
let err = InlineBitpacking::unchunk::<T>(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<u32> = (0..500).map(|i| ((i * 7) % (1 << 12)) as u32).collect();
roundtrip_unchunk(&values, 12);
}
#[test]
fn unchunk_u64_bw23_full() {
let values: Vec<u64> = (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::<u32>(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<dyn Array>,
Arc::new(Int8Array::from(vec![1; 1024])) as Arc<dyn Array>,
Arc::new(Int8Array::from(vec![16; 1024])) as Arc<dyn Array>,
Arc::new(Int8Array::from(vec![-1; 1024])) as Arc<dyn Array>,
Arc::new(Int8Array::from(vec![5; 1])) as Arc<dyn Array>,
];
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<dyn Array>];
// 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<dyn Array>];
// 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::<u32>() as u64 * 8;
let values: Vec<u32> = (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::<u32>(input, bit_width);
let compressed_words = compressed.borrow_to_typed_slice::<u32>().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::<u32>(compressed_block, values.len(), bit_width);
let decoded_values = decoded.data.borrow_to_typed_slice::<u32>();
assert_eq!(decoded_values.as_ref(), values.as_slice());
}
}
vendor/lance-encoding/src/encodings/physical/block.rs
Vendored
+884
View File
@@ -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<i32>,
}
impl CompressionConfig {
/// Create a compression configuration for an encoding mechanism.
pub fn new(scheme: CompressionScheme, level: Option<i32>) -> 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<i32> {
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<CompressionScheme> for pb21::CompressionScheme {
type Error = Error;
fn try_from(scheme: CompressionScheme) -> Result<Self> {
match scheme {
CompressionScheme::Lz4 => Ok(Self::CompressionAlgorithmLz4),
CompressionScheme::Zstd => Ok(Self::CompressionAlgorithmZstd),
_ => Err(Error::invalid_input(format!(
"Unsupported compression scheme: {:?}",
scheme
))),
}
}
}
impl TryFrom<pb21::CompressionScheme> for CompressionScheme {
type Error = Error;
fn try_from(scheme: pb21::CompressionScheme) -> Result<Self> {
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<Self> {
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<u8>) -> Result<()>;
fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> 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<std::result::Result<Mutex<Compressor<'static>>, 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<Compressor<'static>>> {
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<u8>,
) -> 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<u8>) -> 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<u8>) -> 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<u8>) -> 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<u8>) -> 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<u8>) -> Result<()> {
output_buf.extend_from_slice(input_buf);
Ok(())
}
fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> 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<Box<dyn BufferCompressor>> {
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<dyn BlockDecompressor>,
compressor: Box<dyn BufferCompressor>,
}
impl GeneralBlockDecompressor {
pub fn try_new(
inner: Box<dyn BlockDecompressor>,
compression: CompressionConfig,
) -> Result<Self> {
let compressor = GeneralBufferCompressor::get_compressor(compression)?;
Ok(Self { inner, compressor })
}
}
impl BlockDecompressor for GeneralBlockDecompressor {
fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
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<dyn BufferCompressor>,
// 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<Self> {
let compressor = GeneralBufferCompressor::get_compressor(compression_config)?;
Ok(Self {
compressor,
block_compression: compression_config,
})
}
pub fn from_scheme(scheme: pb21::CompressionScheme) -> Result<Self> {
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<T: ArrowNativeType>(
&self,
data: &[u8],
offsets: &[T],
compressed: &mut Vec<u8>,
) -> Result<LanceBuffer> {
let mut new_offsets: Vec<T> = 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<T: ArrowNativeType>(
&self,
data: &[u8],
offsets: &[T],
decompressed: &mut Vec<u8>,
) -> Result<LanceBuffer> {
let mut new_offsets: Vec<T> = 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::<u32>(
data_bytes,
&data.offsets.borrow_to_typed_slice::<u32>(),
&mut compressed,
)?,
64 => self.per_value_compress::<u64>(
data_bytes,
&data.offsets.borrow_to_typed_slice::<u64>(),
&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<DataBlock> {
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::<u32>(),
&mut decompressed,
)?,
64 => self.per_value_decompress(
data_bytes,
&data.offsets.borrow_to_typed_slice::<u64>(),
&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<LanceBuffer>, 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<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
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;
}
}
}
vendor/lance-encoding/src/encodings/physical/byte_stream_split.rs
Vendored
+452
View File
@@ -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<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
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<u64> {
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::<UInt64Type>();
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<f32> = vec![
1.0,
2.5,
-3.7,
4.2,
0.0,
-0.0,
f32::INFINITY,
f32::NEG_INFINITY,
];
let bytes: Vec<u8> = 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<f32> = 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<f64> = vec![
1.0,
2.5,
-3.7,
4.2,
0.0,
-0.0,
f64::INFINITY,
f64::NEG_INFINITY,
];
let bytes: Vec<u8> = 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<f64> = 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);
}
}
vendor/lance-encoding/src/encodings/physical/constant.rs
Vendored
+81
View File
@@ -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<LanceBuffer>,
}
impl ConstantDecompressor {
pub fn new(scalar: Option<LanceBuffer>) -> Self {
Self { scalar }
}
}
impl BlockDecompressor for ConstantDecompressor {
fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
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<DataBlock> {
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()
);
}
}
vendor/lance-encoding/src/encodings/physical/fsst.rs
Vendored
+469
View File
@@ -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<u8>,
}
impl FsstCompressed {
fn fsst_compress(data: DataBlock) -> Result<Self> {
match data {
DataBlock::VariableWidth(variable_width) => {
match variable_width.bits_per_offset {
32 => {
let offsets = variable_width.offsets.borrow_to_typed_slice::<i32>();
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::<i64>();
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<i64>,
}
impl FsstMiniBlockEncoder {
pub fn new(minichunk_size: Option<i64>) -> 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<dyn MiniBlockCompressor>;
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<dyn PerValueCompressor>,
}
impl FsstPerValueEncoder {
pub fn new(inner: Box<dyn PerValueCompressor>) -> 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<dyn VariablePerValueDecompressor>,
}
impl FsstPerValueDecompressor {
pub fn new(
symbol_table: LanceBuffer,
inner_decompressor: Box<dyn VariablePerValueDecompressor>,
) -> Self {
Self {
symbol_table,
inner_decompressor,
}
}
}
impl VariablePerValueDecompressor for FsstPerValueDecompressor {
fn decompress(&self, data: VariableWidthBlock) -> Result<DataBlock> {
// 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::<u8>();
let bytes = bytes.as_ref();
match compressed_variable_data.bits_per_offset {
32 => {
let offsets = compressed_variable_data
.offsets
.borrow_to_typed_slice::<i32>();
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::<i64>();
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<dyn MiniBlockDecompressor>,
}
impl FsstMiniBlockDecompressor {
pub fn new(
description: &pb21::Fsst,
inner_decompressor: Box<dyn MiniBlockDecompressor>,
) -> Self {
Self {
symbol_table: LanceBuffer::from_bytes(description.symbol_table.clone(), 1),
inner_decompressor,
}
}
}
impl MiniBlockDecompressor for FsstMiniBlockDecompressor {
fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
// 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::<i64>();
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::<i32>();
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}");
}
}
vendor/lance-encoding/src/encodings/physical/general.rs
Vendored
+683
View File
@@ -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<dyn MiniBlockCompressor>,
compression: CompressionConfig,
}
impl GeneralMiniBlockCompressor {
pub fn new(inner: Box<dyn MiniBlockCompressor>, 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<dyn MiniBlockDecompressor>,
compression: CompressionConfig,
}
impl GeneralMiniBlockDecompressor {
pub fn new(inner: Box<dyn MiniBlockDecompressor>, compression: CompressionConfig) -> Self {
Self { inner, compression }
}
}
impl MiniBlockDecompressor for GeneralMiniBlockDecompressor {
fn decompress(&self, mut data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
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<u64> {
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<dyn MiniBlockCompressor>,
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<TestCase> {
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<i32>) -> DataBlock {
let array = Int32Array::from(values);
DataBlock::from_array(array)
}
fn create_pattern_i32_block<F>(size: usize, pattern: F) -> DataBlock
where
F: Fn(usize) -> i32,
{
let values: Vec<i32> = (0..size).map(pattern).collect();
let array = Int32Array::from(values);
DataBlock::from_array(array)
}
fn create_pattern_f64_block<F>(size: usize, pattern: F) -> DataBlock
where
F: Fn(usize) -> f64,
{
let values: Vec<f64> = (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<u8> {
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<i32> = (0..1024).collect();
let data = LanceBuffer::from_bytes(
bytemuck::cast_slice(&values).to_vec().into(),
std::mem::align_of::<i32>() 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<i64> = 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<i32>);
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
}
}
}
}
vendor/lance-encoding/src/encodings/physical/packed.rs
Vendored
+1571
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/encodings/physical/rle.rs
Vendored
+3408
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/encodings/physical/value.rs
Vendored
+1396
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/format.rs
Vendored
+773
View File
@@ -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<u8>) -> 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<CompressionConfig>,
) -> 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<u8>) -> 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<ArrayEncoding>,
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<u32>,
) -> 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::BufferCompression>,
) -> 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<bytes::Bytes>,
) -> 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::BufferCompression>,
) -> 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::BufferCompression>,
) -> 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<crate::format::$module::BufferCompression> {
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<crate::format::$module::CompressiveEncoding> {
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<u8>,
) -> 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<crate::format::$module::CompressiveEncoding>,
def_encoding: Option<crate::format::$module::CompressiveEncoding>,
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<Vec<u8>>,
) -> 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<crate::format::pb21::CompressiveEncoding>,
def_compression: Option<crate::format::pb21::CompressiveEncoding>,
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<u64>,
) -> 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::variable_packed_struct::FieldEncoding>,
) -> 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,
),
),
}
}
}
vendor/lance-encoding/src/lib.rs
Vendored
+134
View File
@@ -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<Range<u64>>,
priority: u64,
) -> BoxFuture<'static, Result<Vec<Bytes>>>;
/// 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<u64>,
priority: u64,
) -> BoxFuture<'static, lance_core::Result<bytes::Bytes>> {
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<Arc<dyn EncodingsIo>> {
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<dyn lance_core::utils::io_stats::IoStatsRecorder>,
) -> Option<Arc<dyn EncodingsIo>> {
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<u64>) -> Bytes {
self.data.slice(req.start as usize..req.end as usize)
}
}
impl EncodingsIo for BufferScheduler {
fn submit_request(
&self,
ranges: Vec<Range<u64>>,
_priority: u64,
) -> BoxFuture<'static, Result<Vec<Bytes>>> {
std::future::ready(Ok(ranges
.into_iter()
.map(|range| self.satisfy_request(range))
.collect::<Vec<_>>()))
.boxed()
}
}
vendor/lance-encoding/src/repdef.rs
Vendored
+4083
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/statistics.rs
Vendored
+1269
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/testing.rs
Vendored
+1588
View File
File diff suppressed because it is too large Load Diff
vendor/lance-encoding/src/utils.rs
Vendored
+7
View File
@@ -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;
vendor/lance-encoding/src/utils/accumulation.rs
Vendored
+105
View File
@@ -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<ArrayRef>,
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<ArrayRef>, 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<ArrayRef>, 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,
))
}
}
}
vendor/lance-encoding/src/utils/bytepack.rs
Vendored
+306
View File
@@ -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<u8>,
}
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<u8>,
}
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<u8>,
}
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<u8>,
}
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<u8> {
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<I: Iterator<Item = u8>> {
U8(I),
U16(I),
U32(I),
U64(I),
}
impl<T: Iterator<Item = u8>> ByteUnpacker<T> {
#[allow(clippy::new_ret_no_self)]
pub fn new<I: IntoIterator<IntoIter = T>>(data: I, size: usize) -> impl Iterator<Item = u64> {
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<I: Iterator<Item = u8>> Iterator for ByteUnpacker<I> {
type Item = u64;
fn next(&mut self) -> Option<Self::Item> {
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<_>>(),
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<_>>(),
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<_>>(),
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<_>>(),
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());
}
}