You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "jecsand838 (via GitHub)" <gi...@apache.org> on 2026/01/13 16:33:16 UTC
[PR] Add Avro Reader projection API [arrow-rs]
jecsand838 opened a new pull request, #9162:
URL: https://github.com/apache/arrow-rs/pull/9162
# Which issue does this PR close?
- Closes #8923.
# Rationale for this change
The `arrow-avro` crate's `ReaderBuilder` previously lacked the ability to project (select) specific columns when reading Avro files. This is a common feature in other Arrow readers (like `arrow-csv` and `arrow-ipc`) that enables users to read only the columns they need, improving performance and reducing memory usage.
# What changes are included in this PR?
- Added a `with_projection(projection: Vec<usize>)` method to `ReaderBuilder` that accepts zero-based column indices
- Implemented `AvroSchema::project()` method to create a projected Avro schema with only the selected fields
- The projection supports:
- Selecting a subset of fields
- Reordering fields
- Preserving all record and field metadata (namespace, doc, defaults, aliases, etc.)
- Preserving nested/complex types (records, arrays, maps, unions)
- Added validation for out-of-bounds indices and duplicate indices
# Are these changes tested?
Yes, comprehensive tests have been added:
- Unit tests for `AvroSchema::project()` covering:
- Empty projections
- Single and multiple field selection
- Field reordering
- Metadata preservation (record-level and field-level)
- Nested records and complex types (arrays, maps, unions)
- Error cases (invalid JSON, non-record schemas, out-of-bounds indices, duplicate indices)
- Integration tests in the reader module for end-to-end projection with OCF files
# Are there any user-facing changes?
Yes, this adds a new public API method:
```rust
impl ReaderBuilder {
/// Set a projection of columns to read (zero-based column indices).
pub fn with_projection(self, projection: Vec<usize>) -> Self
}
```
This is consistent with the projection API in `arrow-csv::ReaderBuilder` and `arrow-ipc::FileReaderBuilder`. There are no breaking changes to existing APIs.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Avro Reader projection API [arrow-rs]
Posted by "getChan (via GitHub)" <gi...@apache.org>.
getChan commented on PR #9162:
URL: https://github.com/apache/arrow-rs/pull/9162#issuecomment-3749977475
Thanks for the great feature. I’ll try it out in DataFusion this week, and I’ll leave a comment if I have any further suggestions.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Avro Reader projection API [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #9162:
URL: https://github.com/apache/arrow-rs/pull/9162#discussion_r2696335018
##########
arrow-avro/src/reader/mod.rs:
##########
@@ -1017,8 +1031,33 @@ impl ReaderBuilder {
.ok_or_else(|| {
ArrowError::ParseError("No Avro schema present in file header".into())
})?;
+ let projected_reader_schema = self
+ .projection
+ .as_deref()
+ .map(|projection| {
+ let base_schema = if let Some(reader_schema) = reader_schema {
Review Comment:
I'll do a follow-up PR with comments covering this.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Avro Reader projection API [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #9162:
URL: https://github.com/apache/arrow-rs/pull/9162#discussion_r2696333971
##########
arrow-avro/src/schema.rs:
##########
@@ -3137,4 +3197,546 @@ mod tests {
assert_eq!(union_arr2[1], Value::String("int".into()));
assert_eq!(union_arr2[2], Value::String("string".into()));
}
+
+ #[test]
+ fn test_project_empty_projection() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert!(
+ fields.is_empty(),
+ "Empty projection should yield empty fields"
+ );
+ }
+
+ #[test]
+ fn test_project_single_field() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[1]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 1);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("b"));
+ }
+
+ #[test]
+ fn test_project_multiple_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"},
+ {"name": "d", "type": "boolean"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0, 2, 3]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 3);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("a"));
+ assert_eq!(fields[1].get("name").and_then(|n| n.as_str()), Some("c"));
+ assert_eq!(fields[2].get("name").and_then(|n| n.as_str()), Some("d"));
+ }
+
+ #[test]
+ fn test_project_all_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0, 1]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 2);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("a"));
+ assert_eq!(fields[1].get("name").and_then(|n| n.as_str()), Some("b"));
+ }
+
+ #[test]
+ fn test_project_reorder_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ // Project in reverse order
Review Comment:
Ah good catch, I'll clean the comment up.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Avro Reader projection API [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on PR #9162:
URL: https://github.com/apache/arrow-rs/pull/9162#issuecomment-3757376655
@getChan @nathaniel-d-ef Let me know if we need any further enhancements and I'll be sure to get those up prior to the this month's major release.
> I a merging this in to make it easier to test downstream -- let's make any additional improvements as follow on PRs
100%.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Avro Reader projection API [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9162:
URL: https://github.com/apache/arrow-rs/pull/9162#issuecomment-3756009944
I a merging this in to make it easier to test downstream -- let's make any additional improvements as follow on PRs
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Avro Reader projection API [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9162:
URL: https://github.com/apache/arrow-rs/pull/9162#discussion_r2692148424
##########
arrow-avro/src/reader/mod.rs:
##########
@@ -1119,6 +1178,68 @@ impl ReaderBuilder {
self
}
+ /// Sets an explicit top-level field projection by index.
+ ///
+ /// The provided `projection` is a list of indices into the **top-level record** fields.
+ /// The output schema will contain only these fields, in the specified order.
+ ///
+ /// Internally, this is implemented by pruning the effective Avro *reader schema*:
+ ///
+ /// * If a reader schema is provided via `Self::with_reader_schema`, that schema is pruned.
+ /// * Otherwise, a reader schema is derived from the writer schema and then pruned.
+ /// * For streaming `Decoder` with multiple writer schemas and no reader schema, a projected
+ /// reader schema is derived **per writer schema** in the `SchemaStore`.
+ ///
+ /// # Example
+ ///
+ /// Read only specific columns from an Avro OCF file:
+ ///
+ /// ```
+ /// use std::io::Cursor;
+ /// use std::sync::Arc;
+ /// use arrow_array::{ArrayRef, Int32Array, StringArray, Float64Array, RecordBatch};
+ /// use arrow_schema::{DataType, Field, Schema};
+ /// use arrow_avro::writer::AvroWriter;
+ /// use arrow_avro::reader::ReaderBuilder;
+ ///
+ /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+ /// // Original schema has three fields: id, name, value
+ /// let schema = Schema::new(vec![
+ /// Field::new("id", DataType::Int32, false),
+ /// Field::new("name", DataType::Utf8, false),
+ /// Field::new("value", DataType::Float64, false),
+ /// ]);
+ /// let batch = RecordBatch::try_new(
+ /// Arc::new(schema.clone()),
+ /// vec![
+ /// Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef,
+ /// Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef,
+ /// Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as ArrayRef,
+ /// ],
+ /// )?;
+ ///
+ /// // Write Avro OCF
+ /// let mut writer = AvroWriter::new(Vec::new(), schema)?;
+ /// writer.write(&batch)?;
+ /// writer.finish()?;
+ /// let bytes = writer.into_inner();
+ ///
+ /// // Read only fields at indices 2 and 0 (value, id) — in that order
+ /// let mut reader = ReaderBuilder::new()
+ /// .with_projection(vec![2, 0])
+ /// .build(Cursor::new(bytes))?;
+ ///
+ /// let out = reader.next().unwrap()?;
+ /// assert_eq!(out.num_columns(), 2);
+ /// assert_eq!(out.schema().field(0).name(), "value");
+ /// assert_eq!(out.schema().field(1).name(), "id");
+ /// # Ok(()) }
+ /// ```
Review Comment:
😍
##########
arrow-avro/src/reader/mod.rs:
##########
@@ -1017,8 +1031,33 @@ impl ReaderBuilder {
.ok_or_else(|| {
ArrowError::ParseError("No Avro schema present in file header".into())
})?;
+ let projected_reader_schema = self
+ .projection
+ .as_deref()
+ .map(|projection| {
+ let base_schema = if let Some(reader_schema) = reader_schema {
Review Comment:
It might be worth pointing out in the comments that the projection is relative the to reader schema if set, otherwise it is relative to whatever is in the file
##########
arrow-avro/src/schema.rs:
##########
@@ -3137,4 +3197,546 @@ mod tests {
assert_eq!(union_arr2[1], Value::String("int".into()));
assert_eq!(union_arr2[2], Value::String("string".into()));
}
+
+ #[test]
+ fn test_project_empty_projection() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert!(
+ fields.is_empty(),
+ "Empty projection should yield empty fields"
+ );
+ }
+
+ #[test]
+ fn test_project_single_field() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[1]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 1);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("b"));
+ }
+
+ #[test]
+ fn test_project_multiple_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"},
+ {"name": "d", "type": "boolean"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0, 2, 3]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 3);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("a"));
+ assert_eq!(fields[1].get("name").and_then(|n| n.as_str()), Some("c"));
+ assert_eq!(fields[2].get("name").and_then(|n| n.as_str()), Some("d"));
+ }
+
+ #[test]
+ fn test_project_all_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0, 1]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 2);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("a"));
+ assert_eq!(fields[1].get("name").and_then(|n| n.as_str()), Some("b"));
+ }
+
+ #[test]
+ fn test_project_reorder_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ // Project in reverse order
Review Comment:
nit: it isn't really in reverse order
##########
arrow-avro/src/schema.rs:
##########
@@ -3137,4 +3197,546 @@ mod tests {
assert_eq!(union_arr2[1], Value::String("int".into()));
assert_eq!(union_arr2[2], Value::String("string".into()));
}
+
+ #[test]
+ fn test_project_empty_projection() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert!(
+ fields.is_empty(),
+ "Empty projection should yield empty fields"
+ );
+ }
+
+ #[test]
+ fn test_project_single_field() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[1]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 1);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("b"));
+ }
+
+ #[test]
+ fn test_project_multiple_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"},
+ {"name": "d", "type": "boolean"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0, 2, 3]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 3);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("a"));
+ assert_eq!(fields[1].get("name").and_then(|n| n.as_str()), Some("c"));
+ assert_eq!(fields[2].get("name").and_then(|n| n.as_str()), Some("d"));
+ }
+
+ #[test]
+ fn test_project_all_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0, 1]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 2);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("a"));
+ assert_eq!(fields[1].get("name").and_then(|n| n.as_str()), Some("b"));
+ }
+
+ #[test]
+ fn test_project_reorder_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ // Project in reverse order
+ let projected = schema.project(&[2, 0, 1]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 3);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("c"));
+ assert_eq!(fields[1].get("name").and_then(|n| n.as_str()), Some("a"));
+ assert_eq!(fields[2].get("name").and_then(|n| n.as_str()), Some("b"));
+ }
+
+ #[test]
+ fn test_project_preserves_record_metadata() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "MyRecord",
+ "namespace": "com.example",
+ "doc": "A test record",
+ "aliases": ["OldRecord"],
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ assert_eq!(v.get("name").and_then(|n| n.as_str()), Some("MyRecord"));
+ assert_eq!(
+ v.get("namespace").and_then(|n| n.as_str()),
+ Some("com.example")
+ );
+ assert_eq!(v.get("doc").and_then(|n| n.as_str()), Some("A test record"));
+ assert!(v.get("aliases").is_some());
+ }
+
+ #[test]
+ fn test_project_preserves_field_metadata() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int", "doc": "Field A", "default": 0},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(
+ fields[0].get("doc").and_then(|d| d.as_str()),
+ Some("Field A")
+ );
+ assert_eq!(fields[0].get("default").and_then(|d| d.as_i64()), Some(0));
+ }
+
+ #[test]
+ fn test_project_with_nested_record() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Outer",
+ "fields": [
+ {"name": "id", "type": "int"},
+ {"name": "inner", "type": {
+ "type": "record",
+ "name": "Inner",
+ "fields": [
+ {"name": "x", "type": "int"},
+ {"name": "y", "type": "string"}
+ ]
+ }},
+ {"name": "value", "type": "double"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[1]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 1);
+ assert_eq!(
+ fields[0].get("name").and_then(|n| n.as_str()),
+ Some("inner")
+ );
+ // Verify nested record structure is preserved
+ let inner_type = fields[0].get("type").unwrap();
+ assert_eq!(
+ inner_type.get("type").and_then(|t| t.as_str()),
+ Some("record")
+ );
+ assert_eq!(
+ inner_type.get("name").and_then(|n| n.as_str()),
+ Some("Inner")
+ );
+ }
+
+ #[test]
+ fn test_project_with_complex_field_types() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "arr", "type": {"type": "array", "items": "int"}},
+ {"name": "map", "type": {"type": "map", "values": "string"}},
+ {"name": "union", "type": ["null", "int"]}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0, 2]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 2);
+ // Verify array type is preserved
+ let arr_type = fields[0].get("type").unwrap();
+ assert_eq!(arr_type.get("type").and_then(|t| t.as_str()), Some("array"));
+ // Verify union type is preserved
+ let union_type = fields[1].get("type").unwrap();
+ assert!(union_type.is_array());
+ }
+
+ #[test]
+ fn test_project_error_invalid_json() {
+ let schema = AvroSchema::new("not valid json".to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("Invalid Avro schema JSON"),
+ "Expected parse error, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_not_object() {
+ // Primitive type schema (not a JSON object)
+ let schema = AvroSchema::new(r#""string""#.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("must be a JSON object"),
+ "Expected object error, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_array_schema() {
+ // Array (list) is a valid JSON but not a record
+ let schema = AvroSchema::new(r#"["null", "int"]"#.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("must be a JSON object"),
+ "Expected object error for array schema, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_type_not_record() {
+ let schema_json = r#"{
+ "type": "enum",
+ "name": "Color",
+ "symbols": ["RED", "GREEN", "BLUE"]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("must be an Avro record") && msg.contains("'enum'"),
+ "Expected type mismatch error, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_type_array() {
+ let schema_json = r#"{
+ "type": "array",
+ "items": "int"
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("must be an Avro record") && msg.contains("'array'"),
+ "Expected type mismatch error for array type, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_type_fixed() {
+ let schema_json = r#"{
+ "type": "fixed",
+ "name": "MD5",
+ "size": 16
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("must be an Avro record") && msg.contains("'fixed'"),
+ "Expected type mismatch error for fixed type, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_type_map() {
+ let schema_json = r#"{
+ "type": "map",
+ "values": "string"
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("must be an Avro record") && msg.contains("'map'"),
+ "Expected type mismatch error for map type, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_missing_type_field() {
+ let schema_json = r#"{
+ "name": "Test",
+ "fields": [{"name": "a", "type": "int"}]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("missing required 'type' field"),
+ "Expected missing type error, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_missing_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test"
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("missing required 'fields'"),
+ "Expected missing fields error, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_fields_not_array() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": "not an array"
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("'fields' must be an array"),
+ "Expected fields array error, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_index_out_of_bounds() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[5]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("out of bounds") && msg.contains("5") && msg.contains("2"),
+ "Expected out of bounds error, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_index_out_of_bounds_edge() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ // Index 1 is just out of bounds for a 1-element array
+ let err = schema.project(&[1]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("out of bounds") && msg.contains("1"),
+ "Expected out of bounds error for edge case, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_duplicate_index() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"},
+ {"name": "c", "type": "long"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[0, 1, 0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("Duplicate projection index") && msg.contains("0"),
+ "Expected duplicate index error, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_error_duplicate_index_consecutive() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "a", "type": "int"},
+ {"name": "b", "type": "string"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[1, 1]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("Duplicate projection index") && msg.contains("1"),
+ "Expected duplicate index error for consecutive duplicates, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_with_empty_fields() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "EmptyRecord",
+ "fields": []
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ // Projecting empty from empty should succeed
+ let projected = schema.project(&[]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert!(fields.is_empty());
+ }
+
+ #[test]
+ fn test_project_empty_fields_index_out_of_bounds() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "EmptyRecord",
+ "fields": []
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let err = schema.project(&[0]).unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("out of bounds") && msg.contains("0 fields"),
+ "Expected out of bounds error for empty record, got: {msg}"
+ );
+ }
+
+ #[test]
+ fn test_project_result_is_valid_avro_schema() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "namespace": "com.example",
+ "fields": [
+ {"name": "id", "type": "long"},
+ {"name": "name", "type": "string"},
+ {"name": "active", "type": "boolean"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ let projected = schema.project(&[0, 2]).unwrap();
+ // Verify the projected schema can be parsed as a valid Avro schema
+ let parsed = projected.schema();
+ assert!(parsed.is_ok(), "Projected schema should be valid Avro");
+ match parsed.unwrap() {
+ Schema::Complex(ComplexType::Record(r)) => {
+ assert_eq!(r.name, "Test");
+ assert_eq!(r.namespace, Some("com.example"));
+ assert_eq!(r.fields.len(), 2);
+ assert_eq!(r.fields[0].name, "id");
+ assert_eq!(r.fields[1].name, "active");
+ }
+ _ => panic!("Expected Record schema"),
+ }
+ }
+
+ #[test]
+ fn test_project_non_contiguous_indices() {
+ let schema_json = r#"{
+ "type": "record",
+ "name": "Test",
+ "fields": [
+ {"name": "f0", "type": "int"},
+ {"name": "f1", "type": "int"},
+ {"name": "f2", "type": "int"},
+ {"name": "f3", "type": "int"},
+ {"name": "f4", "type": "int"}
+ ]
+ }"#;
+ let schema = AvroSchema::new(schema_json.to_string());
+ // Select every other field
+ let projected = schema.project(&[0, 2, 4]).unwrap();
+ let v: Value = serde_json::from_str(&projected.json_string).unwrap();
+ let fields = v.get("fields").and_then(|f| f.as_array()).unwrap();
+ assert_eq!(fields.len(), 3);
+ assert_eq!(fields[0].get("name").and_then(|n| n.as_str()), Some("f0"));
+ assert_eq!(fields[1].get("name").and_then(|n| n.as_str()), Some("f2"));
+ assert_eq!(fields[2].get("name").and_then(|n| n.as_str()), Some("f4"));
+ }
+
+ #[test]
Review Comment:
this is quite a collection of tests
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Avro Reader projection API [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #9162:
URL: https://github.com/apache/arrow-rs/pull/9162
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org