Add support for row group pruning on FixedSizeBinary

apache · Mar 16, 2024 · acac9ec · acac9ec
1 parent 156a3f5
commit acac9ec
Show file tree

Hide file tree

Showing 3 changed files with 109 additions and 2 deletions.
diff --git a/datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs b/datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs
@@ -222,6 +222,7 @@ impl PruningStatistics for BloomFilterStatistics {
                 match value {
                     ScalarValue::Utf8(Some(v)) => sbbf.check(&v.as_str()),
                     ScalarValue::Binary(Some(v)) => sbbf.check(v),
+                    ScalarValue::FixedSizeBinary(_size, Some(v)) => sbbf.check(v),
                     ScalarValue::Boolean(Some(v)) => sbbf.check(v),
                     ScalarValue::Float64(Some(v)) => sbbf.check(v),
                     ScalarValue::Float32(Some(v)) => sbbf.check(v),

diff --git a/datafusion/core/src/datasource/physical_plan/parquet/statistics.rs b/datafusion/core/src/datasource/physical_plan/parquet/statistics.rs
@@ -109,17 +109,21 @@ macro_rules! get_statistic {
                     }
                 }
             }
-            // type not supported yet
+            // type not fully supported yet
             ParquetStatistics::FixedLenByteArray(s) => {
                 match $target_arrow_type {
-                    // just support the decimal data type
+                    // just support specific logical data type, there are others each
+                    // with their own ordering
                     Some(DataType::Decimal128(precision, scale)) => {
                         Some(ScalarValue::Decimal128(
                             Some(from_bytes_to_i128(s.$bytes_func())),
                             *precision,
                             *scale,
                         ))
                     }
+                    Some(DataType::FixedSizeBinary(size)) => {
+                        Some(ScalarValue::FixedSizeBinary(*size, Some(s.$bytes_func().to_vec())))
+                    }
                     _ => None,
                 }
             }

diff --git a/datafusion/core/tests/parquet/row_group_pruning.rs b/datafusion/core/tests/parquet/row_group_pruning.rs
@@ -948,6 +948,108 @@ async fn prune_binary_lt() {
         .await;
 }
 
+#[tokio::test]
+async fn prune_fixedsizebinary_eq_match() {
+    RowGroupPruningTest::new()
+        .with_scenario(Scenario::ByteArray)
+        .with_query(
+            "SELECT name, service_fixedsize FROM t WHERE service_fixedsize = ARROW_CAST(CAST('fe6' AS bytea), 'FixedSizeBinary(3)')",
+        )
+        .with_expected_errors(Some(0))
+        // false positive on 'all frontends' batch: 'fe1' < 'fe6' < 'fe7'
+        .with_matched_by_stats(Some(2))
+        .with_pruned_by_stats(Some(1))
+        .with_matched_by_bloom_filter(Some(1))
+        .with_pruned_by_bloom_filter(Some(1))
+        .with_expected_rows(1)
+        .test_row_group_prune()
+        .await;
+
+    RowGroupPruningTest::new()
+        .with_scenario(Scenario::ByteArray)
+        .with_query(
+            "SELECT name, service_fixedsize FROM t WHERE service_fixedsize = ARROW_CAST(CAST('fe6' AS bytea), 'FixedSizeBinary(3)')",
+        )
+        .with_expected_errors(Some(0))
+        // false positive on 'all frontends' batch: 'fe1' < 'fe6' < 'fe7'
+        .with_matched_by_stats(Some(2))
+        .with_pruned_by_stats(Some(1))
+        .with_matched_by_bloom_filter(Some(1))
+        .with_pruned_by_bloom_filter(Some(1))
+        .with_expected_rows(1)
+        .test_row_group_prune()
+        .await;
+
+}
+
+#[tokio::test]
+async fn prune_fixedsizebinary_eq_no_match() {
+    RowGroupPruningTest::new()
+        .with_scenario(Scenario::ByteArray)
+        .with_query(
+            "SELECT name, service_fixedsize FROM t WHERE service_fixedsize = ARROW_CAST(CAST('be9' AS bytea), 'FixedSizeBinary(3)')",
+        )
+        .with_expected_errors(Some(0))
+        // false positive on 'mixed' batch: 'be1' < 'be9' < 'fe4'
+        .with_matched_by_stats(Some(1))
+        .with_pruned_by_stats(Some(2))
+        .with_matched_by_bloom_filter(Some(0))
+        .with_pruned_by_bloom_filter(Some(1))
+        .with_expected_rows(0)
+        .test_row_group_prune()
+        .await;
+}
+
+#[tokio::test]
+async fn prune_fixedsizebinary_neq() {
+    RowGroupPruningTest::new()
+        .with_scenario(Scenario::ByteArray)
+        .with_query(
+            "SELECT name, service_fixedsize FROM t WHERE service_fixedsize != ARROW_CAST(CAST('be1' AS bytea), 'FixedSizeBinary(3)')",
+        )
+        .with_expected_errors(Some(0))
+        .with_matched_by_stats(Some(3))
+        .with_pruned_by_stats(Some(0))
+        .with_matched_by_bloom_filter(Some(3))
+        .with_pruned_by_bloom_filter(Some(0))
+        .with_expected_rows(14)
+        .test_row_group_prune()
+        .await;
+}
+
+#[tokio::test]
+async fn prune_fixedsizebinary_lt() {
+    RowGroupPruningTest::new()
+        .with_scenario(Scenario::ByteArray)
+        .with_query(
+            "SELECT name, service_fixedsize FROM t WHERE service_fixedsize < ARROW_CAST(CAST('be3' AS bytea), 'FixedSizeBinary(3)')",
+        )
+        .with_expected_errors(Some(0))
+        // matches 'all backends' only
+        .with_matched_by_stats(Some(1))
+        .with_pruned_by_stats(Some(2))
+        .with_matched_by_bloom_filter(Some(0))
+        .with_pruned_by_bloom_filter(Some(0))
+        .with_expected_rows(2)
+        .test_row_group_prune()
+        .await;
+
+    RowGroupPruningTest::new()
+        .with_scenario(Scenario::ByteArray)
+        .with_query(
+            "SELECT name, service_fixedsize FROM t WHERE service_fixedsize < ARROW_CAST(CAST('be9' AS bytea), 'FixedSizeBinary(3)')",
+        )
+        .with_expected_errors(Some(0))
+        .with_matched_by_stats(Some(2))
+        .with_pruned_by_stats(Some(1))
+        .with_matched_by_bloom_filter(Some(0))
+        .with_pruned_by_bloom_filter(Some(0))
+        // all backends from 'mixed' and 'all backends'
+        .with_expected_rows(8)
+        .test_row_group_prune()
+        .await;
+}
+
 #[tokio::test]
 async fn prune_periods_in_column_names() {
     // There are three row groups for "service.name", each with 5 rows = 15 rows total