feat: Refactor join execution with enhanced hash join and schema handling, including optimizations for join conditions and index adjustments

Author: longshan.lu

qurious/src/datasource/mod.rs

@@ -1,4 +1,2 @@
-#[cfg(feature = "connectorx")]
-pub mod connectorx;
 pub mod file;
 pub mod memory;

qurious/src/execution/session.rs

@@ -10,6 +10,7 @@ use crate::common::table_relation::TableRelation;
 use crate::datasource::memory::MemoryTable;
 use crate::error::Error;
 use crate::functions::{all_builtin_functions, UserDefinedFunction};
+use crate::internal_err;
 use crate::logical::plan::{
     CreateMemoryTable, DdlStatement, DmlOperator, DmlStatement, DropTable, Filter, LogicalPlan,
 };
@@ -21,7 +22,6 @@ use crate::provider::schema::SchemaProvider;
 use crate::provider::table::TableProvider;
 use crate::utils::batch::make_count_batch;
 use crate::{error::Result, planner::DefaultQueryPlanner};
-use crate::{internal_err, utils};
 
 use crate::execution::providers::CatalogProviderList;
 
@@ -481,30 +481,4 @@ order by
             ],
         );
     }
-
-    #[test]
-    #[cfg(feature = "connectorx")]
-    fn test_postgres() {
-        use crate::datasource::connectorx::postgres::PostgresCatalogProvider;
-
-        let session = ExecuteSession::new().unwrap();
-        let catalog = PostgresCatalogProvider::try_new("postgresql://root:root@localhost:5433/qurious").unwrap();
-
-        session.register_catalog("qurious", Arc::new(catalog)).unwrap();
-
-        let data = session
-            .sql("SELECT * FROM qurious.public.schools WHERE id = 1")
-            .unwrap();
-
-        assert_batch_eq(
-            &data,
-            vec![
-                "+----+--------------------+",
-                "| id | name               |",
-                "+----+--------------------+",
-                "| 1  | BeiJing University |",
-                "+----+--------------------+",
-            ],
-        );
-    }
 }

qurious/src/logical/builder.rs

