Skip to content

Commit ffc3d3a

Browse files
nikhilsinhaparseablenikhilsinhaparseable
committed
deepsource fix
1 parent 5bf0f1a commit ffc3d3a

3 files changed

Lines changed: 266 additions & 175 deletions

File tree

‎src/parseable/staging/reader.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,14 +77,14 @@ impl ForwardFilePlan {
7777
/// Every stream begins with the original schema and receives every
7878
/// dictionary update, preserving Arrow IPC decoder state.
7979
pub(crate) fn reader_lane_count(&self, readers_per_file: usize) -> usize {
80-
readers_per_file.max(1).min(self.record_batches)
80+
readers_per_file.clamp(1, self.record_batches)
8181
}
8282

8383
pub(crate) fn open_lanes(
8484
self,
8585
readers_per_file: usize,
8686
) -> Result<Vec<ForwardReaderLane>, io::Error> {
87-
let lane_count = readers_per_file.max(1).min(self.record_batches);
87+
let lane_count = readers_per_file.clamp(1, self.record_batches);
8888
let (_, schema_offset, schema_size) = self.messages[0];
8989
let schema_range = (schema_offset as u64, schema_size);
9090
let mut lane_ranges = vec![vec![schema_range]; lane_count];

‎src/parseable/streams.rs‎

Lines changed: 160 additions & 145 deletions
Original file line numberDiff line numberDiff line change
@@ -531,10 +531,7 @@ impl MetricForwardRecordIterator {
531531

532532
let total_limit = *METRIC_ARROW_READERS_IN_FLIGHT;
533533
let per_file_limit = (*METRIC_ARROW_READERS_PER_FILE).min(total_limit);
534-
loop {
535-
let Some(plan) = self.plans.as_slice().first() else {
536-
break;
537-
};
534+
while let Some(plan) = self.plans.as_slice().first() {
538535
let reader_lanes = plan.reader_lane_count(per_file_limit);
539536
if !self.active_files.is_empty() && self.readers_in_flight + reader_lanes > total_limit
540537
{
@@ -993,7 +990,7 @@ impl Stream {
993990
/// the same minute.
994991
fn create_inprocess_folder(&self, minute: u128) -> std::io::Result<PathBuf> {
995992
let preferred = Self::inprocess_folder(&self.data_path, minute);
996-
match fs::create_dir(&preferred) {
993+
match fs::DirBuilder::new().create(&preferred) {
997994
Ok(()) => return Ok(preferred),
998995
Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {}
999996
Err(err) => return Err(err),
@@ -1003,7 +1000,7 @@ impl Stream {
10031000
let unique = self
10041001
.data_path
10051002
.join(format!("{INPROCESS_DIR_PREFIX}{minute}_{}", Ulid::new()));
1006-
match fs::create_dir(&unique) {
1003+
match fs::DirBuilder::new().create(&unique) {
10071004
Ok(()) => return Ok(unique),
10081005
Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => continue,
10091006
Err(err) => return Err(err),
@@ -1953,146 +1950,15 @@ impl Stream {
19531950
.truncate(true)
19541951
.open(part_path)
19551952
.map_err(|_| StagingError::Create)?;
1956-
// does pruning help with query?
1957-
let sort_for_metric_pruning = self.is_otel_metrics();
1958-
let time_partition_field = time_partition.map_or_else(
1959-
|| DEFAULT_TIMESTAMP_KEY.to_string(),
1960-
|s| s.as_str().to_string(),
1961-
);
1962-
1963-
if sort_for_metric_pruning {
1964-
let conversion_started = Instant::now();
1965-
let source_arrow_files = record_reader.readable_files().len();
1966-
let source_arrow_bytes = record_reader
1967-
.readable_files()
1968-
.iter()
1969-
.filter_map(|path| path.metadata().ok().map(|metadata| metadata.len()))
1970-
.sum::<u64>();
1971-
let mut timings = MetricParquetPhaseTimings::default();
1972-
let arrow_writer =
1973-
ArrowWriter::try_new(&mut part_file, schema.clone(), Some(props.clone()))?;
1974-
let (mut writer, row_group_factory) = arrow_writer.into_serialized_writer()?;
1975-
let row_group_factory = Arc::new(row_group_factory);
1976-
let column_paths = Arc::new(
1977-
writer
1978-
.schema_descr()
1979-
.columns()
1980-
.iter()
1981-
.map(|column| column.path().string())
1982-
.collect::<Vec<_>>(),
1983-
);
1984-
let target = self.options.row_group_size;
1985-
if target == 0 {
1986-
return Err(ParquetError::General(
1987-
"Parquet row-group target must be greater than zero".to_string(),
1988-
)
1989-
.into());
1990-
}
1991-
1992-
// Producer reads, decodes, concatenates, and sorts the next exact
1993-
// row group while this thread encodes the current row group. The
1994-
// bounded channel limits prepared batches retained in memory.
1995-
let prepared_row_groups = Self::spawn_metric_row_group_pipeline(
1953+
if self.is_otel_metrics() {
1954+
self.write_metric_parquet_part_file(
1955+
part_path,
1956+
&mut part_file,
19961957
record_reader,
1997-
schema.clone(),
1998-
time_partition.cloned(),
1999-
time_partition_field,
2000-
target,
2001-
);
2002-
let mut preparation_complete = false;
2003-
let mut next_row_group_index = writer.flushed_row_groups().len();
2004-
let mut pending_encodings = VecDeque::with_capacity(*METRIC_ROW_GROUP_ENCODE_IN_FLIGHT);
2005-
while !preparation_complete || !pending_encodings.is_empty() {
2006-
if !preparation_complete
2007-
&& pending_encodings.len() < *METRIC_ROW_GROUP_ENCODE_IN_FLIGHT
2008-
{
2009-
let wait_started = Instant::now();
2010-
let message = prepared_row_groups.recv().map_err(|err| {
2011-
StagingError::ObjectStorage(std::io::Error::other(format!(
2012-
"Metric row-group pipeline worker failed: {err}"
2013-
)))
2014-
})?;
2015-
let prepare_wait_duration = wait_started.elapsed();
2016-
match message {
2017-
MetricRowGroupPipelineMessage::Prepared(prepared) => {
2018-
if prepared.batch.num_rows() > target {
2019-
return Err(ParquetError::General(format!(
2020-
"Prepared metric row group has {} rows, exceeding target {target}",
2021-
prepared.batch.num_rows()
2022-
))
2023-
.into());
2024-
}
2025-
let preparation = MetricRowGroupPreparationTimings {
2026-
arrow_read_decode_duration: prepared.arrow_read_decode_duration,
2027-
concat_duration: prepared.concat_duration,
2028-
sort_duration: prepared.sort_duration,
2029-
prepare_wait_duration,
2030-
record_batches: prepared.record_batches,
2031-
};
2032-
let encode = Self::spawn_metric_row_group_encode(
2033-
Arc::clone(&row_group_factory),
2034-
Arc::clone(&column_paths),
2035-
prepared.batch,
2036-
next_row_group_index,
2037-
);
2038-
pending_encodings.push_back((encode, preparation));
2039-
next_row_group_index += 1;
2040-
continue;
2041-
}
2042-
MetricRowGroupPipelineMessage::Failed(err) => return Err(err),
2043-
MetricRowGroupPipelineMessage::Complete => preparation_complete = true,
2044-
}
2045-
}
2046-
2047-
if let Some((encoded, preparation)) = pending_encodings.pop_front() {
2048-
let encoded = encoded.recv().map_err(|err| {
2049-
StagingError::ObjectStorage(std::io::Error::other(format!(
2050-
"Metric row-group encoding worker failed: {err}"
2051-
)))
2052-
})??;
2053-
Self::write_encoded_metric_row_group(
2054-
&self.stream_name,
2055-
part_path,
2056-
&mut writer,
2057-
encoded,
2058-
preparation,
2059-
&mut timings,
2060-
)?;
2061-
}
2062-
}
2063-
let close_started = Instant::now();
2064-
writer.close()?;
2065-
timings.parquet_close_duration = close_started.elapsed();
2066-
let total_duration = conversion_started.elapsed();
2067-
let output_bytes = part_file
2068-
.metadata()
2069-
.map(|metadata| metadata.len())
2070-
.unwrap_or_default();
2071-
info!(
2072-
stream_name = %self.stream_name,
2073-
parquet_path = %part_path.display(),
2074-
source_arrow_files,
2075-
source_arrow_bytes,
2076-
output_bytes,
2077-
record_batches = timings.record_batches,
2078-
rows = timings.rows,
2079-
row_groups = timings.row_groups,
2080-
arrow_readers_per_file = *METRIC_ARROW_READERS_PER_FILE,
2081-
arrow_readers_in_flight = *METRIC_ARROW_READERS_IN_FLIGHT,
2082-
row_group_prepare_in_flight = *METRIC_ROW_GROUP_PREP_IN_FLIGHT,
2083-
row_group_encode_in_flight = *METRIC_ROW_GROUP_ENCODE_IN_FLIGHT,
2084-
columns_per_rayon_task = *PARQUET_COLUMNS_PER_RAYON_TASK,
2085-
total_ms = total_duration.as_secs_f64() * 1000.0,
2086-
arrow_read_decode_ms = timings.arrow_read_decode_duration.as_secs_f64() * 1000.0,
2087-
concat_ms = timings.concat_duration.as_secs_f64() * 1000.0,
2088-
sort_ms = timings.sort_duration.as_secs_f64() * 1000.0,
2089-
prepare_wait_ms = timings.prepare_wait_duration.as_secs_f64() * 1000.0,
2090-
parquet_setup_ms = timings.parquet_setup_duration.as_secs_f64() * 1000.0,
2091-
parquet_encode_ms = timings.parquet_encode_duration.as_secs_f64() * 1000.0,
2092-
parquet_append_ms = timings.parquet_append_duration.as_secs_f64() * 1000.0,
2093-
parquet_close_ms = timings.parquet_close_duration.as_secs_f64() * 1000.0,
2094-
"Metric Arrow-to-Parquet phase timings"
2095-
);
1958+
schema,
1959+
props,
1960+
time_partition,
1961+
)?;
20961962
} else {
20971963
let mut writer =
20981964
ArrowWriter::try_new(&mut part_file, schema.clone(), Some(props.clone()))?;
@@ -2116,6 +1982,155 @@ impl Stream {
21161982
Ok(true)
21171983
}
21181984

1985+
/// Writes a metric conversion group with pipelined row-group preparation
1986+
/// and encoding while preserving row-group output order.
1987+
fn write_metric_parquet_part_file(
1988+
&self,
1989+
part_path: &Path,
1990+
part_file: &mut File,
1991+
record_reader: ConversionRecordReader,
1992+
schema: &Arc<Schema>,
1993+
props: &WriterProperties,
1994+
time_partition: Option<&String>,
1995+
) -> Result<(), StagingError> {
1996+
let conversion_started = Instant::now();
1997+
let source_arrow_files = record_reader.readable_files().len();
1998+
let source_arrow_bytes = record_reader
1999+
.readable_files()
2000+
.iter()
2001+
.filter_map(|path| path.metadata().ok().map(|metadata| metadata.len()))
2002+
.sum::<u64>();
2003+
let mut timings = MetricParquetPhaseTimings::default();
2004+
let arrow_writer =
2005+
ArrowWriter::try_new(&mut *part_file, schema.clone(), Some(props.clone()))?;
2006+
let (mut writer, row_group_factory) = arrow_writer.into_serialized_writer()?;
2007+
let row_group_factory = Arc::new(row_group_factory);
2008+
let column_paths = Arc::new(
2009+
writer
2010+
.schema_descr()
2011+
.columns()
2012+
.iter()
2013+
.map(|column| column.path().string())
2014+
.collect::<Vec<_>>(),
2015+
);
2016+
let target = self.options.row_group_size;
2017+
if target == 0 {
2018+
return Err(ParquetError::General(
2019+
"Parquet row-group target must be greater than zero".to_string(),
2020+
)
2021+
.into());
2022+
}
2023+
2024+
let time_partition_field = time_partition.map_or_else(
2025+
|| DEFAULT_TIMESTAMP_KEY.to_string(),
2026+
|value| value.as_str().to_string(),
2027+
);
2028+
// Producer reads, decodes, concatenates, and sorts the next exact row
2029+
// group while this thread encodes the current row group.
2030+
let prepared_row_groups = Self::spawn_metric_row_group_pipeline(
2031+
record_reader,
2032+
schema.clone(),
2033+
time_partition.cloned(),
2034+
time_partition_field,
2035+
target,
2036+
);
2037+
let mut preparation_complete = false;
2038+
let mut next_row_group_index = writer.flushed_row_groups().len();
2039+
let mut pending_encodings = VecDeque::with_capacity(*METRIC_ROW_GROUP_ENCODE_IN_FLIGHT);
2040+
while !preparation_complete || !pending_encodings.is_empty() {
2041+
if !preparation_complete && pending_encodings.len() < *METRIC_ROW_GROUP_ENCODE_IN_FLIGHT
2042+
{
2043+
let wait_started = Instant::now();
2044+
let message = prepared_row_groups.recv().map_err(|err| {
2045+
StagingError::ObjectStorage(std::io::Error::other(format!(
2046+
"Metric row-group pipeline worker failed: {err}"
2047+
)))
2048+
})?;
2049+
let prepare_wait_duration = wait_started.elapsed();
2050+
match message {
2051+
MetricRowGroupPipelineMessage::Prepared(prepared) => {
2052+
if prepared.batch.num_rows() > target {
2053+
return Err(ParquetError::General(format!(
2054+
"Prepared metric row group has {} rows, exceeding target {target}",
2055+
prepared.batch.num_rows()
2056+
))
2057+
.into());
2058+
}
2059+
let preparation = MetricRowGroupPreparationTimings {
2060+
arrow_read_decode_duration: prepared.arrow_read_decode_duration,
2061+
concat_duration: prepared.concat_duration,
2062+
sort_duration: prepared.sort_duration,
2063+
prepare_wait_duration,
2064+
record_batches: prepared.record_batches,
2065+
};
2066+
let encode = Self::spawn_metric_row_group_encode(
2067+
Arc::clone(&row_group_factory),
2068+
Arc::clone(&column_paths),
2069+
prepared.batch,
2070+
next_row_group_index,
2071+
);
2072+
pending_encodings.push_back((encode, preparation));
2073+
next_row_group_index += 1;
2074+
continue;
2075+
}
2076+
MetricRowGroupPipelineMessage::Failed(err) => return Err(err),
2077+
MetricRowGroupPipelineMessage::Complete => preparation_complete = true,
2078+
}
2079+
}
2080+
2081+
if let Some((encoded, preparation)) = pending_encodings.pop_front() {
2082+
let encoded = encoded.recv().map_err(|err| {
2083+
StagingError::ObjectStorage(std::io::Error::other(format!(
2084+
"Metric row-group encoding worker failed: {err}"
2085+
)))
2086+
})??;
2087+
Self::write_encoded_metric_row_group(
2088+
&self.stream_name,
2089+
part_path,
2090+
&mut writer,
2091+
encoded,
2092+
preparation,
2093+
&mut timings,
2094+
)?;
2095+
}
2096+
}
2097+
2098+
let close_started = Instant::now();
2099+
writer.close()?;
2100+
timings.parquet_close_duration = close_started.elapsed();
2101+
let total_duration = conversion_started.elapsed();
2102+
let output_bytes = part_file
2103+
.metadata()
2104+
.map(|metadata| metadata.len())
2105+
.unwrap_or_default();
2106+
info!(
2107+
stream_name = %self.stream_name,
2108+
parquet_path = %part_path.display(),
2109+
source_arrow_files,
2110+
source_arrow_bytes,
2111+
output_bytes,
2112+
record_batches = timings.record_batches,
2113+
rows = timings.rows,
2114+
row_groups = timings.row_groups,
2115+
arrow_readers_per_file = *METRIC_ARROW_READERS_PER_FILE,
2116+
arrow_readers_in_flight = *METRIC_ARROW_READERS_IN_FLIGHT,
2117+
row_group_prepare_in_flight = *METRIC_ROW_GROUP_PREP_IN_FLIGHT,
2118+
row_group_encode_in_flight = *METRIC_ROW_GROUP_ENCODE_IN_FLIGHT,
2119+
columns_per_rayon_task = *PARQUET_COLUMNS_PER_RAYON_TASK,
2120+
total_ms = total_duration.as_secs_f64() * 1000.0,
2121+
arrow_read_decode_ms = timings.arrow_read_decode_duration.as_secs_f64() * 1000.0,
2122+
concat_ms = timings.concat_duration.as_secs_f64() * 1000.0,
2123+
sort_ms = timings.sort_duration.as_secs_f64() * 1000.0,
2124+
prepare_wait_ms = timings.prepare_wait_duration.as_secs_f64() * 1000.0,
2125+
parquet_setup_ms = timings.parquet_setup_duration.as_secs_f64() * 1000.0,
2126+
parquet_encode_ms = timings.parquet_encode_duration.as_secs_f64() * 1000.0,
2127+
parquet_append_ms = timings.parquet_append_duration.as_secs_f64() * 1000.0,
2128+
parquet_close_ms = timings.parquet_close_duration.as_secs_f64() * 1000.0,
2129+
"Metric Arrow-to-Parquet phase timings"
2130+
);
2131+
Ok(())
2132+
}
2133+
21192134
/// Removes a temporary Parquet output left by a failed conversion.
21202135
fn remove_partial_parquet_file(part_path: &Path) {
21212136
match remove_file(part_path) {

0 commit comments

Comments
 (0)