jaylmiller commented on code in PR #5292:
URL: https://github.com/apache/arrow-datafusion/pull/5292#discussion_r1107743051
##########
datafusion/core/src/physical_plan/sorts/sort.rs:
##########
@@ -905,6 +1133,138 @@ async fn do_sort(
result
}
+/// manages writing potential rows to and from disk
+struct RowWriter {
+ // serializing w/ arrow ipc format for maximum code simplicity... probably
sub-optimal
+ ipc_writer: Option<IPCWriter>,
+ schema: Option<SchemaRef>,
+ num_bytes: usize,
+}
+impl RowWriter {
+ fn try_new(path: Option<impl AsRef<Path>>) -> Result<Self> {
+ use arrow::datatypes::{DataType, Field, Schema};
+ match path {
+ Some(p) => {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "bytes",
+ DataType::Binary,
+ false,
+ )]));
+ Ok(Self {
+ ipc_writer: Some(IPCWriter::new(p.as_ref(),
schema.as_ref())?),
+ schema: Some(schema),
+ num_bytes: 0,
+ })
+ }
+ None => Ok(Self {
+ ipc_writer: None,
+ schema: None,
+ num_bytes: 0,
+ }),
+ }
+ }
+ fn write(&mut self, rows: Option<RowBatch>) -> Result<()> {
+ use arrow::array::BinaryBuilder;
+ if let (Some(writer), Some(schema), Some(rows)) =
+ (self.ipc_writer.as_mut(), self.schema.as_ref(), rows)
+ {
+ let numbytes = rows
+ .iter()
+ .map(|v| {
+ let bytes: &[u8] = v.as_ref();
+ bytes.len()
+ })
+ .sum::<usize>();
+ let mut builder = BinaryBuilder::with_capacity(rows.num_rows(),
numbytes);
+ for r in rows.iter() {
+ let bytes: &[u8] = r.as_ref();
+ builder.append_value(bytes);
+ }
+ let arr = builder.finish();
+ let batch = RecordBatch::try_new(schema.clone(),
vec![Arc::new(arr)])?;
+ writer.write(&batch)?;
+ self.num_bytes += numbytes;
+ Ok(())
+ } else {
+ Ok(())
+ }
+ }
+ fn finish(&mut self) -> Result<()> {
+ self.ipc_writer
+ .as_mut()
+ .map(|v| v.finish())
+ .unwrap_or(Ok(()))
+ }
+}
+
+/// manages reading potential rows to and from disk.
+struct RowReader {
+ /// temporary file format solution is storing it w/ arrow IPC
+ reader: Option<FileReader<BufReader<File>>>,
+ row_conv: RowConverter,
+}
+impl RowReader {
+ fn try_new(
+ path: Option<impl AsRef<Path>>,
+ sort_fields: Vec<SortField>,
+ ) -> Result<Self> {
+ let row_conv = RowConverter::new(sort_fields)?;
+ if let Some(path) = path {
+ let file = BufReader::new(File::open(path)?);
+ let reader = FileReader::try_new(file, None)?;
+
+ Ok(Self {
+ reader: Some(reader),
+ row_conv,
+ })
+ } else {
+ Ok(Self {
+ reader: None,
+ row_conv,
+ })
+ }
+ }
+ fn parse_batch_helper(&self, batch: RecordBatch) ->
Result<Option<RowBatch>> {
+ use arrow::array::BinaryArray;
+ let col = batch.column(0);
+ let bincol = col.as_any().downcast_ref::<BinaryArray>().ok_or_else(|| {
+ DataFusionError::Internal(
+ "unexepected error while parsing spilled row data".to_string(),
+ )
+ })?;
+ let bytes = bincol
+ .into_iter()
+ .map(|v| {
+ v.map(bytes::Bytes::copy_from_slice).ok_or_else(|| {
+ DataFusionError::Internal(
+ "unexepected error while parsing spilled row
data".to_string(),
+ )
+ })
+ })
+ .collect::<Result<Vec<_>>>()?;
+ Ok(Some(
+ RowSelection::from_spilled(self.row_conv.parser(), bytes).into(),
+ ))
+ }
+}
Review Comment:
these 2 impl's will be swapped out with new serialization logic
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
