yordan-pavlov commented on a change in pull request #1849:
URL: https://github.com/apache/arrow-datafusion/pull/1849#discussion_r810675856
##########
File path: datafusion/src/row/writer.rs
##########
@@ -244,89 +346,245 @@ fn write_row(row: &mut RowWriter, row_idx: usize, batch:
&RecordBatch) -> usize
row.row_width
}
+#[cfg(feature = "jit")]
+// TODO: we could remove these functions wrappers below once we enabled call
struct/trait methods
+fn get_array(batch: &RecordBatch, col_idx: usize) -> &Arc<dyn Array> {
+ batch.column(col_idx)
+}
+
+#[cfg(feature = "jit")]
+fn set_null_at(row: &mut RowWriter, col_idx: usize) {
+ row.set_null_at(col_idx)
+}
+
+#[cfg(feature = "jit")]
+fn set_non_null_at(row: &mut RowWriter, col_idx: usize) {
+ row.set_non_null_at(col_idx)
+}
+
+#[cfg(feature = "jit")]
+fn is_null(col: &Arc<dyn Array>, row_idx: usize) -> bool {
+ col.is_null(row_idx)
+}
+
+#[cfg(feature = "jit")]
+fn register_write_functions(asm: &Assembler) -> Result<()> {
+ let reader_param = vec![PTR, I64, PTR];
+ reg_fn!(asm, get_array, vec![PTR, I64], Some(PTR));
+ reg_fn!(asm, set_null_at, vec![PTR, I64], None);
+ reg_fn!(asm, set_non_null_at, vec![PTR, I64], None);
+ reg_fn!(asm, is_null, vec![PTR, I64], Some(BOOL));
+ reg_fn!(asm, write_field_bool, reader_param.clone(), None);
+ reg_fn!(asm, write_field_u8, reader_param.clone(), None);
+ reg_fn!(asm, write_field_u16, reader_param.clone(), None);
+ reg_fn!(asm, write_field_u32, reader_param.clone(), None);
+ reg_fn!(asm, write_field_u64, reader_param.clone(), None);
+ reg_fn!(asm, write_field_i8, reader_param.clone(), None);
+ reg_fn!(asm, write_field_i16, reader_param.clone(), None);
+ reg_fn!(asm, write_field_i32, reader_param.clone(), None);
+ reg_fn!(asm, write_field_i64, reader_param.clone(), None);
+ reg_fn!(asm, write_field_f32, reader_param.clone(), None);
+ reg_fn!(asm, write_field_f64, reader_param.clone(), None);
+ reg_fn!(asm, write_field_date32, reader_param.clone(), None);
+ reg_fn!(asm, write_field_date64, reader_param.clone(), None);
+ reg_fn!(asm, write_field_utf8, reader_param.clone(), None);
+ reg_fn!(asm, write_field_binary, reader_param, None);
+ Ok(())
+}
+
+#[cfg(feature = "jit")]
+fn gen_write_row(
+ schema: &Arc<Schema>,
+ assembler: &Assembler,
+) -> Result<GeneratedFunction> {
+ use DataType::*;
+ let mut builder = assembler
+ .new_func_builder("write_row")
+ .param("row", PTR)
+ .param("row_idx", I64)
+ .param("batch", PTR);
+ let mut b = builder.enter_block();
+ for (i, f) in schema.fields().iter().enumerate() {
+ let dt = f.data_type();
+ let arr = format!("a{}", i);
+ b.declare_as(
+ &arr,
+ b.call("get_array", vec![b.id("batch")?, b.lit_i(i as i64)])?,
+ )?;
+ if f.is_nullable() {
+ b.if_block(
+ |c| c.call("is_null", vec![c.id(&arr)?, c.id("row_idx")?]),
+ |t| {
+ t.call_stmt("set_null_at", vec![t.id("row")?, t.lit_i(i as
i64)])?;
+ Ok(())
+ },
+ |e| {
+ e.call_stmt(
+ "set_non_null_at",
+ vec![e.id("row")?, e.lit_i(i as i64)],
+ )?;
+ let params = vec![
+ e.id("row")?,
+ e.id(&arr)?,
+ e.lit_i(i as i64),
+ e.id("row_idx")?,
+ ];
+ match dt {
+ Boolean => e.call_stmt("write_field_bool", params)?,
+ UInt8 => e.call_stmt("write_field_u8", params)?,
+ UInt16 => e.call_stmt("write_field_u16", params)?,
+ UInt32 => e.call_stmt("write_field_u32", params)?,
+ UInt64 => e.call_stmt("write_field_u64", params)?,
+ Int8 => e.call_stmt("write_field_i8", params)?,
+ Int16 => e.call_stmt("write_field_i16", params)?,
+ Int32 => e.call_stmt("write_field_i32", params)?,
+ Int64 => e.call_stmt("write_field_i64", params)?,
+ Float32 => e.call_stmt("write_field_f32", params)?,
+ Float64 => e.call_stmt("write_field_f64", params)?,
+ Date32 => e.call_stmt("write_field_date32", params)?,
+ Date64 => e.call_stmt("write_field_date64", params)?,
+ Utf8 => e.call_stmt("write_field_utf8", params)?,
+ Binary => e.call_stmt("write_field_binary", params)?,
+ _ => unimplemented!(),
+ }
+ Ok(())
+ },
+ )?;
+ } else {
+ b.call_stmt("set_non_null_at", vec![b.id("row")?, b.lit_i(i as
i64)])?;
Review comment:
this `else` block looks the same as the `|e| {` block above - should it
go in a function?
--
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]
