Multiple row-layout support, part-1: Restructure code for clearness

datafusion/core/src/row/jit/mod.rs

@@ -0,0 +1,224 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Just-In-Time(JIT) version for row reader and writers
+
+mod reader;
+mod writer;
+
+#[macro_export]
+/// register external functions to the assembler
+macro_rules! reg_fn {
+    ($ASS:ident, $FN: path, $PARAM: expr, $RET: expr) => {
+        $ASS.register_extern_fn(fn_name($FN), $FN as *const u8, $PARAM, $RET)?;
+    };
+}
+
+fn fn_name<T>(f: T) -> &'static str {
+    fn type_name_of<T>(_: T) -> &'static str {
+        std::any::type_name::<T>()
+    }
+    let name = type_name_of(f);
+
+    // Find and cut the rest of the path
+    match &name.rfind(':') {
+        Some(pos) => &name[pos + 1..name.len()],
+        None => name,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::error::Result;
+    use crate::row::jit::reader::read_as_batch_jit;
+    use crate::row::jit::writer::write_batch_unchecked_jit;
+    use arrow::record_batch::RecordBatch;
+    use arrow::{array::*, datatypes::*};
+    use datafusion_jit::api::Assembler;
+    use std::sync::Arc;
+    use DataType::*;
+
+    macro_rules! fn_test_single_type {
+        ($ARRAY: ident, $TYPE: expr, $VEC: expr) => {
+            paste::item! {
+                #[test]
+                #[allow(non_snake_case)]
+                fn [<test_single_ $TYPE _jit>]() -> Result<()> {
+                    let schema = Arc::new(Schema::new(vec![Field::new("a", $TYPE, true)]));
+                    let a = $ARRAY::from($VEC);
+                    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(a)])?;
+                    let mut vector = vec![0; 1024];
+                    let assembler = Assembler::default();
+                    let row_offsets =
+                        { write_batch_unchecked_jit(&mut vector, 0, &batch, 0, schema.clone(), &assembler)? };
+                    let output_batch = { read_as_batch_jit(&vector, schema, &row_offsets, &assembler)? };
+                    assert_eq!(batch, output_batch);
+                    Ok(())
+                }
+
+                #[test]
+                #[allow(non_snake_case)]
+                fn [<test_single_ $TYPE _jit_null_free>]() -> Result<()> {
+                    let schema = Arc::new(Schema::new(vec![Field::new("a", $TYPE, false)]));
+                    let v = $VEC.into_iter().filter(|o| o.is_some()).collect::<Vec<_>>();
+                    let a = $ARRAY::from(v);
+                    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(a)])?;
+                    let mut vector = vec![0; 1024];
+                    let assembler = Assembler::default();
+                    let row_offsets =
+                        { write_batch_unchecked_jit(&mut vector, 0, &batch, 0, schema.clone(), &assembler)? };
+                    let output_batch = { read_as_batch_jit(&vector, schema, &row_offsets, &assembler)? };
+                    assert_eq!(batch, output_batch);
+                    Ok(())
+                }
+            }
+        };
+    }
+
+    fn_test_single_type!(
+        BooleanArray,
+        Boolean,
+        vec![Some(true), Some(false), None, Some(true), None]
+    );
+
+    fn_test_single_type!(
+        Int8Array,
+        Int8,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        Int16Array,
+        Int16,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        Int32Array,
+        Int32,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        Int64Array,
+        Int64,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        UInt8Array,
+        UInt8,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        UInt16Array,
+        UInt16,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        UInt32Array,
+        UInt32,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        UInt64Array,
+        UInt64,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        Float32Array,
+        Float32,
+        vec![Some(5.0), Some(7.0), None, Some(0.0), Some(111.0)]
+    );
+
+    fn_test_single_type!(
+        Float64Array,
+        Float64,
+        vec![Some(5.0), Some(7.0), None, Some(0.0), Some(111.0)]
+    );
+
+    fn_test_single_type!(
+        Date32Array,
+        Date32,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        Date64Array,
+        Date64,
+        vec![Some(5), Some(7), None, Some(0), Some(111)]
+    );
+
+    fn_test_single_type!(
+        StringArray,
+        Utf8,
+        vec![Some("hello"), Some("world"), None, Some(""), Some("")]
+    );
+
+    #[test]
+    fn test_single_binary_jit() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", Binary, true)]));
+        let values: Vec<Option<&[u8]>> =
+            vec![Some(b"one"), Some(b"two"), None, Some(b""), Some(b"three")];
+        let a = BinaryArray::from_opt_vec(values);
+        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(a)])?;
+        let mut vector = vec![0; 8192];
+        let assembler = Assembler::default();
+        let row_offsets = {
+            write_batch_unchecked_jit(
+                &mut vector,
+                0,
+                &batch,
+                0,
+                schema.clone(),
+                &assembler,
+            )?
+        };
+        let output_batch =
+            { read_as_batch_jit(&vector, schema, &row_offsets, &assembler)? };
+        assert_eq!(batch, output_batch);
+        Ok(())
+    }
+
+    #[test]
+    fn test_single_binary_jit_null_free() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", Binary, false)]));
+        let values: Vec<&[u8]> = vec![b"one", b"two", b"", b"three"];
+        let a = BinaryArray::from_vec(values);
+        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(a)])?;
+        let mut vector = vec![0; 8192];
+        let assembler = Assembler::default();
+        let row_offsets = {
+            write_batch_unchecked_jit(
+                &mut vector,
+                0,
+                &batch,
+                0,
+                schema.clone(),
+                &assembler,
+            )?
+        };
+        let output_batch =
+            { read_as_batch_jit(&vector, schema, &row_offsets, &assembler)? };
+        assert_eq!(batch, output_batch);
+        Ok(())
+    }
+}

