datafusion-index-provider 0.1.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Defines the [`RecordFetcher`] trait for fetching complete records using row IDs.
//!
//! The [`RecordFetcher`] trait is a key abstraction in the two-phase execution model.
//! After indexes produce row IDs during the index phase, the fetch phase uses this
//! trait to retrieve the actual data records corresponding to those row IDs.

use async_trait::async_trait;
use datafusion::arrow::datatypes::SchemaRef;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::Result;

/// A trait for fetching complete data records based on primary key values produced by index scans.
///
/// This trait abstracts the process of retrieving actual data records from the underlying
/// storage system using primary key identifiers. It serves as the bridge between the index phase
/// (which produces primary key values) and the final query results (which contain complete records).
///
/// ## Implementation Requirements
///
/// Implementations must handle:
/// - **Primary key extraction**: Parse the primary key columns from the input batch.
///   The batch schema matches `Index::index_schema()` and may contain one or more columns
///   forming a composite primary key.
/// - **Efficient lookup**: Retrieve records using the most efficient access pattern for your storage
/// - **Schema consistency**: Return records matching the schema from `schema()`
/// - **Error handling**: Properly propagate storage errors and handle missing records
/// - **Async execution**: Support concurrent fetching for performance
///
/// ## Performance Considerations
///
/// The performance of your `RecordFetcher` implementation directly impacts query performance:
/// - **Batch processing**: Process multiple primary keys together to amortize lookup costs
/// - **Storage locality**: Consider sorting primary keys to improve storage access patterns
/// - **Caching**: Implement appropriate caching strategies for frequently accessed data
/// - **Resource management**: Manage memory and connection pooling efficiently
#[async_trait]
pub trait RecordFetcher: Send + Sync + std::fmt::Debug {
    /// Returns the schema of the complete records that will be fetched.
    ///
    /// This schema represents the full table structure and must match the schema of
    /// the `RecordBatch` returned by `fetch()`. It typically contains all columns
    /// from your table, not just the row ID column.
    fn schema(&self) -> SchemaRef;

    /// Fetches complete data records corresponding to the primary key values in the input batch.
    ///
    /// This method receives a batch containing primary key columns as defined by
    /// `Index::index_schema()` and must return the complete records for those keys.
    /// The order of output records should correspond to the order of input primary keys.
    ///
    /// # Arguments
    /// * `index_batch` - A `RecordBatch` containing primary key values to fetch. The batch
    ///   schema matches `Index::index_schema()` and may contain one or more columns
    ///   forming a composite primary key.
    ///
    /// # Returns
    /// A `RecordBatch` containing the complete records with schema matching `schema()`.
    /// The number of output rows should equal the number of input primary keys unless some
    /// records are missing (which may indicate data consistency issues).
    ///
    /// # Error Handling
    /// Return errors for:
    /// - Storage system failures (connection errors, timeouts)
    /// - Invalid row IDs or data corruption
    /// - Schema mismatches between expected and actual data
    async fn fetch(&self, index_batch: RecordBatch) -> Result<RecordBatch>;
}