You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "rok (via GitHub)" <gi...@apache.org> on 2025/06/29 00:00:15 UTC
[PR] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
rok opened a new pull request, #7818:
URL: https://github.com/apache/arrow-rs/pull/7818
Closes #7359.
# Rationale for this change
This is to enable concurrent column writing with encryption downstream (e.g. with datafusion). See #7359 for more.
See https://github.com/apache/arrow-rs/pull/7111/files#r2015196618
# What changes are included in this PR?
* `ArrowRowGroupWriterFactory.create_row_group_writer` now passes a `FileEncryptor` to `ArrowColumnWriterFactory` that can later be used to write columns concurrently.
* Several methods and properties are now `pub`.
* Minor change to how encryption tests read test data.
# Are these changes tested?
Yes.
# Are there any user-facing changes?
`ArrowRowGroupWriterFactory`, `ArrowRowGroupWriter`, `ArrowRowGroupWriterFactory::new`, `ArrowRowGroupWriterFactory.create_row_group_writer `, .. are now `pub`. Shouldn't be breaking but perhaps needs review.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2201592759
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -755,7 +755,7 @@ impl ArrowColumnWriter {
}
/// Encodes [`RecordBatch`] to a parquet row group
-struct ArrowRowGroupWriter {
+pub struct ArrowRowGroupWriter {
Review Comment:
I do think avoiding exposing ArrowRowGroupWriter and the associated machinery that would be good
The rationale is that @XiangpengHao and @zhuqi-lucas and myself are likely to be reworking it in the next few releases and once it is part of the public API changing it becomes harder as it requires more coordination / backwards compatibility concerns
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2181047250
##########
parquet/tests/encryption/encryption.rs:
##########
@@ -1101,3 +1109,117 @@ fn read_and_roundtrip_to_encrypted_file(
// check re-written example data
verify_encryption_test_file_read(temp_file, decryption_properties);
}
+
+#[tokio::test]
+async fn test_multi_threaded_encrypted_writing() {
+ // Read example data and set up encryption/decryption properties
+ let testdata = arrow::util::test_util::parquet_test_data();
+ let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted");
+
+ let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into())
+ .with_column_key("double_field", b"1234567890123450".into())
+ .with_column_key("float_field", b"1234567890123451".into())
+ .build()
+ .unwrap();
+ let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into())
+ .with_column_key("double_field", b"1234567890123450".into())
+ .with_column_key("float_field", b"1234567890123451".into())
+ .build()
+ .unwrap();
+
+ let (record_batches, metadata) =
+ read_encrypted_file(&path, decryption_properties.clone()).unwrap();
+ let to_write: Vec<_> = record_batches
+ .iter()
+ .flat_map(|rb| rb.columns().to_vec())
+ .collect();
+ let schema = metadata.schema().clone();
+
+ let props = Arc::new(
+ WriterPropertiesBuilder::with_defaults()
+ .with_file_encryption_properties(file_encryption_properties)
+ .build(),
+ );
+ let parquet_schema = ArrowSchemaConverter::new()
+ .with_coerce_types(props.coerce_types())
+ .convert(&schema)
+ .unwrap();
+ let root_schema = parquet_schema.root_schema_ptr();
+
+ // Create a temporary file to write the encrypted data
+ let temp_file = tempfile::NamedTempFile::new().unwrap();
+ let mut file_writer =
+ SerializedFileWriter::new(&temp_file, root_schema.clone(), props.clone()).unwrap();
+
+ let arrow_row_group_writer_factory = ArrowRowGroupWriterFactory::new(
+ &file_writer,
+ parquet_schema,
+ schema.clone(),
+ props.clone(),
+ );
+ let arrow_row_group_writer = arrow_row_group_writer_factory
+ .create_row_group_writer(0)
+ .unwrap();
+
+ // Get column writers with encryptor from ArrowRowGroupWriter
+ let col_writers = arrow_row_group_writer.writers;
+ let num_columns = col_writers.len();
+
+ // Create a channel for each column writer to send ArrowLeafColumn data to
+ let mut col_writer_tasks = Vec::with_capacity(num_columns);
+ let mut col_array_channels = Vec::with_capacity(num_columns);
+ for mut writer in col_writers.into_iter() {
+ let (send_array, mut receive_array) = tokio::sync::mpsc::channel::<ArrowLeafColumn>(100);
+ col_array_channels.push(send_array);
+ let handle = tokio::spawn(async move {
+ while let Some(col) = receive_array.recv().await {
+ let _ = writer.write(&col);
+ }
+ writer.close().unwrap()
+ });
+ col_writer_tasks.push(handle);
+ }
+
+ // Send the ArrowLeafColumn data to the respective column writer channels
+ for (channel_idx, (array, field)) in to_write.iter().zip(schema.fields()).enumerate() {
+ for c in compute_leaves(field, array).into_iter().flatten() {
+ let _ = col_array_channels[channel_idx].send(c).await;
+ }
+ }
Review Comment:
Got it. Changed to:
```rust
let mut worker_iter = col_array_channels.iter_mut();
for (array, field) in to_write.iter().zip(schema.fields()) {
for leaves in compute_leaves(field, array).unwrap() {
worker_iter.next().unwrap().send(leaves).await.unwrap();
}
}
```
Do you think we should also add a struct array test case?
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2181041488
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -199,9 +199,14 @@ impl<W: Write + Send> ArrowWriter<W> {
let max_row_group_size = props.max_row_group_size();
let file_writer =
- SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::new(props))?;
+ SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::new(props.clone()))?;
Review Comment:
Done.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "adamreeve (via GitHub)" <gi...@apache.org>.
adamreeve commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2178885296
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -755,8 +758,9 @@ impl ArrowColumnWriter {
}
/// Encodes [`RecordBatch`] to a parquet row group
-struct ArrowRowGroupWriter {
- writers: Vec<ArrowColumnWriter>,
+pub struct ArrowRowGroupWriter {
+ /// [`ArrowColumnWriter`] for each column in a row group
+ pub writers: Vec<ArrowColumnWriter>,
Review Comment:
Rather than exposing this as a public member, it might be cleaner to add an `into_inner` or `into_column_writers` method that consumes the `ArrowRowGroupWriter` and returns the column writers. Then it's clearer that you probably shouldn't mix writing directly to the column writers with using the other methods on the `ArrowRowGroupWriter`.
##########
parquet/tests/encryption/encryption.rs:
##########
@@ -1101,3 +1109,117 @@ fn read_and_roundtrip_to_encrypted_file(
// check re-written example data
verify_encryption_test_file_read(temp_file, decryption_properties);
}
+
+#[tokio::test]
+async fn test_multi_threaded_encrypted_writing() {
+ // Read example data and set up encryption/decryption properties
+ let testdata = arrow::util::test_util::parquet_test_data();
+ let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted");
+
+ let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into())
+ .with_column_key("double_field", b"1234567890123450".into())
+ .with_column_key("float_field", b"1234567890123451".into())
+ .build()
+ .unwrap();
+ let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into())
+ .with_column_key("double_field", b"1234567890123450".into())
+ .with_column_key("float_field", b"1234567890123451".into())
+ .build()
+ .unwrap();
+
+ let (record_batches, metadata) =
+ read_encrypted_file(&path, decryption_properties.clone()).unwrap();
+ let to_write: Vec<_> = record_batches
+ .iter()
+ .flat_map(|rb| rb.columns().to_vec())
+ .collect();
+ let schema = metadata.schema().clone();
+
+ let props = Arc::new(
+ WriterPropertiesBuilder::with_defaults()
+ .with_file_encryption_properties(file_encryption_properties)
+ .build(),
+ );
+ let parquet_schema = ArrowSchemaConverter::new()
+ .with_coerce_types(props.coerce_types())
+ .convert(&schema)
+ .unwrap();
+ let root_schema = parquet_schema.root_schema_ptr();
+
+ // Create a temporary file to write the encrypted data
+ let temp_file = tempfile::NamedTempFile::new().unwrap();
+ let mut file_writer =
+ SerializedFileWriter::new(&temp_file, root_schema.clone(), props.clone()).unwrap();
+
+ let arrow_row_group_writer_factory = ArrowRowGroupWriterFactory::new(
+ &file_writer,
+ parquet_schema,
+ schema.clone(),
+ props.clone(),
+ );
+ let arrow_row_group_writer = arrow_row_group_writer_factory
+ .create_row_group_writer(0)
+ .unwrap();
+
+ // Get column writers with encryptor from ArrowRowGroupWriter
+ let col_writers = arrow_row_group_writer.writers;
+ let num_columns = col_writers.len();
+
+ // Create a channel for each column writer to send ArrowLeafColumn data to
+ let mut col_writer_tasks = Vec::with_capacity(num_columns);
+ let mut col_array_channels = Vec::with_capacity(num_columns);
+ for mut writer in col_writers.into_iter() {
+ let (send_array, mut receive_array) = tokio::sync::mpsc::channel::<ArrowLeafColumn>(100);
+ col_array_channels.push(send_array);
+ let handle = tokio::spawn(async move {
+ while let Some(col) = receive_array.recv().await {
+ let _ = writer.write(&col);
+ }
+ writer.close().unwrap()
+ });
+ col_writer_tasks.push(handle);
+ }
+
+ // Send the ArrowLeafColumn data to the respective column writer channels
+ for (channel_idx, (array, field)) in to_write.iter().zip(schema.fields()).enumerate() {
+ for c in compute_leaves(field, array).into_iter().flatten() {
+ let _ = col_array_channels[channel_idx].send(c).await;
+ }
+ }
Review Comment:
This doesn't look like the correct logic if you have a nested schema, eg. with struct typed columns. I guess it works fine in this test but should be corrected in case someone tries to use it as an example.
Above, one channel for each leaf level column writer is created, but then `col_array_channels` is indexed using `channel_idx`, which is the index of the top level Arrow schema field.
I'd expect this to look a bit more like the example [here](https://github.com/apache/arrow-rs/blob/52ad7d703acff0e5b4c143c179df206a842bf24e/parquet/src/arrow/arrow_writer/mod.rs#L662-L668).
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -199,9 +199,14 @@ impl<W: Write + Send> ArrowWriter<W> {
let max_row_group_size = props.max_row_group_size();
let file_writer =
- SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::new(props))?;
+ SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::new(props.clone()))?;
Review Comment:
We can avoid cloning all the properties and just clone the Arc, eg:
```rust
let props_ptr = Arc::new(props);
let file_writer =
SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::clone(&props_ptr))?;
```
(then change `props.into()` below to `props_ptr`)
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2201592759
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -755,7 +755,7 @@ impl ArrowColumnWriter {
}
/// Encodes [`RecordBatch`] to a parquet row group
-struct ArrowRowGroupWriter {
+pub struct ArrowRowGroupWriter {
Review Comment:
I do think avoiding exposing ArrowRowGroupWriter and the associated machinery that would be good
The rationale is that @XiangpengHao and @zhuqi-lucas and myself are likely to be reworking it in the next few releases and once it is part of the public API changing it is much harder
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "mbrobbel (via GitHub)" <gi...@apache.org>.
mbrobbel commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2204519641
##########
parquet/src/file/properties.rs:
##########
@@ -457,7 +457,7 @@ pub struct WriterPropertiesBuilder {
impl WriterPropertiesBuilder {
/// Returns default state of the builder.
- fn with_defaults() -> Self {
+ pub fn with_defaults() -> Self {
Review Comment:
Or `impl Default for WriterPropertiesBuilder`?
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "adamreeve (via GitHub)" <gi...@apache.org>.
adamreeve commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2199233416
##########
parquet/src/file/properties.rs:
##########
@@ -457,7 +457,7 @@ pub struct WriterPropertiesBuilder {
impl WriterPropertiesBuilder {
/// Returns default state of the builder.
- fn with_defaults() -> Self {
+ pub fn with_defaults() -> Self {
Review Comment:
We don't need to make this `pub`, you can call `WriterProperties::builder` instead. I think it's best to keep this private rather than exposing two ways to do the same thing.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok closed pull request #7818: Support multi-threaded writing of Parquet files with modular encryption
URL: https://github.com/apache/arrow-rs/pull/7818
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#issuecomment-3018975302
Thanks for the quick review @adamreeve !
I've moved schemas and writer properties to `ArrowRowGroupWriterFactory` as you suggested.
Test also now uses tokio channels.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2174734166
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -141,7 +141,7 @@ pub struct ArrowWriter<W: Write> {
arrow_schema: SchemaRef,
/// Creates new [`ArrowRowGroupWriter`] instances as required
- row_group_writer_factory: ArrowRowGroupWriterFactory,
+ pub row_group_writer_factory: ArrowRowGroupWriterFactory,
Review Comment:
Some of these `pub`s are artifacts of older approaches. Removing.
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -789,44 +790,54 @@ impl ArrowRowGroupWriter {
}
}
-struct ArrowRowGroupWriterFactory {
+/// Factory for creating [`ArrowRowGroupWriter`] instances.
+/// This is used by [`ArrowWriter`] to create row group writers, but can be used
+/// directly for lower level API.
+pub struct ArrowRowGroupWriterFactory {
#[cfg(feature = "encryption")]
file_encryptor: Option<Arc<FileEncryptor>>,
}
impl ArrowRowGroupWriterFactory {
+ /// Creates a new [`ArrowRowGroupWriterFactory`] using provided [`SerializedFileWriter`].
#[cfg(feature = "encryption")]
- fn new<W: Write + Send>(file_writer: &SerializedFileWriter<W>) -> Self {
+ pub fn new<W: Write + Send>(file_writer: &SerializedFileWriter<W>) -> Self {
Self {
file_encryptor: file_writer.file_encryptor(),
}
}
#[cfg(not(feature = "encryption"))]
- fn new<W: Write + Send>(_file_writer: &SerializedFileWriter<W>) -> Self {
+ pub fn new<W: Write + Send>(_file_writer: &SerializedFileWriter<W>) -> Self {
Self {}
}
+ /// Creates a new [`ArrowRowGroupWriter`] for the given parquet schema and writer properties.
#[cfg(feature = "encryption")]
- fn create_row_group_writer(
+ pub fn create_row_group_writer(
&self,
parquet: &SchemaDescriptor,
props: &WriterPropertiesPtr,
arrow: &SchemaRef,
row_group_index: usize,
) -> Result<ArrowRowGroupWriter> {
- let writers = get_column_writers_with_encryptor(
- parquet,
- props,
- arrow,
- self.file_encryptor.clone(),
- row_group_index,
- )?;
+ let mut writers = Vec::with_capacity(arrow.fields.len());
Review Comment:
Removed.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2174738626
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -883,14 +894,21 @@ fn get_column_writers_with_encryptor(
}
/// Gets [`ArrowColumnWriter`] instances for different data types
-struct ArrowColumnWriterFactory {
+pub struct ArrowColumnWriterFactory {
Review Comment:
Removed.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "adamreeve (via GitHub)" <gi...@apache.org>.
adamreeve commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2188109111
##########
parquet/tests/encryption/encryption.rs:
##########
@@ -1101,3 +1109,117 @@ fn read_and_roundtrip_to_encrypted_file(
// check re-written example data
verify_encryption_test_file_read(temp_file, decryption_properties);
}
+
+#[tokio::test]
+async fn test_multi_threaded_encrypted_writing() {
+ // Read example data and set up encryption/decryption properties
+ let testdata = arrow::util::test_util::parquet_test_data();
+ let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted");
+
+ let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into())
+ .with_column_key("double_field", b"1234567890123450".into())
+ .with_column_key("float_field", b"1234567890123451".into())
+ .build()
+ .unwrap();
+ let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into())
+ .with_column_key("double_field", b"1234567890123450".into())
+ .with_column_key("float_field", b"1234567890123451".into())
+ .build()
+ .unwrap();
+
+ let (record_batches, metadata) =
+ read_encrypted_file(&path, decryption_properties.clone()).unwrap();
+ let to_write: Vec<_> = record_batches
+ .iter()
+ .flat_map(|rb| rb.columns().to_vec())
+ .collect();
+ let schema = metadata.schema().clone();
+
+ let props = Arc::new(
+ WriterPropertiesBuilder::with_defaults()
+ .with_file_encryption_properties(file_encryption_properties)
+ .build(),
+ );
+ let parquet_schema = ArrowSchemaConverter::new()
+ .with_coerce_types(props.coerce_types())
+ .convert(&schema)
+ .unwrap();
+ let root_schema = parquet_schema.root_schema_ptr();
+
+ // Create a temporary file to write the encrypted data
+ let temp_file = tempfile::NamedTempFile::new().unwrap();
+ let mut file_writer =
+ SerializedFileWriter::new(&temp_file, root_schema.clone(), props.clone()).unwrap();
+
+ let arrow_row_group_writer_factory = ArrowRowGroupWriterFactory::new(
+ &file_writer,
+ parquet_schema,
+ schema.clone(),
+ props.clone(),
+ );
+ let arrow_row_group_writer = arrow_row_group_writer_factory
+ .create_row_group_writer(0)
+ .unwrap();
+
+ // Get column writers with encryptor from ArrowRowGroupWriter
+ let col_writers = arrow_row_group_writer.writers;
+ let num_columns = col_writers.len();
+
+ // Create a channel for each column writer to send ArrowLeafColumn data to
+ let mut col_writer_tasks = Vec::with_capacity(num_columns);
+ let mut col_array_channels = Vec::with_capacity(num_columns);
+ for mut writer in col_writers.into_iter() {
+ let (send_array, mut receive_array) = tokio::sync::mpsc::channel::<ArrowLeafColumn>(100);
+ col_array_channels.push(send_array);
+ let handle = tokio::spawn(async move {
+ while let Some(col) = receive_array.recv().await {
+ let _ = writer.write(&col);
+ }
+ writer.close().unwrap()
+ });
+ col_writer_tasks.push(handle);
+ }
+
+ // Send the ArrowLeafColumn data to the respective column writer channels
+ for (channel_idx, (array, field)) in to_write.iter().zip(schema.fields()).enumerate() {
+ for c in compute_leaves(field, array).into_iter().flatten() {
+ let _ = col_array_channels[channel_idx].send(c).await;
+ }
+ }
Review Comment:
I don't think this needs an extra test case given the implementation doesn't do anything to specifically handle struct types.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "adamreeve (via GitHub)" <gi...@apache.org>.
adamreeve commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2173944156
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -789,44 +790,54 @@ impl ArrowRowGroupWriter {
}
}
-struct ArrowRowGroupWriterFactory {
+/// Factory for creating [`ArrowRowGroupWriter`] instances.
+/// This is used by [`ArrowWriter`] to create row group writers, but can be used
+/// directly for lower level API.
+pub struct ArrowRowGroupWriterFactory {
#[cfg(feature = "encryption")]
file_encryptor: Option<Arc<FileEncryptor>>,
}
impl ArrowRowGroupWriterFactory {
+ /// Creates a new [`ArrowRowGroupWriterFactory`] using provided [`SerializedFileWriter`].
#[cfg(feature = "encryption")]
- fn new<W: Write + Send>(file_writer: &SerializedFileWriter<W>) -> Self {
+ pub fn new<W: Write + Send>(file_writer: &SerializedFileWriter<W>) -> Self {
Self {
file_encryptor: file_writer.file_encryptor(),
}
}
#[cfg(not(feature = "encryption"))]
- fn new<W: Write + Send>(_file_writer: &SerializedFileWriter<W>) -> Self {
+ pub fn new<W: Write + Send>(_file_writer: &SerializedFileWriter<W>) -> Self {
Self {}
}
+ /// Creates a new [`ArrowRowGroupWriter`] for the given parquet schema and writer properties.
#[cfg(feature = "encryption")]
- fn create_row_group_writer(
+ pub fn create_row_group_writer(
&self,
parquet: &SchemaDescriptor,
props: &WriterPropertiesPtr,
arrow: &SchemaRef,
row_group_index: usize,
) -> Result<ArrowRowGroupWriter> {
- let writers = get_column_writers_with_encryptor(
- parquet,
- props,
- arrow,
- self.file_encryptor.clone(),
- row_group_index,
- )?;
+ let mut writers = Vec::with_capacity(arrow.fields.len());
Review Comment:
Why is this changed? It doesn't look different to `get_column_writers_with_encryptor`
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -141,7 +141,7 @@ pub struct ArrowWriter<W: Write> {
arrow_schema: SchemaRef,
/// Creates new [`ArrowRowGroupWriter`] instances as required
- row_group_writer_factory: ArrowRowGroupWriterFactory,
+ pub row_group_writer_factory: ArrowRowGroupWriterFactory,
Review Comment:
Why is this `pub`?
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -883,14 +894,21 @@ fn get_column_writers_with_encryptor(
}
/// Gets [`ArrowColumnWriter`] instances for different data types
-struct ArrowColumnWriterFactory {
+pub struct ArrowColumnWriterFactory {
Review Comment:
Why is this pub?
##########
parquet/src/encryption/encrypt.rs:
##########
@@ -288,14 +288,14 @@ impl EncryptionPropertiesBuilder {
#[derive(Debug)]
/// The encryption configuration for a single Parquet file
-pub(crate) struct FileEncryptor {
+pub struct FileEncryptor {
Review Comment:
?
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -860,7 +871,7 @@ pub fn get_column_writers(
/// Returns the [`ArrowColumnWriter`] for a given schema and supports columnar encryption
#[cfg(feature = "encryption")]
-fn get_column_writers_with_encryptor(
+pub fn get_column_writers_with_encryptor(
Review Comment:
Why is this pub?
##########
parquet/tests/encryption/encryption.rs:
##########
@@ -1101,3 +1111,109 @@ fn read_and_roundtrip_to_encrypted_file(
// check re-written example data
verify_encryption_test_file_read(temp_file, decryption_properties);
}
+
+#[tokio::test]
+async fn test_multi_threaded_encrypted_writing() {
+ // Read example data and set up encryption/decryption properties
+ let testdata = arrow::util::test_util::parquet_test_data();
+ let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted");
+
+ let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into())
+ .with_column_key("double_field", b"1234567890123450".into())
+ .with_column_key("float_field", b"1234567890123451".into())
+ .build()
+ .unwrap();
+ let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into())
+ .with_column_key("double_field", b"1234567890123450".into())
+ .with_column_key("float_field", b"1234567890123451".into())
+ .build()
+ .unwrap();
+
+ let (record_batches, metadata) =
+ read_encrypted_file(&path, decryption_properties.clone()).unwrap();
+ let to_write: Vec<_> = record_batches
+ .iter()
+ .flat_map(|rb| rb.columns().to_vec())
+ .collect();
+ let schema = metadata.schema().clone();
+
+ let props = Arc::new(
+ WriterPropertiesBuilder::with_defaults()
+ .with_file_encryption_properties(file_encryption_properties)
+ .build(),
+ );
+ let parquet_schema = ArrowSchemaConverter::new()
+ .with_coerce_types(props.coerce_types())
+ .convert(&schema)
+ .unwrap();
+ let root_schema = parquet_schema.root_schema_ptr();
+
+ // Create a temporary file to write the encrypted data
+ let temp_file = tempfile::NamedTempFile::new().unwrap();
+ let mut file_writer =
+ SerializedFileWriter::new(&temp_file, root_schema.clone(), props.clone()).unwrap();
+
+ let arrow_row_group_writer_factory = ArrowRowGroupWriterFactory::new(&file_writer);
+ let arrow_row_group_writer = arrow_row_group_writer_factory
+ .create_row_group_writer(&parquet_schema, &props.clone(), &schema, 0)
Review Comment:
```suggestion
.create_row_group_writer(&parquet_schema, &props, &schema, 0)
```
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -789,44 +790,54 @@ impl ArrowRowGroupWriter {
}
}
-struct ArrowRowGroupWriterFactory {
+/// Factory for creating [`ArrowRowGroupWriter`] instances.
+/// This is used by [`ArrowWriter`] to create row group writers, but can be used
+/// directly for lower level API.
+pub struct ArrowRowGroupWriterFactory {
#[cfg(feature = "encryption")]
file_encryptor: Option<Arc<FileEncryptor>>,
}
impl ArrowRowGroupWriterFactory {
+ /// Creates a new [`ArrowRowGroupWriterFactory`] using provided [`SerializedFileWriter`].
#[cfg(feature = "encryption")]
- fn new<W: Write + Send>(file_writer: &SerializedFileWriter<W>) -> Self {
+ pub fn new<W: Write + Send>(file_writer: &SerializedFileWriter<W>) -> Self {
Self {
file_encryptor: file_writer.file_encryptor(),
}
}
#[cfg(not(feature = "encryption"))]
- fn new<W: Write + Send>(_file_writer: &SerializedFileWriter<W>) -> Self {
+ pub fn new<W: Write + Send>(_file_writer: &SerializedFileWriter<W>) -> Self {
Self {}
}
+ /// Creates a new [`ArrowRowGroupWriter`] for the given parquet schema and writer properties.
#[cfg(feature = "encryption")]
- fn create_row_group_writer(
+ pub fn create_row_group_writer(
Review Comment:
Rather than have to pass the `SchemaDescriptor`, `WriterPropertiesPtr` and `SchemaRef` each time we create a new row group writer, this should probably be refactored so we store those in the `ArrowRowGroupWriterFactory` and set them in `new`. As those shouldn't be allowed to change between row groups.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2174739271
##########
parquet/src/encryption/encrypt.rs:
##########
@@ -288,14 +288,14 @@ impl EncryptionPropertiesBuilder {
#[derive(Debug)]
/// The encryption configuration for a single Parquet file
-pub(crate) struct FileEncryptor {
+pub struct FileEncryptor {
Review Comment:
Removed.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2198388531
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -755,7 +755,7 @@ impl ArrowColumnWriter {
}
/// Encodes [`RecordBatch`] to a parquet row group
-struct ArrowRowGroupWriter {
+pub struct ArrowRowGroupWriter {
Review Comment:
Thanks for the quick response @alamb ! :)
Looking at [the example](https://docs.rs/parquet/latest/parquet/arrow/arrow_writer/struct.ArrowColumnWriter.html#example-encoding-two-arrow-arrays-in-parallel) - we have to pass a single private `FileEncryptor` object (it has random state) to all `ArrowColumnWriter`s, `SerializedRowGroupWriter`s and `SerializedFileWriter`. Perhaps we could hide everything into `ArrowWriter` I'll have to do a quick check. If you have a cleaner idea we'd be most interested!
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2198312443
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -755,7 +755,7 @@ impl ArrowColumnWriter {
}
/// Encodes [`RecordBatch`] to a parquet row group
-struct ArrowRowGroupWriter {
+pub struct ArrowRowGroupWriter {
Review Comment:
Is there a reason we can't use the existing low level APIs documented here: https://docs.rs/parquet/latest/parquet/arrow/arrow_writer/struct.ArrowColumnWriter.html#example-encoding-two-arrow-arrays-in-parallel
In other words, do we really need to make what has previously been an implementation detail part of the public API?
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#issuecomment-3137259655
Thanks for the [ping](https://github.com/apache/arrow-rs/pull/7818#pullrequestreview-3052860514) @alamb! Unfortunately I was distracted and couldn't finish this in time for the release.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#issuecomment-3137234503
Closing in favor of #8029.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2181042123
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -755,8 +758,9 @@ impl ArrowColumnWriter {
}
/// Encodes [`RecordBatch`] to a parquet row group
-struct ArrowRowGroupWriter {
- writers: Vec<ArrowColumnWriter>,
+pub struct ArrowRowGroupWriter {
+ /// [`ArrowColumnWriter`] for each column in a row group
+ pub writers: Vec<ArrowColumnWriter>,
Review Comment:
Good point. Going with `into_column_writers`.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2174960827
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -789,44 +790,54 @@ impl ArrowRowGroupWriter {
}
}
-struct ArrowRowGroupWriterFactory {
+/// Factory for creating [`ArrowRowGroupWriter`] instances.
+/// This is used by [`ArrowWriter`] to create row group writers, but can be used
+/// directly for lower level API.
+pub struct ArrowRowGroupWriterFactory {
#[cfg(feature = "encryption")]
file_encryptor: Option<Arc<FileEncryptor>>,
}
impl ArrowRowGroupWriterFactory {
+ /// Creates a new [`ArrowRowGroupWriterFactory`] using provided [`SerializedFileWriter`].
#[cfg(feature = "encryption")]
- fn new<W: Write + Send>(file_writer: &SerializedFileWriter<W>) -> Self {
+ pub fn new<W: Write + Send>(file_writer: &SerializedFileWriter<W>) -> Self {
Self {
file_encryptor: file_writer.file_encryptor(),
}
}
#[cfg(not(feature = "encryption"))]
- fn new<W: Write + Send>(_file_writer: &SerializedFileWriter<W>) -> Self {
+ pub fn new<W: Write + Send>(_file_writer: &SerializedFileWriter<W>) -> Self {
Self {}
}
+ /// Creates a new [`ArrowRowGroupWriter`] for the given parquet schema and writer properties.
#[cfg(feature = "encryption")]
- fn create_row_group_writer(
+ pub fn create_row_group_writer(
Review Comment:
Done.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2174735192
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -860,7 +871,7 @@ pub fn get_column_writers(
/// Returns the [`ArrowColumnWriter`] for a given schema and supports columnar encryption
#[cfg(feature = "encryption")]
-fn get_column_writers_with_encryptor(
+pub fn get_column_writers_with_encryptor(
Review Comment:
Removed.
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on code in PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#discussion_r2198388531
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -755,7 +755,7 @@ impl ArrowColumnWriter {
}
/// Encodes [`RecordBatch`] to a parquet row group
-struct ArrowRowGroupWriter {
+pub struct ArrowRowGroupWriter {
Review Comment:
Thanks for the quick response @alamb ! :)
Looking at [the example](https://docs.rs/parquet/latest/parquet/arrow/arrow_writer/struct.ArrowColumnWriter.html#example-encoding-two-arrow-arrays-in-parallel) - we have to pass a single private `FileEncryptor` object (it has random state) to all `ArrowColumnWriter`s, `SerializedRowGroupWriter`s and `SerializedFileWriter`. Perhaps we could hide everything into `ArrowWriter`. I'll have to do a quick check. If you have a cleaner idea we'd be most interested!
--
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] Support multi-threaded writing of Parquet files with modular encryption [arrow-rs]
Posted by "rok (via GitHub)" <gi...@apache.org>.
rok commented on PR #7818:
URL: https://github.com/apache/arrow-rs/pull/7818#issuecomment-3057431264
I've opened a [draft PR](https://github.com/apache/datafusion/pull/16738) against datafusion to check if this is usable for multi-threaded writes and it does seem to be. Switching this to ready for review.
--
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