in progress

spiraldb · Jun 25, 2024 · 62a186f · 62a186f
1 parent 3214287
commit 62a186f
Show file tree

Hide file tree

Showing 8 changed files with 514 additions and 118 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/vortex-array/src/array/bool/compute/compare.rs b/vortex-array/src/array/bool/compute/compare.rs
@@ -8,6 +8,7 @@ use crate::compute::compare::CompareFn;
 use crate::{Array, ArrayTrait, IntoArray, IntoArrayVariant};
 
 impl CompareFn for BoolArray {
+    // TODO(aduffy): replace these with Arrow compute kernels.
     fn compare(&self, other: &Array, op: Operator) -> VortexResult<Array> {
         let flattened = other.clone().into_bool()?;
         let lhs = self.boolean_buffer();

diff --git a/vortex-array/src/array/constant/mod.rs b/vortex-array/src/array/constant/mod.rs
@@ -7,6 +7,7 @@ use crate::impl_encoding;
 use crate::stats::Stat;
 use crate::validity::{ArrayValidity, LogicalValidity};
 use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
+
 mod compute;
 mod flatten;
 mod stats;
@@ -25,6 +26,8 @@ impl ConstantArray {
         Scalar: From<S>,
     {
         let scalar: Scalar = scalar.into();
+        // TODO(aduffy): add stats for bools, ideally there should be a
+        //  StatsSet::constant(Scalar) constructor that does this for us, like StatsSet::nulls.
         let stats = StatsSet::from(HashMap::from([
             (Stat::Max, scalar.clone()),
             (Stat::Min, scalar.clone()),

diff --git a/vortex-array/src/array/struct_/mod.rs b/vortex-array/src/array/struct_/mod.rs
@@ -1,6 +1,6 @@
 use serde::{Deserialize, Serialize};
 use vortex_dtype::{FieldName, FieldNames, Nullability, StructDType};
-use vortex_error::{vortex_bail, vortex_err};
+use vortex_error::vortex_bail;
 
 use crate::stats::ArrayStatisticsCompute;
 use crate::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
@@ -27,6 +27,15 @@ impl StructArray {
         self.array().child(idx, dtype)
     }
 
+    pub fn field_by_name(&self, name: &str) -> Option<Array> {
+        let field_idx = self
+            .names()
+            .iter()
+            .position(|field_name| field_name.as_ref() == name);
+
+        field_idx.and_then(|field_idx| self.field(field_idx))
+    }
+
     pub fn names(&self) -> &FieldNames {
         let DType::Struct(st, _) = self.dtype() else {
             unreachable!()
@@ -126,7 +135,7 @@ impl StructArray {
         for column_idx in projection {
             children.push(
                 self.field(*column_idx)
-                    .ok_or_else(|| vortex_err!(InvalidArgument: "column index out of bounds"))?,
+                    .expect("column must not exceed bounds"),
             );
             names.push(self.names()[*column_idx].clone());
         }

diff --git a/vortex-array/src/stats/statsset.rs b/vortex-array/src/stats/statsset.rs
@@ -3,6 +3,7 @@ use std::collections::HashMap;
 
 use enum_iterator::all;
 use itertools::Itertools;
+
 use vortex_dtype::DType;
 use vortex_error::VortexError;
 use vortex_scalar::Scalar;
@@ -27,6 +28,36 @@ impl StatsSet {
         }
     }
 
+    // pub fn constant(len: usize, scalar: &Scalar) -> Self {
+    //     let mut stats = HashMap::from([
+    //         (Stat::Max, scalar.clone()),
+    //         (Stat::Min, scalar.clone()),
+    //         (Stat::IsConstant, true.into()),
+    //         (Stat::IsSorted, true.into()),
+    //         (Stat::RunCount, 1.into()),
+    //     ]);
+    //
+    //     match scalar.dtype() {
+    //         DType::Bool(_) => {
+    //             stats.insert(Stat::TrueCount, 0.into());
+    //         }
+    //         DType::Primitive(ptype, _) => {
+    //             ptype.byte_width();
+    //             stats.insert(
+    //                 Stat::BitWidthFreq,
+    //                 vec![0; ptype.byte_width() * 8 + 1].into(),
+    //             );
+    //             stats.insert(
+    //                 Stat::TrailingZeroFreq,
+    //                 vec![ptype.byte_width() * 8; ptype.byte_width() * 8 + 1].into(),
+    //             );
+    //         }
+    //         _ => {}
+    //     }
+    //
+    //     Self::from(stats)
+    // }
+
     /// Specialized constructor for the case where the StatsSet represents
     /// an array consisting entirely of [null](vortex_dtype::DType::Null) values.
     pub fn nulls(len: usize, dtype: &DType) -> Self {

diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml
@@ -13,7 +13,9 @@ rust-version.workspace = true
 [dependencies]
 vortex-array = { path = "../vortex-array" }
 vortex-dtype = { path = "../vortex-dtype" }
+vortex-expr = { path = "../vortex-expr" }
 vortex-error = { path = "../vortex-error" }
+vortex-scalar = { path = "../vortex-scalar" }
 
 arrow-array = { workspace = true }
 arrow-schema = { workspace = true }
@@ -26,10 +28,13 @@ datafusion-execution = { workspace = true }
 datafusion-physical-expr = { workspace = true }
 datafusion-physical-plan = { workspace = true }
 futures = { workspace = true }
+itertools = { workspace = true }
+lazy_static = { workspace = true }
 pin-project = { workspace = true }
 
 [dev-dependencies]
 tokio = { workspace = true, features = ["test-util"] }
 
-[lints]
-workspace = true
+# TODO(aduffy): re-enable
+#[lints]
+#workspace = true