cobble_table/
scan.rs

1use crate::metadata::TableMetadata;
2use crate::runtime::TableSchemaTransformFactories;
3#[cfg(feature = "ffi")]
4use crate::table::build_scan_options_for_fields;
5use crate::table::{CompiledTable, compile_table, decode_table_scan_row, validate_name};
6use crate::{Result, TableError, TableSchema, Value};
7use bytes::Bytes;
8use cobble::{
9    Config, ReadOnlyDbBuilder, ScanOptions, ScanSplit, ScanSplitScanner, ShardSnapshotRef,
10    VolumeDescriptor, VolumeUsageKind,
11};
12use serde::{Deserialize, Serialize};
13use std::sync::Arc;
14
15const TABLE_SCAN_PLAN_FORMAT: &str = "cobble-table-scan-plan";
16const TABLE_SCAN_PLAN_VERSION: u32 = 1;
17
18/// A serializable, fixed global-snapshot scan plan for one typed table.
19///
20/// A plan does not retain the snapshot against expiration; applications keep
21/// the referenced global and shard snapshots available for their workers.
22#[derive(Clone, Serialize, Deserialize)]
23pub struct TableScanPlan {
24    format: String,
25    version: u32,
26    name: String,
27    global_snapshot_id: u64,
28    total_buckets: u32,
29    metadata: TableMetadata,
30    shards: Vec<ShardSnapshotRef>,
31    source_volumes: Vec<VolumeDescriptor>,
32    #[serde(skip)]
33    auth_source: Option<Config>,
34}
35
36impl TableScanPlan {
37    pub(crate) fn from_global_reader(
38        name: String,
39        metadata: TableMetadata,
40        global_snapshot_id: u64,
41        total_buckets: u32,
42        shards: Vec<ShardSnapshotRef>,
43        config: Config,
44    ) -> Result<Self> {
45        let plan = Self {
46            format: TABLE_SCAN_PLAN_FORMAT.to_string(),
47            version: TABLE_SCAN_PLAN_VERSION,
48            name,
49            global_snapshot_id,
50            total_buckets,
51            metadata,
52            shards,
53            source_volumes: source_volumes(&config),
54            auth_source: Some(config),
55        };
56        plan.validate()?;
57        Ok(plan)
58    }
59
60    /// Return the fixed global snapshot selected for this scan.
61    pub fn snapshot_id(&self) -> u64 {
62        self.global_snapshot_id
63    }
64
65    /// Return the bucket count fixed by the global snapshot.
66    pub fn total_buckets(&self) -> u32 {
67        self.total_buckets
68    }
69
70    /// Return the schema fixed into this plan.
71    pub fn schema(&self) -> &TableSchema {
72        &self.metadata.schema
73    }
74
75    /// Return the total encoded data size reported by all shard snapshots.
76    pub fn data_size_bytes(&self) -> u64 {
77        self.shards.iter().fold(0_u64, |total, shard| {
78            total.saturating_add(shard.data_size_bytes)
79        })
80    }
81
82    /// Produce one independently serializable split per shard snapshot.
83    pub fn splits(&self) -> Result<Vec<TableScanSplit>> {
84        self.validate()?;
85        Ok(self
86            .shards
87            .iter()
88            .cloned()
89            .map(|shard| TableScanSplit {
90                format: self.format.clone(),
91                version: self.version,
92                name: self.name.clone(),
93                global_snapshot_id: self.global_snapshot_id,
94                total_buckets: self.total_buckets,
95                metadata: self.metadata.clone(),
96                split: full_scan_split(shard),
97                source_volumes: self.source_volumes.clone(),
98                auth_source: self.auth_source.clone(),
99            })
100            .collect())
101    }
102
103    fn validate(&self) -> Result<()> {
104        validate_plan(
105            &self.format,
106            self.version,
107            &self.name,
108            self.total_buckets,
109            &self.metadata,
110            &self.source_volumes,
111        )?;
112        if self.shards.is_empty() {
113            return Err(TableError::InvalidSchema(
114                "table scan plan has no shard snapshots".to_string(),
115            ));
116        }
117        Ok(())
118    }
119}
120
121/// A serializable full-table scan assignment for one shard snapshot.
122#[derive(Clone, Serialize, Deserialize)]
123pub struct TableScanSplit {
124    format: String,
125    version: u32,
126    name: String,
127    global_snapshot_id: u64,
128    total_buckets: u32,
129    metadata: TableMetadata,
130    split: ScanSplit,
131    source_volumes: Vec<VolumeDescriptor>,
132    #[serde(skip)]
133    auth_source: Option<Config>,
134}
135
136impl TableScanSplit {
137    /// Return the global snapshot that assigned this shard.
138    pub fn snapshot_id(&self) -> u64 {
139        self.global_snapshot_id
140    }
141
142    /// Return this split's shard snapshot identity.
143    pub fn shard_snapshot_id(&self) -> u64 {
144        self.split.shard.snapshot_id
145    }
146
147    /// Open this shard and return a typed full-scan iterator.
148    pub fn create_scanner(&self, runtime: Config) -> Result<TableScanSplitScanner> {
149        self.create_scanner_with_transforms(runtime, &TableSchemaTransformFactories::default())
150    }
151
152    /// Start configuring a worker-local scanner for this split.
153    pub fn scanner_builder(&self, runtime: Config) -> TableScanSplitScannerBuilder {
154        TableScanSplitScannerBuilder {
155            split: self.clone(),
156            runtime,
157            transforms: TableSchemaTransformFactories::default(),
158        }
159    }
160
161    fn create_scanner_with_transforms(
162        &self,
163        runtime: Config,
164        transforms: &TableSchemaTransformFactories,
165    ) -> Result<TableScanSplitScanner> {
166        let compiled = compile_table(self.metadata.clone(), self.total_buckets)?;
167        let scanner = self.create_raw_scanner_with_transforms(
168            runtime,
169            transforms,
170            ScanOptions::default().with_column_family(self.name.clone()),
171        )?;
172        Ok(TableScanSplitScanner {
173            inner: scanner,
174            compiled,
175        })
176    }
177
178    fn create_raw_scanner_with_transforms(
179        &self,
180        runtime: Config,
181        transforms: &TableSchemaTransformFactories,
182        scan_options: ScanOptions,
183    ) -> Result<ScanSplitScanner> {
184        self.validate()?;
185        let credential_source = self.auth_source.as_ref().unwrap_or(&runtime);
186        let source_volumes = self
187            .source_volumes
188            .iter()
189            .map(|volume| volume.with_credentials_from(credential_source))
190            .collect::<Vec<_>>();
191        let mut config = runtime_read_config(runtime);
192        config.volumes.extend(source_volumes);
193        config.total_buckets = self.total_buckets;
194
195        let builder = transforms.apply_to(ReadOnlyDbBuilder::new(config))?;
196        let scanner = self
197            .split
198            .create_scanner_with_builder(builder, &scan_options)?;
199        Ok(scanner)
200    }
201
202    #[cfg(feature = "ffi")]
203    pub(crate) fn create_projected_raw_scanner(
204        &self,
205        runtime: Config,
206        field_names: &[String],
207        read_ahead_bytes: i64,
208    ) -> Result<ScanSplitScanner> {
209        let compiled = compile_table(self.metadata.clone(), self.total_buckets)?;
210        let mut options = build_scan_options_for_fields(&self.name, compiled, field_names)?;
211        options.read_ahead_bytes = size::Size::from_const(read_ahead_bytes);
212        self.create_raw_scanner_with_transforms(
213            runtime,
214            &TableSchemaTransformFactories::default(),
215            options,
216        )
217    }
218
219    fn validate(&self) -> Result<()> {
220        validate_plan(
221            &self.format,
222            self.version,
223            &self.name,
224            self.total_buckets,
225            &self.metadata,
226            &self.source_volumes,
227        )?;
228        if self.split.shard.ranges.is_empty() {
229            return Err(TableError::InvalidSchema(
230                "table scan split has no bucket ranges".to_string(),
231            ));
232        }
233        if self.split.start.is_some()
234            || self.split.end.is_some()
235            || self.split.start_bucket.is_some()
236            || self.split.start_key_exclusive.is_some()
237            || self.split.end_bucket.is_some()
238            || self.split.end_key_inclusive.is_some()
239        {
240            return Err(TableError::InvalidSchema(
241                "table scan split must cover its complete shard".to_string(),
242            ));
243        }
244        Ok(())
245    }
246}
247
248/// Builder for one worker-local typed scan split scanner.
249pub struct TableScanSplitScannerBuilder {
250    split: TableScanSplit,
251    runtime: Config,
252    transforms: TableSchemaTransformFactories,
253}
254
255impl TableScanSplitScannerBuilder {
256    /// Register a factory for persisted schema transform specifications before opening.
257    pub fn register_schema_transform<F, T>(
258        mut self,
259        transform_type: impl Into<String>,
260        factory: F,
261    ) -> Result<Self>
262    where
263        F: Fn(&[u8]) -> cobble::Result<T> + Send + Sync + 'static,
264        T: Fn(Option<Bytes>) -> cobble::Result<Option<Bytes>> + Send + Sync + 'static,
265    {
266        self.transforms.register(transform_type, factory)?;
267        Ok(self)
268    }
269
270    /// Open this fixed split using the supplied worker runtime configuration.
271    pub fn open(self) -> Result<TableScanSplitScanner> {
272        self.split
273            .create_scanner_with_transforms(self.runtime, &self.transforms)
274    }
275}
276
277/// Typed iterator over every row assigned to one table scan split.
278pub struct TableScanSplitScanner {
279    inner: ScanSplitScanner,
280    compiled: Arc<CompiledTable>,
281}
282
283impl Iterator for TableScanSplitScanner {
284    type Item = Result<Vec<Value>>;
285
286    fn next(&mut self) -> Option<Self::Item> {
287        self.inner.next().map(|row| {
288            let (_, key, columns) = row?;
289            decode_table_scan_row(&self.compiled, &key, &columns)
290        })
291    }
292}
293
294fn full_scan_split(shard: ShardSnapshotRef) -> ScanSplit {
295    ScanSplit {
296        shard,
297        start: None,
298        end: None,
299        start_bucket: None,
300        start_key_exclusive: None,
301        end_bucket: None,
302        end_key_inclusive: None,
303    }
304}
305
306fn validate_plan(
307    format: &str,
308    version: u32,
309    name: &str,
310    total_buckets: u32,
311    metadata: &TableMetadata,
312    source_volumes: &[VolumeDescriptor],
313) -> Result<()> {
314    if format != TABLE_SCAN_PLAN_FORMAT {
315        return Err(TableError::InvalidSchema(format!(
316            "unsupported table scan plan format: {format}"
317        )));
318    }
319    if version != TABLE_SCAN_PLAN_VERSION {
320        return Err(TableError::InvalidSchema(format!(
321            "unsupported table scan plan version: {version}"
322        )));
323    }
324    validate_name(name.to_string())?;
325    metadata.validate()?;
326    if !(1..=u16::MAX as u32 + 1).contains(&total_buckets) {
327        return Err(TableError::InvalidSchema(
328            "table scan plan total_buckets must be in range 1..=65536".to_string(),
329        ));
330    }
331    if !source_volumes
332        .iter()
333        .any(|volume| volume.supports(VolumeUsageKind::Meta))
334        || source_volumes.iter().any(|volume| {
335            volume.supports(VolumeUsageKind::PrimaryDataPriorityHigh)
336                || volume.supports(VolumeUsageKind::PrimaryDataPriorityMedium)
337                || volume.supports(VolumeUsageKind::PrimaryDataPriorityLow)
338                || volume.supports(VolumeUsageKind::Cache)
339                || volume.supports(VolumeUsageKind::Readonly)
340        })
341    {
342        return Err(TableError::InvalidSchema(
343            "table scan plan has invalid source volumes".to_string(),
344        ));
345    }
346    Ok(())
347}
348
349fn source_volumes(config: &Config) -> Vec<VolumeDescriptor> {
350    config
351        .volumes
352        .iter()
353        .filter_map(|source| {
354            let mut volume = source.clone();
355            volume.kinds = 0;
356            for kind in [VolumeUsageKind::Meta, VolumeUsageKind::Snapshot] {
357                if source.supports(kind) {
358                    volume.set_usage(kind);
359                }
360            }
361            (volume.kinds != 0).then_some(volume.without_credentials())
362        })
363        .collect()
364}
365
366fn runtime_read_config(mut runtime: Config) -> Config {
367    runtime.volumes = runtime
368        .volumes
369        .into_iter()
370        .filter_map(|source| {
371            let mut volume = source.clone();
372            volume.kinds = 0;
373            for kind in [
374                VolumeUsageKind::PrimaryDataPriorityHigh,
375                VolumeUsageKind::PrimaryDataPriorityMedium,
376                VolumeUsageKind::PrimaryDataPriorityLow,
377                VolumeUsageKind::Cache,
378                VolumeUsageKind::Readonly,
379            ] {
380                if source.supports(kind) {
381                    volume.set_usage(kind);
382                }
383            }
384            (volume.kinds != 0).then_some(volume)
385        })
386        .collect();
387    runtime
388}
cobble_table/scan.rs

cobble_table/
scan.rs
