Repository: pola-rs/polars
File: crates/polars-stream/src/nodes/io_sinks/partition/max_size.rs (Rust)

Code before:
        Self {
            input_schema,
            max_size,
            base_path,
            file_path_cb,
            create_new,
            ext,
            sink_options,
            num_retire_tasks,
            per_partition_sort_by,
            partition_metrics: Arc::new(Mutex::new(Vec::with_capacity(num_retire_tasks))),
            finish_callback,
        }
    }
}

fn default_file_path_cb(
    ext: &str,
    file_idx: usize,
    _part_idx: usize,
    _in_part_idx: usize,
    _columns: Option<&[Column]>,
    _separator: char,
) -> PolarsResult<String> {
    polars_ensure!(file_idx < 1<<(7*4), ComputeError: "exceeded maximum file count (268,435,455)");
    Ok(format!("{file_idx:07X}.{ext}"))
}

impl SinkNode for MaxSizePartitionSinkNode {
    fn name(&self) -> &str {
        "partition-max-size-sink"
    }

    fn is_sink_input_parallel(&self) -> bool {
        false
    }
    fn do_maintain_order(&self) -> bool {
        self.sink_options.maintain_order
    }

    fn spawn_sink(
        &mut self,
        mut recv_port_recv: Receiver<(PhaseOutcome, SinkInputPort)>,
        state: &StreamingExecutionState,
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        // Main Task -> Retire Tasks
        let (mut retire_tx, retire_rxs) = distributor_channel(self.num_retire_tasks, 1);

        // Whether an error has been observed in the retire tasks.
        let has_error_occurred = Arc::new(AtomicBool::new(false));

Review comment (suggestion):
We should just use 8 hexadecimals here, it's much easier to reason about that way, and also less likely to cause problems down the line.


Code after:
        Self {
            input_schema,
            max_size,
            base_path,
            file_path_cb,
            create_new,
            ext,
            sink_options,
            num_retire_tasks,
            per_partition_sort_by,
            partition_metrics: Arc::new(Mutex::new(Vec::with_capacity(num_retire_tasks))),
            finish_callback,
        }
    }
}

fn default_file_path_cb(
    ext: &str,
    file_idx: usize,
    _part_idx: usize,
    _in_part_idx: usize,
    _columns: Option<&[Column]>,
    _separator: char,
) -> PolarsResult<String> {
    polars_ensure!(file_idx < u32::MAX as usize,
        ComputeError: "exceeded maximum file count within a partition of {}", u32::MAX);
    Ok(format!("{file_idx:08x}.{ext}"))
}

impl SinkNode for MaxSizePartitionSinkNode {
    fn name(&self) -> &str {
        "partition-max-size-sink"
    }

    fn is_sink_input_parallel(&self) -> bool {
        false
    }
    fn do_maintain_order(&self) -> bool {
        self.sink_options.maintain_order
    }

    fn spawn_sink(
        &mut self,
        mut recv_port_recv: Receiver<(PhaseOutcome, SinkInputPort)>,
        state: &StreamingExecutionState,
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        // Main Task -> Retire Tasks
        let (mut retire_tx, retire_rxs) = distributor_channel(self.num_retire_tasks, 1);

        // Whether an error has been observed in the retire tasks.