charlesdong1991 commented on code in PR #578: URL: https://github.com/apache/fluss-rust/pull/578#discussion_r3328533724
########## crates/fluss/src/row/column_vector.rs: ########## @@ -0,0 +1,650 @@ +// 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. + +//! Typed Arrow column vectors built once per `RecordBatch`. Row accessors +//! match on the variant — no per-call downcasting. Nested `ARRAY<T>` / +//! `MAP<K,V>` / `ROW` columns own their element vectors recursively. + +use crate::error::Error::IllegalArgument; +use crate::error::{Error, Result}; +use crate::metadata::{DataType, RowType}; +use crate::row::Decimal; +use crate::row::datum::{Date, Time, TimestampLtz, TimestampNtz}; +use arrow::array::{ + Array, BinaryArray, BooleanArray, Date32Array, Decimal128Array, FixedSizeBinaryArray, + Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, ListArray, MapArray, + RecordBatch, StringArray, StructArray, Time32MillisecondArray, Time32SecondArray, + Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray, + TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, +}; +use arrow::datatypes::{DataType as ArrowDataType, TimeUnit}; +use std::sync::Arc; + +/// One typed Arrow column. +#[derive(Debug, Clone)] +pub(crate) enum TypedColumn { + Boolean(BooleanArray), + TinyInt(Int8Array), + SmallInt(Int16Array), + Int(Int32Array), + BigInt(Int64Array), + Float(Float32Array), + Double(Float64Array), + /// Covers both `STRING` and `CHAR(n)` — both serialize as Arrow `Utf8`. + String(StringArray), + Binary(FixedSizeBinaryArray), + Bytes(BinaryArray), + /// Companion `i64` is the Arrow scale, extracted once at build. + Decimal(Decimal128Array, i64), + Date(Date32Array), + Time(TimeColumn), + /// `TIMESTAMP` and `TIMESTAMP_LTZ` share the Arrow representation; the + /// distinction is purely semantic. + Timestamp(TimestampColumn), + TimestampLtz(TimestampColumn), + Array(ListArray, Box<TypedColumn>), + Map(MapArray, Box<TypedColumn>, Box<TypedColumn>), + Row(StructArray, Arc<TypedBatch>), +} + +/// Arrow time-of-day in whichever unit the writer chose. Read accessors +/// down-convert to milliseconds since midnight per Fluss's `get_time` contract. +#[derive(Debug, Clone)] +pub(crate) enum TimeColumn { + Second(Time32SecondArray), + Millisecond(Time32MillisecondArray), + Microsecond(Time64MicrosecondArray), + Nanosecond(Time64NanosecondArray), +} + +#[derive(Debug, Clone)] +pub(crate) enum TimestampColumn { + Second(TimestampSecondArray), + Millisecond(TimestampMillisecondArray), + Microsecond(TimestampMicrosecondArray), + Nanosecond(TimestampNanosecondArray), +} + +/// Typed Arrow columns sharing a row count. Built once per `ArrowReader` and +/// shared via `Arc` across rows. `record_batch` is `Some` only for the outer +/// batch (the FFI/pyarrow path hands it back); nested batches built from a +/// `StructArray` carry `None`. +#[derive(Debug)] +pub(crate) struct TypedBatch { + pub(crate) columns: Vec<TypedColumn>, + pub(crate) num_rows: usize, + pub(crate) record_batch: Option<RecordBatch>, +} + +impl TypedBatch { + /// Builds a typed batch from `batch` + `row_type`, recursing into nested + /// `ARRAY` / `MAP` / `ROW` columns. + pub(crate) fn build(batch: &RecordBatch, row_type: &RowType) -> Result<Self> { + let fields = row_type.fields(); + if batch.num_columns() != fields.len() { + return Err(IllegalArgument { + message: format!( + "RecordBatch has {} columns but RowType has {} fields", + batch.num_columns(), + fields.len() + ), + }); + } + let mut columns = Vec::with_capacity(fields.len()); + for (i, field) in fields.iter().enumerate() { + columns.push(TypedColumn::build( + batch.column(i).as_ref(), + &field.data_type, + )?); + } + Ok(TypedBatch { + columns, + num_rows: batch.num_rows(), + record_batch: Some(batch.clone()), + }) + } + + fn from_struct(struct_arr: &StructArray, row_type: &RowType) -> Result<Self> { + let fields = row_type.fields(); + if struct_arr.num_columns() != fields.len() { + return Err(IllegalArgument { + message: format!( + "Arrow StructArray has {} columns but RowType has {} fields", + struct_arr.num_columns(), + fields.len() + ), + }); + } + let mut columns = Vec::with_capacity(fields.len()); + for (i, field) in fields.iter().enumerate() { + columns.push(TypedColumn::build( + struct_arr.column(i).as_ref(), + &field.data_type, + )?); + } + Ok(TypedBatch { + columns, + num_rows: struct_arr.len(), + record_batch: None, + }) + } +} + +impl TypedColumn { + /// Recursively types `array` against the Fluss `fluss_type` it represents. + pub(crate) fn build(array: &dyn Array, fluss_type: &DataType) -> Result<Self> { + match fluss_type { + DataType::Boolean(_) => Ok(Self::Boolean(downcast(array, "BooleanArray")?)), + DataType::TinyInt(_) => Ok(Self::TinyInt(downcast(array, "Int8Array")?)), + DataType::SmallInt(_) => Ok(Self::SmallInt(downcast(array, "Int16Array")?)), + DataType::Int(_) => Ok(Self::Int(downcast(array, "Int32Array")?)), + DataType::BigInt(_) => Ok(Self::BigInt(downcast(array, "Int64Array")?)), + DataType::Float(_) => Ok(Self::Float(downcast(array, "Float32Array")?)), + DataType::Double(_) => Ok(Self::Double(downcast(array, "Float64Array")?)), + DataType::Char(_) | DataType::String(_) => { + Ok(Self::String(downcast(array, "StringArray")?)) + } + DataType::Binary(_) => Ok(Self::Binary(downcast(array, "FixedSizeBinaryArray")?)), + DataType::Bytes(_) => Ok(Self::Bytes(downcast(array, "BinaryArray")?)), + DataType::Decimal(_) => { + let arr: Decimal128Array = downcast(array, "Decimal128Array")?; + let arrow_scale = match arr.data_type() { + ArrowDataType::Decimal128(_, s) => *s as i64, + other => { + return Err(IllegalArgument { + message: format!( + "Decimal128Array carries unexpected data type: {other:?}" + ), + }); + } + }; + Ok(Self::Decimal(arr, arrow_scale)) + } + DataType::Date(_) => Ok(Self::Date(downcast(array, "Date32Array")?)), + DataType::Time(_) => Ok(Self::Time(TimeColumn::build(array)?)), + DataType::Timestamp(_) => Ok(Self::Timestamp(TimestampColumn::build(array)?)), + DataType::TimestampLTz(_) => Ok(Self::TimestampLtz(TimestampColumn::build(array)?)), + DataType::Array(at) => { + let list_arr: ListArray = downcast(array, "ListArray")?; + let element = Self::build(list_arr.values().as_ref(), at.get_element_type())?; + Ok(Self::Array(list_arr, Box::new(element))) + } + DataType::Map(mt) => { + let map_arr: MapArray = downcast(array, "MapArray")?; + let key_typed = Self::build(map_arr.keys().as_ref(), mt.key_type())?; + let value_typed = Self::build(map_arr.values().as_ref(), mt.value_type())?; + Ok(Self::Map( + map_arr, + Box::new(key_typed), + Box::new(value_typed), + )) + } + DataType::Row(rt) => { + let struct_arr: StructArray = downcast(array, "StructArray")?; + let inner = TypedBatch::from_struct(&struct_arr, rt)?; + Ok(Self::Row(struct_arr, Arc::new(inner))) + } + } + } + + /// Returns the underlying Arrow array for length / nullability access. + fn outer_array(&self) -> &dyn Array { + match self { + Self::Boolean(a) => a, + Self::TinyInt(a) => a, + Self::SmallInt(a) => a, + Self::Int(a) => a, + Self::BigInt(a) => a, + Self::Float(a) => a, + Self::Double(a) => a, + Self::String(a) => a, + Self::Binary(a) => a, + Self::Bytes(a) => a, + Self::Decimal(a, _) => a, + Self::Date(a) => a, + Self::Time(t) => t.as_array(), + Self::Timestamp(t) | Self::TimestampLtz(t) => t.as_array(), + Self::Array(a, _) => a, + Self::Map(a, _, _) => a, + Self::Row(a, _) => a, + } + } + + pub(crate) fn is_null_at(&self, row_id: usize) -> bool { + self.outer_array().is_null(row_id) + } + + pub(crate) fn get_boolean(&self, row_id: usize) -> Result<bool> { Review Comment: i wonder now since all is centralized in TypedColumn, probably it is cheap to bound check for top level getters? ########## crates/fluss/src/row/columnar.rs: ########## @@ -0,0 +1,364 @@ +// 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. + +//! Zero-copy slice views over a typed Arrow batch. An `ARRAY` cell is +//! `(element column, start, length)`; a `MAP` is two parallel slices over +//! the typed key/value vectors. + +use crate::error::Error::IllegalArgument; +use crate::error::Result; +use crate::row::Decimal; +use crate::row::column::ColumnarRow; +use crate::row::column_vector::TypedColumn; +use crate::row::datum::{Date, Time, TimestampLtz, TimestampNtz}; +use crate::row::view::{ArrayView, MapView, RowView}; +use crate::row::{DataGetters, InternalArray, InternalMap}; +use std::sync::Arc; + +/// A slice of an element `TypedColumn` representing one `ARRAY` cell. +#[derive(Debug, Clone, Copy)] +pub struct ColumnarArray<'a> { + element: &'a TypedColumn, + start: usize, + length: usize, +} + +impl<'a> ColumnarArray<'a> { + pub(crate) fn new(element: &'a TypedColumn, start: usize, length: usize) -> Self { + Self { + element, + start, + length, + } + } + + /// The element-level `TypedColumn` this slice borrows from. + pub(crate) fn element(&self) -> &'a TypedColumn { + self.element + } + + fn absolute_index(&self, pos: usize) -> Result<usize> { + if pos >= self.length { + return Err(IllegalArgument { + message: format!( + "columnar array index out of bounds: pos={pos}, length={}", + self.length + ), + }); + } + Ok(self.start + pos) + } +} + +impl InternalArray for ColumnarArray<'_> { + fn size(&self) -> usize { + self.length + } +} + +impl DataGetters for ColumnarArray<'_> { + fn is_null_at(&self, pos: usize) -> Result<bool> { + let abs = self.absolute_index(pos)?; + Ok(self.element.is_null_at(abs)) + } + + fn get_boolean(&self, pos: usize) -> Result<bool> { + self.element.get_boolean(self.absolute_index(pos)?) + } + fn get_byte(&self, pos: usize) -> Result<i8> { + self.element.get_byte(self.absolute_index(pos)?) + } + fn get_short(&self, pos: usize) -> Result<i16> { + self.element.get_short(self.absolute_index(pos)?) + } + fn get_int(&self, pos: usize) -> Result<i32> { + self.element.get_int(self.absolute_index(pos)?) + } + fn get_long(&self, pos: usize) -> Result<i64> { + self.element.get_long(self.absolute_index(pos)?) + } + fn get_float(&self, pos: usize) -> Result<f32> { + self.element.get_float(self.absolute_index(pos)?) + } + fn get_double(&self, pos: usize) -> Result<f64> { + self.element.get_double(self.absolute_index(pos)?) + } + + fn get_char(&self, pos: usize, _length: usize) -> Result<&str> { + self.element.get_char(self.absolute_index(pos)?) + } + fn get_string(&self, pos: usize) -> Result<&str> { + self.element.get_string(self.absolute_index(pos)?) + } + + fn get_decimal(&self, pos: usize, precision: usize, scale: usize) -> Result<Decimal> { + self.element + .get_decimal(self.absolute_index(pos)?, precision as u32, scale as u32) + } + + fn get_date(&self, pos: usize) -> Result<Date> { + self.element.get_date(self.absolute_index(pos)?) + } + fn get_time(&self, pos: usize) -> Result<Time> { + self.element.get_time(self.absolute_index(pos)?) + } + fn get_timestamp_ntz(&self, pos: usize, _precision: u32) -> Result<TimestampNtz> { + self.element.get_timestamp_ntz(self.absolute_index(pos)?) + } + fn get_timestamp_ltz(&self, pos: usize, _precision: u32) -> Result<TimestampLtz> { + self.element.get_timestamp_ltz(self.absolute_index(pos)?) + } + + fn get_binary(&self, pos: usize, _length: usize) -> Result<&[u8]> { + self.element.get_binary(self.absolute_index(pos)?) + } + fn get_bytes(&self, pos: usize) -> Result<&[u8]> { + self.element.get_bytes(self.absolute_index(pos)?) + } + + fn get_array(&self, pos: usize) -> Result<ArrayView<'_>> { + let abs = self.absolute_index(pos)?; + match self.element { + TypedColumn::Array(list_arr, inner) => { + let start = list_arr.value_offsets()[abs] as usize; + let end = list_arr.value_offsets()[abs + 1] as usize; + Ok(ArrayView::Columnar(ColumnarArray::new( + inner.as_ref(), + start, + end - start, + ))) + } + other => Err(IllegalArgument { + message: format!("expected ARRAY element at position {pos}, got {other:?}"), + }), + } + } + + fn get_map(&self, pos: usize) -> Result<MapView<'_>> { Review Comment: similar here ########## crates/fluss/src/row/columnar.rs: ########## @@ -0,0 +1,364 @@ +// 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. + +//! Zero-copy slice views over a typed Arrow batch. An `ARRAY` cell is +//! `(element column, start, length)`; a `MAP` is two parallel slices over +//! the typed key/value vectors. + +use crate::error::Error::IllegalArgument; +use crate::error::Result; +use crate::row::Decimal; +use crate::row::column::ColumnarRow; +use crate::row::column_vector::TypedColumn; +use crate::row::datum::{Date, Time, TimestampLtz, TimestampNtz}; +use crate::row::view::{ArrayView, MapView, RowView}; +use crate::row::{DataGetters, InternalArray, InternalMap}; +use std::sync::Arc; + +/// A slice of an element `TypedColumn` representing one `ARRAY` cell. +#[derive(Debug, Clone, Copy)] +pub struct ColumnarArray<'a> { + element: &'a TypedColumn, + start: usize, + length: usize, +} + +impl<'a> ColumnarArray<'a> { + pub(crate) fn new(element: &'a TypedColumn, start: usize, length: usize) -> Self { + Self { + element, + start, + length, + } + } + + /// The element-level `TypedColumn` this slice borrows from. + pub(crate) fn element(&self) -> &'a TypedColumn { + self.element + } + + fn absolute_index(&self, pos: usize) -> Result<usize> { + if pos >= self.length { + return Err(IllegalArgument { + message: format!( + "columnar array index out of bounds: pos={pos}, length={}", + self.length + ), + }); + } + Ok(self.start + pos) + } +} + +impl InternalArray for ColumnarArray<'_> { + fn size(&self) -> usize { + self.length + } +} + +impl DataGetters for ColumnarArray<'_> { + fn is_null_at(&self, pos: usize) -> Result<bool> { + let abs = self.absolute_index(pos)?; + Ok(self.element.is_null_at(abs)) + } + + fn get_boolean(&self, pos: usize) -> Result<bool> { + self.element.get_boolean(self.absolute_index(pos)?) + } + fn get_byte(&self, pos: usize) -> Result<i8> { + self.element.get_byte(self.absolute_index(pos)?) + } + fn get_short(&self, pos: usize) -> Result<i16> { + self.element.get_short(self.absolute_index(pos)?) + } + fn get_int(&self, pos: usize) -> Result<i32> { + self.element.get_int(self.absolute_index(pos)?) + } + fn get_long(&self, pos: usize) -> Result<i64> { + self.element.get_long(self.absolute_index(pos)?) + } + fn get_float(&self, pos: usize) -> Result<f32> { + self.element.get_float(self.absolute_index(pos)?) + } + fn get_double(&self, pos: usize) -> Result<f64> { + self.element.get_double(self.absolute_index(pos)?) + } + + fn get_char(&self, pos: usize, _length: usize) -> Result<&str> { + self.element.get_char(self.absolute_index(pos)?) + } + fn get_string(&self, pos: usize) -> Result<&str> { + self.element.get_string(self.absolute_index(pos)?) + } + + fn get_decimal(&self, pos: usize, precision: usize, scale: usize) -> Result<Decimal> { + self.element + .get_decimal(self.absolute_index(pos)?, precision as u32, scale as u32) + } + + fn get_date(&self, pos: usize) -> Result<Date> { + self.element.get_date(self.absolute_index(pos)?) + } + fn get_time(&self, pos: usize) -> Result<Time> { + self.element.get_time(self.absolute_index(pos)?) + } + fn get_timestamp_ntz(&self, pos: usize, _precision: u32) -> Result<TimestampNtz> { + self.element.get_timestamp_ntz(self.absolute_index(pos)?) + } + fn get_timestamp_ltz(&self, pos: usize, _precision: u32) -> Result<TimestampLtz> { + self.element.get_timestamp_ltz(self.absolute_index(pos)?) + } + + fn get_binary(&self, pos: usize, _length: usize) -> Result<&[u8]> { + self.element.get_binary(self.absolute_index(pos)?) + } + fn get_bytes(&self, pos: usize) -> Result<&[u8]> { + self.element.get_bytes(self.absolute_index(pos)?) + } + + fn get_array(&self, pos: usize) -> Result<ArrayView<'_>> { Review Comment: this is public right? should we update api-reference since the returned type has changed ########## crates/fluss/src/record/arrow.rs: ########## @@ -1667,52 +1668,33 @@ impl Iterator for ArrowLogRecordIterator { } pub struct ArrowReader { - record_batch: Arc<RecordBatch>, - row_type: Arc<RowType>, - fluss_row_type: Option<Arc<RowType>>, - row_column_indices: Arc<[usize]>, + batch: Arc<TypedBatch>, } impl ArrowReader { pub fn new(record_batch: Arc<RecordBatch>, row_type: Arc<RowType>) -> Self { - let row_column_indices = arrow_row_column_indices(&record_batch); - ArrowReader { - record_batch, - row_type, - fluss_row_type: None, - row_column_indices, - } + Self::new_with_fluss_row_type(record_batch, row_type, None) } pub fn new_with_fluss_row_type( record_batch: Arc<RecordBatch>, row_type: Arc<RowType>, fluss_row_type: Option<Arc<RowType>>, ) -> Self { - let row_column_indices = match &fluss_row_type { - Some(rt) => fluss_row_column_indices(rt), - None => arrow_row_column_indices(&record_batch), - }; + let schema = fluss_row_type.as_deref().unwrap_or(&row_type); + let typed = TypedBatch::build(&record_batch, schema) Review Comment: so now build is fallible, and a malformed batch panics the scan instead of surfacing error? -- 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]
