Skip to content

refactor(scan): extract reusable scan functions #3093

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xanderbailey wants to merge 5 commits into
base: main
Choose a base branch
from
Open

refactor(scan): extract reusable scan functions #3093

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f8bc7ec
refactor(scan): extract scan planning utilities
xanderbailey Aug 28, 2026
222ebea
comments
xanderbailey Aug 28, 2026
cd35beb
Merge remote-tracking branch 'origin/main' into codex/extract-scan-co…
xanderbailey Aug 28, 2026
5853c0e
comments
xanderbailey Aug 31, 2026
eb33947
remove allocation
xanderbailey Aug 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion crates/iceberg/public-api.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2270,6 +2270,9 @@ pub fn iceberg::spec::NameMapping::eq(&self, other: &iceberg::spec::NameMapping)
impl core::fmt::Debug for iceberg::spec::NameMapping
pub fn iceberg::spec::NameMapping::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
impl core::marker::StructuralPartialEq for iceberg::spec::NameMapping
impl core::str::traits::FromStr for iceberg::spec::NameMapping
pub type iceberg::spec::NameMapping::Err = iceberg::Error
pub fn iceberg::spec::NameMapping::from_str(value: &str) -> iceberg::Result<Self>
impl serde_core::ser::Serialize for iceberg::spec::NameMapping
pub fn iceberg::spec::NameMapping::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer
impl<'de> serde_core::de::Deserialize<'de> for iceberg::spec::NameMapping
Expand Down Expand Up @@ -2730,6 +2733,7 @@ pub fn iceberg::spec::TableMetadata::sort_orders_iter(&self) -> impl core::iter:
pub fn iceberg::spec::TableMetadata::statistics_for_snapshot(&self, snapshot_id: i64) -> core::option::Option<&iceberg::spec::StatisticsFile>
pub fn iceberg::spec::TableMetadata::statistics_iter(&self) -> impl core::iter::traits::exact_size::ExactSizeIterator<Item = &iceberg::spec::StatisticsFile>
pub fn iceberg::spec::TableMetadata::table_properties(&self) -> iceberg::spec::TableProperties<'_>
pub fn iceberg::spec::TableMetadata::unified_partition_type(&self, schema: &iceberg::spec::Schema) -> iceberg::Result<iceberg::spec::StructType>
pub fn iceberg::spec::TableMetadata::uuid(&self) -> uuid::Uuid
pub async fn iceberg::spec::TableMetadata::write_to(&self, file_io: &iceberg::io::FileIO, metadata_location: &iceberg::MetadataLocation) -> iceberg::Result<()>
impl core::clone::Clone for iceberg::spec::TableMetadata
Expand Down Expand Up @@ -2814,6 +2818,7 @@ pub const iceberg::spec::TableProperties<'_>::PROPERTY_DATAFUSION_WRITE_FANOUT_E
pub const iceberg::spec::TableProperties<'_>::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED_DEFAULT: bool
pub const iceberg::spec::TableProperties<'_>::PROPERTY_DEFAULT_FILE_FORMAT: &'static str
pub const iceberg::spec::TableProperties<'_>::PROPERTY_DEFAULT_FILE_FORMAT_DEFAULT: &'static str
pub const iceberg::spec::TableProperties<'_>::PROPERTY_DEFAULT_NAME_MAPPING: &'static str
pub const iceberg::spec::TableProperties<'_>::PROPERTY_DEFAULT_PARTITION_SPEC: &'static str
pub const iceberg::spec::TableProperties<'_>::PROPERTY_DEFAULT_SORT_ORDER: &'static str
pub const iceberg::spec::TableProperties<'_>::PROPERTY_DELETE_DEFAULT_FILE_FORMAT: &'static str
Expand Down Expand Up @@ -2875,6 +2880,7 @@ pub fn iceberg::spec::TableProperties<'properties>::commit_max_retry_wait_ms(&se
pub fn iceberg::spec::TableProperties<'properties>::commit_min_retry_wait_ms(&self) -> iceberg::Result<u64>
pub fn iceberg::spec::TableProperties<'properties>::commit_num_retries(&self) -> iceberg::Result<usize>
pub fn iceberg::spec::TableProperties<'properties>::commit_total_retry_timeout_ms(&self) -> iceberg::Result<u64>
pub fn iceberg::spec::TableProperties<'properties>::default_name_mapping(&self) -> iceberg::Result<core::option::Option<iceberg::spec::NameMapping>>
pub fn iceberg::spec::TableProperties<'properties>::encryption_data_key_length(&self) -> iceberg::Result<usize>
pub fn iceberg::spec::TableProperties<'properties>::encryption_key_id(&self) -> iceberg::Result<core::option::Option<alloc::string::String>>
pub fn iceberg::spec::TableProperties<'properties>::gc_enabled(&self) -> iceberg::Result<bool>
Expand Down Expand Up @@ -3086,7 +3092,6 @@ pub fn iceberg::spec::ViewVersionLog::serialize<__S>(&self, __serializer: __S) -
impl<'de> serde_core::de::Deserialize<'de> for iceberg::spec::ViewVersionLog
pub fn iceberg::spec::ViewVersionLog::deserialize<__D>(__deserializer: __D) -> core::result::Result<Self, <__D as serde_core::de::Deserializer>::Error> where __D: serde_core::de::Deserializer<'de>
pub const iceberg::spec::DEFAULT_SCHEMA_ID: iceberg::spec::SchemaId
pub const iceberg::spec::DEFAULT_SCHEMA_NAME_MAPPING: &str
pub const iceberg::spec::INITIAL_ROW_ID: u64
pub const iceberg::spec::LIST_FIELD_NAME: &str
pub const iceberg::spec::MAIN_BRANCH: &str
Expand Down
173 changes: 87 additions & 86 deletions crates/iceberg/src/scan/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,8 @@ use crate::io::FileIO;
use crate::metadata_columns::{
RESERVED_FIELD_ID_PARTITION, get_metadata_field_id, is_metadata_column_name,
};
use crate::partitioning::compute_unified_partition_type;
use crate::runtime::Runtime;
use crate::spec::{DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, NameMapping, Schema, SnapshotRef};
use crate::spec::{DataContentType, Schema, SchemaRef, SnapshotRef, StructType};
use crate::table::Table;
use crate::util::available_parallelism;
use crate::{Error, ErrorKind, Result};
Expand All @@ -61,6 +60,73 @@ fn resolve_field_id(schema: &Schema, column_name: &str, case_sensitive: bool) ->
}
}

