feat: support asof join

src/query/ast/src/ast/format/syntax/query.rs

@@ -437,6 +437,9 @@ pub(crate) fn pretty_table(table: TableReference) -> RcDoc<'static> {
                 JoinOperator::RightAnti => RcDoc::text("RIGHT ANTI JOIN"),
                 JoinOperator::LeftSemi => RcDoc::text("LEFT SEMI JOIN"),
                 JoinOperator::RightSemi => RcDoc::text("RIGHT SEMI JOIN"),
+                JoinOperator::Asof => RcDoc::text("ASOF JOIN"),
+                JoinOperator::LeftAsof => RcDoc::text("ASOF LEFT JOIN"),
+                JoinOperator::RightAsof => RcDoc::text("ASOF RIGHT JOIN"),
             })
             .append(RcDoc::space().append(pretty_table(*join.right)))
             .append(match &join.condition {

src/query/ast/src/ast/query.rs

@@ -862,6 +862,15 @@ impl Display for TableReference {
                     JoinOperator::CrossJoin => {
                         write!(f, " CROSS JOIN")?;
                     }
+                    JoinOperator::Asof => {
+                        write!(f, " ASOF JOIN")?;
+                    }
+                    JoinOperator::LeftAsof => {
+                        write!(f, " ASOF LEFT JOIN")?;
+                    }
+                    JoinOperator::RightAsof => {
+                        write!(f, " ASOF RIGHT JOIN")?;
+                    }
                 }
                 write!(f, " {}", join.right)?;
                 match &join.condition {
@@ -934,6 +943,10 @@ pub enum JoinOperator {
     RightAnti,
     // CrossJoin can only work with `JoinCondition::None`
     CrossJoin,
+    // Asof
+    Asof,
+    LeftAsof,
+    RightAsof,
 }
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]

src/query/ast/src/parser/query.rs

@@ -626,6 +626,9 @@ pub fn join_operator(i: Input) -> IResult<JoinOperator> {
         value(JoinOperator::RightOuter, rule! { RIGHT ~ OUTER? }),
         value(JoinOperator::FullOuter, rule! { FULL ~ OUTER? }),
         value(JoinOperator::CrossJoin, rule! { CROSS }),
+        value(JoinOperator::LeftAsof, rule! { ASOF ~ LEFT }),
+        value(JoinOperator::RightAsof, rule! { ASOF ~ RIGHT }),
+        value(JoinOperator::Asof, rule! { ASOF }),
     ))(i)
 }
 

src/query/ast/src/parser/token.rs

