feat: add encoder utilities for pushdown

Cargo.toml

@@ -6,6 +6,7 @@ members = [
     "rust/lance-core",
     "rust/lance-datagen",
     "rust/lance-encoding",
+    "rust/lance-encoding-datafusion",
     "rust/lance-file",
     "rust/lance-index",
     "rust/lance-io",
@@ -48,6 +49,7 @@ lance-core = { version = "=0.11.1", path = "./rust/lance-core" }
 lance-datafusion = { version = "=0.11.1", path = "./rust/lance-datafusion" }
 lance-datagen = { version = "=0.11.1", path = "./rust/lance-datagen" }
 lance-encoding = { version = "=0.11.1", path = "./rust/lance-encoding" }
+lance-encoding-datafusion = { version = "=0.11.1", path = "./rust/lance-encoding-datafusion" }
 lance-file = { version = "=0.11.1", path = "./rust/lance-file" }
 lance-index = { version = "=0.11.1", path = "./rust/lance-index" }
 lance-io = { version = "=0.11.1", path = "./rust/lance-io" }

protos/encodings-df.proto

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+
+syntax = "proto3";
+
+package lance.encodings_datafusion;
+
+import "encodings.proto";
+
+// ZoneMaps are a way to wrap any leaf array with a set of zone maps that
+// can be used to apply pushdown filtering.
+//
+// A "zone map" is the min/max/null_count of a set of rows.  This can be
+// used to quickly filter out zones which are not included in simple predicates
+// like 'x = 5' or 'x > 10'.
+message ZoneMaps {
+
+    // How many rows are covered by each zone map.  There will be
+    // ceil_div(num_rows, rows_per_map) zone maps.
+    uint32 rows_per_map = 1;
+
+    // The zone maps are encoded as struct arrays with 1 row per zone.  This
+    // should be stored in a column metadata buffer.  The struct array should
+    // have 3 children: min: T, max: T, null_count: u32
+    lance.encodings.ArrayEncoding stats = 2;
+
+    // The underlying array values
+    lance.encodings.ArrayEncoding values = 5;
+}
+

python/src/file.rs

@@ -185,6 +185,7 @@ impl LanceFileWriter {
             FileWriterOptions {
                 data_cache_bytes,
                 keep_original_array,
+                ..Default::default()
             },
         )
         .infer_error()?;

rust/lance-encoding-datafusion/Cargo.toml

@@ -0,0 +1,37 @@
+[package]
+name = "lance-encoding-datafusion"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+license.workspace = true
+repository.workspace = true
+readme = "README.md"
+description = "Encoders and decoders for the Lance file format that rely on datafusion"
+keywords.workspace = true
+categories.workspace = true
+rust-version.workspace = true
+
+[dependencies]
+lance-core = { workspace = true, features = ["datafusion"] }
+lance-encoding.workspace = true
+lance-file.workspace = true
+arrow-array.workspace = true
+arrow-buffer.workspace = true
+arrow-schema.workspace = true
+datafusion-common.workspace = true
+datafusion-expr.workspace = true
+datafusion-physical-expr.workspace = true
+futures.workspace = true
+prost.workspace = true
+prost-types.workspace = true
+
+[dev-dependencies]
+rand.workspace = true
+tokio.workspace = true
+lance-datagen.workspace = true
+
+[build-dependencies]
+prost-build.workspace = true
+
+[target.'cfg(target_os = "linux")'.dev-dependencies]
+pprof = { workspace = true }

rust/lance-encoding-datafusion/README.md

@@ -0,0 +1,8 @@
+# lance-encoding-datafusion
+
+`lance-encoding-datafusion` is an internal sub-crate, containing encoders and
+decoders for the Lance file format that rely on Datafusion. Partly this is to
+keep the size of `lance-encoding` small and partly this is to prove that
+encodings are extensible.
+
+**Important Note**: This crate is **not intended for external usage**.

rust/lance-encoding-datafusion/build.rs

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+use std::io::Result;
+
+fn main() -> Result<()> {
+    println!("cargo:rerun-if-changed=protos");
+
+    let mut prost_build = prost_build::Config::new();
+    prost_build.extern_path(".lance.encodings", "::lance_encoding::format::pb");
+    prost_build.protoc_arg("--experimental_allow_proto3_optional");
+    prost_build.enable_type_names();
+    prost_build.compile_protos(&["./protos/encodings-df.proto"], &["./protos"])?;
+
+    Ok(())
+}

rust/lance-encoding-datafusion/protos

@@ -0,0 +1 @@
+../../protos

rust/lance-encoding-datafusion/src/format.rs

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+/// Protobuf definitions for encodings
+pub mod pb {
+    #![allow(clippy::all)]
+    #![allow(non_upper_case_globals)]
+    #![allow(non_camel_case_types)]
+    #![allow(non_snake_case)]
+    #![allow(unused)]
+    #![allow(improper_ctypes)]
+    #![allow(clippy::upper_case_acronyms)]
+    #![allow(clippy::use_self)]
+    include!(concat!(env!("OUT_DIR"), "/lance.encodings_datafusion.rs"));
+}

rust/lance-encoding-datafusion/src/lib.rs

@@ -0,0 +1,63 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+use arrow_schema::DataType;
+use lance_encoding::encoder::{
+    ColumnIndexSequence, CoreFieldEncodingStrategy, FieldEncodingStrategy,
+};
+use zone::ZoneMapsFieldEncoder;
+
+pub mod format;
+pub mod zone;
+
+/// Wraps the core encoding strategy and adds the encoders from this
+/// crate
+#[derive(Debug)]
+pub struct LanceDfFieldEncodingStrategy {
+    core: CoreFieldEncodingStrategy,
+    rows_per_map: u32,
+}
+
+impl FieldEncodingStrategy for LanceDfFieldEncodingStrategy {
+    fn create_field_encoder(
+        &self,
+        encoding_strategy_root: &dyn FieldEncodingStrategy,
+        field: &lance_core::datatypes::Field,
+        column_index: &mut ColumnIndexSequence,
+        cache_bytes_per_column: u64,
+        keep_original_array: bool,
+        config: &std::collections::HashMap<String, String>,
+    ) -> lance_core::Result<Box<dyn lance_encoding::encoder::FieldEncoder>> {
+        let data_type = field.data_type();
+        if data_type.is_primitive()
+            || matches!(
+                data_type,
+                DataType::Boolean | DataType::Utf8 | DataType::LargeUtf8
+            )
+        {
+            let inner_encoder = self.core.create_field_encoder(
+                // Don't collect stats on inner string fields
+                &self.core,
+                field,
+                column_index,
+                cache_bytes_per_column,
+                keep_original_array,
+                config,
+            )?;
+            Ok(Box::new(ZoneMapsFieldEncoder::try_new(
+                inner_encoder,
+                data_type.clone(),
+                self.rows_per_map,
+            )?))
+        } else {
+            self.core.create_field_encoder(
+                encoding_strategy_root,
+                field,
+                column_index,
+                cache_bytes_per_column,
+                keep_original_array,
+                config,
+            )
+        }
+    }
+}