comphead commented on code in PR #3540:
URL: https://github.com/apache/iceberg-rust/pull/3540#discussion_r4403563144
##########
crates/iceberg/src/arrow/reader/predicate_visitor.rs:
##########
@@ -655,7 +653,11 @@ fn try_cast_literal(
return Ok(Arc::clone(literal));
}
- let literal_array = cast(literal_array, column_type)?;
+ let options = CastOptions {
+ safe: false,
+ ..Default::default()
+ };
+ let literal_array = cast_with_options(literal_array, column_type,
&options)?;
Review Comment:
Added a one-line comment above the options.
##########
crates/iceberg/src/arrow/reader/row_filter.rs:
##########
@@ -460,6 +478,89 @@ mod tests {
}
}
+ /// Writes `column` as field 1, reads it back as `iceberg_type` under each
predicate, and
+ /// checks the values that survive (`None` when no row does).
+ async fn assert_promoted_reads(
+ column: ArrayRef,
+ iceberg_type: PrimitiveType,
+ cases: Vec<(Predicate, Option<ArrayRef>)>,
+ ) {
+ let tmp = TempDir::new().unwrap();
+ let path = format!("{}/column.parquet", tmp.path().to_str().unwrap());
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![field_with_id(
+ "v",
+ column.data_type().clone(),
+ 1,
+ )]));
+ let batch = RecordBatch::try_new(arrow_schema.clone(),
vec![column]).unwrap();
+ write_row_groups(&path, arrow_schema, vec![batch], false);
+
+ let schema = value_schema(Type::Primitive(iceberg_type));
+ let file_size = std::fs::metadata(&path).unwrap().len();
+ let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(),
Runtime::current()).build();
+ for (predicate, expected) in cases {
+ let task = FileScanTask::builder()
+ .with_file_size_in_bytes(file_size)
+ .with_start(0)
+ .with_length(0)
+ .with_data_file_path(path.clone())
+ .with_data_file_format(DataFileFormat::Parquet)
+ .with_schema(schema.clone())
+ .with_project_field_ids(vec![1])
+ .with_predicate(Some(predicate.bind(schema.clone(),
true).unwrap()))
+ .with_case_sensitive(true)
+ .build()
Review Comment:
The helper now reuses `read_once`, so this code is gone.
##########
crates/iceberg/src/arrow/reader/row_filter.rs:
##########
@@ -62,8 +70,18 @@ impl ArrowReader {
// After collecting required leaf column indices used in the predicate,
// creates the projection mask for the Arrow predicates.
let projection_mask = ProjectionMask::leaves(parquet_schema,
column_indices.clone());
- let predicate_func = visit(&mut converter, predicates)?;
- let arrow_predicate = ArrowPredicateFn::new(projection_mask,
predicate_func);
+ let mut predicate_func = visit(&mut converter, predicates)?;
+ // Predicates are bound to the table schema, so evaluate them on the
filter columns
+ // converted to it, as the reader's output is (e.g. a file written
before an `int` ->
+ // `long` promotion).
+ let mut transformer =
+ RecordBatchTransformerBuilder::new(snapshot_schema,
&field_ids).build();
Review Comment:
Thanks, agreed on both. The description now names the per-batch cast,
including the `timestamptz` relabel, and the per-file plan. I'd like to scope
`generate_batch_transform` to the projected fields in a follow-up, since it
changes the output transformer too and several open PRs touch that file.
##########
crates/iceberg/src/arrow/reader/row_filter.rs:
##########
@@ -460,6 +478,89 @@ mod tests {
}
}
+ /// Writes `column` as field 1, reads it back as `iceberg_type` under each
predicate, and
+ /// checks the values that survive (`None` when no row does).
+ async fn assert_promoted_reads(
+ column: ArrayRef,
+ iceberg_type: PrimitiveType,
+ cases: Vec<(Predicate, Option<ArrayRef>)>,
+ ) {
+ let tmp = TempDir::new().unwrap();
+ let path = format!("{}/column.parquet", tmp.path().to_str().unwrap());
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![field_with_id(
+ "v",
+ column.data_type().clone(),
+ 1,
+ )]));
+ let batch = RecordBatch::try_new(arrow_schema.clone(),
vec![column]).unwrap();
+ write_row_groups(&path, arrow_schema, vec![batch], false);
Review Comment:
Added. `assert_promoted_reads` now writes each file with and without a field
ID. `read_once` takes no name mapping, so the second file goes through the
positional fallback, which feeds the same `with_schema` hint.
##########
crates/iceberg/src/arrow/reader/row_filter.rs:
##########
@@ -460,6 +478,89 @@ mod tests {
}
}
+ /// Writes `column` as field 1, reads it back as `iceberg_type` under each
predicate, and
+ /// checks the values that survive (`None` when no row does).
+ async fn assert_promoted_reads(
+ column: ArrayRef,
+ iceberg_type: PrimitiveType,
+ cases: Vec<(Predicate, Option<ArrayRef>)>,
+ ) {
+ let tmp = TempDir::new().unwrap();
+ let path = format!("{}/column.parquet", tmp.path().to_str().unwrap());
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![field_with_id(
+ "v",
+ column.data_type().clone(),
+ 1,
+ )]));
+ let batch = RecordBatch::try_new(arrow_schema.clone(),
vec![column]).unwrap();
+ write_row_groups(&path, arrow_schema, vec![batch], false);
+
+ let schema = value_schema(Type::Primitive(iceberg_type));
+ let file_size = std::fs::metadata(&path).unwrap().len();
+ let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(),
Runtime::current()).build();
+ for (predicate, expected) in cases {
+ let task = FileScanTask::builder()
+ .with_file_size_in_bytes(file_size)
+ .with_start(0)
+ .with_length(0)
+ .with_data_file_path(path.clone())
+ .with_data_file_format(DataFileFormat::Parquet)
+ .with_schema(schema.clone())
+ .with_project_field_ids(vec![1])
+ .with_predicate(Some(predicate.bind(schema.clone(),
true).unwrap()))
+ .with_case_sensitive(true)
+ .build()
+ .unwrap();
+ let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as
FileScanTaskStream;
+ let batches: Vec<RecordBatch> = reader
+ .clone()
+ .read(tasks)
+ .unwrap()
+ .stream()
+ .try_collect()
+ .await
+ .unwrap();
+ let actual = collapse(&batches).map(|batch|
Arc::clone(batch.column(0)));
+ assert_eq!(actual, expected, "{predicate}");
+ }
Review Comment:
Done, thanks.
##########
crates/iceberg/src/arrow/reader/row_filter.rs:
##########
@@ -460,6 +478,89 @@ mod tests {
}
}
+ /// Writes `column` as field 1, reads it back as `iceberg_type` under each
predicate, and
+ /// checks the values that survive (`None` when no row does).
+ async fn assert_promoted_reads(
+ column: ArrayRef,
+ iceberg_type: PrimitiveType,
+ cases: Vec<(Predicate, Option<ArrayRef>)>,
+ ) {
+ let tmp = TempDir::new().unwrap();
+ let path = format!("{}/column.parquet", tmp.path().to_str().unwrap());
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![field_with_id(
+ "v",
+ column.data_type().clone(),
+ 1,
+ )]));
+ let batch = RecordBatch::try_new(arrow_schema.clone(),
vec![column]).unwrap();
+ write_row_groups(&path, arrow_schema, vec![batch], false);
+
+ let schema = value_schema(Type::Primitive(iceberg_type));
+ let file_size = std::fs::metadata(&path).unwrap().len();
+ let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(),
Runtime::current()).build();
+ for (predicate, expected) in cases {
+ let task = FileScanTask::builder()
+ .with_file_size_in_bytes(file_size)
+ .with_start(0)
+ .with_length(0)
+ .with_data_file_path(path.clone())
+ .with_data_file_format(DataFileFormat::Parquet)
+ .with_schema(schema.clone())
+ .with_project_field_ids(vec![1])
+ .with_predicate(Some(predicate.bind(schema.clone(),
true).unwrap()))
+ .with_case_sensitive(true)
+ .build()
+ .unwrap();
+ let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as
FileScanTaskStream;
+ let batches: Vec<RecordBatch> = reader
+ .clone()
+ .read(tasks)
+ .unwrap()
+ .stream()
+ .try_collect()
+ .await
+ .unwrap();
+ let actual = collapse(&batches).map(|batch|
Arc::clone(batch.column(0)));
+ assert_eq!(actual, expected, "{predicate}");
+ }
+ }
+
+ /// A file written before a type promotion keeps the old physical type, so
the row
+ /// filter must compare its values promoted to the table type.
+ #[tokio::test]
+ async fn test_predicate_on_promoted_column() {
+ let longs = |values: Vec<i64>| Some(Arc::new(Int64Array::from(values))
as ArrayRef);
+ let all = longs(vec![1, 2]);
+ let above = || Datum::long(i64::from(i32::MAX) + 1);
+ let below = || Datum::long(i64::from(i32::MIN) - 1);
+ let v = || Reference::new("v");
+ assert_promoted_reads(
+ Arc::new(Int32Array::from(vec![1, 2])),
+ PrimitiveType::Long,
+ vec![
+ (v().less_than(above()), all.clone()),
+ (v().less_than_or_equal_to(above()), all.clone()),
+ (v().greater_than_or_equal_to(below()), all.clone()),
+ (v().not_equal_to(above()), all),
+ (v().is_in([Datum::long(1), above()]), longs(vec![1])),
+ (v().is_not_in([Datum::long(1), above()]), longs(vec![2])),
+ ],
+ )
+ .await;
+
+ // 0.1f32 promotes to a double that is greater than 0.1, not equal to
it.
+ let promoted = Arc::new(Float64Array::from(vec![f64::from(0.1f32)]))
as ArrayRef;
+ assert_promoted_reads(
+ Arc::new(Float32Array::from(vec![0.1])),
+ PrimitiveType::Double,
+ vec![
+ (v().greater_than(Datum::double(0.1)), Some(promoted)),
+ (v().equal_to(Datum::double(0.1)), None),
+ ],
+ )
+ .await;
+ }
Review Comment:
Added, thanks. I also checked `timestamp_ns` locally, and the description
lists both.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
