iceberg/writer/partitioning/
clustered_writer.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! This module provides the `ClusteredWriter` implementation.
19
20use std::collections::HashSet;
21use std::marker::PhantomData;
22
23use async_trait::async_trait;
24
25use crate::spec::{PartitionKey, Struct};
26use crate::writer::partitioning::PartitioningWriter;
27use crate::writer::{DefaultInput, DefaultOutput, IcebergWriter, IcebergWriterBuilder};
28use crate::{Error, ErrorKind, Result};
29
30/// A writer that writes data to a single partition at a time.
31///
32/// This writer expects input data to be sorted by partition key. It maintains only one
33/// active writer at a time, making it memory efficient for sorted data.
34///
35/// # Type Parameters
36///
37/// * `B` - The inner writer builder type
38/// * `I` - Input type (defaults to `RecordBatch`)
39/// * `O` - Output collection type (defaults to `Vec<DataFile>`)
40pub struct ClusteredWriter<B, I = DefaultInput, O = DefaultOutput>
41where
42    B: IcebergWriterBuilder<I, O>,
43    O: IntoIterator + FromIterator<<O as IntoIterator>::Item>,
44    <O as IntoIterator>::Item: Clone,
45{
46    inner_builder: B,
47    current_writer: Option<B::R>,
48    current_partition: Option<Struct>,
49    closed_partitions: HashSet<Struct>,
50    output: Vec<<O as IntoIterator>::Item>,
51    _phantom: PhantomData<I>,
52}
53
54impl<B, I, O> ClusteredWriter<B, I, O>
55where
56    B: IcebergWriterBuilder<I, O>,
57    I: Send + 'static,
58    O: IntoIterator + FromIterator<<O as IntoIterator>::Item>,
59    <O as IntoIterator>::Item: Send + Clone,
60{
61    /// Create a new `ClusteredWriter`.
62    pub fn new(inner_builder: B) -> Self {
63        Self {
64            inner_builder,
65            current_writer: None,
66            current_partition: None,
67            closed_partitions: HashSet::new(),
68            output: Vec::new(),
69            _phantom: PhantomData,
70        }
71    }
72
73    /// Closes the current writer if it exists, flushes the written data to output, and record closed partition.
74    async fn close_current_writer(&mut self) -> Result<()> {
75        if let Some(mut writer) = self.current_writer.take() {
76            self.output.extend(writer.close().await?);
77
78            // Add the current partition to the set of closed partitions
79            if let Some(current_partition) = self.current_partition.take() {
80                self.closed_partitions.insert(current_partition);
81            }
82        }
83
84        Ok(())
85    }
86}
87
88#[async_trait]
89impl<B, I, O> PartitioningWriter<I, O> for ClusteredWriter<B, I, O>
90where
91    B: IcebergWriterBuilder<I, O>,
92    I: Send + 'static,
93    O: IntoIterator + FromIterator<<O as IntoIterator>::Item> + Send + 'static,
94    <O as IntoIterator>::Item: Send + Clone,
95{
96    async fn write(&mut self, partition_key: PartitionKey, input: I) -> Result<()> {
97        let partition_value = partition_key.data();
98
99        // Check if this partition has been closed already
100        if self.closed_partitions.contains(partition_value) {
101            return Err(Error::new(
102                ErrorKind::Unexpected,
103                format!(
104                    "The input is not sorted! Cannot write to partition that was previously closed: {partition_key:?}"
105                ),
106            ));
107        }
108
109        // Check if we need to switch to a new partition
110        let need_new_writer = match &self.current_partition {
111            Some(current) => current != partition_value,
112            None => true,
113        };
114
115        if need_new_writer {
116            self.close_current_writer().await?;
117
118            // Create a new writer for the new partition
119            self.current_writer = Some(
120                self.inner_builder
121                    .build(Some(partition_key.clone()))
122                    .await?,
123            );
124            self.current_partition = Some(partition_value.clone());
125        }
126
127        // do write
128        self.current_writer
129            .as_mut()
130            .expect("Writer should be initialized")
131            .write(input)
132            .await
133    }
134
135    async fn close(mut self) -> Result<O> {
136        self.close_current_writer().await?;
137
138        // Collect all output items into the output collection type
139        Ok(O::from_iter(self.output))
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use std::collections::HashMap;
146    use std::sync::Arc;
147
148    use arrow_array::{Int32Array, RecordBatch, StringArray};
149    use arrow_schema::{DataType, Field, Schema};
150    use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
151    use parquet::file::properties::WriterProperties;
152    use tempfile::TempDir;
153
154    use super::*;
155    use crate::io::FileIO;
156    use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Type};
157    use crate::writer::base_writer::data_file_writer::DataFileWriterBuilder;
158    use crate::writer::file_writer::ParquetWriterBuilder;
159    use crate::writer::file_writer::location_generator::{
160        DefaultFileNameGenerator, DefaultLocationGenerator,
161    };
162    use crate::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
163
164    #[tokio::test]
165    async fn test_clustered_writer_single_partition() -> Result<()> {
166        let temp_dir = TempDir::new()?;
167        let file_io = FileIO::new_with_fs();
168        let location_gen = DefaultLocationGenerator::with_data_location(
169            temp_dir.path().to_str().unwrap().to_string(),
170        );
171        let file_name_gen =
172            DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
173
174        // Create schema with partition field
175        let schema = Arc::new(
176            crate::spec::Schema::builder()
177                .with_schema_id(1)
178                .with_fields(vec![
179                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
180                    NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
181                    NestedField::required(3, "region", Type::Primitive(PrimitiveType::String))
182                        .into(),
183                ])
184                .build()?,
185        );
186
187        // Create partition spec and key
188        let partition_spec = crate::spec::PartitionSpec::builder(schema.clone()).build()?;
189        let partition_value = Struct::from_iter([Some(crate::spec::Literal::string("US"))]);
190        let partition_key =
191            PartitionKey::new(partition_spec, schema.clone(), partition_value.clone());
192
193        // Create writer builder
194        let parquet_writer_builder =
195            ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone());
196
197        // Create rolling file writer builder
198        let rolling_writer_builder = RollingFileWriterBuilder::new_with_default_file_size(
199            parquet_writer_builder,
200            file_io.clone(),
201            location_gen,
202            file_name_gen,
203        );
204
205        // Create data file writer builder
206        let data_file_writer_builder = DataFileWriterBuilder::new(rolling_writer_builder);
207
208        // Create clustered writer
209        let mut writer = ClusteredWriter::new(data_file_writer_builder);
210
211        // Create test data with proper field ID metadata
212        let arrow_schema = Schema::new(vec![
213            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
214                PARQUET_FIELD_ID_META_KEY.to_string(),
215                1.to_string(),
216            )])),
217            Field::new("name", DataType::Utf8, false).with_metadata(HashMap::from([(
218                PARQUET_FIELD_ID_META_KEY.to_string(),
219                2.to_string(),
220            )])),
221            Field::new("region", DataType::Utf8, false).with_metadata(HashMap::from([(
222                PARQUET_FIELD_ID_META_KEY.to_string(),
223                3.to_string(),
224            )])),
225        ]);
226
227        let batch1 = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
228            Arc::new(Int32Array::from(vec![1, 2])),
229            Arc::new(StringArray::from(vec!["Alice", "Bob"])),
230            Arc::new(StringArray::from(vec!["US", "US"])),
231        ])?;
232
233        let batch2 = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
234            Arc::new(Int32Array::from(vec![3, 4])),
235            Arc::new(StringArray::from(vec!["Charlie", "Dave"])),
236            Arc::new(StringArray::from(vec!["US", "US"])),
237        ])?;
238
239        // Write data to the same partition (this should work)
240        writer.write(partition_key.clone(), batch1).await?;
241        writer.write(partition_key.clone(), batch2).await?;
242
243        // Close writer and get data files
244        let data_files = writer.close().await?;
245
246        // Verify at least one file was created
247        assert!(
248            !data_files.is_empty(),
249            "Expected at least one data file to be created"
250        );
251
252        // Verify that all data files have the correct partition value
253        for data_file in &data_files {
254            assert_eq!(data_file.partition, partition_value);
255        }
256
257        Ok(())
258    }
259
260    #[tokio::test]
261    async fn test_clustered_writer_sorted_partitions() -> Result<()> {
262        let temp_dir = TempDir::new()?;
263        let file_io = FileIO::new_with_fs();
264        let location_gen = DefaultLocationGenerator::with_data_location(
265            temp_dir.path().to_str().unwrap().to_string(),
266        );
267        let file_name_gen =
268            DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
269
270        // Create schema with partition field
271        let schema = Arc::new(
272            crate::spec::Schema::builder()
273                .with_schema_id(1)
274                .with_fields(vec![
275                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
276                    NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
277                    NestedField::required(3, "region", Type::Primitive(PrimitiveType::String))
278                        .into(),
279                ])
280                .build()?,
281        );
282
283        // Create partition spec
284        let partition_spec = crate::spec::PartitionSpec::builder(schema.clone()).build()?;
285
286        // Create partition keys for different regions (in sorted order)
287        let partition_value_asia = Struct::from_iter([Some(crate::spec::Literal::string("ASIA"))]);
288        let partition_key_asia = PartitionKey::new(
289            partition_spec.clone(),
290            schema.clone(),
291            partition_value_asia.clone(),
292        );
293
294        let partition_value_eu = Struct::from_iter([Some(crate::spec::Literal::string("EU"))]);
295        let partition_key_eu = PartitionKey::new(
296            partition_spec.clone(),
297            schema.clone(),
298            partition_value_eu.clone(),
299        );
300
301        let partition_value_us = Struct::from_iter([Some(crate::spec::Literal::string("US"))]);
302        let partition_key_us = PartitionKey::new(
303            partition_spec.clone(),
304            schema.clone(),
305            partition_value_us.clone(),
306        );
307
308        // Create writer builder
309        let parquet_writer_builder =
310            ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone());
311
312        // Create rolling file writer builder
313        let rolling_writer_builder = RollingFileWriterBuilder::new_with_default_file_size(
314            parquet_writer_builder,
315            file_io.clone(),
316            location_gen,
317            file_name_gen,
318        );
319
320        // Create data file writer builder
321        let data_file_writer_builder = DataFileWriterBuilder::new(rolling_writer_builder);
322
323        // Create clustered writer
324        let mut writer = ClusteredWriter::new(data_file_writer_builder);
325
326        // Create test data with proper field ID metadata
327        let arrow_schema = Schema::new(vec![
328            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
329                PARQUET_FIELD_ID_META_KEY.to_string(),
330                1.to_string(),
331            )])),
332            Field::new("name", DataType::Utf8, false).with_metadata(HashMap::from([(
333                PARQUET_FIELD_ID_META_KEY.to_string(),
334                2.to_string(),
335            )])),
336            Field::new("region", DataType::Utf8, false).with_metadata(HashMap::from([(
337                PARQUET_FIELD_ID_META_KEY.to_string(),
338                3.to_string(),
339            )])),
340        ]);
341
342        // Create batches for different partitions (in sorted order)
343        let batch_asia = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
344            Arc::new(Int32Array::from(vec![1, 2])),
345            Arc::new(StringArray::from(vec!["Alice", "Bob"])),
346            Arc::new(StringArray::from(vec!["ASIA", "ASIA"])),
347        ])?;
348
349        let batch_eu = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
350            Arc::new(Int32Array::from(vec![3, 4])),
351            Arc::new(StringArray::from(vec!["Charlie", "Dave"])),
352            Arc::new(StringArray::from(vec!["EU", "EU"])),
353        ])?;
354
355        let batch_us = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
356            Arc::new(Int32Array::from(vec![5, 6])),
357            Arc::new(StringArray::from(vec!["Eve", "Frank"])),
358            Arc::new(StringArray::from(vec!["US", "US"])),
359        ])?;
360
361        // Write data in sorted partition order (this should work)
362        writer.write(partition_key_asia.clone(), batch_asia).await?;
363        writer.write(partition_key_eu.clone(), batch_eu).await?;
364        writer.write(partition_key_us.clone(), batch_us).await?;
365
366        // Close writer and get data files
367        let data_files = writer.close().await?;
368
369        // Verify files were created for all partitions
370        assert!(
371            data_files.len() >= 3,
372            "Expected at least 3 data files (one per partition), got {}",
373            data_files.len()
374        );
375
376        // Verify that we have files for each partition
377        let mut partitions_found = HashSet::new();
378        for data_file in &data_files {
379            partitions_found.insert(data_file.partition.clone());
380        }
381
382        assert!(
383            partitions_found.contains(&partition_value_asia),
384            "Missing ASIA partition"
385        );
386        assert!(
387            partitions_found.contains(&partition_value_eu),
388            "Missing EU partition"
389        );
390        assert!(
391            partitions_found.contains(&partition_value_us),
392            "Missing US partition"
393        );
394
395        Ok(())
396    }
397
398    #[tokio::test]
399    async fn test_clustered_writer_unsorted_partitions_error() -> Result<()> {
400        let temp_dir = TempDir::new()?;
401        let file_io = FileIO::new_with_fs();
402        let location_gen = DefaultLocationGenerator::with_data_location(
403            temp_dir.path().to_str().unwrap().to_string(),
404        );
405        let file_name_gen =
406            DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
407
408        // Create schema with partition field
409        let schema = Arc::new(
410            crate::spec::Schema::builder()
411                .with_schema_id(1)
412                .with_fields(vec![
413                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
414                    NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
415                    NestedField::required(3, "region", Type::Primitive(PrimitiveType::String))
416                        .into(),
417                ])
418                .build()?,
419        );
420
421        // Create partition spec
422        let partition_spec = crate::spec::PartitionSpec::builder(schema.clone()).build()?;
423
424        // Create partition keys for different regions
425        let partition_value_us = Struct::from_iter([Some(crate::spec::Literal::string("US"))]);
426        let partition_key_us = PartitionKey::new(
427            partition_spec.clone(),
428            schema.clone(),
429            partition_value_us.clone(),
430        );
431
432        let partition_value_eu = Struct::from_iter([Some(crate::spec::Literal::string("EU"))]);
433        let partition_key_eu = PartitionKey::new(
434            partition_spec.clone(),
435            schema.clone(),
436            partition_value_eu.clone(),
437        );
438
439        // Create writer builder
440        let parquet_writer_builder =
441            ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone());
442
443        // Create rolling file writer builder
444        let rolling_writer_builder = RollingFileWriterBuilder::new_with_default_file_size(
445            parquet_writer_builder,
446            file_io.clone(),
447            location_gen,
448            file_name_gen,
449        );
450
451        // Create data file writer builder
452        let data_file_writer_builder = DataFileWriterBuilder::new(rolling_writer_builder);
453
454        // Create clustered writer
455        let mut writer = ClusteredWriter::new(data_file_writer_builder);
456
457        // Create test data with proper field ID metadata
458        let arrow_schema = Schema::new(vec![
459            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
460                PARQUET_FIELD_ID_META_KEY.to_string(),
461                1.to_string(),
462            )])),
463            Field::new("name", DataType::Utf8, false).with_metadata(HashMap::from([(
464                PARQUET_FIELD_ID_META_KEY.to_string(),
465                2.to_string(),
466            )])),
467            Field::new("region", DataType::Utf8, false).with_metadata(HashMap::from([(
468                PARQUET_FIELD_ID_META_KEY.to_string(),
469                3.to_string(),
470            )])),
471        ]);
472
473        // Create batches for different partitions
474        let batch_us = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
475            Arc::new(Int32Array::from(vec![1, 2])),
476            Arc::new(StringArray::from(vec!["Alice", "Bob"])),
477            Arc::new(StringArray::from(vec!["US", "US"])),
478        ])?;
479
480        let batch_eu = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
481            Arc::new(Int32Array::from(vec![3, 4])),
482            Arc::new(StringArray::from(vec!["Charlie", "Dave"])),
483            Arc::new(StringArray::from(vec!["EU", "EU"])),
484        ])?;
485
486        let batch_us2 = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
487            Arc::new(Int32Array::from(vec![5])),
488            Arc::new(StringArray::from(vec!["Eve"])),
489            Arc::new(StringArray::from(vec!["US"])),
490        ])?;
491
492        // Write data to US partition first
493        writer.write(partition_key_us.clone(), batch_us).await?;
494
495        // Write data to EU partition (this closes US partition)
496        writer.write(partition_key_eu.clone(), batch_eu).await?;
497
498        // Try to write to US partition again - this should fail because data is not sorted
499        let result = writer.write(partition_key_us.clone(), batch_us2).await;
500
501        assert!(result.is_err(), "Expected error when writing unsorted data");
502
503        let error = result.unwrap_err();
504        assert!(
505            error.to_string().contains("The input is not sorted"),
506            "Expected 'input is not sorted' error, got: {error}"
507        );
508
509        Ok(())
510    }
511}
iceberg/writer/partitioning/clustered_writer.rs

iceberg/writer/partitioning/
clustered_writer.rs
