lance_encoding/encodings/physical/
constant.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Routines for compressing and decompressing constant-encoded data
5
6use crate::{
7    buffer::LanceBuffer,
8    compression::{
9        BlockCompressor, BlockDecompressor, FixedPerValueDecompressor, require_no_block_payload,
10    },
11    data::{AllNullDataBlock, ConstantDataBlock, DataBlock, FixedWidthDataBlock},
12    encodings::physical::{checked_fixed_values, try_vec_with_capacity},
13    format::{EncodingUtils, metadata::CompressiveEncoding},
14};
15
16use lance_core::{Error, Result};
17
18/// Metadata-only compressor for a repeated unsigned `u32` or `u64` value.
19#[derive(Debug)]
20pub struct ConstantEncoder {
21    bits_per_value: u64,
22    value: u64,
23}
24
25impl ConstantEncoder {
26    pub fn new(bits_per_value: u64, value: u64) -> Self {
27        Self {
28            bits_per_value,
29            value,
30        }
31    }
32
33    /// Describes a known constant sequence without materializing its values.
34    /// Zero values have no scalar; non-empty sequences store one typed scalar.
35    pub(crate) fn encoding(&self, num_values: u64) -> Result<CompressiveEncoding> {
36        if !matches!(self.bits_per_value, 32 | 64) {
37            return Err(Error::invalid_input(format!(
38                "Constant block compression only supports 32 or 64-bit values, got {}",
39                self.bits_per_value
40            )));
41        }
42        if num_values == 0 {
43            return Ok(EncodingUtils::constant(None));
44        }
45        let scalar = if self.bits_per_value == 32 {
46            let value = u32::try_from(self.value).map_err(|_| {
47                Error::invalid_input(format!("Constant value {} exceeds u32::MAX", self.value))
48            })?;
49            bytes::Bytes::copy_from_slice(&value.to_le_bytes())
50        } else {
51            bytes::Bytes::copy_from_slice(&self.value.to_le_bytes())
52        };
53        Ok(EncodingUtils::constant(Some(scalar)))
54    }
55}
56
57impl BlockCompressor for ConstantEncoder {
58    fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)> {
59        let DataBlock::FixedWidth(data) = data else {
60            return Err(Error::invalid_input(
61                "Constant block compression requires fixed-width data",
62            ));
63        };
64        if data.bits_per_value != self.bits_per_value {
65            return Err(Error::invalid_input(format!(
66                "Constant codec expects {}-bit values, got {}",
67                self.bits_per_value, data.bits_per_value
68            )));
69        }
70        let encoding = self.encoding(data.num_values)?;
71        match self.bits_per_value {
72            32 => {
73                if checked_fixed_values::<u32>(&data, "Constant input")?
74                    .iter()
75                    .any(|candidate| u64::from(*candidate) != self.value)
76                {
77                    return Err(Error::invalid_input(
78                        "Constant input contains a different value",
79                    ));
80                }
81            }
82            64 => {
83                if checked_fixed_values::<u64>(&data, "Constant input")?
84                    .iter()
85                    .any(|candidate| *candidate != self.value)
86                {
87                    return Err(Error::invalid_input(
88                        "Constant input contains a different value",
89                    ));
90                }
91            }
92            _ => unreachable!("constant width was validated above"),
93        };
94        Ok((None, encoding))
95    }
96}
97
98/// Materializes a metadata-only constant as a typed fixed-width block.
99#[derive(Debug)]
100pub(crate) struct ConstantBlockDecompressor {
101    bits_per_value: u64,
102    value: Option<u64>,
103}
104
105impl ConstantBlockDecompressor {
106    pub(crate) fn new(bits_per_value: u64, value: Option<u64>) -> Self {
107        Self {
108            bits_per_value,
109            value,
110        }
111    }
112}
113
114impl BlockDecompressor for ConstantBlockDecompressor {
115    fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
116        require_no_block_payload(data, "Constant")?;
117        if self.value.is_none() != (num_values == 0) {
118            return Err(Error::invalid_input(
119                "Constant scalar must be absent exactly when cardinality is zero",
120            ));
121        }
122        let value = self.value.unwrap_or(0);
123        let output_len = usize::try_from(num_values)
124            .map_err(|_| Error::invalid_input("Constant output cardinality does not fit usize"))?;
125        let data = match self.bits_per_value {
126            32 => {
127                let value = u32::try_from(value).map_err(|_| {
128                    Error::invalid_input(format!("Constant value {value} exceeds u32::MAX"))
129                })?;
130                let mut values = try_vec_with_capacity::<u32>(num_values, "Constant output")?;
131                values.resize(output_len, value);
132                LanceBuffer::reinterpret_vec(values)
133            }
134            64 => {
135                let mut values = try_vec_with_capacity::<u64>(num_values, "Constant output")?;
136                values.resize(output_len, value);
137                LanceBuffer::reinterpret_vec(values)
138            }
139            bits_per_value => {
140                return Err(Error::invalid_input(format!(
141                    "Constant block decompression only supports 32 or 64-bit values, got {bits_per_value}"
142                )));
143            }
144        };
145        Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
146            bits_per_value: self.bits_per_value,
147            data,
148            num_values,
149            block_info: Default::default(),
150        }))
151    }
152
153    fn requires_payload(&self) -> bool {
154        false
155    }
156}
157
158/// A decompressor for constant-encoded data
159#[derive(Debug)]
160pub struct ConstantDecompressor {
161    scalar: Option<LanceBuffer>,
162}
163
164impl ConstantDecompressor {
165    pub fn new(scalar: Option<LanceBuffer>) -> Self {
166        Self { scalar }
167    }
168}
169
170impl BlockDecompressor for ConstantDecompressor {
171    fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
172        require_no_block_payload(data, "Constant")?;
173        if let Some(scalar) = self.scalar.clone() {
174            Ok(DataBlock::Constant(ConstantDataBlock {
175                data: scalar,
176                num_values,
177            }))
178        } else {
179            Ok(DataBlock::AllNull(AllNullDataBlock { num_values }))
180        }
181    }
182
183    fn requires_payload(&self) -> bool {
184        false
185    }
186}
187
188impl FixedPerValueDecompressor for ConstantDecompressor {
189    fn decompress(&self, _data: FixedWidthDataBlock, num_values: u64) -> Result<DataBlock> {
190        if let Some(scalar) = self.scalar.clone() {
191            Ok(DataBlock::Constant(ConstantDataBlock {
192                data: scalar,
193                num_values,
194            }))
195        } else {
196            Ok(DataBlock::AllNull(AllNullDataBlock { num_values }))
197        }
198    }
199
200    fn bits_per_value(&self) -> u64 {
201        self.scalar
202            .as_ref()
203            .map(|s| s.len() as u64 * 8)
204            .unwrap_or(0)
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[rstest::rstest]
213    #[case(32)]
214    #[case(64)]
215    fn constant_metadata_round_trip(#[case] bits: u64) {
216        let encoder = ConstantEncoder::new(bits, 7);
217        for count in [0, 3] {
218            let expected = if bits == 32 {
219                LanceBuffer::reinterpret_vec(vec![7_u32; count as usize])
220            } else {
221                LanceBuffer::reinterpret_vec(vec![7_u64; count as usize])
222            };
223            let (payload, encoding) = encoder
224                .compress(DataBlock::FixedWidth(FixedWidthDataBlock {
225                    data: expected.clone(),
226                    bits_per_value: bits,
227                    num_values: count,
228                    block_info: Default::default(),
229                }))
230                .unwrap();
231            assert!(payload.is_none());
232            let decoder = crate::encodings::physical::sequence::SequenceDecoder::try_new(
233                &encoding, bits, count,
234            )
235            .unwrap();
236            let decoded = decoder.decode(payload).unwrap().as_fixed_width().unwrap();
237            assert_eq!(decoded.num_values, count);
238            assert_eq!(decoded.data.as_ref(), expected.as_ref());
239        }
240        assert_eq!(
241            encoder.encoding(u64::MAX).unwrap(),
242            encoder.encoding(1).unwrap()
243        );
244    }
245
246    #[test]
247    fn constant_metadata_validates_width_and_scalar() {
248        for encoder in [
249            ConstantEncoder::new(16, 1),
250            ConstantEncoder::new(32, u64::MAX),
251        ] {
252            let error = encoder.encoding(1).unwrap_err();
253            assert!(matches!(error, Error::InvalidInput { .. }));
254            assert!(error.to_string().contains("Constant"));
255        }
256    }
257
258    #[test]
259    fn block_constant_requires_no_payload() {
260        let decompressor = ConstantDecompressor::new(None);
261
262        assert!(!decompressor.requires_payload());
263        assert!(matches!(
264            BlockDecompressor::decompress(&decompressor, None, 3).unwrap(),
265            DataBlock::AllNull(AllNullDataBlock { num_values: 3 })
266        ));
267        assert!(
268            BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::empty()), 3).is_err()
269        );
270    }
271}
lance_encoding/encodings/physical/constant.rs

lance_encoding/encodings/physical/
constant.rs
