data-flow-control 0.1.4__tar.gz → 0.1.6__tar.gz

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/Cargo.lock

@@ -1593,7 +1593,7 @@ dependencies = [
 
 [[package]]
 name = "passant-cli"
-version = "0.1.4"
+version = "0.1.6"
 dependencies = [
  "anyhow",
  "assert_cmd",
@@ -1605,7 +1605,7 @@ dependencies = [
 
 [[package]]
 name = "passant-core"
-version = "0.1.4"
+version = "0.1.6"
 dependencies = [
  "anyhow",
  "criterion",
@@ -1623,7 +1623,7 @@ dependencies = [
 
 [[package]]
 name = "passant-py"
-version = "0.1.4"
+version = "0.1.6"
 dependencies = [
  "passant-core",
  "pyo3",

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/Cargo.toml

@@ -3,7 +3,7 @@ members = ["passant-core", "passant-py"]
 resolver = "2"
 
 [workspace.package]
-version = "0.1.4"
+version = "0.1.6"
 edition = "2024"
 license = "MIT"
 authors = ["Charlie Summers"]

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-flow-control
-Version: 0.1.4
+Version: 0.1.6
 Requires-Dist: boto3>=1.26.0
 Requires-Dist: duckdb>=1.3.0
 Requires-Dist: langchain>=1.3,<2
@@ -11,6 +11,7 @@ Requires-Dist: clickhouse-connect>=0.7.0 ; extra == 'clickhouse'
 Requires-Dist: datafusion>=53.0.0 ; extra == 'datafusion'
 Requires-Dist: pyarrow>=14.0.0 ; extra == 'datafusion'
 Requires-Dist: maturin>=1.7.0 ; extra == 'dev'
+Requires-Dist: pre-commit>=4.0.0 ; extra == 'dev'
 Requires-Dist: pytest>=8.4.2 ; extra == 'dev'
 Requires-Dist: ruff>=0.11.0 ; extra == 'dev'
 Requires-Dist: psycopg[binary]>=3.2.0 ; extra == 'postgres'

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/src/aggregate_registry.rs

@@ -167,6 +167,8 @@ impl AggregateRegistry {
             || lower.contains("countif")
             || lower.contains("approx_count_distinct")
             || lower.contains("approx_distinct")
+            || lower == "arbitrary"
+            || lower == "any_value"
             || is_count_like_name(&lower)
     }
 

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/src/full_push.rs

@@ -18,6 +18,10 @@ fn semiring_allows_full_push(rewriter: &PassantRewriter, semiring: &SemiringAnal
         .all(|aggregate| count_if_style_scan_aggregate(rewriter, aggregate))
 }
 
+fn aggregate_function_name(aggregate: &str) -> &str {
+    aggregate.split('(').next().unwrap_or(aggregate).trim()
+}
+
 fn count_if_style_scan_aggregate(rewriter: &PassantRewriter, aggregate: &str) -> bool {
     let lower = aggregate.to_ascii_lowercase();
     if lower.contains("array_agg")
@@ -27,7 +31,9 @@ fn count_if_style_scan_aggregate(rewriter: &PassantRewriter, aggregate: &str) ->
     {
         return false;
     }
-    rewriter.aggregate_registry.is_scan_transformable(aggregate)
+    rewriter
+        .aggregate_registry
+        .is_scan_transformable(aggregate_function_name(&lower))
 }
 
 /// Full-push engine — mutates the query in place, preserving its overall shape.

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/src/rewrite_strategy.rs

@@ -377,24 +377,71 @@ mod tests {
     }
 
     #[test]
-    fn non_distributive_policy_uses_partial_push() {
+    fn introspected_arbitrary_in_catalog_still_rewrites_kill_insert() {
+        use crate::aggregate_registry::AggregateFunctionSnapshot;
+        use crate::catalog::{CatalogSnapshot, CatalogTableInfo};
+        use std::collections::HashMap;
+
+        let policy_text = "SOURCE Receipts SINK Expenses CONSTRAINT ARBITRARY(Receipts.cat) != 'Meal' OR Expenses.biz_use <= 50.0 ON FAIL KILL";
+        let sql = "INSERT INTO Expenses SELECT 1, 100.0 FROM Receipts";
+
         let mut rewriter = PassantRewriter::new();
-        rewriter.register_policy(non_distributive_policy());
-        let statement = sqlparser::parser::Parser::parse_sql(
-            &sqlparser::dialect::GenericDialect {},
-            "SELECT id FROM foo",
-        )
-        .expect("parse")
-        .pop()
-        .expect("statement");
-        let request = RewriteRequest::analyze(
-            "SELECT id FROM foo",
-            &statement,
-            RewriteOptions::default(),
-            rewriter.policy_store(),
+        rewriter
+            .register_validated_policy(crate::parse_policy_text(policy_text).expect("parse"))
+            .expect("register without catalog");
+        let without_catalog = rewriter.rewrite(sql).expect("rewrite without catalog");
+        assert!(
+            without_catalog.contains("passant_kill"),
+            "baseline rewrite without catalog introspection: {without_catalog}"
+        );
+
+        let mut tables = HashMap::new();
+        tables.insert(
+            "Receipts".to_string(),
+            CatalogTableInfo {
+                columns: vec!["cat".into(), "biz_use".into()],
+                ..CatalogTableInfo::default()
+            },
+        );
+        tables.insert(
+            "Expenses".to_string(),
+            CatalogTableInfo {
+                columns: vec!["biz_use".into()],
+                ..CatalogTableInfo::default()
+            },
+        );
+        let mut rewriter = PassantRewriter::new();
+        rewriter.apply_catalog_snapshot(CatalogSnapshot {
+            dialect: Some("duckdb".to_string()),
+            tables,
+            aggregate_functions: vec![AggregateFunctionSnapshot {
+                name: "arbitrary".to_string(),
+                schema: None,
+                aliases: vec![],
+                classification: None,
+                source: Some("introspected".to_string()),
+            }],
+            ..CatalogSnapshot::default()
+        });
+        rewriter
+            .register_validated_policy(crate::parse_policy_text(policy_text).expect("parse"))
+            .expect("register validated policy");
+        assert!(
+            rewriter
+                .policy_store()
+                .has_enforcement_resolution_policies(),
+            "KILL policy should register as enforcement resolution"
+        );
+        assert_eq!(rewriter.policy_store().active_count(), 1);
+        let with_catalog = rewriter.rewrite(sql);
+        if let Err(err) = &with_catalog {
+            panic!("rewrite with catalog failed: {err}");
+        }
+        let with_catalog = with_catalog.expect("rewrite ok");
+        let summary = rewriter.last_statement_rewrite_summary();
+        assert!(
+            with_catalog.contains("passant_kill"),
+            "expected kill rewrite with introspected arbitrary aggregate, got: {with_catalog}; summary={summary:?}"
         );
-        assert!(!request.semiring.all_distributive);
-        assert!(!FullPushEngine.matches(&rewriter, &request));
-        assert!(PartialPushEngine.matches(&rewriter, &request));
     }
 }

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/src/rewriter/aggregates.rs

@@ -130,6 +130,9 @@ fn count_distinct_inner_column(expr: &Expr) -> Option<Expr> {
 }
 
 fn transform_scan_aggregates_recursive(expr: Expr, registry: &AggregateRegistry) -> Expr {
+    if let Ok(Some(rewritten)) = rewrite_count_distinct_equality(&expr) {
+        return rewritten;
+    }
     if let Expr::BinaryOp {
         left,
         op: BinaryOperator::Divide,

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/src/rewriter/plan/mod.rs

@@ -120,4 +120,85 @@ mod tests {
         assert_eq!(summary.scope_diagnostics.len(), 2);
         assert_eq!(summary.aggregate().candidate_policies, 2);
     }
+
+    #[test]
+    fn kill_policy_with_introspected_arbitrary_emits_tuple_action() {
+        use crate::aggregate_registry::AggregateFunctionSnapshot;
+        use crate::catalog::{CatalogSnapshot, CatalogTableInfo};
+        use std::collections::HashMap;
+
+        let mut tables = HashMap::new();
+        tables.insert(
+            "Receipts".to_string(),
+            CatalogTableInfo {
+                columns: vec!["cat".into()],
+                ..CatalogTableInfo::default()
+            },
+        );
+        tables.insert(
+            "Expenses".to_string(),
+            CatalogTableInfo {
+                columns: vec!["biz_use".into()],
+                ..CatalogTableInfo::default()
+            },
+        );
+        let mut rewriter = PassantRewriter::new();
+        rewriter.apply_catalog_snapshot(CatalogSnapshot {
+            dialect: Some("duckdb".to_string()),
+            tables,
+            aggregate_functions: vec![AggregateFunctionSnapshot {
+                name: "arbitrary".to_string(),
+                schema: None,
+                aliases: vec![],
+                classification: None,
+                source: Some("introspected".to_string()),
+            }],
+            ..CatalogSnapshot::default()
+        });
+        rewriter
+            .register_validated_policy(
+                crate::parse_policy_text(
+                    "SOURCE Receipts SINK Expenses CONSTRAINT ARBITRARY(Receipts.cat) != 'Meal' OR Expenses.biz_use <= 50.0 ON FAIL KILL",
+                )
+                .expect("parse"),
+            )
+            .expect("register");
+        let statement =
+            crate::parser::parse_query("INSERT INTO Expenses SELECT 1, 100.0 FROM Receipts")
+                .expect("parse");
+        let sqlparser::ast::Statement::Insert(insert) = statement else {
+            panic!("expected insert");
+        };
+        let Some(mut query) = insert.source else {
+            panic!("expected insert source");
+        };
+        let sqlparser::ast::SetExpr::Select(select) = query.body.as_mut() else {
+            panic!("expected select");
+        };
+        let analysis = SelectAnalysis::from_select(select);
+        let context = RewriteContext {
+            sink: Some("Expenses".to_string()),
+            sink_expr_by_column: HashMap::from([(
+                "biz_use".to_string(),
+                sqlparser::ast::Expr::Value(sqlparser::ast::Value::Number("100.0".into(), false)),
+            )]),
+            ..RewriteContext::default()
+        };
+        let plan = plan_select_rewrite(
+            rewriter.policy_store(),
+            rewriter.catalog(),
+            None,
+            select.as_mut(),
+            &analysis,
+            &context,
+            &HashSet::new(),
+        )
+        .expect("plan should succeed");
+        assert_eq!(
+            plan.policy_actions.len(),
+            1,
+            "expected one kill action, diagnostics={:?}",
+            plan.diagnostics
+        );
+    }
 }

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/src/rewriter/policy_expr.rs

@@ -433,7 +433,41 @@ pub(crate) fn scan_policy_expr(
             )?,
         ));
     }
+    if let Expr::BinaryOp {
+        left,
+        op: BinaryOperator::Or,
+        right,
+    } = &expr
+    {
+        return Ok(Expr::BinaryOp {
+            left: Box::new(scan_policy_expr(
+                *left.clone(),
+                sources,
+                context,
+                alias_by_base,
+                scan_ready,
+                is_aggregation,
+            )?),
+            op: BinaryOperator::Or,
+            right: Box::new(scan_policy_expr(
+                *right.clone(),
+                sources,
+                context,
+                alias_by_base,
+                scan_ready,
+                is_aggregation,
+            )?),
+        });
+    }
     let registry = &context.aggregate_registry;
+    if !scan_ready {
+        let transformed = transform_scan_aggregates(expr.clone(), registry)?;
+        if !registry.expr_contains_aggregate(&transformed) {
+            let mut transformed = transformed;
+            rewrite_column_qualifiers(&mut transformed, alias_by_base);
+            return Ok(transformed);
+        }
+    }
     let non_distributive = non_distributive_aggregates(&expr, registry)?;
     if non_distributive.is_empty() {
         if scan_ready {

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/src/rewriter/resolution.rs

@@ -60,17 +60,13 @@ fn is_kill_wrapped_select(select: &Select) -> bool {
     if select.projection.len() != 1 {
         return false;
     }
-    if !matches!(
-        select.projection.first(),
-        Some(SelectItem::Wildcard(_))
-    ) {
+    if !matches!(select.projection.first(), Some(SelectItem::Wildcard(_))) {
         return false;
     }
     let Some(TableWithJoins {
-        relation:
-            TableFactor::Derived {
-                alias: Some(alias), ..
-            },
+        relation: TableFactor::Derived {
+            alias: Some(alias), ..
+        },
         joins,
         ..
     }) = select.from.first()
@@ -753,7 +749,10 @@ mod tests {
         remap_query_order_by_after_kill_wrap(&mut query);
 
         let sql = render_statement(&crate::sql::statement_from_query(query), None);
-        assert!(sql.contains("ORDER BY id"), "expected unqualified ORDER BY: {sql}");
+        assert!(
+            sql.contains("ORDER BY id"),
+            "expected unqualified ORDER BY: {sql}"
+        );
         assert!(
             !sql.contains("ORDER BY receipts."),
             "qualified ORDER BY should be remapped: {sql}"

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/tests/gaps/insert_shapes.rs

@@ -88,7 +88,10 @@ fn insert_select_kill_order_by_remapped() {
         reports_catalog(),
     );
     assert!(sql.contains("passant_kill"), "expected kill wrap: {sql}");
-    assert!(sql.contains("ORDER BY id"), "expected unqualified ORDER BY: {sql}");
+    assert!(
+        sql.contains("ORDER BY id"),
+        "expected unqualified ORDER BY: {sql}"
+    );
     assert!(
         !sql.contains("ORDER BY Receipts."),
         "qualified ORDER BY should be remapped: {sql}"

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/tests/rewrite/scan.rs

@@ -269,7 +269,10 @@ fn select_kill_order_by_remapped() {
         .rewrite("SELECT id FROM foo ORDER BY foo.id")
         .expect("query should rewrite");
     assert!(sql.contains("passant_kill"), "expected kill wrap: {sql}");
-    assert!(sql.contains("ORDER BY id"), "expected unqualified ORDER BY: {sql}");
+    assert!(
+        sql.contains("ORDER BY id"),
+        "expected unqualified ORDER BY: {sql}"
+    );
     assert!(
         !sql.contains("ORDER BY foo."),
         "qualified ORDER BY should be remapped: {sql}"

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "data-flow-control"
-version = "0.1.4"
+version = "0.1.6"
 description = "Rust-backed Data Flow Control rewrite engine"
 license = "MIT"
 license-files = ["LICENSE"]
@@ -17,6 +17,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
     "maturin>=1.7.0",
+    "pre-commit>=4.0.0",
     "pytest>=8.4.2",
     "ruff>=0.11.0",
 ]

{data_flow_control-0.1.4 → data_flow_control-0.1.6}/passant-core/src/rewriter/select.rs

@@ -20,8 +20,8 @@ use super::derived_policy::{
 use super::exists::apply_in_semijoin_policy_filters;
 use super::helpers::select_output_column_mapping;
 use super::limit::wrap_limited_policy_query;
-use super::resolution::remap_query_order_by_after_kill_wrap;
 use super::plan::{apply_select_rewrite_plan, plan_select_rewrite};
+use super::resolution::remap_query_order_by_after_kill_wrap;
 use super::types::{PassantRewriter, RewriteContext};
 
 impl PassantRewriter {