fn collect_scan_field_ids(
schema: &Schema,
column_names: Option<&[String]>,
case_sensitive: bool,
) -> Result<Vec<i32>> {
let Some(column_names) = column_names else {
return Ok(schema.as_struct().fields().iter().map(|f| f.id).collect());
};

column_names
.iter()
.map(|column_name| {
if is_metadata_column_name(column_name) {
return get_metadata_field_id(column_name);
}

let field_id = resolve_field_id(schema, column_name, case_sensitive).ok_or_else(|| {
Error::new(
ErrorKind::DataInvalid,
format!("Column {column_name} not found in table. Schema: {schema}"),
)
})?;

schema
.as_struct()
.field_by_id(field_id)
.ok_or_else(|| {
Error::new(
ErrorKind::FeatureUnsupported,
format!(
"Column {column_name} is not a direct child of schema but a nested field, which is not supported now. Schema: {schema}"
),
)
})?;

Ok(field_id)
})
.collect()
}

fn bind_scan_predicate(
schema: &SchemaRef,
predicate: Option<&Predicate>,
case_sensitive: bool,
) -> Result<Option<Arc<BoundPredicate>>> {
predicate
.map(|predicate| predicate.bind(schema.clone(), case_sensitive))
.transpose()
.map(|predicate| predicate.map(Arc::new))
}

fn projected_partition_type(
table: &Table,
schema: &Schema,
field_ids: &[i32],
) -> Result<Option<Arc<StructType>>> {
if !field_ids.contains(&RESERVED_FIELD_ID_PARTITION) {
return Ok(None);
}

table
.metadata()
.unified_partition_type(schema)
.map(Arc::new)
.map(Some)
}

/// Builder to create table scan.
pub struct TableScanBuilder<'a> {
table: &'a Table,
Expand Down Expand Up @@ -233,93 +299,25 @@ impl<'a> TableScanBuilder<'a> {
};

let schema = snapshot.schema(self.table.metadata())?;

let mut field_ids = vec![];
let column_names = self.column_names.clone().unwrap_or_else(|| {
schema
.as_struct()
.fields()
.iter()
.map(|f| f.name.clone())
.collect()
});

for column_name in column_names.iter() {
// Handle metadata columns (like "_file")
if is_metadata_column_name(column_name) {
field_ids.push(get_metadata_field_id(column_name)?);
continue;
}

let field_id =
resolve_field_id(&schema, column_name, self.case_sensitive).ok_or_else(|| {
Error::new(
ErrorKind::DataInvalid,
format!("Column {column_name} not found in table. Schema: {schema}"),
)
})?;

schema
.as_struct()
.field_by_id(field_id)
.ok_or_else(|| {
Error::new(
ErrorKind::FeatureUnsupported,
format!(
"Column {column_name} is not a direct child of schema but a nested field, which is not supported now. Schema: {schema}"
),
)
})?;

field_ids.push(field_id);
}

let snapshot_bound_predicate = if let Some(ref predicates) = self.filter {
Some(predicates.bind(schema.clone(), self.case_sensitive)?)
} else {
None
};