datafusion/core/src/row/jit/reader.rs

@@ -0,0 +1,164 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Accessing row from raw bytes with JIT
+
+use crate::error::{DataFusionError, Result};
+use crate::reg_fn;
+use crate::row::jit::fn_name;
+use crate::row::reader::RowReader;
+use crate::row::reader::*;
+use crate::row::MutableRecordBatch;
+use arrow::array::ArrayBuilder;
+use arrow::datatypes::{DataType, Schema};
+use arrow::record_batch::RecordBatch;
+use datafusion_jit::api::Assembler;
+use datafusion_jit::api::GeneratedFunction;
+use datafusion_jit::ast::{I64, PTR};
+use std::sync::Arc;
+
+/// Read `data` of raw-bytes rows starting at `offsets` out to a record batch
+
+pub fn read_as_batch_jit(
+    data: &[u8],
+    schema: Arc<Schema>,
+    offsets: &[usize],
+    assembler: &Assembler,
+) -> Result<RecordBatch> {
+    let row_num = offsets.len();
+    let mut output = MutableRecordBatch::new(row_num, schema.clone());
+    let mut row = RowReader::new(&schema);
+    register_read_functions(assembler)?;
+    let gen_func = gen_read_row(&schema, assembler)?;
+    let mut jit = assembler.create_jit();
+    let code_ptr = jit.compile(gen_func)?;
+    let code_fn = unsafe {
+        std::mem::transmute::<_, fn(&RowReader, &mut MutableRecordBatch)>(code_ptr)
+    };
+
+    for offset in offsets.iter().take(row_num) {
+        row.point_to(*offset, data);
+        code_fn(&row, &mut output);
+    }
+
+    output.output().map_err(DataFusionError::ArrowError)
+}
+
+fn get_array_mut(
+    batch: &mut MutableRecordBatch,
+    col_idx: usize,
+) -> &mut Box<dyn ArrayBuilder> {
+    let arrays: &mut [Box<dyn ArrayBuilder>] = batch.arrays.as_mut();
+    &mut arrays[col_idx]
+}
+
+fn register_read_functions(asm: &Assembler) -> Result<()> {
+    let reader_param = vec![PTR, I64, PTR];
+    reg_fn!(asm, get_array_mut, vec![PTR, I64], Some(PTR));
+    reg_fn!(asm, read_field_bool, reader_param.clone(), None);
+    reg_fn!(asm, read_field_u8, reader_param.clone(), None);
+    reg_fn!(asm, read_field_u16, reader_param.clone(), None);
+    reg_fn!(asm, read_field_u32, reader_param.clone(), None);
+    reg_fn!(asm, read_field_u64, reader_param.clone(), None);
+    reg_fn!(asm, read_field_i8, reader_param.clone(), None);
+    reg_fn!(asm, read_field_i16, reader_param.clone(), None);
+    reg_fn!(asm, read_field_i32, reader_param.clone(), None);
+    reg_fn!(asm, read_field_i64, reader_param.clone(), None);
+    reg_fn!(asm, read_field_f32, reader_param.clone(), None);
+    reg_fn!(asm, read_field_f64, reader_param.clone(), None);
+    reg_fn!(asm, read_field_date32, reader_param.clone(), None);
+    reg_fn!(asm, read_field_date64, reader_param.clone(), None);
+    reg_fn!(asm, read_field_utf8, reader_param.clone(), None);
+    reg_fn!(asm, read_field_binary, reader_param.clone(), None);
+    reg_fn!(asm, read_field_bool_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_u8_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_u16_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_u32_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_u64_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_i8_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_i16_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_i32_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_i64_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_f32_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_f64_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_date32_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_date64_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_utf8_null_free, reader_param.clone(), None);
+    reg_fn!(asm, read_field_binary_null_free, reader_param, None);
+    Ok(())
+}
+
+fn gen_read_row(
+    schema: &Arc<Schema>,
+    assembler: &Assembler,
+) -> Result<GeneratedFunction> {
+    use DataType::*;
+    let mut builder = assembler
+        .new_func_builder("read_row")
+        .param("row", PTR)
+        .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_mut", vec![b.id("batch")?, b.lit_i(i as i64)])?,
+        )?;
+        let params = vec![b.id(&arr)?, b.lit_i(i as i64), b.id("row")?];
+        if f.is_nullable() {
+            match dt {
+                Boolean => b.call_stmt("read_field_bool", params)?,
+                UInt8 => b.call_stmt("read_field_u8", params)?,
+                UInt16 => b.call_stmt("read_field_u16", params)?,
+                UInt32 => b.call_stmt("read_field_u32", params)?,
+                UInt64 => b.call_stmt("read_field_u64", params)?,
+                Int8 => b.call_stmt("read_field_i8", params)?,
+                Int16 => b.call_stmt("read_field_i16", params)?,
+                Int32 => b.call_stmt("read_field_i32", params)?,
+                Int64 => b.call_stmt("read_field_i64", params)?,
+                Float32 => b.call_stmt("read_field_f32", params)?,
+                Float64 => b.call_stmt("read_field_f64", params)?,
+                Date32 => b.call_stmt("read_field_date32", params)?,
+                Date64 => b.call_stmt("read_field_date64", params)?,
+                Utf8 => b.call_stmt("read_field_utf8", params)?,
+                Binary => b.call_stmt("read_field_binary", params)?,
+                _ => unimplemented!(),
+            }
+        } else {
+            match dt {
+                Boolean => b.call_stmt("read_field_bool_null_free", params)?,
+                UInt8 => b.call_stmt("read_field_u8_null_free", params)?,
+                UInt16 => b.call_stmt("read_field_u16_null_free", params)?,
+                UInt32 => b.call_stmt("read_field_u32_null_free", params)?,
+                UInt64 => b.call_stmt("read_field_u64_null_free", params)?,
+                Int8 => b.call_stmt("read_field_i8_null_free", params)?,
+                Int16 => b.call_stmt("read_field_i16_null_free", params)?,
+                Int32 => b.call_stmt("read_field_i32_null_free", params)?,
+                Int64 => b.call_stmt("read_field_i64_null_free", params)?,
+                Float32 => b.call_stmt("read_field_f32_null_free", params)?,
+                Float64 => b.call_stmt("read_field_f64_null_free", params)?,
+                Date32 => b.call_stmt("read_field_date32_null_free", params)?,
+                Date64 => b.call_stmt("read_field_date64_null_free", params)?,
+                Utf8 => b.call_stmt("read_field_utf8_null_free", params)?,
+                Binary => b.call_stmt("read_field_binary_null_free", params)?,
+                _ => unimplemented!(),
+            }
+        }
+    }
+    Ok(b.build())
+}