@@ -356,6 +356,8 @@ pub enum TokenKind {
     ARRAY,
     #[token("AS", ignore(ascii_case))]
     AS,
+    #[token("ASOF", ignore(ascii_case))]
+    ASOF,
     #[token("AST", ignore(ascii_case))]
     AST,
     #[token("AT", ignore(ascii_case))]
@@ -1644,6 +1646,7 @@ impl TokenKind {
             | TokenKind::WHEN => true,
             | TokenKind::ARRAY
             | TokenKind::AS
+            | TokenKind::ASOF
             | TokenKind::BETWEEN
             | TokenKind::CREATE
             | TokenKind::ATTACH

src/query/expression/src/types.rs

@@ -62,6 +62,10 @@ pub use self::string::StringType;
 pub use self::timestamp::TimestampType;
 pub use self::variant::VariantType;
 use crate::property::Domain;
+use crate::types::date::DATE_MAX;
+use crate::types::date::DATE_MIN;
+use crate::types::timestamp::TIMESTAMP_MAX;
+use crate::types::timestamp::TIMESTAMP_MIN;
 use crate::values::Column;
 use crate::values::Scalar;
 use crate::ColumnBuilder;
@@ -174,6 +178,63 @@ impl DataType {
         }
     }
 
+    pub fn infinity(&self) -> Result<Scalar, String> {
+        match &self {
+            DataType::Timestamp => Ok(Scalar::Timestamp(TIMESTAMP_MAX)),
+            DataType::Date => Ok(Scalar::Date(DATE_MAX)),
+            DataType::Number(NumberDataType::Float32) => Ok(Scalar::Number(NumberScalar::Float32(
+                ordered_float::OrderedFloat(f32::INFINITY),
+            ))),
+            DataType::Number(NumberDataType::Int32) => {
+                Ok(Scalar::Number(NumberScalar::Int32(i32::MAX)))
+            }
+            DataType::Number(NumberDataType::Int16) => {
+                Ok(Scalar::Number(NumberScalar::Int16(i16::MAX)))
+            }
+            DataType::Number(NumberDataType::Int8) => {
+                Ok(Scalar::Number(NumberScalar::Int8(i8::MAX)))
+            }
+            DataType::Number(NumberDataType::Float64) => Ok(Scalar::Number(NumberScalar::Float64(
+                ordered_float::OrderedFloat(f64::INFINITY),
+            ))),
+            DataType::Number(NumberDataType::Int64) => {
+                Ok(Scalar::Number(NumberScalar::Int64(i64::MAX)))
+            }
+            _ => Result::Err(format!(
+                "only support numeric types and time types, but got {:?}",
+                self
+            )),
+        }
+    }
+    pub fn ninfinity(&self) -> Result<Scalar, String> {
+        match &self {
+            DataType::Timestamp => Ok(Scalar::Timestamp(TIMESTAMP_MIN)),
+            DataType::Date => Ok(Scalar::Date(DATE_MIN)),
+            DataType::Number(NumberDataType::Float32) => Ok(Scalar::Number(NumberScalar::Float32(
+                ordered_float::OrderedFloat(f32::NEG_INFINITY),
+            ))),
+            DataType::Number(NumberDataType::Int32) => {
+                Ok(Scalar::Number(NumberScalar::Int32(i32::MIN)))
+            }
+            DataType::Number(NumberDataType::Int16) => {
+                Ok(Scalar::Number(NumberScalar::Int16(i16::MIN)))
+            }
+            DataType::Number(NumberDataType::Int8) => {
+                Ok(Scalar::Number(NumberScalar::Int8(i8::MIN)))
+            }
+            DataType::Number(NumberDataType::Float64) => Ok(Scalar::Number(NumberScalar::Float64(
+                ordered_float::OrderedFloat(f64::NEG_INFINITY),
+            ))),
+            DataType::Number(NumberDataType::Int64) => {
+                Ok(Scalar::Number(NumberScalar::Int64(i64::MIN)))
+            }
+            _ => Result::Err(format!(
+                "only support numeric types and time types, but got {:?}",
+                self
+            )),
+        }
+    }
+
     pub fn is_unsigned_numeric(&self) -> bool {
         match self {
             DataType::Number(ty) => ALL_UNSIGNED_INTEGER_TYPES.contains(ty),

src/query/service/src/pipelines/processors/transforms/hash_join/common.rs

@@ -184,7 +184,7 @@ impl HashJoinState {
     }
 }
 
-pub(crate) fn wrap_true_validity(
+pub fn wrap_true_validity(
     column: &BlockEntry,
     num_rows: usize,
     true_validity: &Bitmap,

src/query/service/src/pipelines/processors/transforms/hash_join/hash_join_probe_state.rs

@@ -160,6 +160,9 @@ impl HashJoinProbeState {
     pub fn probe(&self, input: DataBlock, probe_state: &mut ProbeState) -> Result<Vec<DataBlock>> {
         match self.hash_join_state.hash_join_desc.join_type {
             JoinType::Inner
+            | JoinType::Asof
+            | JoinType::LeftAsof
+            | JoinType::RightAsof
             | JoinType::LeftSemi
             | JoinType::LeftAnti
             | JoinType::RightSemi

src/query/service/src/pipelines/processors/transforms/hash_join/mod.rs

@@ -31,6 +31,7 @@ mod transform_hash_join_probe;
 mod util;
 
 pub use build_spill::BuildSpillState;
+pub use common::wrap_true_validity;
 pub use desc::HashJoinDesc;
 pub use hash_join_build_state::HashJoinBuildState;
 pub use hash_join_probe_state::HashJoinProbeState;

src/query/service/src/pipelines/processors/transforms/range_join/ie_join_state.rs

@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use core::time::Duration;
+use std::sync::atomic::Ordering;
+
 use databend_common_arrow::arrow::bitmap::Bitmap;
 use databend_common_arrow::arrow::bitmap::MutableBitmap;
 use databend_common_catalog::table_context::TableContext;
@@ -28,6 +31,7 @@ use databend_common_expression::DataBlock;
 use databend_common_expression::DataField;
 use databend_common_expression::DataSchemaRef;
 use databend_common_expression::DataSchemaRefExt;
+use databend_common_expression::Scalar;
 use databend_common_expression::ScalarRef;
 use databend_common_expression::SortColumnDescription;
 use databend_common_expression::Value;
@@ -174,6 +178,22 @@ impl IEJoinState {
 
 impl RangeJoinState {
     pub fn ie_join(&self, task_id: usize) -> Result<Vec<DataBlock>> {
+        let partition_count = self.partition_count.load(Ordering::SeqCst) as usize;
+        if task_id < partition_count {
+            let blocks = self.inner_join(task_id);
+            self.completed_pair.fetch_add(1, Ordering::SeqCst);
+            blocks
+        } else {
+            if !self.left_match.read().is_empty() {
+                return Ok(vec![self.fill_left_outer(task_id)?]);
+            } else if !self.right_match.read().is_empty() {
+                return Ok(vec![self.fill_right_outer(task_id)?]);
+            }
+            Ok(vec![DataBlock::empty()])
+        }
+    }
+
+    pub fn inner_join(&self, task_id: usize) -> Result<Vec<DataBlock>> {
         let block_size = self.ctx.get_settings().get_max_block_size()? as usize;
         let tasks = self.tasks.read();
         let (left_idx, right_idx) = tasks[task_id];
@@ -313,6 +333,8 @@ impl RangeJoinState {
         let mut right_buffer = Vec::with_capacity(block_size);
         let mut off1;
         let mut off2 = 0;
+        let mut left_match = self.left_match.write();
+        let mut right_match = self.right_match.write();
         for (idx, p) in p_array.iter().enumerate() {
             if let ScalarRef::Number(NumberScalar::Int64(val)) =
                 unsafe { l1_index_column.index_unchecked(*p as usize) }
@@ -365,27 +387,161 @@ impl RangeJoinState {
         let left_table = self.left_table.read();
         let right_table = self.right_table.read();
         let mut indices = Vec::with_capacity(left_buffer.len());
+        let mut column_builder =
+            NumberColumnBuilder::with_capacity(&NumberDataType::UInt64, left_buffer.len());
         for res in left_buffer.iter() {
             indices.push((0u32, *res as u32, 1usize));
+            if !left_match.is_empty() {
+                column_builder.push(NumberScalar::UInt64((*res + left_offset) as u64));
+            }
         }
         let mut left_result_block =
             DataBlock::take_blocks(&left_table[left_idx..left_idx + 1], &indices, indices.len());
         indices.clear();
         for res in right_buffer.iter() {
             indices.push((0u32, *res as u32, 1usize));
+            if !right_match.is_empty() {
+                column_builder.push(NumberScalar::UInt64((*res + right_offset) as u64));
+            }
         }
         let right_result_block = DataBlock::take_blocks(
             &right_table[right_idx..right_idx + 1],
             &indices,
             indices.len(),
         );
-        // Merge left_result_block and right_result_block
         for col in right_result_block.columns() {
             left_result_block.add_column(col.clone());
         }
+        if !left_match.is_empty() || !right_match.is_empty() {
+            left_result_block.add_column(BlockEntry::new(
+                DataType::Number(NumberDataType::UInt64),
+                Value::Column(Column::Number(column_builder.build())),
+            ));
+        }
         for filter in self.other_conditions.iter() {
             left_result_block = filter_block(left_result_block, filter)?;
         }
+        if !left_match.is_empty() || !right_match.is_empty() {
+            let column = &left_result_block
+                .columns()
+                .last()
+                .unwrap()
+                .value
+                .try_downcast::<UInt64Type>()
+                .unwrap();
+            if let ValueRef::Column(col) = column.as_ref() {
+                for val in UInt64Type::iter_column(&col) {
+                    if !left_match.is_empty() {
+                        left_match.set(val as usize, true);
+                    }
+                    if !right_match.is_empty() {
+                        right_match.set(val as usize, true);
+                    }
+                }
+            }
+            left_result_block.pop_columns(1);
+        }
+        Ok(left_result_block)
+    }
+
+    pub fn fill_left_outer(&self, task_id: usize) -> Result<DataBlock> {
+        let partition_count = self.partition_count.load(Ordering::SeqCst) as usize;
+        let mut completed = self.completed_pair.load(Ordering::SeqCst) as usize;
+        while completed < partition_count {
+            std::thread::sleep(Duration::from_millis(10));
+            completed = self.completed_pair.load(Ordering::SeqCst) as usize;
+        }
+
+        let block_size = self.ctx.get_settings().get_max_block_size()? as usize;
+        let tasks = self.tasks.read();
+        let (left_idx, right_idx) = tasks[task_id];
+        let row_offset = self.row_offset.read();
+        let (left_offset, _right_offset) = row_offset[task_id];
+        let left_table = self.left_table.read();
+        let right_table = self.right_table.read();
+        let mut indices = Vec::with_capacity(block_size);
+        let left_match = self.left_match.read();
+
+        for (i, state) in left_match
+            .iter()
+            .enumerate()
+            .skip(left_offset)
+            .take(left_table[left_idx].num_rows())
+        {
+            if !(state) {
+                indices.push((0u32, (i - left_offset) as u32, 1usize));
+            }
+        }
+        if indices.is_empty() {
+            return Ok(DataBlock::empty());
+        }
+        let mut left_result_block =
+            DataBlock::take_blocks(&left_table[left_idx..left_idx + 1], &indices, indices.len());
+        let nullable_columns = right_table[right_idx]
+            .columns()
+            .iter()
+            .map(|c| BlockEntry {
+                value: Value::Scalar(Scalar::Null),
+                data_type: c.data_type.wrap_nullable(),
+            })
+            .collect::<Vec<_>>();
+        let right_result_block = DataBlock::new(nullable_columns, indices.len());
+        // Merge left_result_block and right_result_block
+        for col in right_result_block.columns() {
+            left_result_block.add_column(col.clone());
+        }
+        Ok(left_result_block)
+    }
+
+    pub fn fill_right_outer(&self, task_id: usize) -> Result<DataBlock> {
+        let partition_count = self.partition_count.load(Ordering::SeqCst) as usize;
+        let mut completed = self.completed_pair.load(Ordering::SeqCst) as usize;
+        while completed < partition_count {
+            std::thread::sleep(Duration::from_millis(10));
+            completed = self.completed_pair.load(Ordering::SeqCst) as usize;
+        }
+
+        let block_size = self.ctx.get_settings().get_max_block_size()? as usize;
+        let tasks = self.tasks.read();
+        let (left_idx, right_idx) = tasks[task_id];
+        let row_offset = self.row_offset.read();
+        let (_, right_offset) = row_offset[task_id];
+        let left_table = self.left_table.read();
+        let right_table = self.right_table.read();
+        let mut indices = Vec::with_capacity(block_size);
+        let right_match = self.right_match.read();
+
+        for (i, state) in right_match
+            .iter()
+            .enumerate()
+            .skip(right_offset)
+            .take(right_table[right_idx].num_rows())
+        {
+            if !(state) {
+                indices.push((0u32, (i - right_offset) as u32, 1usize));
+            }
+        }
+        if indices.is_empty() {
+            return Ok(DataBlock::empty());
+        }
+        let nullable_columns = left_table[left_idx]
+            .columns()
+            .iter()
+            .map(|c| BlockEntry {
+                value: Value::Scalar(Scalar::Null),
+                data_type: c.data_type.wrap_nullable(),
+            })
+            .collect::<Vec<_>>();
+        let mut left_result_block = DataBlock::new(nullable_columns, indices.len());
+        let right_result_block = DataBlock::take_blocks(
+            &right_table[right_idx..right_idx + 1],
+            &indices,
+            indices.len(),
+        );
+        // Merge left_result_block and right_result_block
+        for col in right_result_block.columns() {
+            left_result_block.add_column(col.clone());
+        }
         Ok(left_result_block)
     }
 }