Decouple ParquetWriter and LocationGenerator

[CTTY]

Is your feature request related to a problem or challenge?

Hi icy crustaceans,

While working on adding the [ClusteredWriter](https://github.com/apache/iceberg-rust/pull/1646), I noticed that the existing FileWriterBuilder trait is quite limiting for partition-aware writers. The main issue is that the location generator is coupled with the low-level file writer (such as ParquetWriter), which effectively requires ParquetWriter itself to be partition-aware.

In addition, the current FileWriterBuilder trait cannot pass partition information from outer writers to inner writers, making it very difficult to implement partitioning writers:

fn build(self) -> impl Future<Output = Result<Self::R>> + Send;

A hacky workaround would be to have partitioning writers use the concrete type ParquetWriter as the inner writer, but I think this will lead to extensibility issues fairly quickly.

Describe the solution you'd like

I see a few aspects of this issue:

• Decouple LocationGenerator and ParquetWriter by changing the FileWriterBuilder trait

By changing the FileWriterBuilder interface to the following, we can make ParquetWriter focus solely on writing a file without needing to know which partition it’s writing to:

fn build(self, output_file: OutputFile) -> impl Future<Output = Result<Self::R>> + Send;

We will also need to adjust IcebergWriterBuilder accordingly:

async fn build(self, output_file: OutputFile) -> Result<Self::R>;

• Make the functional writer the outermost layer

After modifying the builder traits, we still need to decide which layer of writers should hold the LocationGenerator. I propose moving functional writers such as RollingFileWriter and the upcoming ClusteredWriter to the outermost layer, and having them own the LocationGenerator.

Currently, the layering looks like this:

// DataFileWriter/IcebergWriter is the outermost layer
DataFileWriter(
  RollingFileWriter( // currently implements FileWriter
    ParquetWriter()
  )
)

My proposal:

SomeFunctionalWriter( // no longer implements FileWriter
  DataFileWriter(
    ParquetWriter()
  )
)

This new design provides several benefits:

1. Clear layering: Functional writers hold the LocationGenerator and generate new OutputFiles for the inner writers, while the inner writers focus only on writing a single file without worrying about path or location.

2. Flexibility: Functional writers no longer need to implement the FileWriter trait, making them easier to use and configure. They don’t necessarily need to implement any trait, though we could define something like:

   FunctionalWriterBuilder {
     async fn build(
       self,
       location_gen: LocationGenerator,
       filename_gen: FileNameGenerator
     ) -> Result<Self::R>;
   }

I’m still exploring and refining this idea, and would appreciate any feedback!

Willingness to contribute

I would be willing to contribute to this feature with guidance from the Iceberg Rust community

[liurenjie1024]

Thanks @CTTY for raising this. I agree that the current writer class hierarchy is somehow cumbersome and over complicated. I plan to do a refactoring after the FileFormat api proposal got merged.

For now I would suggest to make the changes minimum. So here is what's in my mind:

1.Adding one layer: PartitioningWriter, which sould like

pub trait PartitioningWriter {
    fn write(&self, partition_key: PartitionKey, batch: RecordBatch);
}

struct ClusterdPartitioningWriter {...}

impl PartitioningWriter for ClusterdPartitioningWriter {}

struct FanoutParritioningWriter {
   writers: HashMap<Struct, RollingFileWriter>
....
}

impl PartitioningWriter for FanoutParritioningWriter {}

2. Remove the FileWriter trait from RollingFileWriter
3. Move the LocationGenerator and FilenameGenerator to RollingFileWriter
4. Change FileWriterBuilder trait to accept an OutputFile to build FileWriter

WDYT?

[CTTY]

Hi @liurenjie1024 , thanks for the inputs! Your idea sounds good to me and I agree that we should make smaller steps if possible. Next I'll try to make a draft based on it!

One thing I'm not too sure about the PartitioningWriter interface is that the incoming batch may still contain rows from different partitions

pub trait PartitioningWriter {
// if `batch` here contains data from multiple partitions, 
// then the entire batch would still be written to the partition of `partition_key`
    fn write(&self, partition_key: PartitionKey, batch: RecordBatch);
}

I'm thinking of something like this:

pub trait PartitioningWriter {
    // use record batch splitter to split the incoming batch first
    fn write(&self, batch: RecordBatch); 
    // the `batch` here should be splitted only
    // technically this shouldn't be public accessible
    fn do_write(&self, partition_key: PartitionKey, splitted_batch: RecordBatch);
}

Please lmk your thoughts!

[CTTY]

Another issue is the existing DataFileWriter is only aware of files from a single partition (link), I think we will need to remove partition_value and have the underlying PartitioningWriter to find out partition value by scanning through the input record batches

[ZENOTME]

Thanks @CTTY for mention this! I think this feature is also what #891 request. Th e design looks good to me.

[liurenjie1024]

Hi @liurenjie1024 , thanks for the inputs! Your idea sounds good to me and I agree that we should make smaller steps if possible. Next I'll try to make a draft based on it!

One thing I'm not too sure about the PartitioningWriter interface is that the incoming batch may still contain rows from different partitions

pub trait PartitioningWriter {
// if batch here contains data from multiple partitions,
// then the entire batch would still be written to the partition of partition_key
fn write(&self, partition_key: PartitionKey, batch: RecordBatch);
}
I'm thinking of something like this:

pub trait PartitioningWriter {
// use record batch splitter to split the incoming batch first
fn write(&self, batch: RecordBatch);
// the batch here should be splitted only
// technically this shouldn't be public accessible
fn do_write(&self, partition_key: PartitionKey, splitted_batch: RecordBatch);
}
Please lmk your thoughts!

Hi, @CTTY The reason we need to put the partition_key in the interface is that in some frameworks/physical plan, the partition key could be computed in other places, not in the writer node. I think we don't need to change DataFileWriter? The class hierarchy be like:

/// TaskWriter is engine specific.
struct TaskWriter {
   /// Maybe we could use an enum to replace them
   partitioned_datafile_writer: Option<PartitionedWriter>,
   unpartitioned_datafile_writer: Option<UnPartitionedWriter>,
}

struct ClusterdPartitioningDataWriter {...}

impl PartitioningWriter for ClusterdPartitioningWriter {}

struct FanoutParritioningDataWriter {
   writers: HashMap<Struct, DataFileWriter>
....
}

impl PartitioningWriter for FanoutParritioningWriter {}

I did a little change so that PartitionedWriter contains DataFileWriter rather FileWriter

[CTTY]

Hi @liurenjie1024 , thanks for the inputs! I've changed my design in the draft accordingly, ptal!