let field_ids =
collect_scan_field_ids(&schema, self.column_names.as_deref(), self.case_sensitive)?;
let snapshot_bound_predicate =
bind_scan_predicate(&schema, self.filter.as_ref(), self.case_sensitive)?;
let name_mapping = self
.table
.metadata()
.properties()
.get(DEFAULT_SCHEMA_NAME_MAPPING)
.map(|raw| {
serde_json::from_str::<NameMapping>(raw).map_err(|e| {
Error::new(
ErrorKind::DataInvalid,
format!(
"Failed to parse table property {DEFAULT_SCHEMA_NAME_MAPPING} as a NameMapping"
),
)
.with_source(e)
})
})
.transpose()?
.table_properties()
.default_name_mapping()?
.map(Arc::new);

// Compute unified partition type if _partition is projected
let unified_partition_type = if field_ids.contains(&RESERVED_FIELD_ID_PARTITION) {
let partition_type = compute_unified_partition_type(
self.table
.metadata()
.partition_specs_iter()
.map(|s| s.as_ref()),
&schema,
)?;
Some(Arc::new(partition_type))
} else {
None
};
let unified_partition_type = projected_partition_type(self.table, &schema, &field_ids)?;

let plan_context = PlanContext {
snapshot,
table_metadata: self.table.metadata_ref(),
snapshot_schema: schema,
case_sensitive: self.case_sensitive,
predicate: self.filter.map(Arc::new),
snapshot_bound_predicate: snapshot_bound_predicate.map(Arc::new),
snapshot_bound_predicate,
object_cache: self.table.object_cache(),
field_ids: Arc::new(field_ids),
name_mapping,
Expand Down Expand Up @@ -664,11 +662,10 @@ pub mod tests {
};
use crate::scan::FileScanTask;
use crate::spec::{
DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, DataFileBuilder, DataFileFormat, Datum,
FormatVersion, Literal, MAIN_BRANCH, ManifestEntry, ManifestListWriter, ManifestStatus,
ManifestWriterBuilder, NestedField, Operation, PartitionSpec, PrimitiveType, Schema,
Snapshot, Struct, StructType, Summary, TableMetadata, TableMetadataBuilder, Type,
UnboundPartitionSpec,
DataContentType, DataFileBuilder, DataFileFormat, Datum, FormatVersion, Literal,
MAIN_BRANCH, ManifestEntry, ManifestListWriter, ManifestStatus, ManifestWriterBuilder,
NestedField, Operation, PartitionSpec, PrimitiveType, Schema, Snapshot, Struct, StructType,
Summary, TableMetadata, TableMetadataBuilder, TableProperties, Type, UnboundPartitionSpec,
};
use crate::table::Table;
use crate::test_utils::test_runtime;
Expand Down Expand Up @@ -1896,7 +1893,8 @@ pub mod tests {
#[test]
fn test_table_scan_with_name_mapping_property() {
let mapping_json = r#"[{"field-id":1,"names":["id","record_id"]}]"#;
let table = table_with_property(DEFAULT_SCHEMA_NAME_MAPPING, mapping_json);
let table =
table_with_property(TableProperties::PROPERTY_DEFAULT_NAME_MAPPING, mapping_json);

let table_scan = table.scan().build().unwrap();
let mapping = table_scan
Expand All @@ -1917,7 +1915,10 @@ pub mod tests {

#[test]
fn test_table_scan_with_malformed_name_mapping_property() {
let table = table_with_property(DEFAULT_SCHEMA_NAME_MAPPING, "{ not valid json");
let table = table_with_property(
TableProperties::PROPERTY_DEFAULT_NAME_MAPPING,
"{ not valid json",
);

let err = table
.scan()
Expand All @@ -1934,7 +1935,7 @@ pub mod tests {
let mapping_json = r#"[{"field-id":1,"names":["id","record_id"]}]"#;
let mut metadata = fixture.table.metadata().clone();
metadata.properties.insert(
DEFAULT_SCHEMA_NAME_MAPPING.to_string(),
TableProperties::PROPERTY_DEFAULT_NAME_MAPPING.to_string(),
mapping_json.to_string(),
);
let table = Table::builder()
Expand Down
19 changes: 17 additions & 2 deletions crates/iceberg/src/spec/name_mapping/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,13 @@

//! Iceberg name mapping.

use std::str::FromStr;
use std::sync::Arc;

use serde::{Deserialize, Serialize};
use serde_with::{DefaultOnNull, serde_as};

/// Property name for name mapping.
pub const DEFAULT_SCHEMA_NAME_MAPPING: &str = "schema.name-mapping.default";
use crate::{Error, ErrorKind, Result};

/// Iceberg fallback field name to ID mapping.
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
Expand All @@ -44,6 +44,21 @@ impl NameMapping {
}
}

impl FromStr for NameMapping {
type Err = Error;

/// Parses a [`NameMapping`] from its JSON representation.
fn from_str(value: &str) -> Result<Self> {
serde_json::from_str(value).map_err(|error| {
Error::new(
ErrorKind::DataInvalid,
"Failed to parse value as a NameMapping",
)
.with_source(error)
})
}
}

/// Maps field names to IDs.
#[serde_as]
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
Expand Down
Loading
Loading