mito2/sst/parquet/
stats.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Statistics of parquet SSTs.
16
17use std::borrow::Borrow;
18use std::collections::HashSet;
19use std::sync::Arc;
20
21use api::v1::SemanticType;
22use datafusion_common::pruning::PruningStatistics;
23use datafusion_common::{Column, ScalarValue};
24use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt64Array};
25use datatypes::data_type::DataType;
26use parquet::file::metadata::RowGroupMetaData;
27use store_api::metadata::RegionMetadataRef;
28use store_api::storage::ColumnId;
29
30use crate::sst::parquet::flat_format::FlatReadFormat;
31use crate::sst::parquet::format::StatValues;
32
33/// Statistics for pruning row groups.
34pub(crate) struct RowGroupPruningStats<'a, T> {
35    /// Metadata of SST row groups.
36    row_groups: &'a [T],
37    /// Helper to read the SST.
38    read_format: &'a FlatReadFormat,
39    /// The metadata of the region.
40    /// It contains the schema a query expects to read. If it is not None, we use it instead
41    /// of the metadata in the SST to get the column id of a column as the SST may have
42    /// different columns.
43    expected_metadata: Option<RegionMetadataRef>,
44    /// If true, skip columns with Field semantic type during pruning.
45    skip_fields: bool,
46}
47
48impl<'a, T> RowGroupPruningStats<'a, T> {
49    /// Creates a new statistics to prune specific `row_groups`.
50    pub(crate) fn new(
51        row_groups: &'a [T],
52        read_format: &'a FlatReadFormat,
53        expected_metadata: Option<RegionMetadataRef>,
54        skip_fields: bool,
55    ) -> Self {
56        Self {
57            row_groups,
58            read_format,
59            expected_metadata,
60            skip_fields,
61        }
62    }
63
64    /// Returns the column id of specific column name if we need to read it.
65    /// Prefers the column id in the expected metadata if it exists.
66    /// Returns None if skip_fields is true and the column is a Field.
67    fn column_id_to_prune(&self, name: &str) -> Option<ColumnId> {
68        let metadata = self
69            .expected_metadata
70            .as_ref()
71            .unwrap_or_else(|| self.read_format.metadata());
72        let col = metadata.column_by_name(name)?;
73
74        // Skip field columns when skip_fields is enabled
75        if self.skip_fields && col.semantic_type == SemanticType::Field {
76            return None;
77        }
78
79        Some(col.column_id)
80    }
81
82    /// Casts stats values to the expected data type when the SST stores the
83    /// column with a different type (e.g. after `MODIFY COLUMN` changed the
84    /// column type, including widening the time index unit). The pruning
85    /// predicate is built against the expected schema, so raw file-typed
86    /// stats would prune wrongly.
87    ///
88    /// Timestamp stats are raw integers: first reinterpret them in the file's
89    /// type, then convert to the expected type. A single-step cast would
90    /// reinterpret (not rescale) the value. Returns `None` on cast failure so
91    /// the row group is kept (conservative).
92    fn cast_stats_to_expected(&self, column_id: ColumnId, values: ArrayRef) -> Option<ArrayRef> {
93        // Without expected metadata the file metadata is the expected one,
94        // so there is nothing to cast toward.
95        let Some(expected_metadata) = self.expected_metadata.as_ref() else {
96            return Some(values);
97        };
98        let expected_col = expected_metadata.column_by_id(column_id)?;
99        let file_col = self.read_format.metadata().column_by_id(column_id)?;
100        if expected_col.column_schema.data_type == file_col.column_schema.data_type {
101            return Some(values);
102        }
103        let file_arrow_type = file_col.column_schema.data_type.as_arrow_type();
104        let expected_arrow_type = expected_col.column_schema.data_type.as_arrow_type();
105        let values = if values.data_type() == &file_arrow_type {
106            values
107        } else {
108            datatypes::arrow::compute::cast(&values, &file_arrow_type).ok()?
109        };
110        datatypes::arrow::compute::cast(&values, &expected_arrow_type).ok()
111    }
112
113    /// Returns the default value of all row groups for `column` according to the metadata.
114    fn compat_default_value(&self, column: &str) -> Option<ArrayRef> {
115        let metadata = self.expected_metadata.as_ref()?;
116        let col_metadata = metadata.column_by_name(column)?;
117        col_metadata
118            .column_schema
119            .create_default_vector(self.row_groups.len())
120            .unwrap_or(None)
121            .map(|vector| vector.to_arrow_array())
122    }
123}
124
125impl<T: Borrow<RowGroupMetaData>> RowGroupPruningStats<'_, T> {
126    /// Returns the null count of all row groups for `column` according to the metadata.
127    fn compat_null_count(&self, column: &str) -> Option<ArrayRef> {
128        let metadata = self.expected_metadata.as_ref()?;
129        let col_metadata = metadata.column_by_name(column)?;
130        let value = col_metadata
131            .column_schema
132            .create_default()
133            .unwrap_or(None)?;
134        let values = self.row_groups.iter().map(|meta| {
135            if value.is_null() {
136                u64::try_from(meta.borrow().num_rows()).ok()
137            } else {
138                Some(0)
139            }
140        });
141        Some(Arc::new(UInt64Array::from_iter(values)))
142    }
143}
144
145impl<T: Borrow<RowGroupMetaData>> PruningStatistics for RowGroupPruningStats<'_, T> {
146    fn min_values(&self, column: &Column) -> Option<ArrayRef> {
147        let column_id = self.column_id_to_prune(&column.name)?;
148        match self.read_format.min_values(self.row_groups, column_id) {
149            StatValues::Values(values) => self.cast_stats_to_expected(column_id, values),
150            StatValues::NoColumn => self.compat_default_value(&column.name),
151            StatValues::NoStats => None,
152        }
153    }
154
155    fn max_values(&self, column: &Column) -> Option<ArrayRef> {
156        let column_id = self.column_id_to_prune(&column.name)?;
157        match self.read_format.max_values(self.row_groups, column_id) {
158            StatValues::Values(values) => self.cast_stats_to_expected(column_id, values),
159            StatValues::NoColumn => self.compat_default_value(&column.name),
160            StatValues::NoStats => None,
161        }
162    }
163
164    fn num_containers(&self) -> usize {
165        self.row_groups.len()
166    }
167
168    fn null_counts(&self, column: &Column) -> Option<ArrayRef> {
169        let column_id = self.column_id_to_prune(&column.name)?;
170        match self.read_format.null_counts(self.row_groups, column_id) {
171            StatValues::Values(values) => Some(values),
172            StatValues::NoColumn => self.compat_null_count(&column.name),
173            StatValues::NoStats => None,
174        }
175    }
176
177    fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
178        // TODO(LFC): Impl it.
179        None
180    }
181
182    fn contained(&self, _column: &Column, _values: &HashSet<ScalarValue>) -> Option<BooleanArray> {
183        // TODO(LFC): Impl it.
184        None
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use std::sync::Arc;
191
192    use datafusion_common::Column;
193    use datatypes::arrow::array::{Array, Int64Array, TimestampMicrosecondArray};
194    use datatypes::arrow::datatypes::TimeUnit;
195    use datatypes::prelude::ConcreteDataType;
196    use parquet::basic::Type as PhysicalType;
197    use parquet::file::metadata::{ColumnChunkMetaData, RowGroupMetaData};
198    use parquet::file::statistics::Statistics;
199    use parquet::schema::types::{SchemaDescriptor, Type};
200    use store_api::codec::PrimaryKeyEncoding;
201    use store_api::metadata::RegionMetadataRef;
202
203    use super::*;
204    use crate::read::read_columns::ReadColumns;
205    use crate::test_util::sst_util::sst_region_metadata_with_encoding;
206
207    /// Builds one row group whose `ts` column (the time index of the
208    /// `sst_region_metadata_with_encoding` fixture) carries raw Int64
209    /// statistics `ts_min..=ts_max`, like a real parquet file does.
210    fn row_group_with_ts_stats(
211        read_format: &FlatReadFormat,
212        ts_min: i64,
213        ts_max: i64,
214    ) -> RowGroupMetaData {
215        let ts_idx = read_format.arrow_schema().index_of("ts").unwrap();
216        let fields: Vec<Arc<Type>> = read_format
217            .arrow_schema()
218            .fields()
219            .iter()
220            .map(|field| {
221                Arc::new(
222                    Type::primitive_type_builder(field.name(), PhysicalType::INT64)
223                        .build()
224                        .unwrap(),
225                )
226            })
227            .collect();
228        let schema_descr = Arc::new(SchemaDescriptor::new(Arc::new(
229            Type::group_type_builder("schema")
230                .with_fields(fields)
231                .build()
232                .unwrap(),
233        )));
234        let chunks: Vec<_> = (0..schema_descr.num_columns())
235            .map(|i| {
236                let mut builder = ColumnChunkMetaData::builder(schema_descr.column(i));
237                if i == ts_idx {
238                    builder = builder.set_statistics(Statistics::int64(
239                        Some(ts_min),
240                        Some(ts_max),
241                        None,
242                        Some(0),
243                        true,
244                    ));
245                }
246                builder.build().unwrap()
247            })
248            .collect();
249        RowGroupMetaData::builder(schema_descr)
250            .set_num_rows(10)
251            .set_total_byte_size(0)
252            .set_column_metadata(chunks)
253            .build()
254            .unwrap()
255    }
256
257    /// The same region metadata but with the time index unit widened to
258    /// microsecond, as the region looks after a widening alter.
259    fn expected_metadata_us(file_metadata: &RegionMetadataRef) -> RegionMetadataRef {
260        let mut expected = (**file_metadata).clone();
261        for column in expected.column_metadatas.iter_mut() {
262            if column.column_schema.name == "ts" {
263                column.column_schema.data_type = ConcreteDataType::timestamp_microsecond_datatype();
264            }
265        }
266        Arc::new(expected)
267    }
268
269    fn read_format_for(file_metadata: &RegionMetadataRef) -> FlatReadFormat {
270        FlatReadFormat::new(
271            file_metadata.clone(),
272            ReadColumns::new([0, 1, 2, 3]),
273            None,
274            "test",
275            false,
276        )
277        .unwrap()
278    }
279
280    /// Timestamp stats are raw Int64 in the file's unit: when the expected
281    /// type differs, they must be *rescaled* (1000ms -> 1_000_000us), not
282    /// reinterpreted (1000us).
283    #[test]
284    fn test_row_group_stats_cast_to_expected_unit() {
285        let file_metadata: RegionMetadataRef =
286            Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense));
287        let read_format = read_format_for(&file_metadata);
288        let row_group = row_group_with_ts_stats(&read_format, 1_000, 9_000);
289        let column = Column::new_unqualified("ts");
290
291        // No expected metadata: raw file stats pass through as Int64.
292        let groups = [&row_group];
293        let stats = RowGroupPruningStats::new(&groups, &read_format, None, false);
294        let min = stats.min_values(&column).unwrap();
295        let min = min.as_any().downcast_ref::<Int64Array>().unwrap();
296        assert_eq!(1_000, min.value(0));
297
298        // Same type: passthrough, no cast needed.
299        let groups = [&row_group];
300        let stats =
301            RowGroupPruningStats::new(&groups, &read_format, Some(file_metadata.clone()), false);
302        let min = stats.min_values(&column).unwrap();
303        let min = min.as_any().downcast_ref::<Int64Array>().unwrap();
304        assert_eq!(1_000, min.value(0));
305
306        // Widened expected unit: rescaled, not reinterpreted.
307        let expected = expected_metadata_us(&file_metadata);
308        let groups = [&row_group];
309        let stats = RowGroupPruningStats::new(&groups, &read_format, Some(expected), false);
310        let min = stats.min_values(&column).unwrap();
311        assert_eq!(
312            datatypes::arrow::datatypes::DataType::Timestamp(TimeUnit::Microsecond, None),
313            min.data_type().clone()
314        );
315        let min = min
316            .as_any()
317            .downcast_ref::<TimestampMicrosecondArray>()
318            .unwrap();
319        assert_eq!(1_000_000, min.value(0));
320        let max = stats.max_values(&column).unwrap();
321        let max = max
322            .as_any()
323            .downcast_ref::<TimestampMicrosecondArray>()
324            .unwrap();
325        assert_eq!(9_000_000, max.value(0));
326    }
327}
mito2/sst/parquet/stats.rs

mito2/sst/parquet/
stats.rs