@@ -197,7 +197,7 @@ mod tests {
             "SELECT * FROM users a JOIN repos b ON a.id = b.owner_id",
             vec![
                 "Projection: (a.email, a.id, b.id, a.name, b.name, b.owner_id)",
-                "  Inner Join: Filter: users.id = repos.owner_id",
+                "  Inner Join: Filter: a.id = b.owner_id",
                 "    SubqueryAlias: a",
                 "      TableScan: users",
                 "    SubqueryAlias: b",
@@ -212,7 +212,7 @@ mod tests {
             "SELECT * FROM users a LEFT JOIN repos b ON a.id = b.owner_id",
             vec![
                 "Projection: (a.email, a.id, b.id, a.name, b.name, b.owner_id)",
-                "  Left Join: Filter: users.id = repos.owner_id",
+                "  Left Join: Filter: a.id = b.owner_id",
                 "    SubqueryAlias: a",
                 "      TableScan: users",
                 "    SubqueryAlias: b",
@@ -227,7 +227,7 @@ mod tests {
             "SELECT * FROM users a RIGHT JOIN repos b ON a.id = b.owner_id",
             vec![
                 "Projection: (a.email, a.id, b.id, a.name, b.name, b.owner_id)",
-                "  Right Join: Filter: users.id = repos.owner_id",
+                "  Right Join: Filter: a.id = b.owner_id",
                 "    SubqueryAlias: a",
                 "      TableScan: users",
                 "    SubqueryAlias: b",
@@ -242,7 +242,7 @@ mod tests {
             "SELECT * FROM users a FULL JOIN repos b ON a.id = b.owner_id",
             vec![
                 "Projection: (a.email, a.id, b.id, a.name, b.name, b.owner_id)",
-                "  Full Join: Filter: users.id = repos.owner_id",
+                "  Full Join: Filter: a.id = b.owner_id",
                 "    SubqueryAlias: a",
                 "      TableScan: users",
                 "    SubqueryAlias: b",
@@ -257,8 +257,8 @@ mod tests {
             "SELECT * FROM users a JOIN repos b ON a.id = b.owner_id WHERE a.name = 'test'",
             vec![
                 "Projection: (a.email, a.id, b.id, a.name, b.name, b.owner_id)",
-                "  Filter: users.name = Utf8('test')",
-                "    Inner Join: Filter: users.id = repos.owner_id",
+                "  Filter: a.name = Utf8('test')",
+                "    Inner Join: Filter: a.id = b.owner_id",
                 "      SubqueryAlias: a",
                 "        TableScan: users",
                 "      SubqueryAlias: b",
@@ -273,7 +273,7 @@ mod tests {
             "SELECT * FROM users a JOIN repos b ON a.id = b.owner_id AND a.name = b.name",
             vec![
                 "Projection: (a.email, a.id, b.id, a.name, b.name, b.owner_id)",
-                "  Inner Join: Filter: users.id = repos.owner_id AND users.name = repos.name",
+                "  Inner Join: Filter: a.id = b.owner_id AND a.name = b.name",
                 "    SubqueryAlias: a",
                 "      TableScan: users",
                 "    SubqueryAlias: b",

qurious/src/optimizer/pushdown_filter_inner_join.rs

@@ -257,21 +257,21 @@ mod tests {
     #[test]
     fn test_not_valid_join_pair() {
         assert_after_optimizer(
-            "SELECT * FROM users,repos WHERE (users.id = repos.user_id AND users.id = 10) OR (users.name = repos.name AND repos.id = 20)",
+            "SELECT * FROM users,repos WHERE (users.id = repos.owner_id AND users.id = 10) OR (users.name = repos.name AND repos.id = 20)",
             vec![
                 "Projection: (users.email, repos.id, users.id, repos.name, users.name, repos.owner_id)",
-                "  Filter: users.id = repos.user_id AND users.id = Int64(10) OR users.name = repos.name AND repos.id = Int64(20)",
+                "  Filter: users.id = repos.owner_id AND users.id = Int64(10) OR users.name = repos.name AND repos.id = Int64(20)",
                 "    CrossJoin",
                 "      TableScan: users",
                 "      TableScan: repos",
             ],
         );
 
         assert_after_optimizer(
-            "SELECT * FROM users,repos WHERE (users.id = repos.user_id AND users.id = 10) OR (users.id = repos.id OR users.id = 20)",
+            "SELECT * FROM users,repos WHERE (users.id = repos.owner_id AND users.id = 10) OR (users.id = repos.id OR users.id = 20)",
             vec![
                 "Projection: (users.email, repos.id, users.id, repos.name, users.name, repos.owner_id)",
-                "  Filter: users.id = repos.user_id AND users.id = Int64(10) OR users.id = repos.id OR users.id = Int64(20)",
+                "  Filter: users.id = repos.owner_id AND users.id = Int64(10) OR users.id = repos.id OR users.id = Int64(20)",
                 "    CrossJoin",
                 "      TableScan: users",
                 "      TableScan: repos",
@@ -282,11 +282,11 @@ mod tests {
     #[test]
     fn should_not_pushdown_filter_for_inner_join() {
         assert_after_optimizer(
-            "SELECT * FROM users INNER JOIN repos ON users.id = repos.user_id WHERE users.id = 10",
+            "SELECT * FROM users INNER JOIN repos ON users.id = repos.owner_id WHERE users.id = 10",
             vec![
                 "Projection: (users.email, repos.id, users.id, repos.name, users.name, repos.owner_id)",
                 "  Filter: users.id = Int64(10)",
-                "    Inner Join: Filter: users.id = repos.user_id",
+                "    Inner Join: Filter: users.id = repos.owner_id",
                 "      TableScan: users",
                 "      TableScan: repos",
             ],

qurious/src/physical/expr/column.rs

@@ -5,7 +5,7 @@ use arrow::{array::ArrayRef, record_batch::RecordBatch};
 use super::PhysicalExpr;
 use crate::error::{Error, Result};
 
-#[derive(Debug)]
+#[derive(Debug, Hash, PartialEq, Eq)]
 pub struct Column {
     name: String,
     index: usize,

qurious/src/physical/plan/aggregate/hash.rs

@@ -1,35 +1,27 @@
+use crate::arrow_err;
 use crate::error::{Error, Result};
 use crate::physical::expr::Accumulator;
 use crate::physical::{
     expr::{AggregateExpr, PhysicalExpr},
     plan::PhysicalPlan,
 };
-use crate::{arrow_err, hash_array};
+use crate::utils::array::create_hashes;
 use arrow::compute::TakeOptions;
 use arrow::row::{RowConverter, SortField};
-use arrow::{
-    array::{Array, Int32Array, Int64Array, StringArray, UInt8Array},
-    datatypes::*,
-};
 use arrow::{
     array::{ArrayRef, UInt64Array},
     compute,
     datatypes::SchemaRef,
     record_batch::RecordBatch,
 };
-use std::{
-    collections::HashMap,
-    fmt::Display,
-    hash::{DefaultHasher, Hash, Hasher},
-    sync::Arc,
-};
+use std::{collections::HashMap, fmt::Display, hash::DefaultHasher, sync::Arc};
 
 struct GroupAccumulator<'a> {
     /// accumulators each group has a vector of accumulators
     /// Key: group index Value: (group_values, accumulators, indices)
     accumulators: HashMap<usize, (Vec<ArrayRef>, Vec<Box<dyn Accumulator>>)>,
     /// Key: row num Value: hash
-    hashes_buffer: HashMap<usize, DefaultHasher>,
+    hashes_buffer: Vec<DefaultHasher>,
     /// Key: hash Value: group index
     map: HashMap<u64, usize>,
     ///
@@ -43,51 +35,36 @@ impl<'a> GroupAccumulator<'a> {
     {
         Ok(Self {
             accumulators: HashMap::new(),
-            hashes_buffer: HashMap::new(),
+            hashes_buffer: Vec::new(),
             map: HashMap::new(),
             accumlator_factory: f,
         })
     }
 
-    fn update(&mut self, group_by_values: &[ArrayRef], input_values: &[ArrayRef]) -> Result<()> {
+    fn update(&mut self, rows: usize, group_by_values: &[ArrayRef], input_values: &[ArrayRef]) -> Result<()> {
         self.hashes_buffer.clear();
+        self.hashes_buffer.resize(rows, DefaultHasher::new());
 
-        for col in group_by_values {
-            match col.data_type() {
-                DataType::UInt8 => hash_array!(UInt8Array, col, &mut self.hashes_buffer),
-                DataType::Int32 => hash_array!(Int32Array, col, &mut self.hashes_buffer),
-                DataType::Int64 => hash_array!(Int64Array, col, &mut self.hashes_buffer),
-                DataType::Utf8 => hash_array!(StringArray, col, &mut self.hashes_buffer),
-                _ => {
-                    return Err(Error::InternalError(format!(
-                        "[group_indices] unsupported data type {:?}",
-                        col.data_type()
-                    )))
-                }
-            }
-        }
-
+        let hashes = create_hashes(&group_by_values, &mut self.hashes_buffer)?;
         let mut accs_indices = HashMap::new();
-        for (row, target_hash_buffer) in &self.hashes_buffer {
-            let target_hash = target_hash_buffer.finish();
-
+        for (row, target_hash) in hashes.into_iter().enumerate() {
             match self.map.get_mut(&target_hash) {
                 Some(group_index) => {
-                    accs_indices.entry(*group_index).or_insert(vec![]).push(*row as u64);
+                    accs_indices.entry(*group_index).or_insert(vec![]).push(row as u64);
                 }
                 None => {
-                    self.map.insert(target_hash, *row);
+                    self.map.insert(target_hash, row);
 
-                    accs_indices.insert(*row, vec![*row as u64]);
+                    accs_indices.insert(row, vec![row as u64]);
 
                     let accs = (self.accumlator_factory)()?;
-                    let indices = UInt64Array::from_iter(vec![*row as u64]);
+                    let indices = UInt64Array::from_iter(vec![row as u64]);
                     let group_values = group_by_values
                         .iter()
                         .map(|values| compute::take(&values, &indices, None).map_err(|e| arrow_err!(e)))
                         .collect::<Result<Vec<_>>>()?;
 
-                    self.accumulators.insert(*row, (group_values, accs));
+                    self.accumulators.insert(row, (group_values, accs));
                 }
             }
         }
@@ -181,7 +158,7 @@ impl PhysicalPlan for HashAggregate {
                 .map(|e| e.expression().evaluate(&batch))
                 .collect::<Result<Vec<ArrayRef>>>()?;
 
-            group_accumulator.update(&group_by_values, &input_values)?;
+            group_accumulator.update(batch.num_rows(), &group_by_values, &input_values)?;
         }
 
         RecordBatch::try_new(self.schema.clone(), group_accumulator.output(&self.schema)?)