From 67ef1a348e7ff7a57352b9d01326371278ccd4b4 Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 4 Nov 2025 12:10:31 +0100 Subject: [PATCH] feat(scanner): add Scanner::project_from_schema Set the scan projection from a Lance `Schema` instead of column expressions, exposing `ProjectionPlan::from_schema` through the scanner. The schema must be a subset of the dataset schema and, unlike the expression form, can select portions of nested structs while keeping their shape and field ids. Row-id and row-address inclusion carry over from the legacy flags and the configured blob handling is applied, as in `project_with_transform`. --- rust/lance/src/dataset/scanner.rs | 81 +++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 34c305e795f..5d226922c34 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1548,6 +1548,26 @@ impl Scanner { Ok(self) } + /// Set the projection from a schema. + /// + /// This projection will have no complex expressions, the schema must be a subset of the dataset schema. + /// + /// With this approach it is possible to refer to portions of nested fields. + /// + /// See [`ProjectionPlan::from_schema`] for more information. + pub fn project_from_schema(&mut self, projection: &Schema) -> Result<&mut Self> { + self.explicit_projection = true; + self.projection_plan = ProjectionPlan::from_schema(self.dataset.clone(), projection)?; + if self.legacy_with_row_id { + self.projection_plan.include_row_id(); + } + if self.legacy_with_row_addr { + self.projection_plan.include_row_addr(); + } + self.apply_blob_handling(); + Ok(self) + } + /// Should the filter run before the vector index is applied /// /// If true then the filter will be applied before the vector index. This @@ -9158,6 +9178,67 @@ mod test { assert_eq!(taken.schema(), part_schema); } + #[tokio::test] + async fn test_project_from_schema() { + let point_fields: Fields = vec![ + ArrowField::new("x", DataType::Float32, true), + ArrowField::new("y", DataType::Float32, true), + ] + .into(); + let metadata_fields: Fields = vec![ + ArrowField::new("location", DataType::Struct(point_fields), true), + ArrowField::new("age", DataType::Int32, true), + ] + .into(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("metadata", DataType::Struct(metadata_fields), true), + ArrowField::new("idx", DataType::Int32, true), + ])); + let data = lance_datagen::rand(&schema) + .into_ram_dataset(FragmentCount::from(7), FragmentRowCount::from(6)) + .await + .unwrap(); + + // 0 - metadata + // 2 - x + // 4 - age + // A partial struct that keeps `location` but drops `y` from it. + let projection = data.schema().project_by_ids(&[0, 2, 4], false); + + let mut scan = data.scan(); + scan.with_row_id() + .with_row_address() + .blob_handling(BlobHandling::AllBinary) + .project_from_schema(&projection) + .unwrap(); + // The blob handling configured before the projection must survive it, + // as it does for `project_with_transform`. + assert_eq!( + scan.projection_plan.physical_projection.blob_handling, + BlobHandling::AllBinary + ); + let batch = scan.try_into_batch().await.unwrap(); + + // Unlike the expression form, the output keeps the nested shape of the schema. + let part_point_fields = Fields::from(vec![ArrowField::new("x", DataType::Float32, true)]); + let part_metadata_fields = Fields::from(vec![ + ArrowField::new("location", DataType::Struct(part_point_fields), true), + ArrowField::new("age", DataType::Int32, true), + ]); + assert_eq!( + batch.schema().field_names(), + vec!["metadata", ROW_ID, ROW_ADDR] + ); + assert_eq!( + batch["metadata"].data_type(), + &DataType::Struct(part_metadata_fields) + ); + + let row_ids = batch[ROW_ID].as_primitive::().values().to_vec(); + let taken = data.take_rows(&row_ids, projection).await.unwrap(); + assert_eq!(&batch["metadata"], &taken["metadata"]); + } + #[rstest] #[tokio::test] async fn test_limit(